Оптимизация стратегии с управлением границами для эффективного использования памяти при обучении с подкреплением крупных языковых моделей диффузии

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

October 13, 2025

Авторы: Nianyi Lin, Jiajie Zhang, Lei Hou, Juanzi Li

cs.AI

Аннотация

Ключевая проблема при применении обучения с подкреплением (RL) к диффузионным большим языковым моделям (dLLMs) заключается в невычислимости их функций правдоподобия, которые необходимы для RL-цели, что требует соответствующих аппроксимаций на каждом шаге обучения. Хотя существующие методы аппроксимируют логарифмы правдоподобия с помощью их нижних границ доказательства (ELBO) через специализированный метод Монте-Карло (MC), вычислительные графы всех MC-сэмплов необходимо сохранять для вычисления градиентов нелинейных членов в RL-цели, что приводит к значительным затратам памяти. Это ограничение снижает допустимый размер выборки, что ведет к неточным аппроксимациям правдоподобия и, в конечном итоге, искажению RL-цели. Чтобы преодолеть это ограничение, мы предлагаем Boundary-Guided Policy Optimization (BGPO), эффективный по памяти алгоритм RL, который максимизирует специально сконструированную нижнюю границу ELBO-цели. Эта нижняя граница тщательно разработана для удовлетворения двух ключевых свойств: (1) Линейность: она формулируется в виде линейной суммы, где каждый член зависит только от одного MC-сэмпла, что позволяет накапливать градиенты по сэмплам и обеспечивает постоянное использование памяти; (2) Эквивалентность: как значение, так и градиент этой нижней границы равны таковым для ELBO-цели при обучении на политике, что делает её также эффективной аппроксимацией для исходной RL-цели. Эти свойства позволяют BGPO использовать большой размер MC-выборки, что приводит к более точным аппроксимациям правдоподобия и улучшенной оценке RL-цели, что, в свою очередь, повышает производительность. Эксперименты показывают, что BGPO значительно превосходит предыдущие RL-алгоритмы для dLLMs в решении математических задач, генерации кода и задачах планирования.

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.