Learned harmonic mean estimation of the marginal likelihood for multimodal posteriors with flow matching

Résumé Technique : Estimation de la moyenne harmonique apprise avec le Flow Matching pour les postulats multimodaux

Énoncé du problème
La vraisemblance marginale (ou évidence bayésienne), notée $z = p(d|M)$, est une quantité fondamentale pour la comparaison de modèles bayésiens, permettant une évaluation quantitative via le facteur de Bayes. Cependant, le calcul de cette quantité nécessite l'évaluation d'une intégrale multidimensionnelle sur l'ensemble de l'espace des paramètres, ce qui est informatiquement prohibitif pour des modèles complexes et de haute dimension. Bien que les méthodes de Monte Carlo par chaînes de Markov (MCMC) et de Monte Carlo Hamiltonien (HMC) excellent dans la génération d'échantillons de l'a posteriori en haute dimension, elles ne fournissent pas intrinsèquement la vraisemblance marginale. Inversement, l'échantillonnage imbriqué (nested sampling) fournit l'évidence, mais restreint la stratégie d'échantillonnage à une hiérarchie imbriquée spécifique.

L'« estimateur de moyenne harmonique appris » a été introduit précédemment pour estimer la vraisemblance marginale en utilisant uniquement des échantillons de l'a posteriori, quel que soit la méthode d'échantillonnage utilisée pour les obtenir. Cette approche repose sur l'apprentissage d'une densité cible interne $\phi(\theta)$ qui est contenue dans l'a posteriori mais possède des queues plus fines. Cependant, les implémentations précédentes utilisant des flux de normalisation discrets ou des estimateurs de densité sur mesure ont éprouvé des difficultés face à des a posteriori hautement multimodaux. Ces architectures échouaient souvent à capturer des topologies complexes sans modifications heuristiques (par exemple, la construction manuelle de mélanges gausiens pour la distribution de base) ou souffraient de l'attribution d'une masse de probabilité non négligeable aux ponts de densité entre les modes, entraînant une variance infinie de l'estimateur.

Méthodologie
Ce travail traite des limites des estimateurs de densité précédents en intégrant le Flow Matching dans le cadre de la moyenne harmonique apprise. La méthodologie procède comme suit :

1. Estimateur de moyenne harmonique appris : L'estimateur approxime la vraisemblance marginale réciproque $\rho = z^{-1}$ en utilisant des échantillons de l'a posteriori $\theta_i \sim p(\theta|d)$ et une densité cible apprise $\phi(\theta)$ :
   $$\hat{\rho} = \frac{1}{N} \sum_{i=1}^{N} \frac{\phi(\theta_i)}{L(\theta_i)\pi(\theta_i)}$$
   Pour que l'estimateur ait une variance finie, $\phi(\theta)$ doit être concentrée dans les régions de haute probabilité de l'a posteriori et décroître plus rapidement dans les queues de distribution.

2. Architecture de Flow Matching : Au lieu des flux de normalisation discrets, les auteurs emploient des Flux de Normalisation Continus (CNF) entraînés par flow matching.

   • Définition : Un CNF paramètre la transformation entre une distribution de base simple (gaussienne standard) et l'a posteriori cible via un champ de vitesse dépendant du temps $v_t(\theta)$.
   • Objectif d'entraînement : Le réseau est entraîné à l'aide d'un objectif de régression sans simulation. Un chemin de probabilité conditionnelle est défini comme une ligne droite entre un échantillon de bruit $\theta_0$ et un échantillon de l'a posteriori $\theta_1$ : $\theta_t = (1-t)\theta_0 + t\theta_1$. Le réseau minimise l'erreur quadratique moyenne entre le champ de vitesse prédit et le véritable champ vectoriel conditionnel $(\theta_1 - \theta_0)$.
   • Avantages : Cette approche permet d'utiliser des réseaux neuronaux à haute capacité et sans restrictions architecturales (Perceptrons Multicouches avec activation Swish) et évite les simulations d'EDO coûteuses requises dans les méthodes de CNF antérieures.

3. Mise à l'échelle de la température : Pour garantir que la densité cible $\phi(\theta)$ possède des queues plus fines que l'a posteriori, un paramètre de température $T \in (0, 1)$ est appliqué pour mettre à l'échelle la variance de la distribution de base du flux. Cela concentre la distribution du flux, satisfaisant ainsi la condition de confinement requise pour l'estimateur de moyenne harmonique.

Contributions clés

• Nouvelle architecture : L'introduction des flux de normalisation continus entraînés par flow matching comme estimateur de densité interne pour la méthode de la moyenne harmonique apprise.
• Gestion de la multimodalité : Démontrer que cette architecture peut capturer avec précision des structures d'a posteriori complexes et hautement multimodales sans nécessiter de modifications heuristiques de la distribution de base ou de réglages fins.
• Scalabilité : Fournir une méthode agnostique à la stratégie d'échantillonnage (compatible avec MCMC, HMC, NUTS, SBI, etc.) et capable d'opérer dans des espaces de dimension modérée à élevée.

Résultats expérimentaux
Les auteurs ont validé l'approche sur deux tests de référence exigeants :

1. Fonction de Rastrigin (2 dimensions) :

   • Un test standard présentant de nombreux modes locaux.
   • Le modèle de flow matching (6 couches, 256 unités cachées) a capturé avec succès tous les modes avec une température $T=0.98$.
   • Les résultats de 100 séries expérimentales ont montré que les valeurs de vraisemblance marginale estimées étaient précises et correspondaient étroitement à la vérité terrain obtenue par intégration numérique. Les estimations d'erreur étaient fiables.

2. Mélange Gaussien en 20 dimensions :

   • Un mélange de cinq composantes gaussiennes en 20 dimensions avec des modes non chevauchants et des corrélations introduites.
   • Le modèle (10 couches, 256 unités cachées) a été entraîné sur 40 000 échantillons.
   • La distribution cible apprise a montré qu'elle était correctement concentrée dans l'a posteriori, capturant la topologie complexe.
   • Les vraisemblances marginales estimées sur 100 séries s'alignaient avec la vérité terrain analytique, confirmant l'exactitude de la méthode et la robustesse de l'estimation de l'erreur, même dans ce scénario complexe de dimension modérée.

Signification et affirmations
L'article affirme qu'en exploitant l'expressivité des CNF entraînés par flow matching, l'estimateur de moyenne harmonique appris peut désormais estimer de manière robuste l'évidence pour des a posteriori hautement complexes et multimodaux. La principale signification réside dans la suppression de la dépendance aux ajustements heuristiques (tels que les mélanges gausiens choisis manuellement) qui étaient auparavant nécessaires pour gérer la multimodalité. Les auteurs soutiennent que cette approche permet un calcul précis de l'évidence en utilisant des échantillons provenant de n'importe quelle méthode d'échantillonnage accélérée ou scalable, facilitant ainsi la comparaison de modèles bayésiens dans des domaines riches en données comme l'astrophysique, sans les goulots d'étranglement computationnels de l'échantillonnage imbriqué traditionnel ou l'instabilité de l'original estimateur de moyenne harmonique.

Les travaux futurs, comme l'indiquent les auteurs, incluent l'application de cette méthodologie à des problèmes réels tels que l'astrophysique des ondes gravitationnelles et l'exploration de nouveaux raffinements architecturaux pour des tâches de très haute dimension comme l'imagerie et l'inférence cosmologique au niveau des champs.