LSPD: An RL View of On-Policy Distillation

On-policy distillation (OPD) is commonly implemented with policy-gradient updates. But its reverse-KL objective also has an exact interpretation as KL-regularized reinforcement learning, with a reward defined by the teacher relative to a reference policy. This raises a natural question: can value-based RL principles make policy distillation more sample-efficient?

We introduce Least Square Policy Distillation (LSPD), motivated by optimistic reward estimation and KL-regularized policy improvement. Rather than using teacher feedback only to drive policy-gradient updates, LSPD directly matches student and teacher log probabilities through a robust least-squares objective. This formulation enables learning from accumulated experience, bringing a central advantage of value-based RL to LLM distillation. Our theoretical analysis establishes logarithmic reverse-KL regret for an idealized optimistic formulation.

Across six mathematical reasoning benchmarks and three teacher–student settings, LSPD improves reasoning performance and multi-sample solution coverage. Its replay-based variant, LSPD-RB, achieves performance comparable to vanilla OPD with approximately one-quarter of the rollout batches. The broader message: viewing distillation as an RL problem opens up algorithmic choices beyond on-policy policy gradients—and opportunities to learn more from each interaction with the teacher.

📄 Paper · 💻 Code