用强化学习让大模型并行生成更高效又不失准。
LightningRL: Breaking the Accuracy-Parallelism Trade-off of Block-wise dLLMs via Reinforcement Learning
- 用强化学习动态选择高并行且准确的生成路径。
- 平均每轮生成7.32个词,最高达11.10,准确率不降。
- 适合追求高速生成的推理部署场景。
扩散语言模型(dLLMs)作为并行生成文本的新范式受到广泛关注,但现有方法普遍存在准确率与并行度之间的刚性权衡:通过激进并行解码提升每前向传播生成词数(TPF)常导致性能下降和生成不稳定。我们发现根本原因在于模型在高并行状态下难以应对累积的近似误差与局部扰动。为此提出LightningRL,一种面向预训练dLLMs的后训练框架,直接优化速度-质量帕累托前沿。该方法不强制统一并行化,而是利用强化学习识别并强化维持准确性的高并行轨迹。基于组相对策略优化(GRPO)框架,引入三项改进:(1) 通过逐奖励解耦归一化实现训练稳定;(2) 在正确轨迹上施加词级负对数似然正则化以锚定性能;(3) 基于TPF感知的动态采样策略提升训练效率。在数学与编程基准测试中,LightningRL持续突破帕累托前沿,保持竞争性任务准确率的同时显著提升并行度,平均TPF达7.32(MBPP数据集峰值达11.10)。代码已开源。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。