SlowGuess/ACL_2020_OCR huggingface dataset:
A Mixture of  $h - 1$  Heads is Better than  $h$  Heads
Hao Peng\* Roy Schwartz\* Dianqi Li\* Noah A. Smith\*
$\diamond$ Allen Institute for Artificial Intelligence
$\spadesuit$ Paul G. Allen School of Computer Science & Engineering, University of Washington
\*Department of Electrical & Computer Engineering, University of Washington {hapeng, roysch, nasmith}@cs.washington.edu, dianqili@uw.edu
Multi-head attentive neural architectures have achieved state-of-the-art results on a variety of natural language processing tasks. Evidence has shown that they are overparameterized; attention heads can be pruned without significant performance loss. In this work, we instead "reallocate" them—the model learns to activate different heads on different inputs. Drawing connections between multi-head attention and mixture of experts, we propose the mixture of attentive experts model (MAE). MAE is trained using a block coordinate descent algorithm that alternates between updating (1) the responsibilities of the experts and (2) their parameters. Experiments on machine translation and language modeling show that MAE outperforms strong baselines on both tasks. Particularly, on the WMT14 English to German translation dataset, MAE improves over "transformer-base" by 0.8 BLEU, with a comparable number of parameters. Our analysis shows that our model learns to specialize different experts to different inputs.<sup>1</sup>
The transformer architecture and its variants achieve state-of-the-art performance across a variety of NLP tasks, including machine translation (Vaswani et al., 2017; Ott et al., 2018), language modeling (Radford et al., 2018; Baevski and Auli, 2019), semantic role labeling (Strubell et al., 2018), and more (Devlin et al., 2019; Liu et al., 2019b; Yang et al., 2019b). Under the hood, multi-head attention provides the driving force: multiple separately parameterized attention functions act in parallel to contextualize the input representations; their outputs are then gathered by an affine transformation, and fed to onward computation.
Figure 1: Illustration of MAE: a mixture of attentive experts. Each  $\mathbf{H}_i$  box is an attention head in a given layer; there are  $h$  of them in total. Experts are groups of  $h - 1$  attention heads. MAE learns an input-dependent distribution of the experts (g). At each training step, a single expert is selected and updated (solid line); during the evaluation, experts' outputs are linearly combined with weights produced by g.
Recent efforts by Voita et al. (2019) and Michel et al. (2019) suggest that typical transformer networks are overparameterized, in the sense that at test time, many of the heads, or even a full layer (Fan et al., 2020), can be removed without significant loss in performance. In response to this observation, they propose to prune the unimportant attention heads in the model after it is trained, aiming for faster inference.
In this paper, we ask whether, instead of reducing the model capacity, we can use it more effectively. We propose a mixture of attentive experts (MAE). MAE retains all attention heads, and learns to activate different heads on different inputs (see illustration in Figure 1). We start by showing that multi-head attention can be seen as an uniform, input-agnostic mixture of experts (Jacobs et al., 1991), by grouping a subset ofatten
2We do not argue that overparameterization is bad for training. In fact, it may be necessary for successful optimization and good generalization (Neyshabur et al., 2014; Zhang et al., 2016; Soudry and Carmon, 2016, inter alia). Rather, we try to explore more efficient ways to use the modeling capacity, than, e.g., removing part of the model.
1Our implementation is publicly available at https:// github.com/Noahs-ARK/MAE.
Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 6566-6577 July 5 - 10, 2020. ©2020 Association for Computational Linguistics
tion heads as an expert (§2.2). We then introduce MAE, which instead of uniformly weighting the experts, complements the experts with a learned, input-dependent function that assigns their responsibilities (§2.3). To train MAE, we propose a two-step algorithm based on block coordinate descent (§3), which alternates between updating the experts' responsibilities and their parameters.
We evaluate MAE on machine translation and language modeling (§4). Our approach outperforms strong baselines on both; on the WMT14 English to German MT dataset, MAE outperforms transformer-base (Vaswani et al., 2017) by 0.8 BLEU with a negligible increase in the number parameters. Our analysis shows that MAE learns to encourage different experts to specialize on different inputs (§5).
2 MAE: Mixture of Attentive Experts
This section describes MAE in detail. It is inspired by a mixture-of-experts view of multi-head attention, which we present in §2.2. Specifically, we show that multi-head attention can be viewed as a mixture of uniformly weighted experts, each consisting of a subset of attention heads. Based on this observation, we propose MAE, which learns to weight the experts (§2.3) depending on the input. We begin by laying out notation and necessary background in §2.1.
2.1 Background: Mixture of Experts
Mixture of experts is a well-established technique for ensemble learning (Jacobs et al., 1991). It jointly trains a set of expert models  $\{\mathbf{f}_i\}_{i=1}^k$  that are intended to specialize across different input cases. The outputs produced by the experts are aggregated by a linear combination, with a "gating function"  $\mathbf{g} = [g_1, \ldots, g_k]$  determining the importance of each expert in the final decision:
$$
\operatorname {M o E} (\mathbf {x}) = \sum_ {i = 1} ^ {k} g _ {i} (\mathbf {x}) \cdot \mathbf {f} _ {i} (\mathbf {x}). \tag {1}
$$
The gating function can be parameterized by, e.g., a neural network. We will also refer to  $\mathbf{g}$  as the responsibilities or weights of the experts.
2.2 Multi-Head Attention: a Mixture-of-Experts Perspective
Multi-head attention is the key building block for the state-of-the-art transformer architectures (Vaswani et al., 2017). At its core are mul
tiple separately parameterized attention heads. An attention head takes as input a  $n$ -by-  $d$  matrix  $\mathbf{X}$ , with each row being the vector representation of an input element. It contextualizes the input using a dot-product attention mechanism:
$$
\widetilde {\mathbf {H}} _ {i} = \operatorname {s o f t m a x} \left(\mathbf {X Q} _ {i} \mathbf {K} _ {i} ^ {\top} \mathbf {X} ^ {\top}\right) \mathbf {X V} _ {i}, \quad (2)
$$
where  $\mathbf{Q}_i$ ,  $\mathbf{K}_i$ , and  $\mathbf{V}_i$  are learned matrices, and the softmax normalizes row-wise. The outputs of attention heads are then concatenated and fed through a learned affine transformation:
$$
\mathbf {Z} \triangleq \operatorname {M u l t i H e a d} (\mathbf {X}) = \left[ \widetilde {\mathbf {H}} _ {1}; \dots ; \widetilde {\mathbf {H}} _ {h} \right] \mathbf {W} \tag {3}
$$
where  $\mathbf{W}$  is a learned matrix, and  $h$  denotes the number of attention heads.
We now present a different computation equivalent to Eq. 3, aiming for a smoother transition into following sections. Let  $\mathbf{H}_i = \widetilde{\mathbf{H}}_i\mathbf{W}_i$  where  $\mathbf{W}_i$  is a block submatrix of  $\mathbf{W}$ , i.e.,  $\mathbf{W} = [\mathbf{W}_1^\top; \mathbf{W}_2^\top, \ldots; \mathbf{W}_h^\top]^\top$ . Then
$$
\mathbf {Z} = \left[ \widetilde {\mathbf {H}} _ {1}; \dots ; \widetilde {\mathbf {H}} _ {h} \right] \mathbf {W} = \sum_ {i = 1} ^ {h} \mathbf {H} _ {i}. \tag {4}
$$
Eq. 4 provides a different view of the output computation of the multi-head attention: each attention head first projects the contextualized representation with a learned matrix (i.e.,  $\mathbf{H}_i = \widetilde{\mathbf{H}}_i\mathbf{W}_i$ ), then their outputs are gathered with a sum (Eq. 4). We now show that this can be seen as a uniformly weighted mixture of experts.
A mixture-of-experts perspective. Let us take a closer look at Eq. 4 and rewrite it:
$$
\begin{array}{l} \mathbf {Z} = \frac {1}{h - 1} \sum_ {i = 1} ^ {h} (- 1 + h) \mathbf {H} _ {i} \\ = \frac {1}{h - 1} \left(- \sum_ {i = 1} ^ {h} \mathbf {H} _ {i} + \sum_ {i = 1} ^ {h} \sum_ {j = 1} ^ {h} \mathbf {H} _ {j}\right) \tag {5} \\ = \sum_ {i = 1} ^ {h} \underbrace {\frac {1}{h}} _ {\text {g a t e} g _ {i}} \underbrace {\frac {h}{h - 1} \left(- \mathbf {H} _ {i} + \sum_ {j = 1} ^ {h} \mathbf {H} _ {j}\right)} _ {\text {e x p e r t f} _ {i} (\mathbf {X}; \boldsymbol {\theta} _ {i})}. \\ \end{array}
$$
Eq. 5 interprets multi-head attention as a mixture of  $\binom{h}{h-1} = h$  experts. It first constructs a set of  $h$  experts  $\{\mathbf{f}_i(\cdot; \boldsymbol{\theta}_i)\}$ , with  $\boldsymbol{\theta}_i$  denoting  $\mathbf{f}_i$ 's param
<sup>3</sup>Some authors explicitly distinguish queries, keys, and values (Vaswani et al., 2017). These inputs can sometimes differ, e.g., in encoder-decoder attention. We suppress such differences for clarity.
6567
eters.  $\mathbf{f}_i(\cdot ;\pmb {\theta}_i)$  is a parameterized function of the input, which calculates a sum of the outputs by all but the  $i$ th attention head. This is achieved by subtracting  $\mathbf{H}_i$  from  $\sum_{j = 1}^{h}\mathbf{H}_{j}$ , then scaling up the results by  $h / (h - 1)$ . The experts share part of the parameters: any two share  $h - 2$  attention heads. A uniform responsibility of  $1 / h$  is used.
Discussion. Viewing multi-head attention through this MoE lens suggests some interesting consequences. One can replace the input-agnostic responsibility in Eq. 5 with a function over the input. Indeed, we have good reasons for doing so. Voita et al. (2019) and Michel et al. (2019) show that for transformer networks, a handful of important attention heads are sufficient to achieve good test-time performance. They propose to prune the rest using an input-agnostic procedure. Instead of doing so, here we see a potential alternative: keep all the heads, but only activate those that are important to the input. This motivates MAE, which we now introduce.
2.3 MAE: Learning to Weight Experts
MAE is inspired by the connections between MoE and multi-head attention we draw in §2.2. On top of multi-head attention, MAE learns an input-dependent parameterized gating function  $\mathbf{g}(\cdot ;\phi)$  to complement the experts. More formally, the uniform responsibility  $1 / h$  in Eq. 5 is replaced by  $\mathbf{g}(\cdot ;\phi)$ : given input  $\mathbf{X}$ , MAE outputs
$$
\sum_ {i = 1} ^ {h} g _ {i} (\mathbf {X}; \boldsymbol {\phi}) \cdot \mathbf {f} _ {i} (\mathbf {X}; \boldsymbol {\theta} _ {i}). \tag {6}
$$
Experts  $\mathbf{f}_i$  are the same as those in Eq. 5.
$\mathbf{g}(\cdot ;\phi)$  is parameterized with a multi-layer perceptron (MLP) followed by a softmax. It first averages  $\mathbf{X}$  along the row (i.e., the sequence direction), and then feeds the results through a two-layer tanh-MLP.  $\mathbf{g}(\cdot ;\phi)$  outputs a normalized  $h$  dimensional vector using a softmax, indicating the responsibilities of the experts. It can be seen as a learned probability distribution over the experts.
MAE can learn to assign more responsibility to the experts that are more important to the given input, allowing them to contribute more. MAE is applicable wherever multi-head attention is used. For example, in a machine translation experiment (§4.2), we replace with MAE all the multi-head attention in a transformer network, including the self-attention in all encoder and decoder layers, as well as those attending over the encoded source
from the decoder. Each of them is separately treated as a mixture of experts, and has its own gating function. The additional parameter overhead is small: gating functions account for only  $3 - 5\%$  parameters of the full model (Appendix A).
3 Training MAE with Block Coordinate Descent
It is straightforward to jointly train the experts and the gating functions in an MAE model using backpropagation. However, in line with previous observations (Shen et al., 2019), we empirically observe that this is prone to degenerate solutions where the gating functions tend to learn to similarly weight the experts (see §5.1).<sup>4</sup>
As a remedy, we propose a block coordinate descent (BCD) training. At a high level, training is decomposed into two interleaving steps: A G step updates the gating function  $\mathbf{g}(\cdot ;\phi)$ , fixing the experts; an F step fixes the gating function and updates one randomly selected expert  $\mathbf{f}_i(\cdot ;\pmb {\theta}_i)$ . The computations for G and F steps differ:
- In a G step, MAE outputs a linear combination of the experts' outputs, and only updates the gating function's parameters (AlGORITHM 1). No expert is updated.
- An F step computes the experts' responsibilities  $\mathbf{g}(\mathbf{X})$ , according to which an expert  $i$  is then sampled (Algorithm 2). MAE computes the output with  $\mathbf{f}_i$ , which is then updated, without updating the gating function or other experts.
A non-differentiable sampling from  $\mathbf{g}$  is involved in F steps. It does not create difficulties for the
Besides the undesired degeneracy, we also find that the model suffers worse overfitting when  $\theta$  and  $\phi$  are jointly updated (Appendix B). One possible reason is that, compared to the standard multi-head attention, the learned gates give the model additional capacity to compensate for the experts' errors with others' outputs at training time, hurting generalization (Jacobs et al., 1991). Another common degeneracy of MoEs is the "rich get richer" where one of the experts is always picked and others ignored. As observed by Voita et al. (2019), this can happen when the experts are trained to be sparsely weighted. When tuning the hyperparameters, we observe the "rich get richer" degeneracy if the learning rate is set too large.
For clarity, our discussion focuses on  $\theta$  and  $\phi$ . The rest of the model, e.g., the word embeddings in a transformer network, are updated along with  $\theta$ . Training aims to minimize loss  $\mathcal{L}$  over  $\{\theta, \phi\}$ .
In mini-batch training, which we use in the experiments, different experts can be sampled for different instances in a mini-batch. This is because  $\mathbf{g}$  depends on the inputs. This means that multiple experts will be updated in an F step, but each due to a subset of the examples in the mini-batch.
1: procedure MAEG(X)
2:  $\mathbf{Z} \gets \sum_{i=1}^{h} g_i(\mathbf{X}; \boldsymbol{\phi}) \cdot \mathbf{f}_i(\mathbf{X}; \boldsymbol{\theta}_i)$
3: Forwardprop with  $\mathbf{Z}$  and calculate  $\mathcal{L}$ .
4: Calculate  $\nabla_{\phi}\mathcal{L}$  with backprop.
5:  $\phi \gets \phi -\eta \cdot \nabla_{\phi}\mathcal{L}.$
6: end procedure
1: procedure MAEF(X)
2: Draw  $i\sim \operatorname {Cat}(\mathbf{g}(\mathbf{X};\phi))$
3:  $\mathbf{Z} \gets \mathbf{f}_i(\mathbf{X}; \boldsymbol{\theta}_i)$
4: Forwardprop with  $\mathbf{Z}$  and calculate  $\mathcal{L}$ .
5: Calculate  $\nabla_{\theta_i}\mathcal{L}$  with backprop.
6:  $\pmb{\theta}_i\gets \pmb{\theta}_i - \eta \cdot \nabla_{\pmb{\theta}_i}\mathcal{L}.$
7: end procedure
backpropagation, since an F step never calculates the gradients w.r.t.  $\phi$ . At test time, the computation is the same as that in a G step, i.e., MAE outputs a linear combination of the experts, weighted by  $\mathbf{g}$ .
Training time overhead. A straightforward training procedure is to, for each training instance, first take a G step, and then an F step. This doubles the forward propagation computation overhead. In practice, it is not necessary to take G steps as frequently as F steps, since they only update a small portion of the model. In the experiments, we take G steps one fifth as frequently as F steps: we make G updates every 5 epochs while always take F steps. In preliminary experiments, we find this reduces training time overhead without significant impact on the performance.<sup>7</sup>
Algorithm 3 summarizes the block coordinate descent training in a given epoch.
Connections to dropout. In the above block coordinate descent training algorithm, an F step samples an expert to update, and ignores the rest in both forward and backward computation. It is reminiscent of dropout (Srivastava et al., 2014). Specifically, selecting expert  $\mathbf{f}_i$  is equivalent to
dropping head  $i$ . In other words, the F steps (Algorithm 2) can be seen as a structured dropout applied to the attention heads, but with learned input-dependent drop probabilities. When  $g$  is a constant vector with elements  $1 / h$ , it recovers the head dropout, which is also explored by concurrent work (Fan et al., 2020).
So far, we view MAE as a mixture of  $h$  experts, each consisting of  $h - 1$  attention heads. One can, of course, generalize this to other settings, e.g., mixing  $\binom{h}{h-2}$  experts, each containing  $h - 2$  heads. From the dropout view, this translates to dropping more attention heads: dropping  $t$  heads out of  $h$  is equivalent to applying a dropout with drop probability  $t / h$ , in the sense that their expected numbers of dropped units are the same.
Despite the similarity between MAE and dropout, a key difference exists between the two: with the latter, the constant dropout probability is set a priori, while MAE uses a gating function  $\mathbf{g}(\cdot ;\phi)$  to calculate a learned, input-dependent dropout probability.
We empirically evaluate MAE on machine translation (§4.2) and language modeling (§4.3) benchmarks. We first introduce the compared models (§4.1).
4.1 Compared Models
MAE is evaluated under two settings:
- MAE-7 mixes 8 experts each with 7 attention heads.
Recall from Eq. 5 that  $\mathbf{f}_i$  includes all but head  $i$ .
In this way, training time for MAE is roughly 1.2 times longer than that of the transformer network it builds on.
Although we assume supervised learning, we suppress the gold outputs for notational clarity. We slightly overload the notation and denote by  $\mathbf{X}_i$  the training instance, although they cab also be the outputs of intermediate layers.
6569
- MAE-6 is similar to MAE-7, but mixes  $\binom{8}{2} = 28$  experts each with 6 attention heads. $^{10}$
We compare MAE to the following baselines.
- BASE is a sequence-to-sequence model based on the transformer architecture.
- NOBCD is the same model as MAE, but does not use block coordinate descent training. Instead, it jointly updates all experts and the gating function at training time, as discussed at the start of §3.
- UNI-MAE-7 is similar to MAE but does not have parameterized gating functions. It builds on BASE, and mixes 8 experts, each with 7 attention heads. Constant uniform responsibilities are assigned to the experts. At each training step, it updates one uniformly sampled expert; at test time, the outputs of all experts are averaged according to Eq. 5.
- UNI-MAE-6 mixes 28 6-attention-head experts, and is otherwise the same as UNI-MAE-7.
We refer the readers to Appendix A for implementation details.
4.2 Machine Translation
Datasets. We experiment with two machine translation datasets:
- WMT14 EN-DE (Bojar et al., 2014). Following previous practice (Vaswani et al., 2017) we train on WMT14, and designate newstest2013 and newstest2014 as development and test data respectively. Our preprocessing follows that of Vaswani et al. (2017) and Ott et al. (2018). A shared source-target vocabulary is used, with 32k byte pair encoding types (BPE; Sennrich et al., 2016).
- IWSLT14 DE-EN (Cettolo et al., 2014).<sup>12</sup> It is based on TED talks, and is much smaller compared to WMT14. We use the preprocessing from Edunov et al. (2018). Following previous practice, we use separate vocabularies for the source and target, with around 9K and 7K BPE types respectively.
Table 1 summarizes some statistics of the datasets.
Data	Train	Dev.	Test	Vocab.
WMT14	4.5M	3K	3K	32K
IWSLT14	160K	7K	7K	9K/7K
Table 1: Some statistics for WMT14 and IWSLT14 datasets. We use separate source and target vocabularies in IWSLT14 experiments.
Evaluation. The models are evaluated using BLEU (Papineni et al., 2002). A beam search with beam size 5 is used. In the WMT14 experiments, we follow Vaswani et al. (2017), and apply a compound split postprocessing. $^{13}$
Results. Table 2 summarizes WMT14 EN-DE translation test performance. The base and large sized transformer models are due to Vaswani et al. (2017). To control for compounding factors, we additionally compare to our implementation of the base sized model (BASE). It achieves slightly better performance than Vaswani et al. (2017), with a 0.3 BLEU edge. MAE-7 improves over the base transformer by 0.8 BLEU, obtaining similar performance to the large-size transformer of Vaswani et al. (2017) using less than a third as many parameters. Since we do not see similar improvement by UNI-MAE-7, we attribute this gain to input-dependent expert weighting. Having a smaller number of heads for each expert, MAE-6 slightly underperforms MAE-7, and so does UNI-MAE-6 in comparison to UNI-MAE-7. Finally, NOBCD gets worse performance than the transformer baseline, demonstrating the importance of the block coordinate decent training.
We observe similar trends on the IWSLT14 DEEN dataset, summarized in Table 3. The BASE model here is similar to the base-sized transformer in the WMT14 experiment, but with a smaller hidden dimension. MAE-7 outperforms BASE by 0.9 BLEU. Interestingly, UNI-MAE-7 improves over BASE by 0.3 BLEU, possibly because the regularization effect of random expert selection training helps more on this smaller dataset.[14]
4.3 Token-level Language Modeling
Dataset. We experiment with the WikiText-103 dataset (Merit et al., 2016). It contains articles
Preliminary results show that mixing experts with fewer heads leads to overwhelming performance. We conjecture this is due to too strong a regularization effect (§3).
11https://drive.google.com/a/ haopeng.name/uc?export  $\equiv$  download&id  $= 0\mathrm{B}$  _ bZck-ksdkpM25jRUN2X2UxMm8
12http://workshop2014.iwslt.org/.
$^{13}$ https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/util/get_ende_bleu.sh
14Selecting an expert can be seen dropping one attention head in training (3).
Model	BLEU	# Params.
Base Transformer	27.3	65M
Large Transformer	28.4	213M
BASE	27.6	61M
‡NOBCD	27.5	63M
†UNI-MAE-7	27.7	61M
†UNI-MAE-6	27.6	61M
‡MAE-7	28.4	63M
‡MAE-6	28.1	63M
Table 2: WMT14 EN-DE translation test performance on newstest2014. † randomly select an expert to update for each training instance, and ‡ learns a gating function to weight the experts. Transformer performance in the first two rows are due to Vaswani et al. (2017).
Model	BLEU	# Params.
BASE	34.6	39M
‡NOBCD	34.8	41M
†UNI-MAE-7	34.9	39M
†UNI-MAE-6	35.0	39M
†‡MAE-7	35.5	41M
†‡MAE-6	35.4	41M
from English Wikipedia, with a 268K-sized vocabulary. The training/development/test data respectively have 103M/218K/246K tokens.
Setting. Here the BASE model is the strong language model by Baevski and Auli (2019). It is based on a 16-layer transformer network; each multi-head attention layer has 8 heads. It uses different embedding dimensions for the tokens, based on their frequencies. We closely follow Baevski and Auli (2019) in terms of hyperparameters and training procedures. The readers are referred to their paper and Appendix A for further architecture and hyperparameter details.
Notes on context size. Baevski and Auli (2019) study the effect of context window, i.e., the number of history tokens the model attends over. They find that using larger context sizes lead to better performance (Baevski and Auli, 2019, Table 5). Their best setting uses a 3,072 training context size, and 2,048 at test time (i.e., the model has access 2,048 tokens before predicting any token at test time). However, we are not able to train MAE,
Table 3: IWSLT14 GE-DE test set performance. See Table 2 caption for indications of the superscripts.
Model	Perplexity	# Params.
*BASE (B&A, 2019)	18.70	247M
BASE (B&A, 2019)	19.03	247M
‡NOBCD	19.12	249M
†UNI-MAE-7	19.26	247M
†‡MAE-7	18.71	249M