SLIM-RL: Risk-Budgeted Random-Masking RL for Diffusion LLMs Without Trajectory Slicing

Abstract: Reinforcement learning for diffusion LLMs (dLLMs) has largely moved to trajectory-aware methods. The current state of the art, TraceRL, holds that random masking is mismatched with the model's inference trajectory, and it reconstructs that trajectory during training by slicing each rollout into up to K/s trajectory-aligned training samples, a cost that grows with the block size K. We show that this mismatch can be mitigated without reconstructing the trajectory. Our method, SLIM-RL, bounds the commit risk of each rollout step with a tau-budget decoder, reducing aggregate commit risk in the training data. During optimization, SLIM-RL trains on these risk-controlled rollouts with a trace-free random-masking objective that adapts variance-reduction tools, combining sequence-level importance sampling, deterministic quadrature over masking levels under a mean-preserving, monotonically decreasing per-block mask schedule that we introduce. On SDAR-4B, SLIM-RL matches TraceRL's best MATH500 accuracy on only 0.46x its training samples at block size 16, improving over TraceRL by 6.32% on MATH500 and 11.05% on GSM8K under matched dynamic sampling. At block size 4, the 4B SLIM-RL surpasses the larger LLaDA-8B and Dream-7B dLLMs on math, exceeding LLaDA-8B by 10.76% on MATH500 while staying below the autoregressive Qwen2.5-7B. On code, it improves over TraceRL by 4.20% on MBPP and 3.65% on HumanEval. The tau-budget decoder transfers training-free across LLaDA, Dream, and SDAR. The source code is available at https://github.com/laolaorkkkkk/SLIM-RL .

• The paper introduces SLIM-RL, a trace-free RL framework using risk-budgeted random masking to eliminate the need for trajectory slicing.
• It deploys a τ-budget dynamic unmasking rule and variance-reduced objective, enhancing training efficiency and stability in diffusion LLMs.
• Experimental results show significant accuracy gains in mathematical reasoning and code generation tasks while reducing training data requirements.

SLIM-RL: Risk-Budgeted Random-Masking Reinforcement Learning for Diffusion LLMs Without Trajectory Slicing

Introduction

The paper "SLIM-RL: Risk-Budgeted Random-Masking RL for Diffusion LLMs Without Trajectory Slicing" (2607.00208) presents an optimized framework for reinforcement learning (RL) with diffusion LLMs (dLLMs). Unlike prior RL methods such as TraceRL that rely on trajectory-aligned slicing and decoding trace reconstruction, SLIM-RL introduces a trace-free RL pipeline based on risk-budgeted random masking. The method combines a novel $\tau$-budget dynamic-unmasking rule with a variance-reduced RL objective using deterministic quadrature and sequence-level importance sampling, obviating the need for explicit trajectory alignment during training. This design purportedly resolves critical inefficiencies in prior RL approaches for blockwise dLLMs and demonstrates consistent improvements over established baselines on mathematical reasoning and code generation tasks.

Background and Limitations of Trajectory-Aware RL in dLLMs

Diffusion LLMs generate text via iterative denoising, offering block-wise decomposition where blocks of tokens are decoded in parallel under bidirectional attention within each block and causal attention across blocks. RL has become essential for boosting dLLM reasoning ability. The state-of-the-art algorithm, TraceRL, advocates reconstructing the inference trajectory via per-step slicing of rollouts to mitigate the masking-objective and inference-trajectory mismatch. However, the slicing cost scales with block size, and compressing steps via shrinkage leads to degraded trajectory fidelity and persistent trade-offs.

Random masking, although conceptually simpler and inherently parallel, suffers from high gradient variance and trajectory mismatch, limiting its effectiveness unless systemically improved. The key challenge tackled in this work is whether trajectory reconstruction is inherently required to match or surpass trajectory-aware RL in dLLMs, especially in the context of blockwise models designed for scalable parallelization and cache efficiency.

SLIM-RL: Methodological Innovations

$\tau$-Budget Dynamic Unmasking

SLIM-RL introduces a $\tau$-budget decoder that generalizes over the conventional dynamic sampling approach, which relies on fixed per-position confidence thresholds for token commitment during denoising. Instead, the $\tau$-budget mechanism controls collective uncertainty by committing only the largest subset of above-threshold positions such that the total confidence-based uncertainty (sum of $1-p_i$ over committed positions) remains within a per-step budget.

This calibrated commit risk explicitly limits the expected number of decoding errors per step and induces safer rollouts. The method is training-free and is employed both during RL rollout collection and inference, demonstrating robust transferability across architectures (LLaDA, Dream, SDAR). Empirically, the $\tau$-budget decoder significantly reduces the expected number of wrong commitments per step compared to standard dynamic sampling, particularly as block size increases.

Figure 1: Dynamic sampling versus $\tau$-budget unmasking at one denoising step, illustrating selective commitment under a stepwise uncertainty constraint.

Variance-Reduced Random Masking Objective

Building on prior work that identified variance sources in RL for dLLMs, SLIM-RL employs a fully trace-free objective encompassing several variance-reducing strategies:

• Sequence-level Importance Sampling: Rather than token-wise ratios (which lead to high gradient variance proportionally with sequence length), SLIM-RL adopts sequence-level length-normalized ratios, mitigating noise in policy gradient estimation.
• Deterministic Quadrature over Masking Levels: Replacing stochastic sampling of masking level $t$ with deterministic Gauss–Legendre quadrature nodes further suppresses estimator variance for the random masking objective.
• Unnormalized Advantages: SLIM-RL utilizes unnormalized, group-centered advantages, yielding unbiased and robust updates in the presence of grouped rollout rewards.

Together, these components enable efficient and stable RL optimization, allowing trace-free training to match or surpass the best performance of trajectory-based methods under stricter compute and data budgets.

Monotonically Decreasing Per-Block Mask Schedule

To further concentrate the learning signal, SLIM-RL incorporates a mean-preserving, monotonically decreasing mask schedule over blocks. The masking rates for earlier blocks are set higher, focusing learning gradients on positions with maximal conditioning influence. This front-loads the effective RL signal to the most critical positions, empirically increasing reasoning accuracy over uniform or random masking schedules.

Empirical Evaluation

Comprehensive experiments are conducted on SDAR-4B and SDAR-1.7B-Chat dLLMs, evaluated on MATH500, GSM8K, MBPP, and HumanEval. Comparisons include LLaDA-8B, Dream-7B, and autoregressive references such as Qwen2.5-7B.

Figure 2: MATH500 accuracy versus cumulative training samples at block size 16: SLIM-RL achieves TraceRL's best accuracy using less than half as many samples.

Performance: Data Efficiency and Accuracy

SLIM-RL matches or exceeds TraceRL on all evaluated metrics:

• At block size 16 (SDAR-4B), SLIM-RL delivers a 6.32% absolute accuracy gain on MATH500 and 11.05% on GSM8K over TraceRL under matched decoding, and reaches TraceRL’s best result using only 0.46× the training samples.
• At block size 4, the 4B SLIM-RL model outperforms larger LLaDA-8B and Dream-7B models on math benchmarks, exceeding LLaDA-8B by 10.76% on MATH500 while maintaining a performance gap behind autoregressive Qwen2.5-7B.
• On code tasks, SLIM-RL improves over TraceRL by 4.20% on MBPP and 3.65% on HumanEval.

  Figure 3: Scaling data efficiency at 1.7B: SLIM-RL overtakes TraceRL on training accuracy, reaching TraceRL’s best with only 0.76× the training samples.

Parallelism and Generation Dynamics

SLIM-RL produces more parallel generations (higher tokens-per-forward) than TraceRL, with stability sustained throughout training:

Figure 4: Training dynamics at block size 16. SLIM-RL maintains higher tokens-per-forward and mask ratio compared to TraceRL; all methods exhibit generation length reduction over training.

Rollout Commit Risk

Empirically, the $\tau$-budget decoder commits fewer expected wrong tokens per denoising step than dynamic sampling, yielding improved rollout quality and more reliable RL gradients:

Figure 5: Expected wrong commitments per step: SLIM-RL (blue) consistently underperforms TraceRL (red) in error rate at both block sizes.

Ablation and Scaling Analysis

Component ablations reveal all major ingredients of SLIM-RL are vital for optimal performance: quadrature (using 3 vs. 2 points), sequence-level importance ratios, the per-block masking schedule, and risk-budgeted decoding. The advantage of SLIM-RL over TraceRL grows with block size, attesting to the robustness of trace-free variance reduction as parallelism increases. Notably, train-time and inference-time decoders can be decoupled without loss, and the $\tau$-budget rule generalizes seamlessly to other architectures, functioning as a drop-in alternative to dynamic sampling.

Theoretical and Practical Implications

The results demonstrate that explicit reconstruction of the model's decoding trajectory is not necessary to match or surpass trajectory-aware RL objectives at controlled compute and data cost, even for large and complex blockwise dLLMs. The adoption of a risk-budgeted decoding regime during RL rollout collection, in combination with modern variance-reduction techniques and a task-aligned masking schedule, is sufficient for robust RL-based reasoning improvements.

Practically, SLIM-RL enables more efficient RL training for diffusion LLMs—achieving higher accuracy, enhanced parallel decoding, and reduced data requirements. The training-free and architecture-agnostic nature of the $\tau$0-budget decoder further supports deployment across different dLLM designs.

Theoretically, these results reinforce the sufficiency of variance-reduced, trace-free objectives even when inference and training trajectories diverge, provided rollout commit risk is controlled. This weakens prior assumptions regarding the necessity of trajectory fidelity for optimal RL, particularly for blockwise and parallel generative models.

Future Directions

Possible extensions include: scaling SLIM-RL to much larger block sizes and longer outputs relevant for real-world applications, integrating adaptive risk-budgets, exploring more sophisticated advantage estimation or importance weighting, and combining with other policy regularization strategies. The interplay of commit risk, decoding parallelism, and training efficiency is fertile ground for further research, particularly as diffusion-based architectures continue to mature.

Conclusion

SLIM-RL establishes that trajectory reconstruction and slicing are not prerequisites for high-performance RL in diffusion LLMs. Instead, risk-budgeted rollout collection and variance-reduced, trace-free random masking objectives suffice for—indeed, outperform—trajectory-aware RL paradigms, yielding significant improvements in both data efficiency and output accuracy. The approach facilitates scalable, parallel RL for dLLMs with flexible decoding strategies and offers a blueprint for future developments in efficient RL on masked generative architectures.