Boundary-Guided Policy Optimization for Memory-efficient RL of Diffusion
  Large Language Models
Authors: name: Nianyi Lin, name: Jiajie Zhang, name: Lei Hou, name: Juanzi Li
Organization: Z.ai

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