SlowGuess/ACL_2020_OCR huggingface dataset:
Table 4: Language modeling performance on WikiText-103 test set (lower is better).  $\star$ Trains/evaluates with 3,072/2,048 context sizes and therefore not directly comparable to other models which use 512/480 sized ones. See Table 2 caption for the indications of other superscripts. Bold font indicates the best performance using smaller context sizes. The first two rows are due to Table 5 of Baevski and Auli (2019).
nor replicate their results, under this setting—our GPUs have far less memory, and it is impossible to even load a 3,072-token context chunk. $^{15}$  Therefore we train and evaluate MAE and UNI-MAE-7 with smaller 512/480 context sizes, also explored by Baevski and Auli (2019), which allows for a head-to-head comparison.
Results. Table 4 shows the perplexity on WikiText-103 test data. When trained under the same setting, MAE outperforms Baevski and Auli (2019) by more than 0.3 perplexity. Interestingly, despite the much smaller context at both training and test time, MAE matches the best setting by Baevski and Auli (2019). UNI-MAE-7 and NOBCD underperform the baseline (higher perplexity).
This section first empirically confirms that MAE learns to activate different experts on different inputs in §5.1. We then run a synthetic experiment to explore MAE's potential in transfer learning (§5.2).
5.1 Does MAE Learn to Specialize the Experts?
One of the appealing properties of MoE models is that they could learn to activate different experts, depending on what "expertise" is needed for the
$^{15}$  Baevski and Auli (2019) use NVIDIA Tesla V100 GPUs with 32GB memory, while we only have access to GeForce RTX 2080 Ti, with 11GB memory.
6571
Model	BLEU	Diff.
UNI-MAE-7	26.6	-
One random expert	25.8±0.2	↓0.8±0.2
NOBCD	26.7	-
Most specialized expert	26.0	↓0.7
MAE-7	27.1	-
Most specialized expert	26.8	↓0.3
input. Does MAE learn to do so? We empirically study this question, and present evidence indicating that it does, at least in part. We consider the encoders of the UNI-MAE-7, NOBCD, and the MAE-7 models trained on WMT14.[16]
We first study whether BCD training helps drifting MAE away from uniformly weighting the experts agnostic to the inputs. We treat the gating values as probabilities, and calculate their entropies:  $\mathcal{H}(\mathbf{g}) = -\sum_{i=1}^{h} g_i \cdot \log g_i$ , which are then averaged across different layers. The average entropy on the development set for MAE-7 is 1.91, lower than the 2.02 by the NOBCD model trained without BCD. In comparison, UNI-MAE-7 uniformly weights the experts and has the entropy of 2.08. This indicates that gating weights of MAE trained with BCD are more "focused" on one or a subset of experts than trained without.
Second, we study whether MAE learns to specialize different experts for different inputs. To do so we attribute the development instances to the experts that maximize the gating weights. For the first encoder layer of MAE-7, the percentages of instances attributed to each of the 8 experts are relatively balanced:  $13\%$ ,  $14\%$ ,  $9\%$ ,  $16\%$ ,  $10\%$ ,  $15\%$ ,  $10\%$ ,  $12\%$ . This suggests that all experts are assigned a substantial part of the input, and it is not the case that BCD leads to a "rich get richer" outcome.
We then continue and explore whether MAE performs reasonably well when using only the most "specialized" experts. For each development instance, we select those experts maximizing the
Table 5: Performance decrease for different models on WMT14 development set when only one expert is used for each multi-head attention layer (5.1).
Expert 1	Expert 2	Expert 3	Expert 4
neumann	bell	candidacy	veil
debuted	zero	rose	monument
rental	computing	submission	fox
worthy	decentralized	palm	unnerved
landboards	reuters	roles	remainder
Expert 5	Expert 6	Expert 7	Expert 8
spoil	menses	romans	odds
anybody	technological	sticker	heat
endorsed	inevitably	outdated	marvel
reserve	bet	analyst	ornate
pending	punk	venues	anticipating
Table 6: Indicative tokens for each expert (§5.1). Tokens attributed to Expert 2 are mostly computer science terminology; trends for other experts are less clear.
gating weights and ignore the rest, instead of linearly combining them as in Eq. 6. We see from Table 5 a 0.3 BLEU decrease under this setting. In comparison, NOBCD has a larger performance decrease of 0.7 BLEU. NOBCD's performance drop is similar to that of UNI-MAE-7, for which we randomly select an expert at each layer and average the performance over 5 runs. These results support the proposition that MAE specializes better when trained with BCD.
Finally, we search for the tokens that are more likely to activate each expert. We compute the pointwise mutual information (PMI; Church and Hanks, 1990) between tokens and experts:
$$
\operatorname {P M I} \left(\operatorname {t o k e n} _ {i}, \operatorname {e x p e r t} _ {j}\right) = \log \frac {p \left(\operatorname {t o k e n} _ {i} , \operatorname {e x p e r t} _ {j}\right)}{p \left(\operatorname {t o k e n} _ {i}\right) p \left(\operatorname {e x p e r t} _ {j}\right)}.
$$
Table 6 lists the most indicative tokens of each expert, for the first layer. While some of the terms for some experts seem loosely related (e.g., bell, reuters, and computing for expert 2, it is hard to find clear patterns in most of them.
5.2 MAE's Potential in Transfer Learning: A Case Study
We now turn to evaluate another property of MAE: its potential for data-efficient transfer learning, by only updating the gating functions, freezing the experts. We consider the pretrain-then-finetune setting. Due to computation limits, we are unable to explore MAE for pre-training contextual representations (Peters et al., 2018; Devlin et al., 2019). Rather, we focus on the following small-scale machine translation experiments.
Setting. We explore finetuning on IWSLT14 EN-DE data, a MAE model pretrained on the
16The same experiments can be done with the decoders, where the inputs to gating functions are German sentences. The authors lack German expertise, and interpretation of a following analysis would not have been possible for us.
<sup>17</sup>We observe similar trends in other layers. See Appendix C for more details.
6572
Figure 2: IWSLT14 development performance of  $\mathrm{FTG + }$  and FTALL using different amount of training data (§5.2). When trained on less than  $20\%$  subset of the original training data,  $\mathrm{FTG + }$  outperforms FTALL.
much larger WMT14 dataset. $^{18}$  We compare three finetuning methods:
- FTG finetunes the gating functions' parameters (i.e.,  $\phi$ ), keeping the rest frozen.
-  $\mathrm{FTG+}$  updates the parameter matrix  $\mathbf{W}$  in Eq. 4 in addition to  $\phi$ . The rest of the model parameters are fixed.
- FTALL updates all parameters.
As a baseline, NoFT is the out-of-box pretrained model without any finetuning. SCRATCH trains a MAE model from scratch.
Table 7 summarizes the IWSLT14 EN-DE development set performance. Surprisingly, NoFT already outperforms SCRATCH without any fin-tuning. We attribute this improvement to the larger pretraining (WMT14) data. Only updating the gating functions, FTG improves over NoFT by 0.8 BLEU. Yet there is still a significant gap of 1.8 BLEU between FTG and FTALL. Interestingly,  $\mathrm{FTG + }$  almost matches the performance of FTALL, but only updates 1/9 as many parameters. Both FTG and  $\mathrm{FTG + }$  reach the best performance after around 1K gradient updates, i.e., one epoch, significantly less than FTALL or SCRATCH.
We further compare  $\mathrm{FTG + }$  and FTALL where less downstream training data is available. To simulate this, we randomly sample  $[5\%, 10\%, 25\%, 50\%, 75\%]$  subsets of IWSLT14 training data, on which the pretrained model is finetuned. Figure 2 plots their performance. We see a clear trend: as less training data is available, the gap between  $\mathrm{FTG + }$  and FTALL decreases; when less than  $20\%$  of the training data is available,  $\mathrm{FTG + }$  outperforms FTALL. These results suggest that finetuning MAE with  $\mathrm{FTG + }$  can be viable in low-resource transfer learning.
Method	BLEU	# Params.	# Steps.
SCRATCH	28.8	41M	52K
NoFT	29.3	0	0
FTG	30.1	2M	1K
FTG+	31.6	7M	1K
FTALL	31.8	63M	12K
Table 7: IWSLT14 development set performance of different finetuning methods (§5.2). The last two columns indicate the number of parameters to update, and the number of gradient steps needed to achieve the best development performance.
Multi-head attention. An increasing amount of effort has been devoted into developing better attention mechanisms (Malaviya et al., 2018; Deng et al., 2018; Sukhbaatar et al., 2019; Correia et al., 2019; Maruf et al., 2019, inter alia), and improving transformer architectures (Shaw et al., 2018; Dehghani et al., 2019; Hao et al., 2019; Correia et al., 2019; Yang et al., 2019a, inter alia). Closely related, Iida et al. (2019) applies another attention mechanism over the attention heads, allowing a learned reweighting of them. Our work focuses on the connection between multi-head attention and MoE, and the BCD training it suggests and benefits from. Concurrent to our work, (Fan et al., 2020) study structurally pruning transformer layers for more efficient inference.
Another line of work aims to better understand the working of transformer models (Clark et al., 2019; Liu et al., 2019a; Tenney et al., 2019, inter alia).
Mixture of experts. One of the most successful applications of MoE is ensemble learning (Caruana et al., 2004; Liu et al., 2018; Dutt et al., 2017, inter alia). Recent efforts also explore MoE in sequence learning (Shazeer et al., 2017), and to promote diversity in text generation (He et al., 2018; Shen et al., 2019; Cho et al., 2019, inter alia).
We presented MAE. It is inspired by a mixture-of-experts perspective of multi-head attention. With a learned gating function, MAE activates different experts on different inputs. MAE is trained using a block coordinate descent algorithm, which alternates between updating the responsibilities of
18Here we reverse the translation direction of IWSLT14: §4.2 experimented with DE-EN, here we use EN-DE.
the experts and their parameters. Our experiments show that MAE outperforms the transformer baselines on machine translation and language modeling benchmarks. The analysis shows that MAE learns to activate different experts. The code is publicly available at https://github.com/Noahs-ARK/MAE.
We thank the anonymous reviewers, Yoav Artzi, Mandar Joshi, Jungo Kasai, Lingpeng Kong, Kenton Lee, Kelvin Luu, Will Merrill, Phoebe Mulcaire, Mark Neumann, Nikos Pappas, Ofir Press, Lianhui Qin, Swabha Swayamdipta, Vivek Srikumar, Sam Thomson, and Dani Yogatama for their helpful feedback. This work was supported in part by NSF grant 1562364, a Google Fellowship, and NVIDIA Corporation through the donation of a Tesla GPU.
Alexei Baevski and Michael Auli. 2019. Adaptive input representations for neural language modeling. In Proc. of ICLR.
Ondrej Bojar, Christian Buck, Christian Federmann, Barry Haddow, Philipp Koehn, Johannes Leveling, Christof Monz, Pavel Pecina, Matt Post, Herve Saint-Amand, Radu Soricut, Lucia Specia, and Ales Tamchyna. 2014. Findings of the 2014 workshop on statistical machine translation. In Proc. of WMT.
Rich Caruana, Alexandru Niculescu-Mizil, Geoff Crew, and Alex Ksikes. 2004. Ensemble selection from libraries of models. In Proc. of ICML.
Mauro Cettolo, Jan Niehues, Sebastian Stüker, Luisa Bentivogli, and Marcello Federico. 2014. Report on the 11th IWSLT evaluation campaign. In Proc. of IWSLT.
Jaemin Cho, Minjoon Seo, and Hannaneh Hajishirzi. 2019. Mixture content selection for diverse sequence generation. In Proc. of EMNLP.
Kenneth Ward Church and Patrick Hanks. 1990. Word association norms, mutual information, and lexicography. Computational Linguistics, 16(1):22-29.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. 2019. What does BERT look at? an analysis of BERT's attention. In Proc. of BlackBoxNLP.
Gonçalo M. Correia, Vlad Niculae, and André F.T. Martins. 2019. Adaptively sparse transformers. In Proc. of EMNLP.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Łukasz Kaiser. 2019. Universal transformers.
Yuntian Deng, Yoon Kim, Justin Chiu, Demi Guo, and Alexander Rush. 2018. Latent alignment and variational attention. In Proc. of NeurIPS.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proc. of NAACL.
Anuvabh Dutt, Denis Pellerin, and Georges Quenot. 2017. Coupled ensembles of neural networks. arXiv:1709.06053.
Sergey Edunov, Myle Ott, Michael Auli, David Grangier, and Marc'Aurelio Ranzato. 2018. Classical structured prediction losses for sequence to sequence learning. In Proc. of NAACL.
Angela Fan, Edouard Grave, and Armand Joulin. 2020. Reducing transformer depth on demand with structured dropout. In Proc. of ICLR.
Jie Hao, Xing Wang, Baosong Yang, Longyue Wang, Jinfeng Zhang, and Zhaopeng Tu. 2019. Modeling recurrence for transformer. In Proc. of NAACL.
Xuanli He, Gholamreza Haffari, and Mohammad Norouzi. 2018. Sequence to sequence mixture model for diverse machine translation. In Proc. of CoNLL.
Shohei Iida, Ryuichiro Kimura, Hongyi Cui, Po-Hsuan Hung, Takehito Utsuro, and Masaaki Nagata. 2019. Attention over heads: A multi-hop attention for neural machine translation. In Proc. of ACL: Student Research Workshop.
Sergey Ioffe and Christian Szegedy. 2015. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Proc. of ICML.
Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. 1991. Adaptive mixtures of local experts. *Neural Computation*, 3(1):79-87.
Nelson F. Liu, Matt Gardner, Yonatan Belinkov, Matthew E. Peters, and Noah A. Smith. 2019a. Linguistic knowledge and transferability of contextual representations. In Proc. of NAACL.
Xuanqing Liu, Minhao Cheng, Huan Zhang, and Chojui Hsieh. 2018. Towards robust neural networks via random self-ensemble. In Proc. of ECCV.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019b. RoBERTa: A robustly optimized BERT pretraining approach. arXiv:1907.11692.
Chaitanya Malaviya, Pedro Ferreira, and André FT Martins. 2018. Sparse and constrained attention for neural machine translation. In Proc. of ACL.
6574
Sameen Maruf, André FT Martins, and Gholamreza Haffari. 2019. Selective attention for context-aware neural machine translation. In Proc. of NAACL.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2016. Pointer sentinel mixture models. arXiv:1609.07843.
Paul Michel, Omer Levy, and Graham Neubig. 2019. Are sixteen heads really better than one? In Proc. of NeurIPS.
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. 2014. In search of the real inductive bias: On the role of implicit regularization in deep learning. In Proc. of ICLR: Worship Tack.
Myle Ott, Sergey Edunov, David Grangier, and Michael Auli. 2018. Scaling neural machine translation. In Proc. of WMT.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei jing Zhu. 2002. BLEU: a method for automatic evaluation of machine translation. In Proc. of ACL.
Matthew E. Peters, Mark Neumann, Mohit Iyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. 2018. Deep contextualized word representations. In Proc. of NAACL.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2018. Language models are unsupervised multitask learners.
Rico Sennrich, Barry Haddow, and Alexandra Birch. 2016. Neural machine translation of rare words with subword units. In Proc. of ACL.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. 2018. Self-attention with relative position representations. In Proc. of NAACL.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. 2017. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv:1701.06538.
Tianxiao Shen, Myle Ott, Michael Auli, and Marc'Aurelio Ranzato. 2019. Mixture models for diverse machine translation: Tricks of the trade. In Proc. of ICML.
Daniel Soudry and Yair Carmon. 2016. No bad local minima: Data independent training error guarantees for multilayer neural networks. arXiv:1605.08361.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. 2014. Dropout: A simple way to prevent neural networks from overfitting. JMLR, 15(1):1929-1958.
Emma Strubell, Patrick Verga, Daniel Andor, David Weiss, and Andrew McCallum. 2018. Linguistically-informed self-attention for semantic role labeling. In Proc. of EMNLP.
Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. 2019. Adaptive attention span in transformers. In Proc. of ACL.
Ian Tenney, Dipanjan Das, and Ellie Pavlick. 2019. BERT rediscovers the classical NLP pipeline. In Proc. of ACL.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Proc. of NeurIPS.
Elena Voita, David Talbot, Fedor Moiseev, Rico Senrich, and Ivan Titov. 2019. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Proc. of ACL.
Baosong Yang, Longyue Wang, Derek F. Wong, Lidia S. Chao, and Zhaopeng Tu. 2019a. Convolutional self-attention networks. In Proc. of NAACL.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime G. Carbonell, Ruslan Salakhutdinov, and Quoc V. Le. 2019b. XLNet: Generalized autoregressive pretraining for language understanding. arXiv:1906.08237.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. 2016. Understanding deep learning requires rethinking generalization. In Proc. of ICLR.
6575
A Architectures and Implementations
Our model is implemented using the PyTorch toolkit and the fairseq codebase.[19]
Machine translation with WMT'14 Our BASE model in this experiment is the transformer-base by Vaswani et al. (2017). Its encoder and decoder are both of 6 transformer layers. Each multi-head attention layer is of hidden size 512, and uses 8 attention heads; the hidden dimensions for the feed forward networks are 2,048. We follow issue #346 of the fairseq's GitHub repository to replicate the results by Vaswani et al. (2017).20 When training MAE, we mostly use the same hyperparameters, with the only exception being that we warmup the learning rate for 8,000 updates, instead of 4,000.21
At evaluation time, we apply early stopping based on development set loss, and then average the most recent 5 checkpoints of the model, following Vaswani et al. (2017).
Machine translation with IWSLT'14. The BASE model in this experiment is due to the faireq codebase.[22] It mostly follows the transformer-base architecture, but uses a larger dropout rate (0.3 vs. 0.1), a smaller feed forward network hidden size (1,024 vs. 2,048), and a larger weight decay (10 $^{-4}$  vs. 0). We use 8,000 warmup updates.
Language modeling with WikiText-103. For the BASE model, we follow the model by Baevski and Auli (2019). The learning rate is warmed up for 240,000 steps.
For all three experiments, the gating functions in our MAE model and the NOBCD baseline are implemented as tanh-MLPs. They have 256 hidden dimensions. We apply a batch normalization (Ioffe and Szegedy, 2015) to the input to the MLPs. We can see that the gating functions only have a small amount of parameters, accounting for less than  $5\%$  parameters of the full MAE model. A dropout of 0.1 is applied to the output of the first
layer. No weight decay is used.  $\phi$  are updated using SGD with a fixed learning rate 1, separate from the one for the rest part of the models. This aims to avoid using momentum-based optimizing algorithms (e.g., Adam) for the gating functions, which we empirically find helps alleviate the "rich gets richer" degeneracy.[23]
In the language modeling experiment, most recent 100 input vectors are averaged and then fed into the gating functions; while we average all the input vectors in the machine translation as the inputs to  $\mathbf{g}(\cdot ;\phi)$ .
B Learning Curve Comparison for MAE and NOBCD
In §3 (footnote 4) we discuss an overfitting issue by jointly updating the experts and the gating function. This section empirically studies it. We compare the learning curves of BASE, NOBCD, and MAE-7 trained on the IWSLT14 dataset, plotted in Figure 3. The models are described in §4.1. We tune dropout and  $\ell_2$  regularization based on development performance. Other hyperparameters are the same for the compared models.
The training loss for NOBCD decreases much faster than BASE; however, on the development set, it never outperforms BASE, and the development loss starts increasing after epoch 40. MAE-7 finds a nice middle ground in terms of training loss. It outperforms both BASE and NOBCD on the validation set. This provides further evidence for the importance of BCD training.
C Additional Results for §5.1
§5.1 describes a experiment with the MAE-7 model where we attribute the development instances of WMT14 to the experts maximizing the gating weights. Table 8 presents more results. The number of instances each expert receives is relatively balanced, and the trend is consistent across different layers.
$^{19}$ https://torch.org/; https://github. com/pytorch/fairseq  $^{20}$ https://github.com/pytorch/fairseq/ issues/346  $^{21}$ Due to the randomness in random expert selection, we find that warming up learning rate more slowly helps stabilize early training.  $^{22}$ https://github.com/pytorch/fairseq/ tree/master/examples/translation
[23]It is not entirely clear to us why using momentum-based optimization algorithms to learn the gating functions leads to degenerate solutions more often. One possible reason is that the accumulated momentum steers the gating functions to keeping selecting the experts they pick at the early stage of training.
6576
Figure 3: Learning curves of BASE, NOBCD, and MAE-7 (§B), trained on the IWSLT14 EN-DE using the same setup. NOBCD quickly fits the training data, but it does not outperform BASE on validation set. Trained with BCD, MAE finds a nice middle ground. For better readability, x-axis starts at epoch 8.
Layer	E1	E2	E3	E4	E5	E6	E7	E8
1	13.1	13.9	8.9	16.1	10.3	15.3	10.1	11.6
2	13.8	14.5	10.7	10.8	15.4	7.9	16.0	10.9
3	14.0	14.4	12.4	10.6	14.3	9.8	15.4	9.0
4	14.5	13.7	10.4	8.3	15.1	11.8	11.2	15.1
5	11.9	13.8	13.7	15.7	10.1	16.4	6.9	11.5
6	12.9	10.0	12.4	14.6	9.5	15.2	15.7	9.8
Table 8: The percentage of WMT14 development instances attributed to each of the experts in MAE-7's encoder layers (§5.1).
6577