Sur la conception d'algorithmes de gradient de politique régularisés par KL pour le raisonnement des grands modèles de langage

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

May 23, 2025

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

cs.AI

Résumé

Les algorithmes de gradient de politique ont été appliqués avec succès pour améliorer les capacités de raisonnement des grands modèles de langage (LLM). Malgré l'utilisation répandue de la régularisation de Kullback-Leibler (KL) dans les algorithmes de gradient de politique pour stabiliser l'entraînement, l'exploration systématique de la manière dont différentes formulations de divergence KL peuvent être estimées et intégrées dans les fonctions de perte substitutives pour l'apprentissage par renforcement (RL) en ligne représente un espace de conception nuancé et systématiquement explorable. Dans cet article, nous proposons le gradient de politique régularisé (RPG), un cadre systématique pour dériver et analyser les méthodes de gradient de politique régularisées par KL dans le contexte du RL en ligne. Nous dérivons les gradients de politique et les fonctions de perte substitutives correspondantes pour des objectifs régularisés par les divergences KL directe et inverse, en considérant à la fois les distributions de politiques normalisées et non normalisées. De plus, nous présentons des dérivations pour des fonctions de perte entièrement différentiables ainsi que des estimateurs de gradient de style REINFORCE, répondant à divers besoins algorithmiques. Nous menons des expériences approfondies sur le RL pour le raisonnement des LLM en utilisant ces méthodes, montrant des résultats améliorés ou compétitifs en termes de stabilité d'entraînement et de performance par rapport à des bases de référence solides telles que GRPO, REINFORCE++ et DAPO. Le code est disponible à l'adresse 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.