Mixture of Chapters: Scaling Learnt Memory in Transformers

Tasmay Pankaj Tibrewal; Pritish Saha; Ankit Meda; Kunal Singh; Pradeep Moturi

cs.LGMar 22, 2026

Mixture of Chapters: Scaling Learnt Memory in Transformers

Authors: Tasmay Pankaj Tibrewal, Pritish Saha, Ankit Meda, Kunal Singh, Pradeep Moturi

Organizations: IIT Kharagpur · Fractal AI

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

Figures & tables

Panel of Figure 1 — Figure 1: Single transformer block with chapter-based memory routing

Model	Val loss $\downarrow$
Vanilla (backbone-only)	2.92
Vanilla (iso-FLOP)	2.86
Mixture of Chapters (MoC)	2.79

Table 1: Pretraining validation loss (lower is better).

	Vanilla (iso-FLOP)			MoC
Benchmark	Pretrain (%)	$\Delta$ (pp) $\downarrow$	IFT (%)	Pretrain (%)	$\Delta$ (pp) $\downarrow$	IFT (%)
MMLU	26.90	-0.99	25.91	27.87	-0.35	27.52
ARC-C	31.77	-6.69	25.08	31.44	-2.68	28.76
BoolQ	57.13	-6.24	50.89	61.87	+0.24	62.11
OBQA	36.20	-2.00	34.20	37.40	-2.00	35.40

Table 2: Comparison of Vanilla, Mixture of Chapters (MoC) (

\Delta

= IFT - pretrain; lower is better)

Panel of Figure 2 — Figure 2: Loss curves. Left: Pretraining. Right: IFT (Freezing the bank yields identical curves).

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Category	Representative works
Learnable internal memory bank (cross-attn / persistent tokens)	Memformer [ 53 ] , LM2 [ 26 ] , persistent memory vectors [ 44 ] , stateful memory variants [ 50 ] .
Large trainable memory layers (scalable access)	Memory Layers at Scale [ 4 ] , PKM [ 31 ] , EMAT [ 51 ] , Memorizing Transformers [ 52 ] .
Test-time / updatable memory	Titans [ 2 ] , MemoryLLM [ 48 ] , M+ [ 49 ] , Larimar [ 35 ] , EM-LLM [ 20 ] , Engram [ 9 ] .
Recurrence and compressive memory for long context	Transformer-XL [ 14 ] , RMT [ 6 , 7 ] , Compressive Transformers [ 39 ] , Infini-attention [ 38 ] , Cached Transformers [ 57 ] .
Structured/entity memories	Mention Memory [ 15 ] , Entities as Experts [ 19 ] , Facts as Experts [ 47 ] .
Latent-token architectures and context compression	Perceiver/Perceiver IO [ 25 , 24 ] , Flamingo [ 1 ] , Activation Beacon [ 56 ] , ICAE [ 21 ] , compressed-context memory [ 29 ] .

Appendix

Table 3: Literature map (representative, not exhaustive).

Model	Forward FLOPs	Backward FLOPs	Forward+Backward FLOPs
Vanilla (backbone-only)	0.356T	0.713T	1.069T
Vanilla (iso-FLOP)	0.496T	0.992T	1.487T
Mixture of Chapters (MoC)	0.459T	0.918T	1.378T

Appendix

Table 4: Per-sequence FLOPs for one forward pass and one forward+backward pass (

B=1,L=1024

Component	Setting
Backbone	Decoder-only transformer with 16 layers, hidden size 768, 12 attention heads, 4 KV heads, SwiGLU MLP with intermediate size 2304, RMSNorm, RoPE with $\theta=100000$ , tied embeddings, vocabulary size 49152
Memory placement	Memory cross-attention inserted at layers $\{2,6,10,14\}$
Memory bank	262,208 latent memory tokens, shared across memory layers
Chapter routing	4097 total chapters, including 1 shared chapter; top- $k=64$ chapters selected per sequence
Memory attention	12 memory query heads and 12 memory KV heads
Routing details	Sequence-level routing, routed scaling factor 2.5, shared/routed mixing normalized with RMSNorm

Appendix

Table 5: Main MoC architecture used in our experiments.

Setting	Pretraining	IFT
Tokenizer	HuggingFaceTB/SmolLM2-135M	HuggingFaceTB/SmolLM2-135M-Instruct
Dataset	Tasmay-Tib/fineweb-edu-10bt-split	HuggingFaceH4/ultrachat_200k
Data split	train / eval	train_sft / test_sft
Sequence length	1024	2048
Distributed setup	DDP on 8 GPUs	DDP on 4 GPUs
Batch size	32 with gradient accumulation 4	32 with gradient accumulation 1

Appendix

Table 6: Main training settings for pretraining and IFT.

Model	MMLU	ARC-C	BoolQ	OBQA
Vanilla (iso-FLOP)	0.2690	0.3177	0.5713	0.3620
MoC	0.2787	0.3144	0.6187	0.3740
Vanilla (backbone-only)	0.2746	0.3311	0.5446	0.3340

Appendix

Table 7: Benchmark accuracy after pretraining (including backbone-only).

IFT setting	MMLU	ARC-C	BoolQ	OBQA
MoC (bank frozen)	0.2754	0.2843	0.6214	0.3540
MoC (bank 1/10 LR) ∗	0.2752	0.2876	0.6211	0.3540
MoC (bank equal LR)	0.2753	0.2943	0.6211	0.3480

Appendix

Table 8: MoC IFT benchmarks for bank update settings.

Memory-Augmented Language Models

Memory-augmented language models combine persistent or temporary memory with neural computation to retain, retrieve, and use information beyond the active context.

See all 283 papers

Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models
Jan 12, 2026Xin Cheng, Rui Tian, Wangding Zeng +18Memory-Augmented Language Models Language Model Scaling Laws
Peking University · DeepSeek-AI
You Do Not Fully Utilize Transformer's Representation Capacity
Feb 13, 2025Gleb Gerasimov, Yaroslav Aksenov, Nikita Balagansky +2Representation Collapse Memory-Augmented Language Models
1T-Tech · 2Moscow Institute of Physics and Technology · 3HSE University
Learning to Recall with Transformers Beyond Orthogonal Embeddings
Mar 16, 2026Nuri Mert Vural, Alberto Bietti, Mahdi Soltanolkotabi +1Transformer Memorization in Language Models
University of Toronto and Vector Institute · Flatiron Institute · University of Southern California +1

Jan 12, 2026J/K move · Enter open · S save

Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models

Xin Cheng, Rui Tian, Wangding Zeng, Damai Dai+17

Peking University · DeepSeek-AI

Memory-Augmented Language Models Language Model Scaling Laws Language Modeling

Open paper AskPDF arXiv

While Mixture-of-Experts (MoE) scales capacity via conditional computation, Transformers lack a native primitive for knowledge lookup, forcing them to inefficiently simulate retrieval through computation. To address this, we introduce conditional memory as a complementary sparsity axis, instantiated via Engram, a module that modernizes classic

N

-gram embedding for O(1) lookup. By formulating the Sparsity Allocation problem, we uncover a U-shaped scaling law that optimizes the trade-off between neural computation (MoE) and static memory (Engram). Guided by this law, we scale Engram to 27B parameters, achieving superior performance over a strictly iso-parameter and iso-FLOPs MoE baseline. Most notably, while the memory module is expected to aid knowledge retrieval (e.g., MMLU +3.4; CMMLU +4.0), we observe even larger gains in general reasoning (e.g., BBH +5.0; ARC-Challenge +3.7) and code/math domains~(HumanEval +3.0; MATH +2.4). Mechanistic analyses reveal that Engram relieves the backbone's early layers from static reconstruction, effectively deepening the network for complex reasoning. Furthermore, by delegating local dependencies to lookups, it frees up attention capacity for global context, substantially boosting long-context retrieval (e.g., Multi-Query NIAH: 84.2 to 97.0). Finally, Engram establishes infrastructure-aware efficiency: its deterministic addressing enables runtime prefetching from host memory, incurring negligible overhead. We envision conditional memory as an indispensable modeling primitive for next-generation sparse models.

Abstract

Memory-Augmented Language Models

Explore similar work

Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models

You Do Not Fully Utilize Transformer's Representation Capacity

Learning to Recall with Transformers Beyond Orthogonal Embeddings