Knowledge Distillation

The core framework involves a two-model setup:

• Teacher Model: A high-capacity, pre-trained model (e.g., BERT-Large) that produces a rich output distribution, including class probabilities for every token.
• Student Model: A lightweight architecture (e.g., DistilBERT or a BiLSTM) trained to mimic the teacher's behavior. The student learns not just from hard labels but from the soft targets provided by the teacher's output layer.

Instead of training only on one-hot encoded 'hard' labels, the student learns from the teacher's probability distribution.

• A temperature parameter (T) is applied to the softmax function to soften the probability distribution.
• Higher temperatures reveal the dark knowledge—the relative similarities between incorrect classes that the teacher has learned.
• The student is trained with a combined loss: a distillation loss matching the teacher's soft targets and a student loss matching the ground-truth hard labels.

In sequence labeling tasks like NER, distillation is applied at the token level:

• The teacher model outputs a logit vector for each token (e.g., B-PER, I-PER, O).
• The student is trained to minimize the KL divergence between its own softened token probabilities and the teacher's.
• This is particularly effective for transferring knowledge about rare entity types or ambiguous boundary cases, where the teacher's soft distribution provides a richer training signal than a single hard label.

DistilBERT is a prominent example of knowledge distillation applied to the BERT architecture:

• It retains 97% of BERT's language understanding capabilities while being 40% smaller and 60% faster.
• The student is initialized from the teacher by taking one out of every two layers.
• It is trained using a triple loss combining language modeling, distillation, and cosine-distance losses between teacher and student hidden states.
• This proves that a heavily compressed model can retain near-state-of-the-art NER performance.

Knowledge can be transferred at different levels of abstraction:

• Response-Based: The student mimics the final output layer (logits) of the teacher. This is the most common and straightforward method.
• Feature-Based: The student learns to replicate the intermediate representations or hidden state activations of the teacher. A projection layer may be used to match dimensions.
• Relation-Based: The student learns to preserve the relationships (e.g., pairwise distances or angle-wise similarities) between data samples as modeled by the teacher's internal layers.

The primary driver for knowledge distillation is operational efficiency:

• Latency Reduction: Smaller models execute inference in milliseconds, critical for real-time entity extraction in live chat or voice systems.
• Memory Footprint: A distilled student model can fit into the constrained RAM of mobile devices or edge hardware.
• Energy Efficiency: Reduced computational load directly translates to lower power consumption, enabling sustainable, always-on NER services.
• Privacy: Running a capable NER model entirely on-device eliminates the need to transmit sensitive text to a cloud API.

The foundational framework where a large, high-capacity teacher model generates soft probability distributions over output classes, and a compact student model is trained to mimic these distributions rather than hard ground-truth labels. The student learns from both the correct answer and the teacher's dark knowledge—the relative probabilities assigned to incorrect classes, which encode valuable information about inter-class similarities and the teacher's internal feature representations.

A mechanism that controls the smoothness of the teacher's output distribution using a temperature parameter (T). Higher temperatures produce softer probability distributions that reveal more of the teacher's learned similarities between classes. Key behaviors:

• T=1: Standard softmax output
• T>1: Softer distribution, exposing dark knowledge
• T→∞: Uniform distribution, maximum softening The student is typically trained with a combined loss: KL divergence with soft targets plus standard cross-entropy with hard labels.

A landmark application of knowledge distillation to transformer-based language models. DistilBERT is a student model with 40% fewer parameters than BERT-base, trained using a triple loss combining:

• Distillation loss: KL divergence between student and teacher output probabilities
• Masked language modeling loss: Standard pre-training objective
• Cosine embedding loss: Aligning the student's hidden states with the teacher's It retains 97% of BERT's performance on downstream tasks while being 60% faster at inference.

Two primary distillation paradigms:

Response-Based: The student mimics only the final output layer (logits) of the teacher. Simple but loses intermediate representations.

Feature-Based: The student is trained to match the teacher's internal hidden states and attention patterns at multiple layers. This transfers richer representational knowledge but requires careful layer mapping between architectures of different depths.

Modern NER distillation often combines both approaches for optimal span boundary detection.

When distilling for Named Entity Recognition, additional considerations apply beyond standard classification distillation:

• Sequence-level distillation: Matching the teacher's per-token probability distributions across the entire sequence
• CRF layer transfer: Initializing the student's CRF transition matrix from the teacher's learned transition parameters
• Span boundary alignment: Adding auxiliary losses that penalize the student for entity boundary errors the teacher would not make
• Logit distillation on BIO tags: Transferring the teacher's confidence across the full tag set, including rare entity types

Offline Distillation: A pre-trained teacher generates soft targets once, and the student trains on these static targets. This is the standard approach and requires no co-training infrastructure.

Online Distillation: Teacher and student train simultaneously, with the teacher updating its parameters as the student learns. This enables mutual learning where both models improve together.

Self-Distillation: A model distills knowledge into an identical architecture, often using earlier checkpoints as teachers for later training stages—no separate teacher model required.