Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Mixture of Attention Heads: Selecting Attention Heads Per Token

Xiaofeng Zhang, Yikang Shen, Zeyu Huang, Jie Zhou, Wenge Rong, Zhang Xiong

TL;DR

The paper tackles the efficiency-scaling tension in Transformer models by introducing Mixture of Attention Heads (MoA), which fuses mixture-of-experts with multi-head attention. A routing network selects a sparse subset of attention experts per token, and two MoEs handle query and output projections while sharing key/value projections to reduce cost; auxiliary losses balance expert usage and stabilize routing. Empirical results on WMT14 translation and WikiText-103 MLM show MoA achieving strong performance with favorable compute and parameter efficiency, and analyses reveal perceptible head specialization and balanced expert loads. The work demonstrates a scalable, interpretable approach to expanding attention capacity without commensurate increases in computation, setting a path for larger, cheaper-to-run Transformer variants.

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. In addition to the 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.

Mixture of Attention Heads: Selecting Attention Heads Per Token

TL;DR

The paper tackles the efficiency-scaling tension in Transformer models by introducing Mixture of Attention Heads (MoA), which fuses mixture-of-experts with multi-head attention. A routing network selects a sparse subset of attention experts per token, and two MoEs handle query and output projections while sharing key/value projections to reduce cost; auxiliary losses balance expert usage and stabilize routing. Empirical results on WMT14 translation and WikiText-103 MLM show MoA achieving strong performance with favorable compute and parameter efficiency, and analyses reveal perceptible head specialization and balanced expert loads. The work demonstrates a scalable, interpretable approach to expanding attention capacity without commensurate increases in computation, setting a path for larger, cheaper-to-run Transformer variants.

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 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. In addition to the 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.
Paper Structure (33 sections, 24 equations, 4 figures, 7 tables)

This paper contains 33 sections, 24 equations, 4 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Mixture of Experts
  4. Machine Translation Architectures
  5. Specialization of Attention Heads
  6. Preliminaries
  7. Mixture of Experts
  8. Multi-head Attention
  9. Mixture of Attention Heads
  10. Routing Network
  11. Attention Expert
  12. Training Losses
  13. Computational Complexity and Number of Parameters
  14. Experiments
  15. Machine Translation
  16. ...and 18 more sections

Figures (4)

  • 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.
  • Figure 2: Mixture of Attention Heads (MoA) architecture. MoA contains two mixtures of experts. One is for query projection, the other is for output projection. These two mixture of experts select the same indices of experts. One routing network calculates the probabilities for each selected experts. The output of the MoA is the weighted sum of the outputs of each selected experts.
  • Figure 3: Experts' load percentages for encoder layer 4. Experts are indexed by their order of percentages.
  • Figure 4: Experts' load percentages for different encoder layers