PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yanzhe Hu, Yijie Jin, Pengfei Liu, Kai Yu, Zhijie Deng

arXiv:2603.13319v2cs.LG

TL;DR

Existing block-wise dLLMs struggle to increase parallelism without sacrificing accuracy, because aggressive decoding can degrade generation quality. LightningRL post-trains pretrained dLLMs with GRPO-based reinforcement learning and dLLM-specific normalization, NLL anchoring, and TPF-aware sampling. Across mathematical and coding benchmarks, it advances the speed–quality frontier, achieving competitive accuracy with 7.32 average TPF and 11.10 peak TPF on MBPP.

  • Problem

    Existing block-wise dLLMs face an accuracy–parallelism trade-off in which higher TPF often degrades task accuracy and generation quality.

  • Method

    LightningRL uses GRPO-based reinforcement learning to reinforce accurate, highly parallel trajectories with decoupled reward normalization, NLL regularization, and TPF-aware dynamic sampling.

  • Results

    Across math and code benchmarks, LightningRL consistently advances the Pareto frontier, reaching 7.32 average TPF and 11.10 TPF on MBPP while maintaining competitive accuracy.

  • Takeaways & Limitations

    LightningRL reconciles generation speed and accuracy under high-parallelism constraints, supporting practical high-throughput dLLM deployment.

  • Takeaways & Limitations

    The ablation discussion identifies token-level optimization proxies as introducing bias and inconsistency.

Abstract

from arXiv · show

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.

1. Introduction

Block-wise dLLMs enable intra-block parallelism but face an accuracy–parallelism trade-off: higher TPF often degrades quality. LightningRL uses reinforcement learning to optimize accurate, highly parallel trajectories, with GRPO modifications and strong benchmark results.

  • 1. Introduction: Block-wise dLLMs generate sequential blocks while refining tokens within each block in parallel, combining long-range coherence with intra-block parallelism.They are commonly constructed by adapting pretrained autoregressive models for block-wise denoising.
  • 1. Introduction: Increasing tokens per forward through aggressive parallel decoding often lowers task accuracy, leaving existing dLLMs behind speculative decoding when speed and quality are considered together.Training-free sampling and distillation can raise TPF and TPS, but typically incur substantial quality degradation.
  • 1. Introduction: LightningRL formulates post-training as reinforcement learning that shifts probability mass toward trajectories that are both highly parallelizable and accurate, using outcome accuracy and TPF supervision.The framework is implemented on Group Relative Policy Optimization (GRPO).
  • 1. Introduction: LightningRL adds decoupled reward normalization, token-level NLL regularization on correct trajectories, and TPF-aware dynamic sampling to stabilize and improve dLLM training.These modifications address heterogeneous reward scales, anchor model performance, and improve training efficiency.
  • 1. Introduction: The evaluation reports accuracy, TPF as a parallelism measure, and AUP as a summary of the speed–quality trade-off across math and code benchmarks.The experiments apply LightningRL post-training to the representative SDAR block-wise dLLM.
  • 1. Introduction: 437.8 AUP at 7.32 average TPF, reaching 11.10 TPF on MBPP, demonstrates LightningRL’s superior speed–quality profile over established baselines.The LightningRL-8B model with block size 32 substantially surpasses d3LLM and EAGLE-3 while maintaining task accuracy.

2. Preliminaries

The preliminaries model dLLM generation as an MDP over partially masked sequences, with denoising transitions that can decode multiple tokens in parallel. Block-wise dLLMs partition sequences into sequential blocks, while GRPO adapts policy optimization to intermediate noisy states and parallel masked positions.

  • 2. Preliminaries: dLLM generation is represented as an MDP whose states are partially decoded token sequences, beginning from a fully masked state and evolving through denoising transitions.Each state contains tokens or [MASK] symbols, and the model distribution determines the next state.
  • 2. Preliminaries: The Markov property makes the next state depend only on the current state and input prompt, while confidence-driven decoding accepts multiple high-confidence tokens in one iteration.Rejected tokens are reset to [MASK] for subsequent refinement.
  • 2. Preliminaries: Block-wise dLLMs partition sequences into contiguous blocks generated sequentially, with tokens inside each block decoded in parallel.This design addresses vanilla dLLM limitations and can be adapted efficiently from pretrained autoregressive LLMs.
  • 2. Preliminaries: For dLLMs, GRPO maps the policy to the conditional denoising distribution and the trajectory to intermediate noisy states rather than sequential autoregressive token histories.Gradients propagate through parallel masked positions over multiple denoising steps conditioned on intermediate noisy states.
  • 2. Preliminaries: GRPO samples groups of trajectories, computes terminal rewards and advantages, and maximizes advantage-weighted probabilities of valid denoising actions.The dLLM objective sums parallel token predictions across denoising steps and includes a KL penalty against a reference model.
  • 2. Preliminaries: The trajectory length counts the total number of parallel token predictions performed across denoising steps, while ε provides numerical stability in the objective.These quantities connect GRPO’s trajectory optimization to dLLM parallel generation.

3. LightningRL: Breaking the Accuracy–Parallelism Trade-off

LightningRL addresses reward collapse and policy drift in multi-objective dLLM reinforcement learning by co-optimizing accuracy and inference parallelism. Its three modifications stabilize training, anchor correctness, and improve sampling, producing more stable optimization signals.

  • Motivation: Multi-objective dLLM RL can suffer reward collapse when raw accuracy and speed rewards are combined, causing policy drift and degraded generation capability.A coarse discrete accuracy reward can dominate a fine-grained speed reward, making behaviors indistinguishable within groups.
  • Overview: LightningRL combines independent reward normalization, correctness anchoring, and TPF-aware sampling to co-optimize accuracy and inference parallelism.The framework targets the speed–quality frontier rather than uniformly forcing aggressive parallel decoding.
  • 3.1. Decoupled Normalization for Group Rewards: Decoupled normalization reduces advantage collapse and improves optimization stability by independently standardizing accuracy and speed rewards before aggregation.Final batch-wise normalization controls update scale without erasing group-relative distinctions.
  • 3.2. Token-Level NLL for Accuracy Anchoring: Token-level NLL on verifier-correct trajectories converts sparse sequence-level successes into dense supervision that anchors correctness and limits policy drift.The combined objective uses GRPO for relative preference learning and NLL as a self-imitation anchor; when no correct trajectories exist, the NLL term is zero.
  • 3.3. Dynamic Sampling for Efficient Policy Optimization: Dynamic sampling accepts prompts only when sampled trajectories have sufficient TPF diversity and at least one correct outcome, producing denser policy-gradient signals.Filtering near-tie groups improves convergence stability and reduces reward collapse under the same configuration.

4. Experiments

LightningRL is evaluated across math and code benchmarks using accuracy, TPF, and AUP, with results showing improved speed–quality trade-offs and stable training. Ablations indicate that its three components contribute materially to performance and parallelism.

  • Setup: Evaluation uses SDAR models with math and code datasets, and compares accuracy, TPF, and AUP across GSM8K, MATH500, HumanEval, and MBPP.The setup includes SDAR-8B-b32, MATH, GSM8K, and PrimeIntellect for training, with benchmark-specific shot settings for evaluation.
  • Main Results: 412.9 AUP and 11.10 TPF on MBPP surpass TraceRL’s 144.2 AUP and 2.50 TPF, while LightningRL also reaches 507.5 AUP on GSM8K and 409.2 on MATH500.The method outperforms DiRL and TraceRL across four benchmarks, with especially strong gains in coding tasks.
  • Training Stability: LightningRL avoids the training collapse observed for TraceRL, maintaining the accuracy signal while improving speed during optimization.The reported training curves attribute TraceRL’s collapse to speed optimization eroding correctness, whereas LightningRL converges stably.
  • Decoding Behavior: Most LightningRL samples finish within approximately 100 decoding steps, unlike SDAR’s heavy tail extending beyond 850 steps.The shorter decoding behavior is associated with higher throughput during active decoding and fewer long-tail steps.
  • Benchmark Results: 437.8 AUP at 7.32 average TPF and 71.1% accuracy lets LightningRL-8B-b32 outperform SDAR, EAGLE-3, and d3LLM at comparable accuracy.SDAR-8B-b32 achieves 211.2 AUP and 3.12 TPF with nearly identical average accuracy, while EAGLE-3 reaches 252.5 average AUP and 5.63 average TPF.
  • Ablation Study: The full method reaches 90.3% accuracy, 5.58 TPF, and 507.5 AUP on GSM8K, while removing components causes clear performance drops.Removing NLL reduces accuracy to 80.7%, decoupled normalization reduces accuracy to 85.3% and TPF to 4.96, and TPF-aware filtering reduces accuracy to 87.2%.
  • Loss Reduction Ablation: Seq-Tok-Tok achieves the best accuracy and AUP among the tested loss-reduction strategies, whereas pure token-level reduction lowers accuracy to 80.0%.The strategies change implicit sample weighting under varying sequence lengths, substantially affecting training dynamics.

5. Related Work

Prior work accelerates dLLMs through caching, decoding strategies, distillation, and trajectory-level reinforcement learning. LightningRL differs by treating parallelism as a first-class training objective alongside generation quality.

  • dLLM Acceleration: Recent dLLM acceleration methods reduce per-step computation or increase effective parallelism through caching, confidence-aware decoding, unmasking, sampling, and learned parallel decoding.These approaches target inference efficiency through diffusion-compatible caching and improved denoising or sampling strategies.
  • RL for dLLMs: Recent reinforcement-learning approaches optimize dLLM denoising trajectories, but existing frameworks prioritize generation quality and treat parallelism as an inference-time adjustment.The cited directions include step-aligned scheduling and trace-aware reinforcement learning.

6. Conclusion

LightningRL is an RL framework that optimizes the speed–quality trade-off in diffusion language models by mitigating error amplification during aggressive parallel decoding. On math and code benchmarks, it achieves higher AUP under high-parallelism constraints and supports practical high-throughput dLLM deployment.

  • LightningRL mitigates error amplification during aggressive parallel decoding, reconciling generation speed and accuracy in diffusion language models.
  • LightningRL achieves higher AUP under high-parallelism constraints on mathematics and code generation benchmarks.
  • The results support practical, high-throughput deployment of diffusion language models.

Impact Statement

The paper identifies potential societal impacts from improving diffusion language model accuracy and inference efficiency through reinforcement learning but does not discuss them in detail.

  • The work may have societal impacts through more accurate and efficient Diffusion Large Language Models, but it provides no detailed ethical analysis.

A. Discussion on Value Model Incorporation

Adding a learned value model to LightningRL reduces both accuracy and TPF on GSM8K. The authors conjecture that critic errors become unstable under block-wise decoding, especially when rewards are sparse and state transitions are abrupt.

  • The value-model variant shows a clear drop in both accuracy and TPF on GSM8K.Table 5 compares performance with and without value-model incorporation, using each model’s own accuracy at TPF≈1 as ymax for AUP.
  • The authors conjecture that highly non-smooth block-wise state transitions make the critic unreliable, causing advantage sign flips.Approximation errors in the value function can dominate weak intermediate signals and lead to updates that penalize beneficial actions or reinforce suboptimal ones.
  • Confidence-driven block decoding and remasking exacerbate critic instability by abruptly changing many tokens in a single step.These abrupt partial-sequence changes make it harder for a scalar value function to learn a coherent baseline over the visited state manifold.

B.1. Training Details

The experiments use SDAR models with dynamic block-wise decoding, combined accuracy-and-speed rewards, filtered rollout groups, PPO-style policy optimization, and distributed training. Scalability and throughput evaluations compare LightningRL against SDAR and prior baselines across model scales, block sizes, datasets, and H100 inference settings.

  • Training Configuration: Training uses SDAR configurations with block size 32, dynamic decoding at ϕ = 0.9, 32 denoising steps per block, and maximum response length 8192.The rollout setup samples 128 tasks and generates 32 responses per iteration.
  • Reward and Filtering Design: Rewards combine accuracy and speed signals, while filtering retains groups with at least one correct response and sufficient TPF variance.Mathematics rewards use binary answer-equivalence outcomes, and the default TPF-variance threshold is 0.01.
  • Systems Configuration: Distributed training uses AdamW, warmup, gradient clipping, checkpointing, and multi-GPU DeepSpeed configurations.The reported example uses one node with eight H200 GPUs and ZeRO-1 optimization.
  • Optimization Procedure: Policy optimization uses grouped rollouts, decoupled normalization, PPO-style clipped objectives, KL regularization, and NLL regularization.The algorithm initializes grouped data, filters rollout groups, computes group-relative advantages, and updates the policy with the combined loss.
  • Scalability Evaluation: LightningRL’s scalability is evaluated across model scale and block size on four datasets under identical SDAR comparison settings.Table 6 groups results by model scale and block size, with AUP normalized using each model’s own accuracy at TPF≈1.
  • Throughput Evaluation: LightningRL achieves exceptional inference speed and significantly outperforms baselines in single-device H100 TPS benchmarks.Table 7 reports tokens-per-second performance on GSM8K using SGLang with tensor parallel size 1.

C.3. Training Dynamics Details

LightningRL avoids the training collapse seen with TraceRL while preserving accuracy signals and improving decoding speed. On GSM8K, it remains robust as TPF increases, unlike SDAR’s sharp accuracy decline.

  • LightningRL converges stably and maintains a better efficiency–accuracy frontier, whereas TraceRL experiences objective drift, reward collapse, and limited speed gains.TraceRL’s speed optimization progressively erodes correctness until training collapses; LightningRL avoids this behavior.
  • LightningRL sustains higher decoding throughput and more synchronized termination during training than TraceRL.These training dynamics accompany LightningRL’s stable convergence and improved efficiency–accuracy frontier.
  • At TPF = 7, LightningRL retains 87.5% accuracy, while SDAR falls to 71.2% after dropping from 87.4% at TPF = 4.Both methods degrade as TPF rises from 1 to 7, but SDAR is substantially more vulnerable at larger forward steps.

C.5. Hyperparameter Sensitivity

LightningRL remains robust across moderate hyperparameter ranges, with its default configuration achieving the strongest tested accuracy–parallelism trade-off. NLL anchoring, TPF-aware filtering, sampling temperature, and group size each affect the quality of learning signals and final performance.

  • Removing NLL anchoring reduces accuracy to 84.7% and AUP to 385.7, while µ = 0.1 achieves the highest AUP of 507.5.Increasing µ beyond 0.1 gradually lowers accuracy because the NLL term restricts exploration and dominates the policy gradient.
  • Disabling TPF-aware filtering degrades accuracy and AUP, while δ = 0.01 outperforms larger thresholds by preserving more sample diversity.Larger thresholds discard more groups and produce slightly lower AUP despite maintaining reasonable accuracy.
  • Sampling temperature 1.0 gives the best rollout quality and preserves exploration needed for diverse candidates and informative within-group comparisons.Rollout accuracy improves substantially as temperature increases within the tested range.
  • Larger groups increase nonzero advantages, but very large groups raise memory and computational costs despite only mild changes in collapse ratio.A moderately large group enriches effective advantage signals without severe collapse.
  • The default setting µ = 0.1, δ = 0.01, temperature = 1.0, G = 32 achieves the strongest overall accuracy–parallelism trade-off among tested values.The study reports robust performance across a moderate range of all four hyperparameters rather than dependence on one sensitive configuration.
Loading 2603.13319v2…