확산 대형 언어 모델의 메모리 효율적 강화 학습을 위한 경계 지향 정책 최적화

초록

확산 기반 대형 언어 모델(dLLMs)에 강화 학습(RL)을 적용할 때의 주요 과제는 RL 목적 함수에 필수적인 가능도 함수의 계산 불가능성에 있습니다. 이는 각 학습 단계에서 해당 근사가 필요함을 의미합니다. 기존 방법들은 맞춤형 몬테카를로(MC) 샘플링을 통해 증거 하한(ELBO)으로 로그 가능도를 근사하지만, RL 목적 함수의 비선형 항에 대한 그래디언트 계산을 위해 모든 MC 샘플의 순방향 계산 그래프를 유지해야 하므로 상당한 메모리 오버헤드가 발생합니다. 이러한 제약으로 인해 실현 가능한 샘플 크기가 제한되어 부정확한 가능도 근사와 궁극적으로 RL 목적 함수의 왜곡을 초래합니다. 이러한 한계를 극복하기 위해, 우리는 메모리 효율적인 RL 알고리즘인 경계 지향 정책 최적화(BGPO)를 제안합니다. BGPO는 ELBO 기반 목적 함수의 특별히 구성된 하한을 최대화합니다. 이 하한은 두 가지 주요 속성을 충족하도록 신중하게 설계되었습니다: (1) 선형성: 각 항이 단일 MC 샘플에만 의존하는 선형 합으로 공식화되어 샘플 간 그래디언트 누적이 가능하고 일정한 메모리 사용을 보장합니다; (2) 등가성: 이 하한의 값과 그래디언트는 온-정책 학습에서 ELBO 기반 목적 함수의 것과 동일하므로 원래 RL 목적 함수에 대한 효과적인 근사치이기도 합니다. 이러한 속성으로 인해 BGPO는 큰 MC 샘플 크기를 채택할 수 있어 더 정확한 가능도 근사와 개선된 RL 목적 함수 추정이 가능하며, 이는 성능 향상으로 이어집니다. 실험 결과, BGPO는 수학 문제 해결, 코드 생성 및 계획 작업에서 dLLMs를 위한 기존 RL 알고리즘들을 크게 능가하는 것으로 나타났습니다.

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.

확산 대형 언어 모델의 메모리 효율적 강화 학습을 위한 경계 지향 정책 최적화

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

초록

Support