Unbiased Diffusion Variational Inversion via Principled Posterior Matching

Énoncé du problème

Les tâches d'imagerie computationnelle consistent à retrouver un signal inconnu $x$ à partir de mesures bruitées et indirectes $y = A(x) + \eta$, où $A$ est un opérateur direct mal posé. Un défi critique dans ces problèmes inverses réside non seulement dans l'obtention d'une estimation ponctuelle de haute qualité, mais aussi dans la réalisation d'une quantification fiable de l'incertitude par le biais de la récupération de la distribution a posteriori bayésienne complète $p(x|y)$.

Les méthodes existantes basées sur la diffusion pour résoudre ces problèmes inverses relèvent généralement de trois catégories, qui souffrent toutes de limitations théoriques significatives concernant l'optimisation de la divergence de Kullback-Leibler (KL) :

1. Monte Carlo guidé par le gradient (par ex. DPS) : Ces méthodes approximent le terme de vraisemblance dépendant du temps, intraitable, conduisant souvent à des trajectoires d'échantillonnage instables et à des estimations a posteriori biaisées.
2. Inférence variationnelle basée sur l'optimisation (par ex. RED-Diff, SDS) : Ces méthodes supposent généralement une distribution variationnelle dégénérée (un delta de Dirac), éliminant ainsi efficacement le terme d'entropie dans l'objectif variationnel. Cela réduit le problème à une estimation du Maximum A Posteriori (MAP), provoquant un effondrement de mode sévère et une incapacité à capturer la structure multimodale de l'a posteriori.
3. Inférence amortie (par ex. DAVI) : Ces méthodes entraînent des réseaux feed-forward pour prédire directement les a posteriori. Cependant, elles reposent sur la divergence IKL (Integral KL), qui est une approximation heuristique plutôt qu'une borne variationnelle inférieure rigoureuse. Cette déviation, combinée à une dépendance envers des données appariées de vérité terrain pour la supervision, limite leur capacité à généraliser en zéro-shot et à retrouver l'a posteriori bayésien exact.

Méthodologie : Principled Posterior Matching (PPM)

Les auteurs proposent le Principled Posterior Matching (PPM), un cadre qui revient aux fondements de l'inférence variationnelle pour réaliser une optimisation non biaisée de la divergence KL exacte.

1. Fondement théorique : Intégration de la divergence de Fisher
Au lieu de s'appuyer sur des approximations heuristiques ou des hypothèses de distribution dégénérée, le PPM formule rigoureusement l'optimisation exacte de la divergence KL en intégrant la divergence de Fisher le long du processus de diffusion. Les auteurs exploitent un résultat théorique classique stipulant que la divergence KL entre deux distributions peut être développée comme l'intégrale de la divergence de Fisher le long de la trajectoire de diffusion.

• Idée clé : Minimiser l'intégrale de la divergence de Fisher encourage intrinsèquement un comportement de couverture de la masse, atténuant théoriquement l'effondrement de mode observé dans les méthodes précédentes.

2. Théorème d'équivalence des gradients
Pour rendre l'optimisation de la divergence de Fisher traitable, les auteurs dérivent un nouveau Théorème d'équivalence des gradients. Ce théorème fournit un estimateur de gradient exact et non biaisé pour les paramètres variationnels, sans nécessiter de rétropropagation à travers le réseau de score pré-entraîné figé.

• Le gradient est dérivé en utilisant un opérateur « stop-gradient » et un réseau de score auxiliaire, permettant une optimisation efficace par descente de gradient stochastique.

3. Cadre unifié
Le PPM unifie deux paradigmes complémentaires sous un seul objectif d'optimisation :

• Inférence variationnelle (VI) : L'a posteriori variationnel est modélisé comme un ensemble de particules d'images $\{\mu_k\}$. Cela permet une récupération d'a posteriori basée sur des particules, haute fidélité, pour des observations spécifiques.
• Inférence amortie (AI) : L'a posteriori variationnel est modélisé comme un réseau de reconstruction $x = g_\phi(z)$. Cela permet une inférence rapide en une seule étape en apprenant une cartographie globale des observations vers des échantillons d'a posteriori.
• Architecture : Le cadre utilise un réseau de score de diffusion pré-entraîné $s_\theta$ comme a priori et un réseau de score auxiliaire $s_\phi$ (adapté via une adaptation de faible rang) pour approximer le score de la distribution variationnelle $q_\phi$. L'optimisation alterne entre la mise à jour du réseau de score auxiliaire et la mise à jour des paramètres variationnels (particules ou poids du réseau).

4. Généralisabilité
La formulation s'étend naturellement à une famille plus large de mesures de divergence en généralisant l'intégrale de la divergence de Fisher, positionnant le PPM comme une fondation polyvalente pour l'inférence variationnelle basée sur les scores.

Contributions clés

• Optimisation non biaisée : Le PPM est le premier cadre à optimiser rigoureusement la divergence KL exacte pour les problèmes inverses basés sur la diffusion, sans recourir aux approximations biaisées (par ex. hypothèses de delta de Dirac ou heuristiques IKL) qui affectent les méthodes existantes.
• Unification théorique : L'article démontre que la minimisation de la KL exacte via l'intégration de Fisher unifie les méthodes existantes basées sur les scores comme cas particuliers et fournit une explication théorique de la raison pour laquelle les méthodes précédentes souffrent d'effondrement de mode (le « gap d'approximation »).
• Support dual-paradigme : Il fournit une logique d'entraînement unifiée qui supporte à la fois la VI basée sur des particules pour une récupération d'a posteriori haute fidélité et l'inférence amortie pour une génération efficace et non supervisée.
• Quantification robuste de l'incertitude : En évitant l'effondrement de mode, le PPM produit des estimations d'incertitude bien calibrées qui reflètent fidèlement la nature multimodale de l'a posteriori.

Résultats expérimentaux

Les auteurs valident le PPM sur une suite complète de benchmarks d'imagerie computationnelle, incluant des images naturelles et des tâches d'imagerie scientifique.

1. Estimation d'a posteriori 2D (Exemples simples)

• Sur des priors de mélanges gaussiens 2D, le PPM récupère avec précision tous les modes de l'a posteriori. En revanche, RED-Diff présente un effondrement de mode sévère (convergence vers un point unique), et RLSD montre une incohérence de mesure ou une mauvaise adhésion au prior.

2. Photographie computationnelle (FFHQ & ImageNet)

• Tâches : Inpainting d'images, défloutage de mouvement, super-résolution et débruitage gaussien.
• Performance : Le PPM surpasse les références de l'état de l'art (DPS, RED-Diff, RLSD, DAVI) en termes de fidélité de reconstruction (PSNR, SSIM) et de diversité des échantillons.
• Qualitatif : Dans les tâches d'inpainting, le PPM génère des complétions diverses et plausibles (par ex. différentes coiffures), tandis que les références produisent des sorties presque identiques.
• Incertitude : Les cartes de déviation standard du PPM reflètent avec précision l'incertitude positionnelle (par ex. aux bords des objets), tandis que les cartes de RED-Diff indiquent des modes effondrés et celles de RLSD contiennent des artefacts.

3. Applications en imagerie scientifique

• Microscopie à fluorescence en super-résolution : Appliqué aux structures subcellulaires (microtubules, réticulum endoplasmique). Le PPM atteint une fidélité supérieure à RED-Diff et produit des cartes d'incertitude alignées avec la physique de l'imagerie (incertitude plus faible au centre des structures, plus élevée aux bords flous).
• Imagerie par interférométrie radio de trous noirs : Appliqué à des données synthétiques du Event Horizon Telescope (EHT). Le PPM récupère fidèlement la structure en anneau et le croissant brillant de l'ombre du trou noir. Contrairement à RED-Diff (qui s'effondre) et DPS (qui est flou), le PPM fournit une distribution d'incertitude physiquement significative capturant la variance structurelle.

Signification et affirmations

L'article affirme que le PPM établit un cadre robuste pour l'imagerie scientifique en s'attaquant aux limitations théoriques fondamentales des solveurs inverses basés sur la diffusion actuels.

• Fiabilité : En éliminant les biais d'approximation qui conduisent à l'effondrement de mode, le PPM permet une récupération fidèle d'a posteriori multimodaux et des estimations d'incertitude bien calibrées, ce qui est crucial pour les applications scientifiques et médicales où la compréhension de l'espace des solutions complet est aussi importante que la reconstruction elle-même.
• Généralisation en zéro-shot : Contrairement aux méthodes amorties nécessitant des données d'entraînement appariées, le PPM opère de manière non supervisée, le rendant applicable aux domaines scientifiques où les données de vérité terrain sont indisponibles.
• Fondement théorique : Ce travail fournit un pont théorique rigoureux entre la divergence de Fisher et la minimisation de la KL, offrant une alternative fondée sur des principes aux approximations heuristiques qui ont dominé le domaine.

En résumé, le PPM démontre que le retour à l'optimisation exacte de l'objectif variationnel, facilité par l'intégration de la divergence de Fisher, produit une qualité de reconstruction, une diversité et une quantification de l'incertitude supérieures par rapport aux approximations biaisées existantes.