🤖 machine learning

Beyond Mode-Seeking RL: Trajectory-Balance Post-Training for Diffusion Language Models

Ce papier présente TraFL, une méthode de post-entraînement par équilibre de trajectoire pour les modèles de langage par diffusion qui surmonte le mode d'échec de « verrouillage de trajectoire » des approches de maximisation de récompense en alignant la politique sur une distribution cible biaisée par la récompense, permettant ainsi d'obtenir des gains de performance cohérents sur les benchmarks de raisonnement mathématique et de génération de code.

Auteurs originaux : Saba Ahmadi, Prasanna Parthasarathi, Yufei Cui

Publié 2026-05-15

📖 6 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Saba Ahmadi, Prasanna Parthasarathi, Yufei Cui

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

La Vue d'Ensemble : Une Nouvelle Façon d'Enseigner à l'IA

Imaginez que vous avez un artiste très talentueux (le modèle d'IA) capable de peindre des tableaux en partant d'une toile remplie de bruit et de statique, puis en la nettoyant progressivement jusqu'à ce qu'une image claire apparaisse. C'est ainsi que fonctionnent les Modèles de Langage par Diffusion : ils n'écrivent pas les mots un par un comme une machine à écrire (ce qui est la méthode des anciens modèles d'IA) ; au lieu de cela, ils commencent par un enchevêtrement de lettres et les « débruitent » progressivement pour former une phrase cohérente.

Le papier soutient que la façon actuelle dont nous enseignons à ces artistes de s'améliorer pour résoudre des problèmes difficiles (comme les mathématiques ou la programmation) est défectueuse. Les auteurs proposent une nouvelle méthode appelée TraFL (Trajectory Flow baLancing) qui corrige un défaut spécifique appelé « Verrouillage de Trajectoire ».

Le Problème : Le Piège de la « Trajectoire Unique » (Verrouillage de Trajectoire)

Pour comprendre le problème, imaginez un labyrinthe.

L'Objectif : Le labyrinthe possède plusieurs sorties correctes différentes (différentes solutions valides à un problème mathématique).
L'Ancienne Méthode (RL Maximisant la Récompense) : Imaginez que vous entraînez un chien à trouver la sortie. Chaque fois que le chien trouve une sortie, vous lui donnez une friandise.
- Le Défaut : Le chien ne se soucie pas de quel chemin il a emprunté pour obtenir la friandise, seulement qu'il en a obtenu une.
- Le Résultat : Si le chien tombe par hasard sur un chemin spécifique menant à une friandise tôt dans le processus, il obtient une friandise. Il se souvient de ce chemin. La prochaine fois, il tente à nouveau ce chemin. Il obtient une autre friandise. Il devient plus confiant dans ce seul chemin.
- Verrouillage de Trajectoire : Finalement, le chien cesse d'explorer le labyrinthe entièrement. Il ne parcourt que ce chemin étroit, même s'il existe 50 autres sorties valides qu'il aurait pu trouver. Il s'est « verrouillé » sur une seule solution et a oublié comment trouver les autres.

Dans le papier, les auteurs appellent cela le Verrouillage de Trajectoire. Parce que l'IA ne reçoit une « récompense » (un score) que pour la réponse finale, elle ignore le fait qu'il existait de nombreuses façons différentes (chemins) pour y parvenir. Elle effondre toute sa créativité en une seule voie étroite, ce qui la rend mauvaise pour trouver des alternatives de réponses correctes lorsque vous lui demandez d'essayer à nouveau.

La Solution : L'Approche du « Guide » (TraFL)

Les auteurs proposent TraFL, qui modifie les règles d'entraînement. Au lieu de donner une friandise pour n'importe quelle sortie, ils donnent à l'IA un Guide (un modèle de référence figé) et une Carte.

Le Guide (Modèle de Référence) : C'est une version de l'IA qui n'a pas encore été entraînée sur les récompenses spécifiques. Elle représente une façon de penser « saine » et diversifiée. Elle sait qu'il existe de nombreuses façons de résoudre un problème.
La Carte (Cible Inclinée par la Récompense) : Nous disons à l'IA : « Tu veux trouver les sorties qui obtiennent la récompense (les bonnes réponses), MAIS tu dois garder tes schémas de mouvement similaires à ceux du Guide. »

L'Analogie :
Imaginez que vous enseignez à un élève à résoudre un problème de mathématiques.

Ancienne Façon : Vous dites : « Trouve la bonne réponse, peu importe comment ! » L'élève trouve un truc qui fonctionne, le mémorise et refuse d'apprendre toute autre méthode.
Façon TraFL : Vous dites : « Trouve la bonne réponse, mais garde ton processus de pensée diversifié et flexible, comme un excellent élève qui connaît de nombreuses stratégies différentes. »

TraFL force l'IA à équilibrer deux choses :

Obtenir la Récompense : Trouver la bonne réponse.
Rester Diversifié : Ne pas s'effondrer dans un seul chemin répétitif. Elle maintient la « masse de probabilité » (la probabilité de choisir un chemin) répartie sur de nombreuses solutions valides différentes, ancrée par le Guide.

Comment Ils Ont Fait Fonctionner Cela

Le papier note que les modèles de diffusion sont délicats car ils ne montrent pas leur « processus de pensée » étape par étape comme le font les anciens modèles. Pour corriger cela, les auteurs ont inventé deux outils :

Un Score de Substitut : Puisqu'ils ne peuvent pas voir les mathématiques exactes de chaque étape, ils ont créé un score « proxy » qui estime la qualité d'une réponse complète, leur permettant d'entraîner le modèle sans avoir besoin d'une visibilité parfaite sur chaque toute petite étape.
Un Normalisateur Appris : Ils ont enseigné à l'IA un « calculateur » spécial qui ajuste la difficulté de la tâche en fonction de la question spécifique (l'invite), garantissant que l'entraînement reste équitable et équilibré.

Les Résultats : Est-ce Que Cela Aide Vraiment ?

Les auteurs ont testé cette nouvelle méthode sur les Mathématiques (résolution de problèmes de mots) et le Code (écriture de programmes informatiques).

Le Piège de la « Trajectoire Unique » a été brisé : Contrairement à d'autres méthodes qui parfois s'amélioraient pour trouver une solution tout en s'aggravant pour trouver des solutions différentes, TraFL s'est amélioré dans chaque test individuel.
Plus d'Échantillons = Meilleurs Résultats : Lorsqu'ils ont demandé à l'IA de générer 16 réponses différentes au lieu d'une seule, TraFL s'est nettement amélioré. Cela prouve qu'elle a effectivement trouvé plus de solutions correctes différentes, et non pas juste une seule devinette chanceuse.
Cela Fonctionne sur Du Nouveau : L'IA n'a pas seulement mémorisé les questions d'entraînement. Lorsqu'elle a été testée sur de nouveaux problèmes de mathématiques (Minerva Math) et de nouveaux défis de programmation (LiveCodeBench) qu'elle n'avait jamais vus auparavant, TraFL a été le meilleur performant.
Vérification de la Diversité : Ils ont utilisé un « Juge » (une autre IA) pour examiner les réponses. Le Juge a confirmé que TraFL ne donnait pas simplement la même réponse avec des mots différents ; elle utilisait réellement différentes stratégies de raisonnement et différents algorithmes pour résoudre les mêmes problèmes.

Résumé

Le papier identifie un défaut où les modèles d'IA se « coincent » sur une seule façon de résoudre un problème, ignorant d'autres solutions valides. Ils ont corrigé cela avec TraFL, une méthode qui enseigne à l'IA à rechercher des réponses correctes tout en maintenant une « exploration » diversifiée de différents chemins, guidée par un modèle de référence. Le résultat est une IA non seulement plus intelligente, mais aussi plus créative et fiable lorsqu'on lui demande de générer plusieurs solutions.

Résumé technique : Post-entraînement Trajectory-Balance pour les modèles de langage diffusion

Problème : Verrouillage de trajectoire dans le RL diffusion

Les modèles de langage diffusion (dLLM) offrent une alternative prometteuse aux modèles autoregressifs pour le raisonnement et la génération de code. Cependant, les méthodes de post-entraînement existantes pour les dLLM adaptent largement des objectifs d'apprentissage par renforcement (RL) visant à maximiser la récompense (par exemple, PPO, GRPO) conçus pour les modèles autoregressifs. Les auteurs identifient un mode d'échec central dans cette adaptation, appelé verrouillage de trajectoire.

Dans les modèles diffusion, la récompense est attribuée uniquement à la complétion débruitée finale ( $y$ ), tandis que la distribution du modèle est induite sur l'ensemble de la trajectoire de débruitage ( $\tau$ ). Parce que les récompenses terminales sont « indifférentes au chemin » — c'est-à-dire que toute trajectoire aboutissant à la même $y$ correcte reçoit la même récompense — les mises à jour de gradient de politique échantillonnées tendent à renforcer les trajectoires spécifiques qui se trouvent être échantillonnées tôt dans l'entraînement. Cela crée une boucle de rétroaction auto-renforçante où la masse de probabilité s'effondre sur un ensemble restreint de chemins de débruitage. Par conséquent, même si un modèle apprend à produire une réponse correcte, il perd la capacité de générer un ensemble diversifié de solutions valides (modes) lors d'échantillonnages répétés. Les auteurs formalisent cela via le Théorème 1, prouvant que la couverture des modes terminaux est bornée supérieurement par la diversité des trajectoires ; ainsi, le verrouillage de trajectoire limite directement la couverture des solutions du modèle.

Méthodologie : TraFL (Trajectory Flow Balancing)

Pour remédier au verrouillage de trajectoire, les auteurs proposent TraFL, un objectif de post-entraînement fondé sur le principe d'équilibre de trajectoire des Réseaux de Flux Génératifs (GFlowNets). Au lieu de maximiser directement la récompense, TraFL entraîne la politique à correspondre à une distribution cible inclinée par la récompense ancrée à un modèle de référence figé.

L'objectif

La distribution cible $p^*(\tau | x)$ est définie comme suit :
$p^*(\tau | x) = \frac{1}{Z(x)} p_{\text{ref}}(\tau | x) \exp(\beta r(x, y))$
où :

$p_{\text{ref}}$ est un modèle de référence figé.
$r(x, y)$ est la récompense terminale.
$\beta$ contrôle la netteté de la récompense.
$Z(x)$ est un terme de normalisation dépendant de l'invite (fonction de partition).

L'objectif d'entraînement minimise le résidu carré entre le log-probabilité du modèle courant et la cible :
$\delta(\tau, x) = \log p_\theta(\tau | x) - \log p_{\text{ref}}(\tau | x) - \beta r(x, y) + \log Z(x)$
$L_{\text{TraFL}} = \mathbb{E}_{\tau \sim p_\theta} [\delta(\tau, x)^2]$

Mise en œuvre pratique pour la diffusion

Puisque les dLLM n'exposent pas de log-probabilités de trajectoire exactes, les auteurs introduisent deux composants clés pour rendre TraFL pratique :

Surrogé de niveau de séquence compatible diffusion : Ils dérivent un surrogé de log-probabilité traitable ( $\log \hat{p}_\theta(y|x)$ ) basé sur une borne inférieure de reconstruction masquée. Cela implique d'échantillonner un niveau de masquage $l$ , de corrompre $l$ jetons, et de moyenner les log-probabilités des jetons reconstruits. Un facteur $1/l$ est utilisé pour normaliser le score indépendamment du nombre de jetons masqués.
Normalisation apprise dépendante de l'invite : Contrairement aux méthodes qui estiment $Z(x)$ via un softmax sur un lot (ce qui peut induire une concentration), TraFL apprend une fonction de partition $Z_\phi(x)$ via un petit MLP. Ce terme reçoit des signaux de gradient à chaque étape, indépendamment des complétions échantillonnées, empêchant l'échec de normalisation « local au lot » observé dans d'autres approches de correspondance de distribution.

Contributions clés

Identification du verrouillage de trajectoire : L'article identifie formellement et prouve que les objectifs de maximisation de récompense dans les dLLM entraînent un effondrement de la diversité des trajectoires, ce qui limite à son tour la couverture des modes de solutions correctes distincts.
Objectif TraFL : La proposition d'un objectif d'équilibre de trajectoire qui ancre la politique à un modèle de référence, séparant le rôle de la récompense (inclinant les complétions terminales) de celui du modèle de référence (allouant la masse de probabilité sur des trajectoires plausibles).
Surrogé pratique et normalisation : Le développement d'un surrogé de niveau de séquence compatible diffusion et d'un terme de normalisation appris permettant un entraînement stable sans probabilités de trajectoire exactes.
Supériorité empirique : TraFL est la seule méthode de post-entraînement évaluée qui améliore constamment le modèle de base sur tous les réglages de longueur de benchmark (GSM8K, MATH-500, HumanEval, MBPP) et budgets d'échantillonnage.

Résultats

Les auteurs évaluent TraFL sur des benchmarks de raisonnement mathématique et de génération de code par rapport au modèle de base (LLaDA-8B-Instruct) et à des baselines robustes (ESPO, JUSTGRPO).

Performance en distribution : TraFL améliore le Pass@k sur tous les budgets d'échantillonnage ( $k \in \{1, 2, 4, 8, 16\}$ ) et températures. Notamment, il améliore le Pass@1 (précision sur échantillon unique) tout en s'adaptant mieux à l'augmentation des budgets d'échantillonnage, contrairement aux baselines qui sacrifient souvent la précision sur échantillon unique pour des gains multi-échantillons.
Généralisation hors échantillon : Les gains se transfèrent aux benchmarks hors distribution. Sur Minerva Math, TraFL reste au-dessus du modèle de base. Sur LiveCodeBench, TraFL est la méthode la plus performante sur toutes les divisions de difficulté (Facile, Moyen, Difficile), surpassant significativement les baselines sur les problèmes moyens et difficiles.
Analyse de diversité : En utilisant un protocole de juge basé sur un LLM, les auteurs montrent que les gains de TraFL sont associés à une couverture plus large de stratégies de solutions distinctes (par exemple, différentes approches de preuve ou décompositions algorithmiques) plutôt qu'à un simple affinage du score d'un seul mode. Les baselines comme ESPO, bien qu'améliorant la précision, ne montrent pas de gains de diversité significatifs par rapport au modèle de base, ce qui est cohérent avec l'hypothèse du verrouillage de trajectoire.

Importance et affirmations

L'article affirme que TraFL fournit une solution fondée sur des principes au dysfonctionnement de « recherche de mode » du post-entraînement RL standard pour les modèles diffusion. En passant de la maximisation de la récompense à l'équilibre de trajectoire, la méthode préserve la diversité des chemins de débruitage nécessaire pour couvrir plusieurs solutions valides.

Les auteurs déclarent que TraFL est la seule méthode évaluée qui améliore le modèle de base dans chaque réglage de longueur de benchmark testé. Ils soulignent que les améliorations sont robustes à l'augmentation des budgets d'échantillonnage et se généralisent aux jeux de données tenus de côté, suggérant que l'objectif d'équilibre de trajectoire découple efficacement le signal de récompense du mécanisme d'allocation de chemin. Le travail plaide pour des objectifs de correspondance de distribution plutôt que pour une pure maximisation de récompense dans le post-entraînement des dLLM, offrant un cadre pratique qui évite les modes d'échec de concentration inhérents aux approches standard de gradient de politique.

L'article reste modeste concernant les limites, notant que la méthode repose sur un surrogé de niveau de séquence approximatif plutôt que sur des probabilités de trajectoire exactes et utilise des récompenses binaires grossières (correspondance exacte ou réussite du test), ce qui peut limiter l'efficacité de l'apprentissage dans des scénarios avec des succès rares.