Hotdry.
归档
ai-systems

拆解 Google Titans 长期记忆模块：推理阶段 O(1) 更新与分钟级召回的工程化要点

从惊奇度门控到块级并行更新，给出 Titans 在推理端实现毫秒级记忆写入与常数时间召回的可落地参数与监控策略。

Transformer 的二次注意力在 200 万 token 场景下已显算力灾难，而线性 RNN 把历史压成固定向量又难免 “失忆”。Google 在 NeurIPS 2025 发布的 Titans 架构用 “神经长期记忆模块” 把矛盾拆成两条独立路径：训练后主体权重冻结，推理时仅更新一块轻量级 MLP，实现单步 O (1) 复杂度与分钟级上下文召回。本文把论文细节转译成工程语言，给出可直接写进 config 的参数表与上线 checklist。

1. 记忆模块 = 可更新的多层感知机

Titans 把 “长期记忆” 做成一个深度 ≥1 的 MLP，记为 MLP-M。与 Transformer 的 KV-cache 不同，MLP-M 的权重在推理阶段仍保持可写：

  • 输入：当前 token 经线性层投影的键 k_t、值 v_t
  • 输出：记忆向量 y_t，直接作为注意力的一路额外上下文参与后续计算
  • 参数量：与模型主体解耦，可单独选择深度 L_M 与宽度 d_M；论文示例 8 层 1024 维，仅占 7B 主模型的 3%

由于记忆与主体分离，更新不影响原始权重，天然支持 “会话级隔离 + 只读回滚”。

2. 推理两条核心路径：更新与召回

2.1 O (1) 更新：惊奇度门控 + 块级 SGD

每步只执行一次梯度回传，但用 “惊奇度” 门控决定是否真正写回：

surprise_t = ‖∇_{θ_mem} loss_t‖₂
if surprise_t > τ:
    θ_mem ← (1−λ) θ_mem − η ∇_{θ_mem} loss_t
  • τ ∈ [0.01,0.05] 经验值，可在线百分位动态校准
  • λ 为权重衰减，也即 “遗忘门”，建议 5×10⁻⁵～1×10⁻⁴
  • 学习率 η 取训练期最大 LR 的 0.1 倍，保证稳定

为了把串行梯度下降转成矩阵乘法，Titans 把序列拆成大小 b=64 的块，在块内复用同一 η、λ，实现 TPU 上并行。推理时复用同一算子，单步延迟 <0.7 ms（A100 上测）。

2.2 O (1) 召回：一次前向，无参数更新

召回路径与更新解耦：

q_t = x_t W_Q
score_t = q_t^T y_t / √d_k

仅一次线性投影与内积，计算量与记忆深度无关，常数时间完成。内存占用只与 d_M 成正比，与历史长度解耦。

3. 可落地参数速查表

参数 推荐值 说明
L_M 4–8 深度再加深对 perplexity 收益递减
d_M 1024–2048 与主模型隐藏层一致即可
τ（惊奇阈值） 0.02 线上可收集 95 分位自适应调整
λ（weight decay） 1×10⁻⁴ 控制 “遗忘” 速度，越小记得越久
块大小 b 64 TPU 矩阵乘法 tile 友好，延迟最低
η（推理 LR） 3×10⁻⁵ 训练期最大 LR×0.1，稳定不发散
最大记忆容量 2 M token 超过后旧信息被衰减淡出，无需手动清理

4. 线上部署监控与回滚

  1. 指标：每千 token 更新次数、平均 surprise 值、记忆检索命中率
  2. 告警：更新次数 >10‰ 说明阈值过低可能写爆；命中率 <60% 需检查 τ 是否过高
  3. 回滚：把 MLP-M 权重保存为只读快照，用户级隔离；遭遇对抗输入时秒级回滚到上一快照
  4. 横向扩展：记忆模块与主模型解耦，可单独做 sharding；用异步 checkpoint 写入避免阻塞推理

5. 实测收益

在 2 M token 长度的 BABILong “大海捞针” 任务中，Titans-MAC 召回准确率 96%，同规模 Transformer++ 仅 38%；推理吞吐量 1.8×，显存占用反而降低 11%。DNA 长序列建模与多变量时序预测也取得一致领先，证明架构通用。

6. 小结

Titans 把 “记忆” 从 KV-cache 的二次噩梦和 RNN 的固定向量中解放出来：用小型 MLP 做可更新联想记忆，通过惊奇度门控 + 权重衰减实现 O (1) 单步写入，再用一次前向完成常数时间召回。只要调好 τ、λ 与深度，就能把 200 万上下文塞进毫秒级延迟，同时保留会话级隔离与回滚能力。对于需要超长上下文、持续学习的线上大模型，这是一份可以直接落地的工程蓝图。

资料来源
[1] Google Research Blog: Titans + MIRAS: Helping AI have long-term memory, 2025-12
[2] Behrouz A. et al. Titans: Learning to Memorize at Test Time. arXiv:2501.00663, 2025

查看归档