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

翻译：摘要：Transformer架构缺乏显式的机制来存储和组织训练过程中获取的知识。本文提出可学习稀疏记忆库（learnable sparse memory banks）：一组随机初始化并通过端到端训练获得的潜在标记（latent tokens），Transformer层通过交叉注意力机制对这些标记进行查询以检索存储的知识。为了在避免高昂注意力计算成本的前提下扩展记忆容量，我们借鉴混合专家架构（Mixture-of-Experts）的思想，提出基于章节的路由机制（chapter-based routing），将记忆库划分为多个章节（chapters）并训练路由器为每个输入选择相关子集。该方案支持将记忆标记扩展至262K规模，同时保持计算可行性。我们在预训练和指令微调阶段（采用等FLOP设置）将本方法与标准Transformer在相关基准测试上进行对比评估。结果表明，我们的模型超越了等FLOP基线，揭示了模型扩展的新维度——显式关联记忆与隐式参数捕获的知识具有互补性。此外，我们观察到在持续训练过程中知识保留能力显著提升，并在训练阶段过渡（如从预训练到指令微调）时展现出对灾难性遗忘的鲁棒性。