Diffusion Alignment as Variational Expectation-Maximization

Ce papier présente DAV, un cadre d'alignement des modèles de diffusion basé sur l'algorithme variationnel espérance-maximisation qui alterne entre une étape de recherche à l'exécution pour générer des échantillons diversifiés et une étape d'affinement du modèle, permettant d'optimiser les récompenses tout en évitant l'effondrement des modes sur des tâches continues et discrètes.

L'essentiel

🎨 Le Problème : L'Artiste qui perd son âme

Imaginez un artiste génial (le modèle de diffusion) qui sait peindre des paysages magnifiques, des portraits réalistes et des scènes de la vie quotidienne. C'est son "entraînement de base". Il est diversifié, créatif et naturel.

Mais un jour, un directeur de musée (l'objectif de récompense) arrive et dit : "Je veux que tu peignes uniquement des tableaux avec des couleurs très vives et des formes géométriques parfaites !"

Si on force l'artiste à obéir aveuglément :

1. L'Over-optimisation (La dérive) : L'artiste commence à peindre seulement des triangles rouges vifs. Il oublie comment peindre des arbres, des visages ou des nuages. Tout devient bizarre et répétitif. C'est ce qu'on appelle l'effondrement de la diversité.
2. La fragilité : Si le directeur change d'avis et demande des cercles bleus, l'artiste est perdu car il a oublié ses bases.

Les méthodes actuelles pour "aligner" ces modèles (les forcer à suivre un objectif) souffrent souvent de ce problème : elles obtiennent un bon résultat sur le critère demandé, mais tuent la créativité et la diversité de l'œuvre.

💡 La Solution : DAV, le Cycle "Exploration-Entraînement"

Les auteurs proposent une nouvelle méthode appelée DAV. Imaginez que ce n'est pas un simple ordre, mais un processus en deux temps, répété comme un cycle de perfectionnement, inspiré d'une vieille technique mathématique appelée "Espérance-Maximisation" (EM).

On peut le comparer à un Chef Cuisinier qui veut créer le plat parfait selon un goût très spécifique, sans perdre la saveur de ses ingrédients de base.

Étape 1 : L'Exploration (Le "E-step") 🧭

• Le concept : Avant de changer la recette, le Chef sort en cuisine pour tester des milliers de variations.
• L'analogie : C'est comme si le Chef utilisait un GPS intelligent (une recherche au moment de l'exécution) pour explorer la cuisine. Il ne se contente pas de suivre la recette de base. Il essaie des combinaisons audacieuses, guidé par un "radar de récompense" qui lui dit : "Oups, ce mélange est trop salé, mais celui-ci est délicieux !".
• Le but : Trouver les meilleurs plats possibles (les échantillons à haute récompense) tout en essayant de garder une grande variété de saveurs. Il ne choisit pas juste un seul plat, mais un panier plein de bonnes idées.

Étape 2 : L'Amortissement (Le "M-step") 📚

• Le concept : Maintenant que le Chef a goûté à ces excellents plats, il met à jour son livre de recettes officiel.
• L'analogie : Il ne se contente pas de copier un seul plat. Il analyse tout le panier d'échantillons trouvés à l'étape 1 et réécrit son livre de recettes pour que, la prochaine fois, il puisse reproduire ces saveurs délicieuses naturellement, sans avoir besoin du GPS à chaque fois.
• La différence clé : Contrairement aux anciennes méthodes qui se focalisaient sur un seul type de plat (ce qui rendait le livre de recettes rigide), DAV apprend à couvrir toutes les bonnes variations trouvées. Il apprend à être bon dans plusieurs styles à la fois.

🔄 Le Cycle Magique

Ce processus se répète :

1. Le Chef utilise son livre de recettes actuel pour explorer de nouvelles idées (grâce au GPS).
2. Il trouve de nouvelles pépites.
3. Il met à jour son livre de recettes pour intégrer ces pépites.
4. Avec un livre de recettes amélioré, il peut explorer encore mieux la prochaine fois.

C'est un cercle vertueux : le modèle devient meilleur pour trouver ce qu'on veut, tout en restant capable de faire des choses variées et naturelles.

🌍 Pourquoi c'est génial ? (Les Résultats)

Les auteurs ont testé cette méthode sur deux domaines très différents :

1. La Création d'Images (Texte vers Image) :

   • Le défi : Faire des images plus belles (critère esthétique) sans qu'elles deviennent toutes identiques ou bizarres.
   • Le résultat : DAV a créé des images plus belles que les autres méthodes, mais surtout, elles sont restées variées et naturelles. Les autres méthodes ont produit des images "collées" et répétitives.

2. La Biologie (Design d'ADN) :

   • Le défi : Concevoir des séquences d'ADN qui fonctionnent bien (pour activer des gènes) sans créer de "monstres" biologiques qui n'existent pas dans la nature.
   • Le résultat : DAV a conçu des séquences d'ADN très efficaces qui ressemblaient encore à de l'ADN naturel, alors que les autres méthodes créaient des séquences trop "optimisées" et donc biologiquement invalides.

🏆 En Résumé

DAV, c'est comme donner à un artiste un compas pour explorer les meilleures idées, puis lui demander de mettre à jour son style en se basant sur tout ce qu'il a découvert, plutôt que de le forcer à répéter un seul mouvement.

C'est une méthode qui réussit à avoir le beurre et l'argent du beurre : des résultats excellents sur l'objectif visé, tout en préservant la richesse, la diversité et la beauté naturelle de la création.

Résumé technique

1. Problématique

Les modèles de diffusion sont excellents pour générer des échantillons de haute fidélité, mais leur alignement avec des objectifs externes (comme la qualité esthétique d'une image ou l'activité biologique d'une séquence d'ADN) pose des défis majeurs. Les méthodes d'alignement existantes se divisent généralement en deux catégories, chacune présentant des défauts critiques :

• L'ajustement fin par Reinforcement Learning (RL) : Souvent basé sur l'optimisation de la divergence KL inverse, ces méthodes souffrent d'un comportement de "recherche de mode" (mode-seeking). Cela conduit à une convergence prématurée, un effondrement des modes (mode collapse) et une perte de diversité dans les échantillons générés.
• La rétropropagation directe (Direct Backpropagation) : Bien que plus efficace en termes d'échantillons, cette méthode dépend de signaux de gradient provenant de fonctions de récompense apprises. Elle est sujette à une sur-optimisation (over-optimization), où le modèle exploite des failles de la fonction de récompense au détriment de la qualité naturelle et de la diversité des données.

Il existe donc un besoin urgent d'un cadre d'ajustement fin capable de maximiser les récompenses tout en préservant la diversité et la naturalité du modèle pré-entraîné, applicable aussi bien aux données continues (images) qu'aux données discrètes (séquences d'ADN).

2. Méthodologie : DAV (Diffusion Alignment as Variational EM)

Les auteurs proposent DAV, un cadre qui reformule l'alignement des modèles de diffusion comme un processus itératif basé sur l'algorithme Variational Expectation-Maximization (EM). Ce cadre alterne entre deux phases complémentaires :

A. Étape E (Exploration) : Recherche au moment de l'inférence

L'objectif est de découvrir des trajectoires (échantillons) diversifiés et à haute récompense à partir de la distribution postérieure variationnelle.

• Au lieu d'échantillonner directement depuis le modèle actuel, DAV utilise une recherche au moment de l'inférence (test-time search).
• Cette recherche est guidée par une fonction Q douce (soft Q-function) estimée pour identifier les régions prometteuses de l'espace d'état.
• La méthode utilise des techniques telles que la guidance par gradient (si la récompense est différentiable) et l'échantillonnage par importance pour affiner les échantillons et s'approcher de la distribution postérieure optimale $\eta^*$.
• Cette étape permet d'explorer efficacement la structure multimodale de la distribution, évitant ainsi de se coincer dans un seul mode dominant.

B. Étape M (Amortissement) : Mise à jour du modèle

L'objectif est d'incorporer les connaissances acquises lors de l'étape E dans les paramètres du modèle de diffusion.

• Le modèle est mis à jour en minimisant la divergence KL directe (forward KL divergence) entre la distribution des trajectoires découvertes ($\eta^*$) et la distribution du modèle ($p_\theta$).
• Contrairement aux méthodes RL classiques qui minimisent la divergence KL inverse (favorisant un seul mode), la minimisation de la divergence KL directe est une objectif de couverture de modes (mode-covering). Elle encourage le modèle à couvrir tous les modes divers découverts lors de l'étape E.
• Une variante, DAV-KL, ajoute un terme de régularisation KL par rapport à la politique pré-entraînée initiale pour éviter la perte de capacités du modèle original.

Ce cycle itératif crée une synergie : l'étape M affine le modèle vers une distribution alignée et multimodale, ce qui permet à l'étape E suivante de collecter des échantillons encore mieux alignés tout en préservant la diversité.

3. Contributions Clés

1. Cadre Unifié DAV : Introduction d'un algorithme d'alignement basé sur l'EM variationnel qui sépare l'exploration (recherche) de l'amortissement (mise à jour du modèle).
2. Gestion de la Diversité : Utilisation de la minimisation de la divergence KL directe dans l'étape M pour prévenir l'effondrement des modes, un problème récurrent dans les approches RL.
3. Généralité Modale : Le cadre est conçu pour fonctionner nativement avec des modèles de diffusion continus (synthèse d'images) et discrets (conception de séquences d'ADN), sans hypothèse forte sur la différentiabilité de la fonction de récompense (grâce à la recherche au moment de l'inférence).
4. Robustesse à la Sur-optimisation : La méthode atténue le phénomène de sur-optimisation en évitant de dépendre uniquement de gradients fragiles ou de politiques de recherche de mode.

4. Résultats Expérimentaux

Les auteurs ont évalué DAV sur deux tâches distinctes :

• Synthèse Image-to-Text (Continu) :

  • Base : Stable Diffusion v1.5.
  • Objectifs : Maximiser les scores esthétiques, ImageReward et la compressibilité.
  • Résultats : DAV et DAV-KL surpassent les méthodes de base (DDPO, DRaFT, TDPO). Ils atteignent des scores de récompense plus élevés tout en maintenant des scores de diversité (LPIPS) et de cohérence sémantique (CLIPScore/ImageReward) bien supérieurs. Les méthodes de base souffrent d'une dégradation rapide de la diversité et de la qualité humaine (sur-optimisation).
  • Observation : La version "Postérieure" de DAV (recherche au moment de l'inférence sans mise à jour du modèle) obtient les meilleurs scores absolus, confirmant la qualité de l'exploration.

• Conception de Séquences d'ADN (Discret) :

  • Base : Modèle de diffusion masqué (Masked Diffusion Language Model).
  • Objectifs : Concevoir des enhanceurs d'ADN avec une haute activité biologique prédite.
  • Résultats : DAV dépasse les méthodes d'alignement discrètes (DRAKES, VIDD, DDPO). Il parvient à un meilleur compromis entre l'activité prédite (récompense), la validité biologique (ATAC-Acc), la naturalité (corrélation 3-mer) et la diversité (distance de Levenshtein). Les méthodes RL classiques montrent une chute drastique de la validité biologique et de la diversité.

5. Signification et Impact

Ce travail est significatif car il propose une alternative théoriquement fondée et empiriquement robuste aux approches d'alignement actuelles.

• Théorique : Il établit un lien clair entre l'alignement des modèles de diffusion et l'algorithme EM, justifiant l'utilisation de la divergence KL directe pour la couverture de modes.
• Pratique : Il offre une solution viable pour des applications critiques où la diversité est aussi importante que la performance (ex: découverte de médicaments, génération de contenu créatif), là où les méthodes actuelles échouent souvent en produisant des résultats répétitifs ou biologiquement inviables.
• Généralité : La capacité à traiter à la fois des espaces continus et discrets, ainsi que des récompenses non différentiables, fait de DAV un cadre d'alignement universel pour les modèles génératifs de nouvelle génération.

En résumé, DAV réussit à briser le compromis traditionnel entre l'optimisation de la récompense et la préservation de la diversité, offrant une voie prometteuse pour l'alignement robuste des modèles de diffusion.