BERT vs DistilBERT vs TinyBERT: Speed-Accuracy Trade-offs

The 300MB Model That Runs in 40ms

BERT-base sits at 110M parameters and 440MB on disk. DistilBERT cuts that to 66M parameters and 260MB. TinyBERT goes further — 14.5M parameters, 55MB, and inference that actually runs on a Raspberry Pi without thermal throttling.

But here’s the catch: you’re not just shrinking the model. You’re fundamentally changing how it encodes language. DistilBERT uses knowledge distillation, keeping the full 768-dimensional hidden states. TinyBERT applies distillation at every layer AND shrinks the hidden size to 312 dimensions. That architectural difference matters more than the parameter count suggests.

I’m comparing all three on a real-world task: sentiment classification on 50k IMDB reviews. The goal is to see where the accuracy drop actually hurts, and where it’s just noise.

BERT-base: The Baseline You’re Probably Overpaying For

BERT-base (Devlin et al., 2019) uses 12 transformer layers, 768 hidden dimensions, and 12 attention heads. The attention mechanism computes:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

where $d_k = 64$ (768 hidden size / 12 heads). This scaled dot-product attention runs 12 times per layer, across 12 layers. That’s 144 attention operations per forward pass.

Here’s the PyTorch setup:

from transformers import BertTokenizer, BertForSequenceClassification
import torch
import time

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForSequenceClassification.from_pretrained(
    'bert-base-uncased', 
    num_labels=2
)
model.eval()

# Sample text
text = "This movie was surprisingly good despite the weak ending."
inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=128)

# Warmup (first pass is always slower due to CUDA init)
with torch.no_grad():
    _ = model(**inputs)

# Actual timing
start = time.perf_counter()
with torch.no_grad():
    outputs = model(**inputs)
end = time.perf_counter()

print(f"BERT-base inference: {(end - start) * 1000:.2f}ms")
print(f"Model size on disk: {sum(p.numel() for p in model.parameters()) * 4 / 1024 / 1024:.1f}MB")

On my M1 MacBook (CPU inference, no MPS acceleration because I’m simulating a production server), this outputs:

BERT-base inference: 127.34ms
Model size on disk: 418.5MB

That 127ms might not sound bad, but it’s per request. At 100 requests/second, you’re looking at 12.7 CPU cores fully saturated. For a startup API, that’s expensive.

DistilBERT: Same Width, Half the Depth

DistilBERT (Sanh et al., 2019) keeps the 768 hidden size but drops from 12 layers to 6. The distillation loss combines three objectives:

$$L_{\text{distil}} = \alpha L_{\text{ce}} + \beta L_{\text{mlm}} + \gamma L_{\text{cos}}$$

where $L_{\text{ce}}$ is the cross-entropy between student and teacher softmax outputs (with temperature $T=2$), $L_{\text{mlm}}$ is the masked language modeling loss, and $L_{\text{cos}}$ is the cosine embedding loss between hidden states. The temperature softening matters — without it, the student just memorizes hard labels instead of learning the teacher’s uncertainty.

from transformers import DistilBertTokenizer, DistilBertForSequenceClassification

tokenizer_distil = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
model_distil = DistilBertForSequenceClassification.from_pretrained(
    'distilbert-base-uncased',
    num_labels=2
)
model_distil.eval()

inputs_distil = tokenizer_distil(text, return_tensors='pt', padding=True, truncation=True, max_length=128)

# Warmup
with torch.no_grad():
    _ = model_distil(**inputs_distil)

start = time.perf_counter()
with torch.no_grad():
    outputs_distil = model_distil(**inputs_distil)
end = time.perf_counter()

print(f"DistilBERT inference: {(end - start) * 1000:.2f}ms")
print(f"Model size: {sum(p.numel() for p in model_distil.parameters()) * 4 / 1024 / 1024:.1f}MB")

Output:

DistilBERT inference: 68.21ms
Model size: 255.3MB

That’s a 46% speedup with 40% fewer parameters. But here’s what surprised me: on a batch size of 32 (more realistic for production), the gap narrows to 38%. The overhead of tokenization and data movement starts dominating at larger batches.

TinyBERT: Aggressive Compression Meets Layer-wise Distillation

TinyBERT (Jiao et al., 2020) doesn’t just halve the layers — it shrinks the hidden size from 768 to 312 and cuts attention heads from 12 to 12… wait, no, to 12. That’s odd. They keep 12 heads but with only 312 hidden dimensions, so each head operates on $312 / 12 = 26$ dimensions instead of 64.

The distillation happens at three levels:

1. Embedding layer distillation: Match the token embeddings
2. Transformer layer distillation: Match attention matrices AND hidden states
3. Prediction layer distillation: Match softmax logits

The attention distillation loss is:

$$L_{\text{attn}} = \frac{1}{h} \sum_{i=1}^{h} \text{MSE}(A_i^S, A_i^T)$$

where $A_i^S$ and $A_i^T$ are the attention distributions for the $i$-th head in the student and teacher. This forces the student to mimic not just the output, but the internal attention patterns.

# TinyBERT isn't in Hugging Face's main repo, so we use a community upload
# (This is one of those cases where you have to trust a third-party checkpoint)
from transformers import AutoTokenizer, AutoModelForSequenceClassification

tokenizer_tiny = AutoTokenizer.from_pretrained('huawei-noah/TinyBERT_General_4L_312D')
model_tiny = AutoModelForSequenceClassification.from_pretrained(
    'huawei-noah/TinyBERT_General_4L_312D',
    num_labels=2,
    ignore_mismatched_sizes=True  # We're fine-tuning, so output layer mismatch is expected
)
model_tiny.eval()

inputs_tiny = tokenizer_tiny(text, return_tensors='pt', padding=True, truncation=True, max_length=128)

# Warmup
with torch.no_grad():
    _ = model_tiny(**inputs_tiny)

start = time.perf_counter()
with torch.no_grad():
    outputs_tiny = model_tiny(**inputs_tiny)
end = time.perf_counter()

print(f"TinyBERT inference: {(end - start) * 1000:.2f}ms")
print(f"Model size: {sum(p.numel() for p in model_tiny.parameters()) * 4 / 1024 / 1024:.1f}MB")

Output:

TinyBERT inference: 41.87ms
Model size: 54.8MB

Now we’re talking. 41ms is fast enough to run on an AWS Lambda cold start (850ms budget) with room for tokenization overhead. The model fits in a Docker image without bloating it beyond 200MB.

Fine-Tuning on IMDB: Where the Accuracy Drop Appears

I fine-tuned all three on 25k IMDB training samples, validated on 25k test samples. Same hyperparameters: learning rate $2 \times 10^{-5}$, batch size 16, 3 epochs, AdamW optimizer with weight decay DOLLAR_AMOUNT_2$. The loss function is binary cross-entropy:

$$L = -\frac{1}{N} \sum_{i=1}^{N} \left[ y_i \log(\hat{y}_i) + (1 – y_i) \log(1 – \hat{y}_i) \right]$$

Results after 3 epochs:

That 4.8% accuracy drop from BERT to TinyBERT isn’t huge, but look at the F1 gap: 0.049. For sentiment analysis, that’s borderline acceptable. For medical entity extraction? I wouldn’t ship it.

But the training time difference caught me off guard:

• BERT-base: 2h 14min (on a single RTX 3090)
• DistilBERT: 1h 18min
• TinyBERT: 47min

TinyBERT trains 2.8x faster. If you’re iterating on a new dataset and need to run 10 experiments to find the right augmentation strategy, that’s 22 hours saved. Enough time to grab Peet’s Major Dickason’s Blend and actually sleep.

The Real Decision Matrix: When Each Model Wins

Use BERT-base when:
– You have < 1000 requests/day and latency isn’t a bottleneck
– You need the absolute best accuracy for a compliance-critical task
– You’re building a baseline to compare against before optimizing

Use DistilBERT when:
– You need production-ready speed without sacrificing much accuracy
– Your API handles 100-10k requests/day and you want to keep infrastructure costs low
– You’re deploying to a mid-tier cloud instance (2-4 CPU cores)

Use TinyBERT when:
– You’re deploying on mobile or edge devices (Raspberry Pi, Jetson Nano, old Android phones)
– Your latency budget is < 50ms per request
– You have a massive unlabeled dataset and can afford to fine-tune from scratch (TinyBERT benefits a lot from task-specific distillation)

I’m not entirely sure why TinyBERT’s performance drops more on longer sequences (> 256 tokens) compared to DistilBERT. My best guess is that the smaller hidden size (312 vs 768) reduces the representational capacity for distant token interactions. The self-attention scores get noisier when you have fewer dimensions to encode positional relationships.

The Hidden Cost: Fine-Tuning Instability

TinyBERT is more sensitive to learning rate. I initially used the same $2 \times 10^{-5}$ that worked for BERT and DistilBERT, but the validation loss started oscillating after epoch 2:

Epoch 1: train_loss=0.342, val_loss=0.298, val_acc=89.1%
Epoch 2: train_loss=0.198, val_loss=0.281, val_acc=89.7%
Epoch 3: train_loss=0.124, val_loss=0.312, val_acc=88.8%  # Overfitting kicks in

Dropping the learning rate to $1 \times 10^{-5}$ and adding gradient clipping (max_norm=1.0) stabilized it:

from torch.nn.utils import clip_grad_norm_

for epoch in range(3):
    for batch in train_loader:
        optimizer.zero_grad()
        outputs = model_tiny(**batch)
        loss = outputs.loss
        loss.backward()
        clip_grad_norm_(model_tiny.parameters(), max_norm=1.0)  # This matters
        optimizer.step()

After that adjustment:

Epoch 3: train_loss=0.156, val_loss=0.289, val_acc=89.2%

Still not as good as DistilBERT, but stable. Smaller models have less regularization capacity, so they overfit faster. If you’re short on labeled data (< 5k samples), DistilBERT is safer.

Quantization: The 4x Speedup Nobody Mentions

All the numbers above are FP32. If you quantize to INT8 using PyTorch’s dynamic quantization, the speedup stacks:

import torch.quantization

model_tiny_quantized = torch.quantization.quantize_dynamic(
    model_tiny, 
    {torch.nn.Linear}, 
    dtype=torch.qint8
)

start = time.perf_counter()
with torch.no_grad():
    _ = model_tiny_quantized(**inputs_tiny)
end = time.perf_counter()

print(f"TinyBERT (INT8) inference: {(end - start) * 1000:.2f}ms")

Output:

TinyBERT (INT8) inference: 23.14ms

That’s 5.5x faster than BERT-base. The accuracy drop from quantization is negligible (< 0.5% on IMDB), but I’ve seen it hurt more on named entity recognition tasks where rare tokens matter. Always benchmark on your specific task.

FAQ

Q: Can I use DistilBERT embeddings as a drop-in replacement for BERT embeddings in downstream tasks?

Yes, but the hidden size is still 768, so you’re not saving much on the embedding layer. The real speedup comes from the reduced transformer layers (6 vs 12). If you’re using embeddings for semantic search with FAISS, the indexing time is the same — only the encoding step gets faster.

Q: Does TinyBERT work well for non-English languages?

Multilingual TinyBERT exists (TinyBERT_General_6L_768D), but the performance gap widens compared to mBERT. For low-resource languages, I’d stick with DistilBERT or even full mBERT if accuracy is critical. The distillation process assumes the teacher model is already strong, and mBERT’s multilingual capacity is harder to compress without losing cross-lingual transfer.

Q: What if I need even smaller than TinyBERT for mobile deployment?

Look into MobileBERT (Sun et al., 2020) or ALBERT-tiny. MobileBERT has 25M parameters but uses bottleneck structures to keep inference fast on ARM CPUs. ALBERT-tiny (11M parameters) uses parameter sharing across layers, though I’ve found it underperforms TinyBERT on most tasks. For truly aggressive compression, consider switching to a task-specific CNN or LSTM — transformers aren’t always the answer when you’re sub-50MB.

My Pick for a First NLP Project

DistilBERT. It’s the sweet spot.

You get 91-92% of BERT’s accuracy at half the latency and 60% of the model size. The Hugging Face integration is seamless, the pre-trained checkpoints are reliable, and the fine-tuning process is stable enough that you won’t waste a day debugging gradient explosions.

TinyBERT is tempting if you’re resource-constrained, but the fine-tuning instability and accuracy drop make it a poor choice for your first project. You want to prove the concept works before optimizing. Start with DistilBERT, get to 90% accuracy, then decide if the extra 2% is worth the switch to BERT-base or if the extra 10ms latency savings justify TinyBERT.

One thing I’m still curious about: how much does the distillation dataset size matter? The original DistilBERT paper used the full English Wikipedia + BookCorpus (same as BERT pre-training), but what if you distill on a domain-specific corpus? Could you get better task performance by distilling directly on medical papers for a biomedical NER task? I haven’t tested this at scale, but the few experiments I’ve run suggest a 1-2% accuracy boost. Worth exploring if you have the compute budget.