Sobre o Projeto de Algoritmos de Gradiente de Política Regularizados por KL para Raciocínio em LLM

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

May 23, 2025

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

cs.AI

Resumo

Algoritmos de gradiente de política têm sido aplicados com sucesso para aprimorar as capacidades de raciocínio de modelos de linguagem de grande escala (LLMs). Apesar do uso generalizado da regularização de Kullback-Leibler (KL) em algoritmos de gradiente de política para estabilizar o treinamento, a exploração sistemática de como diferentes formulações de divergência KL podem ser estimadas e integradas em funções de perda substitutas para aprendizado por reforço (RL) online apresenta um espaço de design sutil e sistematicamente explorável. Neste artigo, propomos o gradiente de política regularizado (RPG), uma estrutura sistemática para derivar e analisar métodos de gradiente de política regularizados por KL no cenário de RL online. Derivamos gradientes de política e funções de perda substitutas correspondentes para objetivos regularizados por divergências KL direta e reversa, considerando distribuições de política normalizadas e não normalizadas. Além disso, apresentamos derivações para funções de perda totalmente diferenciáveis, bem como estimadores de gradiente no estilo REINFORCE, atendendo a diversas necessidades algorítmicas. Realizamos experimentos extensivos em RL para raciocínio em LLMs utilizando esses métodos, mostrando resultados melhorados ou competitivos em termos de estabilidade de treinamento e desempenho em comparação com baselines robustos como GRPO, REINFORCE++ e DAPO. O código está disponível em 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.