On-policy distillation (GKD), explained

On-policy distillation trains a student model on sequences it generates itself, using a teacher model to score every token of those sequences. The student samples a response, the teacher runs one forward pass over it, and the student is updated to reduce the divergence between its next-token distribution and the teacher's at each position it actually visited. The best-known formulation is Generalized Knowledge Distillation (GKD), which also lets you mix in fixed data and choose the divergence. It combines the two things that make other methods work: like RL, it trains on the student's own distribution, so it fixes the mistakes the student actually makes; like distillation, it gets a dense, per-token training signal instead of one scalar per sequence. In the post-training stack it typically follows an off-policy distillation or SFT stage and either replaces or precedes RL for small models.

The problem and where it came from#

Standard distillation, whether sequence-level (SFT on teacher outputs) or token-level (matching teacher logits on fixed sequences), trains the student only on prefixes written by the teacher or by humans. At inference the student conditions on its own prefixes. When it makes an error the teacher would not have made, it enters a state that never appeared in training, has no idea how to recover, and errors compound over long generations. This is exposure bias, the sequence-model version of a problem imitation learning addressed years earlier: in sequential prediction a learner trained only on expert trajectories makes mistakes that compound over the horizon, and Ross, Gordon and Bagnell (2011) showed with DAgger that repeatedly querying the expert on the states the learner visits yields a policy with guarantees under the learner's own state distribution.

Agarwal, Vieillard, Zhou, Stanczyk, Ramos, Geist and Bachem (2023) brought this idea to LLM distillation as GKD. The student generates sequences, the teacher labels each token with its distribution, and the student minimises a divergence between the two at those positions, with no gradient through the sampling. Across T5 students on summarisation, translation and arithmetic reasoning, on-policy GKD gave average relative gains over the initial student of 2.1×, 1.7× and 1.9× the gains achieved by baseline KD methods. In concurrent work, Gu, Dong, Wei and Huang (2023) proposed MiniLLM, which optimises the sequence-level reverse KL with a policy-gradient method and showed lower exposure bias and better long-text generation. The technique moved into frontier training with the Qwen3 report (Yang et al., 2025): its small models were trained with off-policy distillation followed by on-policy distillation from Qwen3-32B or Qwen3-235B-A22B, and for Qwen3-8B on maths and code, on-policy distillation beat RL from the same checkpoint (74.4 versus 67.6 on AIME'24) with 1,800 GPU hours against 17,920.

The objective#

Let $p_T$ be the teacher and $q_{\theta }$ the student, both autoregressive over a shared vocabulary. For a prompt $x$ and a response $y$ of length $|y|$, define the per-sequence divergence as the average token-level divergence along that sequence:

where $\mathcal{D}$ is a divergence between two next-token distributions over the vocabulary. GKD minimises a mixture of this quantity on fixed data and on student samples:

$\lambda \in [0,1]$ is the student data fraction. $\lambda =0$ with forward KL is ordinary supervised token-level KD; $\lambda =1$ is fully on-policy distillation. The sampling distribution $q_{\theta }$ in the second term is treated as fixed when differentiating: gradients flow through the per-token divergences, not through the choice of which sequences were sampled. The paper notes that this keeps training stable and cheap.

The divergence choice

GKD uses three families. Forward KL $\mathrm{KL}(p_T\Vert q_{\theta })$ is mass-covering. Reverse KL $\mathrm{KL}(q_{\theta }\Vert p_T)$ is mode-seeking: it lets a student with limited capacity concentrate on outputs the teacher finds likely rather than spreading over everything the teacher might say. The generalised Jensen-Shannon divergence interpolates between them:

whose gradients behave like forward KL as $\beta \to 0$ and like reverse KL as $\beta \to 1$. In GKD's instruction-tuning experiments, on-policy GKD with reverse KL substantially outperformed supervised KD and performed much better than forward KL; for combining with RL the authors recommend reverse KL or JSD(0.9).

Why it works: an RL reading

Consider the sequence-level reverse KL, $\mathrm{KL}(q_{\theta }\Vert p_T)={\mathbb{E}}_{y\sim q_{\theta }}[{\sum }_t\log q_{\theta }(y_t\mid \cdot )-\log p_T(y_t\mid \cdot )]$. Minimising it is reinforcement learning with a dense per-token reward

which is positive when the teacher likes the sampled token more than the student does and negative when less. MiniLLM optimises exactly this with policy gradients. GKD's per-position reverse KL, computed over the full vocabulary at each visited state, is a lower-variance version of the same signal: instead of scoring only the token that was sampled, it compares the whole distribution at that position. The KL-regularised RL objective makes the connection explicit. Maximising ${\mathbb{E}}_q[R]-\beta \mathrm{KL}(q\Vert p_{\text{ref}})$ has optimum $q^{\ast }\propto p_{\text{ref}}\exp (R/\beta )$. Set the reference to the teacher and the reward to zero, and the optimum is the teacher itself; on-policy distillation is that RL problem with the KL term doing all the work. Add a task reward back and you get GKD's combination with RL:

The practical difference from RL is information per sample. A verifier reward gives one bit-sized scalar per sequence, and the algorithm must work out which of thousands of tokens deserved credit. The teacher gives a full distribution at every token, which says exactly where the student went wrong and what it should have preferred.

Side by side with its neighbours

Worked example: scoring one student sample#

Suppose the student samples a three-token continuation, and at each position the student and teacher assign these probabilities to the token that was sampled:

The single-sample estimate of the sequence's reverse KL is $-(0-1.386+0.693)=0.693$ nats. Read as rewards, position 2 is where the student went wrong: it put 0.4 on a token the teacher gives 0.1, and the update will push that probability down hard. Position 3 is a token the student undervalues, and it will be pushed up. Position 1 is already right and contributes nothing. A sequence-level verifier would give all three positions the same credit. GKD goes further and compares full distributions at each of these three states, not just the sampled tokens, so position 2 also tells the student which alternatives the teacher preferred.

A toy run: exposure bias, measured exactly#

We built a small Markov model where exposure bias can be computed exactly. There are six tokens: four "on-track" and two "off-track". The teacher stays on track with probability 0.99 and, if it ever goes off track, returns with probability 0.90. The student starts slightly worse on track (0.90) and much worse at recovery (0.30): once off track it tends to stay there, which is the compounding behaviour seen in real small models. Sequences have 20 tokens. Training is SGD (learning rate 0.5) on a per-state logit table with 16 sequences per step, minimising the token-level divergence from the teacher at every visited state, for 300 steps, seed 0. We compared supervised KD on teacher samples ($\lambda =0$), a 50/50 mix, and fully on-policy distillation with forward and reverse KL, and computed the student's off-track rate exactly from its transition matrix.

All runs start at a 22.5% off-track rate (the teacher's is 1.09%). The mechanism is visible in the recovery column. Only 1.1% of states in teacher-sampled sequences are off-track, so supervised KD almost never trains the recovery behaviour: after 300 steps the student's recovery probability had moved from 0.30 to only 0.37. On-policy training visits off-track states as often as the student actually falls into them, which early on is often, and raised recovery to 0.70. All methods improved the on-track transitions similarly (0.965 to 0.976). Reverse KL was the fastest, consistent with GKD's finding. The fully on-policy student was also better on the supervised metric, forward KL measured on teacher sequences (0.283 versus 0.362), because fixing recovery helps everywhere downstream.

The algorithm, step by step#

init student from an SFT or off-policy-distilled checkpoint
for step in range(total_steps):
    prompts = sample_prompts(batch_size)
    if random() < lam:                                           # student data fraction
        ys = student.generate(prompts, temperature=1.0, max_len=L)   # on-policy samples, no grad
    else:
        ys = fixed_dataset_or_teacher_outputs(prompts)
    with no_grad:
        t_logits = teacher.forward(prompts + ys)                     # one prefill pass, no generation
    s_logits = student.forward(prompts + ys)                         # with grad
    loss = mean over response tokens of D(teacher_dist, student_dist)   # reverse KL or JSD(beta)
    loss.backward(); optimizer.step(); optimizer.zero_grad()
    periodically: eval task metrics, response length, entropy, repetition

Data is just prompts. The student writes the responses and the teacher grades them, so no labels or reference responses are needed for the on-policy part, and you can use any prompt set that covers the target distribution: the prompts from your SFT data, from RL training, or unlabelled user queries. Tens of thousands of prompts are typically enough for a narrow task; reasoning distillation at the scale of Qwen3 uses far more. The student should start somewhere reasonable. On-policy training from a base model wastes most early samples on states nothing like the target, so the standard recipe, and the one Qwen3 describes, is off-policy distillation or SFT first, then on-policy distillation.

Hyperparameters that matter#

How to tune: fix the divergence and temperature, then sweep the learning rate on a short run, watching the per-token divergence on held-out student samples. It should fall smoothly. If it falls but task accuracy does not move, the student may be matching the teacher in easy positions and not the decisive ones; try reverse KL if you were on forward, or raise the share of prompts from the hardest part of your distribution.

Compute and memory cost#

Each step has three parts. Student generation, which costs the same as in RL and usually dominates wall-clock time. A teacher prefill over the generated tokens, about $2P_T$ FLOPs per token, with no autoregressive decoding. And a student forward and backward pass, about $6P_S$ FLOPs per token. For a 32B teacher and an 8B student, teacher scoring is $64$ against $48$ for student training per token: comparable to the update, small next to generation. Memory: the teacher must be available for inference (on the same GPUs or a separate pool), the student needs its full training state, and computing full-vocabulary divergences means holding $[B,L,V]$ logits for both models, which should be chunked along the sequence. Compared with RL, there is no reward model or value network, and the dense signal means far fewer steps for the same gain, which is what the Qwen3 GPU-hour comparison measured. Compared with offline distillation, it costs one round of student generation per step, which is the price of training on the right distribution.

Failure modes and how to detect them#

Exploiting the teacher

The teacher is now the reward, and rewards can be gamed. Language models assign high probability to repetition and to generic, safe continuations, so a student optimising reverse KL can drift toward repetitive or bland text that the teacher scores well token by token. Watch n-gram repetition rate, response length and a task metric that does not come from the teacher. If the divergence keeps falling while task accuracy stalls or falls, this is the likely cause.

Diversity collapse

Mode-seeking divergences shrink the student's entropy. Some shrinkage is the goal; too much hurts pass@k and makes subsequent RL harder because the student stops exploring. Track mean token entropy and pass@k alongside pass@1. Qwen3 reported pass@64 improving after on-policy distillation, so a falling pass@k is a warning sign, not an inevitable cost.

Teacher out of its depth

The teacher must give sensible distributions on the student's states, including bad ones. A teacher that has never seen the student's chat template, tool format or reasoning style gives noisy guidance. Check that the teacher's own log-likelihood on student samples is in a sensible range, and use a teacher from the same family and with the same template where possible.

Tokenizer mismatch

Per-token divergences require a shared vocabulary. Across tokenizer families, fall back to sequence-level methods or a dedicated alignment scheme.

Generation-bound training

If GPUs sit idle waiting for student samples, the run is generation-bound, as RL usually is. Use an inference engine for sampling, generate the next batch while training on the current one (slightly off-policy, usually harmless), and cap generation length.

When to use it, and what to use instead#

Use on-policy distillation when you have a stronger teacher with the same tokenizer, your student has been warm-started with SFT or offline distillation, and its remaining errors are the compounding kind: degraded long generations, reasoning that derails, failure to recover from its own mistakes. It is especially strong for small reasoning models, where Qwen3's comparison suggests it beats RL at a tenth of the cost. Do not use it as a first stage from a base model, without a teacher that is clearly better on your target distribution, or when the target behaviour is defined by a verifier rather than by imitation; then RL is the right tool, possibly with a distillation term as a regulariser.

Minimal implementation in PyTorch#

A compact training step with Hugging Face models, using the full-vocabulary reverse KL on student samples. The TRL library includes a GKD trainer with the same structure; in POFE on-policy distillation is a recipe that pairs a student run with a teacher scoring pool.

import torch, torch.nn.functional as F

def per_token_div(s_logits, t_logits, mask, kind="reverse", beta=0.9):
    s_logp, t_logp = F.log_softmax(s_logits, -1), F.log_softmax(t_logits, -1)
    if kind == "forward":
        d = (t_logp.exp() * (t_logp - s_logp)).sum(-1)
    elif kind == "reverse":
        d = (s_logp.exp() * (s_logp - t_logp)).sum(-1)
    else:  # generalised JSD(beta), mixture m = beta*p_T + (1-beta)*q
        m_logp = torch.logsumexp(torch.stack([t_logp + torch.log(torch.tensor(beta)),
                                              s_logp + torch.log(torch.tensor(1 - beta))]), 0)
        d = beta * (t_logp.exp() * (t_logp - m_logp)).sum(-1) + (1 - beta) * (s_logp.exp() * (s_logp - m_logp)).sum(-1)
    return (d * mask).sum() / mask.sum()

def gkd_step(student, teacher, tok, prompts, opt, lam=1.0, max_new=512):
    enc = tok(prompts, return_tensors="pt", padding=True, padding_side="left").to(student.device)
    if torch.rand(()) < lam:
        with torch.no_grad():
            seq = student.generate(**enc, do_sample=True, temperature=1.0, max_new_tokens=max_new)
    else:
        seq = fixed_completions(prompts)                       # dataset or teacher outputs
    attn = (seq != tok.pad_token_id).long()
    resp_mask = torch.zeros_like(seq, dtype=torch.float)
    resp_mask[:, enc["input_ids"].shape[1]:] = 1.0             # supervise only generated tokens
    resp_mask = (resp_mask * attn)[:, 1:]
    with torch.no_grad():
        t_logits = teacher(input_ids=seq, attention_mask=attn).logits[:, :-1].float()
    s_logits = student(input_ids=seq, attention_mask=attn).logits[:, :-1].float()
    loss = per_token_div(s_logits, t_logits, resp_mask, kind="reverse")
    loss.backward()
    torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
    opt.step(); opt.zero_grad(set_to_none=True)
    return loss.item()

Note that gradients flow only through s_logits; the sampled sequence is a constant, as in GKD. For long responses, compute the divergence in chunks along the sequence to avoid materialising two $[B,L,V]$ fp32 tensors at once.

Further reading#

Agarwal et al. (2023) is the GKD paper, including the divergence and student-data-fraction ablations and the combination with RL. Gu et al. (2023) is MiniLLM, the policy-gradient formulation of reverse-KL distillation. Ross, Gordon and Bagnell (2011) is the imitation-learning result behind the whole idea. Ko et al. (2024, DistiLLM) proposes a skew KL and an adaptive off-policy schedule to cut the cost of student generation. The Qwen3 technical report (Yang et al., 2025) documents an off-policy then on-policy distillation pipeline at scale, with the RL comparison quoted above. For the offline methods this builds on, see knowledge distillation.

Frequently asked questions#

References#

1. Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, Olivier Bachem (2023). On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes. arxiv.org/abs/2306.13649
2. Yuxian Gu, Li Dong, Furu Wei, Minlie Huang (2023). MiniLLM: On-Policy Distillation of Large Language Models. arxiv.org/abs/2306.08543
3. Stephane Ross, Geoffrey J. Gordon, J. Andrew Bagnell (2011). A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning. arxiv.org/abs/1011.0686
4. Jongwoo Ko, Sungnyun Kim, Tianyi Chen, Se-Young Yun (2024). DistiLLM: Towards Streamlined Distillation for Large Language Models. arxiv.org/abs/2402.03898
5. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, et al. (2025). Qwen3 Technical Report. arxiv.org/abs/2505.09388
6. Geoffrey Hinton, Oriol Vinyals, Jeff Dean (2015). Distilling the Knowledge in a Neural Network. arxiv.org/abs/1503.02531

Related reading