Self-Distillation, briefly

Post-Training, KL-Divergence, Self-Distillation, Credit Assignment.

By Cristian Gutiérrez

From "On-Policy Self-Distillation for Large Language Models", Zhao (ICML 2026)

Overview of SFT, RLVR, and off-policy and on-policy knowledge distillation

For discrete probability distributions:

$$ D_{KL}(P\mid\mid Q)=\sum_i P(i) \log \frac{P(i)}{Q(i)} $$

For continuous probability distributions:

$$ D_{KL}(P\mid \mid Q)=\int P(x) \log \frac{P(x)}{Q(x)} \, dx $$

Notice that it is not a distance per sé:

$$ D_{KL}(P || Q) \neq D_{KL}(Q||P) $$

In the case of On-Policy Distillation, our probability distribution consists of the next tokens probs over all the next token ids over the vocabulary.

$$ D_{KL}(P\mid\mid Q)=\sum_{i=0}^{248,320} P(i) \log \frac{P(i)}{Q(i)} \qquad \text{where } i \text{ is a next-token id} $$

Notice how each token in our sequence \([1, N]\) will have a prefix, hence, we will start with the next token 1, then the next token 2 (which is conditioned by the previous prefill)…

In self-distillation we do on-policy distillation, but using the model itself with extra context (or hints) as target.

$$ \mathcal{L}(\theta) = D_{KL}(\pi_\theta \mid\mid \pi_{\text{teacher}}) $$

Self-Distillation can be modeled as the classical RL of a policy gradient with token-level advantages.

Make \(\pi_\theta(\cdot \mid x)\) match the improved (conditioned) teacher \(\pi_\theta(\cdot \mid x, f)\).

$$ \text{Minimize} \quad \mathcal{L}_{\text{SDPO}}(\theta) := \frac{1}{|y|} \sum_{i=1}^{|y|} \text{KL}\big(\underbrace{\pi_\theta(\cdot \mid x, y_{<i})}_{\text{student}} \,\|\, \underbrace{\pi_\theta(\cdot \mid x, f, y_{<i})}_{\text{teacher}}\big) $$
$$ \text{Maximize} \quad A_i(x, y, f) := \log \frac{\overbrace{\pi_\theta(y_i \mid x, f, y_{<i})}^{\text{teacher}}}{\underbrace{\pi_\theta(y_i \mid x, y_{<i})}_{\text{student}}} $$

Suppose we have an input prompt asking our student to write a Python function to return all numbers between 1 and n.

The student incorrectly returns \([1, n]\) instead of \([1, n)\).

Then we can condition the teacher, walk the same trajectory and minimize the KL divergence at token-level.

Python range example with feedback and credit assignment by the self-teacher

Notice how with distillation, instead of penalizing the whole trajectory (sparse reward problem in RLVR) we would be just penalizing at a more fine-grained level:

Fine-grained token-level credit assignment compared with trajectory-level rewards

And not only that, but we would be penalizing the log-probs of other possible tokens:

Distillation feedback across the probabilities of alternative next tokens

References