Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
By framing on-policy distillation as KL-regularized policy optimization, the authors develop Least-Square Policy Distillation (LSPD) to integrate optimistic exploration and off-policy data reuse into LLM reasoning distillation. This framework directly tackles the high rollout costs and policy diversity collapse inherent to standard reverse-KL on-policy distillation. Across six mathematical reasoning benchmarks, LSPD achieves an average +1.59 gain in Avg@16, scales robustly up to Pass@64, and matches standard on-policy distillation performance using 75% fewer rollout batches.
Treating on-policy distillation as value-based RL slashes rollout compute by 75% while actively preserving the generation diversity needed for high-k reasoning.
We study on-policy distillation (OPD) through the lens of reinforcement learning, establishing a connection between the reverse-KL objective in OPD and KL-regularized policy optimization. Building on this connection, we introduce Least-Square Policy Distillation (LSPD), an RL-inspired framework that brings optimistic exploration and off-policy data reuse from value-based RL into policy distillation. LSPD preserves policy diversity through exploration while improving rollout efficiency by repeatedly learning from previously collected trajectories. Our theoretical analysis connects LSPD to optimistic value-based learning and shows that its idealized formulation achieves a sharp $\tilde{\mathcal O}(\log K)$ regret bound under online exploration. Empirically, LSPD consistently outperforms existing distillation baselines across six mathematical reasoning benchmarks and diverse teacher-student settings, with average gains of +1.59 points in Avg@16. Remarkably, through Pass@k evaluations up to k=64, we found that LSPD better preserves policy diversity by achieving stronger performance as k grows. Its fully off-policy variant achieves comparable performance to vanilla OPD using only the first 25% of rollout batches. Together, these results provide an RL perspective on OPD that offers both a principled interpretation and a practical route toward more effective and rollout-efficient language model distillation.