PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Qiwei Di, Xuheng Li, Kaixuan Ji, Chenggong Zhang, Heyang Zhao, Quanquan Gu

arXiv:2609.38666v1cs.LGcs.AIstat.ML

TL;DR

The paper asks when on-policy distillation’s benefits and fragility arise relative to off-policy methods. It compares their aggregation targets for multiple teachers, develops algorithms with logarithmic regret guarantees, and finds that reverse KL can preserve a confident expert’s preferences yet remain vulnerable to misleading feedback and long-horizon prefix effects.

  • Problem

    Although prior work identifies conditions where on-policy interaction helps, it does not directly explain the practical fragility of OPD or when it succeeds or fails.

    Source: §1 · p. 2

  • Method

    The paper studies sequential multi-teacher distillation, characterizes forward- and reverse-KL aggregation targets, and develops off-policy and on-policy algorithms to learn them.

    Source: §1 · p. 2

  • Results

    Reverse-KL aggregation can better preserve a confident expert’s preference with uninformative teachers, but misleading feedback can suppress correct responses and long horizons can favor incorrect prefixes.

    Source: §1 · p. 2, §7 · p. 15

  • Takeaways & Limitations

    The distinct aggregation objectives offer possible explanations for both OPD’s benefits and its fragility relative to SFT.

    Source: §7 · p. 15

    The long-horizon prefix reversal is illustrated with a 0.9-weight expert assigning 0.99 probability to the correct first token and a uniform second teacher.

    Source: §6.3 · p. 14

Abstract

from arXiv · show

On-policy distillation (OPD) learns from teacher feedback on student-generated responses and has shown promise in reducing forgetting relative to supervised fine-tuning (SFT). However, its benefits and fragility remain incompletely understood. We study sequential distillation from multiple teachers, where the student minimizes its average divergence from the teachers. Forward Kullback--Leibler (KL) divergence yields a weighted arithmetic mixture, while reverse KL yields a normalized weighted geometric aggregate. We develop algorithms that learn these targets under off-policy and on-policy feedback, respectively, establishing logarithmic regret bounds in the tabular setting and extending the analysis to function approximation. By analyzing these aggregation targets, we identify mechanisms that help explain both the benefits and fragility of OPD. Relative to forward KL, reverse KL can better retain a confident expert's preferences under uninformative feedback, but is more sensitive to teachers that assign very low probabilities to correct responses. Its token-level conditionals also reveal a dependence on continuation distributions that can favor incorrect prefixes over long horizons.

1 Introduction

The paper studies when on-policy distillation’s benefits and fragility arise by comparing the aggregation targets induced by off- and on-policy feedback. It characterizes those targets, develops learning guarantees, and analyzes how teacher confidence and long-horizon continuations affect them.

1 Introduction

  • The paper frames an unresolved question: what pros and cons follow from the distinction between off-policy and on-policy feedback in knowledge distillation?Prior work establishes some conditions where on-policy interaction can outperform offline learning, but does not directly explain practical OPD fragility.Source: §1 · p. 2
  • In sequential multi-teacher distillation, the student minimizes its average divergence from the teachers, linking off- and on-policy distillation to forward- and reverse-KL objectives.The paper derives closed-form aggregation targets to clarify the difference between SFT and OPD.Source: §1 · p. 2
  • Forward KL gives a weighted arithmetic mixture, whereas reverse KL gives a normalized weighted geometric mean, with token-level conditionals derived for autoregressive generation.The paper also develops off-policy and on-policy algorithms for learning these respective targets, with logarithmic regret guarantees in tabular and function-approximation settings.Source: §1 · p. 2
  • Reverse-KL aggregation can preserve a confident expert’s preference against uninformative teachers, but misleading feedback can suppress correct responses and long horizons can favor incorrect prefixes.These comparisons offer possible explanations for both the benefits and fragility of OPD.Source: §1 · p. 2

2 Related Work

Related work examines OPD methods, its capability-retention benefits and failure modes, and the theoretical value of on-policy interaction. This paper’s contribution is to characterize the distinct aggregation targets for multiple teachers and analyze how those targets affect expert preferences.

2 Related Work

  • Prior OPD methods vary training objectives, teacher feedback, data collection, and supervision sources; multiple teachers motivate studying how distillation combines knowledge.Examples include skew-KL or hybrid objectives, outcome-weighted feedback, early rollout termination, self-teachers, and shared students supervised by multiple teachers.Source: §2 · p. 3
  • Studies link OPD to reduced forgetting but also report sensitivity to teacher–student mismatch, long prefixes, unreliable feedback, and token-level reward instability.Prior theoretical work establishes settings where on-policy interaction offers advantages, while the paper identifies distinct optimal targets and analyzes confidence, misleading feedback, and autoregressive continuations.Source: §2 · pp. 3–4
  • The off-policy analysis builds on behavior cloning and log-loss prediction, while the on-policy analysis adapts KL-regularized bandit and reinforcement-learning techniques to sequential distillation.The reverse-KL analysis treats teacher log-probability ratios as rewards and accounts for dependence among rollouts from the same sampled teacher.Source: §2 · p. 4
  • For general function approximation, the paper adapts generalized eluder-dimension analyses to the autoregressive setting.This extends prior eluder-based work on contextual bandits and value-function approximation.Source: §2 · pp. 4–5

3 Preliminaries

The preliminaries define sequential distillation as one student learning from multiple teachers over successive rounds. They specify the teacher-sampling and feedback protocols, then define the aggregate target and regret used to compare student policies with it.

3 Preliminaries

  • The setting uses finite context and token spaces and a finite generation horizon, with responses generated autoregressively as token sequences.The student seeks one policy aggregating all teachers rather than matching only the teacher sampled in the current round, an objective designed to avoid catastrophic forgetting.Source: §3 · p. 5
  • Each round samples a teacher, draws contexts from that teacher’s distribution, and gives the batch a shared teacher identity.Off-policy feedback uses teacher-generated rollouts, optionally with logits; on-policy feedback uses student-generated rollouts with per-token teacher log-probabilities.Source: §3 · p. 5
  • The target policy minimizes average divergence from the teachers, and regret is cumulative divergence between the student’s policies and the corresponding target.The framework focuses on forward KL for off-policy distillation and reverse KL for on-policy distillation; it can also treat the initial reference policy as an additional teacher.Source: §3 · p. 6

4 Off-policy Distillation & Forward KL

Off-policy distillation learns the forward-KL aggregate from teacher-generated data. Its target is a weighted mixture whose prefix-dependent teacher weights have a posterior interpretation, and the proposed per-prefix algorithm has logarithmic regret guarantees.

4 Off-policy Distillation & Forward KL

  • The forward-KL target is a weighted mixture of teacher distributions, with token-level conditionals using weights that depend on the context and generated prefix.These weights are posterior probabilities that each teacher generated the observed context and prefix, so teachers that better explain it receive greater next-token weight.Source: §4 · pp. 6–7

4.1 Algorithm Design

  • The per-prefix algorithm estimates target conditionals from observed next-token frequencies, or uses teacher conditional probability vectors when full logits are available.It accumulates normalized counts with smoothing to avoid zero predicted probabilities and assigns a uniform distribution to unvisited states.Source: §4.1 · pp. 7–8

4.2 Theoretical Guarantees

  • For fixed context, token space, and horizon, the forward-KL algorithm’s average divergence from its target decays as O(log T/T), with bounds under both feedback protocols and with or without teacher logits.The bound’s exponential dependence on horizon reflects growth in the effective state space, which has SA^h state–action pairs at step h.Source: §4.2 · p. 8

5 On-policy Distillation & Reverse KL

This section characterizes the reverse-KL target and develops an on-policy algorithm to learn it from teacher log-probabilities on student-generated responses. It proves a high-probability logarithmic average-divergence guarantee under stated assumptions.

5 On-policy Distillation & Reverse KL

  • The unique reverse-KL target has token conditionals that combine current-token geometric scores with continuation values over possible suffixes.Thus, the target accounts for both the current token and the continuations that follow it.Source: §5 · pp. 8–9

5.1 Algorithm Design

  • The algorithm estimates log-geometric token scores by averaging teacher log-probability feedback at visited student-generated context–prefix–token triples, then uses optimistic scores in a backward recursion.The analysis assumes teacher token log-probabilities are bounded relative to a fixed full-support reference policy; this is needed because log-probabilities can be unbounded below.Source: §5.1 · pp. 9–10

5.2 Theoretical Guarantees

  • For fixed S, A, H, B, and δ, the reverse-KL divergence from the learned policy to the target, averaged over contexts and rounds, decays as eO(log T/T) with high probability.The bound applies under Assumption 5.3; its SAH factor reflects the effective state–action space size.Source: §5.2 · p. 11

6 Comparison of Distinct Aggregation Objectives

The section compares forward and reverse aggregation as ways to combine teacher capabilities, identifying conditions where reverse aggregation can preserve or undermine an expert’s preference.

6.1 Uninformative Teachers

  • At N=10 and α=0.9, the targets cross at q≈0.7617; above the unique threshold, reverse aggregation assigns the correct response higher probability.The comparison assumes one confident expert and otherwise uniform teachers; for N≥3, the threshold is unique.Source: §6.1 · pp. 11–12, Figure 1 · p. 12

6.2 Misleading Teachers

  • A misleading teacher can drive reverse-aggregate correct-response probability arbitrarily close to zero even at small weight, while forward aggregation preserves the expert’s contribution.Forward aggregation sums nonnegative probabilities; reverse aggregation averages log-probabilities, which can become arbitrarily negative.Source: §6.2 · p. 13
  • The authors offer this sensitivity as a possible explanation for OPD failures under teacher–student mismatch; reported remedies include clipping extreme log-ratios or using an SFT cold start.The cited studies report improved stability from clipping or compressing extreme token-level log-ratios, and improved subsequent OPD after an SFT cold start.Source: §6.2 · p. 13

6.3 Long Horizons

  • Although exponent α < 1 preserves complete-response rankings, continuation sums can favor tokens with diffuse continuations over tokens with concentrated continuations.This can steer generation toward a prefix the expert considers less likely, even while the expert’s ranking of complete responses remains unchanged.Source: §6.3 · pp. 13–14
  • The prefix reversal can occur at sufficiently large horizons even when the uniform teacher receives an arbitrarily small aggregation weight.The comparison holds with positive-probability teacher policies and a fixed expert preference for the correct first token.Source: §6.3 · pp. 13–14
  • Longer student-generated prefixes can make teachers less able to recover correct answers, while stopping rollouts early can improve training stability.Related empirical work also reports late-stage collapse on longer rollouts, with instability spreading from later tokens toward earlier positions.Source: §6.3 · p. 14

7 Conclusion

The conclusion connects off- and on-policy feedback to distinct aggregation objectives and reports both advantages and vulnerabilities of reverse-KL aggregation. It presents these findings as possible explanations for OPD’s benefits and fragility relative to SFT.

7 Conclusion

  • The paper connects off- and on-policy feedback protocols to distinct aggregation objectives by characterizing forward- and reverse-KL targets and establishing learning guarantees.Source: §7 · pp. 14–15
  • Reverse-KL aggregation can preserve a confident expert’s preference with uniform teachers, but can suppress correct responses under misleading feedback or favor incorrect prefixes over long horizons.The authors offer these comparisons as possible explanations for OPD’s benefits and fragility relative to SFT.Source: §7 · p. 15

AI use statement

  • The authors used AI models to help articulate and summarize ideas, draft and polish the manuscript, and assist with technical lemmas and proof strategies; they report verifying the proofs and taking responsibility for the final content.Source: AI use statement · p. 15

A Proof of the KL Aggregation Theorems

The appendix proves the KL aggregation results by reducing the objective to contextwise optimization. For forward KL, it identifies the minimizing sequence distribution and derives its autoregressive token conditionals.

A Proof of the KL Aggregation Theorems

  • The proof reduces the objective to a separate optimization at each context, with zero-weight contexts having no effect on the objective.Policies at zero-weight contexts may be chosen arbitrarily.Source: A Proof of the KL Aggregation Theorems · p. 15

A.1 Proof of Theorem 4.1

  • For forward KL, the proof identifies the weighted teacher aggregate as the minimizing sequence distribution and establishes its uniqueness at contexts with positive weight.Source: §A.1 · p. 16
  • The proof derives the forward-KL target’s token conditionals by marginalizing over continuations and taking ratios of prefix probabilities.At a prefix with zero target probability, the conditional can be chosen arbitrarily without changing the sequence distribution.Source: §A.1 · pp. 16–17

A.2 Proof of Theorem 5.1

The proof of Theorem 5.1 constructs the reverse-KL minimizing sequence distribution on the common support of positive-weight teachers. It then derives autoregressive token conditionals and addresses choices at null prefixes.

A.2 Proof of Theorem 5.1

  • The reverse-KL target is supported on the common support of all positive-weight teachers; assigning mass outside it gives an infinite objective.Source: §A.2 · p. 18
  • The proof identifies the reverse-KL target as the unique minimizing sequence distribution.It establishes the target’s normalization through a backward recursion.Source: §A.2 · pp. 17–18
  • The proof obtains token conditionals from prefix probabilities, uniquely determined at positive-probability prefixes.At null prefixes, conditional choices can still induce the minimizing sequence law; uniqueness concerns the sequence distribution and positive-probability conditionals.Source: §A.2 · p. 19

B Proof of Theorem 4.3

The proof establishes the feedback identities used to analyze the forward target and develops a regret bound. It also accounts for rollouts sharing a latent teacher index and distinguishes observed logits from sampled token feedback.

B Proof of Theorem 4.3

C Proof of Theorem 5.5

The proof controls estimation errors for the reverse-KL output policy using a high-probability event and bounds on observed rollouts. The accompanying lemma proofs use concentration and conditional-expectation arguments.

C Proof of Theorem 5.5

  • The Theorem 5.5 proof bounds optimistic token-score estimates on visited and unvisited triples, then controls their errors using observed rollouts and visit counts.It first establishes simultaneous bounds on a high-probability event and concludes with a bound holding with probability at least 1 − 3δ.Source: C Proof of Theorem 5.5 · pp. 28, 30, 31, 32

D.1 Proof of Lemma B.1

  • The Lemma B.1 proof separates the single-action case from the multi-action case and uses beta-function bounds to complete the argument.The excerpt identifies A = 1 as a separate case and states that the proof is complete after the final bound.Source: §D.1 · pp. 33–34

D.2 Proof of Lemma B.2

  • The Lemma B.2 proof uses conditional Jensen’s inequality and an exponential-moment bound for a bounded random variable.It establishes that the conditional mean parameter lies between 0 and L before treating the two signs of y.Source: §D.2 · pp. 34–35

D.3 Proof of Lemma 5.4

  • The Lemma 5.4 proof combines concentration bounds to obtain a simultaneous high-probability bound over rounds, levels, and context-action triples.The argument applies concentration inequalities and union bounds across these indices.Source: §D.3 · pp. 36–37

E Missing Proof in Section 6

These proofs and sensitivity analyses characterize when the aggregation targets differ, including their responses to misleading teachers and long continuation horizons. The function-approximation extension describes a policy-mixture algorithm and its computational cost.

E Missing Proof in Section 6

  • For N ≥ 3, reverse and forward targets cross once as q increases; for N = 2, the reverse target is larger throughout q ∈ (1/2, 1), though both remain below the expert.The crossing for N ≥ 3 separates the ranges where each target retains more correct-response probability.Source: E Missing Proof in Section 6 · p. 39
  • For fixed expert weight and confidence, the reverse target’s correct-response probability can be made arbitrarily small while both teachers retain full support.The construction takes the misleading teacher’s correct-response probability ε toward zero.Source: E Missing Proof in Section 6 · p. 39

F Parameter Sensitivity of Aggregation Targets

  • The parameter-sensitivity curves are deterministic evaluations of closed-form targets, not outcomes from training a student.Each comparison holds other parameters fixed, and dashed lines mark the expert’s probability.Source: F Parameter Sensitivity of Aggregation… · p. 41

Uninformative teachers.

  • With uninformative teachers, the reverse target’s advantage depends on expert weight and response-space size; for N = 2 it is larger throughout q ∈ (1/2, 1).At α = 0.9, increasing N from 10 to 1000 raises the confidence threshold for a reverse-target advantage.Source: Uninformative teachers · p. 42

Misleading teachers.

  • At q = 0.99 and α = 0.9, the reverse probability equals 1/2 near ε = 10^-18, versus 10^-198 at α = 0.99; for every α < 1, it can decline as ε approaches zero.Forward aggregation remains above αq; the extreme values do not estimate how often such predictions occur in practice.Source: Misleading teachers · p. 42

Long horizons.

  • At r = 0.99 and δ = 0.01, the first horizon with reverse first-token probability below 1/2 is H = 10, 68, or 717 for α = 0.5, 0.9, or 0.99, respectively.In this construction, forward aggregation is independent of H; differing continuation distributions drive the reversal, while expert confidence and weight delay it.Source: Long horizons · p. 42

G Forward KL with Function Approximation

  • The function-approximation setting uses a finite class of autoregressive policies whose conditionals share a common representation, under a realizability assumption.The assumption requires a policy in the class to match the specified target across contexts and responses.Source: G Forward KL with Function Approximation · pp. 42–43

G.1 Algorithm Design

  • The algorithm aggregates candidate policies at the token level and updates their weights to favor lower cumulative empirical loss; evaluating the mixture at each decoding step is costly.It uses sampled tokens when logits are unavailable and full teacher probabilities otherwise.Source: §G.1 · pp. 43, 45

G.2 Theoretical Guarantee

This section establishes theoretical guarantees for off-policy and on-policy distillation. It covers regret bounds for finite or finitely coverable policy classes and reverse-KL function approximation, as well as a high-probability cumulative guarantee for on-policy autoregressive rollouts.

G.2 Theoretical Guarantee

  • Under off-policy distillation, Algorithm 3’s regret is controlled by a good candidate’s prior weight and approximation error, which vanishes under realizability.The theorem covers function approximation with and without teacher logits; it also gives a high-probability guarantee.Source: §G.2 · pp. 45–46

G.3 Function Class with Finite Covering Number

  • For finite-covering policy classes, an ε-cover can be used; for every T ≥ 2, choosing ε = 1/T gives expected regret at most eO(dH log T).Under off-policy feedback, the expected-regret bound applies both with and without teacher logits.Source: §G.3 · pp. 51, 54

H Reverse KL with Function Approximation

  • For reverse-KL function approximation, the paper introduces a batched generalized Eluder dimension measuring current-query disagreement against squared disagreement on completed trajectory batches.Because the estimator updates only after each round, the uncertainty measure excludes feedback from the current batch.Source: H Reverse KL with Function Approximation · p. 54

H.1 Algorithm Design

  • Algorithm 4 estimates the score by least squares, adds an optimism bonus based on the batched dimension, and constructs its output policy by backward propagation.The method also establishes pointwise optimism and an error bound for its optimistic score.Source: §H.1 · pp. 55–56

H.2 Theoretical Guarantee

  • Under Assumptions 5.3 and H.1, Algorithm 4 has a regret guarantee parameterized by the batched generalized Eluder dimension, with specified choices of λ and β.The stated parameter choice is λ = B^2 and β = 32B^2 log(2|F|/δ).Source: §H.2 · p. 56
  • Under β ≥ 32B^2 log(2|F|/δ), a high-probability event gives a bound simultaneously for every round t and z.The proof defines this event after applying Lemma H.3.Source: §H.2 · p. 57
  • The on-policy concentration argument uses each rollout’s conditional marginal law and does not require independence within the batch.The proof combines a bounded, measurable quantity with concentration and a union bound.Source: §H.2 · pp. 58–59
  • Summing the per-query bound uses the generalized Eluder dimension because the observed batches form a valid sequence of autoregressive trajectories.Substitution of this dimension-based inequality completes the proof of Theorem H.4.Source: §H.2 · p. 59

H.3 Proof of Lemma H.3

The proof controls estimation errors through conditional concentration and a supermartingale argument, then derives confidence and optimism bounds. The appendix also collects concentration definitions and auxiliary inequalities used in the analysis.

H.3 Proof of Lemma H.3

  • The argument handles within-round dependence without assuming independence among observations in a batch.It uses Jensen’s inequality to handle dependence within a round.Source: §H.3 · pp. 60–61
  • The proof derives a pointwise confidence bound and uses it to establish optimism.It applies conditional expectations and a supermartingale argument, then takes a union bound over candidates to obtain a simultaneous event.Source: §H.3 · pp. 60–62

I.1 Uniform Bernstein Concentration in Bakhtiari et al. (2026)

  • The concentration appendix defines CGF-like and sub-ψ processes and identifies sub-gamma processes as a special case.The sub-gamma definition uses ψ(λ) = λ^2/[2(1−vλ)].Source: §I.1 · p. 62

I.2 Other Auxiliary Lemmas

  • The auxiliary-lemma appendix collects concentration and inequality results, including bounds for nonnegative processes, Stirling’s inequality, and Ville’s inequality.It also states a lemma concerning the KL divergence between distributions proportional to exponentiated functions.Source: §I.2 · pp. 63–64
Loading 2609.38666v1…