LAPO: KL-Free Logit-Averaged Policy Optimization to Leverage RL and SFT in LLM Post-Training

원문 초록

From a classical statistical perspective, model averaging and parameter shrinkage (e.g., Ridge regression with an $l_2$-penalty) are deeply connected mechanisms designed to reduce estimator variance and lower mean squared error. Motivated by this insight, we introduce Logit-Averaged Policy Optimization (LAPO), a KL-free and critic-free post-training framework that replaces explicit KL regularization with logit averaging inside Group Relative Policy Optimization (GRPO); in particular, averaging the logits of the trainable policy and a frozen reference policy functions as an implicit structural regularizer that stabilizes rollout trajectories. For selecting the mixing weight, we propose a compute-efficient adaptive algorithm that dynamically balances exploration and structural anchoring throughout training, evaluated against a static fixed-weight baseline. This adaptive mechanism eliminates explicit KL penalty backpropagation and expensive, fragile hyperparameter sweeps while reducing the incidence of uninformative zero-reward rollout groups. Mathematically, embedding logit averaging within the reinforcement learning loop operates as an in-training Product of Experts (PoE): the multiplicative geometric mean enables the reference anchor to act as a soft logical "AND" that vetoes non-compliant tokens, enforcing schema rigor and safety guardrails without the catastrophic forgetting typical of sequential fine-tuning. Evaluated across Qwen2.5-Instruct models (1.5B, 3B, and 7B), LAPO consistently outperforms canonical KL-regularized GRPO, achieving a 7%-15% accuracy gain on MATH and a 2%-5% gain on cn_k12 and MMLU. Finally, the resulting mixed teacher policy can be distilled back into a single student network, fully restoring standard $1\times$ inference latency.