Distillation Variants › On-Policy Distillation

All three methods above share a common assumption: the student is trained on the same data distribution as the teacher. For discriminative tasks like classification, this is fine. For generative tasks, it creates a fundamental problem.

Consider training a small student language model to distill a large teacher. The standard approach samples a batch of text from a dataset, runs both teacher and student on it, and minimizes the KL divergence between their output distributions at each position. But the text in the dataset was written by humans. The student, when deployed, will generate its own text. That text will follow its own distribution, which is different from the human-written distribution in ways that depend on the student's current capabilities and biases. The student is being trained to handle human-generated inputs, not the inputs it will encounter.

This is the distribution shift problem in sequence generation distillation, and on-policy distillation is the solution. The name comes from reinforcement learning terminology: "on-policy" means training on data generated by the current policy (the model being trained), as opposed to "off-policy" training on a static dataset generated by some other distribution. The key insight is that the model that will be evaluated at inference time is the current student, so it makes sense to train it on the kind of inputs it will produce.

The mismatch between training distribution and inference distribution is called exposure bias. During training on a static dataset, the student always sees gold-standard input prefixes: "The quick brown fox jumps". During inference, it generates its own prefixes: perhaps "The quick brown fox leaps" at step 4. The word "leaps" was never in the training data for position 4, so the student has never learned to continue from it. Small errors compound: each generation step produces a slightly wrong token, which pushes the model further off-distribution, where it has no reliable signal.

The compounding nature of this error is worth dwelling on. After one generation step, the input is slightly off the training distribution. After five steps, it is further off. After twenty steps, the accumulated deviation may be large enough that the student's outputs are qualitatively different from what it would produce on clean inputs. This is visible in practice as repetition loops (where the model gets stuck generating the same phrase repeatedly), incoherence (where later sentences contradict earlier ones), and degradation in quality over long sequences.

Exposure bias affects standard language model training too, not just distillation. Scheduled Sampling (Bengio et al., 2015) is a classical remedy: mix gold tokens with the model's own predictions at training time. The fraction of model-generated tokens starts low and increases over training, progressively exposing the model to its own distribution. On-policy distillation extends this idea specifically for the distillation setting, making it the default rather than a scheduled addition.

In on-policy distillation, the student generates sequences using its own current policy, then uses those generated sequences as training inputs. Let $x_{1:t}$ be the student's own generated prefix. At position $t+1$, the objective minimizes:

$$\mathcal{L}_{\text{on-policy}} = \mathbb{E}_{x_{1:t} \sim \pi_S} \left[ D_{\text{KL}}\!\left( p_T(\cdot \mid x_{1:t}) \,\|\, p_S(\cdot \mid x_{1:t}) \right) \right]$$

where:

• $\pi_S$: the student's own policy (probability distribution over sequences)
• $x_{1:t} \sim \pi_S$: sequences generated by sampling from the student
• $p_T(\cdot \mid x_{1:t})$: teacher's next-token distribution given the student-generated prefix
• $p_S(\cdot \mid x_{1:t})$: student's next-token distribution given the same prefix
• $D_{\text{KL}}$: KL divergence measuring how far the student's distribution is from the teacher's

The key difference from standard distillation is the expectation over $\pi_S$: training inputs come from the student itself, not from a static dataset. The student learns to mimic the teacher on the distribution it encounters during inference.

This objective has a natural interpretation: after each generation step by the student, ask the teacher "given what the student just produced, what would you predict next?" and train the student to match that prediction. The teacher provides a corrective signal at every position in the student's own generated sequences, guiding the student toward teacher-quality continuations even when it starts from its own text.

An important subtlety in on-policy distillation is the direction of the KL divergence. Standard forward KL, $D_{\text{KL}}(p_T \| p_S)$, is mode-covering: it penalizes the student heavily for assigning low probability to anything the teacher assigns high probability to. When the teacher is a large, capable model with a rich, multimodal distribution, this coverage requirement causes the student to spread probability mass across many low-probability teacher tokens. The student tries to cover all the teacher's modes, which leads to averaged-out, overly cautious generations that hedge across many possibilities rather than committing to any one good output.

MiniLLM (Gu et al., 2023) proposes using reverse KL, $D_{\text{KL}}(p_S \| p_T)$, in the on-policy setting:

$$\mathcal{L}_{\text{reverse-KL}} = \mathbb{E}_{x_{1:t} \sim \pi_S} \left[ D_{\text{KL}}\!\left( p_S(\cdot \mid x_{1:t}) \,\|\, p_T(\cdot \mid x_{1:t}) \right) \right]$$

Reverse KL is mode-seeking: it penalizes the student for placing probability mass where the teacher does not. This encourages the student to be precise about the modes it covers rather than spreading evenly. For language generation, this produces sharper, more confident output at the cost of some diversity. The student learns: "focus your probability mass where the teacher is confident." Rather than trying to cover every possible completion the teacher might give, it covers the most likely ones well.

The intuition behind forward versus reverse KL can be grounded in a simple example. Suppose the teacher assigns 40% probability to "cat" and 40% to "dog" as next tokens, with the remaining 20% spread across other words. Forward KL forces the student to assign significant probability to both "cat" and "dog" and all the minor modes. Reverse KL allows the student to pick one of the two dominant modes and concentrate its probability there, accepting that it will miss the other major mode (and all minor modes) in exchange for being precise about the mode it does cover.

Whether this tradeoff is desirable depends on the application. For open-ended generation, diversity matters and forward KL is often preferred. For instruction-following, where there is typically one best response, reverse KL's mode-seeking behavior produces crisper, more decisive outputs that users typically prefer.

The gradient of the reverse KL under the student's policy can be computed using the policy gradient theorem, connecting on-policy distillation to reinforcement learning:

$$\nabla_\theta \mathcal{L}_{\text{reverse-KL}} = \mathbb{E}_{x \sim \pi_S} \left[ \log \frac{p_S(x)}{p_T(x)} \cdot \nabla_\theta \log p_S(x) \right]$$

where $\theta$ are the student's parameters. This is the REINFORCE estimator applied to the distillation objective. In practice, variance reduction techniques like baseline subtraction are needed to make this stable. The log-ratio $\log \frac{p_S(x)}{p_T(x)}$ is a reward signal: if the student's probability exceeds the teacher's, the gradient pushes the student to reduce that probability. If the teacher's probability exceeds the student's, the gradient pushes the student to increase it.

Training with on-policy distillation requires careful engineering. You must:

1. Generate from the student: sample a batch of sequences from the student model using autoregressive decoding
2. Query the teacher: run the teacher forward on the same sequences to get its probability distributions at each position
3. Compute the loss: compute the KL divergence between teacher and student at each generated token position
4. Update the student: backpropagate through the student only (teacher is frozen)
5. Repeat: the student's policy evolves, so the on-policy distribution changes throughout training

This loop is more expensive than offline distillation because every training step requires a generation pass. A standard offline distillation step processes a batch in one forward pass through both models. An on-policy step requires an autoregressive generation loop (one forward pass per token in the sequence) plus the forward pass for the loss computation. For a 256-token sequence, on-policy training is roughly 256 times more expensive per example than offline training. Techniques like caching teacher outputs for recently generated sequences can reduce overhead, but the fundamental cost difference remains significant.

Several engineering strategies have emerged to manage this cost. One is to use short generation prefixes: generate 20-50 tokens from the student, then compute the loss on those tokens rather than generating full sequences of hundreds of tokens. Another is mixed training: alternate between on-policy batches (expensive, for distribution coverage) and offline batches (cheap, for data efficiency). A third is to use a smaller, faster version of the teacher as the on-policy signal generator, with the full large teacher used only for periodic evaluation.

The implementation in production LLM distillation pipelines often uses a variant called SeqKD (sequence-level knowledge distillation, Kim and Rush, 2016), where the teacher generates full sequences, those sequences are added to the training corpus, and the student is trained on them with standard maximum likelihood. This is computationally equivalent to offline training after the one-time generation cost, but ensures the student's training data includes teacher-generated text rather than only human-generated text, partially addressing the distribution shift problem at lower computational overhead than true on-policy generation.