arXiv:2510.11683cs.LGcs.AI2025-10ACL被引 8

提出新算法提升扩散语言模型强化学习效率,内存占用恒定且效果更优。

Boundary-Guided Policy Optimization for Memory-efficient RL of Diffusion Large Language Models

  • 通过线性化采样项设计新型下界,实现梯度累积与固定内存消耗。
  • 支持大规模采样,显著提升似然估计精度与强化学习目标逼近能力。
  • 适用于数学推理、代码生成等任务,适合追求高效训练的AI研究者。

将强化学习(RL)应用于扩散大语言模型(dLLMs)的关键挑战在于其似然函数不可计算,需在训练中进行近似。现有方法通过定制蒙特卡洛(MC)采样使用证据下界(ELBO)近似对数似然,但因需保留所有采样样本以计算非线性项梯度,导致内存开销巨大,限制了采样规模,进而造成似然近似不准确和强化学习目标失真。为此,我们提出边界引导策略优化(BGPO),一种内存高效的强化学习算法,其最大化一个特别构造的基于ELBO目标的下界。该下界满足两个关键性质:(1)线性性:它是单个采样项的线性组合,每项仅依赖一个样本,支持跨样本梯度累积,确保内存消耗恒定;(2)等价性:在在线策略训练中,该下界的值与梯度与原始ELBO目标完全一致,因此也是原强化学习目标的有效近似。这一设计使BGPO可采用大规模采样,提升似然估计与强化学习目标估计质量,从而带来性能提升。实验表明,BGPO在数学问题求解、代码生成和规划任务上显著优于先前dLLM强化学习算法。代码与模型已开源于 https://github.com/THU-KEG/BGPO。

原文摘要 · Abstract (English)

A key challenge in applying reinforcement learning (RL) to diffusion large language models (dLLMs) is the intractability of their likelihood functions, which are essential for the RL objective, necessitating corresponding approximation during training. While existing methods approximate the log-likelihoods by their evidence lower bounds (ELBOs) via customized Monte Carlo (MC) sampling, they incur significant memory overhead due to the need to retain all MC samples for the gradient computation of non-linear terms in the RL objective, and thus restrict feasible sample sizes, leading to imprecise likelihood approximations and distorted RL objective. To address this, we propose \emph{Boundary-Guided Policy Optimization} (BGPO), a memory-efficient RL algorithm that maximizes a specially constructed lower bound of the ELBO-based objective. This lower bound is carefully designed to satisfy two key properties: (1) Linearity: it is a linear sum where each term depends only on a single MC sample, thereby enabling gradient accumulation across samples and ensuring constant memory usage; (2) Equivalence: Both the value and gradient of this lower bound are equal to those of the ELBO-based objective in on-policy training, making it also an effective approximation for the original RL objective. These properties allow BGPO to adopt a large MC sample size, improving likelihood approximations and RL objective estimation, which in turn leads to enhanced performance. Experiments show that BGPO significantly outperforms previous RL algorithms for dLLMs in math problem solving, code generation, and planning tasks. Our codes and models are available at \href{https://github.com/THU-KEG/BGPO}{https://github.com/THU-KEG/BGPO}.

强化学习扩散模型大模型内存优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。