[Paper Review] Multi-Head Attention: Collaborate Instead of Concatenate

Jean-Baptiste Cordonnier, Andreas Loukas+1 more|arXiv (Cornell University)|Jun 29, 2020

Advanced Neural Network ApplicationsComputer Science27 references76 citations

TL;DR

This paper introduces collaborative multi-head attention (collab MHA) to share key/query projections across heads, reducing parameters and maintaining or improving performance across NLP translation, NLP understanding, and vision tasks. It also shows post-hoc reparameterization via tensor decomposition from standard MHA.

ABSTRACT

Attention layers are widely used in natural language processing (NLP) and are beginning to influence computer vision architectures. Training very large transformer models allowed significant improvement in both fields, but once trained, these networks show symptoms of over-parameterization. For instance, it is known that many attention heads can be pruned without impacting accuracy. This work aims to enhance current understanding on how multiple heads interact. Motivated by the observation that attention heads learn redundant key/query projections, we propose a collaborative multi-head attention layer that enables heads to learn shared projections. Our scheme decreases the number of parameters in an attention layer and can be used as a drop-in replacement in any transformer architecture. Our experiments confirm that sharing key/query dimensions can be exploited in language understanding, machine translation and vision. We also show that it is possible to re-parametrize a pre-trained multi-head attention layer into our collaborative attention layer. Collaborative multi-head attention reduces the size of the key and query projections by 4 for same accuracy and speed. Our code is public.

Motivation & Objective

• Characterize redundancy in key/query projections across MHA heads.
• Propose collaborative head attention that shares projections across heads.
• Show how to reparameterize pre-trained transformers to use collaborative attention.
• Demonstrate parameter and computation efficiency in NLP and vision tasks.
• Provide practical guidance for drop-in replacement in transformer architectures.

Proposed method

• Analyze redundancy with PCA on W_Q W_K across heads to show low-rank behavior of concatenated heads.
• Define Collaborative Head Attention (collab MHA) where all heads share W_Q and W_K projected to a shared dimension ˜Dk and use mixing vectors m_i.
• Relate collaborative heads to standard MHA as a special case (˜Dk = Nh dk).
• Use Tucker/CP (canonical) tensor decomposition to reparameterize existing trained attention layers into collaborative attention (derive W̃_Q, W̃_K, and mixing matrix M).
• Discuss biases handling via content/context decomposition and non-contributing bias terms in attention scores.
• Provide complexity/parameter analysis and compare with pruning.
• Validate across NMT, NLU (GLUE), and Vision (ImageNet) tasks, including from-scratch training and post-hoc reparameterization.

Experimental results

Research questions

• RQ1 Do attention heads learn redundant key/query projections allowing shared projections across heads?
• RQ2 Can collaborative MHA match or exceed standard MHA performance while reducing parameters?
• RQ3 Is it possible to convert pre-trained transformers to collaborative MHA without full retraining (post-hoc reparameterization)?
• RQ4 What are the practical efficiency gains in terms of parameters and speed for NLP and vision tasks?
• RQ5 How does collaborative MHA impact model compression and finetuning needs on benchmarks like GLUE and ImageNet?

Key findings

• Heads tend to share projections across subspaces, making W_Q W_K across heads low-rank when concatenated but not individually.
• Collaborative MHA reduces the number of parameters by sharing key/query projections (˜Dk) and a mixing matrix (M) across Nh heads.
• Replacing standard MHA with collab MHA maintains or improves performance while allowing 4× reduction in key/query dimension without accuracy loss in NMT.
• From-scratch training with collab MHA yields up to 35% fewer attention parameters with no BLEU drop and sometimes improved BLEU on WMT14 EN-DE.
• Post-hoc reparameterization of pre-trained transformers to collaborative attention using tensor decomposition preserves or slightly reduces accuracy, with substantial compression possible (e.g., 2×–3× reductions in ˜Dk on GLUE).
• In vision, collab MHA matches or exceeds concatenation-based performance at lower ˜Dk values and enables compression with minimal accuracy loss on ImageNet.