Reinforcement learning improves efficiency of language model reasoning
An RL View of OPD: Least Square Policy Distillation for Sample-Efficient LLM Reasoning
Machine Learning
Summary
Building large language models that can reason well often requires a lot of trial runs, which can be slow and costly. This paper shows how a reinforcement learning perspective can make the process more efficient by reusing past experiences and encouraging the model to explore different approaches. The new method, called Least-Square Policy Distillation, helps models learn faster and maintain diverse reasoning styles. Tests on math reasoning tasks show it outperforms previous methods while using fewer new trial runs.
What this means in practice
- •For machine learning engineers: Improve efficiency and performance of large language models for reasoning by incorporating off-policy data reuse during training.
- •For ai system developers: Develop LLM-based tools that maintain diverse reasoning strategies to enhance robustness and output quality over multiple answer attempts.
Authors
Shangzhe Li, Yuxiao Yang, Tianrun Yu, Kaixiang Zhao, Xiaoyun Wang, Taylor W. Killian, Weitong Zhang
Abstract
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.