안정성-탐험 딜레마를 넘어서: LLM 정책 최적화를 위한 환경 정규화

초록

대규모 언어 모델(LLM)을 위한 정책 최적화(PO)는 안정성-탐험(stability-exploration) 트레이드오프에 직면하며, 현재는 행동 측 정책-KL 정규화기(action-side Policy-KL regularizer)로 조정된다. 이는 실무자를 이중 구속에 빠뜨린다. 정책-KL을 유지하면 응답 행동이 제약되고 행동 측 탐험 예산이 소모되는 반면, 이를 제거하면 명시적인 드리프트 제어 없이 최적화가 진행된다. 우리는 정규화를 입력 측으로 이동함으로써 이 딜레마를 깨는 대안을 제안한다. 학습이 진행됨에 따라 현재 정책에 의해 유도된 학습 쿼리 분포는 RL 이전의 기준 분포로부터 통제되지 않은 채 표류한다. 구체적으로, 환경 정규화 정책 최적화(Environment-Regularized Policy Optimization, ERPO)는 이러한 쿼리 분포 이동을 제한하는 쿼리-KL(Query-KL, QKL) 항을 도입하며, 데이터셋 고정 기준에서 파생된 쿼리별 가중치를 함께 사용하여 각 쿼리별 업데이트가 기준 분포에서 전형적인 쿼리 쪽으로 편향되도록 한다. QKL 기울기는 오직 쿼리 가능도(likelihood)를 통해서만 흐른다. 정책 경사 추정기에 사용되는 응답 점수 함수는 QKL 항에 나타나지 않으므로 QKL은 응답 분포에 직접적인 경사 압력을 가하지 않으며, 따라서 탐험은 보존된다. ERPO는 추가 순전파(forward pass) 없이 GRPO/PPO/REINFORCE 스타일 파이프라인에 통합된다. 여섯 가지 수학적 추론 벤치마크에서 ERPO는 표준 정책-KL 정규화기를 대체하면서 쿼리 분포 드리프트를 효과적으로 제어하고, 고온 디코딩(high-temperature decoding) 및 장기 학습(long-horizon training) 조건에서 더 강력한 정확도와 훨씬 안정적인 동작을 제공한다. 우리의 소스 코드는 https://github.com/alibaba/ERPO 에서 확인할 수 있다.

English

Policy optimization (PO) for Large Language Models faces a stability--exploration trade-off, currently mediated by an action-side Policy-KL regularizer. This puts practitioners in a double bind: keeping Policy-KL constrains response behavior and consumes the action-side exploration budget, while dropping it leaves the optimization without an explicit drift control. We argue for an alternative that breaks the dilemma by moving regularization to the input side. As training progresses, the distribution over training queries induced by the current policy drifts unchecked from its pre-RL reference distribution. Concretely, Environment-Regularized Policy Optimization (ERPO) introduces a Query-KL (QKL) term that bounds this query distribution shift, together with a dataset-static reference-derived per-query weight that biases each per-query update toward queries typical under the reference. The QKL gradient flows strictly through the query likelihood; the response score function used by policy-gradient estimators does not appear in the QKL term, so QKL exerts no direct gradient pressure on the response distribution---exploration is preserved. ERPO plugs into GRPO/PPO/REINFORCE-style pipelines without additional forward passes. On six mathematical reasoning benchmarks, ERPO replaces the standard Policy-KL regularizer while achieving effective control over query distribution drift, delivering stronger accuracy and substantially more stable behavior under high-temperature decoding and long-horizon training.Our source code are available at https://github.com/alibaba/ERPO