Domestication des modèles de langage à diffusion masquée via l'apprentissage par renforcement de trajectoire de cohérence avec moins d'étapes de décodage

Taming Masked Diffusion Language Models via Consistency Trajectory Reinforcement Learning with Fewer Decoding Step

September 28, 2025

papers.authors: Jingyi Yang, Guanxu Chen, Xuhao Hu, Jing Shao

cs.AI

papers.abstract

Les modèles de langage à diffusion masquée (MDLM) ont récemment émergé comme une alternative prometteuse aux modèles de langage autorégressifs (AR), offrant des propriétés telles que le décodage parallèle, des ordres de génération flexibles et la possibilité de réduire le nombre d'étapes d'inférence. Malgré ces avantages, les stratégies de décodage et les algorithmes d'apprentissage par renforcement (RL) adaptés aux MDLM restent peu explorés. Une approche naïve consiste à transférer directement les techniques bien établies pour les modèles AR aux MDLM. Cependant, cela soulève une question immédiate : un tel transfert naïf est-il vraiment optimal ? Par exemple, 1) les stratégies de décodage par blocs et semi-autorégressives ne sont pas utilisées lors de l'entraînement des MDLM, alors pourquoi surpassent-elles le décodage de style diffusion complet lors de l'inférence ? 2) L'application directe des algorithmes RL conçus pour les modèles AR aux MDLM présente une incohérence entre l'entraînement et l'inférence, car le décodage des MDLM est non-causal (parallèle). Cela entraîne des incohérences entre la trajectoire de déploiement et la trajectoire d'optimisation. Pour relever ces défis, nous proposons le rejet précoce EOS (EOSER) et le planificateur de décodage à pas ascendants (ASS), qui permettent aux MDLM de réaliser un décodage de style diffusion complet, atteignant des performances compétitives avec moins d'étapes de décodage. De plus, nous introduisons l'optimisation relative de politique par groupe de trajectoire de cohérence (CJ-GRPO) pour maîtriser les MDLM, qui met l'accent sur la cohérence entre la trajectoire de déploiement et la trajectoire d'optimisation, et réduit les erreurs d'optimisation causées par l'optimisation à pas sautés. Nous menons des expériences approfondies sur des tâches de raisonnement, telles que des benchmarks mathématiques et de planification, en utilisant LLaDA-8B-Instruct. Les résultats démontrent que les mécanismes EOSER et ASS proposés, ainsi que CJ-GRPO, offrent des perspectives significatives pour maîtriser efficacement et efficacement les MDLM. Code : https://github.com/yjyddq/EOSER-ASS-RL.

English

Masked diffusion language models (MDLMs) have recently emerged as a promising alternative to autoregressive (AR) language models, offering properties such as parallel decoding, flexible generation orders, and the potential for fewer inference steps. Despite these advantages, decoding strategies and reinforcement learning (RL) algorithms tailored for MDLMs remain underexplored. A naive approach is to directly transfer techniques well-established for AR models to MDLMs. However, this raises an immediate question: Is such a naive transfer truly optimal? For example, 1) Block-wise and semi-AR decoding strategies are not employed during the training of MDLMs, so why do they outperform full diffusion-style decoding during inference? 2) Applying RL algorithms designed for AR models directly to MDLMs exhibits a training-inference inconsistency, since MDLM decoding are non-causal (parallel). This results in inconsistencies between the rollout trajectory and the optimization trajectory. To address these challenges, we propose EOS Early Rejection (EOSER) and Ascending Step-Size (ASS) decoding scheduler, which unlock the potential of MDLMs to perform full diffusion-style decoding, achieving competitive performance with fewer decoding steps. Additionally, we introduce Consistency Trajectory Group Relative Policy Optimization (CJ-GRPO) for taming MDLMs, which emphasizes the consistency between rollout trajectory and optimization trajectory, and reduces the optimization errors caused by skip-step optimization. We conduct extensive experiments on reasoning tasks, such as mathematical and planning benchmarks, using LLaDA-8B-Instruct. The results demonstrate that the proposed EOSER and ASS mechanisms, together with CJ-GRPO, hold significant promise for effectively and efficiently taming MDLMs. Code: https://github.com/yjyddq/EOSER-ASS-RL.