Paper deep dive
Mixture of Chapters: Scaling Learnt Memory in Transformers
Tasmay Pankaj Tibrewal, Pritish Saha, Ankit Meda, Kunal Singh, Pradeep Moturi
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 96%
Last extracted: 3/26/2026, 2:23:33 AM
Summary
The paper introduces 'Mixture of Chapters' (MoC), a novel architecture for Transformers that incorporates a learnable, sparse memory bank. By partitioning this memory into chapters and using a router to select relevant subsets per input, the model achieves scalable memory capacity (up to 262K tokens) with tractable computation. Experiments demonstrate that MoC improves knowledge retention during continued training and instruction fine-tuning compared to standard dense Transformer baselines.
Entities (5)
Relation Signals (3)
Mixture of Chapters → augments → Transformer
confidence 100% · We augment a standard decoder-only transformer with a memory layer that provides persistent, addressable storage.
Mixture of Chapters → uses → Memory bank
confidence 100% · We introduce learnable sparse memory banks... partitioning the memory bank into chapters
Mixture of Chapters → mitigates → Catastrophic forgetting
confidence 90% · In contrast, the memory model remains stable on these tasks... This pattern is consistent with the memory bank anchoring factual knowledge during continued training.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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).
Tags
Links
- Source: https://arxiv.org/abs/2603.21096v1
- Canonical: https://arxiv.org/abs/2603.21096v1
Trouble viewing inline? Open PDF directly →
Full Text
48,265 characters extracted from source content.
Expand or collapse full text
Mixture of Chapters: Scaling Learnt Memory in Transformers Tasmay Pankaj Tibrewal1,2 Pritish Saha1 Ankit Meda1 Kunal Singh2 Pradeep Moturi2 1IIT Kharagpur 2Fractal AI tasmay.tibrewal@kgpian.iitkgp.ac.in, pritish.saha@kgpian.iitkgp.ac.in, ankitm18@kgpian.iitkgp.ac.in, kunal.singh@fractal.ai, pradeep.moturi@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). 1 Introduction Transformers provide a powerful foundation for sequence modeling [46], but they do not expose an explicit architectural mechanism for persistent, addressable memory. In practice, many systems instead rely on inference-time mechanisms such as segment-level recurrence over prior hidden states [13], KV-cache retention, eviction, or compression methods for long-context generation [54, 58, 8, 27], tool-based updates [42, 55], and retrieval-augmented generation (RAG) [34]. These methods are effective, but they primarily reuse or manage previously processed context, or retrieve external information, rather than providing a learned internal memory store for factual knowledge acquired during training. We study a complementary direction: a transformer augmented with a learned memory bank [53, 45]. The memory is a set of latent tokens, randomly initialized and trained end-to-end, that can be queried via cross-attention, where the token stream produces queries and the memory provides keys and values [53]. This design stores knowledge in a compact latent space rather than as retrieved text, and it can be integrated inside standard transformer blocks [53, 24]. The main bottleneck is scale. Dense cross-attention over a very large bank is expensive. We therefore introduce Mixture of Chapters (MoC)111Code is available at https://github.com/Tasmay-Tibrewal/Memory., which partitions the memory bank into chapters and uses a lightweight router to select a small subset of chapters per input at the sequence level [43, 18, 41]. This enables scaling learned memory capacity to large sizes while keeping computation tractable, and it fits naturally with the goals of associative memory research [4, 32]. 2 Motivation We motivate Mixture of Chapters with three goals. (1) Explicit memory. Transformers store knowledge implicitly in dense parameters [46], making it hard to inspect, edit, or scale memory independently. Prior work adds explicit memory via large key-value tables [31, 3] or attention-based memory modules [53, 16]. We follow this direction by introducing a learned, addressable memory bank that complements parametric capacity. (2) Scalable sparse access. Dense attention over large memory banks is expensive [31]. Inspired by sparse routing in MoE models [43, 17] and routing-based sparsification [40], we partition memory into chapters and route each input to a small subset, enabling predictable scaling at tractable cost. (3) Retention under continued training. Post-training shifts can cause catastrophic forgetting [30], including in LLM instruction tuning [36]. An explicit memory bank can reduce interference during continued training by anchoring factual content, which we observe empirically across training phase transitions [12]. 3 Related Work Most relevant. Our work is closest to learned, scalable internal memory modules integrated into transformers. Product Key Memory (PKM) provides a classic recipe for large trainable key-value memory with efficient sparse lookup [31]. Memory Layers at Scale shows that trainable memory layers can add substantial capacity with near-constant FLOPs and strong factual gains [4]. We build on this direction, but use a learned latent-token memory bank accessed via cross-attention and scale it with sequence-level chapter routing. Other memory mechanisms. A broad set of alternatives externalize memory at inference time or focus on efficiency and long-context handling. Transformer-XL exemplifies recurrent caching to extend effective context [14], and RAG exemplifies retrieval over an external corpus [33]. Conditional computation via routing, as in mixture-of-experts, provides the sparse-activation template we adapt for memory selection [43]. Finally, continual training can induce catastrophic forgetting [30], which motivates our retention experiments. A wider map of related directions is provided in Appendix A.1 (Table 3). 4 Contributions This paper makes the following contributions: • Learned associative memory bank in transformers. We add a persistent bank of latent memory tokens trained end-to-end and accessed via cross-attention, enabling internal retrieval in a compact latent space rather than text-based retrieval [53, 3]. • Mixture of Chapters for scalable sparse access. We partition memory into chapters and train a lightweight router that selects a small subset per input, adopting the conditional computation principle from Mixture of Experts [43, 17]. • Scaling behavior under fixed compute. We evaluate our memory architecture against standard transformers in iso-FLOP settings, complementing compute-optimal scaling discussions in language modeling [23]. • Retention under continued training. We provide evidence that explicit memory improves knowledge retention and reduces forgetting across training phase transitions, connecting to continual learning analyses in LLMs [36, 30]. RMSNormSelf-attn12H, 4KV (GQA); RoPE+RMSNormMem cross-attn12H, 12KV+RMSNormSwiGLU MLPdff=2304d_f=2304+ ResidualRouter inputPool(H)Routersoftmax(Wrr+brW_rr+b_r)Select Top-k=64k=64Memory bank MMNm=262,208,C=4,097,T=64N_m=262,208,\ C=4,097,\ T=64Sh1R1R2⋯·sR644,032 inactive chaptersChapter-based memory routingRouter InputK,VK,V Figure 1: Single transformer block with chapter-based memory routing 5 Method 5.1 Memory layer: learned bank with chapter routing We augment a standard decoder-only transformer [46] with a memory layer that provides persistent, addressable storage. The memory consists of a learned bank of latent tokens that are trained end-to-end, and accessed through cross-attention, following the general memory-augmented attention pattern [45, 53]. Let H∈ℝB×L×dH ^B× L× d be the token hidden states at a layer. We maintain a memory bank M∈ℝNm×d,M ^N_m× d, where each row is a learned memory token. The memory layer reads from M using cross-attention where the token stream produces queries and memory produces keys and values: Q Q =Norm(H)WQ, =Norm(H)W_Q, (1) K K =MWK, =MW_K, (2) V V =MWV, =MW_V, (3) MemRead(H,M) (H,M) =softmax(QK⊤dh)V. =softmax\! ( QK d_h )V. (4) The readout is added back to the token stream through a residual connection: H′=H+MemRead(H,M)H =H+\,MemRead(H,M) 5.2 Scaling the bank with chapters Attending over all NmN_m memory tokens can be expensive for large banks. We therefore partition the bank into C chapters: M=[M1;…;MC],Mc∈ℝT×d,Nm=CT.M=[M_1;…;M_C], M_c ^T× d, N_m=CT. For each input sequence, a lightweight router selects a small subset of chapters. We compute a sequence representation by pooling hidden states (for example, mean pooling): r=Pool(H)∈ℝd,r=Pool(H) ^d, then score chapters and select the top-k: p=softmax(Wrr+br),=TopK(p,k).p=softmax(W_rr+b_r), =TopK(p,k). The memory layer then attends only to the selected chapters: M=Concat(Mc:c∈),H′=H+MemRead(H,M).M_S=Concat(\M_c:c \), H =H+MemRead(H,M_S). This reduces memory attention cost from O(LNm)O(LN_m) to O(LkT)O(LkT) per memory layer while preserving attention-based associative retrieval. Algorithm 1 (see Appendix A.3) summarizes the forward pass of a single MoC memory layer, including sequence-level routing, chapter selection, and routed memory cross-attention. The routing mechanism is inspired by sparse conditional computation in mixture-of-experts models [43, 18] and routing-based sparsification [41]. 6 Experiments 6.1 Setup Models. We compare (i) Vanilla (iso-FLOP): a dense transformer baseline compute-matched to our memory model during pretraining (see Appendix A.2 for the analytic FLOPs calculation), (i) Mixture of Chapters (MoC): our memory-augmented transformer, and (i) Vanilla (backbone-only): the dense transformer backbone of the memory model with memory components removed. Training protocol. We pretrain for 9,600 steps (9.6B tokens) and then instruction fine-tune (IFT) for 2 epochs on 230M tokens (3.2k steps). This is a relatively heavy post-training budget for this model scale and it induces clear forgetting in the dense baseline. During IFT, we increase context length from 1024 to 2048. All plots and reported pretraining results correspond to the 9,600-step run (not a continued run). Evaluation. We report accuracy on MMLU [22], ARC-Challenge [11], BoolQ [10], and OpenBookQA [37]. Random guessing is approximately 0.25 on ARC-Challenge and 0.50 on BoolQ. 6.2 Results Pretraining loss. Table 1 shows validation loss at the end of pretraining. The memory model attains the best loss, and Figure 2 shows it continues to separate late in training, suggesting additional headroom with longer runs. Table 1: Pretraining validation loss (lower is better). Model Val loss ↓ Vanilla (backbone-only) 2.92 Vanilla (iso-FLOP) 2.86 Mixture of Chapters (MoC) 2.79 Benchmarks and forgetting under heavy IFT. Table 2 reports benchmark accuracy after pretraining and after IFT, along with deltas (IFT minus pretrain). Vanilla is strongly affected on knowledge-heavy tasks: ARC-Challenge drops by 0.0669 (0.3177 → 0.2508, approaching random) and BoolQ drops by 0.0624 (0.5713 → 0.5089, approaching random). In contrast, the memory model remains stable on these tasks: BoolQ is effectively unchanged (+0.0024), and ARC-Challenge drops only 0.0268, less than half the vanilla degradation. This pattern is consistent with the memory bank anchoring factual knowledge during continued training. Table 2: Comparison of Vanilla, Mixture of Chapters (MoC) (Δ = IFT - pretrain; lower is better) Vanilla (iso-FLOP) MoC Benchmark Pretrain (%) Δ (p) ↓ IFT (%) Pretrain (%) Δ (p) ↓ 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 IFT loss and freezing the memory bank. Figure 2 (right) compares IFT validation loss for memory variants where the bank is frozen, trained with a very small learning rate, or trained with the same learning rate as the backbone. The curves overlap closely, and post-IFT benchmark scores remain within noise across these settings (Appendix A.4). This indicates the pretrained memory bank can be reused during post-training without further bank updates. Pretraining distribution after IFT. Evaluating IFT checkpoints on the pretraining validation distribution yields higher loss for the memory model (3.2024) than vanilla (3.0574), despite better retention on knowledge benchmarks. This suggests the memory model can shift further toward instruction alignment while preserving factual knowledge, consistent with a separation of roles between backbone adaptation and memory anchoring. (Left) Pretraining eval loss (9,600 steps). (Right) IFT eval loss for bank-freeze and LR variants. Figure 2: Loss curves. Left: Pretraining. Right: IFT (Freezing the bank yields identical curves). 7 Conclusion We introduced Mixture of Chapters (MoC), a memory-augmented transformer that adds a learned latent memory bank and scales access via sequence-level chapter routing, providing an explicit associative memory substrate with sparse, tractable access. In a 9.6B-token pretraining run, the memory model achieves lower validation loss than compute-matched dense baselines and improves knowledge-heavy benchmark performance. Under heavy post-training (2-epoch IFT on 230M tokens with context length increased from 1024 to 2048), the dense baseline exhibits substantial forgetting on ARC-Challenge and BoolQ, while the memory model remains comparatively stable; freezing the bank during IFT performs on par with updating it. Overall, these results point to learned memory as a complementary axis of scaling and motivate further study of larger banks, longer training, and improved routing/organization for specialization and interpretability. References [1] J. Alayrac, J. Donahue, P. Luc, A. Miech, I. Barr, Y. Hasson, et al. (2022) Flamingo: a visual language model for few-shot learning. Note: arXiv preprint arXiv:2204.14198 Cited by: Table 3. [2] A. Behrouz, P. Zhong, and V. Mirrokni (2025) Titans: learning to memorize at test time. External Links: 2501.00663, Link Cited by: §A.1, Table 3. [3] V. Berges, B. Oğuz, D. Haziza, W. Yih, L. Zettlemoyer, and G. Ghosh (2024) Memory layers at scale. arXiv preprint arXiv:2412.09764. Cited by: §2, 1st item. [4] V. Berges, B. Oguz, D. Haziza, W. Yih, L. Zettlemoyer, and G. Ghosh (2025) Memory layers at scale. In Proceedings of the 42nd International Conference on Machine Learning, External Links: Link Cited by: §A.1, Table 3, §1, §3. [5] S. Borgeaud, A. Mensch, J. Hoffmann, T. Cai, E. Rutherford, K. Millican, et al. (2021) Improving language models by retrieving from trillions of tokens. Note: arXiv preprint arXiv:2112.04426 Cited by: Table 3. [6] A. Bulatov, Y. Kuratov, and M. S. Burtsev (2022) Recurrent memory transformer. Note: arXiv preprint arXiv:2207.06881 Cited by: Table 3. [7] A. Bulatov, Y. Kuratov, and M. S. Burtsev (2023) Scaling transformer to 1M tokens and beyond with RMT. Note: arXiv preprint arXiv:2304.11062 Cited by: Table 3. [8] L. Chen, Z. Zhang, Y. Sheng, L. Dong, W. Xu, L. Zheng, B. Zhuang, Y. Cao, Z. Wang, and B. Chen (2024) NaCl: a general and effective KV cache eviction framework for LLMs. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics, External Links: Link Cited by: Table 3, §1. [9] X. Cheng et al. (2026) Engram: conditional memory via scalable lookup: a new axis of sparsity for large language models. External Links: 2601.07372, Link Cited by: Table 3. [10] C. Clark, K. Lee, M. Chang, T. Kwiatkowski, M. Collins, and K. Toutanova (2019) BoolQ: exploring the surprising difficulty of natural yes/no questions. In NAACL-HLT, External Links: Link Cited by: §6.1. [11] P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord (2018) Think you have solved question answering? try ARC, the AI2 reasoning challenge. arXiv preprint arXiv:1803.05457. External Links: Link Cited by: §6.1. [12] A. Cossu, T. Tuytelaars, A. Carta, L. Passaro, V. Lomonaco, and D. Bacciu (2022) Continual pre-training mitigates forgetting in language and vision. arXiv preprint arXiv:2205.09357. Cited by: §2. [13] Z. Dai, Z. Yang, Y. Yang, J. Carbonell, Q. Le, and R. Salakhutdinov (2019) Transformer-XL: attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, External Links: Link Cited by: §1. [14] Z. Dai, Z. Yang, Y. Yang, J. Carbonell, Q. Le, and R. Salakhutdinov (2019) Transformer-XL: attentive language models beyond a fixed-length context. Note: arXiv preprint arXiv:1901.02860 Cited by: §A.1, Table 3, §3. [15] M. de Jong, Y. Zemlyanskiy, N. FitzGerald, F. Sha, and W. Cohen (2021) Mention memory: incorporating textual knowledge into transformers through entity mention attention. External Links: 2110.06176, Link Cited by: Table 3. [16] M. de Jong, Y. Zemlyanskiy, N. FitzGerald, F. Sha, and W. Cohen (2021) Mention memory: incorporating textual knowledge into transformers through entity mention attention. arXiv preprint arXiv:2110.06176. Cited by: §2. [17] W. Fedus, B. Zoph, and N. Shazeer (2021) Switch transformers: scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2101.03961. Cited by: §2, 2nd item. [18] W. Fedus, B. Zoph, and N. Shazeer (2022) Switch transformers: scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research 23 (120), p. 1–39. External Links: Link Cited by: §A.1, §1, §5.2. [19] T. Fevry, L. B. Soares, N. FitzGerald, E. Choi, and T. Kwiatkowski (2020) Entities as experts: sparse memory access with entity supervision. Note: arXiv preprint arXiv:2004.07202 Cited by: Table 3. [20] N. Fountas, N. De Cao, T. Wolff, and M. Peyrard (2025) Human-inspired episodic memory for infinite context LLMs. External Links: 2407.09450, Link Cited by: Table 3. [21] T. Ge, J. Hu, L. Li, X. Qiu, X. Huang, and L. Si (2023) In-context autoencoder for context compression in a large language model. Note: arXiv preprint arXiv:2307.06945 Cited by: Table 3. [22] D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt (2021) Measuring massive multitask language understanding. In International Conference on Learning Representations, External Links: Link Cited by: §6.1. [23] J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de Las Casas, L. A. Hendricks, J. Welbl, A. Clark, et al. (2022) Training compute-optimal large language models. arXiv preprint arXiv:2203.15556. Cited by: 3rd item. [24] A. Jaegle, S. Borgeaud, J. Alayrac, C. Doersch, C. Ionescu, D. Ding, S. Koppula, D. Zoran, A. Brock, E. Shelhamer, O. Hénaff, M. Botvinick, A. Zisserman, O. Vinyals, and J. Carreira (2021) Perceiver io: a general architecture for structured inputs and outputs. arXiv preprint arXiv:2107.14795. External Links: Link Cited by: Table 3, §1. [25] A. Jaegle, F. Gimeno, A. Brock, A. Zisserman, O. Vinyals, and J. Carreira (2021) Perceiver: general perception with iterative attention. Note: arXiv preprint arXiv:2103.03206 Cited by: Table 3. [26] J. Kang, W. Wu, F. Christianos, A. J. Chan, F. Greenlee, G. Thomas, M. Purtorab, and A. Toulis (2025) LM2: large memory models. External Links: 2502.06049, Link Cited by: Table 3. [27] M. Karami, A. Behrouz, P. Kacham, and V. Mirrokni (2025) Trellis: learning to compress key-value memory in attention models. arXiv preprint arXiv:2512.23852. Note: In COLM 2025 External Links: Link Cited by: Table 3, §1. [28] U. Khandelwal, O. Levy, D. Jurafsky, L. Zettlemoyer, and M. Lewis (2019) Generalization through memorization: nearest neighbor language models. arXiv preprint arXiv:1911.00172. External Links: Link Cited by: Table 3. [29] J. H. Kim, J. Ma, L. Lausen, C. Tan, J. Zhang, and Y. Cui (2024) Compressed context memory for online language model interaction. Note: arXiv preprint arXiv:2312.03414 Cited by: Table 3. [30] J. Kirkpatrick, R. Pascanu, N. Rabinowitz, J. Veness, G. Desjardins, A. A. Rusu, K. Milan, J. Quan, T. Ramalho, A. Grabska-Barwinska, D. Hassabis, C. Clopath, D. Kumaran, and R. Hadsell (2017) Overcoming catastrophic forgetting in neural networks. In Proceedings of the National Academy of Sciences, External Links: Link Cited by: §A.1, §2, §3, 4th item. [31] G. Lample, A. Sablayrolles, M. Ranzato, L. Denoyer, and H. Jégou (2019) Large memory layers with product keys. arXiv preprint arXiv:1907.05242. Cited by: §A.1, Table 3, §2, §2, §3. [32] G. Lample, A. Sablayrolles, M. Ranzato, L. Denoyer, and H. Jégou (2019) Large memory layers with product keys. In Advances in Neural Information Processing Systems, External Links: Link Cited by: §1. [33] P. Lewis, E. Perez, A. Piktus, F. Petroni, V. Karpukhin, N. Goyal, H. Küttler, M. Lewis, W. Yih, T. Rocktäschel, S. Riedel, and D. Kiela (2020) Retrieval-augmented generation for knowledge-intensive NLP tasks. arXiv preprint arXiv:2005.11401. Cited by: §A.1, Table 3, §3. [34] P. Lewis, E. Perez, A. Piktus, F. Petroni, V. Karpukhin, N. Goyal, H. Küttler, M. Lewis, W. Yih, T. Rocktäschel, S. Riedel, and D. Kiela (2020) Retrieval-augmented generation for knowledge-intensive NLP tasks. In Advances in Neural Information Processing Systems, External Links: Link Cited by: §1. [35] P. Liu, Y. Zhang, and D. S. Weld (2024) Larimar: large language models with episodic memory control. External Links: 2403.11901, Link Cited by: Table 3. [36] Y. Luo, Z. Yang, F. Meng, Y. Li, J. Zhou, and Y. Zhang (2023) An empirical study of catastrophic forgetting in large language models during continual fine-tuning. arXiv preprint arXiv:2308.08747. Cited by: §2, 4th item. [37] T. Mihaylov, P. Clark, T. Khot, and A. Sabharwal (2018) Can a suit of armor conduct electricity? a new dataset for open book question answering. In EMNLP, External Links: Link Cited by: §6.1. [38] T. Munkhdalai, M. Faruqui, and S. Gopal (2024) Leave no context behind: efficient infinite context transformers with infini-attention. Note: arXiv preprint arXiv:2404.07143 Cited by: Table 3. [39] J. W. Rae, A. Potapenko, S. M. Jayakumar, C. Hillier, and T. P. Lillicrap (2020) Compressive transformers for long-range sequence modelling. Note: arXiv preprint arXiv:1911.05507 Cited by: Table 3. [40] A. Roy, M. Saffar, A. Vaswani, and D. Grangier (2020) Efficient content-based sparse attention with routing transformers. arXiv preprint arXiv:2003.05997. Cited by: §2. [41] A. Roy, M. Saffar, A. Vaswani, and D. Grangier (2021) Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics 9, p. 53–68. External Links: Link Cited by: §1, §5.2. [42] T. Schick, J. Dwivedi-Yu, R. Dessì, R. Raileanu, M. Lomeli, L. Zettlemoyer, N. Cancedda, and T. Scialom (2023) Toolformer: language models can teach themselves to use tools. arXiv preprint arXiv:2302.04761. External Links: Link Cited by: §1. [43] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean (2017) Outrageously large neural networks: the sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538. External Links: Link Cited by: §A.1, §1, §2, §3, 2nd item, §5.2. [44] S. Sukhbaatar, E. Grave, P. Bojanowski, and A. Joulin (2019) Augmenting self-attention with persistent memory. Note: arXiv preprint arXiv:1907.01470 Cited by: Table 3. [45] S. Sukhbaatar, E. Grave, G. Lample, H. Jégou, and A. Joulin (2019) Augmenting self-attention with persistent memory. arXiv preprint arXiv:1907.01470. External Links: Link Cited by: §1, §5.1. [46] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin (2017) Attention is all you need. In Advances in Neural Information Processing Systems, Cited by: §1, §2, §5.1. [47] P. Verga, H. Sun, L. B. Soares, and W. Cohen (2020) Facts as experts: adaptable and interpretable neural memory over symbolic knowledge. Note: arXiv preprint arXiv:2007.00849 Cited by: Table 3. [48] Y. Wang, Y. Gao, X. Chen, H. Jiang, S. Li, J. Yang, Q. Yin, Z. Li, X. Li, B. Yin, J. Shang, and J. McAuley (2024) MEMORYLLM: towards self-updatable large language models. External Links: 2402.04624, Link Cited by: Table 3. [49] Y. Wang et al. (2025) M+: extending MEMORYLLM with scalable long-term memory. External Links: 2502.00592, Link Cited by: Table 3. [50] Q. Wu and Z. Yu (2022) Stateful memory-augmented transformers for efficient dialogue modeling. Note: arXiv preprint arXiv:2209.07634 Cited by: Table 3. [51] Y. Wu, Y. Zhao, B. Hu, P. Minervini, P. Stenetorp, and S. Riedel (2022) An efficient memory-augmented transformer for knowledge-intensive NLP tasks. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, Note: arXiv:2210.16773 Cited by: Table 3. [52] Y. Wu, M. N. Rabe, D. Hutchins, and C. Szegedy (2022) Memorizing transformers. arXiv preprint arXiv:2203.08913. Cited by: Table 3. [53] Y. Wu, M. N. Rabe, D. Hutchins, and C. Szegedy (2020) Memformer: a memory-augmented transformer for sequence modeling. arXiv preprint arXiv:2010.06891. External Links: Link Cited by: §A.1, Table 3, §1, §2, 1st item, §5.1. [54] G. Xiao, Y. Tian, B. Chen, S. Han, and M. Lewis (2024) Efficient streaming language models with attention sinks. In International Conference on Learning Representations, External Links: Link Cited by: Table 3, §1. [55] S. Yao, J. Zhao, D. Yu, N. Du, I. Shafran, K. Narasimhan, and Y. Cao (2023) ReAct: synergizing reasoning and acting in language models. In International Conference on Learning Representations, External Links: Link Cited by: §1. [56] P. Zhang, H. Qian, Q. Ye, and Z. Dou (2024) Long context compression with activation beacon. Note: arXiv preprint arXiv:2401.03462 Cited by: Table 3. [57] Z. Zhang, W. Shao, Y. Ge, X. Wang, J. Gu, and P. Luo (2023) Cached transformers: improving transformers with differentiable memory cache. Note: arXiv preprint arXiv:2312.12742 Cited by: Table 3. [58] Z. Zhang, Y. Sheng, T. Zhou, T. Chen, L. Zheng, R. Cai, Z. Song, Y. Tian, C. Ré, C. Barrett, Z. Wang, and B. Chen (2023) H2O: heavy-hitter oracle for efficient generative inference of large language models. In Advances in Neural Information Processing Systems, External Links: Link Cited by: Table 3, §1. Appendix A Appendix A.1 Extended literature review (concise map) Due to the page limit, the main paper focuses on the closest comparisons (Titans, Memory Layers at Scale, PKM). This appendix provides a compact narrative overview of broader related directions, complemented by Table 3. Inference-time and external memory. A common strategy is to externalize memory at inference time, either by caching past activations to improve efficiency and extend effective context, or by retrieving text from an external index. Transformer-XL is a canonical example of recurrence and caching [14], while RAG exemplifies retrieval-conditioned generation [33]. These approaches are strong baselines, but the stored content remains outside the model and must be fetched and integrated per query. Learned memory inside the model. A complementary line integrates memory as a trainable component. Memformer reads and writes through an internal memory mechanism coupled to token representations [53]. For scaling, PKM provides efficient large key-value lookup [31], and Memory Layers at Scale shows that trainable memory layers can add substantial capacity with near-constant FLOPs and strong factual gains [4]. Our work is closest to this direction, but uses a learned latent-token bank accessed by cross-attention and scales access via sequence-level chapter routing. Updatable and test-time memory. Some methods emphasize explicit test-time updates or state that evolves across interactions, such as Titans [2]. These approaches often target persistent personalization or long-horizon interaction, whereas our focus is on scaling learned associative memory trained end-to-end and analyzing retention under continued training. Sparse routing and conditional computation. Routing and conditional computation provide a general template for scaling by activating only a subset of components. MoE systems route inputs to a small subset of experts [43, 18]. Our method borrows this principle but applies routing to memory selection, with chapters serving as routed memory units. Retention and catastrophic forgetting. Catastrophic forgetting is a classical issue in sequential training [30], and has been measured in continual or sequential tuning of language models. Our experiments connect this literature to architectural memory by testing whether a learned bank helps preserve factual knowledge across the pretraining-to-IFT transition. Table 3: Literature map (representative, not exhaustive). 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]. Retrieval and KV-cache baselines RAG [33], RETRO [5], kNN-LM [28], KV-cache retention/compression (H2O, StreamingLLM, NACL, Trellis) [58, 54, 8, 27]. A.2 Analytic FLOPs calculation In this section, we estimate FLOPs analytically for a single sequence with batch size B=1B=1 and sequence length L=1024L=1024. Our codebase includes a reference implementation, estimate_flops.py, which follows the same counting rules used here for complete training runs. The derivation below mirrors that implementation, while making the layer-by-layer calculation explicit. Linear layers use FLOPslinear=2Ndindout,FLOPs_linear=2\,N\,d_in\,d_out, attention matmuls use FLOPsattnmatmul=4BLqLkd,FLOPs_attn\ matmul=4\,B\,L_q\,L_k\,d, RMSNorm uses FLOPsRMSNorm=N(4d+4),FLOPs_RMSNorm=N(4d+4), and the backward pass is approximated as 2×2× the forward pass, consistent with the estimator. For our backbone, d=768d=768, dff=2304d_f=2304, h=12h=12, hkv=4h_kv=4, so the self-attention KV dimension is dkv=hkv⋅dh=4⋅64=256.d_kv=h_kv· dh=4· 64=256. Standard transformer layer. A standard layer contains self-attention, two RMSNorms, one SwiGLU MLP, RoPE, and two residual additions. Self-attention projections and attention. Q Q :2⋅1024⋅768⋅768=1,207,959,552 :2· 1024· 768· 768=1,207,959,552 (5) K K :2⋅1024⋅768⋅256=402,653,184 :2· 1024· 768· 256=402,653,184 (6) V V :2⋅1024⋅768⋅256=402,653,184 :2· 1024· 768· 256=402,653,184 (7) O O :2⋅1024⋅768⋅768=1,207,959,552 :2· 1024· 768· 768=1,207,959,552 (8) QK⊤,AV QK ,\ AV :4⋅1⋅1024⋅1024⋅768=3,221,225,472. :4· 1· 1024· 1024· 768=3,221,225,472. (9) The estimator also includes attention softmax, masking, and scaling: 12⋅1024⋅1024⋅(5+2)=88,080,384.12· 1024· 1024·(5+2)=88,080,384. Hence total self-attention FLOPs per standard layer are 6,530,531,328.6,530,531,328. RoPE. 3⋅1024⋅(768+256)=3,145,728.3· 1024·(768+256)=3,145,728. Two RMSNorms. 2⋅1024⋅(4⋅768+4)=6,299,648.2· 1024·(4· 768+4)=6,299,648. SwiGLU MLP. Wup W_up :2⋅1024⋅768⋅2304=3,623,878,656 :2· 1024· 768· 2304=3,623,878,656 (10) Wgate W_gate :2⋅1024⋅768⋅2304=3,623,878,656 :2· 1024· 768· 2304=3,623,878,656 (11) Wdown W_down :2⋅1024⋅2304⋅768=3,623,878,656 :2· 1024· 2304· 768=3,623,878,656 (12) SwiGLU activation :1024⋅2304⋅5=11,796,480. :1024· 2304· 5=11,796,480. (13) So total MLP FLOPs per layer are 10,883,432,448.10,883,432,448. Residual adds. 2⋅1024⋅768=1,572,864.2· 1024· 768=1,572,864. Putting these together, one standard transformer layer costs 6,530,531,328+3,145,728+6,299,648+10,883,432,448+1,572,864=17,424,982,0166,530,531,328+3,145,728+6,299,648+10,883,432,448+1,572,864=17,424,982,016 FLOPs. MoC memory layer. An MoC memory layer adds router computation, memory preprocessing, memory cross-attention, an extra RMSNorm, and an extra residual add. The memory bank has 262,208262,208 tokens, partitioned into 40974097 chapters, giving T=262,208/4097=64T=262,208/4097=64 tokens per chapter. With one shared chapter and top-k=64k=64 routed chapters, the selected memory tokens per sequence are Nsel=(1+64)⋅64=4160.N_sel=(1+64)· 64=4160. Router forward. For sequence-level routing, the estimator includes sequence pooling, router projection, softmax, and top-k: pooling :1⋅768⋅(1024−1)+768=786,432 :1· 768·(1024-1)+768=786,432 (14) router linear :2⋅1⋅768⋅4097=6,292,992 :2· 1· 768· 4097=6,292,992 (15) router softmax :4097⋅5=20,485 :4097· 5=20,485 (16) top-k -k :4097⋅⌈log264⌉=24,582. :4097· _264 =24,582. (17) Thus router forward FLOPs are 7,124,491.7,124,491. Router auxiliary-loss compute. The estimator also includes the auxiliary routing-loss terms (load-balancing, z-loss, entropy, and related reductions), which contribute 331,859331,859 FLOPs per memory layer. Memory preprocessing. The estimator includes chapter weighting plus normalization of the selected memory tokens: weighting :1⋅4160⋅768=3,194,880 :1· 4160· 768=3,194,880 (18) RMSNorm :4160⋅(4⋅768+4)=12,796,160. :4160·(4· 768+4)=12,796,160. (19) So memory preprocessing contributes 15,991,04015,991,040 FLOPs. Memory attention. Here hmem=12h_mem=12, hmem,kv=12h_mem,kv=12, so memory KV dimension is 768768. Q Q :2⋅1024⋅768⋅768=1,207,959,552 :2· 1024· 768· 768=1,207,959,552 (20) K K :2⋅4160⋅768⋅768=4,907,335,680 :2· 4160· 768· 768=4,907,335,680 (21) V V :2⋅4160⋅768⋅768=4,907,335,680 :2· 4160· 768· 768=4,907,335,680 (22) O O :2⋅1024⋅768⋅768=1,207,959,552 :2· 1024· 768· 768=1,207,959,552 (23) QK⊤,AV QK ,\ AV :4⋅1⋅1024⋅4160⋅768=13,086,351,360. :4· 1· 1024· 4160· 768=13,086,351,360. (24) The attention softmax, masking, and scaling term is 12⋅1024⋅4160⋅(5+2)=357,826,560.12· 1024· 4160·(5+2)=357,826,560. Hence memory attention contributes 25,674,645,50425,674,645,504 FLOPs. Extra norm and residual. extra RMSNorm=1024⋅(4⋅768+4)=3,149,824,extra residual=1024⋅768=786,432.extra RMSNorm=1024·(4· 768+4)=3,149,824, residual=1024· 768=786,432. Therefore the extra cost added by an MoC memory layer is 7,124,491+331,859+15,991,040+25,674,645,504+3,149,824+786,432=25,702,029,150.7,124,491+331,859+15,991,040+25,674,645,504+3,149,824+786,432=25,702,029,150. Adding the underlying standard transformer layer gives 17,424,982,016+25,702,029,150=43,127,011,16617,424,982,016+25,702,029,150=43,127,011,166 FLOPs per MoC memory layer. Final prediction head and loss. After the transformer stack, the estimator includes a final RMSNorm, LM head, and cross-entropy loss: final RMSNorm :1024⋅(4⋅768+4)=3,149,824 :1024·(4· 768+4)=3,149,824 (25) LM head :2⋅1024⋅768⋅49152=77,309,411,328 :2· 1024· 768· 49152=77,309,411,328 (26) CE loss :(1024−1)⋅49152⋅5=251,412,480. :(1024-1)· 49152· 5=251,412,480. (27) Thus the final prediction/loss block contributes 77,563,973,63277,563,973,632 FLOPs. Model totals. The backbone-only vanilla model has 16 standard layers: Ffwdbackbone=16⋅17,424,982,016+77,563,973,632=356,363,685,888.F_fwd^backbone=16· 17,424,982,016+77,563,973,632=356,363,685,888. The iso-FLOP vanilla baseline has 24 standard layers: Ffwdiso=24⋅17,424,982,016+77,563,973,632=495,763,542,016.F_fwd^iso=24· 17,424,982,016+77,563,973,632=495,763,542,016. The memory model has 12 standard layers and 4 MoC memory layers: FfwdMoC=12⋅17,424,982,016+4⋅43,127,011,166+77,563,973,632=459,171,802,488.F_fwd^MoC=12· 17,424,982,016+4· 43,127,011,166+77,563,973,632=459,171,802,488. Following the estimator, we approximate one backward pass as 2×2× the forward FLOPs: Fbwd≈2Ffwd.F_bwd≈ 2F_fwd. Hence the total for one forward plus one backward pass is 3Ffwd3F_fwd. Table 4: Per-sequence FLOPs for one forward pass and one forward+backward pass (B=1,L=1024B=1,L=1024). 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 Note that these are single-sequence estimates. They differ slightly from dividing the microstep estimates by batch size, because the sequence-level router auxiliary-loss terms do not scale exactly linearly with batch size. A.3 Algorithm: Mixture of Chapters (MoC) memory layer Algorithm 1 MoC memory layer with sequence-level routing 1:Batch hidden states H∈ℝB×L×dH ^B× L× d 2:Chaptered memory bank M=[M1;…;MC]M=[M_1;…;M_C], where Mc∈ℝT×dM_c ^T× d 3:Router parameters Wr,brW_r,b_r, projections WQ,WK,WVW_Q,W_K,W_V, top-k 4:Updated hidden states H′∈ℝB×L×dH ^B× L× d 5:for b=1b=1 to B do 6: rb←Pool(Hb)r_b (H_b) 7: pb←softmax(Wrrb+br)p_b (W_rr_b+b_r) 8: b←TopK(pb,k)S_b (p_b,k) 9: Mb←Concat(Mc:c∈b)M_b (\M_c:c _b\) 10: Qb←Norm(Hb)WQQ_b (H_b)W_Q 11: Kb←MbWKK_b← M_bW_K 12: Vb←MbWVV_b← M_bW_V 13: Ab←softmax(QbKb⊤dh)A_b \! ( Q_bK_b d_h ) 14: MemReadb←AbVbMemRead_b← A_bV_b 15: Hb′←Hb+MemReadbH _b← H_b+MemRead_b 16:end for 17:return H′=[H1′;…;HB′]H =[H _1;…;H _B] A.4 Additional experimental details A.4.1 Model sizes The memory model has a larger total parameter count due to the persistent bank, but uses sparse chapter routing so only a subset of the bank is activated per input. We therefore match pretraining compute to a dense baseline (Vanilla iso-FLOP). For the non-memory baselines, the Vanilla (backbone-only) model uses the same 16-layer backbone architecture with the memory layers removed, while the Vanilla (iso-FLOP) baseline is a denser variant obtained by repeating the standard backbone layers without memory modules and increasing depth to match the pretraining compute budget of the memory model. For reference, the dense baseline has 202.94M parameters, while the memory model has 371.29M total parameters composed of a 147.87M backbone plus a 201.38M bank and 22.04M memory-layer parameters. A.4.2 Model and training configuration Tables 5 and 6 summarize the main Mixture of Chapters (MoC) architecture and the training settings used for pretraining and instruction fine-tuning (IFT). Table 5: Main MoC architecture used in our experiments. 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 θ=100000θ=100000, tied embeddings, vocabulary size 49152 Memory placement Memory cross-attention inserted at layers 2,6,10,14\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=64k=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 Regularization Load-balance loss coefficient 0.01, z-loss coefficient 0.001 Other Memory adapter enabled, memory initialization std. 0.02, memory quantization disabled Table 6: Main training settings for pretraining and IFT. 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 Training duration 9600 steps 2 epochs (3200 steps) Optimizer AdamW AdamW Learning rates base model: 3×10−43× 10^-4 memory layers: 6×10−46× 10^-4 memory bank: 6×10−46× 10^-4 base model: 3×10−53× 10^-5 memory layers: 1.5×10−51.5× 10^-5 memory bank: 5×10−65× 10^-6 Scheduler WSD, warmup 250 steps, min LR ratio 0.1, decay start at step 8160 Cosine, warmup 250 steps Optimization details weight decay 0.1, Adam β1=0.9 _1=0.9, β2=0.95 _2=0.95, grad clip 1.0 weight decay 0.1, Adam β1=0.9 _1=0.9, β2=0.95 _2=0.95, grad clip 1.0 Precision bf16 bf16 Checkpointing gradient checkpointing enabled gradient checkpointing enabled Initialization random initialization initialized from the pretrained MoC checkpoint For comparison, the vanilla backbone-only baseline uses the same 16-layer backbone with memory disabled, while the iso-FLOP vanilla baseline increases depth to 24 layers. A.4.3 Pretraining benchmarks including the backbone-only baseline Table 7 includes the backbone-only benchmark results to contextualize gains from MoC versus simply changing backbone capacity. Table 7: Benchmark accuracy after pretraining (including backbone-only). 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 A.4.4 IFT learning rates and bank update ablations IFT uses a base learning rate of 3×10−53× 10^-5 (one tenth of the pretraining rate). We vary the memory bank learning rate across: (1) Frozen bank (LR = 0), (2) Low bank LR (very small LR), (3) Equal LR (bank LR matches base LR). Across these settings, Figure 2 shows no visible divergence in validation loss curves. Table 8 reports benchmark scores for representative settings. Freezing the bank performs on par with updating it, indicating that the pretrained bank is retained and sufficient for post-training. Table 8: MoC IFT benchmarks for bank update settings. 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 ∗ denotes the main setting reported above. A.4.5 IFT shifts the pretraining distribution loss while preserving knowledge benchmarks When evaluating IFT checkpoints on the pretraining validation distribution, the memory model exhibits higher loss (3.2024) than the vanilla baseline (3.0574), even though it retains benchmark performance substantially better on ARC-Challenge and BoolQ. This suggests that the memory model can adapt more strongly toward the IFT objective while keeping factual knowledge stable through the explicit bank, consistent with a division of labor between backbone adaptation and memory anchoring. A.4.6 Compute and memory implications during post-training Freezing the memory bank during IFT removes bank updates and reduces optimizer state and gradient memory for that component, which lowers VRAM requirements. In addition, increasing context length from 1024 to 2048 increases self-attention cost for all models. Under this longer-context setting, the benefit of freezing the bank is preserved while the dense baseline pays the full quadratic self-attention cost. This combination makes the memory approach attractive for heavier or longer post-training schedules. A.5 Ablation studies and future work This section lists ablations for Mixture of Chapters and outlines directions to clarify scaling, efficiency, and retention. Ablation studies. • Top-k routing sweep. Vary the number of routed chapters, e.g., k∈1,2,4,8,16,32,64k∈\1,2,4,8,16,32,64\, holding bank size fixed to map the quality–compute frontier. • Routing granularity and train/inference mismatch. Compare (i) sequence routing at train and inference (current), (i) token routing at inference only, and (i) token routing at train and inference (if feasible). • Shared chapters. Sweep the number of always-on shared chapters and mixing schemes: nshared∈0,1,2,4,8n_shared∈\0,1,2,4,8\. • Memory placement and sharing. Study where memory layers are inserted and whether the bank is shared across layers vs per-layer banks. • Router regularization. Sweep load-balancing / auxiliary routing losses (and any z-loss) to quantify utilization collapse vs quality tradeoffs. • Initialization. Ablate output-projection initialization (e.g., zero-init) and bank initialization scale to test stability and early training behavior. • Memory size and chapter geometry. Scale total memory tokens, number of chapters, and tokens-per-chapter T to separate “more chapters” from “longer chapters.” • Compression ablations. Evaluate quantized banks and low-rank factorization of bank/projections to trade expressiveness for capacity. • PEFT comparisons: LoRA vs memory adapters vs both. Compare LoRA-only, memory-adapter-only, and combined settings under matched trainable parameter budgets. Future work. • Scaling laws for learned memory. Extend pretraining beyond 9.6B tokens and sweep bank capacity to quantify how performance scales with memory under fixed compute. • Token-level routing with efficient kernels. Investigate kernels that avoid materializing large routed key tensors, enabling token routing during training without large memory overhead. • Dynamic and test-time updates. Add controlled write/update mechanisms to support continual acquisition and personalization while retaining the “freeze bank” benefits. • Interpretability and chapter-level editing. Analyze router selections for specialization, and test targeted chapter edits (zeroing, swapping, pruning) to localize knowledge. • Hybrid memory: shared global + routed private. Explore explicit separation between a small shared global memory and a large routed bank, with different update schedules or learning rates. • Broader evaluations. Add tests for long-context generalization, factual recall, and retention under multiple sequential post-training stages. • Interpretability and analysis of memory usage. Analyze how memory is used across layers and time: (i) chapter selection statistics (entropy, sparsity, stability across prompts), (i) per-layer reliance on memory (attention mass to memory vs self-attn), (i) chapter specialization (clustering chapters by the queries they serve or by input domains), and (iv) causal interventions such as zeroing/swapping/pruning selected chapters or memory tokens to localize which knowledge is stored where. This can be paired with qualitative retrieval probes by decoding nearest-neighbor token projections from memory to inspect what each chapter appears to represent.