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

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.

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

Introduction

The paper "Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned" by Voita et al. presents an in-depth exploration of the roles played by individual attention heads within the Transformer architecture and introduces pruning methodologies. Through rigorous analysis, the paper aims to discern the specific contributions of attention heads, identify their roles, and develop strategies for optimizing model performance by pruning less influential heads.

Key Findings and Methodology

The authors employ layer-wise relevance propagation (LRP) to evaluate the importance of each attention head within the Transformer's encoder. LRP quantifies the relative contributions of neurons to the model's predictions, providing a granular view of each head's significance. The paper predominantly focuses on self-attention within the encoder, scrutinizing:

• The impact of individual encoder heads on translation quality.
• The roles and interpretability of these heads.
• The sensitivity of different model components to head quantity.
• The feasibility of reducing head count without performance degradation.

Initially, the researchers identify the most crucial heads across encoder layers, using layer-wise relevance propagation and a confidence metric based on mean attention weights. They categorize head functions into three primary roles: positional (attending to adjacent tokens), syntactic (focusing on specific syntactic dependencies), and attention to rare words (highlighting less frequent tokens). Through systematic examination, they demonstrate the positional and syntactic heads as predominantly valuable, confirming their relevance through function-specific accuracy benchmarks.

Pruning Methodology

A central contribution of the paper is the introduction of a novel head pruning method based on the principles of stochastic gates and a differentiable relaxation of the $L_0$ penalty. This method facilitates continuous learning, leading to:

• Reduction in Head Count: For the English-Russian WMT dataset, pruning 38 out of 48 heads only reduces BLEU score by 0.15 points.
• Specialized Heads Retention: Specialized heads, particularly those with positional or syntactic functions, are pruned last, underscoring their indispensability.

Numerical Results and Practical Implications

The analysis reveals that merely a subset of heads (approximately 10 out of 48) are vital for maintaining translation quality close to the full model. The paper emphasizes that:

• Translation Quality: A significant reduction in heads results in minimal translation quality loss (0.15 BLEU for WMT and 0.25 BLEU for OpenSubtitles).
• Efficient Utilization: By employing selective sparsification, the model sustains performance while achieving computational efficiency.

Future Developments

The implications of this research extend to both practical and theoretical domains, suggesting several avenues for future exploration:

• Model Compression: Leveraging this pruning methodology can optimize model deployment in resource-constrained settings.
• Head Reassignment: Further research could explore dynamic reassignment of functions across heads during training phases, enhancing model adaptability.
• Broader Applicability: Expanding this analysis to other architectures and languages could generalize these findings, contributing further to the field of neural machine translation.

Conclusion

Voita et al.'s paper offers a comprehensive exploration of the specialized roles of multi-head self-attention in the Transformer model. By identifying and pruning less critical heads, the approach paves the way for more efficient models without substantial performance trade-offs. This research stands as a significant step towards understanding the intricate workings of self-attention mechanisms and optimizing neural machine translation models through informed pruning strategies.