拡散型大規模言語モデルのメモリ効率化RLのための境界誘導型ポリシー最適化

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

要旨

拡散型大規模言語モデル（dLLM）に強化学習（RL）を適用する際の主要な課題は、RLの目的関数に不可欠な尤度関数の計算困難性にあり、各訓練ステップで対応する近似が必要となります。既存の手法では、カスタマイズされたモンテカルロ（MC）サンプリングを通じて対数尤度をその証拠下界（ELBO）で近似していますが、RL目的関数の非線形項の勾配計算のために、すべてのMCサンプルの順方向計算グラフを保持する必要があり、これが大きなメモリオーバーヘッドを引き起こします。この制約により、実行可能なサンプルサイズが制限され、不正確な尤度近似が生じ、最終的にRL目的関数が歪んでしまいます。この制限を克服するため、我々はメモリ効率の良いRLアルゴリズムであるBoundary-Guided Policy Optimization（BGPO）を提案します。BGPOは、ELBOベースの目的関数の特別に構築された下界を最大化します。この下界は、以下の2つの重要な特性を満たすように慎重に設計されています：（1）線形性：各項が単一のMCサンプルにのみ依存する線形和として定式化されており、サンプル間での勾配累積を可能にし、メモリ使用量を一定に保ちます；（2）等価性：この下界の値と勾配は、オン方策訓練においてELBOベースの目的関数のそれらと等しく、元のRL目的関数の効果的な近似としても機能します。これらの特性により、BGPOは大きなMCサンプルサイズを採用することができ、より正確な尤度近似と改善されたRL目的関数の推定が可能となり、結果として性能が向上します。実験結果は、BGPOが数学問題解決、コード生成、および計画タスクにおいて、従来の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.