Boundary-Guided Policy Optimization for Memory-efficient RL of Diffusion
  Large Language Models
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.
68edd3f6de1fee572713a979
Nianyi Lin
68edd3f6de1fee572713a97a
Jiajie Zhang
68edd3f6de1fee572713a97b
Lei Hou
68edd3f6de1fee572713a97c
Juanzi Li
Boundary-Guided Policy Optimization for Memory-efficient RL of Diffusion
  Large Language Models
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.
Boundary-Guided Policy Optimization (BGPO) improves reinforcement learning for diffusion large language models by efficiently approximating likelihoods with a memory-efficient lower bound, enhancing performance in tasks like math problem solving, code generation, and planning.
reinforcement learning diffusion large language models likelihood functions evidence lower bounds Monte Carlo sampling gradient computation memory overhead Boundary-Guided Policy Optimization on-policy training