Grenzgeführte Strategieoptimierung für speichereffizientes Reinforcement Learning von Diffusion Large Language Models

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

October 13, 2025

papers.authors: Nianyi Lin, Jiajie Zhang, Lei Hou, Juanzi Li

cs.AI

papers.abstract

Eine zentrale Herausforderung bei der Anwendung von Reinforcement Learning (RL) auf große Diffusions-Sprachmodelle (dLLMs) liegt in der Unhandhabbarkeit ihrer Wahrscheinlichkeitsfunktionen, die für das RL-Ziel entscheidend sind und entsprechende Approximationen in jedem Trainingsschritt erfordern. Während bestehende Methoden die Log-Likelihoods durch ihre Evidenzuntergrenzen (ELBOs) mittels angepasster Monte-Carlo (MC)-Stichproben approximieren, müssen die Vorwärts-Berechnungsgraphen aller MC-Stichproben für die Gradientenberechnung nichtlinearer Terme im RL-Ziel beibehalten werden, was zu erheblichem Speicheraufwand führt. Diese Einschränkung begrenzt die machbaren Stichprobengrößen, was zu ungenauen Likelihood-Approximationen und letztlich zu einer Verzerrung des RL-Ziels führt. Um diese Beschränkung zu überwinden, schlagen wir Boundary-Guided Policy Optimization (BGPO) vor, einen speichereffizienten RL-Algorithmus, der eine speziell konstruierte Untergrenze des ELBO-basierten Ziels maximiert. Diese Untergrenze ist sorgfältig entworfen, um zwei Schlüsseleigenschaften zu erfüllen: (1) Linearität: Sie wird als lineare Summe formuliert, bei der jeder Term nur von einer einzelnen MC-Stichprobe abhängt, wodurch die Gradientenakkumulation über Stichproben ermöglicht und ein konstanter Speicherverbrauch sichergestellt wird; (2) Äquivalenz: Sowohl der Wert als auch der Gradient dieser Untergrenze sind gleich denen des ELBO-basierten Ziels im On-Policy-Training, was sie auch zu einer effektiven Approximation des ursprünglichen RL-Ziels macht. Diese Eigenschaften ermöglichen es BGPO, eine große MC-Stichprobengröße zu verwenden, was zu genaueren Likelihood-Approximationen und einer verbesserten Schätzung des RL-Ziels führt, was wiederum zu einer gesteigerten Leistung beiträgt. Experimente zeigen, dass BGPO bisherige RL-Algorithmen für dLLMs bei der Lösung von mathematischen Problemen, der Codegenerierung und Planungsaufgaben deutlich übertrifft.

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.