Pit One Against Many: Leveraging Attention-head Embeddings for Parameter-efficient Multi-head Attention (2023)

Abstract: Scaling pre-trained language models has resulted in large performance gains in various natural language processing tasks but comes with a large cost in memory requirements. Inspired by the position embeddings in transformers, we aim to simplify and reduce the memory footprint of the multi-head attention (MHA) mechanism. We propose an alternative module that uses only a single shared projection matrix and multiple head embeddings (MHE), i.e. one per head. We empirically demonstrate that our MHE attention is substantially more memory efficient compared to alternative attention mechanisms while achieving high predictive performance retention ratio to vanilla MHA on several downstream tasks. MHE attention only requires a negligible fraction of additional parameters ( 3nd , where n is the number of attention heads and d the size of the head embeddings) compared to a single-head attention, while MHA requires (3n^2-3n)d^2-3nd additional parameters.

## introduction scaling pre-trained language models (plms) aims to enhance performance by increasing their size and capacity, leading to models with an unprecedented number of parameters @xcite @xcite . just by increasing the size of plms and the pre-training data has yielded state-of-the-art performance on various natural language processing (nlp) tasks @xcite @xcite @xcite @xcite . however, the pursuit of developing larger plms comes with large computational requirements. this has direct environmental implications such as large carbon emissions @xcite @xcite , conflicting with the principles of green artificial intelligence development @xcite . moreover, scaling can hinder researchers with limited access to computing resources to participate in advancing the field @xcite . this results in inequalities, where only a privileged few can actively contribute, potentially impeding diversity and inclusivity @xcite . the backbone of transformers @xcite is the multi-head attention (mha) module that extends the standard single-head attention (sha) proposed by @xcite . mha applies an attention mechanism (i.e. head) multiple times for the same set of queries, keys and values by using a different set of parameters (i.e. projection matrices) for each of them. this results in mha modules with a large memory footprint that increases with the number of layers and attention heads per layer in plms @xcite @xcite . figure 1 shows how the number of parameters of a single attention sublayer increases with its number of attention heads. previous work has attempted to address this issue by proposing to share projection matrices or eliminating them entirely to improve the parameter efficiency of mha. @xcite proposed sharing projection parameters for keys, queries and values across layers, while @xcite introduced a method for sharing the projection matrix between keys and values within each transformer layer. additionally, similar approaches use a multi-query attention approach that uses a pair of global projection matrices for keys and values in each layer @xcite @xcite . furthermore, @xcite eliminate the projection matrices entirely and directly treat the input hidden states as both keys and values. in a different direction, @xcite propose models that replace the attention blocks with token-mixture blocks (i.e. using linear or fourier transformations) that contain fewer or no parameters compared to mha. inspired by the position embeddings in transformers @xcite , we aim to simplify and reduce the memory footprint of the mha mechanism. we achieve this using only a single projection matrix for each of the keys, queries and values respectively shared across all attention heads, and one embedding per head (mhe). our contributions are as follows: ## multiple head embeddings attention inspired by the absolute position embeddings @xcite for distinguishing the representation of the same token in different contexts, we propose multiple head embeddings (mhe) attention. mhe uses a shared 'seed' projection matrix that is subsequently combined with distinct head embeddings to generate multiple attention heads. ## discussion mha enables the model to attend to information from different representation subspaces at different positions @xcite . it uses distinct projection matrices for each attention head and integrates the information from these different representation subspaces. however, @xcite did not explore different methods for performing space transformations per head. previous work has pointed out that overparameterized models might have a low intrinsic dimension. therefore, transforming the projection matrices to smaller low-rank ones usually does not severely harm model predictive performance @xcite . meanwhile, the classic mha approach also does not impose any constraints on the orthogonality of these subspaces during pre-training and fine-tuning. the column vectors in those projection matrices could be highly collinear, i.e. the projection matrices could be rank-deficient. as a result, its inner-working mechanism could be simply understood as introducing levels of variation to the encoded representation of the same token at the same position across different heads. our mhe approach is possible to achieve memory efficiency (similar to sha) together with high prr compared to mha by mimicking the position embeddings for representing different attention heads. on one hand, the addition operation in mhe-add is used for transforming the keys, queries and values. this can be seen as a small distortion of the subspace obtained through projection, followed by rotation. for an input representation, the difference between the projected and injected (i.e. through head embedding addition) queries, keys and values is a constant vector across any pair of heads. on the other hand, the mhe-mul approach employs a multiplication operation, which more aggressively distorts and reshapes the keys, queries and values subspaces. head embeddings in mhe-mul play a role as the scaling factors, respectively stretching each dimension of the input representation. thus, the difference between the keys, queries, and values generated by different heads for the same input representation, is a vector parallel to the projected input. this vector is dependent on the specific input, unlike the constant vector in mhe-add. interestingly, our experimental results consistently show that the multiplication operation outperforms addition in the majority of benchmarks. this corroborates findings of a previous empirical study by @xcite that compared rotary position embeddings (somehow analogous to mhe-mul) with absolute position embeddings (analogous to mhe-add). ## conclusions we have proposed mhe attention that employs a single shared projection matrix along with multiple head embeddings, to simplify and reduce the memory footprint of the mha. our experimental results have demonstrated that mhe attention exhibits superior memory efficiency compared to other memory-efficient attention variants, while achieving high predictive performance ratio to mha on various downstream tasks. compared to a single-head attention, mha requires (3n 2 -3n)d 2 parameters for n attention heads and head dimensionality d, while mhe barely requires a negligible 3nd. for future research, we plan to investigate scaling up mhe models and explore its linguistic capabilities @xcite .