Tiny BERT

Analysis

Why This Paper Matters

Large language models like BERT have revolutionized natural language understanding, but their size (109 million parameters for BERT BASE) makes them impractical for many real-world applications—especially on mobile devices, IoT, or latency-sensitive systems. TinyBERT addresses this gap by compressing BERT into a much smaller model (14.5 million parameters) while retaining most of its accuracy. The key insight is that knowledge distillation can be made more effective by focusing on the internal representations of the teacher, particularly the attention matrices, and by using a two-stage training process that first distills general language knowledge and then task-specific knowledge.

How It Works

TinyBERT builds on standard knowledge distillation, where a small student network learns to mimic a large teacher network. The authors define a behavior function that captures what the student should imitate. For Transformers, this can be the output of multi-head attention (MHA) layers, feed-forward network (FFN) layers, or intermediate representations like the attention matrix A.

The distillation objective minimizes the difference between teacher and student behavior functions over the training data:

A mapping function g(m) selects which teacher layers correspond to each student layer. For example, in TinyBERT4 (4 student layers), g(m) = 3m maps student layer 1 to teacher layer 3, layer 2 to layer 6, etc. The total loss combines several components:

Transformer Layer Distillation includes attention transfer (matching attention matrices A_i for each head) and hidden-state regression (matching the output of MHA and FFN sub-layers).

Embedding Layer Distillation aligns the token embeddings of student and teacher:

Prediction Layer Distillation uses soft targets from the teacher's logits with temperature scaling:

These are combined into a unified loss:

The two-stage learning process first performs general distillation on unlabeled data to transfer broad linguistic knowledge, then task-specific distillation on labeled data to adapt to the target task. This separation is crucial because it allows the student to first learn general language patterns before specializing.

Results

TinyBERT was evaluated on the GLUE benchmark. TinyBERT4 (4 layers, hidden size 312) achieves an average score of 77.6, compared to BERT BASE's 80.5—a drop of only 3.6% while using 7.5x fewer parameters. TinyBERT6 (6 layers, hidden size 768) matches or exceeds BERT BASE on several tasks, including MRPC (85.0 vs 84.8) and CoLA (51.1 vs 52.1).

These results demonstrate that careful distillation of attention and hidden states can preserve most of BERT's performance even with a fraction of the parameters.

Significance

TinyBERT's two-stage distillation framework has become a standard approach for compressing large language models. By separating general and task-specific knowledge transfer, it achieves a practical balance between model size and accuracy. The attention transfer mechanism is particularly influential, as it captures the relational structure that BERT learns. This work has enabled deployment of BERT-like models in production environments with limited compute, such as mobile search, real-time chatbots, and on-device NLP. Future work could extend the approach to even larger models like RoBERTa or T5, or explore adaptive layer mapping strategies.