🤖 machine learning

Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization

Cet article introduit l'optimisation de la politique d'appariement de distribution (DMPO), un nouveau cadre d'apprentissage par renforcement qui améliore considérablement les capacités de raisonnement des grands modèles de langage par diffusion en alignant la distribution de leur politique sur une cible optimale inclinée vers la récompense, atteignant ainsi des améliorations d'exactitude substantielles par rapport aux bases de référence existantes sans ajustement fin supervisé.

Auteurs originaux : Yuchen Zhu, Wei Guo, Jaemoo Choi, Petr Molodyk, Bo Yuan, Molei Tao, Yongxin Chen

Publié 2026-07-14

📖 7 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Yuchen Zhu, Wei Guo, Jaemoo Choi, Petr Molodyk, Bo Yuan, Molei Tao, Yongxin Chen

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 avez un robot super intelligent qui écrit des histoires, résout des problèmes mathématiques et joue à des jeux. La plupart des robots d'aujourd'hui travaillent comme une personne lisant un livre mot par mot, de gauche à droite. Ils ne peuvent pas sauter des étapes ou revenir en arrière ; ils doivent simplement deviner le mot suivant en fonction de ceux qui précèdent. C'est lent, comme essayer de compléter un immense puzzle en ne regardant que la pièce située juste devant soi.

Entrez en scène les Modèles de Langage de Diffusion (dLLM). Ce sont les nouveaux venus sur le marché. Au lieu d'écrire mot par mot, ils partent d'un désordre confus de « jetons mystères » (comme un puzzle où chaque pièce est recouverte de brouillard) et nettoient lentement l'ensemble, révélant la réponse d'un seul coup ou dans n'importe quel ordre. Cela les rend potentiellement beaucoup plus rapides pour réfléchir.

Mais voici le hic : bien que ces robots de diffusion soient rapides, ils ne sont pas toujours les plus intelligents pour les tâches de raisonnement complexes comme les mathématiques avancées ou les énigmes logiques. Pour les rendre plus intelligents, les scientifiques utilisent généralement une technique appelée Apprentissage par Renforcement (RL). Considérez cela comme un jeu vidéo où le robot gagne des points pour une bonne réponse et en perd pour une mauvaise.

Le Problème : Le Piège de la « Recherche de Mode »

L'ancienne méthode pour entraîner ces robots (utilisant des méthodes comme GRPO) était un peu comme un étudiant qui n'étudie que la seule réponse qu'il pense être la bonne. Si le robot trouve un moyen d'obtenir des points qui semble « sûr », il arrête d'explorer. Il s'enferme dans une routine, ignorant d'autres solutions ingénieuses qui pourraient être tout aussi bonnes. Dans l'article, les auteurs appellent cela la « recherche de mode » (mode-seeking). C'est comme un robot qui n'apprendrait à résoudre un problème mathématique que d'une seule manière spécifique, et si cette voie est bloquée, il panique. Il a également tendance à ignorer les réponses « désordonnées » mais correctes qui semblent différentes de la première qu'il a trouvée.

L'article soutient que cette ancienne façon de simplement poursuivre le score le plus élevé est imparfaite pour les modèles de diffusion. Il suggère qu'en se concentrant uniquement sur le meilleur chemin, nous passons à côté de la capacité unique du robot à explorer de nombreux chemins différents à la fois.

La Solution : DMPO (Le Robot « Correspondance de Carte »)

Les auteurs proposent une nouvelle méthode appelée Optimisation de Politique de Correspondance de Distribution (DMPO).

Au lieu de dire au robot : « Trouve simplement le score le plus élevé », DMPO dit : « Voici l'ensemble de la carte de toutes les bonnes réponses que tu pourrais donner. Ton travail est d'apprendre à correspondre à toute cette carte. »

Imaginez que vous essayez d'apprendre à un chien à rapporter un objet.

L'ancienne méthode : Vous lancez une balle, et le chien court vers l'unique endroit où la balle est tombée. Si la balle atterrit dans un buisson, le chien apprend à ne rapporter l'objet que depuis les buissons.
La méthode DMPO : Vous montrez au chien une carte de chaque endroit possible où la balle pourrait atterrir (buissons, herbe, sable, eau) et vous dites : « Apprends à rapporter l'objet de n'importe lequel de ces endroits, en fonction de la qualité de l'endroit. » Le chien apprend à être flexible et à explorer tout le jardin, pas seulement un coin.

Sur le plan technique, DMPO utilise un tour mathématique spécial appelé Entropie Croisée de Débruitage Pondérée (WDCE). Cela permet au robot d'apprendre de ses tentatives passées (même celles qui n'étaient pas parfaites) sans avoir besoin de tout régénérer à partir de zéro à chaque fois. C'est comme avoir un « tampon de relecture » (replay buffer) où le robot peut étudier ses anciens jeux encore et encore, apprenant des bons et des mauvais mouvements sans s'embrouiller.

La Recette Secrète : Le « Baseline de Poids »

Les auteurs ont découvert un problème délicat lors de l'entraînement avec de petits groupes d'exemples (petites tailles de lots/batch sizes). Parfois, le robot se mélangeait les pinceaux et commençait à récompenser de mauvères réponses simplement parce qu'elles étaient les seules qu'il voyait.

Pour corriger cela, ils ont inventé une astuce ingénieuse appelée Soustraction de Baseline de Poids (Weight Baseline Subtraction).
Considérez cela comme un enseignant qui corrige un examen. Si un élève répond correctement à une question, il reçoit une étoile dorée. Mais si l'enseignant n'avait vu que cette question, il pourrait penser que tout ce que l'élève a fait était excellent. La « baseline » est comme un « standard de la moyenne ». L'enseignant soustrait le score « moyen » du score de l'élève.

Si l'élève a fait mieux que la moyenne, il reçoit une grande étoile dorée.
S'il a fait moins bien, il reçoit une « pénalité » (un poids négatif), même s'il a techniquement obtenu quelques points.

Cela garantit que le robot ne s'emballe pas pour des réponses médiocres et continue de viser les réponses véritablement excellentes.

Les Résultats : À quel point est-ce meilleur ?

Les auteurs ont testé cette nouvelle méthode sur des tests de raisonnement très difficiles. Ils ont appliqué DMPO à des modèles pré-entraînés (comme LLaDA-Instruct et Dream-Instruct) de deux manières différentes :

Application Directe : Ils ont appliqué DMPO directement aux modèles de base sans aucun « devoir à la maison » préalable (Ajustement Supervisé ou SFT) sur des ensembles de données de raisonnement. Cette approche « de type R1-Zero » a été utilisée pour démontrer clairement le potentiel brut de DMPO.
Application Améliorée : Ils ont également appliqué DMPO à des modèles qui avaient déjà subi un SFT pour montrer qu'il s'agit d'une mise à niveau puissante pour des modèles déjà entraînés.

Les résultats ont été très impressionnants sur les différents tests de raisonnement :

Sur une énigme mathématique appelée GSM8K, le modèle appliqué directement au modèle de base pré-entraîné (sans SFT) a montré des améliorations massives. Plus précisément, DMPO a atteint jusqu'à 39,63 points de pourcentage d'amélioration de la précision par rapport aux bases de RL non-DMPO précédentes, et une amélioration stupéfiante de 67,97 points de pourcentage par rapport au modèle de base lui-même.
Sur une énigme logique appelée Sudoku, ils ont constaté des gains énormes, avec une version de leur modèle passant d'un taux de réussite de 16,41 % à 80,86 % (une différence de plus de 64 points !).
Crucialement, DMPO a continué de produire des gains de performance significatifs même lorsqu'il était appliqué à des modèles ayant déjà subi un SFT, prouvant qu'il s'agit d'une méthode robuste qui fonctionne que l'on parte de zéro ou que l'on construise sur des entraînements existants.

L'article suggère que DMPO n'est pas seulement un petit ajustement, mais un changement fondamental. Il permet au robot d'être « off-policy », ce qui signifie qu'il peut apprendre de données anciennes efficacement, et « forward-only », ce qui signifie qu'il n'a pas besoin de calculs de rétropropagation coûteux qui ralentissent les autres robots.

L'Essentiel

Les auteurs sont convaincus que cette méthode fonctionne bien pour les modèles spécifiques qu'ils ont testés (comme LLaDA et Dream) sur ces tâches de raisonnement spécifiques. Ils ont mesuré ces améliorations grâce à des expériences rigoureuses sur des ensembles de données standards, montrant que DMPO peut booster les performances, qu'il soit appliqué à un modèle qui vient d'être pré-entraîné ou à un modèle ayant déjà subi un ajustement supervisé. Cependant, ils admettent ne pas avoir encore testé cela sur tous les types de modèles ou de tâches possibles, donc bien que les résultats soient solides, tout le potentiel est encore en cours d'exploration.

En bref, DMPO apprend aux robots de diffusion à arrêter de se focaliser sur une seule réponse « parfaite » et à plutôt apprendre à apprécier tout le paysage des bonnes solutions, ce qui les rend plus intelligents, plus rapides et plus créatifs dans la résolution de problèmes.

Résumé Technique : Améliorer le Raisonnement des LLM de Diffusion via l'Optimisation de l'Appariement de Distribution (DMPO)

Énoncé du Problème

Les Grands Modèles de Langage de Diffusion (dLLM) offrent une alternative prometteuse aux LLM autorégressifs (AR-LLM) en permettant une génération dans n'importe quel ordre et un débit d'inférence potentiellement plus élevé. Cependant, l'adaptation des dLLM pour égaler les capacités de raisonnement des AR-LLM reste un défi majeur. Les méthodes existantes d'ajustement par Apprentissage par Renforcement (RL), telles que le Group Relative Policy Optimization (GRPO), sont mal adaptées aux dLLM pour deux facteurs principaux :

Incompatibilité Computationnelle : Les dLLM reposent sur un traitement bidirectionnel, ce qui rend l'estimation des log-probabilités de séquence (requise pour les méthodes standard de gradient de politique) coûteuse en calcul et souvent intraitable par rapport aux modèles AR.
Limitations Algorithmiques : Les approches de RL actuelles pour les dLLM reposent souvent sur des approximations heuristiques (par exemple, la décomposition de champ moyen) qui ignoreent les corrélations entre les positions des jetons. De plus, elles se concentrent généralement sur la maximisation de la récompense, ce qui tend à optimiser un mode de récompense unique. Ce comportement de « recherche de mode » (mode-seeking) peut conduire au détournement de récompense (reward hacking) et à un échec de l'exploration des divers chemins de raisonnement de haute qualité que les dLLM sont théoriquement capables de générer.

Méthodologie : Optimisation de l'Appariement de Distribution (DMPO)

Le papier propose l'Optimisation de l'Appariement de Distribution (DMPO), un cadre d'ajustement par RL fondé sur le contrôle optimal stochastique (SOC). Contrairement aux méthodes traditionnelles qui maximisent l'espérance de la récompense, la DMPO vise à faire correspondre la distribution de la politique du modèle à une distribution cible optimale, inclinée par la récompense.

Cadre Fondamental

L'objectif est d'apprendre une politique $\pi_\theta$ telle que la distribution de séquence induite $p_\theta(o|q)$ corresponde à la distribution optimale $p^*(o|q)$ , définie par :
$p^*(o|q) = \frac{1}{Z(q)} p_{ref}(o|q) e^{r(q,o)/\alpha}$
où $p_{ref}$ est la politique pré-entraînée, $r$ est la fonction de récompense, et $\alpha$ est un paramètre de température.

Pour y parvenir, les auteurs formulent le problème comme la minimisation de la divergence de Kullback-Leibler directe (Forward KL) (équivalente à l'entropie croisée) entre la cible et la politique actuelle :
$\min_\theta \mathbb{E}_{q \sim D} \mathbb{E}_{p^*(o|q)} \left[ \log \frac{p^*(o|q)}{p_\theta(o|q)} \right]$
Cet objectif de « couverture de masse » (mass-covering) encourage le modèle à explorer des trajectoires de haute récompense diversifiées plutôt qu'à s'effondrer sur un mode unique.

Composantes Techniques Clés

Entropie Croisée de Débruitage Pondérée (WDCE) :
Puisque $p^*$ est intraitable (en raison de la fonction de partition inconnue $Z(q)$ et de l'impossibilité d'échantillonner directement à partir de celle-ci), la DMPO utilise l'Échantillonnage par Importance. Elle traite les échantillons provenant d'une politique de référence $p_v$ (une copie du modèle avec arrêt du gradient) comme des échantillons pondérés de $p^*$ .
La fonction de perte résultante, WDCE, exploite le processus de diffusion (noising) direct unique des dLLM. Elle calcule la perte en utilisant uniquement les échantillons finaux propres et leurs poids d'importance associés, en écartant les trajectoires de déroulement intermédiaires. Cela fait de la WDCE un objectif hors-politique (off-policy) et direct (forward-only).
Soustraction de la Base de Poids (Weight Baseline Subtraction) :
Un défi de mise en œuvre critique identifié est l'instabilité de la WDCE lors de l'utilisation de petites tailles de lots d'entraînement. Dans de petits lots, des poids positifs pour de « mauvaises » réponses (qui ont pourtant une probabilité non nulle sous $p^*$ ) peuvent dominer, entraînant des mises à jour de gradient incorrectes.
Pour remédier à cela, les auteurs proposent la Soustraction de la Base de Poids, où une base de poids $w_{base}$ est soustraite du poids d'importance :
$w_{real} = w - w_{base}$
Trois stratégies pour $w_{base}$ sont proposées :
- Base de Poids de Groupe : Une valeur constante (1) pour encourager la convergence.
- Base de Poids Individuelle : Basée sur une inclinaison de récompense négative.
- Base de Poids du Modèle : Basée sur les propres poids d'importance du modèle actuel.
  Cette technique injecte efficacement des gradients négatifs pour les échantillons à faible récompense, empêchant le modèle de promouvoir de mauvaises réponses dans les scénarios à faible lot.
Optimisation Discriminative Directe Pondérée (WDDO) :
En tant qu'alternative à l'entropie croisée, le papier explore un objectif inspiré des GAN qui incorpore naturellement des gradients négatifs pour les mauvais échantillons, améliorant davantage l'appariement de distribution.

Contributions Clés

Nouveau Cadre de RL : La proposition de la DMPO, un cadre d'appariement de distribution spécifiquement conçu pour les dLLM, déplaçant le paradigme de la maximisation de la récompense vers l'appariement de la distribution entière de la politique inclinée par la récompense.
Fondement Théorique : Dérivation de la perte WDCE via l'échantillonnage par importance et le contrôle optimal stochastique, fournissant une méthode théoriquement solide pour l'entraînement hors-politique sans nécessiter d'estimations coûteuses des log-probabilités de séquence.
Solutions Pratiques : Identification du mode de défaillance des petits lots des objectifs de KL directe et introduction de la Soustraction de la Base de Poids et de la WDDO pour stabiliser l'entraînement et améliorer l'efficacité de l'échantillonnage.
Efficacité : Démonstration que la WDCE permet l'entraînement via le processus direct seul, permettant la réutilisation des déroulements (off-policy) et la compatibilité avec les techniques d'inférence rapide (ex: cache KV), accélérant considérablement l'entraînement par rapport aux méthodes basées sur les trajectoires.

Résultats Expérimentaux

Les auteurs ont évalué la DMPO sur les modèles LLaDA-8B-Instruct et Dream-7B à travers quatre tests de raisonnement : GSM8K, MATH500, Countdown et Sudoku.

Gains de Performance : La DMPO surpasse systématiquement les modèles de base et les autres lignes de base de RL (incluant diffu-GRPO, cGRPO et d1).
- Sur Sudoku, la DMPO-LLaDA a réalisé une amélioration de 39,63 points de pourcentage par rapport à la meilleure ligne de base RL non-DMPO et de 67,97 points de pourcentage par rapport au modèle de base.
- Sur Countdown, des améliorations allant jusqu'à 16,40 points de pourcentage ont été observées par rapport aux bases solides.
Aucun SFT Requis : La méthode a obtenu ces résultats sans ajustement supervisé (SFT) sur des jeux de données curatés, démontant sa capacité à faire passer à l'échelle les capacités de raisonnement directement à partir des modèles pré-entraînés.
Efficacité de l'Entraînement : Grâce à sa nature hors-politique et son calcul direct, la DMPO a atteint une efficacité d'échantillonnage 2 à 3 fois supérieure et jusqu'à 8 fois plus de rapidité de l'échantillonnage de déroulement (lorsqu'elle est combinée avec l'inférence rapide des dLLM) par rapport aux méthodes de type GRPO. Elle a atteint les niveaux de récompense de d1 avec seulement 31 % du budget de calcul d'entraînement.
Stabilité : La méthode est restée stable même avec des données très obsolètes (réutilisation jusqu'à 24 fois du même lot de déroulement), une propriété attribuée à la stabilité des poids d'importance au niveau de la séquence utilisés par la DMPO.

Signification et Revendications

Le papier affirme que la DMPO représente une étape significative dans l'adaptation à l'échelle des dLLM pour les tâches de raisonnement. En s'éloignant du paradigme de maximisation de la récompense et des goulots d'étranglement computationnels du RL basé sur les trajectoires, la DMPO exploite pleinement les caractéristiques uniques des modèles de diffusion :

Diversité : L'objectif de la KL directe empêche l'effondrement de mode, permettant aux dLLM d'explorer des chemins de raisonnement diversifiés.
Efficacité : La nature directe et hors-politique de la WDCE s'aligne sur les capacités de haut débit d'inférence des dLLM, offrant une voie vers un ajustement par RL scalable qui est computationnellement faisable pour de grands modèles.
Généralisabilité : Le cadre est présenté comme une approche fondée sur des principes qui peut être appliquée à divers backbones de dLLM, comme en témoigne son succès sur les architectures LLaDA et Dream.

Les auteurs concluent que, bien que le travail actuel se concentre sur des modèles et des benchmarks spécifiques, le cadre d'appariement de distribution ouvre de nouvelles directions pour améliorer les modèles génératifs et concevoir des techniques de base de poids plus efficaces pour la recherche future.