Аппроксимация логарифмической функции нормализации в методе зеркального спуска для политик индуцирует неявную регуляризацию при пост-обучении больших языковых моделей

Аннотация

Метод зеркального спуска для политик (PMD) предоставляет принципиальную основу для обучения с подкреплением (RL) путем итеративного решения KL-регуляризованных подзадач улучшения политики. Хотя этот подход был применен при обучении продвинутых больших языковых моделей, таких как Kimi K1.5/K2, идеальные PMD-обновления в закрытой форме требуют надежной оценки статистической суммы, что представляет собой серьезную проблему при работе с ограниченными траекториями в обширных пространствах действий БЯМ. Мы исследуем практический алгоритм, названный PMD-mean, который аппроксимирует логарифм статистической суммы математическим ожиданием награды при политике выборки и выполняет регрессию в пространстве логарифмов политик. В частности, мы характеризуем популяционное решение PMD-mean и показываем, что оно неявно оптимизирует подзадачи зеркального спуска с адаптивным смешанным KL–χ² регуляризатором. Эта дополнительная χ²-регуляризация ограничивает большие изменения вероятностей, порождая более консервативные обновления при низких ожидаемых наградах и повышая устойчивость к ошибкам оценки на конечных выборках. Эксперименты на задачах математических рассуждений показывают, что PMD-mean достигает превосходной производительности с улучшенной стабильностью и временной эффективностью. Эти результаты углубляют наше понимание PMD-mean и освещают пути к принципиальным улучшениям RL-алгоритмов для БЯМ. Код доступен по адресу https://github.com/horizon-rl/OpenKimi.

English

Policy mirror descent (PMD) provides a principled framework for reinforcement learning (RL) by iteratively solving KL-regularized policy improvement subproblems. While this approach has been adopted in training advanced LLMs such as Kimi K1.5/K2, the ideal closed-form PMD updates require reliable partition function estimation, a significant challenge when working with limited rollouts in the vast action spaces of LLMs. We investigate a practical algorithm, termed PMD-mean, that approximates the log-partition term with the mean reward under the sampling policy and performs regression in log-policy space. Specifically, we characterize the population solution of PMD-mean and demonstrate that it implicitly optimizes mirror descent subproblems with an adaptive mixed KL--χ^2 regularizer. This additional χ^2 regularization constrains large probability changes, producing more conservative updates when expected rewards are low and enhancing robustness against finite-sample estimation errors. Experiments on math reasoning tasks show that PMD-mean achieves superior performance with improved stability and time efficiency. These findings deepen our understanding of PMD-mean and illuminate pathways toward principled improvements in RL algorithms for LLMs. Code is available at https://github.com/horizon-rl/OpenKimi.

Approximation of Log-Partition Function in Policy Mirror Descent Induces Implicit Regularization for LLM Post-Training

Аннотация

Support