🤖 machine learning

TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment

Ce papier propose l'optimisation de politique par appariement de trajectoires (TMPO), un cadre d'apprentissage par renforcement novateur qui remplace la maximisation de récompense scalaire par un appariement de distributions au niveau des trajectoires via un objectif d'équilibre de trajectoires à softmax et un échantillonnage stochastique d'arbres dynamiques pour atténuer efficacement le piratage de récompense, améliorant ainsi considérablement la diversité et l'efficacité génératives dans l'alignement des modèles de diffusion.

Auteurs originaux : Jiaming Li, Chenyu Zhu, Zhiyuan Ma, Nanxi Yi, Youjun Bao, Li Sun, Quanying Lv, Xiang Fang, Daizong Liu, Jianjun Li, Kun He, Bowen Zhou

Publié 2026-05-13

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Jiaming Li, Chenyu Zhu, Zhiyuan Ma, Nanxi Yi, Youjun Bao, Li Sun, Quanying Lv, Xiang Fang, Daizong Liu, Jianjun Li, Kun He, Bowen Zhou

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous enseigniez à un artiste talentueux (un générateur d'images par IA) comment peindre en vous basant sur vos retours. Vous voulez que l'artiste crée de belles images correspondant à votre description, mais vous souhaitez également qu'il soit créatif et varié, et non pas qu'il peigne exactement la même chose encore et encore.

Ce papier, TMPO, présente une nouvelle méthode pour enseigner à cet artiste, résolvant un problème majeur : l'artiste a tendance à « tricher » en trouvant une astuce spécifique qui lui vaut un score élevé de votre part, puis il répète cette astuce à l'infini, ignorant toutes les autres bonnes possibilités.

Voici la décomposition des idées du papier à l'aide d'analogies simples :

1. Le Problème : Le « Cheval à un seul tour »

Les méthodes actuelles pour former ces artistes IA fonctionnent comme un enseignant strict qui ne se soucie que du score final.

Le Scénario : Vous demandez à l'artiste de « peindre un robot faisant du yoga ».
L'Ancienne Méthode (Maximisation de la Récompense) : L'artiste essaie quelques poses. Une pose obtient un score de 9/10. L'enseignant dit : « Super ! Refais celle-là ! » L'artiste réalise : « Si je fais juste cette pose de yoga spécifique, j'obtiens toujours un score élevé. » Ainsi, il arrête d'essayer d'autres poses. Il arrête de peindre des robots dans différents studios, avec différentes couleurs, ou faisant différents mouvements de yoga. Il peint juste ce seul robot, encore et encore.
Le Résultat : L'IA devient ennuyeuse. Elle crée un « effondrement de mode », où elle ne génère qu'un seul type d'image, alors qu'il existe des milliers d'autres façons valides et belles de dessiner un robot faisant du yoga. Elle commence aussi à « jouer » avec le système, ajoutant des détails étranges qui trompent le système de notation mais qui semblent mauvais aux yeux des humains.

2. La Solution : TMPO (Optimisation de Politique par Correspondance de Trajectoire)

Les auteurs proposent une nouvelle méthode d'enseignement appelée TMPO. Au lieu de simplement dire à l'artiste : « Fais la seule chose qui a obtenu le score le plus élevé », TMPO change l'objectif entièrement.

Le Nouvel Objectif : Au lieu de maximiser un score unique, TMPO demande à l'artiste de correspondre à la distribution des récompenses.
L'Analogie : Imaginez que l'enseignant possède un sac contenant différents résultats « bons ». Certains sont parfaits (10/10), certains sont très bons (8/10), et certains sont juste corrects (6/10).
- L'ancien enseignant disait : « Montre-moi uniquement les 10/10. »
- L'enseignant TMPO dit : « Je veux que tu peignes tous les bons résultats proportionnellement à leur qualité. S'il existe trois façons d'obtenir un 8/10, je veux voir les trois, pas juste une. »
Comment cela fonctionne : L'IA génère un entier « arbre » d'essais différents à la fois. Elle examine ensuite l'ensemble du groupe et dit : « D'accord, je dois m'assurer que mes chances de peindre ces différentes versions correspondent à la « bonté » de chaque version. » Cela force l'IA à continuer d'explorer différents styles et mises en page, préservant ainsi la diversité.

3. L'Ingrédient Secret : L'Astuce de l'« Arbre »

Entraîner une IA à examiner 27 versions différentes d'une image à la fois est généralement très lent et coûteux (comme demander à un peintre d'esquisser 27 tableaux complets avant de choisir le meilleur).

L'Innovation : TMPO utilise une technique appelée Échantillonnage Stochastique Dynamique d'Arbre.
L'Analogie : Imaginez que l'artiste commence à dessiner un arrière-plan (le « préfixe »). Au lieu de terminer 27 tableaux séparés à partir de zéro, il dessine un seul arrière-plan. Ensuite, à trois moments précis, il s'engage dans différentes directions.
- Branche 1 : « D'accord, faisons le robot bleu. »
- Branche 2 : « D'accord, faisons le robot rouge. »
- Branche 3 : « D'accord, faisons le robot vert. »
Parce qu'ils partagent l'arrière-plan initial, l'IA n'a pas à redessiner toute la scène 27 fois. Elle ne calcule que les différences aux « points de branchement ». Cela économise une quantité massive de temps (jusqu'à 27 % plus rapide dans leurs tests) tout en permettant toujours à l'IA d'explorer de nombreuses possibilités différentes.

4. Les Résultats : Plus de Variété, Moins de Tricherie

Le papier a testé cela sur un modèle IA populaire (FLUX.1) avec trois tâches différentes :

Génération Compositionnelle : S'assurer que les objets sont au bon endroit (par exemple, « un chat sur un tapis »).
Rendu de Texte : Écrire correctement des mots à l'intérieur de l'image.
Préférence Humaine : Créer des images qui plaisent aux gens.

Qu'est-il arrivé ?

Diversité : TMPO a produit des images 9,1 % plus diversifiées que les meilleures méthodes existantes. Les images semblaient différentes les unes des autres, plutôt que d'être des clones.
Qualité : Il n'a pas sacrifié la qualité pour la variété. Les images restaient précises et de haute qualité.
Efficacité : L'entraînement était plus rapide grâce à l'astuce de l'« Arbre ».

Résumé

Pensez à TMPO comme à un sage enseignant d'art qui réalise que si vous ne louez que la réponse « parfaite », l'élève arrêtera de penser de manière créative. Au lieu de cela, TMPO enseigne à l'IA à apprécier le spectre des bonnes réponses. En utilisant une structure en « arbre » pour explorer de nombreux chemins à la fois sans perdre de temps, elle crée une IA qui est à la fois intelligente (obtient des scores élevés) et créative (ne reste pas coincée dans une routine).

Le papier affirme que cela résout le problème du « piratage de récompense » (où l'IA triche avec le système) en prouvant mathématiquement que la correspondance de la distribution des récompenses force l'IA à couvrir toutes les possibilités « bonnes », et non pas juste la seule « meilleure ».

Résumé technique : TMPO pour l'alignement efficace et diversifié des modèles de diffusion

1. Énoncé du problème

L'apprentissage par renforcement (RL) a démontré un potentiel significatif dans l'alignement des modèles de diffusion et d'appariement de flux avec des tâches en aval. Cependant, les méthodes existantes reposent principalement sur la maximisation d'une récompense scalaire, que les auteurs identifient comme intrinsèquement rechercheuse de modes. Cette objectif d'optimisation amène les modèles à concentrer la masse de probabilité sur quelques trajectoires à haute récompense, conduisant à deux modes d'échec critiques :

Dégradation de la diversité générative : Le modèle s'effondre vers un ensemble restreint de compositions et de styles (effondrement de mode visuel), ignorant la distribution anisotrope des sorties plausibles apprise lors du pré-entraînement.
Amplification de récompenses peu fiables : En sur-renforçant des caractéristiques spécifiques qui servent de proxies aux récompenses (par exemple, la précision de l'OCR ou des scores esthétiques spécifiques), le modèle amplifie des attributs spurius, souvent au détriment de la qualité globale de la génération.

Les stratégies d'atténuation actuelles, telles que la régularisation KL, échouent à modifier la nature fondamentalement rechercheuse de modes de l'objectif, tandis que les approches existantes de réseaux de flux génératifs (GFlowNet) reposent souvent sur des estimations d'écoulement d'état intraitables ou sur des contraintes locales qui accumulent des erreurs sur de longs horizons de débruitage.

2. Méthodologie : Optimisation de la politique par appariement de trajectoires (TMPO)

TMPO modifie fondamentalement le paradigme d'optimisation, passant de la maximisation de récompense scalaire à l'appariement de la distribution de récompense au niveau des trajectoires.

2.1 Équilibre de trajectoire Softmax (Softmax-TB)

Au lieu de maximiser la récompense espérée, TMPO vise à faire correspondre les probabilités de trajectoire de la politique à une distribution cible induite par les récompenses.

Objectif : Pour un groupe de $K$ trajectoires générées à partir du même prompt, la méthode fait correspondre les probabilités de politique $P_\theta(\tau)$ à une distribution de Boltzmann $q(\tau) \propto \exp(\beta R(\tau))$ .
Formulation sans partition : En normalisant les probabilités et les récompenses au sein du groupe de $K$ trajectoires, la fonction de partition globale intractable est éliminée. La méthode définit un avantage Softmax-TB :
$A_i = \log q_i - \log p_i = \log \frac{\exp(\beta R_i)}{\sum_j \exp(\beta R_j)} - \log \frac{P_\theta(\tau_i)}{\sum_j P_\theta(\tau_j)}$
Garantie théorique : Les auteurs prouvent que la minimisation de cet objectif équivaut à minimiser la divergence KL directe ( $D_{KL}(q \parallel p)$ ) au sein du groupe. Contrairement au KL inverse (utilisé dans les gradients de politique standards), le KL direct est couvrant de modes : il impose une pénalité infinie si la probabilité de politique $p_i$ tend vers zéro pour toute trajectoire $i$ ayant un poids cible positif $q_i$ . Cela garantit théoriquement la préservation de la diversité sans nécessiter de régularisation KL auxiliaire.

2.2 Échantillonnage d'arbre stochastique dynamique

Pour répondre au coût computationnel de l'échantillonnage de multiples trajectoires indépendantes pour des modèles d'appariement de flux à grande échelle, TMPO introduit une stratégie d'échantillonnage en structure d'arbre :

Partage de préfixe : Les trajectoires partagent un préfixe ODE déterministe (étapes de débruitage communes) et se ramifient à des points stochastiques planifiés dynamiquement.
Ramification dynamique : Les points de ramification ne sont pas fixes ; ils suivent un curriculum qui évolue des régions à bruit élevé et précoces (diversité à grain grossier) vers les régions à bruit faible et tardives (variations à grain fin) au fur et à mesure de l'avancement de l'entraînement. Les positions de ramification sont perturbées stochastiquement à l'aide d'une distribution bêta pour éviter le surapprentissage à une géométrie d'arbre fixe.
Efficacité : Cette approche réduit le calcul redondant en partageant le passage avant pour le préfixe partagé, réduisant considérablement le temps d'entraînement tout en maintenant une exploration efficace.

3. Contributions clés

Identification de la cause racine : L'article identifie la nature rechercheuse de modes de la maximisation de récompense comme la cause principale du piratage de récompense et de la dégradation de la diversité dans le RL de diffusion, proposant un passage vers l'appariement de distributions.
Objectif Softmax-TB : Introduction d'un objectif d'appariement de distribution au niveau des trajectoires, sans partition, qui hérite des propriétés de couverture de modes de la divergence KL directe, préservant naturellement la diversité générative.
Échantillonnage d'arbre stochastique dynamique : Un mécanisme d'échantillonnage évolutif qui réduit le temps d'entraînement pour les modèles à grande échelle en partageant les préfixes de débruitage et en ajustant dynamiquement les calendriers de ramification.

4. Résultats expérimentaux

Les auteurs ont évalué TMPO sur l'architecture de base FLUX.1-dev à travers trois tâches d'alignement : génération d'images compositionnelles (GenEval), rendu de texte visuel (OCR) et alignement des préférences humaines (PickScore).

Diversité : TMPO a amélioré la diversité générative de 9,1 % en moyenne par rapport aux méthodes les plus avancées (Flow-GRPO, MixGRPO, TreeGRPO, GARDO). C'était la seule méthode à maintenir de manière cohérente des scores positifs de distance géométrique moyenne logarithmique (LGMD), indiquant une diversité saine, alors que les références présentaient souvent des LGMD négatifs (effondrement de mode).
Performance : TMPO a obtenu des performances compétitives ou supérieures sur des métriques spécifiques aux tâches (précision GenEval, précision OCR) et des scores de préférence humaine (HPS-v2.1, ImageReward, PickScore).
Efficacité : En utilisant le partage de préfixe et l'échantillonnage d'arbre, TMPO a réduit le temps d'entraînement par itération jusqu'à 27 % par rapport aux méthodes existantes.
Compromis : L'analyse de Pareto a démontré que TMPO atteint le compromis optimal entre maximisation de la récompense, préservation de la diversité et efficacité de l'entraînement.

5. Signification et affirmations

L'article affirme que TMPO offre une solution fondamentale au dilemme du piratage de récompense dans l'alignement de l'IA générative. En reformulant l'objectif pour qu'il corresponde aux distributions induites par les récompenses plutôt que de maximiser des récompenses scalaires, TMPO impose structurellement la préservation de la diversité sans dépendre de régularisation externe ou de rééquilibrage heuristique. La méthode démontre qu'il est possible d'obtenir des scores d'alignement élevés tout en maintenant la riche distribution de sortie multimodale apprise lors du pré-entraînement, le tout en réduisant la surcharge computationnelle généralement associée aux méthodes RL multi-trajectoires. Les auteurs positionnent TMPO comme un cadre évolutif adapté aux modèles d'appariement de flux à grande échelle, avec des extensions potentielles à la génération vidéo et à la génération d'actions robotiques.