Feature and Attention Transfer in Distillation

A teacher's output distribution over ten CIFAR classes is ten numbers per image. A single early feature map of 64 channels at 32 by 32 is 65,536. Output distillation, covered in Temperature and Dark Knowledge, throws that internal signal away and asks the student to rediscover how to compute the answer. Feature-based distillation supervises the route as well as the destination, and it immediately runs into a problem the output version never has: the student's internals are a different shape from the teacher's.

Hints and the regressor

FitNets introduced the idea to train students that were thinner and deeper than their teachers, which plain output matching struggled to optimise (Romero et al., 2015, FitNets: Hints for Thin Deep Nets, arXiv:1412.6550). A hint layer in the middle of the teacher supervises a guided layer in the middle of the student. Because the widths differ, a small learned regressor \(r\) maps the student feature into the teacher's space:

where \(u_h\) and \(v_g\) are the teacher and student networks up to their hint and guided layers. Training runs in two stages: first fit the student's lower half and the regressor to the hint, then train the whole student with output distillation. On CIFAR-10 a student with about 10.4 times fewer parameters outperformed its teacher. The authors warned that hints are a regulariser, and that guiding a deep student layer too tightly over-regularises it.

Attention maps remove the width problem

Zagoruyko and Komodakis collapsed the channel dimension instead of regressing it (Zagoruyko and Komodakis, 2017, Paying More Attention to Attention, arXiv:1612.03928). For a convolutional activation \(A \in \mathbb{R}^{C \times H \times W}\) the spatial attention map is \(F^p_{\text{sum}}(A) = \sum_{i=1}^{C}|A_i|^p\), a single \(H\times W\) image of where the layer is active. With \(Q = \mathrm{vec}(F(A))\) for paired student and teacher layers \(j\),

with \(p = 2\) in their experiments. Student and teacher can have any channel counts as long as the spatial resolutions line up, and no regressor parameters are added. The authors stress that this normalisation is important for the student to train successfully.

Transformers: layer maps and relation matrices

TinyBERT carried layer-to-layer transfer into BERT (Jiao et al., 2020, TinyBERT: Distilling BERT for Natural Language Understanding, arXiv:1909.10351). A 4-layer student with hidden size 312 learns from a 12-layer BERT-base with hidden size 768 through the map \(g(m) = 3m\), so student layer \(m\) copies teacher layer \(3m\). Each mapped layer contributes an MSE on attention matrices and an MSE on hidden states through a learned \(W_h \in \mathbb{R}^{312 \times 768}\), which is 239,616 extra parameters per layer. Embeddings and prediction logits are matched too, and distillation runs twice: a general stage on unlabelled text and a task-specific stage. The 14.5M-parameter result kept more than 96.8% of the teacher's GLUE score while being 7.5 times smaller and 9.4 times faster. Its own ablation showed the layer map is a real choice: taking the top teacher layers beat the bottom ones on MNLI and lost on MRPC and CoLA.

MiniLM dropped both the layer map and the projection (Wang et al., 2020, MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers, arXiv:2002.10957). It distils only the teacher's last layer, using KL divergence on the attention distributions and on a value relation \(\mathrm{softmax}(VV^\top/\sqrt{d_k})\). Both are \(|x| \times |x|\) matrices per head whatever the hidden size, so the student's width is unconstrained. A 6-layer, 768-wide student retained more than 99% accuracy on SQuAD 2.0 and several GLUE tasks with half the teacher's Transformer parameters and compute. MiniLMv2 extended the relations to queries and keys and removed the requirement that head counts match.

The distinction worth keeping is between matching features, which needs a dimension alignment such as a regressor or projection, and matching relations, which compares token-by-token or position-by-position structure that is the same size in both models.

When it breaks

Feature transfer can underperform plain distillation. Tian, Krishnan and Isola benchmarked 12 methods on CIFAR-100 and found vanilla output distillation second best overall (Tian et al., 2020, Contrastive Representation Distillation, arXiv:1910.10699). Across architecture families the picture got worse: distilling VGG-13 into MobileNetV2, the student alone reached 64.6%, output distillation 67.37%, FitNets 64.14% and attention transfer 59.40%. A plausible reading is that forcing a mobile network to reproduce a VGG's intermediate maps imposes a computation it was not built for.

The connector can cheat. A regressor or projection with enough capacity can reduce the hint loss without the student's own representation improving much, since the connector absorbs the mismatch. Parameter-free relation targets like MiniLM's avoid that failure by construction.

Layer correspondence is a hyperparameter with no principled default. Uniform, top and bottom maps favour different tasks, and when depths differ by a large ratio the choice is guesswork validated by ablation.

It needs the teacher's internals at training time. Output distillation can use cached logits. Feature and attention transfer need a teacher forward pass with activations exposed for every batch and multiplies memory when the teacher is large. A teacher reachable only through an API exposes outputs, so feature and attention transfer are simply unavailable for it.