Elamur Research⌕
☰
home / papers / 2022 / Mixture of Attention Heads: Selecting Attention Heads Per Token

Mixture of Attention Heads: Selecting Attention Heads Per Token

Paper · 2022 · Xiaofeng Zhang, Yikang Shen, Zeyu Huang et al.

Summary

Sparse Mixture-of-Experts (MoE) layers in Transformers have mostly replaced the feed-forward block. Multi-head attention has received little such treatment, even though earlier analyses showed that individual heads specialize. This paper introduces Mixture of Attention Heads (MoA), in which a pool of attention "experts" replaces the fixed set of heads. For every token, a linear-plus-softmax router picks the top-k experts, and their outputs are combined using the renormalized router weights. Each expert has its own query and output projections, while key and value projections are shared. The authors evaluate MoA on Machine Translation (including WMT14 EN-DE) and Masked Language Modeling. They report that MoA outperforms the standard Transformer, and that a base-scale MoA can match or beat Transformer big, as well as larger or deeper models, at a fraction of the compute. Adding experts increases capacity without raising MACs, because only k heads run per token. A PMI analysis shows that some experts specialize in particular token types. This makes MoA a cheap way to scale attention and a new angle on interpretability.

Key points

  • Routing details: The router uses top-k selection over N experts. The normalizing denominator of the selected weights is detached from the gradient, a trick the authors found improves routing.
  • Scaling experts on WMT14 EN-DE (k=8, head dim 256): Moving from 8 to 32 experts keeps MACs fixed at 841M while parameters rise from 87M to 200M. The 32-expert model reaches 28.8 BLEU. Raising k from 4 to 16 lifts BLEU from 27.5 to 29.4, with MACs growing from 654M to 1220M.
  • Load balance: In encoder layer 4 of the big model (16 of 32 experts active), most experts handle 2–4% of tokens. The extremes are about 5% and 1%.
  • Expert specialization: In the first encoder layer, PMI reveals experts tied to locations, names, technology, computing terms, adverbs, and reaction words. Clear patterns are not visible in deeper layers.
  • Limitations and code: The paper tests at most 64 experts, and the current implementation is slow in wall-clock time because of memory-copy overhead. Code is available at github.com/yikangshen/MoA.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Université de MontréalMontreal · Canada96 works · 92,262 citations🏢TencentShenzhen · China174 works · 21,355 citations🎓Beihang UniversityBeijing · China89 works · 9,227 citations🔬Mila - Quebec Artificial Intelligence InstituteMontreal · Canada73 works · 14,704 citations🔬State Key Laboratory of Software Development EnvironmentBeijing · China3 works · 797 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksArts and Humanities · 17,959 worksLanguage and Linguistics · 17,789 worksComputer Vision and Pattern Recognition · 3,621 works

🧭 Topics 1Research subjects and languages the work is about.Interpretability · 156 works

Cite

Cite as Xiaofeng Zhang, Yikang Shen, Zeyu Huang, Jie Zhou, Wenge Rong and Zhang Xiong. 2022. Mixture of Attention Heads: Selecting Attention Heads Per Token. arXiv:2210.05144.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2210.05144 · arXiv 2210.05144

@misc{zhang2022mixture,
  title = {{Mixture of Attention Heads: Selecting Attention Heads Per Token}},
  author = {Xiaofeng Zhang, Yikang Shen, Zeyu Huang, Jie Zhou, Wenge Rong and Zhang Xiong},
  year = {2022},
  doi = {10.48550/arxiv.2210.05144},
  eprint = {2210.05144},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2210.05144}
}

References

36 references, 18 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. 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 Proceedings of the N…
  2. Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, and Furu Wei. 2022. On the representation collapse of sparse mixture of experts. CoRR, abs/2204.09179. — [arXiv:2204.09179](https://arxiv.org/abs/2204.09179) find
  3. Kenneth Ward Church and Patrick Hanks. 1990. Word association norms, mutual information, and lexicography. Comput. Linguistics, 16(1):22–29.
  4. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies,…
  5. Nan Du, Yanping Huang, Andrew M. Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten Bosma, Zongwei Zhou, Tao Wang, Yu Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kathy Meier-Hellstern, Toju Duke, Lucas Di… — GLaM: Efficient Scaling of Language Models with Mixture-of-Experts
  6. Dheeru Dua, Shruti Bhosale, Vedanuj Goswami, James Cross, Mike Lewis, and Angela Fan. 2021. Tricks for training sparse translation models. CoRR, abs/2110.08246. find
  7. William Fedus, Barret Zoph, and Noam Shazeer. 2021. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. CoRR, abs/2101.03961. find
  8. Sepp Hochreiter and Jürgen Schmidhuber. 1997. Long short-term memory. Neural Comput., 9(8):1735– 1780. find
  9. Phu Mon Htut, Jason Phang, Shikha Bordia, and Samuel R. Bowman. 2019. Do attention heads in BERT track syntactic dependencies? CoRR, abs/1911.12246. find
  10. Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. 1991. Adaptive mixtures of local experts. Neural Comput., 3(1):79–87. find
  11. Michael I. Jordan and Robert A. Jacobs. 1994. Hierarchical mixtures of experts and the EM algorithm. Neural Comput., 6(2):181–214. find
  12. Diederik P. Kingma and Jimmy Ba. 2015. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations.
  13. Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. 2021. Gshard: Scaling giant models with conditional computation and automatic sharding. In 9th International Conference on Learning Representations. OpenReview.net. find
  14. Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. 2021. BASE layers: Simplifying training of large, sparse models. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 6265–6274. PMLR.
  15. Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. 2020a. Understanding the difficulty of training transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, pages 5747–5763. Association for Computational Linguistics.
  16. Xiaodong Liu, Kevin Duh, Liyuan Liu, and Jianfeng Gao. 2020b. Very deep transformers for neural machine translation. CoRR, abs/2008.07772. — [arXiv:2008.07772](https://arxiv.org/abs/2008.07772) find
  17. Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019. Roberta: A robustly optimized BERT pretraining approach. CoRR, abs/1907.11692.
  18. Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2016. Pointer sentinel mixture models.
  19. Xiaonan Nie, Shijie Cao, Xupeng Miao, Lingxiao Ma, Jilong Xue, Youshan Miao, Zichao Yang, Zhi Yang, and Bin Cui. 2021. Dense-to-sparse gate for mixture-of-experts. CoRR, abs/2112.14397. — [arXiv:2112.14397](https://arxiv.org/abs/2112.14397) find
  20. Myle Ott, Sergey Edunov, David Grangier, and Michael Auli. 2018. Scaling neural machine translation. In Proceedings of the Third Conference on Machine Translation: Research Papers, pages 1–9, Brussels, Belgium. Association for Computational Linguistics.
  21. Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. 2002. Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th Annual Meeting of the Association for Computational Linguistics, pages 311–318. ACL.
  22. Hao Peng, Roy Schwartz, Dianqi Li, and Noah A. Smith. 2020. A mixture of h - 1 heads is better than h heads. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 6566–6577. Association for Computational Linguistics. find
  23. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. 2019. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9. find
  24. Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2020. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., 21:140:1–140:67.
  25. Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He. 2022. Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation AI scale. CoRR, abs/2201.05596.
  26. Rico Sennrich, Barry Haddow, and Alexandra Birch. 2016. Neural machine translation of rare words with subword units. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1715– 1725, Berlin, Germany. Association for Computational Li…
  27. Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, Ryan Sepassi, and Blake A. Hechtman. 2018. Mesh-tensorflow: Deep learning for supercomputers. In Advances in Neural Information Processing Systems…
  28. Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and Jeff Dean. 2017. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In 5th International Conference on Learning Representations. OpenReview.net. find
  29. Yikang Shen, Shawn Tan, Alessandro Sordoni, Peng Li, Jie Zhou, and Aaron C. Courville. 2022. Unsupervised dependency graph network. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 4767–4784. Association for Computational Lingu… find
  30. Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. 2019. Megatron-lm: Training multi-billion parameter language models using model parallelism. CoRR, abs/1909.08053. find
  31. Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jonathon Shlens, and Zbigniew Wojna. 2016. Rethinking the inception architecture for computer vision. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, pages 2818–2826. IEEE Computer Society. find
  32. 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 Advances in Neural Information Processing Systems 30, pages 5998–6008.
  33. Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. 2019. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Proceedings of the 57th Conference of the Association for Computational Linguistics, pages 5797–5808. Association for…
  34. Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei. 2022. Deepnet: Scaling transformers to 1, 000 layers. CoRR, abs/2203.00555. find
  35. Felix Wu, Angela Fan, Alexei Baevski, Yann N. Dauphin, and Michael Auli. 2019. Pay less attention with lightweight and dynamic convolutions. In 7th International Conference on Learning Representations. OpenReview.net.
  36. Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. 2022. Designing effective sparse expert models. CoRR, abs/2202.08906. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/zhang-2022-mixture-attention-heads-selecting-attention/