Multi-Head Attention: Collaborate Instead of Concatenate

By Jean-Baptiste Cordonnier, Andreas Loukas, Martin Jaggi. 2020

180 citations recorded by the source corpus.

Attention layers are widely used in natural language processing (NLP) and are beginning to influence computer vision architectures. However, they suffer from over-parameterization. For instance, it was shown that the majority of attention heads could be pruned without impacting accuracy. This work aims to enhance current understanding on how multiple heads interact. Motivated by the observation that trained attention heads share common key/query projections, we propose a collaborative multi-head attention layer that enables heads to learn shared projections. Our scheme improves the computational cost and number of parameters in an attention layer and can be used as a drop-in replacement in any transformer architecture. For instance, by allowing heads to collaborate on a neural machine translation task, we can reduce the key dimension by a factor of eight without any loss in performance. We also show that it is possible to re-parametrize a pre-trained multi-head attention layer into our collaborative attention layer. Even without retraining, collaborative multi-head attention manages to reduce the size of the key and query projections by half without sacrificing accuracy. Our code is public.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, 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 Computer Vision and Pattern Recognition

In Model Architectures

Citations in the corpus

Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth (2021, methodology): “We obtain the final equation by noting that softmax is shift-invariant and disregarding terms that provide a constant contribution across rows [CLJ20a].”
BitFit: Simple Parameter-efficient Fine-tuning for Transformer-based Masked Language-models (2021, neutral): “The ‘key’ bias bk has zero change, consistent with the theoretical observation in (Cordonnier et al., 2020).”
A Survey of Transformers (2021, methodology): “Collaborative Multi-head Attention [21] uses shared query and key projection W 𝑄 and W 𝐾 and a mixing vector m 𝑖 for the 𝑖 -th head to filter from the projection parameters such that Eq.”
Pay Attention when Required (2020, background): “Previous research (Kovaleva et al., 2019; Cordonnier et al., 2020) on transformer layers also indicate that attention layers are severely over parametrized and that they are more prevalent in the beginning of a network (Press et al.”