Attention Distillation

The primary goal is to train a student transformer to replicate the self-attention matrices or attention maps generated by its teacher model. These matrices, calculated in each layer and attention head, represent the model's focus—which input tokens are most relevant to each other. By forcing the student to produce similar attention distributions, it learns the teacher's internal reasoning and feature alignment strategies, often leading to better generalization than just matching final outputs.

Specialized loss functions measure the divergence between teacher and student attention. Common approaches include:

• Mean Squared Error (MSE) Loss: Applied directly to the attention score matrices (QK^T) before the softmax, or to the post-softmax attention probabilities.
• Kullback-Leibler (KL) Divergence Loss: Measures the difference between the probability distributions of the teacher and student's softmax-normalized attention weights.
• Cross-Entropy on Attention Maps: Treats the teacher's attention distribution for a given query token as a target label for the student. These losses are often applied layer-wise, aligning corresponding layers between teacher and student, or adaptively across layers.

Attention distillation operates at a finer granularity than standard logit distillation:

• Layer-Wise Distillation: The student is guided at multiple intermediate layers, not just the final output. This provides a stronger, more continuous learning signal throughout the network's depth.
• Attention Head Alignment: In multi-head attention, distillation can be applied per head, forcing the student to learn the specialized roles (e.g., syntactic, semantic) that different heads develop in the teacher.
• Combined Objectives: Often used in conjunction with logit distillation and hidden state (feature) distillation for a multi-faceted knowledge transfer, summarized as Total Loss = L_attn + L_logits + L_hidden.

Attention distillation is highly flexible regarding model architectures:

• Homogeneous Distillation: Teacher and student share the same transformer architecture but differ in size (e.g., layers, hidden dimensions, attention heads).
• Heterogeneous Distillation: The student can have a fundamentally different architecture (e.g., a CNN student learning from a ViT teacher). The key is defining a projection or alignment method to match the spatial dimensions of their attention maps.
• Attention Mask Transfer: For tasks like masked language modeling, the technique can also transfer knowledge about how the teacher attends to masked versus unmasked tokens.

This technique is pivotal for creating efficient models for deployment:

• Creating Compact Language Models: Models like TinyBERT and MobileBERT use attention distillation to produce small, fast BERT variants for mobile devices.
• Efficient Vision Transformers: DeiT (Data-efficient Image Transformers) uses a distillation token to learn from a CNN teacher's attention-like features, achieving high accuracy with less data.
• Cross-Modal Transfer: Aligning attention patterns from a large, multi-modal teacher (e.g., CLIP) to a smaller, uni-modal student.
• Pruning and Quantization Prep: Softening the model via attention distillation before applying aggressive pruning or quantization can improve the final compressed model's robustness.

Attention distillation offers specific benefits:

• Richer Supervision Signal: Attention maps provide structural and relational information about the data that softened logits alone do not capture.
• Improved Intermediate Representations: By aligning internal mechanisms, the student develops better feature extractors early on, which can accelerate convergence and improve final performance.
• Mitigates Capacity Gap: For a very small student, matching complex final logits may be too difficult. Matching simpler, intermediate attention patterns can be a more achievable learning objective.
• Interpretability: The distilled attention maps in the student are often more interpretable and resemble the teacher's, providing some transparency into the compact model's decision process.

Feature distillation is a knowledge transfer method where the student model is trained to directly match the intermediate feature representations or activations from specific layers of the teacher model. Unlike logits distillation, it provides guidance on how the teacher constructs its internal representations.

• Common Approach: Minimizes the mean squared error (MSE) or cosine similarity between teacher and student feature maps at designated hint layers.
• Advantage: Provides richer, more granular supervisory signals throughout the network's depth, often leading to better student convergence and final performance compared to output-only distillation.

Attention transfer is a specific form of feature distillation pioneered for convolutional neural networks (CNNs) and later adapted for transformers. The student is trained to replicate the spatial attention maps generated by the teacher's intermediate layers.

• For CNNs: Uses the sum of absolute values of feature maps across channels to create 2D attention maps.
• For Transformers: Directly aligns the self-attention matrices from corresponding teacher and student layers, which is the foundation of attention distillation. This forces the student to learn the same contextual relationship priorities as the teacher.

Logits distillation is the original and most common form of knowledge distillation. The student model is trained to match the raw, pre-softmax output logits of the teacher model, in addition to (or instead of) the ground truth labels.

• Temperature Scaling: A temperature parameter (T > 1) is applied to the teacher's softmax to create a softer probability distribution, revealing more inter-class relationships (dark knowledge).
• Loss Function: Typically uses a weighted sum of Kullback-Leibler divergence (for softened teacher targets) and standard cross-entropy loss (for true labels).

The teacher-student framework is the foundational two-model architecture that underpins all knowledge distillation techniques. It consists of a pre-trained, often cumbersome teacher model and a smaller, trainable student model.

• Static vs. Online: In classic distillation, the teacher is static (pre-trained and frozen). In online distillation, both models are co-trained from scratch.
• Variations: Includes self-distillation (same model acts as teacher and student), multi-teacher distillation (knowledge from an ensemble), and mutual learning (multiple peer students teach each other).

Quantization-aware distillation (QAD) is a joint optimization technique where knowledge distillation is performed during training while simulating the effects of quantization. This prepares the student model for efficient low-precision (e.g., INT8) deployment from the outset.

• Process: The forward pass uses quantized weights and activations (simulated in floating-point), while the backward pass updates full-precision weights. The distillation loss is computed using these quantized outputs.
• Outcome: Produces a student model that is inherently robust to the precision loss caused by post-training quantization, achieving better accuracy than applying distillation and quantization sequentially.