Optimización de Políticas Guiada por Límites para el Aprendizaje por Refuerzo Eficiente en Memoria de Modelos de Lenguaje de Difusión a Gran Escala

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

October 13, 2025

Autores: Nianyi Lin, Jiajie Zhang, Lei Hou, Juanzi Li

cs.AI

Resumen

Un desafío clave al aplicar el aprendizaje por refuerzo (RL) a los modelos de lenguaje de difusión a gran escala (dLLMs) radica en la intratabilidad de sus funciones de verosimilitud, las cuales son esenciales para el objetivo de RL, lo que requiere una aproximación correspondiente en cada paso de entrenamiento. Si bien los métodos existentes aproximan los log-verosimilitudes mediante sus límites inferiores de evidencia (ELBOs) a través de muestreo de Monte Carlo (MC) personalizado, los grafos computacionales hacia adelante de todas las muestras de MC deben retenerse para el cálculo del gradiente de los términos no lineales en el objetivo de RL, lo que resulta en un sobrecarga significativa de memoria. Esta restricción limita los tamaños de muestra factibles, llevando a aproximaciones de verosimilitud imprecisas y, en última instancia, distorsionando el objetivo de RL. Para superar esta limitación, proponemos Boundary-Guided Policy Optimization (BGPO), un algoritmo de RL eficiente en memoria que maximiza un límite inferior especialmente construido del objetivo basado en ELBO. Este límite inferior está cuidadosamente diseñado para satisfacer dos propiedades clave: (1) Linealidad: está formulado como una suma lineal donde cada término depende únicamente de una sola muestra de MC, permitiendo así la acumulación de gradientes entre muestras y asegurando un uso constante de memoria; (2) Equivalencia: tanto el valor como el gradiente de este límite inferior son iguales a los del objetivo basado en ELBO en el entrenamiento on-policy, lo que lo convierte también en una aproximación efectiva para el objetivo original de RL. Estas propiedades permiten que BGPO adopte un tamaño de muestra de MC grande, resultando en aproximaciones de verosimilitud más precisas y una mejor estimación del objetivo de RL, lo que a su vez conduce a un rendimiento mejorado. Los experimentos muestran que BGPO supera significativamente a los algoritmos de RL anteriores para dLLMs en la resolución de problemas matemáticos, generación de código y tareas de planificación.

English

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.