Résumé Technique : Align & Invert

Énoncé du Problème

Les problèmes inverses visent à retrouver un signal propre $x_0$ à partir de mesures bruitées ou dégradées $y$. Bien que les modèles de diffusion et basés sur le flux préentraînés soient devenus des priors puissants pour ces tâches, ils éprouvent souvent des difficultés à produire des images hautement détaillées dans des scénarios impliquant une dégradation sévère (par exemple, une super-résolution extrême, un flou de mouvement) ou des scènes complexes. Une limitation spécifique apparaît avec les Modèles de Diffusion Latents (LDM), où les non-linéarités de l'architecture encodeur-décodeur peuvent introduire des artefacts lors du processus d'échantillonnage inverse. De plus, les méthodes existantes nécessitent souvent un grand nombre d'étapes de discrétisation pour obtenir des reconstructions de haute qualité, et elles manquent de contraintes sémantiques fortes lors de l'inférence pour guider la trajectoire vers des solutions perceptuellement réalistes.

Méthodologie

L'article propose Align & Invert, un cadre qui intègre l'Alignement de Représentation (REPA) dans le processus d'inférence des solveurs de diffusion et basés sur le flux pour les problèmes inverses.

Concept Central

La méthode impose un alignement entre les représentations internes du modèle génératif (diffusion/flux) et les caractéristiques d'un encodeur visuel préautorisé et auto-supervisé (spécifiquement DINOv2). Alors que le REPA était précédemment utilisé comme régularisateur d'entraînement pour améliorer la qualité générative, ce travail l'adapte pour l'inférence dans les problèmes inverses où les signaux de vérité terrain sont indisponibles.

Mise en Œuvre Technique

1. Construction d'un Proxy : Puisque la vérité terrain $x_0$ est inconnue, la méthode construit une représentation proxy $c_{proxy}$ pour remplacer les caractéristiques cibles manquantes.

   • Initialement, $c_{proxy}$ est dérivée des caractéristiques DINOv2 de la mesure disponible $y$.
   • Au fur et à mesure que le processus de diffusion inverse progresse, le proxy est progressivement mis à jour avec les caractéristiques DINOv2 de l'estimation débruitée actuelle du modèle, $f_{DINOv2}(D(E[z_0 | z_t]))$.
   • Cette stratégie repose sur la robustesse observée des caractéristiques DINOv2 face aux dégradations d'images courantes.

2. Distribution Penchée : Le processus d'échantillonnage est guidé par une distribution penchée qui intègre le terme d'alignement comme récompense :
   $$\tilde{p}(x_t | y) \propto p(y | x_t) p_t(x_t) \exp(\lambda \cdot \text{REPA}(\bar{x}_0, h_t))$$
   Ici, $p(y | x_t)$ impose la cohérence avec la mesure, $p_t(x_t)$ est le prior non conditionnel, et le terme REPA encourage l'alignement entre les états intermédiaires cachés du modèle de diffusion $h_t$ et les caractéristiques proxy DINOv2.

3. Intégration de l'Algorithme : L'approche est intégrée dans des solveurs existants de l'état de l'art (par exemple, Latent DPS, ReSample) en ajoutant une étape de gradient basée sur l'objectif REPA au cours de la boucle d'échantillonnage inverse. Cette étape aligne l'état de diffusion latente avec les caractéristiques sémantiques de la cible.

4. Adaptabilité : Pour les modèles qui n'ont pas été préentraînés avec le REPA, le cadre apprend une tête de projection légère post-hoc ($g_\phi$) pour mapper les représentations internes du modèle de diffusion vers l'espace de caractéristiques DINOv2. Cette cartographie est optimisée tout en maintenant le modèle de diffusion figé, rendant l'adaptation peu coûteuse en calcul.

Insights Théoriques

L'article fournit deux résultats théoriques principaux caractérisant le régularisateur REPA :

1. Minimisation de la Divergence : Il est démontré que maximiser l'objectif REPA équivaut à minimiser la Discrepancy à Moyenne Maximale (MMD) entre les distributions de caractéristiques DINOv2 de l'image reconstruite et de la vérité terrain, à des termes près rendant compte des erreurs d'alignement des caractéristiques et d'approximation du proxy.
2. Contraction de l'État Latent : Sous des hypothèses de régularité concernant les jacobiens de l'encodeur de diffusion et de la tête de projection, une seule mise à jour REPA induit une contraction dans l'espace d'état de diffusion. Plus précisément, elle oriente l'état latent actuel $z_t$ plus près de la représentation latente de l'image propre $z^*$, à condition que l'état actuel se trouve dans un voisinage local de la solution.

Résultats Expérimentaux

Les auteurs ont évalué la méthode sur quatre problèmes inverses : Super-Résolution (4×, 8×, 16×), Inpainting par Boîte, Défloutage Gaussien et Défloutage par Flou de Mouvement, en utilisant les jeux de données ImageNet et FFHQ.

• Qualité Perceptuelle : L'intégration du REPA a systématiquement amélioré les métriques perceptuelles (LPIPS et FID) sur toutes les tâches et tous les solveurs (Latent DPS, ReSample, Latent DAPS). Par exemple, en Super-Résolution 4× sur ImageNet, Latent DPS + REPA a réduit le LPIPS de 0,238 à 0,217 et le FID de 123,82 à 86,88.
• Métriques Pixel : Les métriques standards au niveau du pixel (PSNR, SSIM) sont restées comparables aux références, reflétant le compromis connu entre perception et distorsion où le REPA privilégie le réalisme sémantique par rapport à la précision parfaite au niveau du pixel.
• Efficacité : Un résultat significatif est que le REPA permet aux solveurs d'atteindre le même niveau de qualité perceptuelle avec nettement moins d'étapes de discrétisation. En super-résolution et en défloutage gaussien, la méthode a obtenu des scores LPIPS comparables avec environ 4 fois moins d'étapes que la référence, entraînant une réduction du temps d'exécution global malgré le calcul supplémentaire par étape.
• Robustesse : La méthode est restée efficace sous des dégradations sévères (par exemple, mise à l'échelle 16×, flou de haute intensité) et lorsqu'elle a été appliquée à des modèles préentraînés indépendamment via la tête de projection apprise.

Importance et Revendications

L'article revendique que l'alignement de représentation constitue un outil efficace pour améliorer les solveurs de problèmes inverses génératifs au moment de l'inférence.

• Guidage Sémantique : En alignant les représentations internes du modèle avec des caractéristiques visuelles robustes (DINOv2), la méthode fournit un prior sémantique fort qui guide la reconstruction vers des solutions perceptuellement significatives, même sans supervision par vérité terrain.
• Justification Théorique : Ce travail élucide pourquoi l'alignement fonctionne, montrant qu'il agit comme un mécanisme pour minimiser la divergence dans l'espace de caractéristiques et contracter l'état latent vers la solution d'image propre.
• Applicabilité Générale : Le cadre est général, compatible avec divers solveurs basés sur la diffusion/le flux, et adaptable aux modèles qui n'ont pas été initialement entraînés avec un alignement via une cartographie post-hoc légère.
• Efficacité : Il adresse le goulot d'étranglement computationnel de l'échantillonnage de diffusion en accélérant la convergence, permettant une reconstruction de haute qualité avec moins d'étapes d'échantillonnage.

Les auteurs concluent que cette approche comble le fossé entre les capacités génératives des modèles de diffusion et la compréhension sémantique des encodeurs auto-supervisés, offrant une manière principielle d'améliorer la qualité de reconstruction dans des contextes de problèmes inverses difficiles.