SLIM-RL: Risk-Budgeted Random-Masking RL for Diffusion LLMs Without Trajectory Slicing
2026-06-30 • Computation and Language
Computation and LanguageArtificial IntelligenceMachine Learning
AI summaryⓘ
The authors introduce SLIM-RL, a new method for training diffusion large language models (dLLMs) with reinforcement learning. Unlike the existing TraceRL method that reconstructs model inference paths and grows costly with block size, SLIM-RL avoids this by controlling risk during training using a special decoder and a new random masking approach. Their experiments show SLIM-RL matches or outperforms TraceRL and other models on math and code tasks while using fewer training samples. The method also works across different dLLM architectures without retraining. The authors provide their code for others to try.
Reinforcement LearningDiffusion Large Language ModelsTraceRLRandom MaskingTrajectory AlignmentVariance ReductionImportance SamplingTau-budget DecoderBlock SizeMathematical Reasoning Tasks
Authors
Ruikang Zhao, Zhenting Wang, Han Gao, Ligong Han
Abstract
Reinforcement learning for diffusion large language models (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 .