PaperlayerView exact revision on arXiv

Source-linked AI summary

Are Sixteen Heads Really Better than One?

Paul Michel, Omer Levy, Graham Neubig

arXiv:1905.10650v3cs.CL

TL;DR

The paper investigates what multiple attention heads contribute in Transformer MT and BERT, where their practical necessity remains unclear. It evaluates post-training head ablations and greedy pruning, finding that many heads can be removed with little performance loss while improving inference efficiency. The analyses also identify component-specific sensitivity and training-time differentiation of head importance.

  • Problem

    The paper asks what multiple attention heads contribute in Transformer-based MT and BERT models, given that their practical value beyond vanilla attention is unclear.

  • Method

    The authors ablate individual heads and entire layers, develop greedy iterative pruning, and analyze head sensitivity across components and training stages.

  • Results

    Most heads can be removed after training with little test-performance impact, some layers can use one head, and pruning yields up to a 17.5% inference-speed increase for BERT.

  • Takeaways & Limitations

    Head redundancy can support post-training model compression, although encoder-decoder attention and many network heads remain important for preserving performance.

  • Takeaways & Limitations

    The ablation results were obtained on specific, rather small test sets, leaving their generalizability to other datasets uncertain.

Abstract

from arXiv · show

Attention is a powerful and ubiquitous mechanism for allowing neural models to focus on particular salient pieces of information by taking their weighted average when making predictions. In particular, multi-headed attention is a driving force behind many recent state-of-the-art NLP models such as Transformer-based MT models and BERT. These models apply multiple attention mechanisms in parallel, with each attention "head" potentially focusing on different parts of the input, which makes it possible to express sophisticated functions beyond the simple weighted average. In this paper we make the surprising observation that even if models have been trained using multiple heads, in practice, a large percentage of attention heads can be removed at test time without significantly impacting performance. In fact, some layers can even be reduced to a single head. We further examine greedy algorithms for pruning down models, and the potential speed, memory efficiency, and accuracy improvements obtainable therefrom. Finally, we analyze the results with respect to which parts of the model are more reliant on having multiple heads, and provide precursory evidence that training dynamics play a role in the gains provided by multi-head attention.

1 Introduction

The paper asks what multiple attention heads contribute and finds that many can be removed after training with little performance cost. It further examines pruning, efficiency, component sensitivity, and training dynamics.

  • 1 Introduction: Transformers use multi-headed attention, computing attention independently through Nh parallel heads.This mechanism is central to Transformer performance across machine translation, question answering, classification, and semantic role labeling.
  • 1 Introduction: Most attention heads can be individually removed after training without significant test-performance loss in Transformer MT and BERT natural language inference.Many attention layers can also be reduced to a single head without affecting test performance.
  • 1 Introduction: A greedy pruning algorithm removes apparently less-contributing heads iteratively, finding that large network portions can be removed while most heads must remain to avoid catastrophic performance drops.The pruning is performed jointly across the network rather than being restricted to one layer.
  • 1 Introduction: 17.5% increase in inference speed is obtained for a BERT-based model after pruning attention heads.The paper reports this as an inference-time efficiency benefit of the pruning approach.
  • 1 Introduction: Encoder-decoder attention layers in machine translation are more sensitive to pruning than self-attention layers.The authors interpret this pattern as suggesting a more critical role for multi-headedness in encoder-decoder attention.
  • 1 Introduction: The distinction between important and unimportant heads increases as training progresses, providing evidence of an interaction between multi-headedness and training dynamics.The paper reports this as an analysis result rather than a complete explanation of the mechanism.

2 Background: Attention, Multi-headed Attention, and Masking

The background defines vanilla and multi-headed attention, then introduces masking as the mechanism for removing heads during ablation experiments.

  • 2 Background: Attention, Multi-headed Attention, and Masking: Vanilla scaled bilinear attention maps a sequence of n d-dimensional vectors and a query vector to a weighted sum parameterized by Wk, Wq, Wv, and Wo.The attention weights are computed from query-key interactions and used to weight value vectors.
  • 2 Background: Attention, Multi-headed Attention, and Masking: In self-attention, each input vector serves as the query, whereas sequence-to-sequence attention typically uses a decoder state as the query and encoder outputs as inputs.This distinguishes the inputs used by the two attention settings.
  • 2 Background: Attention, Multi-headed Attention, and Masking: Multi-headed attention applies Nh independently parameterized attention layers in parallel and combines their outputs.With dh = d, it is more expressive than vanilla attention; with dh = d/Nh, it acts as an ensemble of low-rank vanilla attention layers.
  • 2 Background: Attention, Multi-headed Attention, and Masking: A nonlinear feed-forward network is applied to the multi-headed attention output at each Transformer layer, allowing different heads to interact.This interaction occurs after the parallel attention computations.
  • 2 Background: Attention, Multi-headed Attention, and Masking: A mask variable ξh is assigned to each head, with ξh = 0 removing head h and all ξh = 1 recovering the unmasked formulation.The modified multi-headed attention formula enables test-time head ablation.

3 Are All Attention Heads Important?

The paper tests whether attention heads remain necessary after training in WMT and BERT models. It finds that most individual heads can be removed, and many layers can retain only one head without significant test-performance loss, though some layers are more sensitive and results have dataset-scope limitations.

  • Experimental setup: The experiments remove individual heads or all but one head at test time in trained WMT and BERT models.WMT is a 6-layer, 16-head-per-layer English-to-French Transformer; BERT is a 12-layer, 12-head-per-layer model evaluated on MultiNLI.
  • Ablating one head: Most attention heads can be removed without substantially changing the original BLEU or accuracy, although some removals increase performance.Head importance is assessed by masking a head and comparing the resulting score with the full model.
  • Ablating one head: Only 8 of 96 WMT encoder self-attention heads cause statistically significant performance changes when removed, and half of those removals increase BLEU.This supports the observation that most heads are redundant given the rest of the model.
  • Ablating all heads but one: For most layers, one head is sufficient at test time, reducing attention parameters to 1/16th or 1/12th of a vanilla layer, but some layers require multiple heads.Keeping only one head can severely damage performance in specific layers, including at least a 13.5 BLEU-point drop in the last WMT encoder-decoder attention layer.
  • Ablating all heads but one: Selecting the best head using validation data preserves the finding on held-out evaluation: one head causes no statistically significant performance change in 50% of WMT layers and 100% of BERT layers.The selection uses newstest2013 for WMT and a 5,000-example MNLI training subset for BERT, then evaluates on newstest2014 and MNLI-matched validation data.
  • Are important heads the same across datasets?: The initial ablation results are limited by evaluation on specific, relatively small test sets, motivating tests on out-of-domain datasets.The paper evaluates MNLI mismatched and MTNT English-to-French as contrastive out-of-domain suites and finds positive correlations above 0.5 between head effects across domains.

4 Iterative Pruning of Attention Heads

The paper greedily ranks attention heads by an importance score and prunes them iteratively, finding that substantial pruning can preserve performance while improving efficiency. Pruning up to 20% of WMT heads and 40% of BERT heads has no noticeable negative impact, but further pruning causes sharp degradation.

  • 4.1 Head Importance Score for Pruning: The pruning procedure sorts all attention heads by a proxy importance score and removes them one by one to avoid impractical combinatorial search.The score is estimated from expected sensitivity to mask variables and can be computed with a forward and backward pass.
  • 4.1 Head Importance Score for Pruning: Using Ih to order pruning is faster and yields better results than ordering heads by individual score differences.Figures 3a and 3b incrementally remove 10% of heads at each step for WMT and BERT.
  • 4.2 Effect of Pruning on BLEU/Accuracy: 20% of WMT heads and 40% of BERT heads can be pruned without noticeable negative impact on performance.Performance drops sharply when pruning further, and neither model can become purely single-headed without retraining or substantial losses.
  • 4.3 Effect of Pruning on Efficiency: Each attention head accounts for 6.25% of WMT-layer parameters or approximately 8.34% of BERT-layer parameters, making pruning relevant for memory-constrained deployment.Across both models, roughly one third of total parameters is devoted to multi-headed attention.
  • 4.3 Effect of Pruning on Efficiency: Pruning, rather than merely masking, half of BERT’s attention heads increases inference speed by up to approximately 17.5% at higher batch sizes.The reported speedup vanishes for smaller batch sizes.

5 When Are More Heads Important? The Case of Machine Translation

Machine-translation attention types differ substantially in their reliance on multiple heads. Encoder-decoder attention is especially sensitive to pruning, whereas encoder and decoder self-attention remain usable after much more aggressive reduction.

  • Experimental comparison: The pruning analysis compares Enc-Enc, Enc-Dec, and Dec-Dec attention separately to determine which transformer components depend most on multi-headedness.Figure 4 reports BLEU as heads are incrementally pruned from each attention type.
  • Encoder-decoder attention: Pruning more than 60% of encoder-decoder attention heads causes catastrophic performance degradation.Performance drops much more rapidly for encoder-decoder attention than for the self-attention layers.
  • Self-attention: Encoder and decoder self-attention can still produce reasonable translations with BLEU scores around 30 after retaining only 20% of the original attention heads.This contrasts with the sharper degradation observed in encoder-decoder attention.
  • Conclusion: The results indicate that encoder-decoder attention is much more dependent on multi-headedness than self-attention.The comparison identifies encoder-decoder attention as the most pruning-sensitive component of the translation model.

6 Dynamics of Head Importance during Training

The authors track head pruning throughout training and find that head importance emerges after an initial phase rather than immediately. From epoch 10 onward, some heads can be pruned while retaining most of the original BLEU score.

  • 6 Dynamics of Head Importance during Training: The study repeats incremental pruning at every epoch on a smaller IWSLT German-to-English model.The model has 6 layers and 8 heads per layer and is trained on IWSLT 2014.
  • 6 Dynamics of Head Importance during Training: In epochs 1–2, performance decreases linearly with pruning, indicating that heads are initially similarly important.The relative performance decrease is independent of the number of heads removed.
  • 6 Dynamics of Head Importance during Training: From epoch 10 onward, up to 40% of heads can be pruned while retaining 85–90% of the original BLEU score.This later regime contains a concentration of unimportant heads.
  • 6 Dynamics of Head Importance during Training: The results suggest that important heads are determined early, but not immediately, during training.The paper leaves a more principled investigation of this phenomenon to future work.

7 Related work

Related work situates this study within attention research and neural-network pruning. It contrasts fine-grained weight pruning, structured pruning, and contemporaneous approaches to identifying and removing important attention heads.

  • 7 Related work: Attention mechanisms originated in neural machine translation and were subsequently adapted to tasks including reading comprehension, natural language inference, and summarization.The cited formulations include Bahdanau et al., Cho et al., and Luong et al.
  • 7 Related work: Neural-network pruning includes fine-grained weight-by-weight methods and structured methods that remove entire model components.The paper places its head-pruning approach within this structured-pruning literature.
  • 7 Related work: Concurrent work identified important heads using layer-wise relevance propagation and linguistic properties such as adjacent positions, rare words, and syntactic relations.That work also proposed gradient descent on attention-head mask variables.

8 Conclusion

The paper concludes that trained multi-head attention often contains removable heads, while some components remain more dependent on multiple heads. It also reports that head importance emerges early in training and motivates more parameter-efficient attention models.

  • 8 Conclusion: Several attention heads can be removed from trained Transformer models without statistically significant test-performance degradation, and some layers can be reduced to one head.This pattern holds across a variety of settings examined in the paper.
  • 8 Conclusion: Machine-translation encoder-decoder attention layers rely more heavily on multi-headedness than self-attention layers.The conclusion identifies this component-level difference as a central finding.
  • 8 Conclusion: The relative importance of attention heads is determined in the early stages of training.The authors hope these observations inspire models that use parameters and attention more efficiently.

A Ablating All Heads but One: Additional Experiment.

The additional experiment reports performance changes when each layer retains only one head. The retained head is selected using its standalone performance on a separate dataset, with statistical significance indicated in the table.

  • A Ablating All Heads but One: Additional Experiment.: Tables 5 and 6 report the performance difference for keeping only one head in each layer.The comparison is organized by layer and uses the best single head for that layer.
  • A Ablating All Heads but One: Additional Experiment.: The retained head is selected as the best head based on standalone evaluation on a separate dataset.This selection procedure separates head choice from the reported evaluation.
  • A Ablating All Heads but One: Additional Experiment.: Table 5 reports best delta BLEU by layer on newstest2014 for the WMT model, with the best head evaluated on newstest2013.Underlined values mark statistically significant changes at p < 0.01.

B Additional Pruning Experiments

Additional pruning experiments evaluate importance-driven head removal across four datasets and show that substantial pruning can preserve performance in some cases.

  • The additional experiments cover SST-2, CoLA, MRPC, and IWSLT using fine-tuned BERT or a German-to-English translation model.
  • 60% of heads can be pruned on SST-2 without a noticeable impact on performance.
  • 50% of heads can be pruned on CoLA and MRPC without a noticeable impact on performance.
Loading 1905.10650v3…