通过时空剪枝提升扩散语言模型强化学习的效率与稳定性。
Efficient and Stable Reinforcement Learning for Diffusion Language Models
- 用静态先验压缩探索空间，跳过冗余生成步骤。
- 理论证明可降低似然估计方差，稳定策略更新。
- 在效率和准确率上均超越现有最佳方法。
强化学习（RL）对于激发基于扩散的大语言模型（dLLMs）的复杂推理能力至关重要。然而，将RL应用于dLLMs面临效率与稳定性方面的独特挑战。为此，我们提出时空剪枝（STP）框架，旨在同时提升dLLMs强化学习的效率与稳定性。STP通过以下方式压缩生成过程中的冗余：(1) 空间剪枝，利用静态先验约束探索空间；(2) 时间剪枝，跳过冗余的后期优化步骤。理论分析表明，STP严格降低了对数似然估计的方差，从而确保更稳定的策略更新。大量实验表明，STP在效率和准确性上均优于当前最优基线。代码已开源：https://github.com/Lolo1222/STP。
原文摘要 · Abstract (English)
Reinforcement Learning (RL) is crucial for unlocking the complex reasoning capabilities of Diffusion-based Large Language Models (dLLMs). However, applying RL to dLLMs faces unique challenges in efficiency and stability. To address these challenges, we propose Spatio-Temporal Pruning (STP), a framework designed to simultaneously improve the efficiency and stability of RL for dLLMs. STP compresses the redundancy in the generative process through: (1) \textit{spatial pruning}, which constrains the exploration space using static priors; and (2) \textit{temporal pruning}, which bypasses redundant late-stage refinement steps. Our theoretical analysis demonstrates that STP strictly reduces the variance of the log-likelihood estimation, thereby ensuring more stable policy updates. Extensive experiments demonstrate that STP surpasses state-of-the-art baselines in both efficiency and accuracy. Our code is available at https://github.com/Lolo1222/STP.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。