LLM推論のためのKL正則化ポリシー勾配アルゴリズムの設計について

On the Design of KL-Regularized Policy Gradient Algorithms for LLM Reasoning

May 23, 2025

著者: Yifan Zhang, Yifeng Liu, Huizhuo Yuan, Yang Yuan, Quanquan Gu, Andrew C Yao

cs.AI

要旨

ポリシー勾配アルゴリズムは、大規模言語モデル（LLM）の推論能力を向上させるために成功裏に適用されてきた。ポリシー勾配アルゴリズムにおいて、訓練を安定化させるためにKullback-Leibler（KL）正則化が広く使用されているにもかかわらず、異なるKLダイバージェンスの定式化がどのように推定され、オンライン強化学習（RL）の代理損失関数に統合されるかについての体系的な探求は、微妙で体系的に探求可能な設計空間を提示している。本論文では、オンラインRL設定におけるKL正則化ポリシー勾配法を導出し分析するための体系的なフレームワークである、正則化ポリシー勾配（RPG）を提案する。正規化および非正規化ポリシー分布を考慮し、順方向および逆方向KLダイバージェンスによって正則化された目的関数に対するポリシー勾配と対応する代理損失関数を導出する。さらに、完全微分可能な損失関数とREINFORCEスタイルの勾配推定器の導出を提示し、多様なアルゴリズムニーズに対応する。これらの手法を用いたLLM推論のためのRLに関する広範な実験を行い、GRPO、REINFORCE++、DAPOなどの強力なベースラインと比較して、訓練の安定性と性能において改善または競争力のある結果を示す。コードはhttps://github.com/complex-reasoning/RPGで公開されている。

English

Policy gradient algorithms have been successfully applied to enhance the reasoning capabilities of large language models (LLMs). Despite the widespread use of Kullback-Leibler (KL) regularization in policy gradient algorithms to stabilize training, the systematic exploration of how different KL divergence formulations can be estimated and integrated into surrogate loss functions for online reinforcement learning (RL) presents a nuanced and systematically explorable design space. In this paper, we propose regularized policy gradient (RPG), a systematic framework for deriving and analyzing KL-regularized policy gradient methods in the online RL setting. We derive policy gradients and corresponding surrogate loss functions for objectives regularized by both forward and reverse KL divergences, considering both normalized and unnormalized policy distributions. Furthermore, we present derivations for fully differentiable loss functions as well as REINFORCE-style gradient estimators, accommodating diverse algorithmic needs. We conduct extensive experiments on RL for LLM reasoning using these methods, showing improved or competitive results in terms of training stability and performance compared to strong baselines such as GRPO, REINFORCE++, and DAPO. The code is available at https://github.com/complex-reasoning/RPG.