规模化记忆解码器：一种可预训练的参数化长期记忆 Memory Decoder at Scale: A Pretrained, Parametric Long-Term Memory
把长期记忆做成可独立预训练的参数化模块，410M+6.9B记忆超过12B模型
前置知识
kNN语言模型(kNN-LM)
由 Khandelwal 等人 2019 年提出的方法：先用冻结的基座模型把语料里每个上下文编码成一个向量作为 key、其下一词作为 value，构成一个数据存储(datastore)。推理时用当前上下文的编码向量在该 datastore 里检索最近的 K 个邻居，根据距离做加权投票得到一个下一词分布 $p_{ret}$，再和模型自己的分布做插值。本文的参数化记忆本质上是在用一个小 Transformer 去模仿这个 $p_{ret}$。
整篇论文的监督信号都来自 kNN 分布 $p_{ret}$，不理解它就无法理解记忆模块在学什么。
参数化记忆与非参数化记忆
非参数化记忆把知识存在外部 datastore 里，推理时实时检索（如 kNN-LM、RAG），存储和检索开销大。参数化记忆则把知识蒸馏进神经网络的权重，推理时无需外部检索。本文的 Memory Decoder 属于后者，把 kNN 的检索行为压缩进一个可独立训练、可插拔的 Transformer 模块。
区分两者是理解本文动机（解决检索开销）和创新点（知识进权重、可独立扩展）的基础。
分布对齐 / 知识蒸馏
训练记忆模块时不只学真实的下一词 $x_t$，还要让它输出的分布 $p_\psi$ 去逼近 kNN 的目标分布 $p_{ret}$，用 KL 散度衡量差距。这相当于用更丰富的软标签（一个上下文可能接多个合理词）来监督模型，而不是单一的硬标签。
记忆模块的核心损失函数就是 KL 对齐项加语言建模项，这是方法的灵魂。
Faiss 索引与近似最近邻搜索
Faiss 是 Facebook 开源的高维向量相似度搜索库。IVF（倒排文件）把向量空间划分成多个桶、查询时只搜部分桶；PQ（乘积量化）把向量压成短码来省存储、加速距离计算；HNSW（分层可导航小世界图）是一种高效的近邻图，常用来给 IVF 做 quantizer 做向量分配和路由。
本文最大的工程挑战就是如何在 207B token 规模上构建 kNN 分布，必须理解 Faiss 的这些组件才能看懂分布式流水线。
插值系数与参数效率
推理时最终预测 $p_{final}=(1-\alpha)p_\theta + \alpha p_\psi$，用 $\alpha\in[0,1]$ 控制记忆的贡献大小。参数效率指在达到同样性能时所用总参数量更少，例如 410M 基座加 6.9B 记忆共约 7.3B 参数，性能却超过 12B 模型。
这是衡量本文方法价值的直接指标，也是 motivation 的落脚点。
研究动机
标准 decoder-only 语言模型把长期记忆和推理能力塞在同一套参数里，这种纠缠带来三个具体痛点：第一，记忆容量无法独立扩展，要增加知识就得同时放大整个模型、抬高总参数量和训练成本；第二，领域适应只能靠继续预训练(CPT)或全量微调，要优化全部参数、代价高，还容易灾难性遗忘已有能力；第三，模型没有可独立插拔的记忆模块，无法为不同领域切换记忆、也无法在不同模型间复用记忆。已有工作也没真正解决这个纠缠：处理短上下文的 StreamingLLM、InfLLM、LongRoPE 只管输入上下文内的信息；kNN-LM、RAG 这类长期记忆方案在推理时仍要外部检索，kNN-LM 面临巨大的存储和最近邻搜索开销，RAG 还要额外处理检索到的篇章；Titans 在测试时在线写记忆；而原始 Memory Decoder 只在 1B 参数、百万级 token 的小规模上验证过，能否在真正的预训练规模（百 B token、数十 B 参数）上继续有效、通用记忆能否习得广泛知识、参数预算如何在基座和记忆间分配，都仍是开放问题。
本文的目标是本文要把参数化记忆预训练真正推到语言模型预训练的规模：训练最高达 6.9B 参数的通用记忆模块，在 300B token 上预训练，并在 Pile（207B token）这一规模上构建 kNN 分布。核心目标是验证三件事——记忆预训练能否随模型和数据规模继续扩展、一个通用记忆能否获得广谱知识、以及把参数预算更多地分给记忆是否比单纯放大基座更划算。此外还希望证明这种可插拔记忆能在不同基座规模间迁移、能做成领域记忆（生物/法律/金融）并跨词表跨模型族迁移，同时基座保持冻结、推理时无需任何外部检索。
与已有工作不同的是，本文的独特切入角度是：不再在推理时做外部检索（抛弃 kNN-LM/RAG 的在线检索和 datastore），也不在测试时在线写记忆（区别于 Titans），而是把检索行为在预训练阶段蒸馏进一个独立、可插拔、可独立扩展的参数化模块。与原始 Memory Decoder 相比，本文的关键跨越在于工程和规模——开发了分布式 Faiss 流水线（OPQ 压缩 + IVF 分片 + GPU 并行检索）和稀疏 kNN 分布存储 + 分布式流式加载，才让 207B token 规模的检索监督构建成为可能。这个角度首次让「记忆」成为和「推理」可分离、可独立缩放、可在多模型间复用的一等公民。
核心方法
直觉上，人脑的记忆和推理由不同的功能系统承担，可以分离。本文据此把长期记忆从单一参数集里「拆」出来，做成一个独立的 Transformer 解码器 $M_\psi$，它在预训练时去模仿一个 kNN 检索器输出的下一词分布 $p_{ret}$，从而把检索行为「蒸馏」进权重。技术路线分三层：(1) 用冻结基座 $M_\theta$ 把语料每处上下文编码成 key、其下一词为 value，构成 datastore，对每个训练上下文检索 K 近邻得到稀疏分布 $p_{ret}$；(2) 训练记忆 $M_\psi$，目标同时包含与 $p_{ret}$ 的 KL 对齐和语料语言建模项；(3) 推理时基座和记忆并行处理同一上下文，两者分布按 $\alpha$ 插值得到最终预测。规模化的难点全在第(1)步的 kNN 分布构建上，为此本文设计了分布式 Faiss 流水线、稀疏存储和流式加载。通用记忆用 Pile 训练，领域记忆则在领域语料上用继续预训练后的模型构建 datastore 再训练记忆。
核心创新在于把「记忆」当作可独立预训练、可独立扩展、可热插拔的模块，并通过参数化的方式彻底摆脱推理时的外部检索。与 kNN-LM/RAG 的本质区别是：检索只发生在离线监督构建阶段，蒸馏进 $M_\psi$ 后推理时完全不需要 datastore 或在线最近邻搜索；与 Titans 的区别是记忆来自预训练语料而非测试时在线写入；与原始 Memory Decoder 的区别是规模和工程——把方法从 1B/百万 token 推到 6.9B/300B token，并配套了 OPQ256 压缩、IndexIVFPQ 分片、HNSW 路由的 GPU 并行检索流水线，以及把密集 kNN 分布变成稀疏存储（Pile 上平均每行仅 64.95 个 token-概率对，相对于 50304 词表约 250× 压缩）再流式按需加载，这才让百 B 规模的监督构建可行。这些工程使得「记忆独立缩放」从设想变成可复现的实验。
方法步骤详情
完整步骤如下。第一步，数据准备与 datastore 构建：从训练语料 $\mathcal{D}_{mem}$ 抽出上下文-目标对 $\mathcal{P}_{mem}$，用冻结基座 $M_\theta$ 取最后一层在最后位置的隐状态 $\phi_\theta(c_t)$ 作为 key $k_t$、其真实下一词 $x_t$ 为 value，构成 datastore $\mathcal{R}_{mem}=\{(k_t,x_t)\}$。第二步，规模化 kNN 分布构建（分布式 Faiss）：(a) 训练 OPQ256 把 4096 维隐状态压成 256 维检索向量；(b) 在压缩空间学一个 IVF 划分，用 IndexHNSW 当 quantizer 做向量分配和查询路由，每个分片存一段连续质心范围、实现为 IndexIVFPQ；(c) 检索时分两段——先用 HNSW 把查询路由到质心、按分片分组，再各分片在 GPU 上并行做 IndexIVFPQ 局部检索，写进程按 query id 合并候选得到最终 K 近邻，丢弃 key 等于 query 的自匹配项，按距离温度加权 $\lambda_i \propto \exp(-d(q,k_i)/\tau)$ 投票成稀疏 $p_{ret}$。第三步，稀疏存储与流式加载：用阈值 $\epsilon$ 过滤掉可忽略概率，只存保留的 token id 和概率（FP32 聚合去重、FP16 存概率、INT64/INT32 存偏移和 id），存储复杂度从 $O(N|V|)$ 降到 $O(N+M)$；预处理时把每条打包 LM 样本对应分布的行区间附上，shuffle 后仍能对齐，分布式 worker 用 mmap 只读各 batch 所需切片。第四步，记忆训练：$M_\psi$ 与 $M_\theta$ 共享 tokenizer 和输出词表，损失 $\mathcal{L}_{mem}=\beta D_{KL}(p_{ret}\|p_\psi)+(1-\beta)(-\log p_\psi(x_t|c_t))$，GPU 上保留概率先 FP32 重归一化再算 KL。领域记忆先用领域语料对构建 datastore 的模型做一轮 CPT，再构建该领域 $p_{ret}$、训练独立的领域记忆。第五步，推理：基座和记忆并行编码同上下文，$p_{final}(y|c_t)=(1-\alpha)p_\theta(y|c_t)+\alpha p_\psi(y|c_t)$，$\alpha$ 在验证集上调优后固定用于测试。
技术新颖性
技术新颖性体现在三方面。其一，工程上首次在 207B token 规模构建 kNN 检索监督：OPQ256 压缩（4096→256 维）兼顾存储与距离计算成本、IVF+HNSW quantizer 把超大单索引拆成可独立 GPU 检索的分片、两段式路由+局部检索把 $N$ 查询对 $N$ 索引的二次构造转成压缩+分片+数据并行的流程。其二，稀疏存储把密集分布的 $O(N|V|)$ 降为 $O(N+M)$，Pile 上约 250× 压缩，且分布式流式加载让训练只需读各 batch 所需切片、与标准 LM 预训练接口完全兼容。其三，方法学上首次系统验证记忆与推理可分离并独立缩放：提出参数预算在基座与记忆间分配的新视角，并用通用记忆（Pile）和领域记忆（生物/法律/金融，最高 4.4B token）两套实验、跨基座规模（Pythia 410M–12B、Qwen3 0.6B–14B）和跨词表跨模型族（OLMo-2/3-7B）迁移实验全面刻画其行为。
实验结果
核心发现可以分成几条主线。第一，通用记忆在 17 个任务上对冻结 Pythia 各档基座都带来稳定的平均分提升：1.4B 基座 32.76→34.36，2.8B 基座 33.89→35.49，6.9B 基座 36.30→37.79，且 6.9B 基座+6.9B 记忆(37.79)超过冻结的 12B 基座(37.24)。在等总参数等预算下，「小基座+记忆」优于「大基座」：1.4B 基座+1.4B 记忆(34.36)优于 2.8B 基座(33.89)。最大涨幅集中在知识任务，如 2.8B 基座上 TriviaQA 从 8.30 跳到 17.11、1.4B 基座上 2WikiMultiHopQA 从 16.89 到 22.57。第二（跨基座迁移，Figure 1）：410M 基座+6.9B 记忆把平均分从 29.86 抬到 37.34，超过 12B 基座的 37.24，且总参数少 39%；在达到相同平均分时，Base+Memory 配置分别比 2.8B/6.9B/12B 基座少用 33%/32%/42% 参数。第三（领域记忆，Table 2）：1.7B 领域记忆在 Qwen3 0.6B–14B 各档基座上都拿到三域最高平均，相对冻结基座平均涨 9.88/9.64/10.00/9.09/9.99 分，并比每档最强基线至少高 4.05 分（0.6B 上领先 8.53 分）。在 14B 基座上，BioInst/LawBench/FinEval 分别涨 17.96/8.97/3.04 分，切换领域只需激活对应记忆、基座保持冻结。第四（跨词表迁移，Table 3）：用仅 20% 标准预算把 Qwen3 记忆迁移到 OLMo 词表，OLMo-2-7B 平均 19.57→23.83（+4.26）、OLMo-3-7B 18.67→26.44（+7.77），六项全部提升。第五（消融）：few-shot 下记忆增益依旧（Figure 4，0/3/5-shot 分别 +1.43~1.87/+1.39~1.52/+1.22~1.62）；记忆越大收益越大（Figure 6，6.9B 记忆在 1.4B/2.8B/6.9B 基座上分别 +4.56/+3.67/+1.49）；训练预算越多越好（Table 4，5 epoch 比 1.5 epoch 把通用任务/知识平均从 56.67/13.99 提到 57.33/14.17）；记忆训练目标本身（而非插值接口）带来增益——Table 5 显示 1.7B 记忆比「附加的同规模 CPT 模块」在 BioInst 高 10.21/8.71 分、LawBench 高 1.68/3.93 分。最后两节还证明记忆内容更可提取（Table 6：EM@8,16 从 CPT 的 42.4% 升到记忆的 49.7%、域锚点补全从 22.6% 升到 56.5%）且对检索目标忠实（Table 7：top-token 匹配 86.62%、Pearson r=0.9174、均值 KL=0.1820、TV=0.0823）。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 17 个通用基准(ARC-E/C, LAMBADA, LogiQA, PIQA, SciQ, WinoGrande, MMLU, NQ-Open, TriviaQA, PopQA, 2WikiMultiHopQA, Bamboogle, HotpotQA, GPQA-main, TruthfulQA, HaluEval) 平均分 | AVG(%) | Pythia-6.9B + Mem-6.9B = 37.79 | 冻结 Pythia-6.9B = 36.30；冻结 Pythia-12B = 37.24 | +1.49 分，且超过 12B 基座 |
| 17 基准平均分（参数效率对照） | AVG(%) | Pythia-410M + Mem-6.9B = 37.34（约 7.3B 总参数） | 冻结 Pythia-12B = 37.24（12B 总参数） | 性能反超且总参数少 39% |
| TriviaQA(开放域问答) | EM(%) | Pythia-2.8B + Mem-2.8B = 17.11 | 冻结 Pythia-2.8B = 8.30 | +8.81 分，知识任务涨幅最大之一 |
| 生物/法律/金融三域平均(Qwen3-0.6B~14B) | 三域 Avg(%) | Qwen3-14B-Base + Mem-1.7B = 37.89 | 冻结 Qwen3-14B-Base = 27.90；+RAG = 31.63 | +9.99 分(相对冻结)，且比最强基线高 ≥4.05 分 |
| 跨词表迁移(OLMo-3-7B) | 三域 Avg(%) | OLMo-3-7B + Mem-1.7B(20% 预算) = 26.44 | OLMo-3-7B = 18.67；+CPT = 24.29 | +7.77 分，六项全升 |
| 可提取记忆(EM@8,16) | EM@8,16(%) | 1.7B 记忆模型 = 49.7% | 1.7B CPT 模型 = 42.4% | +7.3 分，记忆更可追溯 |
局限与改进
作者承认的主要局限是：尽管推理时无需外部检索，但训练阶段构建 kNN 目标分布仍有索引与检索开销，即使有了压缩+分片流水线，这一离线成本仍是额外预处理负担，且随语料规模增长。作者还指出未来可把固定插值系数 $\alpha$ 换成依据上下文/模型置信度自适应加权的机制，并探索记忆与基座的联合或分阶段训练以提升协同（代价是训练成本上升、可能削弱跨基座迁移）。我自己的观察还有几点：(1) 稀疏存储里用了固定阈值 $\epsilon$ 截断，对长尾 diffuse 目标分布可能丢信息，Table 7 也显示残余误差主要来自 diffuse 目标；(2) $\alpha$ 在验证集上调优后固定，对知识密集 vs 推理密集任务用同一权重未必最优；(3) 领域记忆需先对建 datastore 的模型做一轮 CPT，仍有领域适配成本，且领域知识覆盖依赖语料质量；(4) TruthfulQA/HaluEval 等幻觉与事实性任务上记忆的增益不稳定甚至略降（如 6.9B 上 TruthfulQA 从 28.49 降到 27.81），说明记忆注入的知识也可能引入新的不实信号；(5) 主要在 Pythia/Qwen3 这类基座上验证，对更大前沿模型（数十 B 以上）的迁移和性价比尚未触及。
独立分析的弱点
独立分析的弱点及改进方向：(1) 离线 kNN 构建成本随语料二次增长——建议把固定 datastore 换成增量式/在线更新的近似索引，或用学习型索引替代 Faiss，进一步压低 207B+ 规模的预处理开销。(2) 固定 $\alpha$ 粗糙——可引入 per-token 或 per-context 的门控网络（类似 MoE 路由）自适应决定记忆权重，尤其能缓解 TruthfulQA 这类记忆反而拉低性能的负面案例。(3) 稀疏 $p_{ret}$ 用阈值 $\epsilon$ 截断损失长尾——可改用 top-k 截断加重要性采样，或在 KL 之外用更能建模不确定性的散度（如 Jensen-Shannon）来减少 diffuse 目标的误差。(4) 跨词表迁移只用 20% 预算且效果依赖 OLMo-3 datastore，词表差异极大时迁移质量未充分验证——建议系统刻画词表重叠度与迁移增益的关系，并提供更稳健的嵌入/head 初始化方案。(5) 实验规模止于 14B 基座和 6.9B 记忆，对当前数十 B 级前沿模型缺乏验证——应补做在更大基座上的扩展规律拟合，给出记忆与基座的最优配比曲线。(6) 评测多为 zero-shot 静态基准，缺少对真实长尾知识、时效性知识和对抗性事实的鲁棒性测试。
未来方向
作者明确提出的方向包括：用自适应权重机制取代固定 $\alpha$、研究记忆与基座的联合或分阶段训练、降低离线 kNN 构建成本。基于本成果可延伸的方向我认为还有：把参数化记忆与 RAG 结合（记忆负责稳定常识、RAG 负责时效/长尾），形成「预训练记忆 + 在线检索」的混合栈；把记忆做成多专家记忆库（Mixture of Memories）并用路由器按领域/任务动态选择，进一步扩展可服务领域数；探索记忆的持续学习与遗忘机制（如何增量写入新知识而避免旧知识被覆盖）；把记忆模块迁移到多模态（图文记忆）和 agent 长期记忆场景；以及从理论角度分析 kNN 蒸馏相比标准 LM 的样本复杂度优势，给出记忆容量与可学知识量的标度律。此外可研究记忆模块的隐私与安全——参数化记忆把训练内容更可提取(Table 6)，如何做记忆层面的差分隐私与遗忘也是现实需求。
复现评估
复现评估：工程与数据门槛较高。资源方面作者用了 256 块 NVIDIA A800 80GB GPU、基于 Megatron-LM 训练，单这一项就超出多数团队承受范围；构建 datastore 用的 Pythia-6.9B、Qwen3-4B-Base 等基座和 Pile(207B token)/领域语料(Biology-Instructions、DISC-Law-SFT、FinTrain，最高 4.4B token)多数为公开数据，但分布式 Faiss 流水线(OPQ256 + IVF/HNSW + IndexIVFPQ 分片 + 并行检索 + 稀疏流式加载)的工程实现细节较多，论文给了架构但未必给出完整可运行代码。训练超参披露较全：记忆学习率 1.4B/2.8B/6.9B 分别为 $3\times10^{-4}/2.5\times10^{-4}/2\times10^{-4}$，AdamW($\beta_1=0.9,\beta_2=0.95$)、weight decay 0.01、余弦衰减到 10%、2000 步 warmup、300B token(约 1.5 epoch)。论文提到代码仓 LUMIA-Group/MemoryDecoder-at-Scale 和 Rubin-wei/MemoryDecoder-at-Scale，若有开放则可显著降低复现门槛；即便如此，207B 规模 kNN 监督的构建仍需相当算力和分布式工程经验。综合看：方法描述清晰、超参完整、有代码仓线索，复现中等工作量与算力的实验（如小基座+小记忆）较可行，而完整复现 6.9B/300B 规模结果难度很高。
论文图表
