Otimização de Políticas Guiada por Limites para RL Eficiente em Memória de Modelos de Linguagem de Difusão em Grande 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

Resumo

Um desafio fundamental na aplicação de aprendizado por reforço (RL) a modelos de linguagem de difusão em larga escala (dLLMs) reside na intratabilidade de suas funções de verossimilhança, que são essenciais para o objetivo do RL, exigindo uma aproximação correspondente em cada etapa de treinamento. Embora os métodos existentes aproximem as log-verossimilhanças por seus limites inferiores de evidência (ELBOs) por meio de amostragem de Monte Carlo (MC) personalizada, os grafos computacionais diretos de todas as amostras MC precisam ser retidos para o cálculo do gradiente dos termos não lineares no objetivo do RL, resultando em um overhead significativo de memória. Essa restrição limita os tamanhos de amostra viáveis, levando a aproximações imprecisas de verossimilhança e, por fim, distorcendo o objetivo do RL. Para superar essa limitação, propomos o Boundary-Guided Policy Optimization (BGPO), um algoritmo de RL eficiente em memória que maximiza um limite inferior especialmente construído do objetivo baseado em ELBO. Esse limite inferior é cuidadosamente projetado para satisfazer duas propriedades-chave: (1) Linearidade: ele é formulado em uma soma linear onde cada termo depende apenas de uma única amostra MC, permitindo assim a acumulação de gradientes entre amostras e garantindo um uso constante de memória; (2) Equivalência: tanto o valor quanto o gradiente desse limite inferior são iguais aos do objetivo baseado em ELBO no treinamento on-policy, tornando-o também uma aproximação eficaz para o objetivo original do RL. Essas propriedades permitem que o BGPO adote um tamanho grande de amostra MC, resultando em aproximações de verossimilhança mais precisas e uma estimativa melhorada do objetivo do RL, o que, por sua vez, leva a um desempenho aprimorado. Experimentos mostram que o BGPO supera significativamente os algoritmos de RL anteriores para dLLMs em tarefas de resolução de problemas matemáticos, geração de código e planejamento.

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.