【arXiv 2025】Titans 论文解读：测试时记忆学习的神经长期记忆架构｜从神经网络记忆机制视角

摘要

本文解读 arXiv 2025 论文《Titans: Learning to Memorize at Test Time》。该论文提出Titans 神经长期记忆架构，通过融合短程注意力、测试时在线更新的神经记忆与可学习持久记忆三分支，让模型在推理时把重要历史持续写入自己的网络权重，其特别之处在于把记忆建模为在线学习问题——用惊奇度量决定记什么、动量保持记忆、权重衰减决定遗忘。实验表明760M 规模下 Wiki 困惑度达到 18.61、有效上下文超过 200 万 token，并在 BABILong 百万级长文档推理上超越 GPT-4，为长上下文序列建模提供了全新范式。

视频讲解：点击观看 B 站视频

论文基本信息

背景与动机

Transformer 的注意力机制本质上是一块联想记忆：模型存储键值关联，并用查询与键的相似度做检索。但注意力对当前上下文窗口之外的依赖无能为力，且计算量随序列长度二次增长——2 万 token 以上的序列就难以承受。另一条路线是线性注意力与状态空间模型（Mamba、Mamba2、GLA、RetNet 等），它们把历史压缩进固定大小的矩阵状态，推理成本线性，但压缩导致精度下降，加性写入还会造成记忆溢出。

从记忆科学视角看，两者都缺少人类记忆系统的关键组件：短期记忆、长期记忆、元记忆（持久记忆）以及它们之间的协同。Titans 的核心主张：长期记忆不该是固定结构，而应该是一个神经网络，在测试时用联想记忆损失在线更新自己的权重——越是出乎意料（惊奇）的信息越值得记住。这正是与 TTT、DeltaNet 一脉相承的测试时学习思想，但 Titans 补上了遗忘门、动量更新与深度记忆三个关键设计。

研究主线：从问题到结论

图 7：Titans 研究主线（Mermaid 流程图）：问题 → 动机 → 设计 → 方法 → 实验 → 结论

基准/方法设计

Titans 架构包含三个分支：核心分支用滑动窗口注意力负责短程精确建模；长期记忆分支是神经记忆模块（LMM），把过去的历史在线压缩进 MLP 权重；持久记忆分支是一组输入无关的可学习参数，编码任务级知识。作者给出三种接入方式：

• MAC（Memory as Context）：把记忆检索结果作为上下文拼进注意力输入，注意力同时决定"需要哪些历史"与"哪些信息值得写入"；
• MAG（Memory as Gate）：用门控机制融合记忆输出与滑动窗口注意力输出；
• MAL（Memory as Layer）：把记忆作为网络中的一层，先压缩再交给注意力。

图 1：Memory as a Context（MAC）架构：核心注意力 + 上下文长期记忆 + 持久记忆三分支

分类全景

图 8：Titans 记忆系统分类全景（Mermaid 流程图）：短程注意力、神经长期记忆 LMM 与持久记忆

方法细节

记忆模块的四个关键设计：

1. 惊奇度量（Surprise）：用输入梯度 $\nabla \ell(\mathcal{M}_{t-1}; x_t)$ 衡量意外程度——梯度越大，说明输入越偏离已有记忆，越值得记住；
2. 动量式更新：$S_t = \eta_t S_{t-1} - \theta_t \nabla \ell$，过去与当前的惊奇共同决定写入，避免一次大惊奇后梯度消失、错过后续重要信息；
3. 自适应遗忘：权重衰减门 $\alpha_t$ 管理有限容量，$\alpha_t \to 0$ 保留记忆、$\alpha_t \to 1$ 清空记忆，等价于数据依赖的遗忘机制；
4. 并行化训练：逐 token 串行梯度下降被重写为分块矩阵乘法 + 并行关联扫描，训练吞吐随序列长度线性扩展。

图 2：神经记忆的并行化训练：分块内用 matmul 与关联扫描，避免串行 O(N) 内循环

理论结果同样重要：Titans 的深度记忆在状态追踪任务上表达力超出 TC⁰，而 Transformer 与多数线性循环模型被限制在 TC⁰ 内。附录中进一步证明，LMM 是 Gated DeltaNet 的严格推广——DeltaNet、Longhorn、RWKV-7、TTT 均可视为其特例（LMM 补上了动量规则、深度记忆与遗忘门）。

图 5：Memory as a Layer（MAL）架构：记忆层先压缩过去与当前上下文，再交给滑动窗口注意力

实验设计与结果

评测覆盖 170M / 340M / 400M / 760M 四档规模，前三个在 FineWeb-Edu 的 15B token 上训练，760M 用 30B token。任务分五类：语言建模（WikiText、LMB）、常识推理（PIQA、HellaSwag、WinoGrande、ARC、SIQA、BoolQ）、长上下文（RULER S-NIAH、BABILong）、时间序列（ETT、ECL、Traffic、Weather）与 DNA 建模（GenomicsBenchmarks）。基线包括 Transformer++、Mamba、Mamba2、GLA、RetNet、DeltaNet、Gated DeltaNet、TTT 及混合模型 Samba、Gated DeltaNet-H2。

760M / 30B tokens 语言建模主表：

Titans 在困惑度与推理两个维度同时压过同规模全部模型：MAG 的 Wiki 困惑度 18.61 比最强混合基线 Gated DeltaNet-H2 低 1.27，平均准确率 52.51（MAC）领先 1.02 个百分点。纯记忆模块 LMM（不带注意力）在简单模型中同样最优，证明记忆机制本身有效。

图 3：BABILong 微调设定：Titans（MAC）在 1M 级长文档推理上超越 GPT-4 等超大模型

图 4：训练吞吐对比：Titans（MAL）吞吐随序列长度保持恒定，快于多数循环基线

长上下文是 Titans 的主场。S-NIAH（RULER 基准）中，序列从 2K 拉到 16K 时 TTT 从 98.4 掉到 88.4、Mamba2 从 98.6 崩到 5.4、DeltaNet 从 96.8 掉到 71.4，而 Titans 三个变体在 16K 仍保持 96 以上（MAC 达 98.4）；BABILong 微调设定下，Titans（MAC）以不到 8 亿参数超越 Llama3.1-8B+RAG、GPT-4 与 Qwen2.5-72B。时间序列预测中，神经记忆在 7 个数据集的 6 个取得最优 MSE（ETTm2 0.261、ECL 0.162、Weather 0.231 等），全面优于 Simba、iTransformer、PatchTST。DNA 建模（GenomicsBenchmarks）上神经记忆在 Enhancer Cohn（75.2）与 Enhancer Ens（89.6）取得最优，与最强 DNA 模型 HyenaDNA 竞争性相当。

消融实验（附录 E）显示所有组件均为正向贡献，影响排序为权重衰减 > 动量 > 卷积 > 持久记忆：去掉权重衰减后长上下文准确率从 92.68 掉到 85.60，印证遗忘机制对记忆容量管理的核心作用；记忆深度从 1 层增加到 4 层，困惑度持续下降且对序列长度更鲁棒（但吞吐随深度线性下降，存在效果与效率权衡）。

图 6：不同记忆深度下 ppl 随序列长度的变化：更深的记忆在更长序列上收益更明显（170M 规模）

结果对比总结

图 9：结果对比总结（Mermaid 流程图）：Titans 在困惑度与准确率双维度领先基线

关键发现

• 记忆即在线学习：长期记忆是一个在测试时做梯度下降的神经网络，联想记忆损失 $\ell = | \mathcal{M}(k_t) - v_t |_2^2$ 驱动写入，这是全文最核心的范式转变。
• 有效上下文超 200 万 token：在 S-NIAH 与 BABILong 上，Titans 在 16K–1M 级序列保持 96% 以上准确率，而 Mamba2 在 16K 时已跌至 5.4%。
• 760M 规模全面领先：Wiki 困惑度 18.61（MAG）与平均准确率 52.51（MAC），同时击败同规模 Transformer++（25.21 / 48.69）与全部混合基线。
• 纯记忆模块即可打满全场：不带注意力的 LMM 在简单模型中困惑度最优（760M 下 20.04），证明记忆机制独立有效。
• 精度与效率兼得：MAL 变体训练吞吐高于多数循环基线，得益于 Flash-Attention 与并行化记忆更新的组合。
• 跨领域泛化：时间序列 6/7 数据集最优、DNA 建模与 SOTA 竞争性相当，机制不止适用于自然语言。

局限性

• 规模验证有限：论文第一版仅报告到 760M / 30B token，更大规模的结论待后续版本。
• 实现开销：深度记忆的测试时更新带来额外计算，吞吐略低于内核高度优化的 Mamba2。
• 可解释性不足：记忆写入权重的语义难以直接审计，不像 KV 缓存那样显式可读。
• 流式稳定性：测试时持续学习在长流式场景下的稳定性仍需更充分验证。
• 作者计划开源 PyTorch / JAX 实现，并接入更复杂的记忆架构（如进化记忆、Muon 类更新）。

常见问题（FAQ）

Titans 和 TTT 有什么区别？

TTT 是 Titans 最接近的前作：两者都把记忆当作梯度更新的模型。但 TTT 没有遗忘机制，长序列会记忆溢出；Titans 引入权重衰减（自适应遗忘门）、动量式惊奇更新与深度记忆，并在消融中验证了每个组件的贡献。

MAC、MAG、MAL 三种变体怎么选？

MAC 长上下文最强（注意力帮记忆筛选信息）、MAG 语言建模困惑度最低（18.61）、MAL 训练吞吐最高（注意力走 Flash-Attention）。追求长程推理选 MAC，追求吞吐选 MAL。

Titans 为什么能超过 200 万 token？

因为长期记忆是参数化的：历史被在线压缩进网络权重而非显式 KV 缓存，推理成本不随历史长度增长；短程注意力只处理当前窗口，两者叠加使有效上下文远超窗口大小。

为什么说 LMM 是 Gated DeltaNet 的推广？

DeltaNet / Gated DeltaNet 的 Delta 规则只考虑瞬时惊奇且限于线性矩阵记忆；LMM 加入动量规则、深度 MLP 记忆与非线性递推，令前者成为 $\eta_t = 0$ 时的特例，RWKV-7、Longhorn 同理。

神经记忆在非语言任务上表现如何？

时间序列预测 7 个数据集中 6 个取得最优 MSE（如 ETTm2 0.261、Weather 0.231），DNA 建模在 GenomicsBenchmarks 与 HyenaDNA 竞争性相当，说明记忆机制对结构化数值与生物学序列同样有效。

有哪些已知局限？

760M 规模封顶、吞吐略慢于 Mamba2、记忆权重不可解释、流式场景稳定性待验证；代码当时尚未开源，社区复现需自行实现。

参考链接

• arXiv 论文页：https://arxiv.org/abs/2501.00663
• Mamba（线性时间序列建模）：https://arxiv.org/abs/2312.00752
• TTT（测试时学习的 RNN）：https://arxiv.org/abs/2407.04620
• Gated Delta Networks：https://arxiv.org/abs/2412.06464
• Agent Memory Survey（Titans 列为参数化记忆代表）：https://arxiv.org/abs/2512.13564

给大家推荐一款自用写文献综述、无虚构文献的 AI：

🌟 复旦大学 FudanNLP 团队自研 切问学术

官网：qiewenpaper.com

覆盖 3.6 亿篇可溯源真实中英文文献，能自动整合文献观点生成规范综述

还能挖掘研究创新点、复现实验，配合视频教学，新手快速上手文献综述写作

🍀后记🍀

博客的关键词集中在编程、算法、机器人、人工智能、数学等等，持续高质量输出中。

🌸讨论QQ群：白拾的小屋 (750365700)

⭐B站账号：白拾的物理AI组会（活跃于知识区和动画区）

✨GitHub主页：YhbCode000（工程文件）