Finding the Pillars of Strength for Multi-Head Attention (2023)

Abstract: Recent studies have revealed some issues of Multi-Head Attention (MHA), e.g., redundancy and over-parameterization. Specifically, the heads of MHA were originally designed to attend to information from different representation subspaces, whereas prior studies found that some attention heads likely learn similar features and can be pruned without harming performance. Inspired by the minimum-redundancy feature selection, we assume that focusing on the most representative and distinctive features with minimum resources can mitigate the above issues and lead to more effective and efficient MHAs. In particular, we propose Grouped Head Attention, trained with a self-supervised group constraint that group attention heads, where each group focuses on an essential but distinctive feature subset. We additionally propose a Voting-to-Stay procedure to remove redundant heads, thus achieving a transformer with lighter weights. Extensive experiments are consistent with our hypothesis. Moreover, our method achieves significant performance gains on three well-established tasks while considerably compressing parameters.

## introduction transformers have shown promising performance across various tasks . however, they have some issues, e.g., redundancy and over-parameterization, which is mainly caused by multi-head attention (mha) @xcite and feed-forward network (ffn) @xcite @xcite of transformer. we aim to mitigate the redundancy and over-parameterization issues by optimizing the mha module. the multi-heads were originally designed to attend to different representation subspaces of input @xcite . however, prior works @xcite have shown that the attention heads are highly redundant and over-parameterized after training because some heads can be switched off with a negligible performance drop. such an issue is probably caused by their parallel design: the heads naturally work in the same way and likely attend to similar features @xcite . the existing redundancy optimization methods are mainly based on homogenization, diversification, and head significance. however, they all have some limits. (1) the homogenization-based methods mitigate redundancy and over-parameterization by making heads similar and removing unnecessary parameters. @xcite homogenized attention heads by sharing most weights between all heads, which reduced the redundant parameters but sacrificed the performance somewhat because of the lack of diversity. (2) the diversificationbased methods diversify the heads to enrich features and reduce the inter-head redundancy. @xcite found that diversifying attention heads by adding a regularization could force mha to reduce inter-head information redundancy, yielding performance gains in machine translation. however, such strategy that retains all feature subsets is suboptimal, because it does not address the issue that mha is over-parameterized. (3) the significancebased methods @xcite @xcite learn significance scores for the heads to prune unimportant ones. however, the retained important heads still remain inter-head redundancy without diversifying them. considering the issues of the above-mentioned methods, we hypothesize that attending to the most representative and distinctive feature subsets with minimum resources leads to more effective and efficient mhas, which is inspired by the minimum-redundancy feature selection @xcite . accordingly, we propose a divide-and-conquer strategy, including group-constrained training (gct) and voting-to-stay (v2s), to achieve the setting of our assumption and mitigate the above-mentioned issues. we illustrate them below. we first propose a strategy to group and distinguish attention heads, where a grouped head attention (gha) is obtained via the self-supervised gct. by encouraging homogenization within a group and diversification between groups, the mha is forced to divide its heads to work in several separate groups, where each group focuses on an essential but unique feature subset, being in line with the setting of our assumption. note that the redundancy exists when the resources deployed by the model are more than enough to process current information @xcite . gha reduces the redundancy in two aspects: next, we show that gha-ps (gha with the pillar of strength), a lighter-weight gha, can be achieved by excluding the redundant heads of gha via the v2s procedure. v2s culls the redundant heads that share similar patterns with the most representative head (ps head) of a group, which is selected by voting on different training batches. note that upon the convergence of the gct, the heads are highly homogenized within a group, thus being redundant because they process similar information. as a result, once the redundant heads are culled, the ps heads can still achieve the essential utility of the original attention layer and yield comparable performance to the unculled model. the lottery ticket hypothesis @xcite argues that subnetworks in an over-parameterized neural network can converge faster and achieve comparable or better performance than the original network. our gha-ps achieving better results is also in line with this hypothesis. such a divide-and-conquer combination resolves the issues of previous redundancy optimization methods: (1) our model achieves better parameter efficiency, resolving the issue of previous diversification-based methods; (2) the feature diversity is guaranteed and the inter-head redundancy is reduced, resolving the problems of previous homogenization-and significance-based methods. we evaluate our method on three benchmarking tasks. we denote the corresponding transformer architectures of gha and gha-ps as grouped head transformers (ght) and grouped head transformers with the pillars of strength (ght-ps), respectively. ght and ght-ps achieve significant improvements over the strong baselines in machine translation (mt) bleu scores (+3.8% and +4.4% averaged on 7 datasets), language modeling (lm) perplexity (-2.8% and -2.9%), and abstractive summarization (as) f1-rouge (+6.7% and +7.0% on average). ght-ps exhibits higher efficiency in model size, inference speed, and floating-point operations (flops). the light architecture of ght-ps reduces 63.6% parameters of the vanilla transformer and yields comparable performance. the key contributions of our work foot_0 are threefold: our method uses fewer parameters and achieves better performance, outperforming the existing mha redundancy/parameter reduction methods. parameter efficiency. different methods were proposed to achieve lightweight transformers: (1) replacing attention with lightweight modules, e.g., convolution modules, such as dynamic conv (wu et al., 2019) and lite transformer (wu et al., 2020); (2) removing or replacing the feed-forward layers, such as sukhbaatar et al. (2019) and wu et al. (2020); (3) pruning the model, such as michel et al. (2019), voita et al. (2019), and li et al. (2021). modified multi-head mechanism. ahmed et al. (2017) learned to weight the projected output of different heads, performing weighted sum over them. li et al. (2019) aggregated the output of different heads by dynamic routing; cui et al. (2019) used different attention mechanisms, e.g., global/local and forward/backward attention for different heads; shazeer et al. (2020) mixed different heads before and after the softmax operation in an attention function to achieve communication between heads. head redundancy optimization. michel et al. (2019) and voita et al. (2019) found that only a subset of the attention heads have significant utilities in transformer, where the important heads could be identified by expected sensitivity and layer-wise relevance propagation (lrp) (ding et al., 2017). upon this, li et al. (2021) learned per-head importance scores and pruned the heads. cordonnier et al. (2020) homogenized the attention heads by sharing a part of the weights between heads, which lowered the number of parameters but sacrificed performance. li et al. (2018) found that diversifying attention heads by adding a regularization can force mha to reduce inter-head redundancy, yielding performance gains for machine translation. however, previous methods either traded performance for efficiency or retained extra parameters. ## methodology there are two core components in our method, namely the group-constrained training (gct) and the voting-to-stay (v2s) procedure. gha (figure 1 ) is developed with gct that removes head redundancy; gha-ps is developed by removing the redundant parameters of gha in v2s. in this section, we detail the process of developing the gha and finding its pillars of strength (ps). ## experimental setup in this section, we detail the key architectural configurations. further training, model, dataset & evaluation setups are detailed in a.1, a.2, & a.3. we follow the transformer of @xcite as a backbone architecture for all datasets and tasks in our experiments. following @xcite , for machine translation and abstractive summarization, we adopt the same 8-head encoderdecoder architecture with 6 layers for both encoder and decoder, where the model dimension d @xmath0. for lm, we adopt the 16-head decoder-only architecture with 16 layers, where the model dimension d @xmath1. the layer normalization is applied before the residual connection of each layer. the parameters of decoder input and output projections are shared. our models are based on fairseq @xcite implementations. we perform the gct as a metric learning task because it does not introduce additional projection layers when the shapes of similarity inputs are identical (eq.4(b)), which makes ght weight-lighter. in addition, it performs better in our experiments compared to the classification-based grouping. ## conclusion in this paper, we assume that only focusing on the most representative and distinctive features with minimum resources may mitigate the redundancy and over-parameterization issues of mha. accordingly, we propose a divide-and-conquer strategy, including gct and v2s to mitigate the issues. the improvements on three tasks and the extensive analysis verify our hypothesis and the effectiveness of our redundancy optimization methods. our study may inspire future mha design and training to achieve higher accuracy and efficiency.