Stochastic Generative Plug-and-Play Priors

1. Problématique

Les problèmes inverses en imagerie (comme la reconstruction IRM, la défloutage, la super-résolution ou l'inpainting) consistent à retrouver un signal inconnu $x$ à partir de mesures bruitées et incomplètes $y = A(x) + e$. La méthode Plug-and-Play (PnP) est une approche populaire qui résout ces problèmes en alternant entre une étape de fidélité aux données et une étape de régularisation via un débruiteur appris (désigné par $D_\theta$).

Cependant, les méthodes PnP traditionnelles rencontrent deux limitations majeures dans les problèmes fortement mal posés (où les données manquantes sont importantes, comme dans l'inpainting avec de grands masques) :

1. Inadéquation de la distribution : Les itérés intermédiaires du PnP contiennent des artefacts structurés et non un bruit gaussien additif simple. Les débruiteurs classiques (entraînés pour un faible niveau de bruit) échouent donc souvent car ils sont utilisés en dehors de leur régime d'entraînement.
2. Limites des débruiteurs classiques : La plupart des débruiteurs utilisés en PnP sont optimisés pour la restauration à faible bruit, ce qui limite leur capacité à générer du contenu sémantique plausible lorsque l'information est fortement manquante.

Bien que les Modèles de Diffusion basés sur le Score (SBDM) aient démontré une grande puissance générative en étant entraînés sur une large gamme de niveaux de bruit, leur intégration systématique dans le cadre PnP sans recourir à un échantillonnage de diffusion inversé complet à chaque étape reste un défi ouvert.

2. Méthodologie Proposée : SGPnP

Les auteurs proposent un cadre appelé Stochastic Generative Plug-and-Play (SGPnP) qui combine l'interprétation par le score des SBDM avec une injection de bruit stochastique.

A. Adaptation par le Score (Score Adaptation)

Pour utiliser directement des SBDM pré-entraînés (qu'ils soient de type Variance-Exploding - VE, ou Variance-Preserving - VP) comme débruiteurs dans un algorithme PnP, les auteurs établissent un lien mathématique direct via la formule de Tweedie.

• Ils définissent une relation entre le score (gradient de la densité de probabilité) et la fonction de débruitage.
• Ils proposent une stratégie d'appariement des paramètres pour mapper le niveau de bruit continu $\sigma$ utilisé par l'algorithme PnP vers l'index de temps discret $t$ utilisé par le modèle de diffusion pré-entraîné, permettant ainsi d'utiliser le SBDM comme un débruiteur conditionné par le bruit.

B. Cadre Stochastique Génératif (SGPnP)

La contribution centrale est l'introduction d'une étape de ré-bruitage (re-noising) explicite avant l'étape de débruitage.

• Mécanisme : Au lieu d'appliquer le débruiteur directement sur l'itéré intermédiaire $z_k$, on injecte un bruit gaussien $\sigma_{inject} n$ avant le débruiteur :
  $$x_{k+1} = D_\theta(z_k + \sigma_{inject} n; \sigma_{cond})$$
• Découplage des niveaux de bruit : Une innovation clé est que le niveau de bruit injecté ($\sigma_{inject}$) peut être différent du niveau de conditionnement du débruiteur ($\sigma_{cond}$). Cela permet de tenir compte à la fois du bruit injecté et des artefacts résiduels introduits par l'étape de fidélité aux données.
• Objectif : Cette injection aligne l'entrée du débruiteur avec la distribution d'entraînement du SBDM (qui s'attend à du bruit gaussien) et introduit une stochasticité nécessaire pour l'optimisation.

3. Contributions Clés et Analyse Théorique

1. Interprétation par le Score du PnP : L'article établit un lien mathématique formel entre les itérations PnP classiques et le débruitage basé sur le score, justifiant l'utilisation directe de SBDM pré-entraînés sans boucles de diffusion inversée complètes.
2. Évasion des points selle stricts (Saddle-point Escape) :
   • Les auteurs prouvent théoriquement que l'injection de bruit induit une optimisation sur un objectif lissé par un noyau gaussien.
   • Sous certaines hypothèses (notamment que le débruiteur est un estimateur MMSE), le bruit injecté crée une dérive stochastique suffisante le long des directions de courbure négative, permettant à l'algorithme d'échapper aux points selle stricts (des points critiques qui ne sont pas des minima locaux). C'est la première analyse théorique de ce type dans la littérature PnP.
3. Convergence vers l'objectif original :
   • En utilisant un programme de recuit (annealing) où le niveau de bruit injecté diminue progressivement vers zéro, les auteurs démontrent que les itérés convergent vers un point stationnaire de l'objectif original (non lissé), garantissant ainsi la cohérence asymptotique avec le problème MAP initial.

4. Résultats Expérimentaux

Les expériences ont été menées sur deux types de données : des images de visages (FFHQ) et des données IRM multi-bobines (fastMRI). Les tâches incluent l'inpainting (masques aléatoires et grands masques rectangulaires), la super-résolution, le défloutage et la reconstruction IRM accélérée.

• Performance Quantitative :
  • SGPnP surpasse systématiquement les méthodes PnP déterministes (DPIR, PnP-ADMM) et stochastiques existantes (SNORE) en termes de PSNR, SSIM et LPIPS.
  • Sur des tâches fortement mal posées (ex: inpainting avec un grand masque central), SGPnP atteint des performances compétitives avec des solveurs basés sur la diffusion pure (comme DPS et DiffPIR), tout en conservant la structure d'optimisation du PnP.
• Performance Qualitative :
  • Contrairement aux méthodes déterministes qui produisent des reconstructions floues ou incomplètes dans les zones masquées, SGPnP réussit à "halluciner" du contenu sémantique réaliste et cohérent.
  • L'ablation montre que l'utilisation d'un SBDM avec une large couverture de niveaux de bruit est cruciale, tout comme le mécanisme d'injection de bruit stochastique.
• Validation Théorique : Les résultats confirment que l'injection de bruit améliore la robustesse là où les méthodes déterministes échouent, validant l'hypothèse d'évasion des points selle.

5. Signification et Impact

Ce travail est significatif car il comble le fossé entre les méthodes d'optimisation classiques (PnP) et les modèles génératifs modernes (SBDM).

• Il permet d'utiliser la puissance générative des modèles de diffusion directement dans des solveurs d'optimisation, sans avoir à exécuter des chaînes de Markov longues à chaque étape.
• Il fournit une garantie théorique sur le comportement de l'optimisation stochastique dans les problèmes inverses, expliquant pourquoi l'injection de bruit fonctionne mieux que les approches purement déterministes pour les problèmes mal posés.
• La méthode SGPnP offre une solution robuste et efficace pour des applications critiques comme l'imagerie médicale (IRM), où la reconstruction de zones manquantes avec une haute fidélité est essentielle.

En résumé, SGPnP transforme le PnP en un cadre stochastique génératif capable de résoudre des problèmes d'imagerie inverse sévèrement mal posés en exploitant pleinement le potentiel des modèles de diffusion pré-entraînés.