mkurman Overview
Knowledge distillation transfers knowledge from a larger teacher model to a smaller student model. Combined with pruning and quantization, it forms the critical middle step in the P-KD-Q compression pipeline.
Installation
uv pip install kd-lib
Logit Distillation
from kd_lib import distill
import torch.nn.functional as F
def kd_loss(student_logits, teacher_logits, labels, temperature=4.0, alpha=0.5):
soft_targets = F.softmax(teacher_logits / temperature, dim=-1)
soft_prob = F.log_softmax(student_logits / temperature, dim=-1)
kd = F.kl_div(soft_prob, soft_targets, reduction="batchmean") * (temperature ** 2)
ce = F.cross_entropy(student_logits, labels)
return alpha * kd + (1 - alpha) * ce
Distillation Methods
| Method | What it transfers | Best For |
|---|---|---|
| Logit distillation | Output probability distribution | Classification, generation |
| Feature distillation | Intermediate hidden states | Transformer layers |
| Relation distillation | Relationships between representations | Structured outputs |
| Self-distillation | Model teaches itself | No teacher needed |
| Online distillation | Teacher & student train jointly | Both models improve |