Attention Transfer

Unlike standard Knowledge Distillation (KD), which primarily transfers knowledge from the final output logits, Attention Transfer enforces alignment at intermediate layers. The student is trained to produce attention maps—visualizations of where the model 'looks'—that are similar to those generated by the teacher's transformer blocks or convolutional layers. This provides a richer, more structured learning signal that guides the student's internal feature formation.

The core objective is to teach the student how the teacher reasons, not just what it predicts. By matching attention maps, the student learns:

• Spatial focus in vision tasks: which regions of an image are salient for classification.
• Contextual focus in language tasks: which tokens in a sequence are most relevant for understanding relationships. This forces the student to develop similar internal feature selectivity, often leading to better generalization and robustness than learning from soft targets alone.

Training requires a loss function that quantifies the difference between teacher and student attention. Common choices include:

• Mean Squared Error (MSE) Loss: Directly minimizes the L2 distance between corresponding attention maps.
• Kullback-Leibler (KL) Divergence Loss: Treats attention maps as probability distributions and minimizes their statistical distance.
• Cosine Similarity Loss: Maximizes the directional alignment of the attention vectors. This attention loss is typically combined with the standard task loss (e.g., cross-entropy) and sometimes a traditional KD loss on logits.

The technique is architecture-agnostic but is prominently applied to:

• Vision Transformers (ViTs): Transferring attention from multi-head self-attention layers.
• Convolutional Neural Networks (CNNs): Using activation-based attention maps from layers like Grad-CAM.
• Large Language Models (LLMs): Distilling contextual attention patterns from teacher decoders. For transformers, attention maps are a natural by-product. For CNNs, they are often generated by summing absolute values of feature map activations across channels.

Attention Transfer is particularly effective when there is a significant model capacity gap between teacher and student. Mimicking intermediate attention provides a structured, step-by-step guide for the smaller network, making the compression process more stable and efficient. It can be seen as a form of intermediate supervision that alleviates the difficulty of the student having to discover useful feature representations from scratch using only final-layer signals.

Attention Transfer is a direct descendant and specialization of the broader Hint Training paradigm introduced by models like FitNets. While hint training can involve matching any intermediate feature activation, Attention Transfer specifically targets the normalized attention weights or derived attention maps. This makes it a more focused and often more efficient method, as attention mechanisms are designed to encode relational importance, providing a high-value signal for distillation.

Knowledge Distillation (KD) is the overarching model compression paradigm where a compact student model is trained to replicate the behavior of a larger, pre-trained teacher model. The core mechanism involves the student learning from the teacher's softened output probabilities (soft targets), which contain dark knowledge about inter-class relationships, rather than just hard labels. This process is governed by a distillation loss, typically a combination of standard cross-entropy and a mimicry loss like Kullback-Leibler Divergence (KL Divergence).

• Primary Goal: Achieve teacher-level accuracy with a smaller, faster model.
• Key Innovation: Utilizing the teacher's output distribution as a richer training signal.

Feature-based distillation is a category of knowledge transfer where the student is trained to mimic the teacher's intermediate feature representations or activations, not just its final outputs. Attention transfer is a prime example of this approach.

• Mechanism: Minimizes a distance metric (e.g., Mean Squared Error) between the feature maps of selected teacher and student layers.
• Rationale: Intermediate features capture the teacher's internal data transformations and spatial focus, providing a stronger learning signal.
• Variants: Include hint training (direct regression on teacher hints) and matching Gram matrices of features.

The teacher model and student model are the two core entities in any distillation pipeline.

• Teacher Model: A large, accurate, and often over-parameterized neural network (e.g., BERT, ResNet-50) that has been pre-trained on a target task. It serves as the source of knowledge, providing targets (logits, features, attention maps) for the student.
• Student Model: A smaller, more efficient neural network architecture (e.g., a shallower or narrower network) designed for deployment. Its parameters are trained from scratch or initialized, with the primary objective of replicating the teacher's predictive behavior and internal representations.

Soft targets and temperature scaling are fundamental to the classic logit-based distillation formulated by Hinton et al.

• Soft Targets: The probability distribution produced by applying the softmax function to the teacher's output logits. These probabilities are 'softer' than one-hot labels, revealing which classes the teacher considers similar (e.g., 'cat' and 'lynx').
• Temperature Scaling (T): A hyperparameter used to control the softness of these targets. The logits are divided by T before softmax: softmax(logits / T). A T > 1 produces a smoother, more informative probability distribution, emphasizing the dark knowledge crucial for the student's learning.

The distillation loss is the composite objective function that trains the student model. It balances learning from ground truth and learning from the teacher.

• Standard Formulation: L_total = α * L_CE(y_true, y_student) + β * L_KD(σ(z_teacher/T), σ(z_student/T))
• L_CE: Cross-entropy loss with hard labels.
• L_KD: The knowledge distillation loss. Most commonly the Kullback-Leibler Divergence Loss, which measures how one probability distribution diverges from another.
• α, β: Weighting coefficients that balance the two objectives. For attention transfer, an additional loss term matching attention maps is added.

These are advanced distillation paradigms that relax the requirement for a static, pre-trained teacher.

• Self-Distillation: The teacher and student are the same model or share an architecture. Knowledge is transferred from the model's own deeper layers to shallower layers, or from a later training checkpoint to an earlier one. Born-Again Networks (BAN) are a key example, where iteratively trained generations of the same model act as teacher and student.
• Online Distillation: The teacher model is not fixed but is updated concurrently with the student during a single training run. This can involve an ensemble of peer students teaching each other or a continuously evolving teacher, eliminating the need for a separate, costly pre-training phase.