Sequence-level distillation

The idea

A language model produces a sequence one token at a time. Ordinary (word-level) distillation trains the student to match the teacher's next-token distribution at every position of a human-written sequence. Kim and Rush argued that what should be copied is the teacher's distribution over whole sequences, and used its mode as a stand-in: run the teacher's beam search on the training inputs and fine-tune the student on those outputs with plain cross-entropy. On-policy distillation (GKD) goes one step further: the student generates the sequences and the teacher grades them.

How it works

Sequence-level KD (Kim & Rush 2016):

1. For each training source x (say, a sentence to translate), decode the teacher with beam search (beam size about 5) and keep the top hypothesis ŷ.
2. Replace the reference target y with ŷ.
3. Train the student by ordinary cross-entropy on (x, ŷ): loss = -Σ_t log p_s(ŷ_t | ŷ_<t, x).
4. Variants: combine it with word-level KD, Σ_t KL(p_t(· | prefix, x) || p_s(· | prefix, x)), or train on the beam hypothesis closest to y instead of the top one (sequence-level interpolation).

On-policy distillation (GKD, Agarwal et al. 2024):

1. With probability λ, sample ŷ ~ p_s(· | x) from the current student; otherwise take a fixed sequence (from the dataset or the teacher).
2. Run the teacher on (x, ŷ) to get its next-token distribution at every position, all in one forward pass.
3. Minimize a token-level divergence between teacher and student along ŷ: forward KL, reverse KL, or generalized Jensen-Shannon. No gradient flows through the sampling.
4. Repeat. The student is corrected on the states it actually reaches at inference time, which removes the train/test mismatch known as exposure bias.

The code, line by line

There is no sequence model here, so lab/train.py:KD is the per-example analogue:

• kind="kl" is word-level KD at a single position: T² · KL(p_t^T || p_s^T), with gradient T · (ps - pt) / N.
• kind="hard" turns the teacher's distribution into a one-hot of its argmax, onehot(p.argmax(1), ...), and trains with cross-entropy. That is sequence-level KD with sequences of length one: the teacher's mode, used as a label.
• kind="rkl" is reverse KL, KL(p_s || p_t), a mode-seeking divergence that GKD found useful on some generation tasks.

When it breaks

• Beam outputs are lower-entropy than human text. The student inherits that blandness and repetitiveness, and loses calibration.
• Teacher errors become hard labels, with no signal about how confident the teacher was.
• On-policy sampling is expensive: every step needs student generation and a teacher forward pass.
• A weak student's early samples are garbage, so the teacher's feedback on them teaches little. Mixing in fixed sequences (λ < 1) helps.

Going further

To run this you need an autoregressive teacher and student that share a tokenizer, a decoding loop (beam search or sampling), and a corpus of source texts. A small pair such as a 1B and a 7B open model, with the transformers library, is enough. This repo has none of these: its models classify a single 12-D input. The closest analogue is the loss_zoo experiment (lab/experiments/distill.py), which compares hard pseudo-labels, soft KL and reverse KL on the toy world. The per-position KD loss can be applied unchanged at each token of a teacher- or student-generated sequence.

Source code

This method doesn't run here. Below is the closest code the lab does run, read from the repository when this page was built.

class KD:
    """Match the teacher's outputs.

    kind:
      kl    -- Hinton KD: T^2 * KL(p_t^T || p_s^T)            (mode-covering)
      rkl   -- reverse KL: T^2 * KL(p_s^T || p_t^T)           (mode-seeking)
      mse   -- squared error on raw logits (T ignored)
      hard  -- CE on the teacher's argmax (pseudo-labels; no dark knowledge)
      topk  -- teacher only reveals its top-k probabilities (an API that returns
               top-k logprobs); the rest of the mass is spread uniformly
    group (C_t x C_s matrix): sum teacher probabilities into the student's
      coarser label space, e.g. 20 classes -> {bird, other} for a specialist.
    """

    def __init__(self, T: float = 4.0, weight: float = 1.0, kind: str = "kl", k: int = 3,
                 group: np.ndarray | None = None):
        self.T, self.w, self.kind, self.k, self.G = T, weight, kind, k, group

    def _teacher_probs(self, tl):
        T = 1.0 if self.kind == "hard" else self.T
        if self.kind == "topk":
            # An API reveals the top-k probabilities at T=1 and nothing else.
            # Spread the unknown remainder uniformly, *then* apply temperature.
            p = softmax(tl)
            thr = -np.sort(-p, 1)[:, self.k - 1:self.k]
            keep = p >= thr
            rest = (1 - (p * keep).sum(1, keepdims=True)) / (~keep).sum(1, keepdims=True).clip(1)
            p = np.where(keep, p, np.maximum(rest, 1e-12))
            if T != 1:
                p = p ** (1 / T)
                p = p / p.sum(1, keepdims=True)
        else:
            p = softmax(tl, T)
        if self.G is not None:
            p = p @ self.G
        if self.kind == "hard":
            p = onehot(p.argmax(1), p.shape[1])
        return p.astype(F)

    def __call__(self, z, hid, b: Batch):
        N, T = len(z), self.T
        if self.kind == "mse":
            t = b.t_logits if self.G is None else np.log(softmax(b.t_logits) @ self.G + 1e-8)
            r = z - t
            return self.w * (r ** 2).mean(), self.w * 2 * r / r.size, {}
        pt = self._teacher_probs(b.t_logits)
        if self.kind == "hard":
            loss = -(pt * log_softmax(z)).sum() / N
            return self.w * loss, self.w * (softmax(z) - pt) / N, {}
        ls = log_softmax(z, T)
        ps = np.exp(ls)
        # exact teacher log-probs when we can (log(p + eps) distorts tiny probabilities)
        lt = log_softmax(b.t_logits, T) if self.kind == "kl" and self.G is None \
            else np.log(pt + 1e-12)
        if self.kind == "rkl":
            per = (ps * (ls - lt)).sum(1, keepdims=True)
            loss = T * T * per.sum() / N
            dz = T * ps * (ls - lt - per) / N
        else:  # kl / topk
            loss = T * T * (pt * (lt - ls)).sum() / N
            dz = T * (ps - pt) / N
        return self.w * loss, self.w * dz, {}

    params = staticmethod(lambda: [])

Simulation figures

No simulation in the lab covers this method yet.

Papers

• Kim & Rush. Sequence-Level Knowledge Distillation (2016)
• Agarwal et al. On-Policy Distillation of Language Models (GKD) (2024)