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

1 INTRODUCTION

• Full diffusion can fall into an “<EOS> Trap” when early high EOS confidence leads to premature termination; direct AR-style RL also mismatches MDLMs’ non-causal rollout and optimization trajectories.The introduction also notes that storing intermediate denoising states makes consistent online RL memory-intensive.Source: §1 · p. 2
• EOSER suppresses <EOS> confidence early and gradually restores it during denoising, significantly improving full-diffusion decoding performance.The method is intended to alleviate the <EOS> Trap.Source: §1 · p. 2
• ASS replaces uniform step sizes with cautious early decoding and more aggressive later decoding, reducing decoding steps from L to log2 L.This schedule follows the observed rise in average token confidence as denoising progresses.Source: §1 · p. 3
• CJ-GRPO aligns rollout and optimization denoising trajectories, reducing skip-step optimization errors and showing potential for considerable performance with fewer decoding steps alongside EOSER and ASS.It is designed to account for full diffusion decoding’s non-causal nature.Source: §1 · p. 3

2.1 MASKED DIFFUSION LANGUAGE MODELS

• MDLMs have reached the billion-parameter scale, with reported performance comparable to similarly sized autoregressive models.The section situates masked diffusion among other diffusion language-model approaches and notes its parallel decoding, flexible order, and potential for fewer inference steps.Source: §2.1 · p. 3

2.2 ONLINE REINFORCEMENT LEARNING FOR DIFFUSION LANGUAGE MODELS

• Prior MDLM work explores reasoning tasks, while d1 applies GRPO and LLaDOU explores online RL approaches for diffusion language models.Source: §2.2 · p. 3
• Policy-gradient RL designed for causal, token-by-token autoregressive generation is not directly applicable to MDLMs’ non-causal attention.Source: §2.2 · p. 3

3.1 MASKED DIFFUSION LARGE LANGUAGE MODELS

• MDLMs train a mask predictor on partially masked sequences to recover clean text, with instruction tuning leaving prompts unchanged while masking responses.The training objective minimizes a negative evidence lower bound that upper-bounds negative log-likelihood.Source: §3.1 · pp. 3–4

3.2 EOS EARLY REJECTION

• EOSER suppresses EOS confidence during early denoising, then restores it later to prevent premature sequence termination in full-diffusion decoding.The EOS trap arises because early-step EOS confidence is disproportionately high relative to non-EOS confidence.Source: §3.2 · pp. 4–5

3.3 ASCENDING STEP-SIZE SCHEDULER FOR FEWER STEPS INFERENCE

• ASS reduces empirically optimal uniform steps from L/2 to log2 L, lowering time complexity from O(L) to O(log2 L).The scheduler uses a power-of-2 schedule, decoding fewer tokens early and more later as token confidence rises.Source: §3.3 · pp. 5–6

3.4 CONSISTENCY TRAJECTORY GROUP RELATIVE POLICY OPTIMIZATION

• CJ-GRPO records intermediate denoising states and uses them to align optimization with rollout, addressing confidence mismatches from masked-context MDLM trajectories.Storing intermediate states increases memory overhead with the number of denoising steps; using ASS is described as reducing both time and space complexity.Source: §3.4 · p. 6

4.1 IMPLEMENTATION DETAILS

• LLaDA-8B-Instruct is evaluated on GSM8K, MATH500, Countdown, and Sudoku with zero-shot prompting and greedy decoding.The evaluation compares SFT, diffu-GRPO, and CJ-GRPO.Source: §4.1 · p. 7

4.2 MAIN RESULTS

• CJ-GRPO+Semi-AR beats SFT and diffu-GRPO across all tasks and 16 settings; CJ-GRPO+EOSER excels in planning but lags on mathematics.On Sudoku, CJ-GRPO+EOSER reaches approximately twice the performance of CJ-GRPO+Semi-AR and four times that of diffu-GRPO.Source: Table 1 · p. 8, §4.2 · pp. 8–9
• With only 7/8 (log2 L) steps, CJ-GRPO+EOSER+ASS significantly surpasses other methods in all settings and surpasses original diffu-GRPO on planning tasks.Table 2 reports fewer-step results for Countdown, GSM8K, MATH500, and Sudoku.Source: Table 2 · p. 8, §4.2 · p. 9
• Training-free full diffusion underperforms semi-AR and EOSER across tasks; EOSER beats semi-AR in planning but not mathematics, while delaying EOS spatially and temporally.The EOS arrival pattern is shown in Figure 4.Source: §4.2 · p. 8, Figure 4 · p. 8

4.3 ABLATION STUDY

• Optimizing over each intermediate denoising step leads to significantly better performance; skipping intermediate steps introduces substantial optimization errors in the actual rollout trajectory and some optimization bias, while the near-final alternative yields the worst performance.Source: §4.3 · p. 9

4.4 ANALYSIS

• EOSER suits planning's parallel reasoning patterns, whereas semi-AR is more suitable for mathematics' sequential reasoning patterns.The analysis describes Countdown and Sudoku as involving parallel reasoning patterns, while GSM8K and MATH500 typically involve sequential reasoning.Source: §4.4 · p. 9

5 CONCLUSION

• EOSER alleviates the full-diffusion <EOS> trap by suppressing early <EOS> generation, while ASS reduces decoding steps by adapting to token-confidence changes.The conclusion describes confidence as low early and sharply increasing later.Source: §5 · p. 10
• CJ-GRPO aligns rollout and optimization trajectories for MDLM reinforcement learning.The conclusion says the algorithm ensures trajectory consistency.Source: §5 · p. 10
• With EOSER and ASS, CJ-GRPO achieves comparable performance to L steps using only log2 L steps.The conclusion characterizes CJ-GRPO as effective and efficient for taming MDLMs on reasoning tasks.Source: §5 · p. 10

6 ETHICS STATEMENT

• The study uses publicly available datasets without sensitive or private information and involves no human subjects.The authors report compliance with relevant legal regulations and no conflicts of interest or sponsorship to disclose.Source: §6 · p. 10

7 REPRODUCIBILITY STATEMENT

• The authors provide experimental configurations, open-source scripts and instructions, and publicly accessible datasets to support replication.The experiments apply reinforcement-learning and decoding strategies to LLaDA-8B-Instruct.Source: §7 · p. 10

• The appendix covers supplementary experiments and discussions, limitations, and the authors’ use of large language models.Source: §8 · p. 13

• The ablations vary denoising steps from 1 to 256 across Countdown, GSM8K, MATH500, and Sudoku.Source: A SUPPLEMENTARY EXPERIMENTS AND… · p. 13, Table 4 · p. 13, Table 5 · p. 13
• Semi-AR and EOSER often perform best at L/2 steps, while full diffusion usually performs best below L/2.Source: A SUPPLEMENTARY EXPERIMENTS AND… · p. 14
• Performance differences across denoising-step settings of 64, 128, and 256 are relatively small; for Semi-AR decoding, block length may be a sensitive hyperparameter that results in large performance variance.Source: A SUPPLEMENTARY EXPERIMENTS AND… · p. 14
• At L/2 steps, full diffusion generates two tokens per step, and more decoding steps may not make MDLM generation faster than AR under the same infrastructure.Source: A SUPPLEMENTARY EXPERIMENTS AND… · p. 14

B LIMITATIONS

• CJ-GRPO’s memory overhead grows proportionally with denoising steps, constraining training speed; fewer-step EOSER+ASS remains behind Semi-AR approaches on mathematics.Source: B LIMITATIONS · p. 14

C THE USE OF LARGE LANGUAGE MODELS (LLMS) STATEMENT

• The authors used LLMs to improve manuscript clarity, structure, and grammar, but not for research ideation or core scientific content.Source: C THE USE OF LARGE LANGUAGE MODELS… · p. 14

• The reasoning-pattern analysis is intuitive, loosely analogical to human cognition, and preliminary rather than rigorously formalized.Source: D SHOWCASE ON SEQUENTIAL AND PARALLEL… · p. 15

Sequential Reasoning Pattern

• Mathematical examples show sequential substeps: MATH500 combines component ratios, while the GSM8K example sums exercise counts before calculating a percentage.Source: D SHOWCASE ON SEQUENTIAL AND PARALLEL… · pp. 15–16

Parallel Reasoning Pattern: Countdown

• Countdown permits parallel reasoning because inferring an operator between operands need not depend on resolving earlier pairs first.Source: D SHOWCASE ON SEQUENTIAL AND PARALLEL… · p. 17

Parallel Reasoning Pattern: Sudoku

• Sudoku is presented as benefiting from parallel or bidirectional reasoning, with constraints enumerated before the answer is deduced.Source: D SHOWCASE ON SEQUENTIAL AND PARALLEL… · p. 18