Zur Gestaltung von KL-regularisierten Policy-Gradient-Algorithmen für das Reasoning mit großen Sprachmodellen (LLM)

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

May 23, 2025

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

cs.AI

Zusammenfassung

Policy-Gradient-Algorithmen wurden erfolgreich eingesetzt, um die Fähigkeiten zur logischen Schlussfolgerung von großen Sprachmodellen (LLMs) zu verbessern. Trotz der weit verbreiteten Verwendung von Kullback-Leibler (KL)-Regularisierung in Policy-Gradient-Algorithmen zur Stabilisierung des Trainings bietet die systematische Untersuchung, wie verschiedene KL-Divergenz-Formulierungen geschätzt und in Surrogate-Loss-Funktionen für Online-Reinforcement-Learning (RL) integriert werden können, einen differenzierten und systematisch erforschbaren Designraum. In diesem Artikel schlagen wir Regularized Policy Gradient (RPG) vor, einen systematischen Rahmen zur Ableitung und Analyse von KL-regularisierten Policy-Gradient-Methoden im Online-RL-Setting. Wir leiten Policy-Gradienten und entsprechende Surrogate-Loss-Funktionen für Ziele ab, die sowohl durch die Vorwärts- als auch die Rückwärts-KL-Divergenz regularisiert werden, wobei sowohl normalisierte als auch nicht normalisierte Policy-Verteilungen berücksichtigt werden. Darüber hinaus präsentieren wir Ableitungen für vollständig differenzierbare Loss-Funktionen sowie REINFORCE-artige Gradientenschätzer, die unterschiedliche algorithmische Anforderungen berücksichtigen. Wir führen umfangreiche Experimente zu RL für LLM-Schlussfolgerungen mit diesen Methoden durch und zeigen verbesserte oder wettbewerbsfähige Ergebnisse in Bezug auf Trainingsstabilität und Leistung im Vergleich zu starken Baselines wie GRPO, REINFORCE++ und DAPO. Der Code ist verfügbar unter 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.