Résumé Technique : FlowADMM : ADMM Plug-and-Play avec des Priors de Rébruitage-Débruitage Basés sur les Flots

Énoncé du Problème

Les problèmes inverses en imagerie, tels que le débruitage, le défloutage, la super-résolution et l'inpainting, sont typiquement mal posés et formulés comme la récupération d'une image $x$ à partir de mesures dégradées $y = Ax + \nu$. Bien que l'apprentissage profond offre des mappings directs puissants, ils manquent souvent d'interprétabilité et de garanties théoriques. Les approches variationnelles classiques formulent cela comme la minimisation d'un terme de fidélité aux données plus un terme de régularisation, souvent résolu via des algorithmes proximaux comme la séparation avant-arrière (FBS).

Les méthodes Plug-and-Play (PnP) améliorent cela en remplaçant l'opérateur proximal d'un régularisateur conçu à la main par un débruiteur appris. Récemment, les modèles de flux génératifs et de diffusion ont fourni des familles dépendantes du temps de débruiteurs $\{D_t\}$ adaptés à des niveaux de bruit spécifiques. Cependant, un défi majeur surgit : lors de la reconstruction itérative, les itérés intermédiaires s'écartent souvent de la distribution de bruit $p_t$ pour laquelle le débruiteur a été entraîné. Les méthodes PnP basées sur les flots existantes (par exemple, PnP-Flow, Flower) abordent cela en incorporant des étapes explicites de rébruitage-débruitage, où l'itéré courant est perturbé vers le niveau de bruit approprié avant le débruitage. Ces méthodes moyennent souvent sur plusieurs échantillons de bruit pour stabiliser le processus.

Malgré leur succès empirique, les méthodes PnP de rébruitage-débruitage existantes sont formulées comme des algorithmes stochastiques, rendant leur comportement de convergence difficile à analyser théoriquement. Plus précisément, l'opérateur déterministe implicitement induit par la moyenne de la sortie du débruiteur sur la distribution de bruit latente n'a pas été caractérisé, laissant les propriétés de convergence de ces méthodes mal comprises.

Méthodologie

L'Opérateur Moyen de Rébruitage-Débruitage

Les auteurs identifient et formalisent l'opérateur déterministe sous-jacent aux méthodes PnP stochastiques basées sur les flots. Étant donné un réseau de vitesse de flot préentraîné $v_\theta$ définissant un débruiteur MMSE dépendant du temps $D_t = I + (1-t)v_\theta^t$, l'étape stochastique standard consiste à rébruite un itéré $x$ avec un bruit $\epsilon \sim \mathcal{N}(0, I)$ pour obtenir $tx + (1-t)\epsilon$, suivi d'un débruitage.

L'article définit l'opérateur moyen de rébruitage-débruitage $\bar{S}_t(x)$ comme l'espérance de ce processus sur la distribution de bruit :
$$\bar{S}_t(x) = \mathbb{E}_{\epsilon \sim \mathcal{N}(0, I)} [D_t(tx + (1-t)\epsilon)]$$
Cet opérateur représente la limite déterministe des mises à jour stochastiques utilisées par des méthodes comme PnP-Flow. Les auteurs analysent sa forme résiduelle $R_t(x) = \bar{S}_t(x) - x$, l'interprétant comme un champ de vecteurs fournissant des corrections locales vers la variété des données. Sous des hypothèses MMSE idéales, les points sur la variété des données sont des points fixes de cet opérateur.

Algorithme FlowADMM

S'appuyant sur cette idée, les auteurs proposent FlowADMM, qui intègre l'opérateur moyen de rébruitage-débruitage dans le cadre de la méthode du multiplicateur de direction alternée (ADMM). Contrairement à FBS, ADMM permet plusieurs opérateurs de prior et gère les termes de données non différentiables sans restrictions de pas dépendantes de l'opérateur.

L'algorithme résout le problème contraint $\min_{x,z} F_y(x) + R(z)$ sous la contrainte $x=z$ en minimisant le lagrangien augmenté. Les mises à jour sont :

1. Cohérence des données : $x^{k+1} = \text{prox}_{\tau F_y}(z^k - u^k)$
2. Mise à jour du prior : $z^{k+1} = \bar{S}_{t_k}(x^{k+1} + u^k)$
3. Mise à jour duale : $u^{k+1} = u^k + x^{k+1} - z^{k+1}$

Puisque l'espérance dans $\bar{S}_t$ est intraitable, elle est approximée par échantillonnage de Monte Carlo. Les auteurs proposent une stratégie d'échantillonnage dépendante du pas de temps : utiliser moins d'échantillons dans les premières itérations (où des corrections grossières sont apportées et où l'opérateur est moins sensible au bruit) et augmenter le nombre d'échantillons dans les itérations ultérieures (où les détails fins sont affinés et la précision est critique).

Analyse de Convergence

L'article établit des garanties de convergence pour FlowADMM sous des conditions de Lipschitz faibles sur le réseau de flot sous-jacent :

• Calendrier de temps fixe : Pour un pas de temps constant $t$, la convergence est garantie si la constante de Lipschitz $L_v$ du réseau de flot satisfait $L_v < \frac{1}{1-t}$. Cette condition est significativement plus faible que la contraction stricte ($L < 1$) requise par de nombreuses autres méthodes PnP, surtout lorsque $t \to 1$.
• Calendriers de temps non stationnaires : L'analyse est étendue aux calendriers de temps $(t_k)_k$ qui convergent vers une limite $t_{max} < 1$. Les auteurs prouvent la convergence vers un point fixe de l'opérateur limite $T_{t_{max}}$ à condition que la suite d'opérateurs approche la limite suffisamment rapidement et que le réseau de flot soit lipschitzien par rapport au temps $t$.

Contributions Clés

1. Caractérisation Théorique : Les auteurs identifient et analysent l'opérateur moyen déterministe de rébruitage-débruitage sous-jacent aux méthodes PnP stochastiques basées sur les flots, caractérisant sa structure résiduelle et son comportement de Lipschitz.
2. Proposition d'Algorithme : Ils proposent FlowADMM, un algorithme PnP qui intègre directement cet opérateur moyen dans le cadre ADMM.
3. Garanties de Convergence : Ils dérivent des critères de convergence pour FlowADMM sous des hypothèses de Lipschitz faibles, couvrant à la fois les calendriers de pas de temps constants et non constants (variant dans le temps).
4. Stratégie d'Échantillonnage : Ils proposent une stratégie d'échantillonnage Monte Carlo de grossier à fin qui augmente le nombre d'échantillons dans les itérations ultérieures, motivée par le comportement de l'opérateur moyen.
5. Performance Empirique : Ils démontrent des performances de pointe sur plusieurs problèmes inverses tout en nécessitant moins d'évaluations de cohérence des données que les approches antérieures de moyennage stochastique.

Résultats Expérimentaux

Les auteurs ont évalué FlowADMM sur les ensembles de données CelebA et AFHQ-Cat à travers cinq tâches : débruitage, défloutage, super-résolution, inpainting aléatoire et inpainting par boîte.

• Performance : FlowADMM a obtenu des résultats de pointe dans les métriques de distorsion (PSNR et SSIM) sur 7 tâches sur 10 par rapport aux méthodes PnP récentes basées sur les flots et la diffusion (incluant PnP-Flow, Flower, OT-ODE, D-Flow et DiffPIR). Des améliorations notables ont été observées dans les tâches de super-résolution et d'inpainting.
• Efficacité : Tout en correspondant au nombre total d'évaluations de flot (passes avant du réseau) du concurrent de premier plan Flower5-OT, FlowADMM a nécessité significativement moins d'évaluations d'opérateurs proximaux de terme de données. Cela s'explique par le fait que Flower5-OT moyenne des trajectoires de reconstruction complètes, tandis que FlowADMM moyenne uniquement l'étape de prior au sein de chaque itération.
• Études d'Ablation : Les expériences ont confirmé que le calendrier d'échantillonnage "lourd en fin" (augmentation des échantillons dans les itérations ultérieures) améliore significativement les performances par rapport à un échantillonnage constant. De plus, l'application de ce calendrier à PnP-Flow a amélioré ses résultats, mais FlowADMM l'a toujours surpassé, suggérant que la formulation ADMM elle-même apporte un avantage.
• Limitations : L'analyse de convergence repose sur des constantes de Lipschitz du réseau préentraîné qui ne sont pas explicitement pénalisées pendant l'entraînement. De plus, les garanties théoriques s'appliquent à l'opérateur moyen déterministe, tandis que l'implémentation pratique utilise des approximations stochastiques de Monte Carlo.

Importance et Revendications

L'article revendique que, en formalisant explicitement l'opérateur déterministe sous-jacent aux méthodes stochastiques de rébruitage-débruitage, il comble le fossé entre le succès empirique et la compréhension théorique dans les méthodes PnP basées sur les flots. Le travail fournit :

• Convergence Prouvable : Contrairement aux méthodes PnP stochastiques précédentes, FlowADMM est accompagné de garanties de convergence sous des conditions de Lipschitz réalistes sur le réseau de flot.
• Efficacité Algorithmique : L'intégration dans ADMM permet une utilisation plus efficace des ressources de calcul (moins d'étapes de cohérence des données) par rapport aux méthodes de moyennage de trajectoires.
• Insight Théorique Opérateur : L'analyse suggère que le comportement "de grossier à fin" de l'opérateur moyen peut être exploité via des calendriers d'échantillonnage adaptatifs, un principe qui bénéficie à la fois aux approches basées sur ADMM et FBS.

Les auteurs concluent que FlowADMM offre un cadre robuste et théoriquement fondé pour résoudre des problèmes inverses avec des priors basés sur les flots, bien qu'ils reconnaissent que des travaux futurs pourraient se concentrer sur l'imposition de contraintes de Lipschitz pendant l'entraînement pour renforcer les garanties de stabilité.