DistilBERT

The core compression technique is a 50% reduction in the number of transformer layers. The original BERT-base model has 12 encoder layers, while DistilBERT uses only 6. This is the primary driver for the model being 40% smaller and 60% faster at inference. The student model's layers are initialized by copying every other layer from the teacher (BERT), providing a strong starting point for training.

DistilBERT is trained using a triple-loss objective that combines:

• Distillation Loss (L_ce): A Kullback-Leibler Divergence loss that forces the student's output logits to match the softened probability distribution (soft targets) from the BERT teacher model. Temperature scaling (T=2.0) is applied to the teacher's softmax to create richer, softer targets containing dark knowledge.
• Masked Language Modeling Loss (L_mlm): The standard BERT pre-training objective, predicting masked tokens.
• Cosine Embedding Loss (L_cos): Aligns the direction of the student and teacher's hidden state vectors, encouraging similar contextual representations.

To further streamline the architecture, DistilBERT removes two specific components from BERT:

• Token-Type Embeddings: These embeddings distinguish between sentence pairs (e.g., for Next Sentence Prediction). DistilBERT omits them, as it is trained solely on the Masked Language Modeling objective, simplifying the input processing.
• Pooler Layer: The fully-connected pooler layer used for Next Sentence Prediction classification is eliminated. For downstream tasks requiring a pooled representation, the standard practice is to use the embedding of the [CLS] token or mean/max pooling over the sequence.

While the depth is halved, DistilBERT retains the same hidden dimension size (768) and vocabulary size (30522) as BERT-base. This preservation of width is crucial. It allows the student model to maintain a representational capacity similar to the teacher in each layer, facilitating more effective knowledge transfer. The feed-forward network dimension within each transformer block also remains at 3072.

DistilBERT's training begins not from random initialization, but from a smart initialization derived from the BERT teacher. Specifically, the 6 layers of the student are initialized by copying weights from alternating layers of the 12-layer BERT teacher (e.g., student layer 1 from teacher layer 2, student layer 2 from teacher layer 4, etc.). This provides the student with a functional, pre-trained starting point, significantly accelerating and stabilizing the distillation process.

DistilBERT is pre-trained on the same corpus as BERT: Wikipedia and BookCorpus, using the Masked Language Modeling objective only. The training leverages:

• Large Batch Sizes: Effective batch size of 4000 sequences.
• Dynamic Masking: A new masking pattern is generated for each sequence each time it is fed to the model, unlike BERT's static masking.
• No Next Sentence Prediction: Simplifies the pre-training task, aligning with the removed architectural components. This focused training, combined with distillation, allows it to retain 97% of BERT's language understanding capability (as measured on the GLUE benchmark) despite the architectural reduction.

Knowledge distillation is the overarching model compression technique where a smaller student model is trained to replicate the behavior and outputs of a larger, pre-trained teacher model. The core idea is to transfer the teacher's learned representations and decision boundaries, often encapsulated in its softened output probabilities (soft targets), which contain richer information than hard class labels. This process allows the student to achieve comparable performance with a fraction of the parameters and computational cost.

The teacher-student framework is the foundational two-model architecture for knowledge distillation. It consists of:

• A teacher model: A large, accurate, and often cumbersome pre-trained model (e.g., BERT-base). Its role is to provide supervisory signals.
• A student model: A smaller, more efficient architecture (e.g., DistilBERT) designed for deployment. It is trained to mimic the teacher. The framework defines the direction of knowledge flow and the training objectives used to align the student's predictions with the teacher's.

Logits distillation is the most common form of distillation, where the student is trained to match the teacher's raw, pre-softmax output logits. The primary loss function used is the Kullback-Leibler (KL) Divergence, which measures how one probability distribution diverges from a second, reference distribution. The training objective combines:

• Distillation Loss (KL Divergence): Minimizes the difference between the softened logit distributions of teacher and student.
• Standard Cross-Entropy Loss: Ensures the student also learns from the original hard labels. Temperature scaling (parameter T) is applied to the softmax to soften the teacher's output distribution, making the "dark knowledge" between classes more accessible.

Beyond logits, knowledge can be transferred from intermediate representations. Feature distillation (or hint learning) trains the student to match the teacher's intermediate feature maps or activations from specific layers, often using a mean squared error loss. For transformer models like BERT, attention distillation is particularly effective. Here, the student is trained to replicate the self-attention matrices (or patterns) from the teacher's layers, forcing it to learn similar contextual relationships between tokens. DistilBERT employs a combination of these losses alongside logits distillation.

TinyBERT is a prominent, task-specific distilled version of BERT that achieves even greater compression than DistilBERT. Its key innovation is a two-stage distillation process:

1. General Distillation: Distills BERT's knowledge during pre-training on a large corpus.
2. Task-Specific Distillation: Further distills the model during fine-tuning on a downstream task (e.g., sentiment analysis). This approach applies transformer-layer distillation (attention and hidden states) at both stages, resulting in a model that is highly optimized for a specific use case and extremely efficient for on-device deployment.

These are variants of the standard distillation paradigm:

• Online Distillation: The teacher and student models are co-trained simultaneously from scratch, rather than using a static, pre-trained teacher. This can be more efficient and allows both models to improve collaboratively.
• Self-Distillation: A model serves as both its own teacher and student. A common technique is born-again networks, where a student with the same architecture as the teacher is trained via distillation to outperform its teacher, often through iterative cycles. Another form distills knowledge from deeper layers to shallower layers within the same network.