Knowledge distillation

Temperature T softens the teacher's output probability distribution by dividing logits before the softmax: p_i = exp(z_i/T) / sum_j(exp(z_j/T)). Higher T produces softer distributions that reveal more inter-class similarity (dark knowledge), giving the student richer gradient signal beyond what hard labels provide. At T=1 the distribution is the standard softmax; as T→∞ it approaches uniform.

import torch
import torch.nn.functional as F

def soft_targets(logits, T=4.0):
    return F.softmax(logits / T, dim=-1)

The total loss is L = (1-α)·CE(student_logits, hard_labels) + α·T²·KL(soft_teacher || soft_student), where both soft distributions are computed at temperature T. The T² factor compensates for the gradient magnitude reduction caused by dividing logits by T. α is a hyperparameter balancing task loss vs distillation loss.

def kd_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.9):
    soft_s = F.log_softmax(student_logits / T, dim=-1)
    soft_t = F.softmax(teacher_logits / T, dim=-1)
    distill = F.kl_div(soft_s, soft_t, reduction='batchmean') * T * T
    hard = F.cross_entropy(student_logits, labels)
    return alpha * distill + (1 - alpha) * hard

Dark knowledge refers to the information encoded in the teacher's soft output probabilities beyond the argmax class—specifically the relative probabilities assigned to wrong classes. For example, a model distinguishing 'dog' from 'cat' and 'automobile' differently reveals structural knowledge about class similarity. This inter-class relationship information is invisible in one-hot labels but transferable through soft targets.

Response-based KD transfers the teacher's final output layer (logits/soft probabilities). Feature-based KD matches intermediate layer activations between teacher and student (e.g., FitNets). Relation-based KD captures structural relationships between sample pairs or feature maps in the teacher (e.g., pairwise distances, correlation matrices) and forces the student to replicate these relational structures rather than individual activations.

FitNets address the difficulty of training very thin but deep student networks directly from teacher outputs. They add a 'hint layer' in the student and a regression layer to match an intermediate 'guided layer' in the teacher via MSE, pre-training the student bottom-half before standard KD fine-tuning. This staged approach guides gradient flow into early layers and avoids vanishing gradients that plague direct end-to-end distillation of deep thin nets.

# Hint loss: match student hint to teacher guided layer
hint_loss = F.mse_loss(regressor(student_hint), teacher_guided.detach())

When the student is far smaller than the teacher, soft targets can be too complex for the student to fit, degrading performance below what direct training achieves. Strategies include: using an intermediate-capacity teaching assistant (Mirzadeh et al., 2020) to bridge the gap progressively; applying curriculum distillation that starts with easy examples; or reducing temperature to limit soft-target complexity early in training.

AT transfers spatial attention maps derived from intermediate feature maps of the teacher to the student. Attention maps are computed by summing squared (or p-norm) activations across channels, then normalized. An L2 loss between teacher and student attention maps is added to the training objective. This forces the student to 'look at' the same image regions as the teacher, transferring visual reasoning rather than raw activations.

def attention_map(feat, p=2):
    return F.normalize(feat.pow(p).mean(dim=1).view(feat.size(0), -1))

at_loss = (attention_map(student_feat) - attention_map(teacher_feat.detach())).pow(2).mean()

Instance-level feature matching (e.g., FitNets) penalizes the distance between individual sample embeddings of teacher and student. RKD instead penalizes differences in pairwise distances (distance-wise) and triplet angles (angle-wise) between groups of samples in embedding space. This relational transfer is invariant to absolute embedding positions and better preserves the global geometry of the teacher's representation.

# Distance-wise RKD
def dist_loss(s_emb, t_emb):
    t_d = torch.pdist(F.normalize(t_emb, dim=-1))
    s_d = torch.pdist(F.normalize(s_emb, dim=-1))
    return F.smooth_l1_loss(s_d, t_d / t_d.mean())

BAN (Furlanello et al. 2018) trains a student with identical architecture to the teacher using KD, then uses that student as the new teacher for the next generation. Ensembling multiple generations further boosts performance. The student outperforms the teacher because soft targets act as stronger regularizers than hard labels—they encode uncertainty that prevents overconfident memorization, effectively functioning like data augmentation in label space.

DML trains a cohort of student networks simultaneously without a pre-trained teacher: each network learns from the soft targets of all its peers via KL divergence, in addition to task loss. Unlike traditional offline KD where a fixed large teacher teaches a small student, DML is online and symmetric—all networks improve together, and even small networks can outperform those trained independently, because the ensemble of peers acts as a collective teacher that evolves during training.

DistilBERT reduces BERT-base from 12 to 6 Transformer layers (every other layer initialized from BERT's), achieving ~40% size reduction with ~97% performance retention on GLUE. It uses three losses: (1) soft cross-entropy on MLM logits between teacher and student, (2) cosine embedding similarity on hidden state vectors, and (3) masked language model loss on hard labels. No next-sentence prediction head is used since it adds little benefit.

DFKD transfers knowledge from a trained teacher to a student without access to the original training data, which is important for privacy-sensitive or proprietary datasets. A common approach trains a GAN-like generator to produce synthetic 'pseudo-data' that maximally activates or confuses the teacher. The student trains on these generated samples with teacher soft labels. Accuracy degrades versus data-driven KD but is often acceptable for deployment constraints.

Noisy Student (Xie et al. 2020) is a form of self-training / iterative KD on unlabeled data. A teacher generates pseudo-labels on unlabeled images; a student trains on labeled + pseudo-labeled data with strong noise (RandAugment, dropout, stochastic depth). The student then becomes the new teacher for the next iteration. Unlike standard KD, the student is at least as large as the teacher and noise forces robustness, enabling students to surpass teachers.

Label smoothing replaces one-hot labels with (1-ε)·one_hot + ε/K (uniform), which is equivalent to a KD-style mixture where the 'teacher' is the uniform distribution. This is a special case of KD where the teacher has no structure—it provides entropy but no dark knowledge about inter-class relationships. True KD with a trained teacher provides structured soft targets that carry richer similarity information beyond this uniform smoothing.

Detection outputs are spatially structured (feature maps, anchors, region proposals) rather than a single class vector, making naive logit-level KD inadequate. Mimicking matches intermediate feature maps between teacher and student detectors via L2 loss on full feature grids, focusing on foreground regions (where objects are detected) to avoid the large negative-region imbalance dominating the loss. Later works like DeFeat use separate distillation for foreground vs background.

Progressive KD chains multiple intermediate teachers of decreasing size between the original large teacher and the final small student, distilling step-by-step. Each step has a smaller capacity gap, avoiding the difficulty of jumping from a very large teacher to a tiny student directly. It is preferred when the final student is extremely compact (e.g., mobile-scale) and the capacity gap is large enough that direct distillation degrades below independent training.

Offline distillation pre-trains the teacher to convergence, then trains the student using fixed teacher outputs. The teacher does not update during student training. Examples: Hinton KD, FitNets. Online distillation trains teacher and student simultaneously, allowing the teacher to co-evolve. Examples: Deep Mutual Learning (DML), ON-THE-FLY Native Ensemble. Online methods avoid separate teacher training cost but can be less stable since the teacher is itself imperfect during training.

CRD frames feature distillation as a contrastive learning problem: for each sample, the student embedding should be similar to the teacher embedding of the same sample (positive pair) and dissimilar from teacher embeddings of other samples (negative pairs). The InfoNCE loss is used. This maximizes a lower bound on mutual information between student and teacher representations, enabling cross-architecture transfer without requiring dimensional alignment layers.

KL(p_T || p_S) is the forward KL (mean-seeking): minimizing it forces the student to cover all modes of the teacher distribution, producing averaged but complete approximations. KL(p_S || p_T) is reverse KL (mode-seeking): it causes the student to concentrate on one mode, ignoring others. In KD we want the student to reflect all teacher confidences, so forward KL is standard. Using reverse KL would make the student ignore low-probability teacher signals (dark knowledge).

# PyTorch F.kl_div computes KL(target || input) when using reduction='batchmean'
# kl_div(log_student, teacher) = KL(teacher || student)  <- correct for KD
loss = F.kl_div(F.log_softmax(s_logits/T, -1), F.softmax(t_logits/T, -1), reduction='batchmean')

Self-distillation uses the network's own deeper (later) predictions as soft teachers for its shallower intermediate classifiers, with no external teacher. BYOT attaches auxiliary classifiers to intermediate layers; the deepest classifier's soft outputs supervise the intermediate ones via KD loss. This creates an internal gradient shortcut that improves early-layer representations and the final prediction simultaneously, acting as an implicit ensemble regularizer.

In Once-for-All (OFA) and similar NAS frameworks, a large supernet is trained once, then sub-networks of varying architecture sizes are distilled from this supernet teacher as student networks without retraining from scratch. Progressive shrinking trains sub-networks at decreasing sizes, using the supernet as soft-label teacher at each size step. This decouples training cost from the number of deployment-time architecture configurations.

A large LLM generates soft output distributions (token-level probabilities) or reasoning traces on a task dataset. A smaller student LM is trained to match these soft outputs (sequence-level KL divergence), chain-of-thought rationales (sequence-level supervision), or both. Examples: Alpaca distills text-davinci-003 (GPT-3.5-class) instruction-following data; WizardLM distills GPT-4 via evolved instructions. LLM-Distillation papers match next-token distributions. Task-specific distillation dramatically reduces inference cost while retaining most task accuracy.

Word-level KD applies KD token-by-token, minimizing KL divergence between teacher and student distributions at each decoding step conditioned on ground-truth prefix (teacher-forcing). Sequence-level KD generates full output sequences from the teacher (beam search) and uses them as hard targets for the student. Sequence-level avoids exposure bias since the student trains on likely teacher outputs, and is cheaper because teacher probabilities per-step are not required at student training time.

GNN distillation must transfer knowledge across irregular, variable-size graph structures rather than fixed-size tensors. Challenges include: node/edge feature alignment when teacher and student have different hidden dimensions; graphs may have different receptive field depths; and there is no canonical spatial ordering for feature map matching. Methods like RKD-on-graphs transfer pairwise node embedding similarities, while local structure preserving (LSP) distillation transfers local topology information.

PKT (Passalis & Tefas 2018) models the teacher's feature space as a probability distribution over pairwise sample similarities using a Gaussian kernel. The student is trained to match this probability distribution (using KL divergence between teacher and student kernel similarity matrices) rather than matching raw feature vectors. This is more robust to dimensional differences and preserves the relational structure of the embedding space rather than absolute feature magnitudes.

In FedDF / Ensemble Distillation for federated learning, each client trains a local model on private data. Instead of averaging weights (FedAvg), the server collects only soft predictions from all client models on a public unlabeled dataset. A global model is then distilled from this ensemble of soft predictions. This preserves privacy (only predictions shared, not weights), allows heterogeneous client architectures, and is robust to systems heterogeneity.

Activation boundary distillation targets the decision boundaries defined by ReLU neurons. Rather than matching activation magnitudes, it matches the binary activation pattern (which neurons fire/don't fire) of the teacher. The loss encourages the student to replicate the same sign pattern of pre-ReLU activations as the teacher. This captures structural/combinatorial information about how the teacher partitions feature space, which is more informative than just magnitude similarity.

# Connector projects student features to teacher dimension
# Loss encourages matching ReLU activation sign patterns
loss = (1 + torch.exp(-connector(student_feat) * teacher_feat.sign())).log().mean()

Quantizing a model (reducing weights/activations to INT8 or lower) typically degrades accuracy. Quantization-aware distillation (QAD) trains a quantized student to match a full-precision teacher's soft outputs and/or intermediate features during quantization-aware training (QAT). The teacher's rich soft targets compensate for quantization noise, recovering more accuracy than QAT alone. This is critical for deploying quantized models without large accuracy gaps.

Dataset distillation compresses a large training dataset into a tiny set of synthetic images such that a model trained only on the synthetic images achieves performance close to training on the full dataset. The synthetic images are meta-learned via bi-level optimization (inner loop: train model on synthetic data; outer loop: minimize validation loss on real data). This is distinct from model KD, which compresses the model; dataset distillation compresses the data, and the teacher's knowledge is baked into the synthetic images themselves.

Task-agnostic distillation compresses the general pre-trained model (e.g., distilling BERT during pre-training on MLM) so the student can be fine-tuned on any downstream task—examples: DistilBERT, TinyBERT stage 1. Task-specific distillation distills after fine-tuning on a target task, so the student captures both language understanding and task decision boundaries—examples: PKD-BERT, TinyBERT stage 2. Task-specific typically achieves higher accuracy on that task but produces a single-task student.

TinyBERT defines a layer mapping function from student to teacher layers and distills at three levels per mapped layer: (1) attention matrix distillation—MSE between teacher and student self-attention weight matrices (A^T and A^S), (2) hidden state distillation—MSE on hidden representations (after linear projection if dimensions differ), and (3) prediction-layer distillation using soft logit KL divergence. This granular layer-wise supervision yields state-of-the-art compression.

# Attention matrix distillation for one layer
attn_loss = F.mse_loss(student_attn, teacher_attn.detach())
# Hidden state distillation (with projection W_h)
hidden_loss = F.mse_loss(student_hidden @ W_h, teacher_hidden.detach())

When teacher and student intermediate feature maps have incompatible spatial or channel dimensions, a small learned linear layer (or 1x1 conv) called a hint regressor (or adapter) projects the student's feature map into the teacher's dimension space before computing the distillation loss. This allows MSE or cosine loss to be applied despite architectural mismatches. The regressor is typically discarded after distillation since it is not part of the final deployed student.

# 1x1 conv adapter to match channels
adapter = nn.Conv2d(student_channels, teacher_channels, kernel_size=1)
hint_loss = F.mse_loss(adapter(student_feat), teacher_feat.detach())

Speculative decoding uses a small draft model (often a distilled version of the large target model) to auto-regressively generate multiple candidate tokens cheaply. The large model then verifies all candidates in a single forward pass in parallel and accepts/rejects them. KD ensures the draft model's distribution closely matches the target model's, maximizing acceptance rate. The speedup is approximately proportional to accepted tokens per large-model forward pass, typically 2-4x.

Beyond top-1/top-5 accuracy: (1) KL divergence between teacher and student output distributions on a held-out set measures knowledge transfer fidelity; (2) calibration metrics (ECE—expected calibration error) check if the student's confidence is as well-calibrated as the teacher's; (3) throughput/latency improvements on target hardware measure compression payoff; (4) out-of-distribution robustness gap between teacher and student reveals if dark knowledge generalizes; (5) intermediate feature alignment loss tracks feature-level distillation quality.

Ensemble distillation trains a single student to match the averaged soft predictions of multiple independently trained teacher models. The ensemble soft targets are smoother and better calibrated than any single teacher, capturing more stable inter-class relationships (averaged dark knowledge). The student achieves accuracy closer to the full ensemble than to any single member, while deploying at single-model cost. This is used in BAN, Snapshot Distillation, and production systems like Google's Knowledge Distillation for TTS.

# Average soft predictions from N teachers at temperature T
ensemble_soft = torch.stack([F.softmax(t/T, -1) for t in teacher_logits_list]).mean(0)
loss = F.kl_div(F.log_softmax(student_logits/T, -1), ensemble_soft, reduction='batchmean') * T**2