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.