Multi-Head Attention: Collaborate Instead of Concatenate

TL;DR

This paper introduces a collaborative multi-head attention mechanism that shares key/query projections among heads, reducing parameters and maintaining performance across NLP and vision tasks.

Contribution

It proposes a novel shared projection scheme for multi-head attention, enabling parameter reduction and compatibility with existing transformer models.

Findings

Sharing key/query projections maintains accuracy while reducing parameters.

The method is effective in language understanding, translation, and vision tasks.

Pre-trained models can be re-parametrized into the collaborative attention form.

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.