О разработке алгоритмов градиента политики с регуляризацией KL для рассуждений в крупных языковых моделях

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

May 23, 2025

Авторы: Yifan Zhang, Yifeng Liu, Huizhuo Yuan, Yang Yuan, Quanquan Gu, Andrew C Yao

cs.AI

Аннотация

Алгоритмы градиента политики успешно применяются для улучшения способностей к рассуждению у крупных языковых моделей (LLM). Несмотря на широкое использование регуляризации по Кулбеку-Лейблеру (KL) в алгоритмах градиента политики для стабилизации обучения, систематическое исследование того, как различные формулировки расхождения KL могут быть оценены и интегрированы в суррогатные функции потерь для онлайн-обучения с подкреплением (RL), представляет собой тонкий и систематически исследуемый дизайн-пространство. В данной работе мы предлагаем Regularized Policy Gradient (RPG) — систематическую структуру для вывода и анализа методов градиента политики с KL-регуляризацией в условиях онлайн RL. Мы выводим градиенты политики и соответствующие суррогатные функции потерь для целей, регуляризованных как прямым, так и обратным расхождениями KL, учитывая как нормализованные, так и ненормализованные распределения политик. Кроме того, мы представляем выводы для полностью дифференцируемых функций потерь, а также для оценщиков градиента в стиле REINFORCE, что позволяет учитывать разнообразные алгоритмические потребности. Мы проводим обширные эксперименты по RL для рассуждений LLM с использованием этих методов, демонстрируя улучшенные или конкурентоспособные результаты с точки зрения стабильности обучения и производительности по сравнению с сильными базовыми методами, такими как GRPO, REINFORCE++ и DAPO. Код доступен по адресу 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.