Mixture of Chapters: Scaling Learnt Memory in Transformers

2026-03-22

AI Summary

This work addresses the limitation of Transformers in lacking an explicit knowledge storage mechanism, which hinders efficient retention and retrieval of learned information. To overcome this, the authors propose a chaptered sparse memory bank, where learnable memory tokens are queried by the Transformer via cross-attention. Inspired by Mixture-of-Experts, a dynamic chapter routing strategy selectively activates relevant subsets of memory, enabling scalable knowledge access while maintaining computational efficiency. The approach expands memory capacity to 262K tokens—introducing a new scaling dimension beyond model parameters—without incurring prohibitive computational overhead. Experiments demonstrate that, under matched FLOPs, the proposed model outperforms standard Transformers in both pretraining and instruction fine-tuning tasks, while also exhibiting substantially improved knowledge retention and robustness against catastrophic forgetting in continual learning scenarios.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Natural Language Processing: Learning & Optimization for NLP
• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

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).

Problem

• Transformer
• memory
• knowledge storage
• scaling
• forgetting

Innovation

• learnable memory
• chapter-based routing
• sparse memory banks
• associative memory
• Mixture-of-Experts