像人一样学会记忆，谷歌新架构Titans打破Transformer上下文限制

作者
现有的架构，如Hopfield网络、LSTM和Transformer会带来了二次方的时间和内存复杂度，限制了模型处理长序列的能力。
人类的记忆不是单一的过程，而是由短期记忆、工作记忆和长期记忆等不同系统组成，每个系统都有不同的功能和神经结构。
基于此，Google提出一种新型的神经网络架构——Titans，它通过引入长期记忆模块来增强模型在处理长序列数据时的能力，并且能够扩展到超过200万的上下文窗口大小，超越GPT-4、Llama3.1-70B等模型。

受人类记忆中“违反预期的事件更难忘”这一观点的启发，定义了一个模型的惊讶度为其对输入的梯度。梯度越大，输入数据与过去数据的差异越大。提出了一个改进的惊讶度量，将过去惊讶度和瞬间惊讶度结合起来，以更好地处理有限记忆。

基于惊讶度量，提出了一个记忆更新规则，结合了过去惊讶度和瞬间惊讶度，通过这种方式，模型可以更好地管理有限的记忆资源。

遗忘机制

  • 适应性遗忘：为了管理大量序列数据中的过去信息，提出了一种适应性遗忘机制，允许模型在需要时忘记不再需要的信息。

  • 遗忘门控：遗忘机制通过一个门控机制实现，该机制可以灵活地控制记忆的更新，决定多少信息应该被遗忘。这种遗忘机制与现代循环模型中的遗忘门控机制有关。

记忆架构

  • 简单MLP：选择了简单的多层感知器（MLP）作为长期记忆的架构，希望专注于长期记忆的设计以及如何将其整合到架构中。这种选择也为未来设计更有效的记忆架构提供了新的研究方向。

如何整合记忆

提出了Titans架构，包括三种变体：Memory as a Context (MAC)、Memory as a Gate (MAG) 和 Memory as a Layer (MAL)。这些变体展示了如何将长期记忆模块有效地整合到深度学习架构中。

  • MAC架构：将记忆作为当前信息的上下文，通过注意力机制决定是否需要长期记忆信息，并帮助记忆存储有用的信息。

  • MAG架构：使用滑动窗口注意力作为短期记忆，长期记忆模块作为模型的渐忘记忆，通过门控机制结合两者。

  • MAL架构：将神经记忆模块作为深度神经网络的一层，负责压缩过去和当前上下文，然后通过注意力模块进行处理。

在语言建模、常识推理、基因组学和时间序列任务上评估了Titans架构及其变体的性能。结果表明，Titans在这些任务上比现代循环模型和线性Transformer更有效，并且能够扩展到超过200万的上下文窗口大小，同时在haystack 任务中比基线模型具有更高的准确性，比如：GPT4和Llama3.1 70B等。

https://arxiv.org/pdf/2501.00663Titans: Learning to Memorize at Test Time

（文：PaperAgent）

欢迎分享

发表评论