Source-linked AI summary
Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, Ivan Titov
TL;DR
The paper asks how individual encoder heads affect translation and what roles they play. It uses relevance analysis and stochastic-gate pruning with a differentiable L0 relaxation. Specialized heads are pruned last, while most heads can be removed without seriously affecting performance.
Problem
The paper investigates how translation quality depends on individual encoder heads and whether their roles are consistent and interpretable.
Method
The paper combines layer-wise relevance propagation with stochastic-gate head pruning using a differentiable relaxation of the L0 penalty.
Results
Specialized heads are pruned last, and the vast majority of heads, especially encoder self-attention heads, can be removed without seriously affecting performance.
Takeaways & Limitations
Important heads tend to have specialized, interpretable functions such as attending to adjacent words or tracking syntactic relations.
Takeaways & Limitations
The results do not establish whether a model trained from scratch with the same small number of heads can match the performance of a pruned model.
Abstract
from arXiv · showhide
Multi-head self-attention is a key component of the Transformer, a state-of-the-art architecture for neural machine translation. In this work we evaluate the contribution made by individual attention heads in the encoder to the overall performance of the model and analyze the roles played by them. We find that the most important and confident heads play consistent and often linguistically-interpretable roles. When pruning heads using a method based on stochastic gates and a differentiable relaxation of the L0 penalty, we observe that specialized heads are last to be pruned. Our novel pruning method removes the vast majority of heads without seriously affecting performance. For example, on the English-Russian WMT dataset, pruning 38 out of 48 encoder heads results in a drop of only 0.15 BLEU.
1 Introduction
The paper examines how individual encoder attention heads contribute to translation quality and whether their roles are consistent and interpretable. It finds that a small subset of specialized heads is especially important, while many heads can be pruned.
- Previous analyses averaged attention across heads or focused on maximum weights, obscuring differences in head importance and function.
- The study asks how much translation quality depends on individual encoder heads and whether many heads can be removed while preserving quality.
- Layer-wise relevance propagation identifies important heads whose functions include positional, syntactic, and rare-word attention.
- The paper introduces head pruning based on stochastic gates and a differentiable relaxation of the L0 penalty to test whether remaining heads are redundant.
- Important heads are specialized and interpretable, attending to neighbouring words or tokens in specific syntactic dependency relations.
2 Transformer Architecture
The Transformer is an encoder-decoder architecture built from stacked attention and feed-forward layers. Multi-head attention forms multiple representations, processes them separately, and combines their outputs.
- The Transformer uses stacked self-attention and fully connected layers in both its encoder and decoder.
- Each encoder layer contains a multi-head self-attention sub-layer and a feed-forward network.
- Multi-head attention creates h representations of queries, keys, and values, computes attention for each, concatenates the results, and projects them.
- Each decoder layer adds decoder-encoder attention, which receives the encoder output as its keys and values.
- The paper focuses primarily on encoder self-attention among the Transformer’s three uses of multi-head attention.
3 Data and setting
The experiments use English as the source language and Russian, German, and French as targets, controlling training-data size and model parameter counts. Additional English-Russian experiments examine domain effects.
- The study evaluates English-to-Russian, English-to-German, and English-to-French translation.
- Each language pair uses 2.5m WMT sentence pairs and Transformer models with the same numbers of parameters.
- English-Russian experiments additionally use OpenSubtitles2018 to evaluate the impact of domain.
- The English-Russian analysis concentrates on the WMT and OpenSubtitles domains.
4 Identifying Important Heads
The analysis ranks encoder attention heads by LRP relevance and compares that ranking with head confidence. Relevance generally agrees with confidence, while a small number of heads in each layer stand out as especially important.
- Identifying important heads: Head confidence is the average maximum attention weight over evaluation-set tokens, excluding the end-of-sentence symbol.
- Identifying important heads: LRP evaluates how different attention heads contribute to the model’s top-1 predicted logit.Heads with higher relevance values are judged more important to the model’s predictions.
- Identifying important heads: In each layer, LRP ranks a small number of heads as much more important than the others.
- Identifying important heads: LRP relevance agrees reasonably with head confidence, except for the most relevant head in the first layer, whose average maximum attention weight is low.
5 Characterizing heads
The analysis identifies positional, syntactic, and rare-word functions among important encoder heads, using attention patterns and dependency-based comparisons. These functions are associated with heads that are confident or highly relevant, and specialized functions persist as heads are pruned.
- The analysis identifies three head functions: positional attention to adjacent tokens, syntactic attention to specific dependency relations, and attention to rare words.These categories are used to characterize heads ranked highly by layer-wise relevance propagation.
- Positional heads: Positional heads assign their maximum attention weight to a specific relative position at least 90% of the time, usually an adjacent token.The criterion uses relative position -1 or +1.
- Positional heads: Every positional head across the considered language pairs has an average maximum attention weight above 0.8 and corresponds substantially to highly confident and important heads.
- Syntactic heads: Syntactic heads are evaluated by comparing their maximum-attention tokens with dependency structures and are defined as at least 10% more accurate than a positional baseline.The analysis considers nominal subjects, direct objects, adjectival modifiers, and adverbial modifiers.
- Syntactic heads: Certain heads detect syntactic relations more accurately than the positional baseline, supporting some syntactic disambiguation in the encoder; multiple heads can share a relation.
- Rare-word heads: One first-layer head is much more important than the others in that layer across WMT and OpenSubtitles models, and it attends to the least frequent sentence tokens.For OpenSubtitles, it selects the rarest token in 66% of qualifying cases and one of the two least frequent tokens in 83%; on WMT, the latter exceeds 50%.
6 Pruning Attention Heads
The paper introduces stochastic Hard Concrete gates to prune entire attention heads and evaluates how translation quality and head functions change as heads are removed. Specialized, interpretable functions tend to survive pruning, while many heads can be removed with little BLEU loss.
- 6.1 Method: Hard Concrete gates stochastically relax L0 regularization so entire attention heads can be switched off during optimization.Each head representation is multiplied by a head-specific gate, and the regularized objective combines translation cross-entropy with a gate penalty.
- 6.1 Method: Varying the regularization coefficient produces models with different numbers of retained heads after fine-tuning a converged model.The decoder is fixed while encoder parameters and gates are fine-tuned, preventing pruned encoder functions from migrating into the decoder.
- 6.2.1 Quantitative results: BLEU score: 0.25 BLEU is lost on OpenSubtitles when pruning all but 4 of 48 encoder heads, while 10 encoder heads remain within 0.15 BLEU of the full WMT model.These results show that encoder translation performance remains effective with only a few retained heads.
- 6.2.2 Functions of retained heads: A 17-head model retains all previously identified functions even after pruning 2⁄3 of the heads.With fewer heads, some functions drift to other heads, including positional heads that begin tracking syntactic dependencies.
- 6.3.1 Quantitative results: BLEU score: More than half of all attention heads can be pruned with a loss of no more than 0.25 BLEU, including substantial pruning in encoder and decoder attention.For WMT, almost 3⁄4 of encoder heads and more than 1⁄3 of decoder self-attention and decoder-encoder heads can be pruned without noticeable loss.
- 6.3.2 Heads importance: The model prunes encoder self-attention heads first, retains decoder-encoder heads most strongly, and distributes decoder functions across layers.Decoder self-attention importance varies with domain and sentence length; decoder self-attention is retained lower, while decoder-encoder attention is retained higher.
7 Related work
Prior work analyzes NMT representations, grammatical sensitivity, and attention weights, while this paper focuses on the importance and specialization of entire attention heads. It relates its findings to target-language effects and distinguishes its head-level analysis from prior syntactic evaluations.
- Representation analysis: NMT representation analyses have examined morphology, syntax, and semantics using probing-style linguistic tasks.These studies evaluate how informative learned representations are for different levels of linguistic analysis.
- Target-language effects: Target language influences which information is encoded, consistent with heads tracking syntactic relations requiring complex target-language agreement.The paper reports this connection for English-Russian translation across domains.
- Syntactic analysis: Grammatical-error sensitivity studies show that NMT models learn some syntactic structures, whereas this work provides further insight into multi-head attention’s role.The comparison concerns the analytical focus rather than a replacement of grammatical-sensitivity methods.
- Attention analysis: Prior attention-weight analyses include inducing encoder trees, but do not evaluate specific labeled dependency relations or head specialization.This paper addresses those distinctions by examining specialized attention heads.
- Head-level analysis: Unlike neuron-level importance analysis, this work identifies and verifies the importance of entire attention heads using relevance propagation and pruning.The verification criterion is which heads remain when the model is sparsified.
8 Conclusions
The paper evaluates individual encoder-head contributions with LRP and finds that a small subset has interpretable functions, while most heads can be pruned with limited performance impact.
- Head analysis: Important heads exhibit interpretable positional, syntactic, and rare-word attention functions.Examples include attending to adjacent tokens, tracking dependency relations, and pointing to the least frequent tokens.
- Pruning: Specialized heads are pruned last, while the vast majority of heads—especially encoder self-attention heads—can be removed without seriously affecting performance.This directly supports the importance of specialized heads and the redundancy of many others.
- LRP method: LRP propagates relevance backward through the network using a conservation principle rather than gradients and stops at the encoder layer of interest.The total contribution is treated as constant across layers during redistribution.
- Head analysis: LRP estimates each head’s relative contribution to the model’s top-1 predicted logit and averages relevance across generation steps.Head relevance is summed over its neurons and normalized over heads within a layer.
B.1 Data preprocessing
The experiments use byte-pair encoded translation data with approximately 32,000-token source and target vocabularies, following the Transformer-base architecture with six encoder and decoder layers.
- Data preprocessing: Sentences are encoded with byte-pair encoding and source and target vocabularies of about 32000 tokens.OpenSubtitles pairs additionally require at least 0.9 relative subtitle-frame time overlap to reduce noise.
- Data preprocessing: Training batches group translation pairs by approximate sequence length and contain approximately 1600010 source tokens.
- Model setup: The model follows Transformer base with N = 6 encoder layers, N = 6 decoder layers, and h = 8 attention heads.Its model dimension is dmodel = 512 and feed-forward dimension is dff = 2048.
B.3 Optimizer
The optimizer follows the Transformer-base training setup, using Adam with specified hyperparameters and a learning rate that varies during training.
- Optimizer: Adam is used with β1 = 0.9, β2 = 0.98, and ε = 10−9.
- Optimizer: The learning rate varies over the course of training according to the stated formula.
- Optimizer: Gradient accumulation or several GPUs can be used to reach the approximately 10-million-token training-batch target.