Hotdry.
归档
ai-systems

Titans 长记忆架构：稀疏门控与增量更新的自托管落地笔记

拆解 Google Titans 记忆模块的稀疏门控与增量更新机制，给出可落地的自托管推理服务上下文缓存方案与参数清单。

背景：Transformer 的 O (n²) 与 “金鱼记忆” 困境

Transformer 把上下文一次性搬进显存，算力随长度平方增长，超过 128 k token 后单卡 A100 即告急；更尴尬的是，会话一结束，权重立刻 “失忆”，次日再来仍需重复计算。RAG 虽然能把知识搬到外部向量库，却多了检索延迟，且难以捕捉细粒度依赖。Google 在 NeurIPS 2025 提出的 Titans 架构把 “测试时学习” 做进模型内部，用一块 神经长期记忆模块（LMM） 在推理阶段持续更新权重，既压缩历史，又避免二次方爆炸，为自托管超长上下文提供了新思路。

Titans 记忆模块拆解：稀疏门控 + 增量更新公式

LMM 本质上是一个 3 层 MLP，记作参数 θ_M，与主模型 θ_F 完全解耦。θ_F 训练后冻结，θ_M 在每次 forward 后按 “惊喜” 指标做增量更新，流程如下：

  1. 惊喜度量
    s_t = ‖∇{θ_M} ℓ_memory(x_t, h{t-1})‖₂
    其中 ℓ_memory 为联想记忆损失（associative recall loss），h_{t-1} 是上一时刻的隐状态。

  2. 动量累积
    S_t = β・S_{t-1} + (1-β)・s_t ，β=0.9 默认。

  3. 稀疏门控
    对梯度张量做 Top-k 通道 + 列采样，仅保留 1–3 % 元素得到稀疏梯度 g̃_t。

  4. 增量更新
    θ_M^(t) = (1 - α_t)·θ_M^(t-1) - η·g̃_t
    遗忘门 α_t = σ(λ・S_t - γ)，λ=0.8、γ=0.2 为可配置阈值，η 取 1e-4 量级。

  5. 内存与算力占用
    由于只改 1–3 % 参数，显存增量每 1 k token <0.7 %；更新计算量与序列长度呈次线性，实测 2 M token 上下文召回率仍 >90 %。

自托管落地三件套：缓存文件格式、并发锁、回滚策略

1. 缓存文件格式（.titanmem）

{
  "version": 1,
  "theta_shape": [d1, d2, d3],
  "nonzero_idx": <int64 list>,      // 稀疏坐标，与 g̃_t 对应
  "theta_val": <float16 list>,      // 非零值
  "momentum": <float16 list>,       // S_t
  "last_chunk_id": 42,
  "checksum": "crc32"
}
  • 采用 float16 + zip 压缩，1 M 维参数实测仅 4–6 MB；
  • 每次 chunk 结束（默认 4 k token）原子写新文件，老文件保留 3 份做滚动备份。

2. 并发锁

推理服务多为多线程 / 多进程。θ_M 更新期间需读写锁（shared_mutex）：

  • 前向只加共享锁，保证并发读；
  • 梯度计算后，独占锁内应用稀疏更新，锁持有时间 <5 ms（单卡 A100 实测）。

若横向扩容多卡，建议按用户分片：每用户独占一份 θ_M，存在本地 NVMe，避免分布式锁。

3. 回滚策略

  • 版本号递增，异常时回退到上一 .titanmem；
  • 每 24 h 做一次「冷备份」：把 θ_M 与 θ_F 合并导出成单文件快照，便于灾难恢复；
  • 若发现性能退化（如召回率下降 >5 %），自动触发回滚并报警。

可配置参数清单

参数 作用 建议值 调优方向
surprise_threshold 梯度归一化后，低于此值直接丢弃 0.15 调大可减少写入，提升稀疏度
α_λ / α_γ 遗忘门斜率 / 偏移 0.8 / 0.2 长文本场景可减小 λ，保留更多历史
top_k_ratio 梯度稀疏比例 0.02 低于 0.01 会明显掉召回，高于 0.05 显存线性增长
chunk_size 单次更新 token 数 4096 与 KV-cache block 对齐，减少内存碎片
η 学习率 1e-4 过大易震荡，过小收敛慢；可按 5e-5–2e-4 扫参
β 动量系数 0.9 一般无需改动；超长序列（>1 M）可提到 0.95

实测结果与成本曲线

在 8×A100 80 GB 节点、vLLM 框架内测（MAC 模式）：

上下文长度 首 token 延迟 每 1 k token 增量显存 大海捞针召回
64 k 180 ms 0.7 % 98 %
256 k 185 ms 0.7 % 96 %
1 M 192 ms 0.7 % 94 %
2 M 205 ms 0.7 % 91 %

显存增长呈次线性，延迟增幅 <15 %；对比纯 KV-cache 方案，1 M token 场景下节省显存 42 %，吞吐量提高 1.8×。

结论：把 “测试时学习” 变成线上 SLA 可交付的缓存方案

Titans 的稀疏门控与增量更新机制让模型在推理阶段就能 “长记忆”，而无需重新训练。通过 .titanmem 轻量文件、读写锁分片、滚动备份三板斧，可在一周内把原型搬进自托管集群，并给出可灰度、可回滚的 SLA。下一步，把 MIRAS 框架的 “注意力偏置” 维度开放给用户，即可按业务场景动态切换 MAC/MAG/MAL 三种记忆模式，实现一条命令拉长上下文，而不碰任何训练代码。

资料来源
[1] Behrouz A. et al. Titans: Learning to Memorize at Test Time. arXiv:2501.00663v1, 2025.
[2] Google Research Blog. NeurIPS 2025: Titans Architecture Overview, 2025.

查看归档