대규모 언어 모델 사후 학습에서 정책 미러 하강의 로그 분할 함수 근사가 유도하는 암묵적 정규화

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

February 5, 2026

저자: Zhenghao Xu, Qin Lu, Changlong Yu, Tuo Zhao

cs.AI

초록

정책 미러 디센트(PMD)는 KL 정규화된 정책 개선 하위 문제를 반복적으로 해결함으로써 강화 학습(RL)에 대한 원칙적인 프레임워크를 제공합니다. 이러한 접근법은 Kimi K1.5/K2와 같은 고도화된 대규모 언어 모델(LLM) 학습에 채택되었지만, 이상적인 폐쇄형 PMD 업데이트는 신뢰할 수 있는 분할 함수 추정을 요구하며, 이는 LLM의 방대한 행동 공간에서 제한된 롤아웃으로 작업할 때 큰 도전 과제입니다. 본 연구에서는 샘플링 정책 하의 평균 보상으로 로그 분할 항을 근사하고 로그 정책 공간에서 회귀를 수행하는 PMD-mean이라는 실용적인 알고리즘을探究합니다. 구체적으로, 우리는 PMD-mean의 모집단 해를 규명하고 이 알고리즘이 적응형 혼합 KL-χ^2 정규화를 사용한 미러 디센트 하위 문제를 암묵적으로 최적화함을 입증합니다. 이 추가적인 χ^2 정규화는 큰 확률 변화를 제약하여 예상 보상이 낮을 때 더 보수적인 업데이트를 생성하고 유한 표본 추정 오차에 대한 견고성을 향상시킵니다. 수학적 추론 작업에 대한 실험 결과, PMD-mean이 향상된 안정성과 시간 효율성으로 우수한 성능을 달성함을 보여줍니다. 이러한 발견은 PMD-mean에 대한 우리의 이해를 심화하고 LLM을 위한 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.