论文测试时学习 ★ 必读

Titans: Learning to Memorize at Test Time

Titans:用"惊讶度"在测试时把上下文写进权重的神经长期记忆

Ali Behrouz, Peilin Zhong, Vahab Mirrokni · Google Research · arXiv · 2024-12 · 被引 313

一句话Google Research 提出神经长期记忆模块 LMM:一个小 MLP 在推理时也用带动量和遗忘门的梯度下降在线更新权重,把历史上下文"写进参数"而非堆 KV cache;与 attention 组成 MAC/MAG/MAL 三种架构,760M 规模下语言建模平均分超过 Gated DeltaNet 与 Samba,BABILong 上微调后在 10M token 长度仍保持约 70% 准确率、胜过 GPT-4 和 Llama3.1-70B——但论文承诺的官方代码 20 个月后仍未放出,战绩无人完整复现。

这是什么

长上下文的两条老路各有死穴:attention 对窗口内依赖建模精确,但代价是 O(N^2) 且窗口有限,本质上是短期记忆;Mamba、DeltaNet 这类线性递归模型把全部历史压进固定大小的 hidden state,序列一长信息就被挤掉。Titans 提出第三条路:在模型里内置一个"神经长期记忆模块"(LMM)——一个 1-4 层的小 MLP,它的权重在测试时也随输入序列用梯度下降在线更新,相当于把上下文压缩进参数本身,而不是外挂向量数据库或 RAG。

驱动更新的信号借用认知科学里的"惊讶度":对当前 token 的重建损失梯度越大,说明这个输入越出乎意料、越值得记住。加上动量(过去的惊讶会延续)和自适应权重衰减(遗忘门),整个记忆更新恰好等价于带 momentum 和 weight decay 的在线 SGD,和现代线性 RNN 的 gating 机制在数学上接得上。这条"test-time learning"路线直接承接 Sun et al. 的 TTT(2024),Titans 把更新规则做得更完整,之后 Google 又沿这条线推出 MIRAS 和 ATLAS。

论文 2024 年 12 月 31 日挂 arXiv(2501.00663),作者 Ali Behrouz、Peilin Zhong、Vahab Mirrokni 三人,最终以 NeurIPS 2025 poster 接收,Semantic Scholar 引用已 313 次,是"把长期记忆做进架构"这一波讨论的引爆点。注意:至今没有官方开源代码,arXiv 也停在 v1。

MAC(Memory as Context)架构总览,也是理解 Titans 三层记忆分工的关键图:蓝色 Core 分支是滑窗/分段 attention(短期记忆),黄色 Contextual Memory 是测试时仍在梯度更新的神经长期记忆(先检索、attention 后再写入),红色 Persistent Memory 是任务级的固定可学习 token。最右侧标注了测试时各部分的状态:长期记忆在学、attention 做 in-context learning、persistent 冻结。

机制与做法

记忆即在线学习:惊讶度、动量与遗忘门

LMM 学的是联想记忆(associative memory):对每个 token x_t 用两个线性投影得到 k_t、v_t,记忆模块 M 的内层目标是最小化 ||M(k_t) - v_t||^2。更新规则为 M_t = (1-α_t)·M_{t-1} + S_t,其中 S_t = η_t·S_{t-1} - θ_t·∇ℓ(M_{t-1}; x_t):梯度项是"瞬时惊讶",动量项 S 是"过去惊讶"的记忆(避免大惊讶之后的相关信息因梯度变小而漏记),α_t 是数据依赖的遗忘门(权重衰减),η_t、θ_t 也都是 token 的函数。检索则是纯前向:y_t = M*(q_t),不更新权重。

W_K、W_V、投影层等属于外层参数,由正常训练(outer loop)学出;M 的权重在 inner loop 里持续被序列本身"训练",训练和测试时都在更新——这就是标题里的 learning to memorize at test time。记忆本体选深 MLP 而非矩阵是关键设计:矩阵记忆等价于在线线性回归,只能拟合线性依赖;实验中 L_M≥2 的深记忆在长序列上 perplexity 明显更稳。

论文的头牌战绩:BABILong 微调设置下,约 760M 的 Titans (MAC)-FT(红星)在 10K-1M token 长度保持 95%+ 准确率,10M token 仍约 70%;GPT-4、Qwen2.5-72B、Llama3.1-70B 在 100K 左右就掉到 60% 以下,Llama3.1-8B+RAG 全程约 40-48%。注意:此图无第三方复现,横轴是对数刻度。

并行化:chunk 内 mini-batch 梯度 + associative scan

逐 token 做梯度更新看似天然串行,论文沿用 TTT 的技巧把序列切成 chunk,在 chunk 内做 mini-batch 梯度下降,把梯度和权重衰减项重写成纯 matmul 形式;动量项 S_t = η_t·S_{t-1} - θ_t·u_t 是标准线性递推,可用 parallel associative scan 并行计算。代价是 chunk 内的梯度都对 chunk 起点的权重求(近似),换来 TPU/GPU 友好的实现。

实测训练吞吐:LMM 与 Mamba、Mamba2、Gated DeltaNet 同量级(略慢,归因于深记忆和对方的优化 kernel),且随序列长度保持常数吞吐(线性复杂度);Transformer++ 在 16K 时吞吐已腰斩。

记忆深度消融(170M 规模,Pile 子集):记忆 MLP 从 1 层加深到 3-4 层,各序列长度下 perplexity 持续下降,且 Mamba 在 32K 处明显上翘(遗忘失控)而 LMM 保持平稳——支撑"深的、带遗忘门的权重记忆优于线性状态压缩"这一核心论点。

三种组装:MAC / MAG / MAL,外加 persistent memory

除了 LMM(上下文记忆),Titans 还有一组可学习但与输入无关的 persistent memory token,拼在序列开头,存任务级知识、顺带缓解 attention 对开头 token 的偏置。围绕"长期记忆放哪"给了三种变体:MAC(Memory as Context)把序列分段,用当前段查询记忆、把检索结果和 persistent token 拼进 attention 的上下文,attention 反过来决定哪些信息写入记忆;MAG(Memory as Gate)不分段,滑窗 attention 与记忆输出做逐元素 gating;MAL(Memory as Layer)最平凡,把记忆当一层堆在 attention 前面——也是现有 hybrid 模型(Samba 等)的通用做法。

结论:MAC 和 MAG 全面好于 MAL,说明"怎么接"本身有信息量,常见的层堆叠是三者里最差的;MAC 长上下文最强,MAL 靠 Flash-Attention 吞吐最高,取舍明确。论文还给了一个(v1 中无证明的)定理:Titans 的表达力超出 TC0,强于 Transformer、对角线性 RNN 和 DeltaNet 的 state tracking 能力。

训练吞吐对比(10^3 tokens/s vs 序列长度):所有线性复杂度模型(含 LMM 和三种 Titans)吞吐随长度基本不变,LMM 略慢于 Mamba2 和 Gated DeltaNet;Transformer++(粉线)从 2K 的 50K tokens/s 一路跌到 16K 的 24K。说明测试时梯度更新的开销工程上可接受,但也没有速度优势。

消融:遗忘门贡献最大

以 LMM 为基线逐项去除:去掉 weight decay(遗忘门)伤害最大(ppl 27.01→29.04),其次是动量(28.98)、卷积(28.73)、深记忆换线性记忆(28.49)、persistent memory(27.63)。这组数字是论文对 TTT 的核心差异化论据——TTT 没有遗忘机制,而"能忘"恰恰是长序列下管住有限记忆容量的关键。

关键结果

语言建模+常识推理(760M / 30B tokens,FineWeb-Edu):纯递归的 LMM 平均分 51.56,超过 Gated DeltaNet 49.69 和 Transformer++ 48.69;混合变体 MAC 52.51、MAG 52.50,高于 Samba 51.08 和 Gated DeltaNet-H2 51.49。
BABILong few-shot:MAC 以远小于基线的参数量,在 1M token 长度上保持约 52-55% 准确率,曲线全程压过 GPT-4、GPT4o-mini、Llama3.1-8B、RecurrentGemma-9B;微调设置下 10M token 长度仍约 70%,而参数多约 70 倍的 Llama3.1-8B+RAG 只有约 40%。
S-NIAH(RULER,2K-16K):neural memory 在所有长度上最好且不随长度掉点;对比项的失败模式各有归因——TTT 无遗忘机制、Mamba2 无法擦除记忆(16K 时显著崩)。
记忆深度:L_M 从 1 加到 3-4,32K 序列上 perplexity 持续改善且对长度更鲁棒,但训练吞吐随深度线性下降,是效果/效率的直接权衡。
消融:遗忘门 > 动量 > 卷积 > 深记忆 > persistent memory,全部正贡献;LMM 基线 ppl 27.01,去掉遗忘门恶化到 29.04。
所有实验最大只到 760M 参数 / 30B tokens;论文脚注承诺"下一版报告更大模型结果",至今没有下一版。

实证核查

有水分思路真、影响力真(NeurIPS 2025 接收、313 引用、带火 test-time memory 一波),但代码承诺 20 个月未兑现、实现细节不全,BABILong 上"小模型打赢 GPT-4"的头牌战绩至今无人独立复现;实验规模停在 760M,承诺的更大规模结果也从未出现。核心机制(遗忘门+动量的在线记忆)可信,头条数字建议打折看。

结论部分明确承诺:"Titans are implemented in Pytorch and JAX and we intend to make the code we used to train and evaluate our models available soon."(source/MainText/Conclusion.tex 第 5 行)

截至 2026-08,没有任何官方仓库;Google Research 后来的 Titans+MIRAS 官方博客(research.google/blog/titans-miras-helping-ai-have-long-term-memory)依旧未附代码,HN 相关讨论(item 46181231)直接以"It's not open source"开场。社区事实标准是 lucidrains/titans-pytorch——README 自标 unofficial,1981 stars,维护到 2026-07。

BABILong 微调设置下,约 760M 的 Titans (MAC) 胜过 GPT-4、Qwen2.5-72B、Llama3.1-70B,10M token 长度仍约 70% 准确率(论文 Fig. BABILong-FT)。

无 checkpoint、无评测脚本,该数字只能听信论文。第三方的一致抱怨是细节不足:Reddit r/mlscaling 热帖(1i291tw)"Not enough specific details to replicate. Needs more ablations";alphaxiv 评论区甚至猜测"implementation details are intentionally left out to hinder independent reproduction"。unofficial 实现的 issues 印证了复现之难:#32(MAC 里记忆更新与 attention 的先后顺序按论文描述对不上)、#53(MAC 训练出 NaN)、#29(长序列不稳定)、#51(没有任何预训练权重)。至今没有公开的完整复现报告确认 BABILong 战绩。

实验部分脚注:"We are working on finalizing the results of larger models and will report them in the next version."

arXiv 至今只有 v1(2024-12-31 提交,curl arxiv.org/abs/2501.00663 只见 [v1]),更大规模结果从未出现;全部实验停留在 170M-760M / 15-30B tokens。"取代 Transformer"的媒体叙事与实际验证规模之间有明显缺口。

正文给出定理:Titans 表达力超越 TC0,强于 Transformer 与 DeltaNet(Methods.tex 第 278-280 行)。

v1 全文(含 Appendix)搜不到任何 proof;OpenReview 页面(forum 8GjSf9Rh7Z)可见评审批评"Despite promising formal proofs in the ICML rebuttal, the revised paper lacks the theoretical analysis"——说明该文先投 ICML 未中、理论承诺在修改稿中仍未兑现,最终以 NeurIPS 2025 poster 接收(neurips.cc/virtual/2025/poster/119639)。

与我们方向的关系

对做 continual learning 的同学,Titans 是"把持续学习做进架构"的代表:记忆不是外挂数据库,而是一组在部署后仍持续被梯度更新的权重,遗忘由可学习的 weight decay 门控——这正是 stability-plasticity 权衡的一个在线、可微版本。它的消融给了一个可直接借鉴的结论:在权重级记忆里,"会忘"(遗忘门)比"记得深"(MLP 深度)贡献更大,做任何 test-time 权重更新方案时应优先设计遗忘/容量管理机制。

工程上的可用资产是 lucidrains/titans-pytorch(注意它与论文在 MAC 更新顺序等细节上存在已知出入,见 issue #32),以及后续同一作者线的 ATLAS(2505.23735)与 MIRAS(2504.13173)。若要在组里实验,建议从 LMM 单模块(不带 attention)在小规模语言建模上复现 ablation 表开始,那部分机制清晰、最可能复现;BABILong 级别的战绩不要作为立项依据。

阅读笔记

读 Methods 时注意 inner/outer loop 的分工:W_K/W_V/W_Q 和门控网络是外层参数(训练后冻结),只有记忆 MLP 的权重在测试时更新。MAC 的段内流程是检索->拼接->attention->用 attention 输出更新记忆,lucidrains 实现里顺序不同,自己实现时以论文公式 (21)-(25) 为准。另外论文正文没写清 chunk size、门控网络结构等超参,这是所有复现者共同的坑。