Elamur Research
☰
home / papers / 2020 / Efficient Content-Based Sparse Attention with Routing Transformers

Efficient Content-Based Sparse Attention with Routing Transformers

Written for the Research library (claude-opus-5-5) from the paper text, in our own words: the license "arXiv — non-exclusive licence to distribute" does not allow redistributing the text itself. Original — arxiv.org/pdf/2003.05997

Summary

Standard self-attention scales quadratically with sequence length, which makes long documents and large images costly to model. Earlier remedies took one of two routes. Some limit each position to fixed local or strided windows regardless of content. Others pick relevant positions by content (sparsemax/entmax) but must still build the full dense attention matrix. The Routing Transformer combines the strengths of both. Using online mini-batch spherical k-means, it groups queries and keys into clusters, and each position attends only to context items in its own cluster. These routing heads are mixed with ordinary local attention heads. The design reduces attention complexity from O(n²d) to O(n^1.5 d).

The authors evaluated the model on language modeling (enwik-8, Wikitext-103, PG-19) and unconditional image generation (CIFAR-10, ImageNet-64). It outperformed comparable sparse-attention models, often with fewer layers, and set a new best result on PG-19 when trained on 8192-token sequences. The work offers a memory-efficient way to handle long contexts in generative models, and the TensorFlow code is publicly released.

Key points

  • Wikitext-103: perplexity of 15.8 versus 18.3 for the strongest comparable sparse model. ImageNet-64: 3.43 bits/dim versus 3.44, compared with roughly 3.65 for the LSH-based Reformer.
  • PG-19: a 22-layer model trained on sequences of length 8192 reached a test perplexity of 33.2, beating the previous 33.6. enwik-8 results were close to the state of the art (0.99 vs 0.98).
  • Unlike Reformer, which keeps its random hashing hyper-planes fixed, this model learns its cluster centroids during training.
  • On a TPUv3, the Local Transformer trains about 1.7× faster on PG-19 (1.231 vs 0.7236 steps/sec). The authors attribute this to missing sparse-operation support on TPUs and treat memory savings, not speed, as the main goal.
  • Jensen-Shannon divergence analysis on Wikitext-103 shows that routing heads attend to different positions than local heads, so the two head types are complementary.

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

Fields: Other

Cited by

157 papers of our library cite it (of the 140673 whose bibliographies we read); 9 of them have a page here, 148 not yet.

Cite

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

@misc{roy2020efficient,
  title = {{Efficient Content-Based Sparse Attention with Routing Transformers}},
  author = {Roy, Aurko and Saffar, Mohammad and Vaswani, Ashish and Grangier, David},
  year = {2020},
  doi = {10.48550/arxiv.2003.05997},
  eprint = {2003.05997},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2003.05997}
}

References

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

  1. Rami Al-Rfou, Dokook Choe, Noah Constant, Mandy Guo, and Llion Jones. 2019. Character-level language modeling with deeper self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 3159–3166. find
  2. Alex Auvolat, Sarath Chandar, Pascal Vincent, Hugo Larochelle, and Yoshua Bengio. 2015. Clustering is efficient for approximate maximum inner product search. arXiv preprint arXiv:1507.05910. — [arXiv:1507.05910](https://arxiv.org/abs/1507.05910) find
  3. Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. 2016. Layer normalization. arXiv preprint arXiv:1607.06450. find our card
  4. Alexei Baevski and Michael Auli. 2019. Adaptive input representations for neural language modeling. In International Conference on Learning Representations. find
  5. Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. 2015. Neural machine translation by jointly learning to align and translate. In 3rd International Conference on Learning Representations, ICLR 2015. find our card
  6. Arindam Banerjee and Joydeep Ghosh. 2004. Frequency-sensitive competitive learning for scalable balanced clustering on high-dimensional hyperspheres. IEEE Transactions on Neural Networks, 15(3):702–719. find
  7. Yoshua Bengio, Nicholas Léonard, and Aaron Courville. 2013. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432. find
  8. Mathieu Blondel, André F. T. Martins, and Vlad Niculae. 2019. Learning classifiers with fenchel-young losses: Generalized entropies, margins, and algorithms. In The 22nd International Conference on Artificial Intelligence and Statistics, AIS-TATS 2019, 16-18 April 2019, Naha, Okinawa, Japan, pages 6… find
  9. Leon Bottou and Yoshua Bengio. 1995. Convergence properties of the k-means algorithms. In Advances in neural information processing systems, pages 585–592. find
  10. Xi Chen, Nikhil Mishra, Mostafa Rohaninejad, and Pieter Abbeel. 2018. Pixelsnail: An improved autoregressive generative model. In International Conference on Machine Learning, pages 864–872. find
  11. Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. 2019. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509. find
  12. Chung-Cheng Chiu* and Colin Raffel*. 2018. Monotonic chunkwise attention. In International Conference on Learning Representations. find
  13. Kyunghyun Cho and Yoshua Bengio. 2014. Exponentially increasing the capacity-to-computation ratio for conditional computation in deep learning. arXiv preprint arXiv:1406.7362. — [arXiv:1406.7362](https://arxiv.org/abs/1406.7362) find
  14. Kyunghyun Cho, Bart van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. 2014. Learning phrase representations using rnn encoder– decoder for statistical machine translation. In Proceedings of the 2014 Conference on Empirical Methods in Natural Langu… find
  15. Jan K Chorowski, Dzmitry Bahdanau, Dmitriy Serdyuk, Kyunghyun Cho, and Yoshua Bengio. 2015. Attention-based models for speech recognition. In Advances in neural information processing systems, pages 577–585. find our card
  16. Gonçalo M Correia, Vlad Niculae, and André FT Martins. 2019. Adaptively sparse transformers. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 2174–2184. find
  17. Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov. 2019. Transformer-xl: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 2978–2988. find
  18. Ludovic Denoyer and Patrick Gallinari. 2014. Deep sequential neural network. arXiv preprint arXiv:1410.0510. — [arXiv:1410.0510](https://arxiv.org/abs/1410.0510) find
  19. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. Bert: Pre-training of deep bidirectional transformers for language understanding. In NAACL-HLT (1). find
  20. David Eigen, Marc’Aurelio Ranzato, and Ilya Sutskever. 2013. Learning factored representations in a deep mixture of experts. arXiv preprint arXiv:1312.4314. — [arXiv:1312.4314](https://arxiv.org/abs/1312.4314) find
  21. Trevor Gale, Matei Zaharia, Cliff Young, and Erich Elsen. 2020. Sparse gpu kernels for deep learning. arXiv preprint arXiv:2006.10901. find
  22. Edouard Grave, Armand Joulin, and Nicolas Usunier. 2017. Improving neural language models with a continuous cache. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net. find
  23. Alex Graves, Greg Wayne, and Ivo Danihelka. 2014. Neural turing machines. arXiv preprint arXiv:1410.5401. find
  24. Karol Gregor, Ivo Danihelka, Alex Graves, Danilo Jimenez Rezende, and Daan Wierstra. 2015. DRAW: A recurrent neural network for image generation. In Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conferen… find
  25. Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. 2020. The curious case of neural text degeneration. In International Conference on Learning Representations. find
  26. Cheng-Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Ian Simon, Curtis Hawthorne, Noam Shazeer, Andrew M Dai, Matthew D Hoffman, Monica Dinculescu, and Douglas Eck. 2018. Music transformer: Generating music with long-term structure. In International Conference on Learning Representations. find
  27. Sathish Reddy Indurthi, Insoo Chung, and Sangha Kim. 2019. Look harder: A neural machine translation model with hard attention. In Proceedings of the 57th Conference of the Association for Computational Linguistics, pages 3037–3043. find
  28. Navdeep Jaitly, Quoc V Le, Oriol Vinyals, Ilya Sutskever, David Sussillo, and Samy Bengio. 2016. An online sequence-to-sequence model using partial conditioning. In Advances in Neural Information Processing Systems, pages 5067–5075. find
  29. Diederik P. Kingma and Jimmy Ba. 2015. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings. find
  30. Durk P Kingma and Prafulla Dhariwal. 2018. Glow: Generative flow with invertible 1x1 convolutions. In Advances in Neural Information Processing Systems, pages 10215–10224. find
  31. Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. 2020. Reformer: The efficient transformer. find our card
  32. In International Conference on Learning Representations. find
  33. Guillaume Lample, Alexandre Sablayrolles, Marc’Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. 2019. Large memory layers with product keys. In Advances in Neural Information Processing Systems, pages 8548–8559. find
  34. Peter J Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. 2018. Generating wikipedia by summarizing long sequences. In International Conference on Learning Representations. find our card
  35. Xiaodong Liu, Pengcheng He, Weizhu Chen, and Jianfeng Gao. 2019. Multi-task deep neural networks for natural language understanding. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4487–4496. find our card
  36. Minh-Thang Luong, Hieu Pham, and Christopher D Manning. 2015. Effective approaches to attention-based neural machine translation. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pages 1412–1421. find
  37. Matt Mahoney. 2011. Large text compression benchmark. URL: http://www. mattmahoney. net/text/text. html. find
  38. Chaitanya Malaviya, Pedro Ferreira, and André F. T. Martins. 2018. Sparse and constrained attention for neural machine translation. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 370–376, Melbourne, Australia. Association fo… find
  39. Mikko I Malinen and Pasi Fränti. 2014. Balanced k-means for clustering. In Joint IAPR International Workshops on Statistical Techniques in Pattern Recognition (SPR) and Structural and Syntactic Pattern Recognition (SSPR), pages 32– 41. Springer. find
  40. André F. T. Martins and Julia Kreutzer. 2017. Learning what’s easy: Fully differentiable neural easy-first taggers. In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, pages 349–362, Copenhagen, Denmark. Association for Computational Linguistics. find
  41. Jacob Menick and Nal Kalchbrenner. 2018. Generating high fidelity images with subscale pixel networks and multidimensional upscaling. In International Conference on Learning Representations. find
  42. Stephen Merity, Nitish Shirish Keskar, and Richard Socher. 2018. An analysis of neural language modeling at multiple scales. arXiv preprint arXiv:1803.08240. find
  43. Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2017. Pointer sentinel mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net. find our card
  44. Aaron Van den Oord, Nal Kalchbrenner, Lasse Espeholt, Oriol Vinyals, Alex Graves, et al. 2016. Conditional image generation with pixelcnn decoders. In Advances in neural information processing systems, pages 4790–4798. find
  45. Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. 2018. Image transformer. In International Conference on Machine Learning, pages 4055–4064. find
  46. Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. 2018. Improving language understanding by generative pre-training. URL https://s3-us-west-2. amazonaws. com/openai-assets/research-covers/languageunsupervised/language understanding paper. pdf. find
  47. Jack Rae, Jonathan J Hunt, Ivo Danihelka, Timothy Harley, Andrew W Senior, Gregory Wayne, Alex Graves, and Timothy Lillicrap. 2016. Scaling memory-augmented neural networks with sparse reads and writes. In Advances in Neural Information Processing Systems, pages 3621– 3629. find
  48. Jack Rae and Ali Razavi. 2020. Do transformers need deep long-range memory? In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 7524–7529, Online. Association for Computational Linguistics. find
  49. Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. 2020. Compressive transformers for long-range sequence modelling. In International Conference on Learning Representations. find our card
  50. Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. 2018. Self-attention with relative position representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers), pages 464–468. find our card
  51. Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. find
  52. Hinton, and Jeff Dean. 2017. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net. find
  53. Noam Shazeer and Mitchell Stern. 2018. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596–4604. find
  54. Sainbayar Sukhbaatar, Édouard Grave, Piotr Bojanowski, and Armand Joulin. 2019. Adaptive attention span in transformers. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 331–335. find
  55. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in neural information processing systems, pages 5998–6008. find
  56. Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron C. Courville, Ruslan Salakhutdinov, Richard S. Zemel, and Yoshua Bengio. 2015. Show, attend and tell: Neural image caption generation with visual attention. In Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille… find our card
  57. Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. 2019. Xlnet: Generalized autoregressive pretraining for language understanding. In Advances in neural information processing systems, pages 5753–5763. find our card