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

2026-03-04

arXiv.org

AI Summary

This work addresses the challenge of balancing generation quality and parallelism in block-wise speculative decoding for large language models (dLLMs), where increased parallel generation often leads to degraded accuracy and instability. To overcome this trade-off, the authors propose LightningRL, a novel framework that leverages reinforcement learning to optimize the speed–quality Pareto frontier of dLLMs for the first time. Built upon Group Relative Policy Optimization (GRPO), LightningRL introduces throughput-per-fidelity (TPF)-aware sampling, decoupled reward normalization, and token-level negative log-likelihood regularization to dynamically reinforce accurate generation trajectories under high parallelism. Experiments demonstrate that LightningRL significantly advances the Pareto frontier on mathematical and code generation tasks, achieving an average of 7.32 tokens per forward pass (peaking at 11.10 on MBPP) while maintaining competitive accuracy.

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.

Problem

• accuracy-parallelism trade-off
• diffusion LLMs
• parallel token generation
• generation instability
• block-wise decoding

Innovation

• reinforcement learning
• diffusion LLMs
• parallel token generation
• Pareto frontier optimization
• GRPO