How DistilBERT was distilled from BERT

DistilBERT remains the reference implementation for anyone asking how DistilBERT was distilled from BERT, and for good reason: the paper (Sanh et al., 2019) laid out a reproducible recipe that compresses BERT-base by 40% while retaining 97% of its GLUE score. The thesis is straightforward — knowledge distillation works best when you combine soft-target supervision with the original training objective and add a structural constraint that forces the student to mimic the teacher’s hidden-state geometry. Below I walk through each component, show the loss formulation, and call out where the approach succeeds or falls short in production.

The teacher-student setup

BERT-base (110M parameters, 12 layers, 768 hidden) serves as the teacher. DistilBERT (66M parameters, 6 layers, 768 hidden) is the student. The architectural changes are deliberate:

• Half the layers: 6 transformer blocks instead of 12, keeping hidden size constant so the student can attend to the same representation space.
• No token-type embeddings: The segment embeddings (used for next-sentence prediction) are dropped. DistilBERT was trained without the NSP task, which the original BERT paper later showed was unnecessary.
• No pooler: The [CLS] pooler head is removed; downstream tasks add their own.

These choices aren’t arbitrary. Keeping hidden size at 768 means the student’s hidden states live in the same vector space as the teacher’s, which makes the cosine embedding loss (covered below) well-defined without projection layers.

# Simplified DistilBERT config vs BERT-base
bert_config = {
    "num_hidden_layers": 12,
    "hidden_size": 768,
    "num_attention_heads": 12,
    "type_vocab_size": 2,          # token-type embeddings
}

distilbert_config = {
    "num_hidden_layers": 6,
    "hidden_size": 768,
    "num_attention_heads": 12,
    "type_vocab_size": 0,          # removed
}

The triple loss function

The core contribution is a weighted sum of three losses. Let $T$ be the teacher, $S$ the student, $x$ the input tokens, $y$ the masked positions, and $h^l$ the hidden state at layer $l$.

1. Distillation loss (soft targets)

Standard knowledge distillation with temperature $\tau$:

$$\mathcal{L}_{distil} = \tau^2 \cdot \text{CE}\left(\text{softmax}\left(\frac{z_T}{\tau}\right), \text{softmax}\left(\frac{z_S}{\tau}\right)\right)$$

where $z_T, z_S$ are the teacher and student logits over the vocabulary at masked positions. The $\tau^2$ scaling preserves gradient magnitude when $\tau > 1$. DistilBERT uses $\tau = 2$ (the paper ablates this; $\tau \in [1, 10]$ works).

def distillation_loss(student_logits, teacher_logits, temperature=2.0, mask=None):
    # student_logits, teacher_logits: [batch, seq_len, vocab]
    # mask: [batch, seq_len] boolean, True for masked positions
    if mask is not None:
        student_logits = student_logits[mask]
        teacher_logits = teacher_logits[mask]
    
    teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)
    
    loss = F.kl_div(student_log_probs, teacher_probs, reduction='batchmean')
    return loss * (temperature ** 2)

2. Masked language modeling loss (hard targets)

The student still trains on the original MLM objective with ground-truth labels:

$$\mathcal{L}_{mlm} = \text{CE}(y, \text{softmax}(z_S))$$

This prevents the student from collapsing to a degenerate solution that matches the teacher’s distribution but ignores the input.

3. Cosine embedding loss (hidden-state alignment)

This is the structural constraint. For each layer $l$ in the student, align its output to the corresponding teacher layer $2l$ (since the student has half the layers):

$$\mathcal{L}_{cos} = 1 - \cos(h^l_S, h^{2l}_T)$$

Summed over all 6 student layers. This forces the student’s intermediate representations to occupy the same directional cones as the teacher’s, not just match the final output distribution.

def cosine_embedding_loss(student_hidden, teacher_hidden):
    # student_hidden: list of 6 tensors [batch, seq_len, 768]
    # teacher_hidden: list of 12 tensors [batch, seq_len, 768]
    loss = 0.0
    for i, h_s in enumerate(student_hidden):
        h_t = teacher_hidden[2 * i]  # map student layer i -> teacher layer 2i
        # Flatten batch and seq_len dims
        h_s_flat = h_s.view(-1, h_s.size(-1))
        h_t_flat = h_t.view(-1, h_t.size(-1))
        loss += 1 - F.cosine_similarity(h_s_flat, h_t_flat, dim=-1).mean()
    return loss / len(student_hidden)

Combined objective

$$\mathcal{L} = \alpha \mathcal{L}{distil} + \beta \mathcal{L}{mlm} + \gamma \mathcal{L}_{cos}$$

The paper uses $\alpha = 0.5, \beta = 0.5, \gamma = 2.0$. The higher weight on cosine loss reflects its role as the primary structural regularizer — without it, the student matches logits but learns different internal representations, which hurts transfer to downstream tasks.

Training procedure details

The distillation runs on the same data as BERT: English Wikipedia + BookCorpus (~16GB text). Key hyperparameters:

The teacher is frozen. Only the student updates. This is critical — if you fine-tune the teacher simultaneously, you lose the fixed target that makes distillation a well-posed optimization problem.

One practical detail: the paper uses dynamic masking (mask pattern regenerated each epoch) rather than static masking. This matters because the student sees more diverse masking patterns over 500k steps, which acts as additional regularization.

Results and what they actually mean

The headline numbers from the paper:

The 1.6× speedup is measured on GPU (V100) with batch size 32. On CPU the speedup is closer to 2× because the reduced layer count dominates when memory bandwidth isn’t saturated.

But the GLUE average masks variance across tasks. DistilBERT loses more on tasks requiring deep reasoning (MNLI: -2.7, QQP: -1.8) and less on lexical tasks (SST-2: -0.4, QNLI: -0.9). This pattern holds across distillation recipes: the student mimics the teacher’s output distribution but lacks the depth to compose multi-hop reasoning.

Tradeoffs you should weigh

When DistilBERT wins

• Latency-constrained serving: 1.6–2× throughput improvement with minimal quality drop on classification and extraction tasks.
• Memory-constrained environments: 40% smaller model fits on edge devices or allows larger batch sizes.
• Fine-tuning speed: Fewer layers means fewer activation checkpoints, faster backward passes.

Where it falls short

• Complex reasoning: Tasks requiring chaining multiple pieces of evidence (multi-hop QA, long-context NLI) degrade disproportionately.
• Domain shift: The cosine loss aligns representations on the pretraining distribution. On out-of-domain data (biomedical, legal, code), the alignment can hurt because the teacher’s hidden states encode domain knowledge the student never saw.
• Quantization stacking: DistilBERT + INT8 quantization often degrades more than BERT-base + INT8 because the student has less capacity to absorb quantization noise.

# Rough latency comparison on V100, batch=32, seq_len=128
# (measured with torch.compile, fp16)
# BERT-base:  ~45 ms / batch
# DistilBERT: ~28 ms / batch
# Speedup:    ~1.6x

Variations and follow-ups worth knowing

The original recipe has been extended in several directions:

TinyBERT (Jiao et al., 2020) adds attention-map distillation and uses a two-stage process (general distillation on large corpus → task-specific distillation). It gets closer to BERT-base on GLUE (78.5 vs 79.5) with 4 layers and 312 hidden size.

MobileBERT (Sun et al., 2020) uses inverted-bottleneck blocks and progressive knowledge transfer from a specially designed teacher (IB-BERT). It targets mobile NPUs, not GPUs.

DistilBERT-distilled (Hugging Face’s distilbert-base-uncased-distilled-squad) shows that task-specific distillation after general distillation recovers most of the SQuAD gap. The pattern: general distillation for representation alignment, then task distillation for output alignment.

Practical recommendations

If you’re deciding whether to use DistilBERT or distill your own model:

1. Start with the off-the-shelf checkpoint (distilbert-base-uncased). It’s battle-tested, converts cleanly to ONNX/TensorRT, and the tokenizer is identical to BERT’s.

2. If you need domain adaptation, continue pretraining DistilBERT on your corpus with MLM loss only (no distillation). The cosine loss requires a teacher, which you may not have for your domain.

3. If you’re distilling a custom teacher, copy the triple loss exactly but tune $\gamma$ (cosine weight). For deeper teachers (24 layers), map student layer $l$ to teacher layer $4l$ or use a learned projection.

4. Don’t skip the cosine loss. Ablation in the paper shows removing it drops GLUE by 1.5 points — more than removing the distillation loss (0.8 points). The hidden-state alignment is the secret sauce.

5. Benchmark on your hardware. The 1.6× speedup assumes GPU with large batch. On CPU with batch=1, the speedup shrinks to ~1.3× because kernel launch overhead dominates. On Apple Silicon with ANE, the smaller model may not map efficiently to the neural engine.

The decisive takeaway

How DistilBERT was distilled matters because it proved that a carefully designed triple loss — soft targets, hard targets, and hidden-state geometry — can compress a transformer by 40% with negligible quality loss on the tasks that dominate production workloads (classification, NER, extractive QA). The recipe is reproducible, the checkpoint is reliable, and the tradeoffs are well-characterized.

For most teams, the right move is not to re-distill but to start from distilbert-base-uncased, fine-tune on your task, and only invest in custom distillation if you have a proprietary teacher that encodes knowledge no public model captures. The marginal gain from tuning $\alpha, \beta, \gamma$ rarely justifies the engineering cost compared to collecting better task data.