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

Abstract

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.

arXiv commentACL 2019 (camera-ready)

References in corpus (6)

• Categorical Reparameterization with Gumbel-Softmax
• The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables
• The State of Sparsity in Deep Neural Networks
• What do Neural Machine Translation Models Learn about Morphology?
• Training Tips for the Transformer Model
• Discrete Autoencoders for Sequence Models

Cited by in corpus (41)

• Transformers in Vision: A Survey
• VisualBERT: A Simple and Performant Baseline for Vision and Language
• Reducing Transformer Depth on Demand with Structured Dropout
• The NLP Cookbook: Modern Recipes for Transformer based Deep Learning Architectures
• DynaBERT: Dynamic BERT with Adaptive Width and Depth
• Structured Pruning of a BERT-based Question Answering Model
• Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers
• BERT Loses Patience: Fast and Robust Inference with Early Exit
• Graph Transformer for Graph-to-Sequence Learning
• AdaBERT: Task-Adaptive BERT Compression with Differentiable Neural Architecture Search
• Medical Image Segmentation Using Squeeze-and-Expansion Transformers
• DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference
• EarlyBERT: Efficient BERT Training via Early-bird Lottery Tickets
• Blockwise Self-Attention for Long Document Understanding
• Pre-Trained Models: Past, Present and Future
• Low-Rank Bottleneck in Multi-head Attention Models
• KDLSQ-BERT: A Quantized Bert Combining Knowledge Distillation with Learned Step Size Quantization
• SANST: A Self-Attentive Network for Next Point-of-Interest Recommendation
• Let's measure run time! Extending the IR replicability infrastructure to include performance aspects
• Controlling Computation versus Quality for Neural Sequence Models
• Understanding and Improving Encoder Layer Fusion in Sequence-to-Sequence Learning
• Patch Slimming for Efficient Vision Transformers
• Understanding Multi-Head Attention in Abstractive Summarization
• Do Transformer Attention Heads Provide Transparency in Abstractive Summarization?
• Rethinking Network Pruning -- under the Pre-train and Fine-tune Paradigm
• Learning on heterogeneous graphs using high-order relations
• On the Dynamics of Training Attention Models
• Block-wise Dynamic Sparseness
• Shatter: An Efficient Transformer Encoder with Single-Headed Self-Attention and Relative Sequence Partitioning
• Attend What You Need: Motion-Appearance Synergistic Networks for Video Question Answering
• On the Importance of Local Information in Transformer Based Models
• Super Tickets in Pre-Trained Language Models: From Model Compression to Improving Generalization
• Soft Attention: Does it Actually Help to Learn Social Interactions in Pedestrian Trajectory Prediction?
• WaLDORf: Wasteless Language-model Distillation On Reading-comprehension
• Generating Diverse Translation from Model Distribution with Dropout
• A Dynamic Head Importance Computation Mechanism for Neural Machine Translation
• Pruning Attention Heads of Transformer Models Using A* Search: A Novel Approach to Compress Big NLP Architectures
• How BPE Affects Memorization in Transformers
• Unbiased Sentence Encoder For Large-Scale Multi-lingual Search Engines
• Rotate King to get Queen: Word Relationships as Orthogonal Transformations in Embedding Space
• Modularized Transfomer-based Ranking Framework