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

2026-06-30

AI Summary

This work addresses the high-cost trajectory slicing problem in reinforcement learning for diffusion-based large language models, which arises from the mismatch between stochastic masking during training and deterministic inference trajectories. To circumvent the need for expensive trajectory reconstruction, the authors propose an efficient training method featuring three key innovations: a τ-budget decoder that controls commitment risk at each decoding step, a monotonically decreasing block masking schedule, and an optimized objective combining sequence-level importance sampling with deterministic mask-level integration. Evaluated on MATH500 and GSM8K, the approach significantly outperforms TraceRL, and it also achieves performance gains on code generation benchmarks MBPP and HumanEval—surpassing larger baselines despite using a smaller model size.

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 .

Problem

• diffusion LLMs
• reinforcement learning
• random masking
• trajectory mismatch
• risk budgeting

Innovation

• risk-budgeted decoding
• random-masking RL
• trajectory-free training
• variance reduction
• diffusion LLMs