Stage-wise Distortion-Perception Traversal in Zero-shot Inverse Problems with Diffusion Models

Résumé technique : Parcours de distorsion-perception par étapes dans les problèmes inverses sans échantillonnage avec des modèles de diffusion

Énoncé du problème
Les problèmes inverses bayésiens, qui visent à reconstruire des signaux originaux à partir d'observations dégradées, sont fondamentalement contraints par le compromis distorsion-perception (D–P). La performance de distorsion mesure la fidélité de la reconstruction au niveau de l'échantillon (par exemple, via l'erreur quadratique moyenne ou le rapport signal sur bruit de crête), tandis que la performance de perception reflète la qualité visuelle et la similarité distributionnelle par rapport aux données réelles. Des travaux théoriques ont établi que ces objectifs sont intrinsèquement en tension : améliorer l'un nécessite souvent de dégrader l'autre. Bien que les modèles de diffusion aient obtenu des résultats à l'état de l'art dans la résolution de problèmes inverses sans échantillonnage sans réentraînement spécifique à la tâche, les méthodes existantes manquent de stratégies fondées sur des principes et efficaces sur le plan computationnel pour parcourir la courbe D–P au moment de l'inférence. Les approches actuelles, telles que l'ajustement des étapes d'échantillonnage ou la moyenne d'échantillons, sont souvent heuristiques ou coûteuses en calcul, et les mécanismes permettant de contrôler ce compromis dans les algorithmes sans échantillonnage basés sur la diffusion restent insuffisamment caractérisés.

Méthodologie
Les auteurs proposent un cadre en deux étapes, nommé MAP-RPS (Maximum A Posteriori - Re-noised Posterior Sampling), pour permettre un parcours flexible D–P en utilisant un seul modèle de diffusion préentraîné. La méthode est ensuite étendue à l'espace latent sous le nom de LMAP-RPS pour exploiter des architectures de diffusion latente à grande échelle.

1. Étape 1 : Initialisation à faible distorsion (Estimation MAP)
   La première étape approxime la solution d'erreur quadratique moyenne minimale (MMSE), qui correspond au point de distorsion optimal sur la courbe D–P. Les auteurs analysent des distributions a posteriori fortement log-concaves et établissent que l'estimateur MAP sert de substitut valide à la solution MMSE avec une erreur d'approximation bornée. Ils dérivent un gradient traitable pour le log-prior en utilisant la fonction de score du modèle de diffusion préentraîné à un instant spécifique $t_1$, permettant une estimation MAP efficace via une optimisation par gradient stochastique. Cela produit un point de départ à faible distorsion ($x_{MAP}$).

2. Étape 2 : Amélioration de la qualité perceptive (Échantillonnage de l'a posteriori rebruité)
   La deuxième étape parcourt la courbe D–P en améliorant progressivement la qualité perceptive. L'estimation MAP est rebruitée jusqu'à un instant prescrit $t_0$, et un échantillonnage de l'a posteriori est effectué à partir de cette initialisation bruitée. Théoriquement, les auteurs montrent que la distance de Wasserstein-2 ($W_2$) (une métrique pour l'erreur de perception) entre la distribution générée et la distribution des données réelles est contrôlée par l'instant de rebruitage $t_0$. En faisant varier $t_0$, la méthode permet une interpolation continue entre la solution MAP à faible distorsion (où $t_0=0$) et l'échantillonnage standard de l'a posteriori (où $t_0=T$).

Contributions clés

• Cadre théorique : L'article fournit des analyses théoriques pour les deux étapes. Il établit une borne d'erreur pour l'approximation de la solution MMSE par l'estimateur MAP sous des hypothèses de forte log-concavité et dérive une borne supérieure sur la distance $W_2$ pour l'étape d'échantillonnage de l'a posteriori rebruité, prouvant que l'erreur de perception peut être contrôlée en ajustant $t_0$.
• Conception de l'algorithme : L'introduction de MAP-RPS et de sa variante dans l'espace latent, LMAP-RPS, offre un mécanisme efficace sur le plan computationnel pour le parcours D–P. Contrairement aux méthodes nécessitant un échantillonnage répété ou un réglage complexe des hyperparamètres, cette approche ne requiert qu'une procédure d'estimation MAP et un processus d'échantillonnage de l'a posteriori.
• Extension à l'espace latent : En étendant le cadre aux modèles de diffusion latente (LMAP-RPS), les auteurs exploitent des architectures préentraînées puissantes (par exemple, Stable Diffusion) pour une applicabilité plus large à des tâches réelles à haute résolution.

Résultats expérimentaux
Des expériences extensives ont été menées sur le jeu de données FFHQ (pour l'analyse de la courbe D–P) et le jeu de données MS-COCO (pour des problèmes inverses proches du monde réel).

• Parcours D–P : Sur les tâches FFHQ (débruitage, super-résolution, inpainting, défloutage), MAP-RPS et LMAP-RPS ont tous deux obtenu de manière cohérente des courbes plus proches du coin inférieur gauche idéal (distorsion et perception optimales) par rapport aux méthodes existantes telles que VSDPS et PSCGAN.
• Performance en conditions réelles : Sur MS-COCO, LMAP-RPS a obtenu des performances à l'état de l'art sur presque toutes les tâches et métriques (PSNR, SSIM, LPIPS, FID) par rapport à neuf récentes références basées sur la diffusion latente.
• Efficacité : LMAP-RPS a démontré des avantages computationnels significatifs. À $t_0=0$, il était nettement plus rapide que toutes les références. Même à $t_0=600$, il restait efficace sur le plan computationnel, surpassant souvent d'autres méthodes en vitesse tout en maintenant une qualité de reconstruction supérieure.

Signification et affirmations
L'article affirme que MAP-RPS et LMAP-RPS offrent une perspective nouvelle pour réaliser le parcours D–P dans les problèmes inverses sans échantillonnage. Les auteurs soulignent que leur méthode offre une alternative fondée sur des principes, théoriquement étayée et efficace sur le plan computationnel aux stratégies heuristiques existantes. En découplant les objectifs de distorsion et de perception en deux étapes distinctes, le cadre permet aux praticiens de sélectionner flexiblement une reconstruction qui correspond le mieux aux exigences spécifiques de l'application sans réentraîner les modèles. Les auteurs soulignent également le fort potentiel pratique de LMAP-RPS, notant sa capacité à fournir des résultats à l'état de l'art sur des problèmes inverses complexes et réels avec une complexité computationnelle inférieure à celle des algorithmes actuels à la pointe de la technologie.

Limites et travaux futurs
Les auteurs reconnaissent que leurs garanties théoriques reposent sur l'hypothèse de distributions a posteriori fortement log-concaves (approximativement unimodales). Bien que les résultats empiriques suggèrent que la méthode reste efficace dans des contextes plus généraux, le comportement sous des distributions a posteriori multimodales complexes nécessite une investigation plus approfondie. De plus, l'implémentation actuelle repose sur des algorithmes spécifiques d'échantillonnage de l'a posteriori (comme DPS) pour l'étape 2, ce qui peut introduire des erreurs d'estimation ; des travaux futurs pourraient intégrer des méthodes d'échantillonnage de l'a posteriori asymptotiquement cohérentes pour affiner davantage le parcours D–P.