Optimisation de Politique Guidée par les Frontières pour un Apprentissage par Renforcement Économe en Mémoire des Modèles de Langage à Grande Échelle de Diffusion

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

Un défi majeur dans l'application de l'apprentissage par renforcement (RL) aux modèles de langage à grande échelle basés sur la diffusion (dLLMs) réside dans l'intractabilité de leurs fonctions de vraisemblance, qui sont essentielles pour l'objectif de RL, nécessitant ainsi une approximation correspondante à chaque étape d'entraînement. Bien que les méthodes existantes approchent les log-vraisemblances par leurs bornes inférieures de l'évidence (ELBOs) via un échantillonnage Monte Carlo (MC) personnalisé, les graphes de calcul directs de tous les échantillons MC doivent être conservés pour le calcul du gradient des termes non linéaires dans l'objectif de RL, ce qui entraîne une surcharge mémoire significative. Cette contrainte limite les tailles d'échantillons réalisables, conduisant à des approximations de vraisemblance imprécises et, finalement, à une distorsion de l'objectif de RL. Pour surmonter cette limitation, nous proposons l'Optimisation de Politique Guidée par les Bornes (BGPO), un algorithme de RL économe en mémoire qui maximise une borne inférieure spécialement construite de l'objectif basé sur ELBO. Cette borne inférieure est soigneusement conçue pour satisfaire deux propriétés clés : (1) Linéarité : elle est formulée sous forme de somme linéaire où chaque terme dépend uniquement d'un seul échantillon MC, permettant ainsi l'accumulation du gradient à travers les échantillons et garantissant une utilisation mémoire constante ; (2) Équivalence : la valeur et le gradient de cette borne inférieure sont égaux à ceux de l'objectif basé sur ELBO dans l'entraînement on-policy, en faisant également une approximation efficace de l'objectif de RL original. Ces propriétés permettent à BGPO d'adopter une grande taille d'échantillon MC, résultant en des approximations de vraisemblance plus précises et une meilleure estimation de l'objectif de RL, ce qui conduit à une performance améliorée. Les expériences montrent que BGPO surpasse significativement les algorithmes de RL précédents pour les dLLMs dans la résolution de problèmes mathématiques, la génération de code et les tâches de planification.

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.