Échantillonnage MAP Local pour Modèles de Diffusion

Informations Fondamentales

• ID de l'article: 2510.07343
• Titre: Local MAP Sampling for Diffusion Models
• Auteurs: Shaorong Zhang (UC Riverside), Rob Brekelmans (Vector Institute), Greg Ver Steeg (UC Riverside)
• Classification: cs.GR cs.AI eess.IV
• Date de publication/Conférence: Préimpression (En révision)
• Lien de l'article: https://arxiv.org/abs/2510.07343

Résumé

L'échantillonnage postérieur par diffusion (DPS) fournit une approche bayésienne rigoureuse pour les problèmes inverses en échantillonnant à partir de $p(x_0 \mid y)$. Cependant, en pratique, l'objectif de la résolution de problèmes inverses n'est pas de couvrir la distribution postérieure, mais de récupérer la reconstruction la plus précise. Les solveurs de diffusion basés sur l'optimisation excellent généralement dans cette tâche, bien qu'ils manquent de fondements probabilistes clairs. Cet article introduit l'échantillonnage MAP local (LMAPS), un nouveau cadre d'inférence qui résout itérativement des sous-problèmes MAP locaux le long de la trajectoire de diffusion. Cette perspective clarifie les connexions avec l'estimation MAP globale et DPS, fournissant une interprétation probabiliste unifiée pour les méthodes basées sur l'optimisation. Sur cette base, nous développons des algorithmes pratiques avec des approximations de covariance probabilistiquement interprétables, des objectifs reformulés pour la stabilité et l'interprétabilité, ainsi que des approximations de gradient pour les opérateurs non différentiables.

Contexte de Recherche et Motivation

Contexte du Problème

Les modèles de diffusion font face à deux défis majeurs dans la résolution de problèmes inverses :

1. Inadéquation des objectifs: DPS vise à échantillonner à partir de la distribution postérieure $p(x_0|y)$, mais l'objectif réel de la résolution de problèmes inverses est d'obtenir la meilleure reconstruction, non une diversité d'échantillons
2. Absence de fondements théoriques: Les solveurs de diffusion basés sur l'optimisation (tels que Resample, DiffPIR, DCDP, etc.) excellent en performance mais manquent de fondements théoriques probabilistes clairs

Motivation de la Recherche

• Orientation pratique: Les protocoles d'évaluation des problèmes inverses comparent généralement avec une seule référence de vérité, sans récompenser la couverture ou la diversité
• Unification théorique: Nécessité de fournir une interprétation probabiliste pour les méthodes d'optimisation, clarifiant leurs relations avec l'estimation MAP et DPS
• Amélioration des performances: Réaliser de meilleures performances de reconstruction tout en maintenant les fondements théoriques

Contributions Principales

1. Contribution théorique: Proposition du cadre d'échantillonnage MAP local (LMAPS), résolvant itérativement des sous-problèmes MAP locaux le long de la trajectoire de diffusion, analyse des relations avec MAP global et DPS, unification de TMPD et des méthodes d'optimisation pour les problèmes inverses
2. Contribution méthodologique:
   • Fourniture d'approximations de covariance probabilistiquement interprétables, remplaçant les choix heuristiques dans les solveurs existants
   • Introduction de reformulations d'objectifs pour des paramètres interprétables et une stabilité améliorée
   • Développement de stratégies d'approximation de gradient pour les opérateurs non différentiables
3. Contribution expérimentale: Validation sur 10 tâches de restauration d'images et 3 problèmes inverses scientifiques, obtenant les meilleurs résultats dans 46/60 cas FFHQ/ImageNet, avec des améliorations PSNR ≥2dB sur le flou de mouvement, la restauration JPEG et les tâches de quantification

Détails de la Méthode

Définition de la Tâche

Le problème inverse vise à récupérer une image ou un signal inconnu $x_0 \in \mathbb{R}^n$ à partir d'une distribution a priori $\pi(x_0)$ et d'une mesure bruitée $y \in \mathbb{R}^m$ : $y = H(x_0) + z$ où $H(\cdot): \mathbb{R}^n \to \mathbb{R}^m$ est l'opérateur direct et $z \sim \mathcal{N}(0, \sigma_y^2 I)$ est le bruit de mesure.

Cadre Théorique Principal

MAP Global vs MAP Local

MAP Global optimise directement la postérieure complète : $x_0^{\text{MAP}} := \arg\max_{x_0} p(x_0|y)$

MAP Local résout un problème d'optimisation conditionnelle à chaque pas de temps $t$ : $x_0^*(t, x_t, y) := \arg\max p(x_0|x_t, y)$$x_{t-\Delta t} = g(x_0^*, x_t, \epsilon), \quad \epsilon \sim \mathcal{N}(0,I)$

Relations avec DPS et DAPS

• DPS: Utilise la moyenne conditionnelle $E[x_0|x_t, y]$
• DAPS: Échantillonne à partir de $p(x_0|x_t, y)$
• LMAPS: Utilise le mode conditionnel $\arg\max p(x_0|x_t, y)$

Insight clé : DPS et LMAPS sont équivalents uniquement lorsque $p(x_0|x_t, y)$ suit une distribution gaussienne.

Conception d'Algorithmes Pratiques

Approximation de Covariance

Adoption d'une approximation isotrope : $\Sigma_{0|t} \approx \frac{k}{\text{SNR}} I, \quad \text{SNR} := \frac{\alpha_t^2}{\sigma_t^2}$

Reformulation de la Fonction Objectif

Reformulation de l'objectif original : $x_0^* = \arg\min \left\{\frac{\text{SNR}}{k}\|x_0 - m_{0|t}\|^2 + \frac{1}{\sigma_y^2}\|y - H(x_0)\|^2\right\}$

en : $x_0^* = \arg\min \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 = \frac{\sigma_t^2}{\sigma_t^2 + k_1^2} \in (0,1)$, réalisant :

• Interprétation de combinaison convexe: Poids $(1-\mu_t)$ et $\mu_t$
• Recuit automatique: Transition de l'entraînement par mesure à l'entraînement par a priori à mesure que $\sigma_t^2$ diminue
• Stabilité numérique: Évite les mises à l'échelle SNR extrêmes

Traitement des Opérateurs Non Différentiables

Pour les tâches non différentiables comme la restauration JPEG et la quantification, utilisation de gradients de substitution : $\nabla_{x_0}\|y - H(x_0)\|^2 \approx 2J_{H'}(x_0)^T(H(x_0) - y)$

Pour la quantification, adoption de $H'(x_0) = x_0$, simplifiant en : $\nabla_{x_0}\|y - H(x_0)\|^2 \approx 2(H(x_0) - y)$

Configuration Expérimentale

Ensembles de Données

• Restauration d'images: FFHQ 256×256 et ImageNet 256×256, utilisant chacun 100 images de test
• Problèmes inverses scientifiques: Utilisation de l'ensemble de données InverseBench, incluant des images de microscopie de fluorescence (diffusion inverse linéaire), données GRMHD (imagerie de trou noir), données fastMRI du genou (IRM par détection comprimée)

Métriques d'Évaluation

• Restauration d'images: PSNR, SSIM, LPIPS
• Problèmes inverses scientifiques: Principalement PSNR, complétés par des métriques spécifiques aux tâches

Méthodes de Comparaison

Incluant DDNM, DDRM, ΠGDM, DPS, LGD, PnP-DM, FPS, MCG-diff, RedDiff, DAPS, DiffPIR, DCDP, DMPlug et 12 autres méthodes

Détails d'Implémentation

• Pas de diffusion: Généralement 200 pas
• Pas de mise à jour de gradient: 20-200 pas (dépendant de la tâche)
• Taux d'apprentissage: 0,01-1,0 (dépendant de la tâche)
• Paramètres $k_1$: 0-10, $k_2$: 0,01-30000

Résultats Expérimentaux

Résultats Principaux

Tâches de Restauration d'Images

Dans le tableau 1, LMAPS obtient les meilleures performances dans 49 résultats sur 60 :

• Défloutage de mouvement: 32,62 dB sur FFHQ vs 29,66 dB pour DAPS (+2,96 dB)
• Restauration JPEG: 27,25 dB sur FFHQ vs 25,04 dB pour ΠGDM (+2,21 dB)
• Quantification: 29,51 dB sur FFHQ vs 25,82 dB pour ΠGDM (+3,69 dB)

Problèmes Inverses Scientifiques

Dans le tableau 2, LMAPS obtient le meilleur PSNR sur toutes les tâches :

• Diffusion inverse linéaire (NR=360): 38,07 dB vs 36,56 dB pour RED-diff (+1,51 dB)
• Diffusion inverse linéaire (NR=180): 37,19 dB vs 35,41 dB pour RED-diff (+1,78 dB)
• Diffusion inverse linéaire (NR=60): 30,75 dB vs 27,07 dB pour RED-diff (+3,68 dB)

Études d'Ablation

La figure 4 montre le compromis entre le nombre d'étapes d'optimisation et le nombre d'étapes de diffusion :

• Les meilleures performances sont généralement observées pour NFE=200-500
• L'augmentation du nombre d'étapes d'optimisation par pas de diffusion améliore significativement les performances
• Par rapport à SITCOM (600 NFEs), LMAPS atteint des performances similaires avec moins de ressources de calcul

Efficacité Computationnelle

Le tableau 3 montre le temps d'échantillonnage de LMAPS sur la tâche de défloutage :

• LMAPS (200 pas de diffusion, 100 pas d'optimisation): 61 secondes/image, 30,88 dB
• DAPS (200 pas de diffusion, 100 pas d'optimisation): 110 secondes/image, 29,19 dB
• SITCOM (600 pas): 73 secondes/image, 29,93 dB

Travaux Connexes

Échantillonnage Postérieur par Diffusion

DPS et ses variantes résolvent les problèmes inverses en échantillonnant directement à partir de la distribution postérieure $p(x_0|y)$, incluant les méthodes TMPD, DDNM, ΠGDM, etc.

Méthodes Basées sur l'Optimisation

Les méthodes Resample, DiffPIR, DCDP, DMPlug, etc. résolvent les problèmes inverses par débruitage, optimisation et rééchantillonnage alternés, avec d'excellentes performances mais manquant de fondements théoriques.

Méthodes d'Estimation MAP

Les travaux récents commencent à se concentrer sur l'estimation MAP sous les a priori de diffusion, mais sont principalement limités aux problèmes inverses linéaires.

Conclusions et Discussion

Conclusions Principales

1. Unification théorique: LMAPS fournit une interprétation probabiliste unifiée pour les méthodes de diffusion basées sur l'optimisation
2. Amélioration des performances: Obtient des améliorations PSNR significatives sur plusieurs tâches, particulièrement sur les tâches non linéaires et non différentiables difficiles
3. Efficacité computationnelle: Meilleure efficacité computationnelle par rapport aux méthodes existantes

Limitations

1. Convergence: La séquence MAP locale ne converge pas nécessairement vers le MAP global
2. Diversité: Par rapport à DPS, LMAPS peut produire moins de diversité de sortie
3. Sensibilité aux hyperparamètres: Nécessite l'ajustement des paramètres $k_1$ et $k_2$ pour différentes tâches

Directions Futures

L'article souligne que le rôle clé du MAP global en inférence bayésienne est largement négligé, et la résolution efficace du MAP global sous les a priori de diffusion reste un défi ouvert. Bien que le MAP puisse réduire la diversité de sortie en se concentrant sur un seul mode, il offre une plus grande certitude et un meilleur alignement avec les données observées.

Évaluation Approfondie

Avantages

1. Contribution théorique significative: Première fourniture d'une interprétation probabiliste claire pour les méthodes de diffusion basées sur l'optimisation
2. Expériences complètes: Couvre 10 tâches de restauration d'images et 3 problèmes inverses scientifiques, avec une configuration expérimentale complète
3. Améliorations de performance évidentes: Obtient des améliorations significatives de plus de 2dB sur plusieurs tâches difficiles
4. Méthode pratique: Fournit des stratégies efficaces pour traiter les opérateurs non différentiables
5. Rédaction claire: L'analyse théorique et la description des méthodes sont claires

Insuffisances

1. Profondeur de l'analyse théorique: Bien que fournissant une interprétation probabiliste, l'analyse de la convergence et des garanties théoriques est relativement limitée
2. Complexité des hyperparamètres: Nécessite l'ajustement de plusieurs hyperparamètres pour chaque tâche, pouvant affecter la généralisation de la méthode
3. Surcharge computationnelle: Bien que plus efficace que certaines méthodes, nécessite toujours plusieurs mises à jour de gradient à chaque pas de temps
4. Limitations d'évaluation: Se concentre principalement sur la qualité de reconstruction, avec une évaluation insuffisante de la quantification de l'incertitude

Impact

1. Valeur académique: Fournit une nouvelle perspective théorique pour l'application des modèles de diffusion aux problèmes inverses
2. Valeur pratique: Possède une valeur d'application directe dans les domaines de la restauration d'images et du calcul scientifique
3. Nature inspirante: Peut inspirer davantage de recherches sur l'estimation MAP des modèles de diffusion

Scénarios Applicables

• Problèmes inverses nécessitant une reconstruction de haute qualité plutôt que de la diversité
• Tâches impliquant des opérateurs directs non différentiables (comme la restauration JPEG, la quantification)
• Résolution de problèmes inverses en calcul scientifique
• Applications en temps réel avec certaines exigences d'efficacité computationnelle

Références

L'article cite des travaux importants dans les domaines des modèles de diffusion, de la résolution de problèmes inverses et de l'inférence bayésienne, incluant les articles originaux des méthodes clés telles que DPS, DAPS, TMPD, fournissant une bonne base bibliographique pour la recherche connexe.