🤖 machine learning

Mixture of Chapters: Scaling Learnt Memory in Transformers

该论文提出了一种名为“章节混合（Mixture of Chapters）”的架构，通过引入可学习的稀疏记忆库和基于路由的章节划分机制，使 Transformer 能够在保持计算可行的同时扩展至 26 万记忆令牌，从而在等算力设置下显著提升了模型性能、知识保留能力及抗遗忘性。

原作者： Tasmay Pankaj Tibrewal, Pritish Saha, Ankit Meda, Kunal Singh, Pradeep Moturi

发布于 2026-03-24

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Tasmay Pankaj Tibrewal, Pritish Saha, Ankit Meda, Kunal Singh, Pradeep Moturi

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

这篇论文提出了一种让 AI 变得更聪明、更不容易“忘事”的新方法。为了让你轻松理解，我们可以把这篇论文的核心思想想象成给一个只有“短期记忆”的超级大脑，装上了一个带“智能索引”的巨型图书馆。

以下是用通俗语言和生动比喻对这篇论文的解读：

1. 核心问题：大脑太“健忘”，且没有“记事本”

目前的 AI（Transformer 模型）就像是一个过目不忘但记不住重点的超级学霸。

现状：它在学习时，把所有知识都压缩进了自己的“大脑参数”里（就像把书的内容强行塞进脑子里）。
缺点：
1. 很难修改：如果你想让它忘掉一个错误知识或学习一个新知识，往往需要重新训练整个大脑，这就像为了改一个错别字要把整本书重写一遍。
2. 容易“失忆”：当它学习新任务（比如从写文章变成回答医疗问题）时，它很容易把旧知识（比如写文章的技巧）给忘了，这叫“灾难性遗忘”。
3. 没有外部存储：它没有像人类那样的“笔记本”或“图书馆”来专门存放事实性知识。

2. 解决方案：给 AI 装一个“智能图书馆” (Learned Memory Bank)

作者们给 AI 加了一个可学习的记忆库。

比喻：想象 AI 的大脑旁边有一个巨大的图书馆，里面存放着成千上万张“知识卡片”（Latent Tokens）。
运作方式：当 AI 遇到一个问题时，它不会只靠脑子里的模糊印象，而是会去图书馆里检索相关的卡片，把卡片上的信息读出来辅助思考。
优势：这些卡片是专门训练出来的，可以独立更新。如果知识过时了，只需换掉卡片，不用重写整个大脑。

3. 核心挑战与突破：如何管理这个巨大的图书馆？

挑战：如果图书馆有 26 万张卡片，每次 AI 思考都要把 26 万张卡片全翻一遍，那速度会慢到无法接受（计算成本太高）。

创新方案：Mixture of Chapters (MoC) —— “章节式路由”
作者想出了一个绝妙的办法，把图书馆分成了4000 多个“章节”（Chapters），并雇佣了一个超级图书管理员（Router）。

比喻：
- 图书馆：被分成了很多个独立的“章节区”（比如“历史区”、“科学区”、“文学区”）。
- 图书管理员：当 AI 输入一个问题（比如“苹果是谁发明的？”），图书管理员不会去翻整个图书馆。它只需看一眼问题，就能迅速判断：“哦，这个问题属于‘科学区’和‘历史区’"。
- 精准检索：管理员只打开这两个特定的“章节区”，从中挑选出最相关的几十张卡片给 AI。
效果：
- 快：AI 只需要处理一小部分卡片，速度极快。
- 大：虽然每次只查一小部分，但整个图书馆可以无限扩大（论文中做到了 26 万张卡片），容量巨大。
- 省：就像你不需要为了找一本书而跑遍整个城市，只需要去特定的几个街区。

4. 实验结果：它真的有用吗？

作者把这种新 AI 和传统的 AI 进行了对比，发现：

学得更扎实：在同样的计算资源下，带“图书馆”的 AI 考试成绩（验证损失）更好，说明它更聪明。
不再“失忆”：这是最精彩的发现。
- 传统 AI：当它开始学习新任务（指令微调）时，就像大脑被“格式化”了一样，以前学的知识（比如常识问答）大量丢失，成绩暴跌。
- 新 AI：因为它的知识都稳稳地存在“图书馆”的卡片里，大脑只负责处理逻辑和语言结构。所以，即使在学习新任务时，它依然能完美保留旧知识，几乎没有“失忆”。
冻结也能用：有趣的是，作者发现甚至可以在学习新任务时完全冻结图书馆（不让卡片更新），AI 依然表现很好。这意味着，只要把知识存好，以后随时可以拿出来用，不需要反复折腾。

5. 总结：这意味什么？

这篇论文提出了一种新的扩展 AI 能力的方法。

以前，我们想让 AI 变强，只能拼命增加它的“大脑皮层”（参数），这既费钱又容易让它“忘事”。
现在，我们可以在大脑旁边加一个外挂的、可管理的“知识硬盘”。

对未来的意义：
- 更聪明的 AI：可以存储海量事实知识。
- 更稳定的 AI：学习新技能时不会丢掉旧技能。
- 更灵活的 AI：我们可以像整理书架一样，直接编辑、删除或更新 AI 的特定知识，而不需要重新训练整个模型。

一句话总结：
这就好比给一个只有短期记忆的超级大脑，配上了一个带智能索引的无限容量图书馆。它让 AI 在保持高速思考的同时，拥有了过目不忘的长期记忆，并且在学习新东西时，再也不会把旧知识给弄丢了。

这篇论文提出了一种名为**“章节混合”（Mixture of Chapters, MoC）**的新架构，旨在解决 Transformer 模型缺乏显式、可寻址的持久化记忆机制的问题。该研究通过引入可学习的稀疏记忆库，并结合类似混合专家（MoE）的路由机制，实现了在保持计算可行的前提下，将记忆容量扩展至数十万 token 级别。

以下是对该论文的详细技术总结：

1. 研究背景与问题 (Problem)

隐式记忆的局限性：传统的 Transformer 将知识隐式地存储在密集的参数中。这种机制使得记忆难以被独立检查、编辑或扩展。
现有方案的不足：目前的长上下文处理主要依赖推理时的机制（如 KV Cache 压缩、检索增强生成 RAG 等），这些方法主要复用或检索外部信息，而非在训练过程中学习并存储一个内部的、结构化的事实知识库。
核心挑战：如何在 Transformer 中引入显式的记忆库，同时解决大规模记忆库带来的计算成本过高（密集注意力机制复杂度为 $O(L \times N_m)$ ）以及持续训练中的灾难性遗忘问题。

2. 方法论 (Methodology)

2.1 可学习的记忆层 (Learned Memory Layer)

记忆库构建：引入一组随机初始化并端到端训练的潜在（latent）记忆 token，构成记忆库 $M$ 。
交叉注意力机制：Transformer 层通过交叉注意力（Cross-Attention）查询记忆库。输入序列生成查询（Query），记忆库提供键（Key）和值（Value）。
- 公式： $MemRead(H, M) = \text{softmax}(\frac{QK^\top}{\sqrt{d_h}})V$
- 输出通过残差连接加回原始 token 流。

2.2 章节混合路由 (Mixture of Chapters, MoC)

为了解决大规模记忆库的扩展性问题，作者提出了 MoC 架构：

分块（Partitioning）：将巨大的记忆库 $M$ 划分为 $C$ 个“章节”（Chapters），每个章节包含 $T$ 个 token。
稀疏路由（Sparse Routing）：
- 引入一个轻量级的路由器（Router），基于输入序列的隐藏状态（通过池化得到）计算每个章节的得分。
- 对于每个输入，路由器仅选择得分最高的 $k$ 个章节（Top-k）。
- 模型仅对选中的章节进行注意力计算。
计算复杂度优化：将注意力复杂度从 $O(L \times N_m)$ 降低到 $O(L \times k \times T)$ ，使得在保持计算量可控（Tractable）的同时，能够支持高达 262,208 个记忆 token 的容量。

2.3 架构细节

模型配置：基于 Decoder-only Transformer，在特定层（如第 2, 6, 10, 14 层）插入记忆层。
参数规模：记忆库包含约 26 万个 token，分为 4097 个章节，每次推理选择 64 个章节。

3. 主要贡献 (Key Contributions)

Transformer 中的可学习关联记忆库：提出了一种基于潜在 token 的持久化记忆机制，通过交叉注意力在紧凑的潜在空间内进行内部检索，而非依赖外部文本检索。
可扩展的稀疏访问机制 (MoC)：受 MoE 启发，通过章节路由实现了大规模记忆库的稀疏访问，平衡了容量与计算成本。
固定计算下的扩展行为：在等 FLOPs（iso-FLOP）设置下，证明了引入记忆库的模型优于标准密集 Transformer 基线，开辟了新的扩展维度。
持续训练下的知识保留：证明了显式记忆库能有效缓解从预训练到指令微调（IFT）过程中的灾难性遗忘，起到“事实锚点”的作用。

4. 实验结果 (Results)

4.1 预训练表现

验证损失：在 96 亿 token 的预训练任务中，MoC 模型的验证损失（2.79）优于等计算量的密集基线（2.86）和仅骨干网络基线（2.92）。
基准测试：在 MMLU、ARC-Challenge、BoolQ 和 OBQA 等知识密集型任务上，MoC 模型在预训练阶段即展现出更高的准确率。

4.2 指令微调与遗忘分析 (IFT & Forgetting)

实验设置：在 2.3 亿 token 上进行 2 轮指令微调，并将上下文长度从 1024 增加到 2048。
遗忘现象：
- 密集基线 (Vanilla)：在知识密集型任务上出现严重遗忘。例如，ARC-Challenge 准确率从 31.77% 降至 25.08%（接近随机猜测），BoolQ 从 57.13% 降至 50.89%。
- MoC 模型：表现出极强的稳定性。BoolQ 准确率几乎不变（+0.24%），ARC-Challenge 仅下降 2.68%（远低于基线的 6.69% 下降）。
冻结记忆库实验：在指令微调期间，如果冻结记忆库（不更新记忆 token），模型性能与更新记忆库相当。这表明预训练好的记忆库可以作为稳定的事实锚点，无需在微调阶段重新学习。

4.3 分布偏移

在微调后的模型上评估预训练分布，MoC 模型的损失略高于基线，但这表明模型能够更灵活地适应指令对齐任务，同时通过记忆库保留事实知识，实现了“骨干适应”与“记忆锚定”的分工。

5. 意义与结论 (Significance & Conclusion)

新的扩展维度：该研究证明了除了增加模型参数或计算量之外，显式记忆容量是语言模型扩展的另一个重要维度。
解决遗忘问题：MoC 为持续学习（Continual Learning）提供了一种架构级的解决方案，通过显式存储事实知识来减少微调过程中的干扰。
效率与容量的平衡：通过章节路由机制，成功将记忆容量扩展至 26 万 token 级别，同时保持了计算的可控性，为未来构建更大规模、更具可解释性的记忆系统奠定了基础。
未来方向：论文建议进一步研究更大规模的记忆库、动态更新机制、路由的细粒度优化（如 token 级路由）以及记忆内容的可解释性分析。

总结：这篇论文通过引入“章节混合”机制，成功将可学习的显式记忆集成到 Transformer 中，不仅提升了模型在知识密集型任务上的表现，更重要的是显著增强了模型在持续训练过程中的抗遗忘能力，为构建更智能、更持久的语言模型提供了新的架构范式。