Addomesticare i Modelli Linguistici a Diffusione Mascherata tramite Apprendimento per Rinforzo della Traiettoria di Consistenza con Meno Passaggi di Decodifica

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

September 28, 2025

Autori: Jingyi Yang, Guanxu Chen, Xuhao Hu, Jing Shao

cs.AI

Abstract

I modelli linguistici a diffusione mascherata (MDLM) sono emersi recentemente come una valida alternativa ai modelli linguistici autoregressivi (AR), offrendo proprietà come la decodifica parallela, ordini di generazione flessibili e il potenziale per un minor numero di passaggi di inferenza. Nonostante questi vantaggi, le strategie di decodifica e gli algoritmi di apprendimento per rinforzo (RL) specifici per gli MDLM rimangono poco esplorati. Un approccio ingenuo consiste nel trasferire direttamente tecniche ben consolidate per i modelli AR agli MDLM. Tuttavia, ciò solleva una domanda immediata: un tale trasferimento ingenuo è davvero ottimale? Ad esempio, 1) le strategie di decodifica a blocchi e semi-AR non vengono utilizzate durante l'addestramento degli MDLM, quindi perché superano la decodifica in stile diffusione completa durante l'inferenza? 2) L'applicazione diretta degli algoritmi RL progettati per i modelli AR agli MDLM mostra un'incongruenza tra addestramento e inferenza, poiché la decodifica degli MDLM è non-causale (parallela). Ciò si traduce in incongruenze tra la traiettoria di rollout e la traiettoria di ottimizzazione. Per affrontare queste sfide, proponiamo il meccanismo di Early Rejection di EOS (EOSER) e lo scheduler di decodifica Ascending Step-Size (ASS), che sbloccano il potenziale degli MDLM di eseguire una decodifica in stile diffusione completa, raggiungendo prestazioni competitive con un minor numero di passaggi di decodifica. Inoltre, introduciamo l'ottimizzazione della politica relativa per gruppi di traiettorie di consistenza (CJ-GRPO) per addestrare gli MDLM, che enfatizza la coerenza tra la traiettoria di rollout e la traiettoria di ottimizzazione, riducendo gli errori di ottimizzazione causati dall'ottimizzazione a passi saltati. Abbiamo condotto esperimenti estesi su compiti di ragionamento, come benchmark matematici e di pianificazione, utilizzando LLaDA-8B-Instruct. I risultati dimostrano che i meccanismi proposti EOSER e ASS, insieme a CJ-GRPO, mostrano un potenziale significativo per addestrare efficacemente ed efficientemente gli MDLM. Codice: 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.