Résumé technique : Guidance amortie pour l'inpainting d'images avec des modèles de diffusion préentraînés (AID)

Énoncé du problème

L'inpainting d'images est une tâche de génération conditionnelle visant à reconstruire les régions manquantes d'une image tout en maintenant la cohérence avec les pixels visibles. Bien que les modèles de diffusion préentraînés aient considérablement fait progresser cette capacité, les approches existantes font face à une dichotomie entre efficacité et adaptabilité :

1. Entraînement spécifique à la tâche : Des méthodes comme LaMa, MAT et BrushNet entraînent des modèles dédiés ou de grands modules d'adaptation. Bien que cela amortisse le coût de l'inférence, cela nécessite un temps d'entraînement substantiel et ajoute une complexité significative au modèle.
2. Optimisation par instance : Des méthodes comme RePaint, DPS et MCG maintiennent le backbone de diffusion préentraîné fixe, mais l'adaptent séparément pour chaque image masquée au moment du test. Cela préserve le modèle préentraîné, mais engendre des coûts d'optimisation répétés lors du déploiement, traitant chaque image comme un nouveau problème d'inférence.

Il existe un besoin pour une solution « intermédiaire » qui amortisse l'inférence sans entraîner un grand modèle spécifique à la tâche ni effectuer d'optimisation par instance au moment du test.

Méthodologie : Inpainting amorti avec diffusion (AID)

Les auteurs proposent AID, un cadre qui maintient le backbone de diffusion préentraîné fixe et entraîne un petit module de guidance réutilisable hors ligne. Ce module est ensuite réutilisé pour toutes les images masquées au moment du test, sans optimisation supplémentaire.

1. Formulation du contrôle

Le problème est formulé comme un problème de contrôle déterministe.

• Dynamique : Le processus de diffusion inverse est régi par une équation différentielle ordinaire (ODE) de flot de probabilité pilotée par une fonction de score préentraînée fixe $\hat{S}$.
• Guidance : Un terme de contrôle dépendant de la tâche $u(t, x; \xi)$ est ajouté à la dérive, où $\xi = (M, y)$ représente le masque observable et les pixels visibles. Le contrôle ne dépend pas de l'image propre cachée $x^\dagger$.
• Objectif : Le but est de minimiser une perte terminale $\Psi$ (équilibrant la fidélité aux pixels visibles et à la région manquante) plus un coût de fonctionnement pénalisant l'amplitude du champ de guidance. Cela conduit à une équation de Hamilton-Jacobi-Bellman (HJB) où le contrôle optimal $u^*$ est déterministe.

2. Pont théorique : du déterministe au stochastique

Résoudre directement l'équation HJB déterministe de haute dimension est intraitable. Les auteurs introduisent un problème de contrôle aléatoire auxiliaire utilisant des politiques gaussiennes pour rendre le problème apprenable via l'apprentissage par renforcement en temps continu (CT-RL).

• Problème auxiliaire : Au lieu d'un contrôle déterministe, le système utilise une politique gaussienne $\pi^{(\lambda)} \sim \mathcal{N}(\mu, \sigma^2)$, où la moyenne $\mu$ est la loi de guidance candidate.
• Pont préservant l'optimiseur : L'article prouve deux théorèmes critiques :
  1. Suffisance : La guidance déterministe optimale $u^*$ peut être relevée vers une politique gaussienne optimale en fixant la moyenne de la politique à $u^*$.
  2. Nécessité : Toute politique gaussienne optimale pour le problème auxiliaire doit avoir une moyenne qui récupère exactement la guidance déterministe optimale $u^*$.
• Signification : Ce pont garantit que l'apprentissage d'une politique stochastique (qui permet des mises à jour basées sur le gradient) n'altère pas la guidance déterministe cible utilisée lors du déploiement. La variance de la politique gaussienne est fixe et sert uniquement d'outil d'apprentissage, et non de régularisateur d'entropie.

3. Algorithme Acteur-Critic amorti

Sur la base de ce pont théorique, les auteurs développent un algorithme acteur-critic entièrement piloté par les données :

• Paramétrisation : L'acteur (moyenne de guidance $\hat{\mu}_\phi$) et le critique (fonction de valeur $\hat{V}_\theta$) sont paramétrés par des réseaux de neurones.
• Apprentissage : L'algorithme utilise des conditions d'orthogonalité de martingale dérivées de l'équation HJB auxiliaire. Il met à jour l'acteur et le critique en utilisant une approximation stochastique basée sur des trajectoires échantillonnées à partir de la distribution des tâches.
• Déploiement : Une fois entraîné, le module de guidance est déployé de manière déterministe. Pour une nouvelle image masquée, le système utilise la guidance moyenne apprise $\hat{\mu}_\phi$ au sein du solveur en temps inverse, en rejetant le bruit de la politique.

Contributions clés

1. Méthodologie : AID introduit un cadre de guidance contrôlée qui amortit l'inférence en entraînant hors ligne un module de guidance léger, qui est ensuite réutilisé pour toutes les tâches sans optimisation par instance.
2. Théorie : L'article établit un « pont d'équivalence de politique » rigoureux entre le problème de guidance déterministe et une formulation gaussienne auxiliaire. Il prouve que le problème auxiliaire préserve l'optimiseur, justifiant l'utilisation de méthodes acteur-critic stochastiques pour apprendre un champ de contrôle déterministe.
3. Efficacité : L'approche ajoute moins de 1 % de paramètres entraînables par rapport au backbone préentraîné tout en améliorant significativement le compromis qualité-vitesse.

Résultats empiriques

Les auteurs ont évalué AID sur AFHQv2, FFHQ (en utilisant le pipeline EDM dans l'espace des pixels) et ImageNet (en utilisant le pipeline EDM2 dans l'espace latent).

• Performance par rapport aux bases à backbone fixe : AID surpasse systématiquement de fortes bases à backbone fixe (Unguided, Replacement, MCG, DPS) sur plusieurs métriques (PSNR, SSIM, LPIPS) et types de masques (libre forme, centre, bande).
• Performance par rapport à RePaint : AID dépasse RePaint (une forte base d'optimisation itérative) en qualité tout en n'utilisant que ~10 % du budget d'échantillonnage (par exemple, 35 étapes contre 350 étapes pour RePaint).
• Performance par rapport à LatentPaint : Sur ImageNet, AID surpasse LatentPaint (une méthode d'adaptation réutilisable dans l'espace latent) sur toutes les métriques tout en utilisant un module de guidance avec moins de 45 % des paramètres entraînables.
• Latence : AID prend en charge un point de fonctionnement à faible latence (12 étapes) qui réduit davantage le temps d'inférence tout en maintenant une performance supérieure aux bases non-RePaint et restant compétitive avec RePaint.
• Généralisation : Le même module de guidance entraîné se transfère efficacement à des types de masques non vus et à différents budgets d'échantillonnage sans réentraînement.

Signification et revendications

L'article revendique qu'AID est le premier cadre de guidance contrôlée amortie pour l'inpainting d'images qui :

1. Maintient le backbone de diffusion préentraîné gelé.
2. Entraîne uniquement un petit module de guidance réutilisable.
3. Fournit un pont exact préservant l'optimiseur de l'apprentissage acteur-critic stochastique vers le champ de guidance déterministe utilisé lors du déploiement.

Les auteurs positionnent ce travail comme une étape vers l'utilisation de cadres basés sur le contrôle et l'apprentissage par renforcement pour fournir des méthodes théoriquement fondées pour étudier et améliorer les modèles de diffusion. Ils soulignent que l'apprentissage par renforcement en temps continu sous-tend à la fois l'analyse rigoureuse et la conception pratique de l'algorithme, offrant une alternative de principe à l'optimisation par instance heuristique ou au réentraînement spécifique à la tâche lourd.