# Mixture of Chapters: Scaling Learnt Memory in Transformers

arXiv: https://arxiv.org/abs/2603.21096 · HTML: https://mlsanity.com/paper/2603.21096/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 46 of its 48 references that are in the map.
- First posted on arXiv 22 Mar 2026.

## Builds on (most cited first, 30 of 46)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [Overcoming catastrophic forgetting in neural networks](https://mlsanity.com/paper/1612.00796/) (2016)
- [Flamingo: a Visual Language Model for Few-Shot Learning](https://mlsanity.com/paper/2204.14198/) (2022)
- [Toolformer: Language Models Can Teach Themselves to Use Tools](https://mlsanity.com/paper/2302.04761/) (2023)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [Perceiver: General Perception with Iterative Attention](https://mlsanity.com/paper/2103.03206/) (2021)
- [H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models](https://mlsanity.com/paper/2306.14048/) (2023)
- [Generalization through Memorization: Nearest Neighbor Language Models](https://mlsanity.com/paper/1911.00172/) (2019)
- [Compressive Transformers for Long-Range Sequence Modelling](https://mlsanity.com/paper/1911.05507/) (2019)
- [An Empirical Study of Catastrophic Forgetting in Large Language Models During Continual Fine-tuning](https://mlsanity.com/paper/2308.08747/) (2023)
- [Perceiver IO: A General Architecture for Structured Inputs & Outputs](https://mlsanity.com/paper/2107.14795/) (2021)
- [Efficient Content-Based Sparse Attention with Routing Transformers](https://mlsanity.com/paper/2003.05997/) (2020)
- [Titans: Learning to Memorize at Test Time](https://mlsanity.com/paper/2501.00663/) (2024)
- [Memorizing Transformers](https://mlsanity.com/paper/2203.08913/) (2022)
- [Recurrent Memory Transformer](https://mlsanity.com/paper/2207.06881/) (2022)
- [Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention](https://mlsanity.com/paper/2404.07143/) (2024)
- [In-context Autoencoder for Context Compression in a Large Language Model](https://mlsanity.com/paper/2307.06945/) (2023)
- [Large Memory Layers with Product Keys](https://mlsanity.com/paper/1907.05242/) (2019)
- [Long Context Compression with Activation Beacon](https://mlsanity.com/paper/2401.03462/) (2024)

## Built on by (most cited first, 1 of 1)

- [Towards In-Parameter Memory Augmentation for Large Language Models](https://mlsanity.com/paper/2610.08630/) (2026)

## Abstract (arXiv)

Transformers lack an explicit architectural mechanism for storing and organizing knowledge acquired during training. We introduce learnable sparse memory banks: a set of latent tokens, randomly initialized and trained end-to-end, that transformer layers query via cross-attention to retrieve stored knowledge. To scale memory capacity without prohibitive attention costs, we propose chapter-based routing inspired by Mixture-of-Experts architectures, partitioning the memory bank into chapters and training a router to select relevant subsets per input. This enables scaling to 262K memory tokens while maintaining tractable computation. We evaluate our approach against standard transformers (in iso-FLOP settings) on pre-training and instruction fine-tuning across relevant benchmarks. Our models surpass iso-FLOP baselines suggesting scope for a new axis of scaling, demonstrating that explicit associative memory provides complementary capacity to what is captured implicitly in model parameters. Additionally, we observe improved knowledge retention under continued training, with robustness to forgetting when transitioning between training phases (e.g., pretraining to instruction fine-tuning).