11. Memorizing Transformers:把 KNN 检索塞进 Transformer 的记忆层

论文:Memorizing Transformers
本地原文:

![[papers/11-Memorizing-Transformers-2203.08913.pdf]]

论文定位与阅读路径

阅读边界

这篇笔记不是逐字搬运论文全文,而是“原文内嵌 + 逐节精读”。你在 Obsidian 里打开本文,上面就是 PDF 原文,下面是中文拆解;遇到不懂的段落,可以直接上下对照。

核心定位

这篇论文把语言模型的记忆从“只能靠参数和当前上下文”扩展成“模型运行时还能查一个外部 KNN 记忆库”。它更偏算法层,不是应用层的 Mem0/Zep。

如果用一句大二能懂的话说:

Transformer 原本只能看当前上下文;Memorizing Transformer 让它还能按相似度翻以前见过的 token 隐状态。

原文结构导读

证据链阅读

研究背景与问题定义

背景:普通 Transformer 的记忆在哪里

普通 LLM 的记忆主要有两种:

Memorizing Transformers 加了第三种:

这里的关键是:它记的不是人类可读文本,而是模型中间层的向量表示。

方法框架与核心机制

核心结构

这个图要这样理解:

1. 当前 token 进入模型。
2. 在某一层,模型产生 query。
3. query 去外部记忆库里找相似 key。
4. 找到的 value 参与 attention。
5. 当前序列自己的 attention 和外部记忆 attention 融合。
6. 当前输入也会变成新的 key-value,继续写回记忆库。

key、value、query 到底是什么

先别被公式吓住。你可以把它理解成搜索系统:

普通 RAG 是:

问题文本 -> embedding -> 查文档 chunk

Memorizing Transformer 更像:

模型中间层 query -> 查历史 hidden states -> 融入 attention

所以它比应用层 RAG 更贴近模型内部。

Method 逐层拆解

1. Memorizing layer

论文不是在每一层都查外部记忆,而是在 Transformer 里选某些层作为 memorizing layer。

这样做有两个原因:

1. 每层都检索成本太高。
2. 中间层的表示往往比原始 token 更适合相似性匹配。

对你来说,重点不是背具体层数,而是理解这个设计原则:

记忆检索发生在模型内部表示空间,而不是只发生在输入文本空间。

2. 外部 memory 是非可微的

论文中的 memory store 不是训练参数的一部分。它更像一个外部数据库。

这带来两个结果:

这就是它和 MemoryLLM、Titans 这类“可学习记忆模块”的区别。

3. 近似 KNN

真正的记忆库会很大,暴力比较所有 key 太慢,所以需要 approximate nearest neighbor。

学习者视角:

暴力搜索：一本书一页页翻。
近似 KNN：先按章节/索引缩小范围，再找最像的几页。

工程上类似 FAISS、ScaNN、HNSW 这类向量检索结构。

4. 记忆与注意力融合

普通 attention 只看当前上下文内的 key-value。论文加了外部 memory key-value。

可以理解成:

当前上下文 attention：看眼前这段话。
记忆 attention：翻以前见过的相似片段。
最终输出：两边信息合并后预测下一个词。

这也是为什么它能处理比上下文窗口更远的信息。

相关工作与方法对比

这篇和 RAG 的关系

所以你现在做 Java 后端 Agent,不需要真的实现 Memorizing Transformer。但你要吸收它的思想:

记忆不一定只是文本数据库,也可以是模型中间状态的可检索缓存。

实验论证与评价方式

实验阅读方法

论文在多类数据上评估,包括网页文本、论文、书籍、代码、形式化定理。它想证明的不是“某个任务赢了”,而是:

1. 外部记忆对不同类型语言建模都有帮助。
2. memory size 变大时,效果可以继续提升。
3. 检索到的历史片段能补足有限上下文。

你读实验时不要死抠每个数字,先抓三个问题:

1. baseline 是什么。
2. 加 memory 后 perplexity 有没有下降。
3. 记忆库变大时趋势是否稳定。

perplexity 可以粗暴理解为:

模型对下一个词越不困惑，perplexity 越低。

概念校准

你考试做编程题时,脑子里能记住当前题目,这相当于上下文。你以前写过很多类似代码,但不可能全部背下来。于是你有一个代码片段库,看到当前题目后检索最相似的旧代码,再辅助当前推理。

Memorizing Transformer 做的就是这个,只不过它检索的不是代码文本,而是模型内部向量。

工程映射与实践启发

工程落地启发

你暂时不用改模型,但可以借鉴这个分层:

一个简化实现:

flowchart TD
    A[用户问题] --> B[生成 embedding]
    B --> C[向量库 KNN 检索]
    C --> D[取回相似记忆]
    D --> E[重排 rerank]
    E --> F[拼入 prompt]
    F --> G[LLM 回答]

这不是论文原方法,但工程精神是一致的:把“过去经验”变成可检索的外部记忆。

局限、对比与阅读检查

局限与边界

1. 需要改模型结构,普通 API 用户做不了。
2. 记忆是向量,不如文本记忆可解释。
3. KNN 检索会带来额外延迟和存储成本。
4. 记住旧数据可能带来隐私与版权问题。
5. 它更擅长“找相似历史”,不等于真正理解长期因果关系。

阅读检查清单

读这篇后,你至少要能说清:

1. 什么是外部 KNN 记忆。
2. 为什么它不是普通 RAG。
3. key/value/query 在记忆检索里分别是什么。
4. 为什么 memory size 能影响模型表现。
5. 为什么它属于模型驱动派,而不是应用驱动派。

和后续论文的关系

Memorizing Transformer 是“模型内部记忆”路线的早期代表。后面的 MemoryLLM、Memory3、WISE、Titans 都在继续问同一个问题:

新知识到底应该放在外部数据库、模型参数、激活缓存,还是一个专门设计的记忆模块里?

深度精读:KNN memory 是怎么接到 Transformer 里的

1. 它和 kNN-LM 的关系

这篇不是凭空冒出来的。它继承了一个直觉:

语言模型在预测下一个 token 时，如果能找到过去很相似的上下文，就能参考过去那个位置后面出现了什么。

kNN-LM 更像在输出层外接一个检索器;Memorizing Transformer 更进一步,把检索接进 Transformer 的中间层。

差别:

所以它更“内生”,但也更难工程复现。

2. key/value 不是数据库字段,而是模型内部状态

在普通后端里 key/value 可能是:

key = user_id
value = user_profile

在这篇里:

key = 过去某个 token 在某层的隐藏向量
value = 与这个位置相关的 value 向量

模型当前处理 token 时生成 query,去找相似 key。找到后,把对应 value 当成额外注意力材料。

这就是为什么它不容易解释:取回的不是文本,而是向量。

3. 为什么要 approximate KNN

如果记忆库有几千万个 hidden states,暴力搜索就是:

当前 query 和每个 key 算一次距离

这不可接受。所以要近似搜索。

工程上你可以类比:

这和向量数据库里的 HNSW、IVF、ScaNN 是同一类问题。

4. 为什么 memory size 越大可能越好

模型参数固定时,参数里能记的知识有限。外部记忆库变大,相当于能查的历史经验更多。

但它不是无限好:

1. 检索越大,延迟越高。
2. 相似但错误的记忆会干扰。
3. 旧数据可能过时。
4. 隐私和数据隔离更复杂。

所以真实系统一定要有 memory eviction / filtering / scope。

5. 它对应用层 RAG 的启发

虽然你不会改 Transformer,但可以学到三点:

例如编程 Agent:

当前 bug 栈信息 -> 检索相似历史 bug -> 取回修复步骤 -> 注入上下文

这就是“外部经验增强推理”。

6. 为什么它不是万能长上下文

KNN memory 更擅长“局部相似”。但有些任务需要:

1. 全局结构理解。
2. 长距离因果推理。
3. 多个证据合成。
4. 时间变化判断。

这些不能只靠相似 hidden state 解决。

所以它和 Zep/HippoRAG 是互补关系:

KNN memory: 找相似模式
Graph memory: 找实体关系
Temporal memory: 找时间有效事实
Reflection memory: 找经验教训

7. 读公式的最低要求

你不需要推完整公式,但要知道:

query q：当前 token 表示
memory keys K：过去 token 表示
memory values V：过去可被 attention 读取的信息
KNN(q, K)：找最相似的 k 个历史位置
attention(q, K_mem, V_mem)：把取回的历史信息融合进当前表示

把这几行翻成人话,你就不会被公式吓住。

8. 如果写成工程系统

record MemoryVector(
    String id,             // 向量记忆 ID。
    float[] keyEmbedding,  // 用于 KNN 检索的 key 表示。
    String payloadRef,     // 指向原始 token 片段、文档块或外部内容。
    Instant createdAt,     // 写入时间，可用于过期或分层索引。
    String scope           // 作用域，避免不同用户/任务的记忆混检。
) {}

// retrieve = 给当前 query 表示找最近的 k 条历史记忆。
List<MemoryVector> retrieve(float[] query, String scope, int k) {
    // vectorIndex.search 应同时使用向量相似度和 scope 过滤。
    return vectorIndex.search(query, scope, k);
}

注意:应用层的 payloadRef 可以指向文本、代码、工具日志;论文里的 payload 是模型内部 value。思想相同,层级不同。

论文精读主线:问题-方法-验证-局限-启发

1. 核心结论

Memorizing Transformers 把历史 hidden states 存入外部 KNN memory,让 Transformer 在生成时能检索更远历史的内部表示。

本篇的阅读重心是:在 Transformer 中间层检索历史 hidden states,扩展模型内部可访问记忆。读完后要能说清它改变的是 runtime、存储、检索、更新还是评测层;否则只是记住了论文名。

2. 问题背景

普通 Transformer 的上下文有限,长距离信息离开窗口后就无法直接注意到。外部文本 RAG 检索的是人类可读文本,但模型内部其实可以记住隐藏状态表示;问题是如何把这种内部记忆接回注意力机制。

把问题具体化为失败样例:模型需要回忆远早于当前窗口的局部模式,文本级 RAG 太粗,输出层 kNN 又介入太晚。用这个样例反推论文动机,比泛泛说“上下文不够”或“记忆不准”更有判断力。

3. 核心方法

模型在特定层把 key/value hidden states 写入外部 memory。生成时当前 query 去 KNN memory 找相似 key,取回 value,并和普通 attention 结果融合。

flowchart TD
    H[历史 token hidden states] --> K[写入 KNN key/value memory]
    Q[当前 query 表示] --> N[KNN 检索近邻]
    K --> N
    N --> V[取回历史 value]
    V --> F[和 attention 融合]
    F --> O[预测下一个 token]

方法部分按数据流读:训练或推理时保存 hidden state key/value,后续在中间层按相似 hidden state 检索并融合。每一步都要能标出输入、输出、控制点和失败后该记录的日志。

4. 技术细节

这里的 key/value 不是数据库字段,而是模型内部向量。memory 通常非可微,依赖近似最近邻检索。它更像模型层面的长程缓存,不是应用层知识库。

final class KnnTransformerMemory {
    private final HiddenStateWriter stateWriter; // 从指定层抽取 key/value hidden states 写入外部索引
    private final ApproxKnnIndex knnIndex; // 近似最近邻检索，控制大规模 memory 的延迟
    private final MemoryAttentionMixer mixer; // 把 KNN 取回的 value 和普通 attention 输出融合
    private final MemoryEvictionPolicy evictionPolicy; // 控制 memory 容量、过期和采样策略

    MemoryResult handle(MemoryInput input) {
        // 1. 从指定层抽取 key/value hidden states 写入外部索引
        // 2. 近似最近邻检索，控制大规模 memory 的延迟
        // 3. 把 KNN 取回的 value 和普通 attention 输出融合
        // 4. 控制 memory 容量、过期和采样策略
        // 返回前附带 evidenceIds、scoreBreakdown、traceId，避免只看到答案看不到记忆来源。
        return MemoryResult.withTrace(input.traceId());
    }
}

字段检查重点是 layerId、hiddenKey、valueState、knnIndex、retrievalK、mixingWeight 是理解重点。这些字段不是装饰,它们决定后续能否检索、审计、回滚、限权和做离线评测。

5. 实验验证

实验看语言建模困惑度、长上下文任务和 memory size 消融。要确认收益来自 KNN memory,而不是更多参数或更长窗口。随着 memory 变大性能是否稳定提升,是这类方法的重要证据。

实验阅读重点:看长序列语言建模、稀有模式回忆和不同层接入消融;不要把它和普通文档 RAG 混为一谈。如果论文没有直接给出这层证据,复现时就要自己补机制级指标。

6. 局限与启发

它难解释、难用闭源 API 复现,也不适合作为用户可控记忆。启发是应用层 RAG 可以借鉴 key/value 分离思想:检索 key 用于匹配,value 保存更完整证据。

迁移到你的项目时,结论是:这类方法暂不适合作为 Java 后端首版,但能启发你区分“文本证据记忆”和“模型表示记忆”。先把边界做对,再追求更复杂的模型或更大的记忆规模。

7. 关键术语表

8. 复现难点

实现需要改模型结构,训练和推理耦合强。memory 太大延迟高,太小收益低。内部状态无法人工审计,写入噪声也难发现。

复现入口应从失败集开始:最小复现可用 toy Transformer + FAISS 模拟 hidden state 检索,观察 k、层位置和缓存大小的影响。只有先证明旧方案会坏,才知道新机制是否真的解决了问题。

9. 代码框架建议

如果只做工程类比,可实现 RetrievalKey(textEmbedding) 与 EvidenceValue(rawChunk) 分离:key 负责匹配,value 保留原文、来源、时间和权限。

代码验收门槛:任何内部记忆方案都要回答污染、回滚、权限和解释问题;回答不了就只能做研究实验。这条门槛应写进测试或 trace 字段里,避免实现只停留在 demo。

10. 相关论文对比

它和 RAG 都是检索增强,但检索对象不同:RAG 检索文本,Memorizing Transformer 检索 hidden states。Titans 进一步引入可学习神经长期记忆,MemLong 更偏文本历史缓存。

横向比较时抓这个角度:它比 MemLong 更贴近模型内部,比 MemoryLLM 更偏检索缓存,比 WISE 少了参数编辑语义。比较对象不是谁“更先进”,而是谁负责哪一层失败。

11. 批判性思考

这类方法证明模型内部记忆有潜力,但产品可控性弱。对个人知识 Agent,更适合作为思想来源,而不是第一阶段技术路线。

批判边界:hidden state 记忆难以解释给用户看;一旦记错,很难像外部数据库那样定位和删除。这不是否定论文,而是把论文结论放回真实系统的成本、权限和恢复约束里。