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

📄 arXiv: 2510.11683v2 📥 PDF

作者: Nianyi Lin, Jiajie Zhang, Lei Hou, Juanzi Li

分类: cs.LG, cs.AI, cs.CL

发布日期: 2025-10-13 (更新: 2025-10-14)

🔗 代码/项目: GITHUB

💡 一句话要点

提出边界引导策略优化(BGPO)，解决扩散大语言模型RL训练中的内存瓶颈问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 扩散模型 强化学习 大语言模型 策略优化 内存效率

📋 核心要点

  1. 现有方法在扩散大语言模型强化学习中，使用蒙特卡洛采样近似似然函数，但梯度计算需要保留大量计算图，导致内存开销大。
  2. BGPO通过构建ELBO的下界，保证了线性性和等价性，从而允许使用更大的蒙特卡洛样本量，提高似然近似的精度。
  3. 实验表明，BGPO在数学问题解决、代码生成和规划任务中，显著优于之前的强化学习算法，验证了其有效性。

📝 摘要（中文）

将强化学习(RL)应用于扩散大语言模型(dLLMs)的一个关键挑战在于其似然函数难以处理，这对于RL目标至关重要，需要在每个训练步骤中进行相应的近似。现有方法通过定制的蒙特卡洛(MC)采样，利用证据下界(ELBO)来近似对数似然，但所有MC样本的前向计算图需要被保留，以便计算RL目标中非线性项的梯度，导致显著的内存开销。这种约束限制了可行的样本大小，导致不精确的似然近似，最终扭曲了RL目标。为了克服这个限制，我们提出了边界引导策略优化(BGPO)，这是一种内存高效的RL算法，它最大化了基于ELBO的目标的特殊构造的下界。这个下界经过精心设计，以满足两个关键属性：(1)线性性：它被公式化为一个线性总和，其中每一项仅取决于单个MC样本，从而能够跨样本进行梯度累积，并确保恒定的内存使用；(2)等价性：在on-policy训练中，这个下界的值和梯度都等于基于ELBO的目标的值和梯度，使其也成为原始RL目标的有效近似。这些属性允许BGPO采用大的MC样本大小，从而产生更精确的似然近似和改进的RL目标估计，进而提高性能。实验表明，BGPO在数学问题解决、代码生成和规划任务中显著优于以前的dLLM的RL算法。

🔬 方法详解

问题定义：论文旨在解决将强化学习应用于扩散大语言模型时，由于似然函数难以处理而导致的内存瓶颈问题。现有方法通过蒙特卡洛采样近似对数似然，但为了计算梯度，需要保留所有样本的前向计算图，导致内存占用随样本数量线性增长，限制了样本数量，降低了似然近似的精度，最终影响强化学习的效果。

核心思路：论文的核心思路是构建一个基于证据下界(ELBO)的RL目标的新下界。这个下界被设计成具有两个关键性质：线性性和等价性。线性性保证了梯度计算可以跨样本累积，从而实现恒定的内存占用。等价性保证了这个下界在on-policy训练中与原始ELBO目标具有相同的值和梯度，从而保证了优化方向的一致性。通过这种方式，可以使用更大的蒙特卡洛样本量，提高似然近似的精度，同时避免内存溢出。

技术框架：BGPO算法的整体框架如下：首先，使用扩散大语言模型生成多个样本。然后，使用精心设计的下界目标函数来近似原始的RL目标。该下界目标函数具有线性性和等价性，允许使用大的蒙特卡洛样本量进行梯度估计，同时保持内存占用恒定。最后，使用梯度上升等优化算法来更新策略，提高模型的性能。

关键创新：BGPO的关键创新在于构建了一个既具有线性性又具有等价性的ELBO下界。线性性保证了内存效率，等价性保证了优化方向的正确性。这种巧妙的设计使得BGPO能够克服现有方法的内存瓶颈，提高似然近似的精度，从而提升强化学习的效果。

关键设计：BGPO的关键设计在于下界目标函数的具体形式。该函数被设计成一个线性总和，其中每一项只依赖于一个蒙特卡洛样本。此外，该函数还包含一些边界项，用于保证其与原始ELBO目标在on-policy训练中的等价性。具体的函数形式和边界项的选择需要根据具体的扩散大语言模型和强化学习任务进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，BGPO在数学问题解决、代码生成和规划任务中显著优于之前的强化学习算法。例如，在数学问题解决任务中，BGPO的性能提升了XX%，在代码生成任务中，BGPO生成的代码质量提升了YY%。这些结果验证了BGPO的有效性和优越性。

🎯 应用场景

该研究成果可广泛应用于需要利用强化学习优化扩散大语言模型的场景，例如文本生成、图像生成、代码生成、数学问题求解、机器人控制等。通过提高训练效率和模型性能，可以降低开发成本，提升用户体验，并推动相关领域的发展。

📄 摘要（原文）

A key challenge in applying reinforcement learning (RL) to diffusion large language models (dLLMs) lies in the intractability of their likelihood functions, which are essential for the RL objective, necessitating corresponding approximation in each training step. While existing methods approximate the log-likelihoods by their evidence lower bounds (ELBOs) via customized Monte Carlo (MC) sampling, the forward computational graphs of all MC samples need to be retained for the gradient computation of non-linear terms in the RL objective, resulting in significant memory overhead. This constraint restricts feasible sample sizes, leading to imprecise likelihood approximations and ultimately distorting the RL objective. To overcome this limitation, 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 formulated in 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, resulting in more accurate likelihood approximations and improved 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}.