Timestep Rescheduling in Diffusion Inversion

Résumé Technique : Reprogrammation des Pas de Temps dans l'Inversion de Diffusion

Énoncé du Problème
L'inversion de diffusion est le processus consistant à mapper une image réelle $z_0$ vers l'espace latent gaussien ($z_T$) d'un modèle de diffusion, permettant des tâches comme la reconstruction et l'édition d'images. Bien que les modèles de diffusion implicites de type DDIM (Denoising Diffusion Implicit Models) fournissent une formulation déterministe et efficace pour cela, le processus d'inversion est intrinsèquement approximatif. L'approche standard consiste à inverser analytiquement l'étape de débruitage, ce qui introduit des écarts qui s'accumulent au fil du temps, entraînant des erreurs de reconstruction notables et une dégradation de la fidélité de l'édition.

Les méthodes existantes tentent principalement d'atténuer ces erreurs en résolvant des problèmes de point fixe pour affiner l'équation d'inversion. Cependant, ces approches négligent largement un facteur critique : l'influence de la programmation des pas de temps (la sélection et l'espacement des pas de temps de diffusion) sur la fidélité de l'inversion. La plupart des pipelines utilisent des pas de temps espacés uniformément, ce qui, selon les auteurs, est sous-optimal. Le papier postule que la distribution des pas de temps impacte significativement l'accumulation des erreurs d'inversion globales, mais que cette relation reste sous-explorée par rapport à son rôle dans l'échantillonnage ou l'entraînement de la diffusion.

Méthodologie : TRDI (Timestep Rescheduling in Diffusion Inversion)
Les auteurs proposent TRDI, un programmateur de pas de temps non uniforme conçu pour minimiser l'erreur d'inversion globale sans nécessiter de paramètres supplémentaires ou de surcharge de calcul lors de la phase d'inférence. La méthodologie est fondée sur une analyse théorique de l'erreur d'inversion et implémentée via une stratégie "du grossier au fin" (coarse-to-fine) en deux étapes.

1. Insight Théorique :
   Les auteurs démontrent que l'erreur d'inversion induite par un saut de pas de temps $\Delta t$ peut être reformulée comme un problème de point fixe mis à l'échelle. L'amplitude de l'erreur $\delta$ est proportionnelle à un coefficient de mise à l'échelle $c_\alpha(t, \Delta t)$ multiplié par l'écart entre les prédictions consécutives du modèle ($\Delta \epsilon_\theta$).

   • Observation Clé : Le coefficient de mise à l'échelle présente une tendance parabolique par rapport à l'indice du pas de temps. Les erreurs sont concentrées aux petits pas de temps (étapes de diffusion précoces) et aux grands pas de temps (étapes tardives), avec un creux au milieu. De plus, des tailles de pas plus grandes ($\Delta t$) induisent intrinsèquement des erreurs plus importantes.

2. Rescaling Global (Ajustement Grossier) :
   Pour traiter la distribution parabolique de l'erreur, les auteurs introduisent un mécanisme de mise à l'échelle globale. Ils modifient la séquence originale de pas de temps uniformes $t_k$ en utilisant une transformation de loi de puissance contrôlée par un hyperparamètre $\gamma$ :
   $$t_k = t_1 + (t_K - t_1) \left( \frac{k-1}{K-1} \right)^\gamma$$

   • Lorsque $\gamma > 1$, les pas de temps sont étirés vers les étapes précoces (allouant plus de résolution à la région précoce à haute erreur).
   • Lorsque $\gamma < 1$, les pas de temps deviennent plus denses vers la fin.
     Cette étape prépare le programme pour une optimisation fine en ajustant globalement l'allocation des étapes.

3. Rescheduling Local via la Programmation Dynamique (Ajustement Fin) :
   Pour minimiser davantage l'erreur accumulée, les auteurs emploient une approche de programmation dynamique (DP) par fenêtre glissante.

   • Carte de Coût : Une carte de coût $E[k, t]$ est initialisée pour suivre l'erreur accumulée minimale pour atteindre l'étape $t$ à la $k$-ième étape d'inversion.
   • Récursion : Dans une fenêtre glissante de largeur $2d+1$ autour de l'étape actuelle, l'algorithme recherche récursivement la transition localement optimale entre les pas de temps adjacents qui minimise la somme de l'erreur précédente accumulée et du coefficient d'erreur de l'étape actuelle.
   • Backtracking : La séquence optimale des pas de temps est récupérée par un retour en arrière (backtracking) de l'étape finale vers la première.

L'ensemble du processus de reprogrammation est effectué hors ligne (ou comme étape de pré-traitement) sur la base du programme de bruit connu, ce qui signifie qu'il n'ajoute qu'une surcharge de calcul négligeable au processus réel d'inversion ou d'édition d'image.

Contributions Clés

• Formulation Théorique : Le papier dérive une connexion théorique entre les sauts de pas de temps et l'erreur d'inversion, formulant l'erreur comme un problème de point fixe mis à l'échelle dépendant du programme de bruit.
• Programmateur Adaptatif : Il propose un cadre de reprogrammation du grossier au fin (TRDI) qui alloue de manière adaptative des pas plus grands aux régions de faible erreur et des pas plus fins aux régions de haute erreur, équilibrant efficacement efficacité et précision.
• Amélioration Plug-and-Play : La méthode est conçue comme une amélioration "prête à l'emploi" qui ne nécessite aucun paramètre supplémentaire, entraînement ou surcharge de calcul significative, la rendant compatible avec les pipelines d'inversion existants.

Résultats Expérimentaux
Les auteurs valident TRDI via des expériences approfondies sur la reconstruction d'images (jeu de données MSCOCO) et l'édition d'images (PIE-Bench) en utilisant divers modèles de diffusion (Stable Diffusion v1.5, SDXL, et SDXL Turbo) et diverses méthodes d'inversion de référence (DDIM, ReNoise, NPI, GNRI).

• Reconstruction : Sur SD v1.5, l'intégration de TRDI a systématiquement amélioré les métriques pour toutes les méthodes de référence. Par exemple, avec la base DDIM, le PSNR est passé de 20,07 à 20,21, et l'LPIPS a diminué (améliorant la qualité perceptuelle) de 193,97 à 187,85. Les résultats qualitatifs ont montré des corrections dans les détails structurels comme les fenêtres de trains et les formes de fruits que les bases n'avaient pas réussi à reconstruire avec précision.
• Édition : Sur le PIE-Bench, TRDI a considérablement amélioré la préservation de la structure et de l'arrière-plan tout en maintenant une haute similitude CLIP. Notamment, pour les modèles à peu d'étapes comme SDXL Turbo (4 étapes), où la dégradation de l'arrière-plan est typiquement sévère, TRDI a obtenu des améliorations substantielles de la préservation de l'arrière-plan (par exemple, en réduisant la distance de structure de 17,43 % pour DDIM).
• Études d'Ablation : L'étude a confirmé qu'un facteur d'échelle $\gamma \approx 1,05$ et une largeur de fenêtre $d=8$ offraient le meilleur compromis entre fidélité structurelle et éditabilité sémantique.

Signification et Revendications
Le papier affirme que la programmation des pas de temps est un facteur critique, bien que précédemment négligé, de la fidélité de l'inversion de diffusion. En révélant la relation parabolique entre la taille/l'indice du pas de temps et l'accumulation d'erreur, les auteurs démontrent qu'une allocation stratégique et non uniforme des étapes de calcul peut réduire considérablement les erreurs d'inversion globales.

Les auteurs positionnent TRDI non pas comme un remplacement pour des formulations d'inversion plus puissantes (comme les itérations de point fixe ou l'inversion stochastique), mais comme un complément léger de type "rétrofit" au niveau du programmateur. Il permet aux pipelines d'inversion déterministes existants d'obtenir des résultats de reconstruction et d'édition supérieurs sans augmenter les coûts de calcul ni nécessifier de réentraînement du modèle. Ce travail souligne que optimiser quand effectuer les étapes est aussi important que de savoir comment les effectuer.