2025-11-20T08:25:14.880374

Titans: Learning to Memorize at Test Time

Behrouz, Zhong, Mirrokni
Over more than a decade there has been an extensive research effort on how to effectively utilize recurrent models and attention. While recurrent models aim to compress the data into a fixed-size memory (called hidden state), attention allows attending to the entire context window, capturing the direct dependencies of all tokens. This more accurate modeling of dependencies, however, comes with a quadratic cost, limiting the model to a fixed-length context. We present a new neural long-term memory module that learns to memorize historical context and helps attention to attend to the current context while utilizing long past information. We show that this neural memory has the advantage of fast parallelizable training while maintaining a fast inference. From a memory perspective, we argue that attention due to its limited context but accurate dependency modeling performs as a short-term memory, while neural memory due to its ability to memorize the data, acts as a long-term, more persistent, memory. Based on these two modules, we introduce a new family of architectures, called Titans, and present three variants to address how one can effectively incorporate memory into this architecture. Our experimental results on language modeling, common-sense reasoning, genomics, and time series tasks show that Titans are more effective than Transformers and recent modern linear recurrent models. They further can effectively scale to larger than 2M context window size with higher accuracy in needle-in-haystack tasks compared to baselines.
academic

Titans: Learning to Memorize at Test Time

基本信息

  • 论文ID: 2501.00663
  • 标题: Titans: Learning to Memorize at Test Time
  • 作者: Ali Behrouz, Peilin Zhong, Vahab Mirrokni (Google Research)
  • 分类: cs.LG cs.AI cs.CL
  • 发表时间: 2024年12月31日
  • 论文链接: https://arxiv.org/abs/2501.00663

摘要

本文提出了一种新的神经长期记忆模块，能够学习记忆历史上下文并帮助注意力机制在利用长期过去信息的同时关注当前上下文。作者从记忆角度论证，注意力机制由于有限的上下文但准确的依赖建模表现为短期记忆，而神经记忆由于其记忆数据的能力，充当长期、更持久的记忆。基于这两个模块，作者引入了一个新的架构家族Titans，并提出三种变体来有效地将记忆融入架构中。实验结果显示，Titans在语言建模、常识推理、基因组学和时间序列任务上比Transformers和现代线性循环模型更有效，能够有效扩展到超过2M上下文窗口大小。

研究背景与动机

核心问题

现有的序列建模架构面临着效率与性能的权衡问题：

  1. Transformers：虽然能够准确建模依赖关系，但计算复杂度为O(n²)，限制了上下文长度
  2. 线性Transformers/RNNs：虽然效率高，但将信息压缩到固定大小的状态中，导致长序列性能下降
  3. 记忆系统缺失：现有架构缺乏类似人脑的多层次记忆系统（短期记忆、长期记忆、元记忆等）

研究动机

作者受人类记忆系统启发，认为有效的学习范式需要：

  • 不同但相互连接的模块，每个模块负责学习过程中的关键组件
  • 能够主动从数据中学习并记忆过去历史抽象的能力
  • 在测试时仍能持续学习和适应的机制

核心贡献

  1. 神经长期记忆模块：提出了一个深度神经网络作为元模型，在测试时学习如何记忆/存储数据到其参数中
  2. 记忆管理机制：设计了基于"惊讶度"的记忆更新机制和自适应遗忘机制
  3. Titans架构家族：提出三种将记忆融入深度学习架构的方式：Memory as Context (MAC)、Memory as Gate (MAG)、Memory as Layer (MAL)
  4. 并行化训练算法：提供了快速并行化的训练算法，使深度记忆模块训练高效
  5. 广泛实验验证：在多个任务上验证了Titans的有效性，包括语言建模、常识推理、基因组学和时间序列预测

方法详解

任务定义

本文研究序列建模任务，输入为序列 x∈RN×dinx \in \mathbb{R}^{N \times d_{in}}，目标是学习一个能够有效处理长序列的模型，该模型需要：

  • 在测试时能够持续学习和记忆
  • 平衡短期和长期记忆的使用
  • 具有线性复杂度但保持高表达能力

神经长期记忆模块

核心设计思想

受人类长期记忆启发，违反预期（令人惊讶）的事件更容易被记住。作者使用神经网络相对于输入的梯度来衡量"惊讶度"。

记忆更新机制

基本更新规则：

M_t = M_{t-1} - θ_t ∇ℓ(M_{t-1}; x_t)

改进的更新规则（引入动量）：

M_t = M_{t-1} + S_t
S_t = η_t S_{t-1} - θ_t ∇ℓ(M_{t-1}; x_t)

其中：

  • S_t：惊讶度动量，包含过去惊讶度和瞬时惊讶度
  • η_t：数据依赖的惊讶度衰减参数
  • θ_t：控制瞬时惊讶度融入程度的参数

遗忘机制

为处理长序列，引入自适应遗忘机制：

M_t = (1 - α_t)M_{t-1} + S_t

其中 α_t ∈ [0,1] 是门控机制，控制遗忘程度。

目标函数

使用关联记忆损失函数：

ℓ(M_{t-1}; x_t) = ||M_{t-1}(k_t) - v_t||²₂

其中 k_t = x_t W_K, v_t = x_t W_V

Titans架构变体

1. Memory as Context (MAC)

  • 将记忆作为当前信息的上下文
  • 序列分块处理，每个块查询长期记忆获取相关历史信息
  • 注意力机制决定是否需要长期记忆信息

2. Memory as Gate (MAG)

  • 一个分支使用滑动窗口注意力作为短期记忆
  • 另一个分支使用神经记忆模块作为长期记忆
  • 通过门控机制结合两个分支的输出

3. Memory as Layer (MAL)

  • 将神经记忆作为深度网络的一层
  • 顺序处理：先通过记忆层，再通过注意力层
  • 类似现有混合模型的层级设计

并行化训练

通过将训练过程重新表述为使用矩阵乘法和求和操作，实现了高效的并行化训练：

  • 将序列分割成大小为b的块
  • 使用并行关联扫描计算动量项
  • 通过张量化小批量梯度下降实现快速训练

实验设置

数据集

  • 语言建模：FineWeb-Edu数据集，15B/30B tokens
  • 常识推理：PIQA, HellaSwag, WinoGrande, ARC-easy/challenge, SIQA, BoolQ
  • 长上下文任务：RULER benchmark (S-NIAH), BABILong benchmark
  • 时间序列：ETT, ECL, Traffic, Weather数据集
  • 基因组学：GenomicsBenchmarks数据集

模型规模

  • 170M, 340M, 400M, 760M参数的模型
  • 训练长度：4K tokens
  • 上下文窗口：可扩展至2M+ tokens

对比方法

  • Transformers：Transformer++
  • 线性循环模型：RetNet, GLA, Mamba, Mamba2, DeltaNet, TTT, Gated DeltaNet
  • 混合模型：Samba, Gated DeltaNet-H2
  • 大模型：GPT-4, Llama3, RecurrentGemma, Mistral

实验结果

语言建模性能

在340M参数模型上：

  • Titans (LMM)：困惑度26.18 (Wiki), 29.97 (LMB)
  • 最佳基线TTT：困惑度27.44 (Wiki), 34.19 (LMB)
  • **混合模型中Titans (MAG)**表现最佳：困惑度25.07 (Wiki), 28.72 (LMB)

长上下文任务

在S-NIAH任务上（16K序列长度）：

  • Titans (MAC)：S-NIAH-PK 98.4%, S-NIAH-N 97.4%, S-NIAH-W 95.2%
  • Mamba2：S-NIAH-PK 5.4%, S-NIAH-N 0.0%, S-NIAH-W 0.0%
  • TTT：S-NIAH-PK 88.4%, S-NIAH-N 4.4%, S-NIAH-W 0.0%

BABILong基准测试

  • Titans在few-shot设置下超越了所有基线，包括GPT-4
  • 在fine-tuning设置下，小规模Titans超越了大70倍参数的GPT-4

时间序列预测

神经记忆模块在所有数据集上均优于基线方法，包括基于Mamba、Transformer和线性模型的方法。

消融实验

各组件贡献度（按重要性排序）：

  1. 权重衰减（遗忘机制）
  2. 动量机制
  3. 卷积层
  4. 持久记忆
  5. 深度记忆 vs 线性记忆

相关工作

线性循环模型

  • 第一代：RetNet, LRU, RWKV, S4/S5 - 使用数据无关的转移矩阵
  • 第二代：Griffin, Mamba系列 - 引入门控机制
  • 第三代：DeltaNet, TTT, Longhorn - 基于元学习/在线学习的更新规则

Transformer变体

  • 效率优化：稀疏注意力、线性注意力、I/O感知实现
  • 分段Transformer：RMT等使用简单向量记忆在块间传递信息

测试时训练

  • 受早期局部学习算法启发
  • 与MNM和TTT-layer最相关，但Titans具有遗忘机制和动量更新

结论与讨论

主要结论

  1. 记忆系统的重要性：多层次记忆系统（短期+长期+持久记忆）对序列建模至关重要
  2. 测试时学习的有效性：在测试时持续学习和记忆能显著提升长序列处理能力
  3. 架构设计的影响：MAC和MAG架构优于传统的MAL层级设计
  4. 扩展性验证：Titans能有效扩展到2M+上下文窗口

局限性

  1. 计算开销：深度记忆模块比简单矩阵状态需要更多计算资源
  2. 记忆深度权衡：更深的记忆模块效果更好但训练更慢
  3. 参数敏感性：需要仔细调节惊讶度相关参数
  4. 理论分析不足：缺乏对记忆容量和遗忘策略的理论保证

未来方向

  1. 记忆架构优化：探索更高效的神经记忆架构设计
  2. 理论分析：提供记忆容量、遗忘策略的理论分析
  3. 大规模验证：在更大规模模型上验证方法的有效性
  4. 应用扩展：探索在更多领域的应用潜力

深度评价

优点

  1. 概念创新性强：从人类记忆系统角度重新思考序列建模，提出了新颖的多层次记忆架构
  2. 技术贡献全面：不仅提出了神经记忆模块，还设计了三种融入架构的方式和高效的并行化算法
  3. 实验验证充分：在多个领域（NLP、时间序列、基因组学）进行了全面实验，结果令人信服
  4. 理论基础扎实：将记忆更新与梯度下降、动量、权重衰减联系起来，提供了理论解释

不足

  1. 计算复杂度分析不够深入：虽然声称线性复杂度，但深度记忆模块的实际计算开销分析不够详细
  2. 超参数敏感性：多个数据依赖参数(α_t, θ_t, η_t)的设置可能较为复杂
  3. 与人类记忆的类比有限：虽然受人类记忆启发，但"惊讶度"的定义相对简单
  4. 大规模模型验证不足：最大模型仅760M参数，缺乏十亿级参数模型的验证

影响力

  1. 学术价值：为序列建模提供了新的记忆视角，可能启发更多相关研究
  2. 实用价值：在长序列处理任务上表现出色，有实际应用潜力
  3. 可复现性：作者承诺开源代码，有助于方法的推广和验证

适用场景

  1. 长文档处理：适用于需要处理长文档的NLP任务
  2. 时间序列分析：特别适合需要长期历史信息的预测任务
  3. 在线学习场景：适用于需要在测试时持续适应的应用
  4. 记忆密集型任务：如问答系统、对话系统等需要记忆大量信息的任务

参考文献

论文引用了138篇相关文献，涵盖了Transformer、循环神经网络、注意力机制、记忆网络、测试时训练等多个相关领域的重要工作，为本研究提供了坚实的理论基础。