JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference

Résumé Technique : JADAI – Amortissement Conjoint de la Conception Adaptative et de l'Inférence Bayésienne

Énoncé du Problème
De nombreux domaines scientifiques et d'ingénierie impliquent l'inférence de paramètres non observés $\theta$ à partir d'observations $x$ générées par des simulateurs complexes. Bien que l'Inférence Basée sur la Simulation (SBI) ait fait progresser la capacité d'approximer les distributions postérieures $p(\theta|x)$ sans évaluation explicite de la vraisemblance, ces méthodes supposent généralement une conception expérimentale fixe. Inversement, la Conception Expérimentale Bayésienne (BED) et la Conception Adaptative (BAD) se concentrent sur la sélection de variables de conception $\xi$ pour maximiser le gain d'information (par exemple, le Gain d'Information Attendu, EIG), mais traitent souvent l'inférence comme une étape distincte et coûteuse en calcul, effectuée après la collecte des données. Les approches existantes déconnectent fréquemment ces tâches : les méthodes BAD reposent souvent sur une inférence non amortie et lente ou sur de simples familles paramétriques, tandis que les méthodes SBI ne cherchent pas à optimiser l'acquisition de données. Cette séparation conduit à un gain d'information sous-optimal et à une inefficacité computationnelle, particulièrement dans des contextes de haute dimension ou multimodaux.

Méthodologie
Les auteurs introduisent JADAI (Jointly Amortizing Adaptive Design and Bayesian Inference), un cadre qui unifie la conception adaptative et l'inférence postérieure en un système unique, entraînable de bout en bout. L'innovation centrale est l'amortissement conjoint de trois composants :

1. Un Réseau de Politique ($\pi_\phi$) : propose la conception suivante $\xi_t$ basée sur l'historique expérimental actuel.
2. Un Réseau d'Historique/Résumé ($\eta_\omega$) : encode la séquence des paires conception-observation passées $\{(\xi_k, x_k)\}_{k=1}^t$ en une statistique de résumé de dimension fixe $h_t$.
3. Un Réseau d'Inférence ($q_\psi$) : approxime la distribution postérieure $p(\theta|h_t)$ à chaque étape de la séquence expérimentale.

Objectif d'Optimisation
JADAI minimise une fonction de perte générique qui agrège les réductions incrémentielles de l'erreur postérieure le long des séquences expérimentales. Au lieu de maximiser directement le Gain d'Information Attendu (EIG) — ce qui nécessite des évaluations de densité — le cadre utilise une somme télescopique des pertes postérieures.

• L'objectif d'entraînement est dérivé d'une borne variationnelle inférieure sur le Gain d'Information Total Attendu (TEIG).
• La perte à l'étape $t$, $\ell_t$, est définie par l'estimateur de la distribution postérieure (par exemple, la perte de diffusion ou de flow matching).
• La fonction d'utilité $u_T$ est construite comme une somme des différences entre les pertes consécutives : $u_T = \sum (\ell^*_{t-1} - \ell_t)$, où $\ell^*$ désigne une version détachée (avec arrêt du gradient) de la perte précédente.
• Cette formulation garantit que le réseau de politique et le réseau de résumé sont entraînés pour maximiser la réduction de la perte postérieure à chaque étape, maximisant ainsi l'apport d'information sans nécessiter d'évaluations explicites de la densité de la distribution postérieure.

Choix Architecturaux Clés

• Modèles de Diffusion : Le réseau d'inférence $q_\psi$ est instancié à l'aide de modèles de diffusion (ou de flow matching), permettant au cadre d'approximer des distributions postérieures de haute dimension et multimodales à chaque étape expérimentale, une capacité souvent absente des méthodes BAD précédentes limitées aux flux de normalisation (normalizing flows).
• Stabilité de l'Entraînement : Pour éviter la rétropropagation par le temps (BPTT) imbriquée et assurer un entraînement stable, le cadre emploie :
  • États d'Historique Détachés : Les conceptions sont générées à partir d'un état de résumé détaché pour empêcher les chemins de gradient cycliques entre le réseau de politique et le réseau d'historique.
  • Apprentissage par Curriculum : La longueur de la séquence (rollout) est augmentée progressivement durant l'entraînement, et une stratégie de "mélange de priorité de conception" (design prior mix-in) est utilisée au début de l'entraînement pour exposer les réseaux à des conceptions diverses avant de dépendre uniquement de la politique apprise.

Résultats
JADAI a été évalué sur trois tests représentatifs d'une complexité croissante de la politique et de la structure de la distribution postérieure :

1. Localisation de Source (LF) : Une tâche consistant à inférer des localisations de sources à partir de mesures d'intensité bruitées. JADAI a obtenu des scores sPCE (Estimation Contrastive de Prior Séquentielle) supérieurs ou compétitifs par rapport aux meilleures méthodes de référence (ex: DAD, iDAD, RL-BOED), particulièrement pour des horizons longs ($T > 10$). La politique apprise a efficacement équilibré l'exploration et l'exploitation, et la distribution postérieure basée sur la diffusion a correctement capturé la nature multimodale et échangeable des sources.
2. Élasticité de Substitution Constante (CES) : Un problème de conception difficile impliquant l'inférence de préférences où les conceptions informatives se situent dans un intervalle très étroit ("sweet spot"). JADCI a surpassé toutes les méthodes de référence existantes à l'horizon standard $T=10$, démontant une robustesse dans des scénarios où les conceptions aléatoires ne sont pas informatives.
3. Découverte d'Images MNIST : Une tâche de haute dimension où l'objectif est de reconstruire une image à partir de mesures partielles. JADAI a atteint des scores d'indice de similitude structurelle (SSIM) et d'erreur quadratique moyenne normalisée (NRMSE) de pointe, surpassant la base de référence CoDiff. La méthode a réussi à identifier les classes de chiffres précocement et à affiner les détails locaux, maintenant ses performances même sous un bruit de mesure additif.

Signification et Revendications
L'article affirme que JADAI représente une étape significative en unifiant la conception et l'inférence, éliminant ainsi le besoin d'étapes d'inférence distinctes et coûteuses durant le processus adaptatif. Les principales contributions incluent :

• Amortissement Conjoint : C'est le premier cadre qui amortit conjointement la conception adaptative et l'inférence postérieure en utilisant une utilité générale qui agrège les réductions incrémentielles de l'erreur postérieure, évitant ainsi le besoin d'évaluations de densité.
• Scalabilité aux Postérieurs Complexes : En exploitant les modèles de diffusion, JADAI peut gérer des distributions postérieures de haute dimension et multimodales à chaque étape expérimentale, dépassant les limitations d'expressivité des méthodes précédentes reposant sur les flux de normalisation ou les mélanges gaussiens.
• Performance : Le cadre atteint des performances supérieures ou compétitives à travers divers tests, incluant la reconstruction d'images de haute dimension, tout en maintenant une efficacité de calcul lors du déploiement (les séquences sont complétées en quelques millisecondes).

Les auteurs notent que, bien que JADAI dépende actuellement de simulateurs différentiables, l'approche est parfaitement adaptée au paysage croissant des cadres de simulation différentiable. Ils reconnaissent que l'échantillonnage de la distribution postérieure reste un goulot d'étranglement pour un déploiement en temps réel (sous la seconde) dans les cas de haute dimension, mais suggèrent la distillation post-hoc comme une direction future potentielle. Ce travail ne prétend pas résoudre l'optimisation de simulateurs "boîte noire" sans gradients, mais se positionne comme une solution robuste pour les contextes différentiables où les pipelines classiques sont trop coûteux.