Mixture of Attention Heads: Selecting Attention Heads Per Token (2022)

Abstract: Mixture-of-Experts (MoE) networks have been proposed as an efficient way to scale up model capacity and implement conditional computing. However, the study of MoE components mostly focused on the feedforward layer in Transformer architecture. This paper proposes the Mixture of Attention Heads (MoA), a new architecture that combines multi-head attention with the MoE mechanism. MoA includes a set of attention heads that each has its own set of parameters. Given an input, a router dynamically selects a subset of k attention heads per token. This conditional computation schema allows MoA to achieve stronger performance than the standard multi-head attention layer. Furthermore, the sparsely gated MoA can easily scale up the number of attention heads and the number of parameters while preserving computational efficiency. Despite performance improvements, MoA also automatically differentiates heads’ utilities, providing a new perspective to discuss the model’s interpretability. We conducted experiments on several important tasks, including Machine Translation and Masked Language Modeling. Experiments have shown promising results on several tasks against strong baselines that involve large and very deep models.

## introduction in recent years, large models have become a popular trend in the research of natural language processing, especially large-scale transformer @xcite . the model's capacity has increased from millions of parameters @xcite , to billions of parameters @xcite @xcite , even to trillions of parameters @xcite . howthe output is a weighted sum of the selected attention heads given the confidence calculated by the router. a popular trend is to utilize conditional computation with a sparsely activated model to seek greater computational efficiency. thus, only a part of the model's parameters is used for a specific input during the forward computation, which alleviates the computational load. among these attempts, the mixture of experts (moe) @xcite is an essential technique. since first applying the mixture of experts to transformer architecture @xcite , researchers have mainly focused on combining the feed-forward network layer and the mixture of experts. recent works have discussed how to get a better routing strategy @xcite @xcite or how to scale up the mixture of experts on different nodes of gpus @xcite . however, few attempts have explored the possibility of combining moe with the multi-head attention (mha) mechanism. since the mha is another compulsory module in the transformer architecture, combining moe and the attention mechanism could also help achieve better performance while restraining the computational cost. besides, previous research has investigated the utility of different attention heads. @xcite found that the combination (reallocation) of a subset of attention heads helps the translation task since they prune the useless attention heads. in the field of dependency parsing, researchers have unveiled that some attention heads in bert-like language models @xcite model individual dependency types @xcite and syntactic functions @xcite . @xcite claimed that the attention heads have different functions that could be categorized into three types. there is no need to pass through all multiple attention heads for an input token if we could select some relevant attention heads whose function is proper. thus, we conceive an attention mechanism that selects different attention heads per token. based on the above discussion, we proposed mixture of attention heads (moa) (section 4), an attention mechanism that selects different attention heads for different inputs. a simple illustration of this idea is shown in figure 1 . moa includes a set of of attention heads with different parameters. given an input, a routing network dynamically selects a subset of k attention heads for each token. the output is a weighted sum of the selected attention heads given the confidence calculated by the routing network. we conducted experiments on two tasks: machine translation and masked language modeling (section 5). experiments shown promising results against several strong baselines. in all tasks, our proposed mixture of attention heads outperforms the original transformer architecture @xcite . our model surpasses many large models or achieves comparable results with only a half computational cost. our contributions can be summarized in three folds: 1) we proposed a new attention mechanism called mixture of attention heads, combining the idea of mixture of experts with the attention mechanism. 2) moa can improve the model's performance without substantially adding parameters and computational cost. 3) moa is easy to scale up while maintaining with a restrained computation complexity, resulting in a further performance amelioration. various routing strategies @xcite @xcite have been investigated for stabilizing the moe training and balancing the expert loads. @xcite pointed out the representation collapse issue in the sparse mixture of experts models and solved by a two-stage routing strategy. machine translation architectures with original transformer architecture @xcite , @xcite found that training with reduced precision and large batch could improve the translation performance. some models get better performance on translation by using larger scale of transformer. @xcite deepened the encoder and decoder of the transformer by adequately initializing the model. deepnet @xcite scaled transformers up to 1,000 layers by introducing a new normalization function. however, these methods require a great amount of computational cost. some models make changes to the self-attention module. @xcite proposed mae model. the reallocation of attention heads got better performance on translation, since the model prune useless multi-head attention heads. however, their method is difficult to scale up and get further improvement of the results because it needs to use all the attention heads in the model rather than sparsely activate them. it also requires the complicated block coordinate descent training steps. @xcite proposed dynamicconv and lightconv by replacing self-attention mechanism with a lightweight convolution. ## mixture of attention heads in this work, we propose a variant of multi-head attention for transformer called mixture of attention heads (moa), illustrated in figure 2 . moa consists of two major components, the routing network g and a group of n attention experts {e 1 , ..., e n }. similar to standard multi-head self-attention, the input of moa includes three sequences, query sequence q, key sequence k, and value sequence v . we note q t as the query vector at time step t. for each q t , the routing network g selects a subset of k experts g(q t ) ⊆ {e i } based on q t and assign a weight w i to each selected expert. then, these selected experts take q t , k, and v as inputs and compute an output e i (q t , k, v ). the output of the moa is the weighted sum of the selected experts' outputs. formally, the moa output at time step t can be written as: ## conclusion this work introduces the mixture of attention heads (moa). moa contains a set of attention ex-perts and a routing network. given an input, the moa attributes a probability to each expert by the routing network and selects the top-k experts. the output of moa is a weighted sum of the selected attention experts. the weighting mechanism allows different tokens to focus on different experts, thus improving the model's parameter and computation efficiency. experimental results show that a base-scale moa model could achieve comparable or better performance than a transformer big model. furthermore, moa could improve its performance by adding more attention experts while maintaining a relatively small computational complexity. in this way, moa can achieve comparable performance with deeper and computationally more expensive models. the interpretability analysis shows that different attention experts tend to specialize in a specific type of input tokens.