d2: Improving Reasoning in Diffusion Language Models via Trajectory Likelihood Estimation

L'article présente d2, un nouveau cadre d'apprentissage par renforcement pour les modèles de langage de diffusion masqués qui utilise des estimateurs de vraisemblance de trajectoire spécialisés (d2-AnyOrder et d2-StepMerge) pour améliorer considérablement les capacités de raisonnement sur les benchmarks logiques et mathématiques, atteignant de nouvelles performances de l'état de l'art.

L'essentiel

Imaginez que vous avez un artiste très talentueux capable de peindre un chef-d'œuvre en partant d'une toile vierge et en ajoutant progressivement des détails jusqu'à ce que l'image soit complète. C'est ainsi que fonctionnent les Modèles de Langage de Diffusion (DLM). Contrairement aux écrivains IA traditionnels qui construisent des phrases mot après mot (comme l'empilement de briques), ces modèles partent d'un fouillis de « bruit » (comme de la neige sur une vieille télévision) et le « débruitent » lentement, révélant une histoire ou une solution cohérente étape par étape.

Le papier présente une nouvelle méthode d'entraînement appelée d2 pour apprendre à ces artistes comment résoudre des énigmes difficiles (comme des problèmes de mathématiques ou des jeux de logique) bien mieux. Voici la décomposition de son fonctionnement, en utilisant des analogies simples.

Le Problème : Le Critique « Aveugle »

Pour apprendre à une IA à raisonner, les chercheurs utilisent l'Apprentissage par Renforcement (RL). Voyez cela comme un jeu où l'IA tente de résoudre une énigme, et un « Critique » (un système de récompense) lui donne un score. Si l'IA obtient un bon score, elle apprend à refaire cela.

Cependant, il y a un piège avec les modèles de diffusion. Pour apprendre efficacement, l'IA doit savoir exactement à quel point la séquence spécifique de mots qu'elle a générée était probable.

• Pour l'IA traditionnelle (Autorégressive) : C'est comme écrire une lettre. Vous savez exactement ce que vous avez écrit en premier, en deuxième, en troisième. Calculer la « probabilité » est facile et rapide.
• Pour l'IA de Diffusion : C'est comme regarder une sculpture émerger d'un bloc de pierre. L'IA effectue des milliers de petits ajustements. Pour connaître la probabilité exacte de la statue finale, il faudrait théoriquement rembobiner la vidéo et rejouer chaque image du processus de sculpture. Faire cela pour chaque étape d'entraînement est informatiquement impossible (trop lent et trop coûteux).

Parce qu'ils ne pouvaient pas calculer ce « score » avec précision, les méthodes précédentes faisaient des estimations, ce qui entraînait de faibles capacités de raisonnement.

La Solution : d2 (Le Scorekeeper Intelligent)

Les auteurs ont créé d2, un cadre qui agit comme un scorekeeper super efficace. Il détermine la « probabilité » du parcours de l'IA sans avoir à rejouer tout le film à chaque fois. Ils ont construit deux outils différents pour deux types d'artistes différents :

Outil 1 : d2-AnyOrder (Le « Miroir Magique »)

Certains modèles de diffusion sont spéciaux ; ils peuvent révéler l'image finale dans n'importe quel ordre. Vous pourriez peindre les yeux d'abord, puis le nez, puis l'arrière-plan, ou l'inverse, et cela fonctionne toujours.

• L'Analogie : Imaginez que vous avez un miroir magique. Au lieu de regarder la peinture se faire étape par étape, vous pouvez regarder la peinture finie et voir instantanément ce que l'artiste aurait dit à chaque étape, tout d'un coup.
• Le Résultat : Cet outil, appelé d2-AnyOrder, calcule le score parfait en un seul regard (un seul passage informatique). Il est incroyablement précis et permet à l'IA d'apprendre très vite.

Outil 2 : d2-StepMerge (La « Stratégie de Fractionnement »)

La plupart des modèles de diffusion populaires (comme le célèbre LLaDA) ne peuvent pas peindre dans n'importe quel ordre ; ils doivent suivre une séquence spécifique. Ils n'ont pas le « miroir magique ».

• L'Analogie : Puisque nous ne pouvons pas voir tout le film d'un coup, nous divisons le film en segments. Au lieu de regarder chaque image, nous regardons le film par segments de 10 secondes. Nous estimons le score pour chaque segment et les additionnons.
• Le Résultat : Cet outil, appelé d2-StepMerge, est une approximation. Ce n'est pas parfait, mais les auteurs ont prouvé mathématiquement que l'erreur est faible et contrôlable. Il échange une petite partie de précision contre un gain massif de vitesse, ce qui le rend pratique pour les modèles standards.

Les Résultats : De Novice à Grand Maître

Les chercheurs ont testé ces outils sur certaines des énigmes les plus difficiles pour l'IA :

1. Jeux de Logique : Sudoku et Countdown (un jeu mathématique).
2. Problèmes de Mathématiques : GSM8K et MATH500 (benchmarks standards pour le raisonnement mathématique).

Le Résultat :

• Sans aucune « fiche de triche » supplémentaire (Fine-tuning supervisé) : Le cadre d2 a appris aux modèles à mieux raisonner que toute méthode précédente.
• Battre la Compétition : Sur les tests de Sudoku et Countdown, les modèles entraînés par d2 ont atteint des scores proches de 92 % et 56 % respectivement, écrasant les meilleures méthodes précédentes qui étaient bloquées autour de 22 % et 42 %.
• Maîtrise des Mathématiques : Sur les benchmarks mathématiques complexes, d2 a établi un nouveau record de « State-of-the-Art », prouvant que les modèles de diffusion peuvent être aussi bons que les modèles de génération de texte traditionnels pour le raisonnement, avec les avantages supplémentaires d'être plus rapides et plus contrôlables.

Résumé

Le papier affirme : « Nous avons trouvé un moyen d'apprendre aux modèles d'IA de diffusion comment penser clairement. »

• Ils ont réalisé que l'ancienne façon de noter le travail de l'IA était trop lente et imprécise.
• Ils ont inventé d2, un nouveau système de notation qui est soit parfaitement précis (pour les modèles spéciaux), soit intelligemment efficace (pour les modèles standards).
• En utilisant ce nouveau système, les modèles d'IA ont appris à résoudre des énigmes logiques et mathématiques bien mieux qu'auparavant, sans avoir besoin qu'on leur pré-enseigne les réponses.

En bref, d2 donne aux modèles de diffusion la capacité de « regarder en arrière » sur leur propre processus de pensée de manière précise, leur permettant d'apprendre de leurs erreurs et de devenir de véritables experts en raisonnement.

Résumé technique

Résumé technique : d2 : Améliorer le raisonnement des modèles de langage par diffusion via l'estimation de la vraisemblance de trajectoire

Énoncé du problème

Les modèles de langage par diffusion (DLM), en particulier les DLM masqués, sont apparus comme des alternatives compétitives aux modèles autorégressifs (AR) en raison de leur contrôlabilité et de leurs capacités de génération parallèle. Cependant, l'amélioration de leurs capacités de raisonnement via l'apprentissage par renforcement (RL) reste un défi. Si le RL est devenu la norme pour induire le raisonnement dans les modèles AR, l'appliquer aux DLM est non trivial. La difficulté centrale réside dans la formulation du gradient de politique : contrairement aux modèles AR où les vraisemblances se factorisent proprement sur les positions de jetons, la vraisemblance exacte d'une trajectoire de diffusion est de calcul intractable. Le calcul naïf de ces vraisemblances nécessite $T$ passages avant (où $T$ est le nombre d'étapes de diffusion), rendant les méthodes de gradient de politique standard comme le Group Relative Policy Optimization (GRPO) extrêmement coûteuses en termes de calcul. Les approximations existantes introduisent souvent un biais significatif, menant à des mises à jour de politique sous-optimales.

Méthodologie

Les auteurs introduisent d2, un cadre de raisonnement adapté aux DLM masqués. Le cadre repose sur un nouvel algorithme de gradient de politique dérivé de l'objectif GRPO, qui nécessite explicitement une estimation précise des vraisemblances de trajectoire d'échantillonnage. Pour lever le goulot d'étranglement computationnel, l'article propose une famille d'estimateurs adaptés à différentes classes de modèles :

1. Fondement théorique

Les auteurs dérivent un objectif de type GRPO pour les DLM masqués en marginalisant la vraisemblance sur les latents temporels et en employant l'échantillonnage par importance. La décomposition de gradient résultante (Théorème 3.1) dépend du ratio des vraisemblances de trajectoire entre la politique actuelle $\pi_\theta$ et une politique obsolète $\pi_{old}$. Le défi clé est d'évaluer $\pi(x_{0:T})$ efficacement.

2. Les estimateurs

Le cadre d2 offre deux estimateurs spécifiques basés sur l'architecture du modèle :

• d2-AnyOrder (Estimateur exact) :

  • Cible : Les DLM supportant le décodage dans n'importe quel ordre (AO-dLLMs). Ces modèles permettent de décoder les jetons dans n'importe quel ordre, à condition que le masque d'attention respecte certaines contraintes de causalité (masques indépendants et causalité d'ordre).
  • Mécanisme : Il construit une séquence de longueur $2L$ (en concaténant les jetons propres et les jetons de masque) avec des encodages positionnels et des masques d'attention spécifiques. Cela permet au modèle de calculer les vraisemblances conditionnelles exactes de tous les jetons en un seul passage avant.
  • Limitation : Les DLM masqués standards (comme LLaDA) ne supportent pas nativement le décodage dans n'importe quel ordre ; leurs objectifs d'entraînement ne garantissent pas les propriétés d'attention nécessaires pour que cet estimateur soit sans biais.

• d2-StepMerge (Estimateur approximatif) :

  • Cible : Les modèles de diffusion masqués (MDM) standards qui ne supportent pas le décodage dans n'importe quel ordre.
  • Mécanisme : Inspiré par la vraisemblance composite par blocs, cet estimateur approxime la vraisemblance de la trajectoire en divisant la trajectoire en $T$ étapes en $N$ segments temporels contigus. Au lieu d'évaluer chaque étape, il évalue la vraisemblance des jetons non masqués au sein de chaque segment à l'aide d'un seul passage avant par segment.
  • Compromis : Il réduit le nombre de passages avant de $T$ à $N$. L'article fournit une borne supérieure théorique (Théorème 4.1) sur l'erreur d'approximation (divergence KL), montrant que l'erreur diminue de manière monotone à mesure que $N$ augmente, quantifiant ainsi le compromis calcul-biais.

Contributions clés

1. Cadre d2 : Un cadre de post-entraînement par RL principled pour les DLM masqués qui dérive un objectif GRPO dépendant d'une estimation précise de la vraisemblance de trajectoire.
2. d2-AnyOrder : Un estimateur de vraisemblance exact en un seul passage pour les AO-dLLMs, démontrant qu'une estimation de vraisemblance sans biais surpasse significativement les bases existantes dans les contextes de RL.
3. d2-StepMerge : Un estimateur déterministe à plusieurs passages pour les MDM standards avec une erreur d'approximation analytiquement bornée, offrant un compromis efficacité-précision supérieur par rapport aux méthodes précédentes comme diffu-GRPO.
4. Validation empirique : Des expériences approfondies montrant que d2 atteint des performances de pointe sur les benchmarks de raisonnement sans dépendre du fine-tuning supervisé de chaîne de pensée (CoT).

Résultats expérimentaux

Les auteurs ont évalué d2 sur des benchmarks de raisonnement logique (Countdown, Sudoku) et de raisonnement mathématique (GSM8K, MATH500), ainsi que sur des tâches de codage et de pilotage de toxicité.

• Performance sur les AO-dLLMs : Appliqué à des modèles supportant le décodage dans n'importe quel ordre (ex. : Eso-LM, Set Diffusion, et un LLaDA custom Any-Order), d2-AnyOrder a nettement surpassé les bases comme DDPO et diffu-GRPO. Par exemple, sur GSM8K avec une base Qwen3-1.7B, d2-AnyOrder a atteint 67 % de précision contre 63 % pour diffu-GRPO.
• Performance sur les MDM standards : Appliqué à LLaDA-8B-Instruct (un DLM masqué standard), d2-StepMerge (avec $N=16$) a établi de nouveaux records de performance :
  • Sudoku : 91,9 % (contre 22,1 % pour d1 avec SFT).
  • Countdown : 56,6 % (contre 42,2 % pour d1).
  • GSM8K : 85,0 % (contre 82,1 % pour d1).
  • MATH500 : 41,6 % (contre 39,0 % pour wd1).
    Notamment, ces résultats ont été obtenus sans fine-tuning supervisé sur des données CoT supplémentaires, alors que certaines bases reposaient sur de telles données.
• Généralisation : La méthode s'est généralisée à d'autres architectures (Dream 7B) et tâches (HumanEval, MBPP), montrant une meilleure généralisation de la longueur par rapport à diffu-GRPO.
• Ablation : Les expériences ont confirmé que les vraisemblances exactes (d2-AnyOrder) fournissent des signaux de plus haute fidélité que les approximations, et que d2-StepMerge avec un $N$ optimal (ex. : 16) équilibre mieux le coût de calcul et la précision que les valeurs extrêmes.

Signification et affirmations

L'article affirme qu'une estimation précise de la vraisemblance de la trajectoire est centrale au succès du RL pour le raisonnement dans les DLM. Les auteurs soutiennent que les méthodes précédentes n'ont pas pu atteindre des performances optimales car elles reposaient sur des estimations de vraisemblance biaisées ou coûteuses en calcul.

• Contribution théorique : Ce travail fournit la première dérivation rigoureuse d'un objectif GRPO pour les DLM masqués et établit des bornes théoriques sur l'erreur d'approximation des stratégies de fusion d'étapes (step-merging).
• Impact pratique : d2 démontre que les modèles de diffusion peuvent atteindre des capacités de raisonnement compétitives, voire supérieures, aux modèles autorégressifs lorsqu'ils sont correctement entraînés avec le RL, même sans données de CoT supervisées étendues.
• Perspectives : Les auteurs notent que bien que d2-AnyOrder soit supérieur, son applicabilité est actuellement limitée aux modèles spécifiquement entraînés pour le décodage dans n'importe quel ordre. Ils considèrent le développement de modèles AO-dLLM plus puissants et polyvalents comme une étape naturelle pour exploiter pleinement l'estimateur exact.

L'article conclut que d2 établit un nouveau niveau de référence pour les DLM sur les tâches de raisonnement logique et mathématique, validant l'efficacité de leurs stratégies d'estimation de vraisemblance pour débloquer le potentiel de raisonnement des modèles de langage basés sur la diffusion.