Over het Ontwerp van KL-Geregulariseerde Policy Gradient Algoritmen voor Redeneren met Grote Taalmodellen

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

Samenvatting

Policy gradient-algoritmen zijn met succes toegepast om de redeneervaardigheden van grote taalmodelen (LLMs) te verbeteren. Ondanks het wijdverbreide gebruik van Kullback-Leibler (KL)-regularisatie in policy gradient-algoritmen om de training te stabiliseren, biedt de systematische verkenning van hoe verschillende KL-divergentieformuleringen kunnen worden geschat en geïntegreerd in surrogaatverliesfuncties voor online reinforcement learning (RL) een genuanceerd en systematisch verkenningsgebied. In dit artikel stellen we regularized policy gradient (RPG) voor, een systematisch raamwerk voor het afleiden en analyseren van KL-geregulariseerde policy gradient-methoden in de online RL-setting. We leiden policy gradients en bijbehorende surrogaatverliesfuncties af voor doelen die zijn geregulariseerd door zowel voorwaartse als achterwaartse KL-divergenties, waarbij we zowel genormaliseerde als niet-genormaliseerde policy-distributies overwegen. Bovendien presenteren we afleidingen voor volledig differentieerbare verliesfuncties, evenals REINFORCE-stijl gradient-schatters, die diverse algoritmische behoeften accommoderen. We voeren uitgebreide experimenten uit op RL voor LLM-redenering met behulp van deze methoden, waarbij verbeterde of competitieve resultaten worden getoond op het gebied van trainingsstabiliteit en prestaties in vergelijking met sterke baselines zoals GRPO, REINFORCE++ en DAPO. De code is beschikbaar op 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.