Elamur Research⌕
☰
home / papers / 2024 / Associative Recurrent Memory Transformer

Associative Recurrent Memory Transformer

Paper · 2024 · Ivan Rodkin, Yuri Kuratov, Aydar Bulatov et al.

Summary

The authors address a core difficulty in long-context modeling: building a network whose cost for each new chunk of input stays fixed, however long the sequence grows. They propose the Associative Recurrent Memory Transformer (ARMT), which builds on the segment-level recurrent model RMT. Ordinary self-attention handles each segment locally. In addition, every layer gets its own associative key-value matrix. Memory tokens from the previous segment are written into this matrix, and current tokens read from it. The update resembles linear attention with the DPFP-3 feature map. A correction to the normalization vector prevents forgetting after many overwrite operations. On synthetic associative retrieval tests, ARMT holds more key-value pairs than Mamba, despite having an equal recurrent state size, and it copes with repeated rewrites. Built on GPT-2 (137M) and trained on only 16k-token inputs, it sets a record on BABILong by answering questions at contexts of up to 50 million tokens. This shows that compact recurrent memory can generalize far beyond its training length.

Key points

  • Architecture and ablations:
    • ARMT is RMT plus a layerwise associative block, using the DPFP-3 nonlinearity.
    • An ablation shows that the normalization fix is needed for hundreds of erase-insert operations.
    • The PRMT variant has layerwise memory but no associative block. It gives no gain over RMT, which indicates that the associative mechanism drives the improvement.
  • Rewrite task: ARMT was trained on 50 key-value pairs and still recalls perfectly with 500 updates, a 10x generalization.
  • BABILong:
    • On QA1, ARMT reaches 79.9% accuracy at 50M tokens.
    • On the multi-hop tasks QA2–QA5, it performs strongly up to 10M tokens.
  • Length generalization: ARMT generalizes about 60x beyond its training length (1M/16k), compared with roughly 8x for Mamba-130m (128k/16k).
  • Baselines: The comparisons cover Mamba and RMT. The authors could not train RWKV successfully. The paper appeared at an ICML 2024 workshop, and the code is released on GitHub.

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

🎓Moscow Institute of Physics and TechnologyDolgoprudny · Russia8 works · 2,761 citations🤝AIRI - Artificial Intelligence Research InstituteMoscow · Russia5 works · 29 citations🎓London Institute for Mathematical SciencesLondon · United Kingdom2 works · 8 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksComputer Vision and Pattern Recognition · 2,869 works

Cite

Cite as Ivan Rodkin, Yuri Kuratov, Aydar Bulatov and Mikhail Burtsev. 2024. Associative Recurrent Memory Transformer. arXiv:2407.04841.

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

@misc{rodkin2024associative,
  title = {{Associative Recurrent Memory Transformer}},
  author = {Ivan Rodkin, Yuri Kuratov, Aydar Bulatov and Mikhail Burtsev},
  year = {2024},
  doi = {10.48550/arxiv.2407.04841},
  eprint = {2407.04841},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2407.04841}
}

References

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

  1. Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, et al. Improving language models by retrieving from trillions of tokens. In International conference on machine learning, p… find
  2. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.
  3. Aydar Bulatov, Yuri Kuratov, and Mikhail S. Burtsev. Recurrent memory transformer, 2022. find
  4. Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models. In The Twelfth International Conference on Learning Representations, 2023.
  5. Alexis Chevalier, Alexander Wettig, Anirudh Ajith, and Danqi Chen. Adapting language models to compress contexts, 2023.
  6. Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov. 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, 2019. find
  7. Soham De, Samuel L. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Desjardins, Arnaud Doucet, David Budden, Yee Whye Teh, Razvan Pascanu, Nando De Freitas, and Caglar Gulcehre. Griffin: Mixing g… find
  8. Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models, 2023. find
  9. Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling, 2020. find
  10. Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2023. find
  11. Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2021. find
  12. Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. Retrieval augmented language model pre-training. In International conference on machine learning, pages 3929–3938. PMLR, 2020. — REALM: Retrieval-Augmented Language Model Pre-Training
  13. Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Comput., 9 (8):1735–1780, November 1997. ISSN 0899-7667. doi: 10.1162/neco.1997.9.8.1735. URL https://doi.org/10.1162/neco.1997.9.8.1735. ↗
  14. Samy Jelassi, David Brandfonbrener, Sham M. Kakade, and Eran Malach. Repeat after me: Transformers are better than state space models at copying, 2024. find
  15. Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention, 2020. find
  16. Yuri Kuratov, Aydar Bulatov, Petr Anokhin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. In search of needles in a 11m haystack: Recurrent memory finds what llms miss, 2024. find
  17. OpenAI, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, Red Avila, Igor Babuschkin, Suchir Balaji, Valerie Balcom, Paul Baltescu, Haiming Bao, Mohammad Bavarian, Jeff Belgum, Irwan Bello, Ja… — GPT-4 Technical Report
  18. Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, F… find
  19. Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. In The Twelfth International Conference on Learning Representations, 2023. find
  20. Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint, 2019. URL https://arxiv.org/abs/1911.05507. ↗
  21. Ohad Rubin and Jonathan Berant. Long-range language modeling with self-retrieval. arXiv preprint arXiv:2306.13421, 2023. — [arXiv:2306.13421](https://arxiv.org/abs/2306.13421) find
  22. Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers, 2021. find
  23. Weijia Shi, Sewon Min, Michihiro Yasunaga, Minjoon Seo, Rich James, Mike Lewis, Luke Zettlemoyer, and Wen-tau Yih. Replug: Retrieval-augmented black-box language models. arXiv preprint arXiv:2301.12652, 2023.
  24. Sainbayar Sukhbaatar, Arthur Szlam, Jason Weston, and Rob Fergus. End-to-end memory networks, 2015. find
  25. Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models, 2023.
  26. Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Tim-othée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
  27. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Advances in neural information processing systems, pages 5998–6008, 2017. URL http://papers.nips.cc/paper/7181-attention-is-all-you-need. ↗
  28. Jason Weston, Sumit Chopra, and Antoine Bordes. Memory networks. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http://arxiv.org/abs/1410.3916. ↗
  29. Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=TrjbxzRcnf-. ↗
  30. Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks, 2023. find
  31. Peitian Zhang, Zheng Liu, Shitao Xiao, Ninglu Shao, Qiwei Ye, and Zhicheng Dou. Soaring from 4k to 400k: Extending llm’s context with activation beacon, 2024. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/rodkin-2024-associative-recurrent-memory-transformer/