Distillation Variants › Attention Transfer

Attention mechanisms, covered in depth in the self-attention chapters of this book, produce one of the most interpretable outputs in a transformer: a distribution over positions indicating which tokens the model attends to when processing each token. These attention patterns contain structural knowledge that goes beyond the final logit distribution.

Attention transfer teaches the student to replicate the teacher's attention distributions, along with its output probabilities. The mechanism is direct: if the teacher has learned that the verb "is" should attend strongly to its grammatical subject, that syntactic relationship is implicit in the attention weights. Transferring it gives the student a scaffold for building correct representations.

This is distinct from feature distillation in an important way. Feature distillation constrains the student's hidden representations, which are the outputs of entire transformer blocks after multi-head attention, residual connections, and feed-forward layers. Attention transfer constrains the attention weights themselves, which are computed inside the multi-head attention operation before the subsequent processing stages. Attention weights reflect a more focused kind of structural knowledge: specifically, the token-to-token relevance relationships that the model has learned to track.

For a given layer $l$ and attention head $h$, let $A^T_{l,h} \in \mathbb{R}^{n \times n}$ be the teacher's attention weight matrix over a sequence of length $n$, and $A^S_{l,h} \in \mathbb{R}^{n \times n}$ be the student's counterpart. The standard formulation minimizes the mean squared error:

$$\mathcal{L}_{\text{attn}} = \frac{1}{L \cdot H} \sum_{l=1}^{L} \sum_{h=1}^{H} \| A^T_{l,h} - A^S_{l,h} \|_F^2$$

where:

• $L$: number of layers with attention transfer (usually all layers or a subset)
• $H$: number of attention heads per layer
• $A^T_{l,h}, A^S_{l,h}$: teacher and student attention matrices at layer $l$, head $h$
• $\| \cdot \|_F^2$: squared Frobenius norm, summing squared differences over all $n^2$ position pairs

The Frobenius norm sums across all $n \times n$ entries of the attention matrix. This means the loss penalizes discrepancies at every (query, key) position pair, including positions outside the attention maximum. A student that broadly approximates the teacher's attention pattern but misaligns on a few critical structural relationships will receive gradient pressure to correct those misalignments at every training step.

When the teacher and student have different numbers of heads, you have two options. You can aggregate across heads first by averaging into a single per-layer attention matrix:

$$\bar{A}_l = \frac{1}{H} \sum_{h=1}^{H} A_{l,h}$$

and then align $\bar{A}^T_l$ to $\bar{A}^S_l$. Alternatively, you can match a subset of student heads to specific teacher heads, which requires a head correspondence mapping. Averaging is simpler and more stable in practice: individual heads often specialize differently across runs, so one-to-one head correspondence can be brittle. The averaged matrix captures the collective structural intent of all heads, which is the more stable transfer target.

A practical note on normalization: attention weights are already normalized to sum to 1.0 within each row (they are output of softmax over query-key dot products). So unlike hidden state features, there is no need for an explicit normalization step in the attention transfer loss. The MSE directly compares probability distributions over keys for each query position.

Consider a 4-token sequence: "The dog barks." Suppose the teacher at layer 0 has learned the following averaged attention pattern (each row sums to 1.0):

The teacher has learned that "barks" attends strongly to "dog" (0.65), capturing the grammatical subject-verb relationship. The student, early in training, might produce a more uniform pattern:

The Frobenius norm between these two matrices captures the structural gap. The attention transfer loss provides a gradient that pushes the student toward concentrating attention on "dog" when processing "barks," effectively teaching it the subject-verb relationship that the teacher has already discovered. This is knowledge that might take the student many more training steps to independently acquire from the task loss alone.

Attention transfer provides the most benefit for tasks that rely on precise structural relationships: parsing, coreference resolution, question answering requiring multi-hop reasoning, and translation. For these tasks, the teacher has developed clear attention patterns that encode syntactic and semantic structure. In contrast, for tasks like sentence classification where the answer is largely determined by a few key tokens and the overall theme, attention transfer adds less over standard output-level distillation.

One notable finding from TinyBERT is that attention transfer is particularly useful in the lower layers, where the teacher's heads specialize for syntactic dependencies. Upper-layer attention patterns are more diffuse and task-specific, making them harder for the student to replicate exactly. This aligns with the general observation about layer specialization in transformers: lower layers handle structural parsing, upper layers handle semantic aggregation. The structural parsing patterns are more consistent and transferable; the semantic aggregation patterns depend heavily on the specific task and are harder to match precisely.

Attention transfer incurs substantial computational cost. Computing $\mathcal{L}_{\text{attn}}$ requires storing and comparing all attention matrices at all layers, in addition to the final output. For a sequence of length 512, each attention matrix is $512 \times 512 = 262{,}144$ entries. For 12 layers and 12 heads, that is roughly 37 million entries per example. Compared to the vocabulary-size soft-label comparison (which involves maybe 50,000 entries per position), the memory and compute cost of attention transfer is substantially higher. This reinforces the practice of applying attention transfer selectively to a subset of layers rather than all of them.

An advanced variant uses not raw per-layer attention maps but attention rollout, which propagates attention through the network to estimate the effective attention from input tokens to each layer's output. This captures how information flows across multiple layers rather than within a single layer, giving a transfer target that integrates information across layers.

The motivation for rollout is that raw per-layer attention weights can be misleading. A head in layer 5 might attend strongly to position 3, but position 3's representation in layer 5 is already a mixture of all input positions due to the attention layers below it. The effective influence of token 3 at layer 5 is determined by more than the direct attention weight to position 3 at that layer; it is the product of the attention paths that connect the original input token 3 through all intermediate layers. Attention rollout computes this propagated influence by taking the matrix product of attention matrices across layers, accounting for the residual connections that also pass information forward directly.

Transferring attention rollout rather than raw attention weights gives the student a more accurate picture of the teacher's true information routing, but at higher computational cost: you must compute the matrix products across layers for every example in every training batch.