Zähmung maskierter Diffusionssprachmodelle durch Konsistenz-Trajektorien-Verstärkungslernen mit weniger Dekodierungsschritten

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

Masked Diffusion Language Models (MDLMs) haben sich kürzlich als vielversprechende Alternative zu autoregressiven (AR) Sprachmodellen etabliert, da sie Eigenschaften wie parallele Dekodierung, flexible Generierungsreihenfolgen und das Potenzial für weniger Inferenzschritte bieten. Trotz dieser Vorteile bleiben Dekodierungsstrategien und Reinforcement-Learning (RL)-Algorithmen, die speziell für MDLMs entwickelt wurden, weitgehend unerforscht. Ein naiver Ansatz besteht darin, Techniken, die sich für AR-Modelle bewährt haben, direkt auf MDLMs zu übertragen. Dies wirft jedoch die unmittelbare Frage auf: Ist eine solche naive Übertragung tatsächlich optimal? Beispielsweise 1) werden blockweise und semi-autoregressive Dekodierungsstrategien während des Trainings von MDLMs nicht eingesetzt – warum übertreffen sie dann die vollständige Diffusions-Dekodierung während der Inferenz? 2) Die direkte Anwendung von RL-Algorithmen, die für AR-Modelle entwickelt wurden, auf MDLMs zeigt eine Trainings-Inferenz-Inkonsistenz, da die MDLM-Dekodierung nicht-kausal (parallel) ist. Dies führt zu Inkonsistenzen zwischen dem Rollout-Pfad und dem Optimierungspfad. Um diese Herausforderungen zu bewältigen, schlagen wir den EOS Early Rejection (EOSER) und den Ascending Step-Size (ASS) Dekodierungsplaner vor, die das Potenzial von MDLMs freisetzen, eine vollständige Diffusions-Dekodierung durchzuführen und dabei mit weniger Dekodierungsschritten eine wettbewerbsfähige Leistung zu erzielen. Zusätzlich führen wir die Consistency Trajectory Group Relative Policy Optimization (CJ-GRPO) zur Steuerung von MDLMs ein, die die Konsistenz zwischen Rollout-Pfad und Optimierungspfad betont und die Optimierungsfehler, die durch Skip-Step-Optimierung verursacht werden, reduziert. Wir führen umfangreiche Experimente zu Denkaufgaben, wie mathematischen und Planungs-Benchmarks, mit LLaDA-8B-Instruct durch. Die Ergebnisse zeigen, dass die vorgeschlagenen EOSER- und ASS-Mechanismen zusammen mit CJ-GRPO ein erhebliches Potenzial für die effektive und effiziente Steuerung von MDLMs besitzen. 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.