给Transformer加可学习的外部记忆库，提升知识存储与保持能力。
Mixture of Chapters: Scaling Learnt Memory in Transformers
- 用可训练的稀疏记忆令牌池，通过交叉注意力检索知识。
- 支持26.2万记忆项，计算开销可控，优于同等算力的普通Transformer。
- 适合需要长期记忆和持续学习的任务，如知识密集型对话系统。
Transformer缺乏显式的知识存储与组织机制。本文引入可学习的稀疏记忆池：一组随机初始化并端到端训练的潜在标记，通过跨注意力机制被变压器层查询以获取存储的知识。为在不增加显著注意力开销的前提下扩展记忆容量，受Mixture-of-Experts启发，提出基于章节的路由策略，将记忆池划分为多个章节，并训练路由器根据输入选择相关子集。该方法使记忆规模扩展至262,000个记忆令牌，同时保持计算效率。我们在预训练与指令微调任务中，在等FLOP设置下对比标准Transformer模型，结果表明所提模型超越基线，说明显式关联记忆能为模型参数隐式容量提供补充。此外，模型在持续训练中表现出更强的知识保留能力，跨训练阶段（如预训练到指令微调）时对遗忘具有鲁棒性。
原文摘要 · Abstract (English)
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).
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。