Ottimizzazione delle Politiche Guidata dai Confini per un Apprendimento per Rinforzo Efficiente in Memoria di Modelli Linguistici di Diffusione di Grandi Dimensioni

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

October 13, 2025

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

cs.AI

Abstract

Una delle principali sfide nell'applicare l'apprendimento per rinforzo (RL) ai modelli linguistici di diffusione su larga scala (dLLM) risiede nell'intrattabilità delle loro funzioni di verosimiglianza, che sono essenziali per l'obiettivo dell'RL, rendendo necessaria un'approssimazione corrispondente in ogni fase di addestramento. Mentre i metodi esistenti approssimano i log-verosimiglianze tramite i loro limiti inferiori di evidenza (ELBO) utilizzando un campionamento Monte Carlo (MC) personalizzato, i grafici computazionali in avanti di tutti i campioni MC devono essere mantenuti per il calcolo del gradiente dei termini non lineari nell'obiettivo dell'RL, risultando in un significativo sovraccarico di memoria. Questo vincolo limita le dimensioni dei campioni fattibili, portando a approssimazioni imprecise della verosimiglianza e, in ultima analisi, distorcendo l'obiettivo dell'RL. Per superare questa limitazione, proponiamo l'ottimizzazione della politica guidata dai confini (BGPO), un algoritmo RL efficiente in termini di memoria che massimizza un limite inferiore appositamente costruito dell'obiettivo basato su ELBO. Questo limite inferiore è progettato con cura per soddisfare due proprietà chiave: (1) Linearità: è formulato come una somma lineare in cui ogni termine dipende solo da un singolo campione MC, consentendo così l'accumulo del gradiente tra i campioni e garantendo un uso costante della memoria; (2) Equivalenza: sia il valore che il gradiente di questo limite inferiore sono uguali a quelli dell'obiettivo basato su ELBO nell'addestramento on-policy, rendendolo anche un'approssimazione efficace per l'obiettivo originale dell'RL. Queste proprietà permettono a BGPO di adottare una dimensione del campione MC ampia, risultando in approssimazioni più accurate della verosimiglianza e in una migliore stima dell'obiettivo dell'RL, che a sua volta porta a prestazioni migliorate. Gli esperimenti dimostrano che BGPO supera significativamente i precedenti algoritmi RL per i dLLM nella risoluzione di problemi matematici, nella generazione di codice e nelle attività di pianificazione.

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.