Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned

By Elena Voita, David Talbot, F. Moiseev, Rico Sennrich, Ivan Titov. 2019

1,560 citations recorded by the source corpus.

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 to the overall performance of the model and analyze the roles played by them in the encoder. 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.

Subfields: Artificial Intelligence, Information Systems, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Information Systems

In Model Architectures

Citations in the corpus

Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers (2020, background): “It also may be possible to remove entire attention heads in addition to zeroing out weights (Michel et al., 2019; Voita et al., 2019).”
TinyBERT: Distilling BERT for Natural Language Understanding (2019, background): “Recent studies (Kovaleva et al., 2019; Michel et al., 2019; Voita et al., 2019) demonstrate that there is redundancy in PLMs.”
The Bottom-up Evolution of Representations in the Transformer: A Study with Machine Translation and Language Modeling Objectives (2019, background): “In this work, we study how the choice of ob-jective affects the process by which information is encoded in token representations of the Trans-former (Vaswani et al., 2017), as this architecture achieves state-of-the-art results on tasks with very different objectives such as machine translation (MT)”
VisualBERT: A Simple and Performant Baseline for Vision and Language (2019, background): “Syntactic Grounding Given that many have observed that the attention heads of BERT can discover syntactic relationships (Voita et al., 2019; Clark et al., 2019), we also analyze how grounding information is passed through syntactic relationships that VisualBERT may have discovered.”
Data-Efficient Structured Pruning via Submodular Optimization (2022, methodology): “Other methods employ sparsity and group-sparsity regularisation during training, to induce sparse weights [Collins and Kohli, 2014, Voita et al., 2019].”
Foundation Posteriors for Approximate Probabilistic Inference (2022, neutral): “Prior works [6, 63, 1, 14] using transformer networks in natural language have re-purposed the attention weights in the later layers as an mechanism to introspect model logic.”
Reducing Transformer Depth on Demand with Structured Dropout (2019, background): “Similarly, various papers have shown evidence that Transformers are over-parameterized, especially that most self-attention heads can be dropped at test time (Michel et al., 2019; Voita et al., 2019).”
Knowledge Neurons in Pretrained Transformers (2021, methodology): “It is motivated by FFNs’s analogy with the self-attention mechanism (as described in Section 2), because self-attention scores are usually used as a strong attribution baseline (Kovaleva et al., 2019; Voita et al., 2019; Hao et al., 2021).”
Transformer Feed-Forward Layers Are Key-Value Memories (2020, background): “The study of the transformer architecture has focused on the role and function of self-attention layers (Voita et al., 2019; Clark et al., 2019; Vig and Belinkov, 2019) and on inter-layer differences (i.e. lower vs. upper layers) (Tenney et al., 2019; Jawahar et al., 2019).”
The Right Tool for the Job: Matching Model and Instance Complexities (2020, background): “More sophisticated approaches induce structured sparsity, which removes full blocks (Michel et al., 2019; Voita et al., 2019; Dodge et al., 2019b).”
What Does BERT Look at? An Analysis of BERT’s Attention (2019, background): “Concurrently with our work Voita et al. (2019) identify syntactic, positional, and rare-word-sensitive attention heads in machine translation models.”
Are Sixteen Heads Really Better than One? (2019, background): “Concurrently to our own work, Voita et al. (2019) have made to a similar observation on multi-head attention.”
Efficient Transformers: A Survey (2020, background): “An example is to prune attention heads during inference (Voita et al., 2019; Michel et al., 2019).”
A Primer in BERTology: What We Know About How BERT Works (2020, background): “To mitigate that, Michel et al. (2019) prune heads in the order set by a proxy importance score, and Voita et al. (2019) fine-tune the pretrained Transformer with a regularized objective that has the head-disabling effect.”
BERTology Meets Biology: Interpreting Attention in Protein Language Models (2020, background): “It has been suggested that attention weights represent a model’s conﬁdence in detecting certain features [17, 85].”