💬 NLP

Mask-Aware Policy Gradients for Diffusion Language Models

Cet article introduit les Mask-Aware Policy Gradients, un cadre d'apprentissage par renforcement qui formalise la génération par modèle de langage de diffusion masqué comme un processus de décision en deux étapes afin d'optimiser conjointement la sélection de jetons et le masquage de position, surmontant ainsi l'intractabilité de la vraisemblance logarithmique et atteignant des performances de pointe sur les benchmarks de raisonnement mathématique et de codage.

Auteurs originaux : Haran Raajesh, Kulin Shah, Adam Klivans, Philipp Krähenbühl

Publié 2026-07-17

📖 3 min de lecture☕ Lecture pause café

CC BY 4.0

Auteurs originaux : Haran Raajesh, Kulin Shah, Adam Klivans, Philipp Krähenbühl

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 essayiez d'apprendre à un robot comment écrire une histoire ou résoudre un problème de mathématiques. Pendant longtemps, la meilleure façon de faire cela était de faire écrire au robot un mot à la fois, comme un humain tapant sur un clavier. C'est ce qu'on appelle la génération « autorégressive ». C'est fiable, mais c'est aussi lent, comme un escargot traversant une autoroute. Récemment, des scientifiques ont découvert une méthode plus rapide appelée « Diffusion ». Au lieu d'écrire en partant de zéro, le robot commence avec une page remplie de masques vides (comme [MASK]) et les remplit progressivement, révélant les mots un par un jusqu'à ce que la phrase soit complète. C'est comme un puzzle où l'on commence avec une image recouverte et où l'on lève lentement les couvertures pour voir l'image.

Cependant, apprendre à ces robots de « diffusion » rapides à devenir plus intelligents grâce à des récompenses (une technique appelée Apprentissage par Renforcement) a été un cauchemar. Dans l'ancienne méthode « mot par mot », il est facile de calculer exactement quelle était la probabilité que le robot fasse un choix spécifique. Mais dans la méthode de « diffusion », le robot fait deux choix à chaque étape : il décide de quel mot mettre dans un espace vide, et il décide également de quel espace vide révéler ensuite. Les méthodes précédentes essayaient d'enseigner au robot en ne regardant que le premier choix (le mot) et en ignorant le second choix (l'ordre de révélation). C'est comme essayer d'apprendre à un chef à cuisiner un repas parfait en critiquant uniquement les ingrédients qu'il a choisis, tout en ignorant complètement l'ordre dans lequel il a haché et mélangé.

Cet article, présenté à la conférence COLM 2026, suggère que nous avons manqué une partie cruciale du puzzle. Les auteurs, Haran Raajesh, Kulin Shah et leur équipe de l'Université du Texas à Austin, soutiennent que pour véritablement maîtriser l'art de la diffusion, il faut récompenser le robot pour à la fois choisir les bons mots et choisir le bon ordre pour les révéler. Ils ont développé une nouvelle méthode appelée « Mask-Aware Policy Gradients » (Gradients de politique sensibles au masque). Au lieu de simplement deviner le mot suivant, leur système traite la décision de « quel masque lever ensuite » comme une partie cruciale de la stratégie du robot. En décomposant mathématiquement le processus d'apprentissage en ces deux parties distinctes, ils ont constaté que le robot apprend beaucoup plus vite et commet moins d'erreurs.

Les résultats sont assez impressionnants. Lorsqu'ils ont testé leur nouvelle méthode sur des problèmes de mathématiques difficiles et des défis de programmation, le robot s'est nettement amélioré pour les résoudre. Sur un test de mathématiques célèbre appelé GSM8K, leur méthode a atteint une précision de 87,1 %, et sur un test de codage nommé MBPP, elle a atteint 53,4 %. Ces chiffres sont plus élevés que n'importe quelle méthode précédente ayant tenté d'enseigner les modèles de diffusion à l'aide de récompenses. L'article montre qu'en prêtant attention aux décisions de « masquage » — l'ordre dans lequel le robot révèle ses pensées — nous pouvons débloquer un nouveau niveau d'intelligence pour ces modèles d'IA rapides et flexibles, les rendant non seulement plus rapides, mais aussi plus intelligents.

Résumé technique : Gradients de politique sensibles au masquage pour les modèles de langage par diffusion

Énoncé du problème

L'apprentissage par renforcement (RL) a amélioré avec succès les capacités de raisonnement des grands modèles de langage (LLM) autorégressifs en optimisant les gradients de politique basés sur des log-vraisemblances tractables. Cependant, étendre ces méthodes aux modèles de langage par diffusion masqués (MDLM) reste un défi. Contrairement aux modèles autorégressifs, qui génèrent des jetons séquentiellement de gauche à droite, les MDLM génèrent du texte en démasquant de manière itérative des jetons à partir d'une séquence entièrement masquée.

La difficulté centrale réside dans l'intractabilité de l'estimation de la log-vraisemblance pour les MDLM. Les approches existantes approximent cette vraisemblance en utilisant des variantes de la borne inférieure de l'évidence (ELBO) ou en modélisant uniquement la composante de prédiction de jetons de la trajectoire de génération. Ces méthodes ignorent un aspect critique de la génération MDLM : l'ordre dans lequel les positions sont démasquées. Dans les MDLM, la trajectoire de génération implique deux décisions distinctes à chaque étape : (1) prédire quel jeton placer aux positions masquées, et (2) décider lesquelles de ces positions conserver démasquées (et lesquelles remettre en état de masquage). Les méthodes antérieures qui ignorent cette seconde décision ne parviennent pas à utiliser le signal complet disponible dans le processus de génération pour l'optimisation de la politique.

Méthodologie

Les auteurs formalisent la génération MDLM comme un processus de décision de Markov (MDP) à deux étapes d'action. À chaque étape de débruitage $t$ , le modèle effectue :

Prédiction de jetons : Prédire les jetons pour toutes les positions actuellement masquées.
Sélection de position (Remasquage) : Décider de l'ensemble de sous-ensembles de positions à révéler et de celles à retourner à l'état [MASK].

Pour permettre l'optimisation par gradient de politique sur ce processus, les auteurs introduisent le Remasquage Probabiliste. L'inférence standard des MDLM utilise souvent une stratégie top- $K$ déterministe et gourmande pour sélectionner les positions à démasquer en fonction des scores de confiance. Cette sélection gourmande est non différentiable, empêchant l'optimisation par gradient de l'ordre de démasquage. La méthode proposée remplace cela par une variante probabiliste :

Au lieu de sélectionner déterministement les top- $K$ positions, le modèle échantillonne un sous-ensemble de positions $U_t$ à partir d'une distribution de Plackett–Luce.
Les probabilités de cette distribution sont dérivées directement des log-vraisemblances (logits) des jetons du modèle lui-même, mises à l'échelle par un paramètre de température $\tau$ .
Cette formulation crée une distribution différentiable sur les sous-ensembles de positions, permettant d'inclure la log-probabilité de la décision de démasquage dans la vraisemblance de la trajectoire.

Décomposition du gradient de politique

En modélisant la vraisemblance de la trajectoire complète $\pi_\theta(\hat{z}|c)$ comme un produit des probabilités de prédiction de jetons et des probabilités de démasquage, le gradient de politique se décompose naturellement en deux termes :
$\nabla_\theta J(\theta) = \mathbb{E}_{\hat{z}} \left[ R(c, z_T) \sum_{t=1}^T \left( \underbrace{\nabla_\theta \log \pi_\theta(\hat{z}_t | c, z_{t-1})}_{\text{Gradient de jeton}} + \underbrace{\nabla_\theta \log p_{\text{unmask}}(U_t | \hat{z}_t, z_{t-1}, c)}_{\text{Gradient de masquage}} \right) \right]$

Terme de jeton : Renforce la correction des jetons prédits.
Terme de masquage : Renforce la sélection de l'ordre de démasquage.
Efficacité : Le terme de masquage est calculé à partir des logits existants du modèle utilisés pour la prédiction de jetons, ne nécessitant aucun paramètre supplémentaire, changement architectural ou passage avant (forward pass) additionnel.

Les auteurs implémentent cela en utilisant l'algorithme Group Sequence Policy Optimization (GSPO), en plafonnant (clipping) les rapports d'importance pour les termes de jeton et de position séparément. Pour gérer le coût computationnel de l'évaluation des trajectoires complètes, ils emploient l'approximation StepMerge, regroupant les étapes de débruitage en segments pour réduire le nombre de passages avant tout en préservant la structure de la trajectoire.

Contributions clés

Formalisation du MDLM comme un MDP à deux étapes : L'article fournit un cadre théorique montrant que la génération MDLM implique à la fois la sélection de jetons et la sélection de positions, et que le gradient de politique se décompose de la sorte.
Remasquage Probabiliste : La dérivation d'une distribution de démasquage différentiable (Plackett–Luce) basée sur les propres logits du modèle, permettant l'optimisation de l'ordre de démasquage sans modifications architecturales.
Insight Théorique : Démonstration que l'omission de la composante de position du gradient de politique peut entraîner la perte de directions qui améliorent le rendement attendu, même lorsque les probabilités de jetons restent inchangées.
Performance de l'état de l'art : Validation empirique montrant que l'optimisation conjointe des deux termes produit des améliorations constantes par rapport aux méthodes existantes basées sur les trajectoires ou sur l'ELBO.

Résultats expérimentaux

La méthode a été évaluée sur des tâches de raisonnement mathématique (GSM8K, MATH500) et de génération de code (HumanEval, MBPP), en utilisant LLaDA-8B-Instruct comme modèle de base.

Gains de performance : La méthode proposée a atteint de nouveaux résultats de l'état de l'art sur tous les benchmarks.
- GSM8K : 87,1 % de précision (à une longueur de génération de 128), soit une amélioration de +2,5 % par rapport au précédent meilleur résultat (StepMerge).
- MATH500 : 53,4 % de précision, soit une amélioration de +4,0 %.
- HumanEval : 47,1 % de précision, soit une amélioration de +2,2 %.
- MBPP : 53,4 % de précision, soit une amélioration de +2,5 %.
Études d'ablation :
- Taille de bloc : L'écart de performance entre la méthode proposée et les bases s'est creusé à mesure que la taille du bloc d'inférence augmentait (de 32 à la séquence complète), suggérant que le bénéfice de l'optimisation de la sélection de position croît avec la complexité de l'espace de décision de démasquage.
- Efficacité : La méthode a convergé plus rapidement en temps réel par rapport à la forte base ELBO (SPG), atteignant la précision finale de SPG en environ 15 heures contre environ 18 heures, et atteignant une précision finale plus élevée (81,0 % contre 78,5 % sur GSM8K).
- Généralisation : L'approche s'est généralisée à un modèle de base différent (Dream-7B) et à des tâches de planification (Sudoku, Countdown), surpassant systématiquement les bases.
Comparaison avec DCoLT : Comparée à LLaDOU (qui utilise une tête séparée pour la sélection de position), la méthode proposée a obtenu une précision comparable ou meilleure avec des coûts d'entraînement nettement inférieurs (~160 heures GPU contre ~800), car elle ne nécessite aucun module supplémentaire.

Signification et affirmations

L'article affirme que la flexibilité des MDLM s'accompagne d'une structure de trajectoire plus riche que celle des modèles autorégressifs, spécifiquement la décision de quand démasquer les positions. En traitant cette décision comme faisant partie de la politique et en l'optimisant via un cadre probabiliste différentiable, les auteurs démontrent que l'optimisation conjointe des prédictions de jetons et de la sélection de position est cruciale pour maximiser la performance des modèles de langage par diffusion dans des contextes de RL.

Les auteurs soulignent que leur approche est pratiquement efficace, car elle exploite les logits existants du modèle sans nécessiter de nouveaux paramètres ou de changements architecturaux, ce qui la rend compatible avec les infrastructures d'entraînement MDLM standards comme StepMerge. Les résultats suggèrent que le "gradient de masquage" fournit un signal d'entraînement significatif que les méthodes antérieures, qui se concentraient uniquement sur les log-vraisemblances de jetons, n'ont pas réussi à capturer.