Teacher-Guided Policy Optimization for LLM Distillation

TL;DR

This paper introduces Teacher-Guided Policy Optimization (TGPO), an on-policy method for LLM distillation that improves upon Reverse KL by providing dense guidance, leading to better performance and robustness.

Contribution

TGPO is a novel on-policy algorithm that incorporates teacher guidance conditioned on student rollouts, enhancing LLM distillation efficiency without extra data annotation.

Findings

TGPO outperforms standard baselines on reasoning benchmarks.

TGPO is robust to different teacher models.

TGPO integrates seamlessly with existing RL frameworks.

Abstract

The convergence of reinforcement learning and imitation learning has positioned Reverse KL (RKL) as a promising paradigm for on-policy LLM distillation, aiming to unify exploration with teacher supervision. However, we identify a critical limitation: when the student and teacher distributions diverge significantly, standard RKL often fails to yield meaningful improvement due to uninformative negative feedback. To address this inefficiency, we propose Teacher-Guided Policy Optimization (TGPO), an on-policy algorithm that incorporates dense directional guidance by leveraging teacher predictions conditioned on the student's rollout. Because TGPO remains on-policy, the algorithm integrates seamlessly with existing RLVR frameworks without requiring additional data annotation. Experiments on complex reasoning benchmarks demonstrate that TGPO significantly outperforms standard baselines and is robust to different teachers.