Notation used throughout:
x= prompt (document + instruction),y = y_1..y_T= target summary tokensz^T_t ∈ R^V= teacher logits at positiont;z^S_t= student logitsp_t = softmax(z^T_t / τ)= teacher distribution at temperature τq_t = softmax(z^S_t / τ)= student distributionV= vocab size (151,936),τ= distillation temperatureM= set of unmasked (supervised) positions;|M|= its size
The total objective is a weighted sum:
L = α · L_CE + (1 − α) · τ² · L_KD + λ_h · L_hidden + λ_a · L_attn
Every term is explained below, including why τ² is there.
L_CE = − (1/|M|) Σ_{t∈M} log q̂_t[y_t] where q̂ = softmax(z^S_t) (τ=1 always)
What it is: negative log-likelihood of the target token. Equivalently, KL(one_hot(y_t) ‖ q̂_t).
Intuition: "make the correct token probable." Only one entry of the 151,936-vector receives gradient signal directly; all others are pushed down through the softmax normalizer.
Units: nats. L_CE = 2.3 means perplexity e^2.3 ≈ 10 — the model is effectively choosing
among ~10 equally-likely tokens.
Typical values for summarization:
- Start (pretrained student, before tuning): 1.8–2.6
- After good distillation: 0.8–1.4
- Below 0.4: you are memorizing, or you have the off-by-one bug from
docs/04 - Exactly 0.0 or NaN: bug (see §9)
In code: F.cross_entropy(logits.view(-1,V), labels.view(-1), ignore_index=-100).
Where it comes from in distillation: the "target" can be a human reference (rare here) or the teacher's generated text (sequence-level KD). In practice the teacher text is your ground truth.
L_FKD = (1/|M|) Σ_{t∈M} Σ_{v=1}^{V} p_t[v] · log( p_t[v] / q_t[v] )
Expanded: = Σ p log p − Σ p log q = −H(p) − Σ p log q. Since H(p) doesn't depend on the student,
minimizing forward KL ≡ minimizing the cross-entropy against soft targets. That is why you'll
see implementations use F.kl_div(log_q, p) and others use -(p * log_q).sum(); they differ by a
constant and have identical gradients.
Gradient (the thing to actually understand):
∂L_FKD / ∂z^S_t = (1/τ) · (q_t − p_t)
Beautifully simple: push the student's probability vector toward the teacher's, elementwise. The
magnitude at token v is q[v] − p[v].
Behaviour — mode-covering / mass-covering. Look at the term p·log(p/q). Wherever the teacher
has mass (p > 0) but the student has none (q → 0), the term → +∞. So forward KL severely
punishes the student for missing any mode the teacher has. It does not punish the student for
putting mass where the teacher has none (if p=0, the whole term is 0 regardless of q).
Result for a small student that cannot represent the teacher's full distribution: it spreads its probability to cover everything, producing a flatter, hedging distribution. In generation this shows up as bland, generic, sometimes incoherent text ("mean-seeking blur").
When to use: default choice. Stable, cheap, well-conditioned. Use it for Stage 2.
Typical values: at τ=1 on in-domain data, a same-family pretrained student vs teacher starts
around 0.5–3.0 nats and should fall to 0.1–0.6. If it starts near ln(V) ≈ 11.9 your alignment is
broken.
L_RKD = (1/|M|) Σ_{t∈M} Σ_v q_t[v] · log( q_t[v] / p_t[v] )
Gradient: requires care because q appears in the expectation and inside the log:
∂L_RKD/∂z^S = q ⊙ ( log(q/p) − Σ_v q_v log(q_v/p_v) )
i.e. each logit is pushed by how much worse than average its log-ratio is.
Behaviour — mode-seeking / zero-forcing. Now the q log(q/p) term blows up when the student
puts mass where the teacher has none. So the student is punished for hallucinating options the
teacher rejects, and is free to ignore teacher modes it can't fit. A small student under reverse
KL picks one mode of the teacher and commits to it.
Why this matters for summarization: the failure mode you care about is the student saying things the teacher wouldn't. Reverse KL directly penalizes exactly that. MiniLLM (2023) showed reverse KL beats forward KL for instruction-following distillation for precisely this reason.
Costs:
- Lower diversity. The student collapses onto one style. For a summarizer, usually fine — even desirable.
- Higher variance / less stable. Especially with truncated top-K, since
pin the denominator can be tiny or unavailable for tokens outside top-K. You must floorp(p.clamp_min(1e-8)) or use the residual bucket, or you get NaNs. This is the single most common reverse-KL bug. - Strictly, on-policy reverse KL requires sampling from
q(policy gradient, as in MiniLLM). The cheap version here — computing reverse KL on teacher-generated sequences — is an approximation. It still helps, but the full benefit needs on-policy sampling (§6).
Verdict: worth trying as a Stage-3 objective. Don't start here.
Jensen–Shannon divergence, with mixture m = ½(p+q):
JSD(p,q) = ½ KL(p‖m) + ½ KL(q‖m)
Bounded (≤ ln 2), symmetric, and finite even when supports don't overlap. That boundedness is
the practical benefit: no exploding terms, no NaNs.
Generalized / skewed JSD (from GKD), with β ∈ [0,1]:
m_β = β·p + (1−β)·q
JSD_β = β·KL(p ‖ m_β) + (1−β)·KL(q ‖ m_β)
β → 1recovers forward KL behaviour (mode-covering)β → 0recovers reverse KL behaviour (mode-seeking)β = 0.5is standard JSD
This is the most useful single knob in the whole loss design. It gives you a continuous dial
between the two regimes with none of the numerical fragility. Recommended sweep: β ∈ {0.1, 0.5, 0.9}. Empirically β ≈ 0.1–0.3 (reverse-leaning) is often best for instruction distillation.
Implementation note: compute in log-space.
log m = logaddexp(log β + log p, log(1−β) + log q).
| Loss | Formula | Notes |
|---|---|---|
| MSE on logits | ‖z^S − z^T‖² / V |
No softmax, no temperature. Preserves all logit info including the shift-invariant part (which is meaningless) — so it wastes capacity. Occasionally more stable at low precision. Rarely better. |
| TVD (total variation) | `½ Σ | p − q |
| Wasserstein / ULD | OT cost between sorted prob vectors | The cross-tokenizer option. Vocabulary-agnostic because it only compares sorted distributions. |
| Top-K CE | CE restricted to teacher's top-K | Equivalent to forward KL under top-K renormalization. |
| Rank-based (list-wise) | e.g. RankNet over top-K | Matches the ordering rather than the values. Robust to teacher miscalibration. Niche. |
Independent of which divergence you use, there is the question of on what sequences you compute it. This matters more than the divergence choice.
| Data source | Name | Property |
|---|---|---|
| Fixed teacher-generated summaries | off-policy / SeqKD | Cheap, precomputable logits. Suffers exposure bias. |
| Student-generated, teacher-scored | on-policy / GKD | Student sees its own error distribution and learns to recover. 2–4x cost, teacher must be in memory. |
Mixture, λ fraction on-policy |
GKD(λ) | Best of both. λ=0.25–0.5 typical. |
Exposure bias, concretely: during training the student always conditions on a perfect prefix
written by the teacher. At inference it conditions on its own prefix, which contains its own
mistakes — a state distribution it never trained on. Errors compound. On-policy KD is the direct
fix: generate with the student (temperature≈1, no grad), score with the teacher, apply the
divergence.
GKD training step:
with torch.no_grad():
y = student.generate(x, do_sample=True, temperature=1.0, max_new_tokens=cfg.gkd_max_new)
t_logits = teacher(cat(x,y)).logits
s_logits = student(cat(x,y)).logits # WITH grad
loss = jsd_beta(s_logits, t_logits, beta=0.1)Note: the generation is no_grad, and gradient flows only through the re-forward. This is
not REINFORCE — GKD treats it as a supervised loss on self-generated data.
Hidden-state distillation
L_hidden = (1/|M|) Σ_{t∈M} Σ_{l∈map} 1 − cos( W_l · h^S_{l,t} , h^T_{map(l),t} )
W_l ∈ R^{d_T × d_S}is a learned projection (896 → 3584 for 0.5B→7B). Its parameters are trained with the student and discarded at the end — they are scaffolding.- Cosine over MSE: hidden-state magnitudes differ wildly between models and across layers, and RMSNorm makes absolute scale semantically irrelevant. Matching direction is the meaningful part. MSE forces the student to match a scale it has no reason to match, and the loss is then dominated by a few high-norm outlier dimensions.
- Layer mapping:
map(l) = round(l · L_T / L_S). Or the cheaper "last + middle" variant. - λ_h: start at 0.0. Add it only if logit KD alone underfits. Typical
0.1–1.0, and you must normalize it — hidden losses have a different natural scale than KL. Log both terms separately.
Edge case: models with residual-stream outlier dimensions (a handful of dims with 100x the typical magnitude — real in Qwen and Llama). Cosine similarity gets dominated by those dims. Mitigation: LayerNorm both sides before comparing, or drop the top-k outlier dims.
L_attn = (1/|M|) Σ_l Σ_h KL( A^T_{l,h} ‖ A^S_{l,h} )
Only valid when head counts correspond. Qwen2.5-7B has 28 heads, 0.5B has 14 — you'd need to pool
teacher heads pairwise, which is an arbitrary choice. MiniLM's trick avoids this: instead of
matching attention maps, match relations (softmax(QKᵀ/√d), softmax(VVᵀ/√d)) computed on a
relation-head-count of your choosing, so head counts need not match. That's the version to use
if you want attention transfer at all.
Verdict: skip both for Saransh v1. Logit KD dominates; feature losses add complexity and two more hyperparameters for a few points at best. Revisit if the 0.5B student underfits badly.
Soften with temperature and the gradient shrinks:
∂/∂z^S [ KL(p_τ ‖ q_τ) ] = (1/τ)(q_τ − p_τ)
And q_τ − p_τ itself shrinks roughly like 1/τ for small logit differences (Taylor-expand: for
z/τ small, softmax(z/τ) ≈ 1/V + (z − z̄)/(τV)). So the gradient scales as 1/τ². Multiplying the
KD term by τ² restores its magnitude, which means α keeps the same meaning when you change
τ. Without it, raising τ silently turns the KD term off.
Practical: always include τ². It costs nothing and decouples two hyperparameters.
| τ | Effect |
|---|---|
| 1.0 | Native distribution. Use when teacher and student are close in size, or with reverse KL. |
| 2.0 | Good default. Exposes tail structure without washing out the mode. |
| 3–4 | Strongly softened. Helps large compression ratios. Risk: the student learns to be under-confident, generation becomes rambly. |
| >5 | Distribution approaches uniform; you're training on noise |
Interaction with top-K: higher τ moves mass into the tail, so the mass captured by your top-64
drops. If you use τ=4 with K=32, you may be truncating away 10% of the target. Measure captured
mass at the τ you train with, not at τ=1.
Also: apply τ to both sides. Softening only the teacher is a common bug and makes the target systematically flatter than anything the student can produce, so it never converges.
L = α·L_CE + (1−α)·τ²·L_KD
| α | Regime |
|---|---|
| 1.0 | Pure SFT — no distillation |
| 0.7 | CE-dominant. Safe start; the KD term acts as a regularizer |
| 0.3–0.5 | Recommended. KD carries most of the signal |
| 0.1 | KD-dominant. Best final quality when teacher data is clean |
| 0.0 | Pure distillation. Works, but you lose the anchor to actual token targets; the student can drift if the teacher is miscalibrated |
Schedule it. Start α=0.7 (learn the format from hard targets), anneal to α=0.2 over the
first 30% of training (then refine on soft targets). losses.py supports alpha_schedule: linear.
- Always work in log-space.
F.log_softmax, neverlog(softmax(x)). F.kl_divexpectsinput= log-probs,target= probs (unlesslog_target=True). Getting this backwards gives a loss that is silently wrong and often negative. A negative KL is always a bug — KL ≥ 0 by Jensen. Assert it.reduction='batchmean', not'mean'.'mean'divides byB*T*V, making your loss ~150,000x too small and your KD term effectively zero. This is the most common KD implementation bug in the wild.- Clamp the denominator in reverse KL:
p.clamp_min(1e-8). - fp16 overflow: logits can exceed 65,504 in fp16 after scaling. Cast to fp32 before the
softmax:
logits.float(). With bf16 (range like fp32) this is less critical but still do it — bf16 has only 8 mantissa bits, and the log-sum-exp reduction over 151,936 terms accumulates real error. - Empty batches: if every label in a micro-batch is
-100,|M| = 0→ division by zero → NaN. Guard:denom = mask.sum().clamp_min(1). - Padding leaking into the softmax: if you didn't slice out padded vocab columns, mass goes to undecodable tokens.
torch.no_grad()around the teacher. Forgetting this doesn't error — it silently doubles memory and, if the teacher shares any module with the student, corrupts gradients.
loss/total loss/ce loss/kd loss/hidden
kd/temperature kd/alpha
diag/teacher_entropy # nats; if ~0 the teacher is greedy-degenerate
diag/student_entropy # if it collapses toward 0, mode collapse (reverse KL too strong)
diag/topk_mass_captured # should stay >0.98
diag/agreement_top1 # fraction of positions where argmax_s == argmax_t; the money metric
diag/grad_norm # pre-clip
diag/lr
diag/agreement_top1 is the metric to watch. It should climb from ~0.55 to 0.80–0.92. If total
loss falls but agreement doesn't move, you're optimizing something that isn't imitation.
Next: 08 — the training loop