Titans: Learning to Memorize at Test Time

十多年来，关于如何有效利用循环模型和注意力机制的研究一直在广泛进行。循环模型旨在将数据压缩到固定大小的记忆单元（称为隐藏状态），而注意力机制则允许关注整个上下文窗口，捕捉所有标记之间的直接依赖关系。然而，这种更精确的依赖关系建模伴随着二次成本，限制了模型处理固定长度的上下文。我们提出了一种新的神经长期记忆模块，该模块学习记住历史上下文，并帮助注意力机制在利用远期信息的同时关注当前上下文。我们展示了这种神经记忆的优势在于快速并行化训练，同时保持快速推理速度。从记忆的角度来看，我们认为由于注意力机制具有有限的上下文但准确的依赖关系建模，它表现得像短期记忆；而神经记忆由于其能够记住数据，充当更持久的长期记忆。基于这两种模块，我们引入了一个新的架构系列，称为泰坦（Titans），并提出了三种变体，以探讨如何有效地将记忆整合到此架构中。我们在语言建模、常识推理、基因组学和时间序列任务上的实验结果表明，泰坦比Transformer和最近的现代线性循环模型更有效。此外，它们可以在超过200万的时间窗口大小上有效扩展，在“大海捞针”任务中比基线模型具有更高的准确性。

该论文试图解决如何有效结合长期记忆和注意力机制，以克服现有模型在处理长上下文时的局限性。具体来说，它旨在解决注意力机制由于其固定长度上下文窗口带来的二次成本问题，以及循环模型在压缩数据到固定大小隐藏状态时的信息丢失问题。这并不是一个全新的问题，但该研究提出了一个新的方法来改进现有的解决方案。

关键思路是引入一种新的神经长期记忆模块，该模块能够学习存储历史上下文，并帮助注意力机制关注当前上下文的同时利用更久远的信息。与仅能处理有限上下文但依赖关系建模准确的注意力机制（作为短期记忆）相比，这种神经记忆模块充当了长期、持久的记忆。基于这两种模块，研究者提出了一种新的架构系列——Titans，并展示了三种变体，用于有效地将记忆融入架构中。这一思路的新颖之处在于它不仅保留了注意力机制的优点，还通过引入长期记忆解决了其局限性。

论文的亮点包括：1) 提出了神经长期记忆模块，使得模型可以处理超过200万的上下文窗口大小；2) 在语言建模、常识推理、基因组学和时间序列任务上的实验表明，Titans比Transformer和最近的线性递归模型更有效；3) 特别是在针头寻觅草堆（needle-in-haystack）任务中，展示了更高的准确性；4) 研究提供了开源代码，鼓励进一步的研究和发展。值得继续深入研究的方向包括优化神经记忆模块的设计及其与其他类型神经网络的集成。

近期在这个领域中的相关研究还包括：《Longformer: The Long-Document Transformer》, 《Big Bird: Transformers for Longer Sequences》 和 《Performer: Generalized Attention with RFF Kernels for Long Sequence Modeling》。这些研究都致力于解决处理长序列数据的问题，但采用了不同的方法和技术路径。