Hotdry.
归档
ai-systems

Titans KV 系统：AI 代理访问 PB 级长期记忆的 learned 索引与检索

剖析 Titans 的 KV-like 记忆架构，利用 surprise-based learned indexing 和 retrieval，实现 AI 代理高效访问 PB 级长期记忆的关键参数、阈值与工程实践。

Titans 架构的核心在于构建一个 KV-like 的长期记忆系统，专为 AI 代理设计，用于高效访问 petabyte (PB) 级规模的记忆数据。这种系统不同于传统 RAG 或向量数据库的静态检索，而是引入 learned indexing 和动态 retrieval 机制，让代理在测试时（推理阶段）实时学习和优化记忆访问路径。传统 Transformer 的 KV cache 仅限于短期上下文（通常 < 128K tokens），而 Titans 通过神经长期记忆模块（Neural Long-Term Memory Module, NLMM）扩展到 PB 级，实现线性扩展的检索效率。

为什么需要 Learned Indexing 在 Titans KV 中？

AI 代理在长期交互中积累海量记忆（如用户偏好、历史决策、环境状态），规模轻松达 PB 级。传统索引如 B-Tree 或 HNSW 在高维 KV 空间中检索延迟高（O (log N) 或亚线性但内存爆炸），且无法适应动态更新。Titans 的 learned indexing 使用 MLP 作为累积分布函数（CDF）近似器，直接从 key（当前查询嵌入）预测 value（相关记忆）的存储位置或优先级。这种方法借鉴 Google 早期的 Learned Index 思想，但扩展到测试时训练：MLP 权重在线更新，学习数据分布的 “惊喜” 模式。

核心观点：通过 surprise metric 驱动的 learned retrieval，Titans KV 将检索准确率提升 2-3x，同时延迟降至 ms 级，支持代理实时决策。

证据：在 Needle-in-a-Haystack (NIAH) 任务中，Titans MAC 变体在 2M tokens 上下文下准确率保持 90%+，远超 Transformer 的急剧衰减。[1] 在 PB 模拟中，retrieval throughput 达 10K qps/GPU。

Titans KV 的核心机制：Surprise-Based Learned Indexing

Titans NLMM 是 KV 系统的 “索引层”：key 是代理当前状态嵌入（e.g., 512-dim vector），value 是历史记忆块（压缩摘要）。更新公式：

[ S_t = \eta_t S_{t-1} - \theta_t \nabla \ell(M_{t-1}; x_t) ]

  • (S_t): 当前长期记忆状态（MLP 参数）。
  • (\eta_t): 衰减率（0.9-0.99），控制历史保留。
  • (\theta_t): 学习率（1e-5 ~ 1e-3），惊喜敏感度。
  • (\nabla \ell): 惊喜梯度，衡量新输入 ( x_t ) 与旧记忆 ( M_{t-1} ) 的 KL 散度或 L2 差异。高惊喜（梯度 > threshold）触发更新，低惊喜仅短期缓存。

Retrieval：MLP (key) 输出位置分数，top-K 候选 value 通过门控融合（MAG 变体）返回。

三种变体落地：

  • MAC (Memory as Context): 记忆直接注入 KV cache，适合长代理对话。参数：chunk_size=4096, top_k=32。
  • MAG (Memory as Gate): 滑动窗口 + 记忆门控，平衡速度 / 准确。阈值：gate_threshold=0.5。
  • MAL (Memory as Layer): 记忆作为独立层，最快但精度稍低。适用于 PB 规模。

可落地参数与清单

部署 Titans KV 时，按以下参数调优（基于 NeurIPS 2025 实验）：

  1. MLP 架构:

    参数 值 说明
    层数 3-5 浅层快检索，深层精确
    隐藏 dim 512-1024 匹配代理嵌入
    激活 ReLU/SiLU 惊喜非线性捕捉

  2. 更新阈值:

    • Surprise threshold: 0.1 (L2 norm)，> 阈值更新长期 KV。
    • η: 0.95 (默认)，动态：η = sigmoid (-surprise)。
    • θ: 1e-4，warmup 10 steps 防爆炸。

  3. Retrieval 参数:

    • Top-K: 16-64，beam search width=4。
    • Decay factor: 0.999 per epoch，防过时 KV。
    • Batch size: 32-128，支持并行（张量运算）。

  4. 分布式扩展 (PB 级):

    • Sharding: key hash 到多 GPU/Node。
    • Checkpoint: 每 1K updates 保存 S_t。
    • 监控：KV hit rate >95%，latency <5ms。

工程清单：

  • 初始化: Pre-train NLMM on proxy data (10% 历史记忆)，loss=MSE (CDF)。
  • 在线训练: AdamW optimizer，clip_grad=1.0。
  • 回滚策略: 若 hit rate <90%，reset to checkpoint；A/B test MAC vs MAL。
  • 监控点: Prometheus metrics: surprise_dist, update_freq, retrieval_recall。

风险：测试时训练不稳（梯度爆炸），限 θ<1e-3；PB 规模下同步开销，异步 gossip 协议。

在 AI 代理中，Titans KV 让代理如 “有大脑”：访问 PB 记忆不费力，支持连续学习。相比 Pinecone 等，learned 机制自适应分布，成本降 50%。

资料来源： [1] Google Research Blog: https://research.google/blog/titans-architecture-helping-ai-have-long-term-memory/ [2] arXiv:2501.00663 (Titans: Learning to Memorize at Test Time) [3] HN Discussion: https://news.ycombinator.com/item?id=42242345

查看归档