Amortized Factor Inference Networks for Posterior Inference

Résumé Technique : Réseaux d'Inférence par Facteur Amortis (AFIN)

1. Énoncé du Problème

L'inférence amortie vise à apprendre un unique réseau de neurones qui mappe les observations vers des distributions a posteriori, offrant une inférence rapide au moment du test après un coût d'entraînement unique. Cependant, les méthodes amorties existantes sont généralement liées à des modèles probabilistes fixes. Si l'a priori, la vraisemblance, la dimensionnalité latente ($d$) ou le nombre d'observations ($N$) changent, le réseau d'inférence nécessite généralement un réentraînement ou un affinage coûteux au moment du test. Bien que certaines travaux récents se généralisent à des $N$ variables ou à des changements de paramètres spécifiques, ils manquent souvent de la capacité à gérer des combinaisons arbitraires de familles d'a priori/vraisemblance variables, de dimensions et de nombres d'observations simultanément, sans sacrifier la précision ou nécessiter une adaptation spécifique au modèle.

La question centrale posée par ce travail est : Un unique réseau d'inférence peut-il se généraliser à toute une famille de problèmes d'inférence bayésienne distincts, gérant des a priori, des vraisemblances, des dimensions et des nombres d'observations variables, tout en égalant la précision des méthodes non amorties comme NUTS ?

2. Méthodologie : Réseaux d'Inférence par Facteur Amortis (AFIN)

Les auteurs introduisent les Réseaux d'Inférence par Facteur Amortis (AFIN), une architecture encodeur-fusionneur-décodeur conçue pour mapper une représentation typée d'une spécification de modèle et de ses observations directement vers les paramètres d'une postérieure variationnelle.

A. Représentation par Facteur Typé
Au lieu de traiter le modèle comme un bloc monolithique, l'AFIN représente la densité conjointe $p(z, y_{1:N})$ comme une collection de $N+1$ facteurs : un facteur a priori $p_0(z)$ et $N$ facteurs de vraisemblance $p_n(y_n|z)$. Chaque facteur est représenté par :

• Type ($t_n$) : Une étiquette indiquant la famille de distribution (par exemple, "vraisemblance gaussienne", "vraisemblance de Student-t").
• Paramètres ($\theta_n$) : Paramètres spécifiques au facteur (par exemple, échelle, degrés de liberté) nécessaires pour instancier le facteur.
• Observations ($y_n$) : Les données associées aux facteurs de vraisemblance.

Cette représentation permet au réseau de gérer des modèles hétérogènes où différentes observations peuvent suivre différentes familles de vraisemblance.

B. Architecture Indépendante de la Dimension
Pour supporter des dimensions latentes variables ($d$) et des nombres de facteurs variables ($N$), l'AFIN emploie des modules indépendants de $d$ et équivariants aux permutations des coordonnées latentes :

1. BoxMLP : Un Perceptron Multicouche à poids liés où les poids appris ne dépendent pas de $d$. Il traite des caractéristiques au niveau des coordonnées et des résumés mis en commun (par exemple, la moyenne sur les dimensions) pour garantir que la sortie soit équivariante aux permutations du vecteur latent $z \in \mathbb{R}^d$.
2. BoxTransformer : Un analogue indépendant de la dimension du Transformer standard. Il opère sur une séquence d'encodages de facteurs. Crucialement, les cartes de projection (Requête, Clé, Valeur, Sortie, FFN) sont implémentées comme des BoxMLP appliqués à des tokens indexés par coordonnées. Cela permet au réseau de modéliser les interactions entre les facteurs sans aplatir la dimension latente, préservant la structure $O(d)$ ou $O(d^2)$ des paramètres.

C. Pipeline Encodeur-Fusionneur-Décodeur

1. Encodage : Chaque facteur est traité indépendamment. Un adaptateur spécifique au type mappe les paramètres et les observations du facteur vers un espace d'encodage commun, suivi d'un encodeur BoxMLP partagé. L'encodage inclut à la fois des caractéristiques de nœud (au niveau des coordonnées) et des caractéristiques de paires (interactions de paires de coordonnées, utiles pour les matrices de covariance/précision).
2. Fusion : La séquence d'encodages de facteurs est traitée par $M$ blocs BoxTransformer. Ces blocs utilisent l'attention sur l'axe des facteurs (traitant les facteurs comme la dimension de séquence) pour agréger l'information. Le mécanisme d'attention combine les scores de nœud et de paire pour mettre à jour les représentations.
3. Décodage : Les représentations fusionnées sont mises en commun par somme sur l'axe des facteurs pour produire un unique encodage au niveau de la tâche. Un décodeur BoxMLP final mappe cet encodage vers les paramètres variationnels $\phi$ (par exemple, moyenne et précision pour une Gaussienne, ou paramètres de flot pour un flot de normalisation).

D. Objectif d'Entraînement
L'AFIN est entraîné par inférence basée sur la simulation. Le réseau est entraîné pour maximiser la densité logarithmique attribuée aux tirages latents $z$ échantillonnés à partir du modèle conjoint $p(z, y)$, conditionné aux observations générées $y$. Cela équivaut à minimiser la divergence KL forward attendue ($KL(p(z|y) || q_\phi(z))$) entre la vraie postérieure et l'approximation variationnelle. L'objectif KL forward est noté comme bénéfique pour utiliser l'approximation comme proposition pour l'échantillonnage par importance.

3. Contributions Clés

• Généralisation : Un unique AFIN entraîné peut effectuer une inférence postérieure à travers des familles d'a priori variables, des familles de vraisemblance variables, des dimensions latentes ($d$) et des nombres d'observations ($N$) sans réentraînement.
• Indépendance de la Dimension : L'introduction des modules BoxMLP et BoxTransformer permet au réseau d'opérer sur des espaces latents de dimension $d$ arbitraire tout en maintenant un nombre de paramètres fixe, une déviation significative par rapport aux réseaux amortis standards qui nécessitent des tailles d'entrée/sortie fixes.
• Abstraction au Niveau du Facteur : En opérant sur des facteurs typés plutôt que sur des graphes de modèles fixes, l'AFIN prend en charge des vraisemblances hétérogènes (par exemple, mélanger des observations Gaussiennes et Bernoulli) au sein d'une seule exécution d'inférence.
• Efficacité : La méthode atteint une haute précision avec un calcul au moment du test significativement inférieur par rapport aux méthodes itératives.

4. Résultats Expérimentaux

Les auteurs ont évalué un unique point de contrôle AFIN de 5 millions de paramètres sur des tâches synthétiques et réelles sans affinage.

• Performance Synthétique : Sur 48 configurations (16 combinaisons a priori-vraisemblance $\times$ 3 niveaux de difficulté), l'AFIN a atteint une précision postérieure comparable à NUTS et à diverses bases de référence d'Inférence Variationnelle (VI) (Gaussienne Pleine Rang, IAF, MAF, NSF).
  • Vitesse : L'inférence en un seul tirage de l'AFIN est 2 à 4 ordres de grandeur plus rapide que NUTS et les méthodes VI basées sur des flots.
  • Raffinement : Lorsqu'utilisé comme proposition pour l'Échantillonnage par Importance Auto-Normalisé (SNIS), l'AFIN+SNIS égalait ou dépassait la précision de NUTS avec un temps d'exécution nettement inférieur.
• Extrapolation : L'AFIN a démontré une généralisation zero-shot à des configurations hors distribution (OOD), incluant des dimensions latentes ($d=32$) et des nombres d'observations ($N=512$) supérieurs à ceux rencontrés pendant l'entraînement ($d \le 16, N \le 256$).
• Transfert Réel : Sur 12 jeux de données UCI (incluant des configurations de vraisemblance hétérogènes), l'AFIN+SNIS a atteint une précision compétitive par rapport aux références NUTS à long terme presque immédiatement, tandis que les bases de référence itératives nécessitaient un temps d'optimisation significatif pour converger.
• Variantes de Décodeur : Les décodeurs Gaussien et Flot ont été testés. Le décodeur Flot offrait de légères améliorations de précision par rapport au décodeur Gaussien tout en maintenant l'avantage de vitesse de l'amortissement.

5. Signification et Revendications

L'article revendique que l'AFIN comble avec succès le fossé entre la vitesse de l'inférence amortie et la flexibilité des méthodes non amorties. En découplant le réseau d'inférence des spécifications de modèles fixes grâce à une interface de facteur typé et à des modules indépendants de la dimension, l'AFIN permet une inférence amortie réutilisable.

Les auteurs positionnent l'AFIN comme une étape vers un "modèle de fondation" pour l'inférence bayésienne, capable de gérer des programmes probabilistes divers. Ils reconnaissent des limitations, telles que la nécessité d'un ensemble prédéfini de types de facteurs et la mise à l'échelle $O(d^2)$ des encodages de composants de paires, mais soulignent que l'approche permet à un unique réseau de résoudre une large famille de problèmes d'inférence qui nécessitaient auparavant des modèles distincts et réentraînés. Le travail suggère que l'amortissement peut égaler la qualité des méthodes d'échantillonnage de pointe (NUTS) tout en offrant des économies computationnelles massives au moment du test.