Source-linked AI summary
Reverse-KL Policy Distillation: Statistical Guarantees and Limits
Haolin Liu, Chen-Yu Wei, Julian Zimmert
TL;DR
The statistical guarantees of reverse-KL distillation remain unclear, including whether full teacher feedback helps and whether offline teacher trajectories can match online student rollouts. The paper analyzes these settings and develops learning methods with optimal guarantees; its results show an online sample-efficiency advantage for full feedback, while offline learning can attain a comparable guarantee under a trajectory likelihood-ratio bound.
Problem
The statistical guarantees of reverse-KL learning remain unclear, including the effects of teacher-feedback richness and whether offline teacher-generated data can match online student-rollout guarantees.
Source: §1 · p. 1Method
The paper analyzes online and offline reverse-KL learning under full-distribution or observed-action-probability feedback, and develops algorithms including online prefix reweighting.
Source: §1 · p. 2Results
Full feedback can significantly improve online sample efficiency over action feedback, while offline learning under a trajectory likelihood-ratio bound can attain a comparable guarantee.
Source: §1 · p. 2, §5 · p. 11Takeaways & Limitations
The results show that optimal reverse-KL learning need not require online interaction, although online action feedback can entail additional sample cost.
Source: §1 · p. 2, §5 · p. 11The guarantees’ boundedness parameters may grow with the horizon, so their lack of an explicit horizon factor does not imply horizon independence.
Source: §2 · p. 4
Abstract
from arXiv · showhide
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.
1 Introduction
Reverse-KL distillation lets a limited-capacity student focus on teacher behaviors it can represent, but the statistical guarantees and limits of this approach remain unclear. The paper compares online and offline learning under full and action-only teacher feedback, showing that feedback and trajectory source affect sample efficiency.
1 Introduction
- Reverse KL lets a limited-capacity student concentrate on teacher-supported behaviors it can represent, unlike forward KL, which encourages coverage of the teacher’s full response range.Prior empirical work reports promise for reverse KL in on-policy LLM distillation [Gu et al., 2024, Agarwal et al., 2024, Lu and Thinking Machines Lab, 2025].Source: §1 · p. 1
- The study asks whether full teacher distributions improve statistical efficiency and whether offline teacher-generated trajectories can match guarantees from online student rollouts.These questions arise because reverse KL is evaluated under the student’s distribution, while some methods observe only the teacher’s probability for a sampled action.Source: §1 · pp. 1–2
- The framework compares online student-selected and offline teacher-generated trajectories, each paired with either full-distribution or observed-action-probability feedback.Across these four settings, the target is to match the best student in the policy class using a given number of training trajectories.Source: §1 · p. 2
- Under full feedback, the paper establishes minimax-optimal reverse-KL guarantees and develops online and offline algorithms that attain the optimal guarantee.The online method reweights observed prefixes to account for each candidate policy’s generation behavior; the offline algorithm uses teacher-generated trajectories.Source: §1 · p. 2
- Online action feedback can require substantially more samples than full feedback, while offline learning can attain a comparable guarantee under a trajectory likelihood-ratio bound.The paper also identifies common online objectives that can be suboptimal when candidate policies are evaluated on states generated by past rollouts.Source: §1 · p. 2
2 Problem formulation
The paper studies a finite student policy class imitating a fixed teacher in a reward-free MDP, using trajectory reverse KL as its risk. It distinguishes online and offline trajectory sources, two teacher-feedback models, and boundedness assumptions used for guarantees.
2 Problem formulation
- The target risk is trajectory reverse KL, equal by the KL chain rule to the expected cumulative local reverse KL under the student’s trajectory distribution.The best-in-class risk is L⋆; the setting is realizable when the teacher belongs to the student class and agnostic otherwise.Source: §2 · pp. 3–4
- Online trajectories come from learner-selected student policies, whereas offline trajectories are independent samples from the teacher policy.At visited states, full feedback reveals the teacher’s action distribution and action feedback reveals only its probability for the observed action; neither provides feedback at unvisited states.Source: §2 · p. 4
- Full-feedback guarantees use a relative second-moment condition, while action-feedback guarantees use a stronger trajectory likelihood-ratio bound.The likelihood-ratio condition implies the relative second-moment condition with B = 4(B∞ + 1).Source: §2 · p. 4
- The relative second-moment assumption controls the second moment of cumulative local KL linearly by its mean under each candidate policy’s own trajectory distribution.It is weaker than several uniform boundedness assumptions, including a uniform local KL bound.Source: §2 · p. 4
- Although the guarantees have no explicit horizon factor, their boundedness parameters can grow with H, so they do not imply horizon independence.A uniform local KL bound kh(s; π) ≤ G implies a trajectory likelihood-ratio bound with exponent HG.Source: §2 · p. 4
- The guarantee separates approximation cost from estimation error, with the estimation term scaling as log(|Π|/δ)/T.The analysis fixes the estimation coefficient at its minimax-optimal order and characterizes the smallest attainable approximation coefficient.Source: §2 · p. 5
3 Online learning from student trajectories
The section shows that standard online refitting can fail because candidate policies are evaluated on states visited by rollout policies. It develops clipped prefix comparisons and policy-space learning to obtain optimal full-feedback guarantees, while showing that online action feedback can require more samples.
3.1 Why common objectives can fail
- After Θ(B^3) episodes, common weighted or subset refitting can return risk Ω(B) with constant probability because rollout states make πgood look worse.The failure persists with adaptive selection and weighting of past trajectories.Source: §3.1 · pp. 5–6
3.2 Correcting visitation mismatch
- Prefix clipping limits candidate-versus-rollout weights; when rollout risk is small, truncated loss still controls sequence risk, with an O(B) clipping cost.Clipped prefix comparisons retain shared prefix probability mass rather than amplifying rare observations.Source: §3.2 · pp. 6–7
3.3 Online learning for reverse KL minimization
- The learner updates candidate weights from clipped prefix comparisons against the sampled rollout policy and returns the candidate with the largest final weight.The loss comparison uses full teacher feedback at visited states.Source: §3.3 · p. 7
- C_est = Θ(B) is minimax-optimal, while achieving C_est = O(B) forces worst-case C_app = Ω(B) under Assumption 2.1.Algorithm 1 attains the optimal estimation rate and the smallest possible order of approximation coefficient at that rate.Source: §3.3 · p. 8
- Online full-feedback learning for Hellinger distance achieves constant approximation and estimation coefficients, matching the offline teacher-trajectory guarantee.Under Assumption 2.2, this guarantee transfers to reverse KL at cost B∞; the analogous reduction does not hold under Assumption 2.1.Source: §3.3 · p. 9
3.4 The Limitation of Action Feedback in the Online Setting
- Online action feedback requires an additional min{|A|, |Π|, e^B∞} estimation factor relative to full feedback, even in realizable instances.The lower bound holds with probability at least 1/4.Source: §3.4 · p. 9
4 Offline learning from teacher trajectories
Offline learning uses teacher-generated trajectories and compares full teacher-distribution feedback with action-only feedback. The section shows that clipping repairs a misspecification problem for full-feedback ERM, while offline action feedback can retain a comparable guarantee under a trajectory likelihood-ratio bound.
4 Offline learning from teacher trajectories
- The offline setting supplies independent teacher trajectories, and the full-feedback analysis estimates average accumulated local KL over those trajectories before connecting the estimate to target risk.Clipping controls the change of measure between teacher-trajectory averages and risk under the candidate policy.Source: §4 · p. 10, §4.1 · pp. 10–11
4.1 Offline learning with full feedback
- Unclipped offline ERM has an optimal estimation rate when the teacher belongs to the policy class, but under misspecification its uniform guarantee requires e^Ω(B²) approximation cost.The failure arises because teacher-visited states that a candidate rarely reaches can have large local KL losses and dominate the objective.Source: §4.1 · p. 10
- Clipping preserves the optimal realizable estimation rate and improves the pathwise misspecified approximation coefficient from e^Ω(B²) to O(B), matching the lower bound.The unclipped estimator requires e^Ω(B²) approximation error under misspecification, even under the stronger pathwise bound.Source: §4.1 · p. 10
4.2 Offline learning with action feedback
- Under Assumption 2.2, offline action-feedback learning is minimax optimal and can require exponentially fewer episodes than online action-feedback learning in the worst case.The method uses an offline Hellinger estimator and converts its bound to reverse-KL risk; the lower bounds also hold in this setting.Source: §4.2 · p. 11
5 Conclusion
The paper compares reverse-KL learning across data sources and feedback models under misspecification. Its conclusion is that online learning can suffer from weak objectives and action-only feedback, while offline teacher trajectories can avoid the additional action-feedback sample cost under a trajectory likelihood-ratio bound.
5 Conclusion
- Online learning has failures of common empirical algorithms, while the paper gives an optimal full-feedback algorithm under a relative second-moment condition.The paper also reports that full feedback can significantly improve online sample efficiency over action feedback.Source: §5 · p. 11
- Under a trajectory likelihood-ratio bound, offline teacher trajectories can avoid the additional sample cost of online learning with action feedback.The conclusion identifies practical evaluation of prefix reweighting and trajectory clipping, and optimal guarantees for computationally efficient learners, as future work.Source: §5 · p. 11
A Related work
The related work places reverse-KL distillation at the intersection of offline imitation-learning theory, interactive imitation, and language-model distillation. The paper distinguishes its own comparison of online and offline reverse-KL rates from prior results built on different losses, targets, and assumptions.
A Related work
- Offline imitation work studies learning from fixed expert demonstrations, sometimes with known dynamics, realizability, or auxiliary transition data.This paper’s offline learners use teacher trajectories and specified probability feedback without auxiliary transition data.Source: A Related work · p. 16
- Rohatgi et al. [2025] study offline trajectory Hellinger learning, while Xu et al. [2026b] study forward KL; this work instead studies reverse KL under teacher coverage.The related-work discussion contrasts teacher coverage with the assumption that students cover the teacher, noting that teacher coverage suits limited-capacity distillation.Source: A Related work · p. 16
- Interactive-imitation guarantees use different comparators or assumptions, so small online regret alone does not imply this paper’s reverse-KL oracle bound.The cited settings include exploratory samples, completeness, offline information bottlenecks, noisy experts, and realizable expert Q-functions.Source: A Related work · p. 16
- Language-model distillation uses teacher- or student-generated responses and several losses, while feedback choices range from sampled-token probabilities to full distributions and teacher-preferred alternatives.Prior methods also include adaptive or skew KL, contrastive objectives, and policy-gradient updates for sequence-level reverse KL.Source: A Related work · pp. 16–17
B.1 Failure of exact empirical fitting
This appendix analyzes why exact empirical fitting on collected online trajectories can fail and develops clipping and prefix truncation tools. The construction shows how empirical comparisons can persist with a poor policy, while the auxiliary-distribution lemmas support controlling clipped losses.
B.1 Failure of exact empirical fitting
- The fitting rule uses history-dependent trajectory weights shared across candidates and retains the current policy whenever it is an empirical minimizer.Consequently, changing policies requires a strict decrease in the empirical objective.Source: §B.1 · pp. 17, 19
- For a two-policy, two-step instance, every specified exact empirical fitting update can return a policy with poor risk after ⌈λ³⌉ episodes, despite the best policy having risk ≲λe^-λ.The proposition assumes uniform initialization and the stated budgets; persistence follows when the distinguishing state sequence is not observed.Source: §B.1 · pp. 17–19
B.2.1 Clipping and prefix truncation
- A bounded-loss change of measure and a shared-mass bound provide tools for controlling how clipped quantities behave across distributions.The stated lemmas relate bounded expectations under two laws and lower-bound shared mass when KL is bounded.Source: §B.2.1 · pp. 20, 22
- The appendix’s clipped-loss analysis represents each trajectory through retained prefixes or a complete-trajectory outcome, with prefix retention determined by the candidate policy.This auxiliary distribution supplies a framework for analyzing truncated losses and likelihood ratios.Source: §B.2.1 · pp. 20–21
B.2.2 Exponential-weights analysis
The analysis uses exponential weights and bounds the cumulative weighted risk through conditional loss moments, concentration, and final-weight selection. For η = 1/256, it obtains ST ≲ TL⋆ + log(|Π|/δ).
The second moments satisfy
- The conditional loss moments and clipping bounds are used to relate the cumulative weighted risk to the comparator’s risk.The analysis then applies the resulting bound to the final output policy.Source: §B.2.2 · p. 24, The second moments satisfy · pp. 25, 27
- The proof controls weighted loss, quadratic terms, and comparator sampling error using simultaneous bounds and a union bound.The passages state that these bounds hold with probability at least 1 − δ/2.Source: The second moments satisfy · pp. 25–26
- For η = 1/256, ST ≲ TL⋆ + log(|Π|/δ), giving the cumulative weighted-risk bound used in the analysis.The positive coefficient on ST follows by choosing a sufficiently small fixed λ > 0.Source: The second moments satisfy · pp. 26–27
B.2.2 Exponential-weights analysis
- Exponential-weights analysis bounds cumulative observed loss for candidates retaining at least their initial weight, then concentration converts that bound to a risk bound.The final maximum-weight policy has weight at least 1/|Π|, so it satisfies the selection condition.Source: The second moments satisfy · p. 27
B.3 Online Hellinger learning from full feedback
This section develops online Hellinger learning with full teacher feedback, using clipped, prefix-weighted losses and exponential weights. Its guarantee is used to recover an online reverse-KL guarantee under Assumption 2.2.
B.3 Online Hellinger learning from full feedback
- Algorithm 2 replaces Algorithm 1’s local-KL score with a prefix-weighted local Hellinger score before clipping.It samples a candidate policy, collects a full episode, computes clipped increments, updates candidate weights, and returns a maximum-weight policy.Source: §B.3 · p. 28
- Theorem B.6 gives an online Hellinger guarantee for Algorithm 2, and Corollary B.9 uses it to obtain a reverse-KL guarantee under Assumption 2.2.The stated corollary recovers an approximation term B∞L⋆.Source: §B.3 · pp. 28, 32
- The analysis defines auxiliary distributions to express clipped Hellinger loss and establish identities used for clipping and prefix truncation.The stated proof uses c = 1 for comparison and c = 1/4 to bound prefix truncation.Source: §B.3 · pp. 28–29
- The clipping and truncation lemmas relate the auxiliary-distribution distances to Hellinger losses, supporting the guarantee’s transfer to reverse KL.The proof applies these relations to the conditional means and uses the resulting bounds in the online analysis.Source: §B.3 · pp. 29–31
B.4 Estimation and approximation lower bounds
The lower bounds characterize estimation and approximation limits, and show that online action feedback can impose an additional sample cost. Offline action feedback retains the stated lower bounds under the ratio-bound assumption.
B.4 Estimation and approximation lower bounds
- Under Assumption 2.1, full-feedback learners in either trajectory setting face an estimation lower bound, and under Assumption 2.2 the same bound extends to offline action feedback.The proof states that restricting feedback preserves the offline action-feedback lower bound.Source: §B.4 · p. 32, §B.4.1 · p. 34
B.4.2 Approximation lower bound
- At optimal estimation, any learner with Cest ≤ κB on the specified H = 2 instances must have Capp = Ω(B), for both online and offline full feedback.Under Assumption 2.2, the conclusion uses B∞ in place of B and also holds for offline action feedback.Source: §B.4 · p. 32
- The approximation lower bound uses an identification construction in which a trajectory reaching the relevant state reveals an equality test about the teacher’s index.The stated construction says one equality test also simulates an offline teacher trajectory.Source: §B.4.2 · p. 35
B.5 The additional cost of online action feedback
- For online action feedback, querying an action can reveal only whether its label matches the unknown teacher label, leaving other teacher preferences unobserved.The theorem concerns arbitrary causal exploration with output in Π and gives a lower-bound instance under Assumption 2.2.Source: §B.5 · p. 36
C.1 Full feedback
The full-feedback section establishes a minimax-optimal estimation rate in the realizable setting and gives a counterexample to selection by an untruncated objective. The action-feedback section shows that offline observations can be indistinguishable even when reported action probabilities are available.
C.1 Full feedback
- For a teacher in Π, both empirical objectives have minimum zero, and nonnegativity of local KL losses supports the stated comparison.Source: §C.1 · p. 37
- For realizable instances, the estimation rate established for full-feedback learning is minimax optimal.The proof says the empirical objectives have the same minimizers and the lower bound makes the estimation rate minimax optimal.Source: §C.1 · p. 37
- The untruncated objective ranks πbad ahead of πgood, despite πgood being the better candidate under the population risks described in the construction.The population objective of πbad is smaller by at least 3/16.Source: §C.1 · p. 38
C.2 Action feedback
- In the offline action-feedback construction, two teachers produce identical observations on event E, including the reported probability of the observed action.Because the conditional output distributions coincide on E, the learner cannot select the correct policy for both teachers on that event.Source: §C.2 · pp. 39–40
D.1 Second moments and clipping
The appendix gives sufficient conditions for second-moment control and clipping, relates bounded likelihood ratios to KL and Hellinger quantities, and collects concentration and testing bounds.
D.1 Second moments and clipping
- A pathwise accumulated-KL bound Kπ(τ) ≤ B for B ≥ 1 suffices for Assumption 2.1, while Assumption 2.2 implies it with B = 4(B∞ + 1).Source: §D.1 · p. 41
- Under Assumption 2.1, every policy has L(π) ≤ B, and the clipping argument uses f = min{1, K}.Source: §D.1 · pp. 42–43
D.2 KL and Hellinger inequalities
- With P ≪ Q and P(z)/Q(z) ≤ e^b, Lemma D.5 converts Hellinger distance to a KL bound and gives KL(P∥Q) ≤ b.Source: §D.2 · p. 44
D.3 Concentration and testing bounds
- D.3 gives fixed-time exponential and empirical-minimization bounds for bounded scores, and formulates adaptive equality tests that query whether j = Θ.Source: §D.3 · pp. 45–46