Inférence parallèle pour un langage réactif probabiliste

Abstract

ProbZelus est un langage synchrone probabiliste qui permet de décrire des modèles probabilistes réactifs en interaction avec un environnement observable. Des méthodes d'inférences réactives permettent d'apprendre en ligne les distributions associées aux paramètres non-observés du modèle à partir d'observations statistiques. Ce problème n'a, en général, pas de solution analytique simple. Pour obtenir des estimations précises, une méthode classique consiste à analyser les résultats obtenus par de nombreuses exécutions indépendantes du modèle. Ces méthodes sont coûteuses, mais ont l'avantage de pouvoir être massivement parallélisées. Nous proposons d'utiliser JAX pour paralléliser les méthodes d'inférence réactive de ProbZelus. JAX est une bibliothèque récente qui permet de paralléliser automatiquement du code Python qui peut ensuite être exécuté de manière efficace et transparente sur CPU, GPU, ou TPU. Dans cet article, nous décrivons un nouveau moteur d'inférence réactive parallèle implémenté en JAX et le nouveau backend JAX associé pour ProbZelus. Nous montrons sur des exemples existants que notre nouvelle implémentation surpasse l'implémentation séquentielle originale pour un nombre de particules élevé.