Momentum Guidance: Plug-and-Play Guidance for Flow Models

Ce papier présente la Guidance par Momentum (MG), une méthode plug-and-play pour les modèles de flux qui améliore la qualité des échantillons générés en extrapolant la vitesse via une moyenne mobile exponentielle, sans augmenter le coût d'inférence ni réduire la diversité, tout en surpassant les performances existantes sur des benchmarks comme ImageNet-256 et des modèles de pointe tels que Stable Diffusion 3.

L'essentiel

Imaginez que vous essayez de dessiner un portrait très détaillé à partir d'une idée vague dans votre tête. C'est ce que font les modèles de génération d'images (comme ceux qui créent des images à partir de texte).

Le problème, c'est que ces "artistes numériques" ont tendance à être un peu trop prudents. Pour éviter de faire une erreur, ils lissent trop leurs traits. Le résultat ? Des images qui ressemblent à des photos floues, où les détails fins (comme les cils, les textures de la peau ou les feuilles d'un arbre) sont effacés, comme si quelqu'un avait passé un linge humide sur le dessin.

Voici comment les chercheurs ont résolu ce problème avec une technique appelée Momentum Guidance (ou "Guidage par l'Élan").

1. Le problème : L'artiste qui a peur de se tromper

Les modèles actuels fonctionnent comme un voyageur qui avance pas à pas vers une destination (l'image finale). À chaque pas, le modèle regarde la direction à prendre. Mais comme il a été entraîné à être "moyen" pour ne pas faire d'erreur, il a tendance à choisir la direction la plus sûre, ce qui lisse trop le chemin. L'image devient floue.

2. La solution classique (CFG) : Demander à deux personnes

Pour corriger cela, on utilise souvent une technique appelée "Guidage sans classifieur" (CFG). Imaginez que vous demandez à deux personnes de vous guider :

• Une personne vous dit : "Va vers l'endroit exact où tu veux aller !" (La version précise).
• L'autre dit : "Va vers n'importe où, tant que c'est vague." (La version floue).
• Le guide prend la différence entre les deux et vous pousse plus fort vers la direction précise.

Le hic ? Cela demande de faire le calcul deux fois à chaque pas. C'est comme si vous deviez payer deux fois le prix du billet pour le même voyage. C'est lent et coûteux en énergie.

3. La nouvelle idée : Le "Momentum Guidance" (MG)

C'est ici que l'astuce de ce papier devient géniale. Au lieu de faire appel à une deuxième personne (ou un deuxième modèle), l'artiste utilise son propre souvenir.

Imaginez que vous marchez dans un brouillard. Au lieu de regarder seulement où vous êtes maintenant, vous vous souvenez de la direction que vous aviez il y a un instant.

• L'idée clé : Les pas précédents (quand l'image était encore très floue) contiennent déjà une version "lissée" et sûre de la trajectoire.
• Le mécanisme : Le modèle se dit : "Attends, je suis en train de faire ce mouvement, mais si je regarde où j'étais il y a un instant, je vois que je suis en train de trop me calmer. Je vais donc ajouter un petit coup de pouce dans la direction opposée à mon 'souvenir' pour me rendre plus vif."

C'est comme l'élan d'un skieur. Si vous glissez trop doucement sur la neige, vous ne progresserez pas. Mais si vous vous souvenez de votre vitesse précédente et que vous vous donnez un petit coup de bâton pour accélérer, vous gagnez en vitesse et en précision sans avoir besoin de quelqu'un d'autre pour vous pousser.

Pourquoi c'est génial ?

1. C'est gratuit (en temps de calcul) : Le modèle n'a pas besoin de faire un calcul supplémentaire. Il réutilise simplement les informations qu'il a déjà calculées à l'instant d'avant. C'est comme si vous obteniez une image HD sans payer le supplément "qualité supérieure".
2. C'est simple : On peut ajouter cette fonctionnalité à n'importe quel modèle existant sans le réentraîner. C'est un "plug-and-play" (brancher et jouer).
3. Le résultat : Les images deviennent nettes, les détails apparaissent (les reflets sur une voiture, les plis d'un vêtement), et l'image a plus de "vie", tout en restant fidèle à la demande.

En résumé

Si la génération d'images était une course, les modèles actuels courraient prudemment pour ne pas tomber, ce qui rendait le trajet lent et flou.

• L'ancienne méthode (CFG) consistait à courir deux fois plus vite en doublant l'effort.
• Momentum Guidance, c'est comme apprendre au coureur à utiliser son propre élan pour accélérer naturellement. Il arrive à la ligne d'arrivée plus vite, avec une image plus nette, et sans avoir besoin de courir deux fois plus fort.

C'est une petite astuce mathématique qui permet d'obtenir de superbes images avec moins d'effort, un peu comme trouver le secret pour faire un gâteau parfait sans avoir besoin d'un four plus puissant.

Résumé technique

1. Problématique

Les modèles génératifs basés sur le flux (Flow-based models), y compris les modèles de flux rectifiés (Rectified Flow) et les modèles de diffusion, sont devenus des références pour la synthèse d'images de haute qualité. Cependant, une limitation pratique majeure persiste : les échantillons générés directement par ces modèles pré-entraînés (sans guidage) apparaissent souvent diffus, manquant de détails fins et présentant des textures floues.

Ce phénomène est attribué à l'effet de lissage inhérent aux réseaux de neurones et à l'utilisation courante de la moyenne mobile exponentielle (EMA) des paramètres du modèle pendant l'entraînement, qui tend à supprimer les structures haute fréquence.

Pour corriger cela, des techniques de guidage comme le Classifier-Free Guidance (CFG) sont utilisées. Le CFG améliore la fidélité en extrapolaant la prédiction conditionnelle loin d'une prédiction inconditionnelle (plus lisse). Toutefois, cette méthode présente deux inconvénients majeurs :

1. Coût de calcul doublé : Elle nécessite deux passes avant (forward passes) par étape d'échantillonnage (une conditionnelle, une inconditionnelle).
2. Perte de diversité : L'augmentation de la force de guidage améliore la netteté mais réduit souvent la diversité des échantillons (baisse du Recall).

Des méthodes alternatives comme l'Autoguidance nécessitent des checkpoints supplémentaires ou des réseaux auxiliaires, ce qui les rend peu pratiques pour les grands modèles ouverts.

2. Méthodologie : Momentum Guidance (MG)

Les auteurs proposent une nouvelle méthode de guidage appelée Momentum Guidance (MG). Cette technique exploite directement la trajectoire de l'Équation Différentielle Ordinaire (ODE) du modèle, sans nécessiter de modèles supplémentaires ni de passes de calcul supplémentaires.

Principe de base

L'idée centrale repose sur l'observation que les vitesses calculées aux étapes antérieures du processus de débruitage (plus bruyantes) correspondent à des distributions marginales plus lisses. Au lieu de calculer une vitesse inconditionnelle coûteuse, MG utilise une moyenne mobile exponentielle (EMA) des vitesses passées pour créer une référence lisse.

Algorithme

Le processus d'inférence modifié fonctionne comme suit :

1. Initialisation : On échantillonne un bruit initial $Z_0$ et on initialise le momentum $m_{t_0}$ avec la vitesse initiale $v_\theta(Z_0, t_0)$.
2. Boucle d'échantillonnage : À chaque étape $t_i$ :
   • On calcule la vitesse actuelle $v_{t_i} = v_\theta(Z_{t_i}, t_i)$.
   • On met à jour le momentum (EMA) :
     $$m_{t_{i+1}} = (1 - \beta) v_{t_i} + \beta m_{t_i}$$
     où $\beta$ contrôle la décroissance des vitesses historiques.
   • On met à jour l'échantillon en extrapolant la vitesse actuelle par rapport au momentum :
     $$Z_{t_{i+1}} = Z_{t_i} + \Delta t \left[ v_{t_i} + \alpha (v_{t_i} - m_{t_i}) \right]$$
     où $\alpha$ est le poids du guidage (force d'extrapolation).

Avantages techniques

• Coût nul supplémentaire : MG réutilise la vitesse déjà calculée pour l'étape courante. Le nombre d'évaluations de réseau (NFE) reste identique à celui d'un échantillonnage standard (1 évaluation par étape).
• Compatibilité : MG fonctionne indépendamment ou en combinaison avec le CFG.
• Mémoire négligeable : Seule la vectorisation du momentum (de même dimension que l'état latent) doit être stockée, ce qui est minime par rapport aux paramètres du modèle.

3. Contributions Clés

1. Nouveau paradigme de guidage : Introduction d'une méthode qui utilise l'historique des vitesses du modèle lui-même comme signal de référence lisse, éliminant le besoin de branches inconditionnelles ou de modèles auxiliaires.
2. Efficacité computationnelle : Réduction du coût d'inférence de 50 % par rapport au CFG standard (car MG ne nécessite qu'une seule passe de réseau, contrairement aux deux du CFG).
3. Amélioration de la qualité et de la diversité : MG permet d'obtenir des images plus nettes et structurées tout en préservant (voire en améliorant) la diversité des échantillons, contrairement au CFG qui tend à réduire le Recall à forte intensité.
4. Plug-and-Play : La méthode est simple à intégrer dans n'importe quel pipeline de flux existant sans réentraînement du modèle.

4. Résultats Expérimentaux

Les expériences ont été menées sur plusieurs benchmarks et modèles de pointe :

• ImageNet-256 (Rectified Flow) :

  • Sans CFG : MG améliore le FID (Fréchet Inception Distance) de 36,68 % en moyenne par rapport à l'échantillonnage de base.
  • Avec CFG : MG améliore le FID de 25,52 % par rapport au CFG seul.
  • Meilleur résultat : Un FID de 1,597 atteint en 64 étapes d'échantillonnage.
  • MG maintient un meilleur compromis Précision-Rappel (Precision-Recall) que le CFG, évitant l'effondrement de la diversité observé avec des guidages forts.

• Modèles à grande échelle (Text-to-Image) :

  • Tests sur Stable Diffusion 3 (SD3) et FLUX.1-dev.
  • MG démontre des gains constants sur les métriques de préférence humaine (HPSv2.1) et de récompense (ImageReward).
  • Qualitativement, MG produit des images avec des contours plus nets, des textures plus riches (ex: détails des ailes d'anges, structures de corail), moins d'artefacts de flou et une géométrie d'objets plus stable, même à des niveaux de CFG élevés où le CFG seul produit des images sursaturées ou déformées.

5. Signification et Conclusion

Momentum Guidance représente une avancée significative pour l'inférence des modèles génératifs basés sur le flux. En résolvant le compromis classique entre la qualité de l'image, la diversité et le coût de calcul, MG offre une solution pratique et évolutive.

• Impact économique : En éliminant la nécessité de doubles passes de calcul, MG permet de réduire considérablement les coûts d'inférence pour les applications commerciales et la recherche.
• Qualité visuelle : Elle permet d'exploiter pleinement le potentiel des modèles pré-entraînés en "dé-lissant" leurs prédictions sans sacrifier la richesse des détails.
• Futur : La méthode est compatible avec d'autres techniques de guidage (comme CFG++ ou ADG), suggérant qu'elle peut être combinée pour des gains encore plus importants.

En résumé, MG transforme la trajectoire de l'ODE elle-même en un outil de guidage intelligent, offrant une amélioration de la qualité "plug-and-play" sans pénalité de performance.