Boundary-Guided Policy Optimalisatie voor Geheugenefficiënte Reinforcement Learning van Diffusion Grote Taalmodellen

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

October 13, 2025

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

cs.AI

Samenvatting

Een belangrijke uitdaging bij het toepassen van reinforcement learning (RL) op diffusion large language models (dLLMs) ligt in de onhanteerbaarheid van hun waarschijnlijkheidsfuncties, die essentieel zijn voor het RL-doel, wat een corresponderende benadering in elke trainingsstap noodzakelijk maakt. Hoewel bestaande methoden de log-waarschijnlijkheden benaderen via hun evidence lower bounds (ELBOs) met behulp van aangepaste Monte Carlo (MC) steekproeven, moeten de voorwaartse rekenkundige grafieken van alle MC-steekproeven worden behouden voor de gradientberekening van niet-lineaire termen in het RL-doel, wat resulteert in aanzienlijk geheugengebruik. Deze beperking beperkt de haalbare steekproefgroottes, wat leidt tot onnauwkeurige waarschijnlijkheidsbenaderingen en uiteindelijk het RL-doel vervormt. Om deze beperking te overwinnen, stellen we Boundary-Guided Policy Optimization (BGPO) voor, een geheugenefficiënt RL-algoritme dat een speciaal geconstrueerde ondergrens van het ELBO-gebaseerde doel maximaliseert. Deze ondergrens is zorgvuldig ontworpen om twee belangrijke eigenschappen te voldoen: (1) Lineariteit: het is geformuleerd in een lineaire som waarbij elke term alleen afhankelijk is van een enkele MC-steekproef, waardoor gradientaccumulatie over steekproeven mogelijk is en constant geheugengebruik wordt gegarandeerd; (2) Gelijkwaardigheid: zowel de waarde als de gradient van deze ondergrens zijn gelijk aan die van het ELBO-gebaseerde doel in on-policy training, waardoor het ook een effectieve benadering is voor het oorspronkelijke RL-doel. Deze eigenschappen stellen BGPO in staat om een grote MC-steekproefgrootte te gebruiken, wat resulteert in nauwkeurigere waarschijnlijkheidsbenaderingen en verbeterde RL-doelschatting, wat op zijn beurt leidt tot betere prestaties. Experimenten tonen aan dat BGPO aanzienlijk beter presteert dan eerdere RL-algoritmen voor dLLMs in het oplossen van wiskundige problemen, codegeneratie en plannings taken.

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.