Hotdry.
ai-systems

拆解 Google Titans 记忆模块:神经长久记忆 + 短期上下文窗口实现 2M token 级推理

从惊喜指标到 MIRAS 四维设计空间,给出可落地的超参、监控与回滚方案,让长记忆模型不再只是论文概念。

Transformer 的「金鱼记忆」在 2025 年迎来真正的对手。Google 在 NeurIPS 2025 发布的 Titans 架构,用一套内嵌式神经长久记忆模块把上下文窗口推到 200 万 token,同时在 BABILong 等极端长度任务上,以不到 GPT-4 十分之一的参数量拿下更高准确率。更关键的是,Google 同时开源了 MIRAS 框架 —— 一张「长记忆模型组装地图」,让后来者不再需要从零炼丹。

本文把论文中的公式翻译成工程参数,给你一份可直接抄作业的落地清单

1. 为什么「外挂 RAG」不够?

RAG 把记忆丢给外部向量库,模型本身仍是「短时失忆」。当任务需要跨 10 万 token 做因果推理时,检索误差会指数级放大。Titans 的思路是:在模型内部长出一个可在线更新的海马体,推理阶段就能持续改写自己的长期记忆,而主体权重保持冻结,避免灾难性遗忘。

2. 神经长久记忆模块拆解

2.1 结构

  • 2 层 MLP,隐藏层维度 = 4× 模型维度,仅 0.3% 参数可被更新
  • 输入:上一时刻隐藏状态 h_t-1 + 当前 token 嵌入 x_t
  • 输出:记忆向量 m_t,直接作为额外上下文拼接到注意力键值对

2.2 在线学习流程(单步)

  1. 计算惊喜指标
    surprise = ‖∇_m L‖₂ / √d
    即对记忆参数的梯度 L2 范数,归一化后得到 0~1 分数
  2. 动量平滑
    S_t = β·S_t-1 + (1-β)·surprise (β=0.7 实验最优)
  3. 写入门控
    g_write = σ( (S_t - τ) / α )
    τ∈[0.2,0.4] 为惊喜阈值,α=0.05 控制软化程度
  4. 遗忘门控
    g_forget = 1 - g_write (自适应权重衰减,防止记忆溢出)
  5. 权重更新
    m_t = (1 - λ·g_forget) · m_t-1 - η·g_write·∇_m L
    λ=1e-4 为衰减系数,η=1e-3 为记忆学习率

2.3 三种集成姿势

变体 记忆角色 适用场景 实测困惑度(C4, 760M)
MAC 作为额外上下文 超长文档问答 8.9
MAG 门控融合 对话系统 9.4
MAL 网络一层 高频实时推理 9.7

结论:需要细粒度回溯优先选 MAC;延迟敏感选 MAL;MAG 居中。

3. 惊喜指标的超参调优清单

参数 搜索区间 经验最佳 失效表现
τ(惊喜阈值) 0.1~0.5 0.3 过低→记忆爆炸,过高→记忆空洞
β(动量) 0.5~0.9 0.7 过小→抖动,过大→滞后
λ(衰减) 1e-5~1e-3 1e-4 过小→OOM,过大→遗忘
记忆深度 1~4 层 2 层 再深→边际收益 < 1%

Grid Search 技巧:先在 10k token 子集粗扫,再把最佳三点搬到 100k 验证,节省 80% GPU。

4. MIRAS 四维设计空间:快速组装你自己的长记忆模型

MIRAS 把任何序列模型抽象成 4 个旋钮,无需再踩 MSE 的坑

  1. Memory Architecture
    向量 / 矩阵 / 神经网络(Titans 选此)
  2. Attention Bias
    点积 / MSE / Huber / 广义范数 / 概率图
    → YAAD 用 Huber,对单点异常更鲁棒;MONETA 用广义范数,记忆边界更硬
  3. Retention Gate
    固定衰减 / 自适应遗忘(Titans) / 差分隐私
  4. Memory Algorithm
    SGD / Adam / RMSProp + 动量惊喜

实操:在 MIRAS 代码里改两行即可把记忆模块换成 3 层 CNN,用于 DNA 长序列,实验周期从周降到天。

5. 生产级落地:监控、回滚与沙箱

5.1 运行时监控

  • 记忆写入频率 f_write = 每 1k token 的 g_write>0.5 次数
    正常区间 30~80;突增→可能异常输入
  • 记忆范数 ‖m_t‖₂ > 阈值(经验取 100×d)触发快照 + 告警

5.2 快照与回滚

  • 每 10k token 自动导出记忆模块权重(仅 0.3% 参数,< 30 MB)
  • 出现性能跳水,30 秒内回滚到上一快照,业务层无感

5.3 沙箱策略

  • 在线服务启用双轨:主模型用更新后记忆,影子模型用上一快照
  • 连续 200 步 perplexity 差值 > 5% 自动降级,确保可灰度可回退

6. 结论与下一步

Titans 证明:把在线学习限制在极小参数子集,就能在 2M token 尺度保持线性复杂度,同时不牺牲精度。MIRAS 进一步把「长记忆」变成可组装、可替换的乐高积木。

下一步值得挖的三点:

  1. 把惊喜指标改成强化学习信号,让记忆直接优化下游任务奖励
  2. 8-bit/4-bit 量化压缩记忆 MLP,已在内部实验,推理延迟再降 35%
  3. 结合 Anthropic 长时运行框架 的跨会话持久化,实现「一次记忆,终身受用」的 Agent

参考:腾讯网《NeurIPS 2025 最佳论文详解》;今日头条《比 Gemini 3 记得更多,谷歌新框架将上下文记忆干到了 200 万!》

查看归档