Mixture of Attention Heads: Selecting Attention Heads Per Token

Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 4150 - 4162 December 7-11, 2022 ©2022 Association for Computational Linguistics Mixture of Attention Heads: Selecting Attention Heads Per Token Xiaofeng Zhang1,2∗ Yikang Shen3,4∗ Zeyu Huang1,2 Jie Zhou4 Wenge Rong1 Zhang Xiong1 1 State Key Laboratory of Software Development Environment, School of Computer Science and Engineering, Beihang University, China 2 Sino-French Engineer School, Beihang University, China 3 Mila, University of Montreal, Canada 4 Wechat AI, Tencent, China Abstract Mixture-of-Experts (MoE) networks have been proposed as an efficient way to scale up model capacity and implement conditional comput- ing. 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 dynami- cally 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 eas- ily scale up the number of attention heads and the number of parameters while preserving computational efficiency. In addition to the per- formance improvements, MoA also automati- cally differentiates heads’ utilities, providing a new perspective to discuss the model’s inter- pretability. We conducted experiments on sev- eral important tasks, including Machine Trans- lation and Masked Language Modeling. Exper- iments have shown promising results on several tasks against strong baselines that involve large and very deep models1. 1 Introduction In recent years, large models have become a popular trend in the research of Natural Lan- guage Processing, especially large-scale Trans- former (Vaswani et al., 2017). The model’s capac- ity has increased from millions of parameters (De- vlin et al., 2019; Liu et al., 2019), to billions of pa- rameters (Shoeybi et al., 2019; Raffel et al., 2020; Wang et al., 2022), even to trillions of parame- ters (Du et al., 2021; Fedus et al., 2021). How- ∗ Equal contribution. xiaofeng_z@buaa.edu.cn, yikang.shn@gmail.com 1The code can be found at https://github.com/ yikangshen/MoA. Input Top-K Router H1 H2 H3 H4 H5 H6 MoA Output Experts: Figure 1: Simple illustration of MoA. MoA consists of a set of attention heads named attention experts. For each token in the input, a Router selects k attention heads among all attention experts with different confidences. The output is a weighted sum of the selected attention heads given the confidence calculated by the Router. ever, these large-scale models demand substan- tially more computations than small-scale models. A popular trend is to utilize conditional computa- tion 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 dur- ing the forward computation, which alleviates the computational load. Among these attempts, the Mixture of Experts (MoE) (Jacobs et al., 1991; Jordan and Jacobs, 1994) is an essential technique. Since …

Content truncated to 500 words