Elamur Research
☰
home / papers / 2020 / Reformer: The Efficient Transformer

Reformer: The Efficient Transformer

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/2001.04451

Summary

Transformers get very expensive to train on long inputs, for three reasons. First, every layer keeps its activations for backpropagation, so memory grows with depth. Second, the wide feed-forward layers take up a lot of space. Third, attention needs time and memory that grow with the square of sequence length. The authors propose the Reformer, which makes three changes to fix this. Attention is approximated with locality-sensitive hashing. Queries and keys come from one shared projection and are hashed using random rotations. They are then sorted by bucket, and each position attends only within small fixed-size chunks. This cuts the cost from O(L²) to O(L log L). Several hashing rounds lower the chance of separating similar vectors. Reversible residual layers let activations be recomputed instead of stored for every layer, and feed-forward layers are processed in chunks. The authors tested the model on a synthetic task, on enwik8 with 64K-token sequences and on imagenet-64 generation with 12K-token sequences. In these tests the Reformer matches a full Transformer's quality while using much less memory and running faster on long inputs. This makes long-context models practical on modest hardware.

Key points

  • LSH attention: each vector is hashed as h(x) = argmax([xR; −xR]) with a random matrix R (Andoni et al., 2015). After sorting, each chunk of m positions attends to itself and the chunk before it.
  • Shared-QK Transformer: queries and keys use the same linear layer, and the authors report no loss in quality. A token is blocked from attending to itself unless it has no other valid target.
  • Reversible layers (Gomez et al., 2017) remove the factor N (number of layers) from activation memory. Chunking the feed-forward layers removes the d_ff factor and gives numerically identical results.
  • The memory problem in numbers: full attention over 64K tokens needs a 64K × 64K matrix, about 16GB in 32-bit floats even with batch size 1.
  • Number of hash rounds: more rounds bring results closer to full attention but cost more compute. The authors pick a value that is both efficient and close to full attention.

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

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

Cite

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

@misc{kitaev2020reformer,
  title = {{Reformer: The Efficient Transformer}},
  author = {Kitaev, Nikita and Kaiser, Łukasz and Levskaya, Anselm},
  year = {2020},
  doi = {10.48550/arxiv.2001.04451},
  eprint = {2001.04451},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2001.04451}
}

References

26 references, 6 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. Character-level language modeling with deeper self-attention. CoRR, abs/1808.04444, 2018. URL http: //arxiv.org/abs/1808.04444. find
  2. Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, Ilya P. Razenshteyn, and Ludwig Schmidt. Practical and optimal LSH for angular distance. CoRR, abs/1509.02897, 2015. URL http://arxiv. org/abs/1509.02897. — [arXiv:1509.02897](https://arxiv.org/abs/1509.02897) find
  3. Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016. URL http://arxiv.org/abs/1607.06450. ↗ our card
  4. Antoine Bordes, Nicolas Usunier, Sumit Chopra, and Jason Weston. Large-scale simple question answering with memory networks. CoRR, abs/1506.02075, 2015. URL http://arxiv.org/abs/1506.02075. ↗
  5. Sarath Chandar, Sungjin Ahn, Hugo Larochelle, Pascal Vincent, Gerald Tesauro, and Yoshua Bengio. Hierarchical memory networks. arXiv preprint arXiv:1605.07427, 2016. — [arXiv:1605.07427](https://arxiv.org/abs/1605.07427) find
  6. Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. URL https://openai.com/blog/sparse-transformers, 2019. ↗
  7. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. CoRR, abs/1810.04805, 2018. URL http://arxiv.org/abs/1810.04805. ↗
  8. Aidan N Gomez, Mengye Ren, Raquel Urtasun, and Roger B Grosse. The reversible residual network: Backpropagation without storing activations. In Advances in neural information processing systems, pp. 2214–2224, 2017. find
  9. Felix Hill, Antoine Bordes, Sumit Chopra, and Jason Weston. The goldilocks principle: Reading children’s books with explicit memory representations. CoRR, abs/1511.02301, 2015. URL http://arxiv.org/abs/1511.02301. ↗ our card
  10. Cheng-Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Noam Shazeer, Curtis Hawthorne, Andrew M Dai, Matthew D Hoffman, and Douglas Eck. Music transformer: Generating music with long-term structure. arXiv preprint arXiv:1809.04281, 2018. — [arXiv:1809.04281](https://arxiv.org/abs/1809.04281) find
  11. Guillaume Lample, Alexandre Sablayrolles, Marc’Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys. CoRR, abs/1907.05242, 2019. URL http: //arxiv.org/abs/1907.05242. — [arXiv:1907.05242](https://arxiv.org/abs/1907.05242) find
  12. Peter J. Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. Generating wikipedia by summarizing long sequences. CoRR, abs/1801.10198, 2018. URL http://arxiv.org/abs/1801.10198. ↗ our card
  13. Myle Ott, Sergey Edunov, David Grangier, and Michael Auli. Scaling neural machine translation. In Proceedings of the Third Conference on Machine Translation: Research Papers, pp. 1–9, Brussels, Belgium, October 2018. Association for Computational Linguistics. doi: 10.18653/v1/W18-6301. URL https://w… find our card
  14. Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, and Alexander Ku. Image transformer. CoRR, abs/1802.05751, 2018. URL http://arxiv.org/abs/1802.05751. — [arXiv:1802.05751](https://arxiv.org/abs/1802.05751) ↗
  15. Matt Post. A call for clarity in reporting BLEU scores. In Proceedings of the Third Conference on Machine Translation: Research Papers, pp. 186–191, Belgium, Brussels, October 2018. Association for Computational Linguistics. URL https://www.aclweb.org/anthology/W18-6319. ↗ our card
  16. Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019. find
  17. Jack W Rae, Jonathan J Hunt, Tim Harley, Ivo Danihelka, Andrew Senior, Greg Wayne, Alex Graves, and Timothy P Lillicrap. Scaling memory-augmented neural networks with sparse reads and writes. In Advances in Neural Information Processing Systems, (NIPS), 2016. find
  18. Prajit Ramachandran, Niki Parmar, Ashish Vaswani, Irwan Bello, Anselm Levskaya, and Jonathon Shlens. Stand-alone self-attention in vision models. CoRR, abs/1906.05909, 2019. URL http: //arxiv.org/abs/1906.05909. — [arXiv:1906.05909](https://arxiv.org/abs/1906.05909) find
  19. Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy P. Lillicrap. One-shot learning with memory-augmented neural networks. CoRR, abs/1605.06065, 2016. URL http://arxiv.org/abs/1605.06065. ↗
  20. Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. CoRR, abs/1804.04235, 2018. URL http://arxiv.org/abs/1804.04235. ↗
  21. Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, Ryan Sepassi, and Blake Hechtman. Mesh-tensorflow: Deep learning for supercomputers. CoRR, abs/1811.02084, 2018. URL http://arxiv.org/abs/1811.0208… ↗ our card
  22. Nimit Sharad Sohoni, Christopher Richard Aberger, Megan Leszczynski, Jian Zhang, and Christopher Ré. Low-memory neural network training: A technical report. CoRR, abs/1904.10631, 2019. URL http://arxiv.org/abs/1904.10631. — [arXiv:1904.10631](https://arxiv.org/abs/1904.10631) ↗
  23. Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. CoRR, abs/1905.07799, 2019a. URL http://arxiv.org/abs/1905.07799. ↗
  24. Sainbayar Sukhbaatar, Edouard Grave, Guillaume Lample, Hervé Jégou, and Armand Joulin. Augmenting self-attention with persistent memory. CoRR, abs/1907.01470, 2019b. URL http: //arxiv.org/abs/1907.01470. — [arXiv:1907.01470](https://arxiv.org/abs/1907.01470) find
  25. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. CoRR, 2017. URL http: //arxiv.org/abs/1706.03762. find
  26. Jason Weston, Sumit Chopra, and Antoine Bordes. Memory networks. CoRR, abs/1410.3916, 2014. URL http://arxiv.org/abs/1410.3916. ↗