Source-linked AI summary
Multi-Head Attention: Collaborate Instead of Concatenate
Jean-Baptiste Cordonnier, Andreas Loukas, Martin Jaggi
TL;DR
Multi-head attention contains redundant key/query representations, raising whether its heads should remain independently projected. The paper replaces duplicated projections with collaborative shared projections and reports a fourfold key/query-dimension reduction without performance loss in NMT and vision.
Problem
The roles and interactions of attention heads are poorly understood, and many heads or key/query dimensions may be redundant.
Method
Collaborative multi-head attention lets heads share common key/query projections, with tensor decomposition enabling post-hoc conversion of pretrained transformers.
Results
Collaborative attention divides the key/query dimension by 4 without performance drop when used in NMT and vision.
Takeaways & Limitations
The approach provides a drop-in transformer replacement and a principled alternative to post-hoc head pruning for reducing attention inefficiency.
Takeaways & Limitations
The analyzed attention operator omits key/query biases, although some transformer re-implementations include them.
Abstract
from arXiv · showhide
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.
1 Introduction
Multi-head attention improves accuracy, but the roles and redundancy of individual heads remain poorly understood. This work analyzes shared key/query features and proposes collaborative attention as a re-parametrized alternative to concatenated heads.
- Multiple heads can improve accuracy, yet some heads can be pruned without affecting performance, while other evidence finds multiple heads necessary for convolution-like self-attention.
- The paper investigates whether independent heads learn overlapping or distinct concepts and finds redundant key/query projected dimensions across concatenated heads.
- The proposed collaborative attention shares common key/query projections instead of duplicating them across heads.
- PCA characterizes cross-head key/query redundancy, motivating a re-parametrization in which heads learn common projections.
- The experiments evaluate collaborative heads across neural machine translation, language understanding, and image classification.
2 Multi-Head Attention
Attention maps query and key inputs into a shared space, computes scaled dot-product scores, and combines values; multi-head attention replicates this mechanism across heads. The paper also analyzes how biases and shared column-space structure affect the mechanism.
- 2.1 Attention: An attention layer projects queries, keys, and values using WQ, WK, and WV, then maps each query token from Din to Dout dimensions.
- 2.1 Attention: Self-attention applies the attention layer to the same sequence by setting X = Y.
- 2.2 Content vs. Context: Some transformer re-implementations add key/query biases bQ and bK, unlike the bias-free attention operator defined in the paper.
- 2.2 Content vs. Context: The concatenated key/query matrices are low rank across heads, indicating that heads share common projections in their column-space.
- 2.3 Multi-Head Attention: In standard multi-head attention, each head has distinct value parameters, and an output matrix projects the concatenated head outputs into the output space.
3 Improving the Multi-Head Mechanism
The paper finds that attention heads learn redundant key/query subspaces and replaces independent concatenation with collaborative shared projections. This reparameterization supports adaptive head expressiveness, parameter reduction, and post-hoc conversion of pretrained attention layers.
- 3.1 How much do heads have in common?: Many heads focus on common key/query subspaces: one third of the dimensions captures almost all energy in their product.Individual key/query matrices are not low rank on average, but their concatenation is low rank.
- 3.2 Collaborative Multi-Head Attention: Collaborative attention learns shared key/query projections and lets each head re-weight those projections with a mixing vector.The shared matrices have dimension Din × ˜Dk, while each head defines a custom dot product over the shared projected dimensions.
- 3.2 Collaborative Multi-Head Attention: Learning mixing vectors increases each head’s expressive power with a negligible parameter increase and allows heads to attend to smaller or larger subspaces.Head size becomes adaptive rather than fixed at the usual dk.
- 3.3 Head Collaboration as Tensor Decomposition: Canonical tensor decomposition extracts shared projections and mixing weights, enabling conversion of pretrained standard attention without retraining.If ˜Dk ≥ Dk, the decomposition is exact and collaborative attention can express any concatenation-based attention layer.
- 3.2 Collaborative Multi-Head Attention: At ˜Dk = Dk/2, collaborative and traditional attention layers have similar inference speed for the tested setting.The comparison uses T = 128 tokens, batch size 32, and a V100 GPU.
4 Experiments
Experiments across translation, vision, and language understanding show that collaborative multi-head attention can replace concatenation-based attention while reducing key/query dimensions and preserving performance. Post-hoc tensor-decomposition re-parametrization also compresses pretrained models with limited accuracy changes.
- Experimental scope: Collaborative MHA serves as a drop-in replacement across NMT, NLU, and image classification transformer experiments.Experiments cover WMT14 EN-DE, GLUE, and ImageNet, including training from scratch and pretrained-model re-parametrization.
- NMT: 4× reduction in shared key/query dimension preserves BLEU on WMT14 EN-DE, whereas reducing classic head dimensions degrades performance.At Dk = 64, classic MHA loses 0.6 BLEU points, while collaborative MHA remains within 0.1 point of the concatenation baseline.
- NMT: Collaborative attention consistently improves BLEU and permits reducing Dk from 512 to 128 without a performance drop.The comparison uses an encoder-decoder transformer-base on WMT14 EN-DE.
- Vision: On ImageNet, collaborative attention improves performance at matched key/query dimensions and matches DeiT-B with four times fewer dimensions.Collaborative models with Dk = 384 and 192 match the baseline, with an 8% gain in total parameter efficiency.
- Vision: Post-hoc Tucker re-parametrization of pretrained DeiT-B changes Acc@1 by only 0.1% when compressing Dk from 768 to 512.Compressing to Dk = 256 changes ImageNet Acc@1 by 1%.
- NLU: On GLUE, BERT and DistilBERT key/query dimensions can be reduced by 2× and 3× respectively without sacrificing more than 1.5%.ALBERT tolerates a smaller 1.5× reduction with minor performance degradation; compression losses up to 3× can be recovered through second fine-tuning.
5 Conclusion
Collaborative multi-head attention replaces concatenated heads to address redundant query/key representations, while supporting compression and post-hoc transformation of trained networks.
- 4× smaller key/query dimensions preserve performance in NMT and vision when collaborative MHA replaces standard MHA.The conclusion reports no performance drop under this compression.
- Figure 4 measures the time required to decompose BERT-base from Dk = 768 to ˜Dk.
- Post-hoc decomposition can compress a fine-tuned BERT-base model, with additional fine-tuning recovering original performance after drastic compression.The figure caption describes performance across MNLI, MRPC, and STS-B as compressed dimension varies.
- Collaborative MHA offers a shared-projection alternative to post-hoc head pruning for reducing multi-head attention inefficiency.The authors characterize the approach as a principled replacement for pruning-based processing.
A Hyperparameters for Neural Machine Translation Experiments
The NMT experiments use the Fairseq implementation and document their experimental settings in a dedicated hyperparameter table.
- The NMT implementation is based on Fairseq.
- Table 6 presents the hyperparameters for the NMT experiment.
B Hyperparameters for ImageNet Experiments
The ImageNet experiments use a DeiT-based implementation with hyperparameters reported separately.
- The ImageNet implementation is based on DeiT.
- Table 7 reports the hyperparameters used for the ImageNet experiments.
C Hyperparameters for Natural Language Understanding Experiments
The natural-language-understanding experiments use standard HuggingFace models and specified GLUE fine-tuning settings, with optimization details organized in the accompanying tables.
- The experiments use standard models downloaded from HuggingFace with configurations presented in Table 8.
- GLUE fine-tuning uses a learning rate of 2·10^-5 for 3 epochs, except SST-2 and RTE, which use 10 epochs.
- The tensor decomposition tolerance was tested at 10^-6, 10^-7, and 10^-8 without significant improvement, so the authors retained the default.
- Table 9 reports the optimization hyperparameters for the natural-language-understanding experiments.