Sobre el Diseño de Algoritmos de Gradiente de Política Regularizados con KL para el Razonamiento en Modelos de Lenguaje de Gran Escala

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

Resumen

Los algoritmos de gradiente de políticas han sido aplicados con éxito para mejorar las capacidades de razonamiento de los modelos de lenguaje grandes (LLMs). A pesar del uso generalizado de la regularización de Kullback-Leibler (KL) en los algoritmos de gradiente de políticas para estabilizar el entrenamiento, la exploración sistemática de cómo diferentes formulaciones de divergencia KL pueden ser estimadas e integradas en funciones de pérdida sustitutas para el aprendizaje por refuerzo (RL) en línea presenta un espacio de diseño matizado y sistemáticamente explorable. En este artículo, proponemos el gradiente de políticas regularizado (RPG), un marco sistemático para derivar y analizar métodos de gradiente de políticas regularizados por KL en el entorno de RL en línea. Derivamos gradientes de políticas y las correspondientes funciones de pérdida sustitutas para objetivos regularizados tanto por divergencias KL directas como inversas, considerando distribuciones de políticas normalizadas y no normalizadas. Además, presentamos derivaciones para funciones de pérdida completamente diferenciables, así como estimadores de gradiente al estilo REINFORCE, adaptándonos a diversas necesidades algorítmicas. Realizamos experimentos extensivos en RL para el razonamiento de LLMs utilizando estos métodos, mostrando resultados mejorados o competitivos en términos de estabilidad de entrenamiento y rendimiento en comparación con líneas base sólidas como GRPO, REINFORCE++ y DAPO. El código está disponible en 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.