Résumé technique : Échantillonnage MAP local pour les modèles de diffusion

Énoncé du problème

L'échantillonnage postérieur de diffusion (DPS) s'est imposé comme un cadre puissant pour résoudre des problèmes inverses en échantillonnant à partir de la distribution postérieure $p(x_0 | y)$, offrant une quantification rigoureuse de l'incertitude et une multimodalité. Cependant, de nombreuses applications pratiques en imagerie et en sciences privilégient une estimation ponctuelle précise plutôt qu'un échantillonnage complet de la distribution. L'estimateur du Maximum A Posteriori (MAP), défini comme $\arg \max p(x_0 | y)$, demeure un objectif de reconstruction standard.

Les méthodes existantes basées sur l'optimisation (par exemple, Resample, DiffPIR, DCDP) tentent d'améliorer les performances de reconstruction en alternant des étapes de débruitage et d'optimisation. Toutefois, le lien théorique entre ces procédures itératives, la solution MAP globale et l'échantillonnage postérieur reste flou. Il n'est pas établi si ces méthodes convergent vers le MAP global ou si elles sont cohérentes avec le DPS. De plus, la nature spécifique des étapes d'optimisation dans ces méthodes — qu'elles résolvent un objectif global ou local — nécessite d'être clarifiée.

Méthodologie : Échantillonnage MAP local (LMAPS)

Les auteurs introduisent l'Échantillonnage MAP local (LMAPS), un nouveau cadre d'inférence qui résout itérativement des sous-problèmes de maximum a posteriori local le long de la trajectoire de diffusion.

1. Cadre théorique

LMAPS fonctionne en résolvant une séquence de problèmes MAP locaux plutôt qu'un problème global unique. À chaque étape de diffusion $t$, étant donné un état bruité $x_t$ et une mesure $y$, la méthode cherche :
$$x^*_0(t, x_t, y) := \arg \max_{x_0} p(x_0 | x_t, y)$$
Ceci est suivi d'une étape de transition DDIM standard :
$$x_{t-\Delta t} = g(x^*_0, x_t, \epsilon)$$

Distinctions par rapport aux méthodes existantes :

• Vs. MAP global : Le MAP global cherche le mode de la postérieure complète $p(x_0 | y)$. LMAPS cherche le mode de la postérieure conditionnelle $p(x_0 | x_t, y)$ à chaque étape. Les auteurs soutiennent que la séquence de maximiseurs locaux ne converge pas nécessairement vers le maximiseur global de $p(x_0 | y)$.
• Vs. DPS : Le DPS utilise la moyenne conditionnelle $E[x_0 | x_t, y]$ dans la mise à jour. LMAPS remplace cette moyenne par le mode conditionnel. Ces deux approches sont équivalentes uniquement si la postérieure conditionnelle est gaussienne ; sinon, LMAPS introduit un biais de recherche de mode, évitant les régions de faible densité entre les modes.
• Vs. DMAP : Bien que DMAP utilise également une perspective MAP locale, il optimise dans l'espace des variables de diffusion latentes ($x^*_{t-\Delta t} = \arg \max p(x_{t-\Delta t} | x_t, y)$). LMAPS pose le problème directement dans l'espace du signal propre ($x^*_0$), offrant une vue complémentaire qui connecte explicitement à la vraisemblance de la mesure $p(y|x_0)$.

2. Algorithme pratique et approximations

Pour rendre le problème MAP local traitable pour des problèmes inverses non linéaires généraux, les auteurs proposent des approximations et des reformulations spécifiques :

• Approximation de la covariance : Les auteurs approximent la covariance conditionnelle $\Sigma_{0|t}$ de la prior $p(x_0 | x_t)$ en utilisant une forme isotrope motivée par une hypothèse de prior gaussien. À mesure que le niveau de bruit $\sigma_t \to 0$, la covariance postérieure devient asymptotiquement isotrope. Ils approximent $\Sigma_{0|t} \approx \frac{1}{SNR} I$, où $SNR = \alpha_t^2 / \sigma_t^2$.

• Reformulation de l'objectif : L'objectif MAP local est formulé comme la minimisation d'une somme pondérée de la distance au prior et de la fidélité aux données :
  $$x^*_0 = \arg \min_{x_0} \left\{ \frac{SNR}{k} \|x_0 - m_{0|t}\|^2 + \frac{1}{\sigma_y^2} \|y - H(x_0)\|^2 \right\}$$
  Pour améliorer la stabilité numérique et l'interprétabilité, ceci est reformulé en une combinaison convexe invariante d'échelle :
  $$x^*_0 = \arg \min_{x_0} \left\{ (1 - \mu_t) \frac{1}{2} \|x_0 - m_{0|t}\|^2 + \mu_t k_2 \|y - H(x_0)\|^2 \right\}$$
  où $\mu_t = \sigma_t^2 / (\sigma_t^2 + k_1^2)$. Cette formulation permet un recuit automatique : à mesure que $t$ diminue ($\sigma_t \to 0$), l'objectif passe d'une orientation guidée par la mesure à une orientation guidée par le prior.

• Linéaire vs Non linéaire : Pour les problèmes inverses linéaires, la solution MAP locale peut être dérivée sous forme fermée, retrouvant la diffusion projetée par le moment de Tweedie (TMPD) comme cas particulier. Pour les problèmes non linéaires, la solution est obtenue par descente de gradient.

Contributions clés

1. Unification théorique : L'article formule LMAPS pour unifier les méthodes de résolution de problèmes inverses basées sur l'optimisation et la diffusion projetée par le moment de Tweedie (TMPD) dans un cadre probabiliste unique. Il clarifie que les méthodes basées sur l'optimisation résolvent intrinsèquement des problèmes MAP locaux plutôt que des MAP globaux ou un échantillonnage complet de la postérieure.
2. Avancées méthodologiques : Les auteurs introduisent une approximation de covariance motivée par la gaussienne et une fonction objectif reformulée qui améliore la stabilité numérique et l'interprétabilité grâce à un recuit automatique.
3. Validation empirique : Le cadre est validé sur 10 tâches de restauration d'images (incluant des opérateurs linéaires, non linéaires et non différentiables) et 3 problèmes inverses scientifiques.

Résultats expérimentaux

Les auteurs ont évalué LMAPS sur les ensembles de données FFHQ et ImageNet pour la restauration d'images, et sur des benchmarks pour des problèmes inverses scientifiques (Diffusion Inverse Linéaire, IRM par Compression Sensée, et Imagerie de Trou Noir).

• Restauration d'images : LMAPS a atteint des performances de pointe dans 43 cas sur 60 selon les métriques PSNR, SSIM et LPIPS sur FFHQ et ImageNet. Il a systématiquement surpassé la base de référence DMAP, montrant des améliorations de plus de 2 dB en PSNR dans des tâches telles que le défloutage de mouvement, la restauration JPEG et la quantification.
• Problèmes inverses scientifiques : LMAPS a constamment obtenu le PSNR le plus élevé sur toutes les tâches scientifiques. Notamment, il a réalisé des gains supérieurs à 1,5 dB en PSNR sur trois tâches de Diffusion Inverse Linéaire (LIS) par rapport aux méthodes précédentes.
• Efficacité : LMAPS a démontré une efficacité supérieure par rapport à DAPS, obtenant de meilleurs résultats avec moins de ressources computationnelles (par exemple, moins d'étapes d'optimisation par étape de diffusion).

Signification et affirmations

L'article affirme que LMAPS fournit une interprétation probabiliste rigoureuse pour les méthodes basées sur l'optimisation, comblant le fossé entre l'échantillonnage postérieur (DPS) et l'estimation ponctuelle (MAP). En résolvant explicitement des sous-problèmes MAP locaux, la méthode offre une vue unifiée qui explique le comportement des solveurs basés sur l'optimation existants.

Les auteurs soulignent que, bien que LMAPS n'échantillonne pas à partir de la distribution postérieure complète (et ne capture donc pas la multimodalité de la même manière que le DPS), il est spécifiquement conçu pour des scénarios où l'estimation ponctuelle précise est l'objectif principal. Le cadre offre une alternative robuste qui équilibre la fidélité de la reconstruction avec les priors structurels des modèles de diffusion, obtenant des résultats de pointe tant dans la restauration d'images standard que dans des tâches de reconstruction scientifique complexes. Ce travail suggère que la recherche future pourrait explorer davantage l'utilité de l'estimation MAP globale avec des priors de diffusion, ce qui reste un défi ouvert.