Résumé Technique : Problèmes Inverses avec les Modèles de Diffusion via une Perte de Recherche de Mode (Mode-Seeking Loss)

Définition du Problème

L'article traite du défi consistant à résoudre des problèmes inverses en utilisant des modèles de diffusion inconditionnels pré-entraînés. Un problème inverse consiste à estimer un échantillon de données original $x_0$ étant donné une mesure partiellement dégradée $y$, régie par l'équation $y = A(x_0) + \eta$, où $A$ est un opérateur de dégradation connu et $\eta$ est un bruit gaussien.

Bien que les modèles de diffusion pré-entraînés puissent résoudre ces problèmes de manière zero-shot (plug-and-play), les approches existantes font face à des limitations significatives :

1. Méthodes d'Échantillonnage de l'A posteriori (Posterior Sampling) : Elles visent à tirer des échantillons proportionnels à la probabilité a posteriori $p(x_0|y)$. Cependant, elles reposent souvent sur des approximations de modélisation (par exemple, des approximations gaussiennes du score conditionnel) ou des méthodes de Monte-Carlo séquentiel, ce qui peut être coûteux en termes de calcul et ne respecte pas strictement l'a posteriori réelle sans validation par vérité terrain.
2. Méthodes d'Estimation MAP : Elles visent à trouver l'estimation du Maximum A Posteriori (MAP) (le mode de l'a posteriori). Les solveurs existants utilisent souvent des approximations heuristiques, reposent sur la séparation de variables (par exemple, HQS, ADMM), ou nécessitent des modèles de cohérence. Ces approches impliquent fréquemment des calculs de gradients coûteux et des hypothèses de modélisation qui peuvent ne pas s'appliquer à tous les opérateurs linéaires.

Les auteurs soutiennent que de nombreuses tâches de restauration d'image privilégient la génération d'une image plausible cohérente avec les mesures plutôt que la couverture de l'ensemble de la distribution a posteriori. Cela s'aligne davantage sur l'estimation MAP, motivant le besoin d'une stratégie MAP plus principled et exempte d'approximations.

Méthodologie : Perte de Recherche de Mode Variationnelle (VML)

La contribution centrale est la Perte de Recherche de Mode Variationnelle (VML), une nouvelle fonction objectif conçue pour guider les échantillons générés vers l'estimation MAP pendant le processus de diffusion inverse.

Fondement Théorique

Les auteurs dérivent la VML d'une perspective variationnelle en minimisant la divergence de Kullback-Leibler (KL) entre la distribution de diffusion $p(x_0|x_t)$ et la distribution a posteriori de la mesure $p(x_0|y)$ à chaque étape de diffusion inverse $t$ :
$$\text{VML}(x_t, t) = D_{KL}(p(x_0|x_t) \parallel p(x_0|y))$$

L'intuition est qu'à mesure que $t \to 0$, la distribution $p(x_0|x_t)$ devient de plus en plus pointue autour de $x_t$. En trouvant un $x_t^*$ optimal qui minimise cette divergence KL à chaque étape, le processus guide l'échantillon vers les modes de l'a posteriori $p(x_0|y)$.

Dérivation en Forme Fermée pour les Opérateurs Linéaires

Une réussite théorique clé est la dérivation d'une expression en forme fermée pour la VML dans le cas des problèmes inverses linéaires ($A(x) = Hx$) sans aucune approximation de modélisation. L'expression complète de la VML inclut :

1. Un terme de cohérence de la mesure : $\|y - HD(x_t, t)\|^2 / (2\sigma_y^2)$.
2. Des termes de priorité (prior) : impliquant le débruiteur (denoiser) $D(x_t, t)$, le planning de bruit $\sigma_t$, et la covariance de l'a posteriori $Cov[x_0|x_t]$.

Objectifs Simplifiés et Repondérés

La VML complète contient des termes d'ordre supérieur impliquant la matrice de covariance $Cov[x_0|x_t]$, qui sont coûteux à calculer et à différencier. Les auteurs prouvent que ces termes d'ordre supérieur convergent uniformément vers zéro (ou une constante) lorsque $t \to 0$. Par conséquent, ils proposent une VML simplifiée ($VML_S$) qui supprime ces termes, ne conservant que les termes de cohérence de la mesure et de priorité.

Pour améliorer davantage la stabilité de l'optimisation, ils introduisent une VML simplifiée repondérée ($VML_S^\rho$) :
$$VML_S^\rho(x_t, t) = \frac{\|y - HD(x_t, t)\|^2}{2\sigma_y^2} - \rho_t \log p(x_t) - \rho_t \frac{\|D(x_t, t) - x_t\|^2}{2\sigma_t^2} + C$$
où $\rho_t$ est un planning de poids de priorité tel que $\rho_t \to 1$ quand $t \to 0$. Cette formulation permet un calcul de gradient efficace en utilisant la différenciation automatique standard.

Algorithmes

Sur la base de la minimisation de la VML, les auteurs proposent :

1. VML-MAP : Un algorithme utilisant la descente de gradient pour minimiser $VML_S$ à chaque étape de diffusion inverse.
2. VML-MAP$_{pre}$ : Une variante préconditionnée pour les opérateurs linéaires mal conditionnés. Elle utilise la décomposition en valeurs singulières (SVD) de la matrice de dégradation $H$ pour construire un préconditionneur $P$, accélérant la convergence pour des tâches difficiles comme le défloutage (deblurring).
3. VML-MAP$^\tau$ : Une stratégie pour la restauration d'images bruyantes. Pour éviter les artefacts nets causés par les défis d'optimisation à de faibles niveaux de bruit, l'algorithme minimise la VML uniquement jusqu'à un seuil de temps de diffusion $\tau$, puis bascule vers l'échantillonnage DDIM standard pour générer l'image finale.
4. LatentVML-MAP : Une extension aux modèles de diffusion latents (LDM), approximant la VML dans l'espace latent.

Contributions Clés

1. Nouvelle Fonction de Perte : Introduction de la VML, qui guide théoriquement les échantillons vers les estimations MAP en minimisant la divergence KL entre les a posteriori de la diffusion et de la mesure.
2. Dérivation Sans Approximation : Pour les problèmes inverses linéaires, la VML est dérivée sous forme fermée sans approximations de modélisation heuristiques, fournissant une base théorique formelle pour l'estimation MAP avec les modèles de diffusion.
3. Analyse de Convergence : Preuve formelle que les objectifs VML simplifiés et repondérés convergent vers le log-a posteriori négatif (l'objectif MAP) lorsque $t \to 0$ sous des hypothèses modérées.
4. Algorithmes Pratiques : Développement de VML-MAP et de sa variante préconditionnée, offrant une solution empiriquement efficace pour diverses tâches de restauration d'image.
5. Validation Étendue : Démonstration de performances supérieures par rapport aux bases de référence existantes (ex: DDRM, $\Pi$GDM, MAPGA, DAPS) sur plusieurs jeux de données (ImageNet, FFHQ, CelebA) et tâches (inpainting, super-résolution, défloutage).

Résultats Expérimentaux

Les auteurs ont évalué leurs méthodes sur des tâches de restauration d'image en utilisant des modèles pré-entraînés sur ImageNet256, FFHQ256 et CelebA.

• Métriques : L'article met l'accent sur les métriques perceptuelles (LPIPS et FID) plutôt que sur les métriques de pixel (PSNR, SSIM), arguant que ces dernières échouent souvent à capturer la qualité perceptuelle des reconstructions plausibles dans les problèmes inverses mal posés.
• Performance :
  • VML-MAP a systématiquement surpassé les méthodes d'échantillonnage a posteriori (DDRM, $\Pi$GDM) et les autres estimateurs MAP (MAPGA, DAPS) en termes de LPIPS et FID pour les tâches d'inpainting, de super-résolution et de défloutage.
  • VML-MAP$_{pre}$ a montré des améliorations significatives par rapport à la version non préconditionnée pour les tâches mal conditionnées (ex: défloutage), atteignant des scores LPIPS et FID plus bas.
  • Efficacité Computationnelle : VML-MAP a obtenu une meilleure qualité perceptuelle avec un temps d'exécution moindre (moins d'évaluations de fonctions neurales) par rapport aux méthodes nécessitant de nombreuses étapes de diffusion inverse.
• Restauration Bruyante : La stratégie VML-MAP$^\tau$ a géré avec succès les mesures bruyantes, produisant des reconstructions de haute qualité là où le VML-MAP standard peinait avec des artefacts.
• Modèles Latents : LatentVML-MAP a démontré son efficacité sur CelebA, bien que les auteurs notent que l'optimisation reste plus difficile dans l'espace latent en raison de la non-linéarité du décodeur.

Signification et Revendications

L'article affirme que la VML représente la première modélisation sans approximation de l'estimation MAP pour les problèmes inverses linéaires en exploitant uniquement un modèle de diffusion inconditionnel pré-entraîné, appuyée par une analyse théorique formelle.

Les auteurs soulignent que leur approche diffère fondamentalement de l'échantillonnage a posteriori (qui cible la couverture de la distribution) et des solveurs MAP existants (qui reposent souvent sur des heuristiques ou des modèles de cohérence). En minimisant la VML, la méthode cible efficacement les "modes" de l'a posteriori, ce qui, selon les auteurs, est le comportement souhaité pour générer des restaurations d'images plausibles et de haute qualité.

Le travail est modeste quant à ses limites, reconnaissant que :

• L'efficacité dépend de la disponibilité d'optimiseurs pratiquement efficaces pour l'objectif VML.
• Des défis d'optimisation persistent pour les opérateurs mal conditionnés ou non linéaires, bien que le préconditionnement aide.
• Des travaux futurs sont nécessaires pour développer des optimiseurs d'ordre supérieur afin de mieux atténuer les problèmes d'optimisation sans coûts de calcul prohibitifs.

L'article conclut que VML-MAP fournit un cadre principiel, théoriquement fondé et empiriquement efficace pour résoudre les problèmes inverses avec les modèles de diffusion, comblant ainsi le fossé entre l'estimation MAP théorique et la restauration d'image pratique.