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

TL;DR

BGPO addresses memory bottlenecks in RL for diffusion-based LLMs by optimizing a linear, equivalent lower bound $\hat{\mathcal{J}}_{lb}(\theta)$ to the ELBO-based objective, enabling large Monte Carlo sample sizes $n_t$ without increasing memory usage. The authors prove equivalence to the original RL objective in on-policy training and demonstrate reduced gradient bias/variance, yielding improved performance on math, coding, and planning benchmarks with LLaDA-8B-Instruct. Empirical results show BGPO outperforming diffu-GRPO and VRPO-OL across tasks, with notable gains on Sudoku and Countdown and only modest increases in step-time. The work provides a practical, scalable approach to RL for dLLMs and releases open-source code for replication and further research.

Abstract

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}.

Figures (5)

• Figure 1: Left: Comparison of memory usage of previous ELBO-based RL method (VRPO-OL) and our BGPO using different Monte Carlo sample size $n_t$ for the RL objective approximation. The max response length is set to 512. Middle and right: Performance of LLaDAs with different RL algorithms on mathematical tasks.
• Figure 2: Training reward dynamics of BGPO, diffu-GRPO, and VRPO-OL across different tasks.
• Figure 3: Standard deviation (std) of gradients of different RL algorithms with different MC sampling size $n_t$. The std is normalized by the absolute value of each parameter to unify the scale.
• Figure 4: Gradient bias (normalized by the absolute value of parameter) with different MC sampling size $n_t$.
• Figure 5: Training speed comparison between baselines.

Theorems & Definitions (2)

• Lemma 1
• Lemma 2