Reverse-KL Policy Distillation: Statistical Guarantees and Limits

Haolin Liu; Chen-Yu Wei; Julian Zimmert

cs.LGOct 22, 2026

Reverse-KL Policy Distillation: Statistical Guarantees and Limits

Authors: Haolin Liu, Chen-Yu Wei, Julian Zimmert

Organizations: University of Virginia · Google Research · National Taiwan University

Abstract

Reverse Kullback-Leibler (KL) distillation has shown empirical promise for training language models, but its statistical guarantees and limits remain poorly understood. In particular, the effects of two design choices remain unclear: Does access to the teacher's full action distribution provide an advantage over observing only the probability it assigns to the sampled action? Must training trajectories come from the student (online distillation), or can teacher-generated trajectories (offline distillation) suffice? Our work takes a first step toward understanding the theory underlying these questions. We study the sample complexity of online reverse-KL distillation under two feedback models: the teacher reveals either its full action distribution or only the probability of the student-generated action. We establish minimax bounds for the former and show that substantially more data may be required for the latter. In contrast, offline reverse-KL distillation can attain the same minimax rate whether the teacher reveals its full distribution or only its probability of the teacher-generated action. The key distinction is that teacher-generated actions themselves reveal the teacher's preferences, while queries at student-selected actions can leave those preferences unobserved. This highlights an additional statistical challenge for online distillation compared to offline distillation. These statistical guarantees do not follow simply from optimizing standard distillation objectives: our lower bounds show that common online objectives can fail when candidates are evaluated only on past student-generated trajectories. We develop trajectory-loss and importance-ratio truncation techniques together with policy-space learning to achieve optimal online distillation guarantees.

Figures & tables

Data source	Feedback	Assumption	$C_{\rm app}$	$C_{\rm est}$
Online / Offline	Full	Relative second moment [2pt] $\mathbb{E}_{P^{\pi}}[K_{\pi}^{2}]\leq B\,\mathbb{E}_{P^{\pi}}[K_{\pi}]$	$\Theta(B)$	$\Theta(B)$
Offline	Action	$P^{\pi}(\tau)\leq e^{B_{\infty}}P^{\pi^{\star}}(\tau)$	$\Theta(B_{\infty})$	$\Theta(B_{\infty})$
Online	Action	$P^{\pi}(\tau)\leq e^{B_{\infty}}P^{\pi^{\star}}(\tau)$	—	$\Omega\!\left(\frac{B_{\infty}m}{\log(em)}\right)$

Table 1: Reverse-KL learning by data source and teacher feedback. Approximation coefficients are compared at the optimal estimation scale.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

g_{h}(\xi_{t,h};\pi)=\sqrt{\prod_{v=1}^{h-1}\frac{\pi_{v}^{\star}(a_{t,v}\mid s_{t,v})}{\pi_{v}(a_{t,v}\mid s_{t,v})}}D_{\mathrm{H}}^{2}\!\left(\pi_{h}(\cdot\mid s_{t,h}),\pi_{h}^{\star}(\cdot\mid s_{t,h})\right).

Appendix

Algorithm 2 Online Hellinger learning with full teacher feedback

Language Model Distillation

Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.

See all 508 papers

A Token-Level Analysis of Sampled-Token Reverse-KL On-Policy Distillation
Aug 26, 2026Bing Shao, Jiazheng Zhang, Long Ma +13Language Model Distillation On-Policy Distillation
College of Computer Science and Artificial Intelligence, Fudan University
Understanding Off- vs On-Policy Distillation: A Tale of Distinct Training Objectives
Sep 29, 2026Qiwei Di, Xuheng Li, Kaixuan Ji +3Multi-Teacher Knowledge Distillation On-Policy Distillation
Department of Computer Science, University of California, Los Angeles, CA 90095, USA
Decoupling KL and Trajectories: A Unified Perspective for SFT, DAgger, Offline RL, and OPD in LLM Distillation
May 16, 2026Anhao Zhao, Haoran Xin, Yingqi Fan +3Reasoning Distillation On-Policy Distillation
Eastern Institute of Technology, Ningbo · Thrust of Artificial Intelligence, The Hong Kong University of Science and Technology (Guangzhou) · Shanghai Jiao Tong University +1

Aug 26, 2026J/K move · Enter open · S save

A Token-Level Analysis of Sampled-Token Reverse-KL On-Policy Distillation

Bing Shao, Jiazheng Zhang, Long Ma, Yujiong Shen+12

College of Computer Science and Artificial Intelligence, Fudan University

Language Model Distillation On-Policy Distillation Knowledge Distillation Kullback-Leibler Divergence

Open paper AskPDF arXiv

On-policy distillation (OPD) supervises a student on its own trajectories with token-level signals from a frozen teacher, yet how a sampled loss allocates updates across tokens remains poorly understood. We analyze the gradient of the per-token K2 estimator of reverse KL with respect to the student logits. The

\ell_1

norm of this gradient factorizes into the absolute teacher--student log-probability gap and a student-side softmax factor that grows as the sampled token becomes less likely under the student. In our math-distillation runs, these per-token norms are highly non-uniform: low-student-probability tokens account for a disproportionate share of their sum and are also enriched in large teacher--student gaps. As a lightweight intervention suggested by this analysis, we study Surprise-aware Reweighting (SuRe), a detached, bounded weighting rule that further amplifies this existing allocation. Across two Qwen3 student scales, SuRe improves several math metrics over vanilla OPD and shows no clear degradation on the selected out-of-domain benchmarks. Our primary contribution is therefore a gradient-level characterization of reverse-KL OPD trained with the K2 estimator, with SuRe as one empirical instantiation.

Abstract

Language Model Distillation

Explore similar work

A Token-Level Analysis of Sampled-Token Reverse-KL On-Policy Distillation

Understanding Off- vs On-Policy Distillation: A Tale of Distinct Training Objectives

Decoupling KL and Trajectories: A Unified Perspective for SFT, DAgger, Offline RL, and OPD in LLM Distillation