Het temmen van gemaskeerde diffusie-taalmodelle via consistentie-traject versterkend leren met minder decodeerstappen

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

September 28, 2025

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

cs.AI

Samenvatting

Gemaskerde diffusie-taalmodelen (MDLM's) zijn recent naar voren gekomen als een veelbelovend alternatief voor autoregressieve (AR) taalmodelen, met eigenschappen zoals parallel decoderen, flexibele generatievolgordes en het potentieel voor minder inferentiestappen. Ondanks deze voordelen blijven decodeerstrategieën en reinforcement learning (RL)-algoritmen die specifiek zijn afgestemd op MDLM's onderbelicht. Een naïeve aanpak is om technieken die goed zijn ingeburgerd voor AR-modellen direct over te dragen naar MDLM's. Dit roept echter een directe vraag op: Is zo'n naïeve overdracht wel echt optimaal? Bijvoorbeeld: 1) Bloksgewijze en semi-AR-decodeerstrategieën worden niet gebruikt tijdens de training van MDLM's, dus waarom presteren ze beter dan volledige diffusie-stijl decodering tijdens inferentie? 2) Het direct toepassen van RL-algoritmen die zijn ontworpen voor AR-modellen op MDLM's vertoont een inconsistentie tussen training en inferentie, aangezien MDLM-decodering niet-causaal (parallel) is. Dit resulteert in inconsistenties tussen de rollout-trajectorie en de optimalisatietrajectorie. Om deze uitdagingen aan te pakken, stellen we EOS Early Rejection (EOSER) en Ascending Step-Size (ASS) decodeerscheduler voor, die het potentieel van MDLM's ontsluiten om volledige diffusie-stijl decodering uit te voeren, met competitieve prestaties en minder decodeerstappen. Daarnaast introduceren we Consistency Trajectory Group Relative Policy Optimization (CJ-GRPO) voor het temmen van MDLM's, wat de consistentie tussen rollout-trajectorie en optimalisatietrajectorie benadrukt en de optimalisatiefouten veroorzaakt door skip-step optimalisatie vermindert. We voeren uitgebreide experimenten uit op redeneertaken, zoals wiskundige en planningsbenchmarks, met behulp van LLaDA-8B-Instruct. De resultaten tonen aan dat de voorgestelde EOSER- en ASS-mechanismen, samen met CJ-GRPO, veelbelovend zijn voor het effectief en efficiënt temmen van MDLM's. 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.