Domando Modelos de Linguagem de Difusão Mascarada via Aprendizado por Reforço de Trajetória de Consistência com Menos Passos de Decodificação

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

September 28, 2025

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

cs.AI

Resumo

Modelos de linguagem de difusão mascarada (MDLMs, na sigla em inglês) surgiram recentemente como uma alternativa promissora aos modelos de linguagem autoregressivos (AR), oferecendo propriedades como decodificação paralela, ordens de geração flexíveis e o potencial para menos etapas de inferência. Apesar dessas vantagens, estratégias de decodificação e algoritmos de aprendizado por reforço (RL, na sigla em inglês) adaptados para MDLMs ainda são pouco explorados. Uma abordagem ingênua é transferir diretamente técnicas bem estabelecidas para modelos AR para MDLMs. No entanto, isso levanta uma questão imediata: essa transferência ingênua é realmente ideal? Por exemplo, 1) Estratégias de decodificação em blocos e semi-AR não são empregadas durante o treinamento de MDLMs, então por que elas superam a decodificação completa no estilo de difusão durante a inferência? 2) A aplicação direta de algoritmos de RL projetados para modelos AR em MDLMs exibe uma inconsistência entre treinamento e inferência, já que a decodificação de MDLMs é não-causal (paralela). Isso resulta em inconsistências entre a trajetória de execução e a trajetória de otimização. Para enfrentar esses desafios, propomos o EOS Early Rejection (EOSER) e o Ascending Step-Size (ASS) como agendadores de decodificação, que desbloqueiam o potencial dos MDLMs para realizar decodificação completa no estilo de difusão, alcançando desempenho competitivo com menos etapas de decodificação. Além disso, introduzimos o Consistency Trajectory Group Relative Policy Optimization (CJ-GRPO) para domar MDLMs, que enfatiza a consistência entre a trajetória de execução e a trajetória de otimização, e reduz os erros de otimização causados pela otimização de etapas puladas. Realizamos experimentos extensos em tarefas de raciocínio, como benchmarks matemáticos e de planejamento, utilizando o LLaDA-8B-Instruct. Os resultados demonstram que os mecanismos propostos EOSER e ASS, juntamente com o CJ-GRPO, mostram grande potencial para domar MDLMs de forma eficaz e eficiente. Código: 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.