Memory Decoder at Scale: A Pretrained, Parametric Long-Term Memory
Rubin Wei, Jiaqi Cao, Jiarui Wang, Junming Zhang, Qipeng Guo, Bowen Zhou, Zhouhan Lin
cs.CL
2026-07-30
把参数化记忆模块扩到 69 亿参数、3000 亿 token,冻结的 Pythia-410M 配 6.9B 记忆体在 17 项基准上反超 Pythia-12B,总参数还少 39%。
标准的 decoder-only 大模型把长期记忆和推理能力塞在同一套参数里。这带来几个麻烦:记忆没法单独扩,要把模型「记得多」就得把整套参数一起变大;做领域适应得全量微调,训练贵,还容易灾难性遗忘;也没有一个能换、能跨模型复用的独立记忆模块。
已有的解法各有短板。RAG 在推理时检索外部段落当上下文,kNN-LM 把预测跟外部数据仓的最近邻分布做插值,两者推理时都要外部检索,要么多吃上下文,要么要建大索引、做近邻搜索。Memory Decoder(前作)提出了参数化记忆模块,但只做到约 10 亿参数、几百万 token 的规模,再往大能不能继续 work 没人验证。这篇就是把它推到大模型预训练的量级。
记忆模块是一个 transformer decoder,训练目标是去模仿一个 kNN 检索器的输出分布 pret,而不是只学单一的下一个 token。每个位置的损失是 β 倍的 KL(pret 和 pψ 的散度)加 (1-β) 倍的交叉熵语料语言模型项:KL 项让记忆去逼近检索分布,语言模型项防止它跑偏离真实语料。检索只在离线构造训练监督时用,推理时不需要。
推理时,冻结的主干模型和记忆模块并行处理同一段上下文,两者的下一个 token 分布按 pfinal 等于 (1-α) 倍主干分布加 α 倍记忆分布插值,α 控制记忆权重。
真正的工程难点在规模。要把记忆训到 69 亿参数、3000 亿 token,就得在去重的 Pile(2070 亿 token)上为每个训练上下文构造 kNN 分布。2.07 亿条记录下,索引和检索的联合成本让标准 Faiss 管线跑不动。作者的解法是一套分布式 Faiss 管线:嵌入压缩(OPQ256,把 4096 维压到 256 维)、索引分片、并行检索,再配合稀疏 kNN 分布存储和按批分布式流式加载,只读每个 batch 需要的条目。通用记忆用去重 Pile 训,数据仓用 Pythia-6.9B 建;领域记忆有生物、法律、金融三套,1.7B,用 Qwen3-4B-Base 的数据仓训。计算用 256 张 A800。
通用记忆在冻结的 Pythia 主干上、17 项基准上,等规模记忆在每个尺度都抬高了均分:1.4B 主干 32.76 提到 34.36,2.8B 33.89 到 35.49,6.9B 36.30 到 37.79。更关键的是同等总参数下的对比:1.4B 主干配 1.4B 记忆(34.36)高于 2.8B 单干(33.89),两者总参数和训练预算一样;6.9B 配 6.9B(37.79)超过冻结的 12B 主干(37.24)。51 个任务乘尺度组合里,记忆在 47 个上提升、1 个持平。知识类任务涨得最猛:TriviaQA 在 2.8B 主干上从 8.30 涨到 17.11,2WikiMultiHopQA 在 1.4B 上 16.89 到 22.57。
| 配置 | 总参数 | AVG(17 任务) |
| Pythia-410M(单独) | 410M | 29.86 |
| Pythia-410M + Mem-6.9B | 7.3B | 37.34 |
| Pythia-12B(单独) | 12B | 37.24 |
最抓人的是这条:410M 主干配 6.9B 记忆做到 37.34,反超 12B 主干的 37.24,总参数还少 39%。在同等均分下,主干加记忆的配置比 2.8B、6.9B、12B 单干分别少用 33%、32%、42% 的参数,而且主干和记忆能并行跑,推理延迟可以比同等总参数的单模型更低。领域记忆上,1.7B 记忆让 Qwen3 从 0.6B 到 14B 在生物、法律、金融三域均分各涨 9 到 10 分,每个尺度都赢过 CPT、LoRA、RAG 里最强的那个至少 4.05 分;切领域只需换记忆,主干冻结不动。跨词表迁到 OLMo 上,用 20% 预算,OLMo-2-7B 和 OLMo-3-7B 也分别涨 4.26 和 7.77 分。
它给出了一个跟「把主干越做越大」不同的提效路径:记忆和推理可以解耦,记忆单独预训练、单独扩、按域换。对做领域适应的人,换领域不用动主干、不用全量微调,换个记忆模块就行,绕开了灾难性遗忘。参数效率的账也漂亮,小主干配大记忆在同等分数下省下三到四成参数。
作者自己说,虽然推理不用检索了,但离线构造 kNN 目标分布的索引和检索成本还在,这是额外的预处理负担,随语料规模增长。推理时记忆权重 α 是固定的,没做到按输入自适应。记忆和主干的联合或分阶段训练还没探,这种训练可能让两者配合更好,但会抬高训练成本,还可能削弱跨主干迁移性。另外一个论文没正面回答的问题是:记忆模块把外部知识压进了参数,跟 RAG 比知识更新和可审计性如何,论文侧重大模型基准,没在时效性强的场景里检验。