LightningRL: Breaking the Accuracy-Parallelism Trade-off of Block-wise dLLMs via Reinforcement Learning

Abstract

Diffusion Large Language Models (dLLMs) have emerged as a promising paradigm for parallel token generation, with block-wise variants garnering significant research interest. Despite their potential, existing dLLMs typically suffer from a rigid accuracy-parallelism trade-off: increasing the number of tokens per forward (TPF) via aggressive parallel decoding often leads to performance degradation and increased generation instability. We identify that this limitation stems from the model's inability to navigate high-parallelism regimes where approximation errors and local corruptions accumulate, ultimately undermining the reliability of parallel generation. To address this, we propose LightningRL, a post-training framework designed to directly optimize the speed-quality Pareto frontier of pre-trained dLLMs. Instead of forcing uniform parallelization, our approach leverages reinforcement learning to identify and reinforce high-parallelism trajectories that maintain generation accuracy. Built upon the Group Relative Policy Optimization (GRPO) framework, LightningRL introduces several enhancements tailored for dLLMs: (1) stabilized training via per-reward decoupled normalization; (2) token-level negative log-likelihood (NLL) regularization on correct trajectories to anchor model performance; and (3) a dynamic sampling strategy with TPF-aware filtering to enhance training efficiency. Experimental results across mathematical and coding benchmarks demonstrate that LightningRL consistently advances the Pareto frontier, achieving competitive task accuracy while significantly increasing parallelism, reaching an average TPF of 7.32 (with a peak of 11.10 on the MBPP dataset). Our code is available at https://github.com/SJTU-DENG-Lab/LightningRL.

Figures (8)

• Figure 1: Evaluation results of LightningRL and the baselines. LightningRL achieves superior accuracy on math and code benchmarks while accelerating parallel decoding to an average of 7.32 tokens per forward (TPF) and 497.9 accuracy under parallelism (AUP) d3llm, significantly outperforming baselines including Eagle-3 eagle3 and Fast-dLLM-v2 fast-dllm-v2.
• Figure 2: Overview of LightningRL. LightningRL samples a group of decoding trajectories per prompt, applies per-reward decoupled normalization to preserve within-group ranking under heterogeneous scales. The policy is optimized with a GRPO-style objective plus a token-level NLL anchor. The bottom panel shows the resulting shift toward the fastest correct trajectory, improving TPF without degrading accuracy.
• Figure 3: Per-reward decoupled normalization improves training stability. It reduces signal collapse (a) and yields more stable reward optimization (b) under the same training setup.
• Figure 4: Token-level NLL loss anchors the accuracy objective on GSM8K. Compared with training without the token-level NLL term, it maintains a higher accuracy reward and mitigates late-stage drift in the accuracy signal under the same setup.
• Figure 5: Dynamic sampling improves robustness of total-reward optimization on GSM8K. Relative to training without dynamic sampling, it reduces reward collapse under the same setup.