Titans：在测试时间学习记忆

在过去二十年里，关于如何有效利用循环模型和注意力机制的研究 effort 持续进行。虽然循环模型旨在将数据压缩到固定大小的记忆（称为隐藏状态）中，但注意力机制允许对整个上下文窗口进行注意，捕捉所有 token 的直接依赖关系。这种更准确的依赖建模虽然带来了优势，但伴随二次计算成本，限制了模型仅能处理固定长度的上下文。我们提出一种新的神经长期记忆模块，学习记忆历史上下文，帮助注意力关注当前上下文，同时利用长期信息。我们表明，这种神经记忆具有快速可并行训练的优势，同时保持快速推理速度。从记忆角度来看，我们认为注意力因其有限的上下文但准确的依赖建模表现为短期记忆，而神经记忆因其数据记忆能力，充当长期、更持久的记忆。基于这两个模块，我们引入一种新的架构家族，称为 Titans，通过三个变体来解决如何有效地将记忆整合到该架构中的问题。我们的实验结果在语言建模、常识推理、基因组学和时间序列任务上表明，Titans 在 Transformer 和最近的现代线性循环模型中表现更有效。它们进一步可以有效扩展到超过 200 万上下文窗口大小，在 needle-in-haystack 任务中比基线模型具有更高的准确率。