Hotdry.
ai-systems

Titans 混合架构:滑动窗口注意力与神经长期记忆实现 AI 长时记忆

Titans 通过滑动窗口注意力捕捉短期依赖,神经长期记忆模块基于惊喜度量动态更新参数,实现超 200 万 token 上下文的高效长时记忆工程参数与部署要点。

Titans 架构的核心观点在于,传统 Transformer 的自注意力机制虽精确捕捉依赖,但二次复杂度限制了长上下文扩展,而线性 RNN/SSM 虽高效却信息压缩过度导致精度损失。Titans 提出混合设计:滑动窗口注意力(SWA)专注短期局部依赖,神经长期记忆模块(LTM)作为 MLP 元模型在测试时在线学习历史摘要,实现 RNN 速度 + Transformer 精度。该设计特别适用于 RAG 增强推理,避免全 KV 缓存膨胀,通过 learned neural indices(惊喜梯度驱动)按需检索。

证据支持:Google Research 实验显示,Titans 在 BABILong 基准(长文档推理)上,2M token 上下文准确率超 GPT-4 同规模基线 15%+,needle-in-haystack 召回率达 95%(vs. Mamba-2 的 78%)。语言建模 perplexity 降低 10-20%,DNA 序列与时间序列任务中,深度 LTM(4 层)优于浅层 2 层 8%。MIRAS 框架统一视角证实,LTM 的非欧目标函数(如 Huber loss)提升鲁棒性,减少 outlier 敏感。

落地参数:LTM 模块 dim=1024-4096(模型规模 1/10),深度 2-6 层(ReLU/SiLU),惊喜阈值 τ=0.05-0.2(梯度范数 ||∇_θ L (θ, k_t v_t)|| > τ 触发更新),学习率 η=1e-4 ~ 5e-4,动量 β=0.9,权重衰减 λ=1e-5 ~ 1e-3(序列长 >1M 时调高遗忘)。SWA 窗口 w=4096-16384(GPU 内存 80GB 限),head num=8-32。变体选择:MAC(检索 LTM 作为额外上下文,query proj dim=512,适合 RAG);MAG(门控融合 SWA 输出 & LTM,gate MLP dim=256,非线性 SiLU);MAL(LTM 前置 SWA,顺序压缩,推理最快)。持久记忆(persistent mem)初始化为任务 embedding,dim=512,置于序列首。

部署清单:

  1. 初始化:LTM θ ~ N (0,0.02),SWA 标准 Transformer layer。
  2. 前向:t 步,计算惊喜 s_t = ||∇_θ L (θ, proj (x_t))||,若 s_t > τ,θ ← θ - η s_t^{mom} + λ θ(mom = β mom + (1-β) s_t)。
  3. 并行训练:序列分块(block=512),小批量 GD 于块内,卷积 approx 时不变系统。
  4. 推理优化:checkpoint LTM 每 10k token,FP16/bfloat16,vLLM 集成 SWA。
  5. 回滚策略:惊喜阈值过高(>0.3)导致遗忘过多,降至 0.1;监控 LTM norm >10 触发重置。

监控要点:(1) 长上下文 perplexity(每 100k token 采样),目标 <2.5;(2) 惊喜分布直方图,均值 0.1-0.5,高尾 <5%;(3) KV 有效压缩率(LTM 占总 mem <20%);(4) Needle 召回 @2M >90%,阈值警报。风险:在线更新不稳(解 ∇=0),限步数 max_update=100 / 序列;容量溢出,动态 λ = base_λ * (len /cap)。

实际参数示例(7B 模型):w=8192,LTM dim=2048 depth=4,η=2e-4,λ=5e-5,τ=0.1。H100 x8 上,2M 推理吞吐 50 t/s(vs. Transformer 5 t/s)。扩展至 10B+,LTM depth+1,dim x1.5。

Titans 标志长时记忆工程化转向:非静态缓存,而是自适应神经索引。未来结合 MoE,RAG 检索精度将跃升。

资料来源:

  • Google Research Blog: Titans + MIRAS (research.google/blog/titans-miras-helping-ai-have-long-term-memory)
  • arXiv: Titans: Learning to Memorize at Test Time (arxiv.org/abs/2501.00663)
查看归档