Structured Pruning › Head Pruning

Transformer attention layers contain multiple parallel attention heads. Each head applies a separate set of queries, keys, and values to the input, then projects the results back to the model dimension. The motivation for multiple heads is that different heads can attend to different aspects of the input: one head might track syntactic subject-verb dependencies while another tracks positional relationships, and another might focus on coreference.

The multi-head attention computation concatenates the outputs of all heads and applies a final linear projection. The full attention output for a layer is:

$$\text{MultiHead}(\mathbf{X}) = \text{Concat}(\text{head}_1, \text{head}_2, \ldots, \text{head}_H) \mathbf{W}_O$$

where each head computes:

$$\text{head}_h = \text{softmax}\!\left(\frac{\mathbf{X}\mathbf{W}_Q^{(h)} (\mathbf{X}\mathbf{W}_K^{(h)})^\top}{\sqrt{d_k}}\right) \mathbf{X}\mathbf{W}_V^{(h)}$$

where:

• $H$ is the number of heads
• $d_k$ is the head dimension
• $\mathbf{W}_Q^{(h)}, \mathbf{W}_K^{(h)}, \mathbf{W}_V^{(h)} \in \mathbb{R}^{d_{\text{model}} \times d_k}$ are the query, key, and value projections for head $h$
• $\mathbf{W}_O \in \mathbb{R}^{(H \cdot d_k) \times d_{\text{model}}}$ is the output projection that assembles all heads

Because each head's contribution is additive after the output projection, you can set any head's output to zero and the remaining heads still produce a valid (though potentially degraded) result. This additive structure is precisely what makes head pruning possible: it is a well-defined operation with predictable effects.

The most direct importance score for a head is its sensitivity to removal. You can estimate this by computing how much the loss increases when you zero out that head's output. Formally, for head $h$ in layer $l$, define the importance score as:

$$I_h^{(l)} = \left| \mathcal{L}(\theta) - \mathcal{L}(\theta \setminus h_l) \right|$$

where:

• $\mathcal{L}(\theta)$ is the loss with all heads present
• $\mathcal{L}(\theta \setminus h_l)$ is the loss with head $h$ in layer $l$ zeroed out
• $|\cdot|$ is the absolute value, giving a non-negative importance score

Computing this exactly requires a separate forward pass for each head, which is expensive. A first-order approximation uses gradient information instead. For a head parameterized by its output projection weights $\mathbf{W}_O^{(h,l)}$, the gradient-based importance score is:

$$I_h^{(l)} = \left\| \frac{\partial \mathcal{L}}{\partial \mathbf{W}_O^{(h,l)}} \odot \mathbf{W}_O^{(h,l)} \right\|_F$$

where:

• $\frac{\partial \mathcal{L}}{\partial \mathbf{W}_O^{(h,l)}}$ is the gradient of the loss with respect to the head's output projection weights
• $\odot$ is element-wise multiplication
• $\|\cdot\|_F$ is the Frobenius norm, giving the overall magnitude of the score

This gradient-weight product, sometimes called the Taylor expansion criterion, approximates the change in loss from zeroing out the head. Intuitively, a head with large weights and large gradients is one where the model is heavily relying on that head and adjusting it actively. Large values indicate that the weights are large and their gradients are large, meaning the head contributes substantially to the loss. Small values indicate the head could be removed with little effect. This approximation is derived from a first-order Taylor expansion of the loss around the current parameters, treating the removal of the head as a perturbation.

A simpler but surprisingly effective approach uses attention head activation norms. For each head, compute the mean magnitude of its output across a calibration dataset:

$$I_h^{(l)} = \mathbb{E}_{x \sim \mathcal{D}} \left[ \left\| \text{head}_h^{(l)}(x) \right\|_2 \right]$$

where $\text{head}_h^{(l)}(x)$ is the output vector produced by head $h$ in layer $l$ for input $x$, before it is passed through the output projection.

Activation norm-based importance has the advantage that it requires no gradient computation and can be collected during normal inference. The disadvantage is that it measures how active a head is, not how important it is to the final output. A very active head whose output is largely cancelled by the output projection could score high by this criterion yet contribute little to the model's predictions. In practice, gradient-based scores tend to be more accurate, but activation norms are a useful fast approximation.

Once importance scores are computed, pruning proceeds by selecting the heads to remove and modifying the weight matrices accordingly. An attention layer in a transformer has query, key, value, and output projection matrices. With $H$ heads each of dimension $d_k$:

$$\mathbf{Q} \in \mathbb{R}^{d_{\text{model}} \times (H \cdot d_k)}, \quad \mathbf{K} \in \mathbb{R}^{d_{\text{model}} \times (H \cdot d_k)}, \quad \mathbf{V} \in \mathbb{R}^{d_{\text{model}} \times (H \cdot d_k)}$$

Each head $h$ uses columns $(h-1) \cdot d_k$ through $h \cdot d_k$ of these matrices. The output projection $\mathbf{W}_O \in \mathbb{R}^{(H \cdot d_k) \times d_{\text{model}}}$ assembles the contributions from all heads. To remove head $h$, you delete the corresponding columns from $\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$, and the corresponding rows from $\mathbf{W}_O$.

After removing $k$ heads, the model dimension of the attention projections shrinks from $H \cdot d_k$ to $(H - k) \cdot d_k$. This directly reduces the memory footprint of the attention layer and the number of FLOPs in every forward pass. The change is structural and hardware-visible: the new, smaller weight matrices are stored in contiguous memory, and standard matrix multiplication routines operate on the smaller tensors with no overhead.

Pruning decisions can be made at the layer level (removing a fixed number of heads from each layer independently) or globally (ranking all heads across all layers together and removing the globally least important ones). Global pruning is generally preferable because it finds that some layers have more redundant heads than others. In BERT, for example, researchers consistently found that middle layers had more redundant heads than early or late layers. A layer-level strategy that forces uniform pruning across layers would retain unnecessary heads in middle layers while potentially over-pruning early or late layers.

Removing heads causes an immediate quality drop. The model was trained with all heads present, and the surviving heads were not optimized to compensate for the removed ones. A brief fine-tuning pass, sometimes called recovery fine-tuning, allows the remaining heads to adapt. In practice, this recovery phase needs far fewer steps than the original training because you are not learning from scratch but adjusting weights that already encode useful representations.

The amount of quality you can recover depends on how aggressively you pruned. Removing 10-20% of heads typically recovers fully within a small number of fine-tuning steps. Removing 50% or more often results in permanent quality degradation, though the model may still be useful for applications that tolerate some accuracy loss in exchange for speed. The recovery fine-tuning learning rate should be substantially lower than the original training learning rate to avoid disrupting the representations that the surviving heads have already learned.

One practical consideration is the order of operations. If you plan to prune both layers and heads, the recommended order is: compute layer importance and remove layers first, then compute head importance on the reduced model and remove heads, then fine-tune. This order matters because removing layers changes the distribution of attention patterns, which can shift head importance scores. Computing head importance on the original model and then removing both layers and heads simultaneously tends to produce worse results.