Mixture of Chapters: Scaling Learnt Memory in Transformers

• URL: http://arxiv.org/abs/2603.21096v1
• Date: Sun, 22 Mar 2026 07:16:58 GMT
• Title: Mixture of Chapters: Scaling Learnt Memory in Transformers
• Abstract summary: Transformers lack an explicit architectural mechanism for storing and organizing knowledge acquired during training.<n>We introduce learnable sparse memory banks: latent tokens, randomly and trained end-to-end.<n>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.
• Score: 3.479712135769016
• License: http://creativecommons.org/licenses/by/4.0/
• 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).

Related papers

• MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup [35.92693145413963]
  We introduce Mixture of Memory Embeddings (MoME), a context-aware memory mechanism.<n>MoME replaces each token's single memory row with a mixture of M slots and uses a learned gate over the hidden state to choose which slots to read at each position.<n>In pretraining experiments across nanochat, Llama-3/MobileLLM, and Qwen3 backbones, MoME improves over Value Embedding, Bigram, and STEM baselines.
  arXiv  Detail & Related papers  (2026-09-14T07:01:52Z)
• NextMem: Towards Latent Factual Memory for LLM-based Agents [58.35585202907478]
  NextMem is a latent factual memory framework that utilizes an autoregressive autoencoder to efficiently construct latent memory.<n>Extensive experiments demonstrate that NextMem achieves superior performance.
  arXiv  Detail & Related papers  (2026-02-26T14:35:27Z)
• The AI Hippocampus: How Far are We From Human Memory? [77.04745635827278]
  Implicit memory refers to the knowledge embedded within the internal parameters of pre-trained transformers.<n>Explicit memory involves external storage and retrieval components designed to augment model outputs with dynamic, queryable knowledge representations.<n>Agentic memory introduces persistent, temporally extended memory structures within autonomous agents.
  arXiv  Detail & Related papers  (2026-01-14T03:24:08Z)
• Memo: Training Memory-Efficient Embodied Agents with Reinforcement Learning [53.72709564555407]
  Memo is a transformer-based architecture and training recipe for reinforcement learning.<n>It incorporates the creation and retrieval of memory by interleaving periodic summarization tokens with the inputs of a model during training.<n>We demonstrate Memo's effectiveness on a gridworld meta-RL benchmark and a multi-object navigation task in photo-realistic indoor settings.
  arXiv  Detail & Related papers  (2025-10-22T16:24:47Z)
• Mem-α: Learning Memory Construction via Reinforcement Learning [20.916677456417464]
  Large language model (LLM) agents are constrained by limited context windows.<n>Current memory-augmented agents depend on pre-defined instructions and tools for memory updates.<n>Mem-alpha is a reinforcement learning framework that trains agents to effectively manage complex memory systems.
  arXiv  Detail & Related papers  (2025-09-30T08:02:34Z)
• Parameter-Efficient and Memory-Efficient Tuning for Vision Transformer: A Disentangled Approach [87.8330887605381]
  We show how to adapt a pre-trained Vision Transformer to downstream recognition tasks with only a few learnable parameters. We synthesize a task-specific query with a learnable and lightweight module, which is independent of the pre-trained model. Our method achieves state-of-the-art performance under memory constraints, showcasing its applicability in real-world situations.
  arXiv  Detail & Related papers  (2024-07-09T15:45:04Z)
• PARMESAN: Parameter-Free Memory Search and Transduction for Dense Prediction Tasks [5.5127111704068374]
  This work addresses flexibility in deep learning by means of transductive reasoning.<n>We propose PARMESAN, a scalable method which leverages a memory module for solving dense prediction tasks.<n>Our method is compatible with commonly used architectures and canonically transfers to 1D, 2D, and 3D grid-based data.
  arXiv  Detail & Related papers  (2024-03-18T12:55:40Z)
• A Memory Transformer Network for Incremental Learning [64.0410375349852]
  We study class-incremental learning, a training setup in which new classes of data are observed over time for the model to learn from. Despite the straightforward problem formulation, the naive application of classification models to class-incremental learning results in the "catastrophic forgetting" of previously seen classes. One of the most successful existing methods has been the use of a memory of exemplars, which overcomes the issue of catastrophic forgetting by saving a subset of past data into a memory bank and utilizing it to prevent forgetting when training future tasks.
  arXiv  Detail & Related papers  (2022-10-10T08:27:28Z)
• Mesa: A Memory-saving Training Framework for Transformers [58.78933015299703]
  We present Mesa, a memory-saving training framework for Transformers. Mesa uses exact activations during forward pass while storing a low-precision version of activations to reduce memory consumption during training. Experiments on ImageNet, CIFAR-100 and ADE20K demonstrate that Mesa can reduce half of the memory footprints during training.
  arXiv  Detail & Related papers  (2021-11-22T11:23:01Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.