Sulla Progettazione di Algoritmi di Policy Gradient con Regolarizzazione KL per il Ragionamento nei Modelli Linguistici di Grande Dimensione

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

May 23, 2025

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

cs.AI

Abstract

Gli algoritmi a gradiente della politica sono stati applicati con successo per potenziare le capacità di ragionamento dei grandi modelli linguistici (LLM). Nonostante l'uso diffuso della regolarizzazione di Kullback-Leibler (KL) negli algoritmi a gradiente della politica per stabilizzare l'addestramento, l'esplorazione sistematica di come diverse formulazioni della divergenza KL possano essere stimate e integrate nelle funzioni di perdita surrogate per l'apprendimento per rinforzo (RL) online rappresenta uno spazio di progettazione sfumato e sistematicamente esplorabile. In questo articolo, proponiamo il gradiente della politica regolarizzato (RPG), un framework sistematico per derivare e analizzare metodi a gradiente della politica regolarizzati con KL nel contesto del RL online. Deriviamo i gradienti della politica e le corrispondenti funzioni di perdita surrogate per obiettivi regolarizzati sia dalla divergenza KL diretta che inversa, considerando sia distribuzioni di politica normalizzate che non normalizzate. Inoltre, presentiamo derivazioni per funzioni di perdita completamente differenziabili nonché stimatori di gradiente in stile REINFORCE, adattandoci a diverse esigenze algoritmiche. Condividiamo esperimenti estesi sul RL per il ragionamento con LLM utilizzando questi metodi, mostrando risultati migliorati o competitivi in termini di stabilità dell'addestramento e prestazioni rispetto a baseline robusti come GRPO, REINFORCE++ e DAPO. Il codice è disponibile all'indirizzo 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.