Transformer继任者「Titans」来了，上下文记忆瓶颈被打破

长江后浪推前浪。

原标题：Transformer继任者「Titans」来了，上下文记忆瓶颈被打破
文章来源：Founder Park
内容字数：10921字

谷歌推出Titans架构：Transformer的继任者？

谷歌在Transformer架构发布八年后，推出了全新的Titans架构，旨在解决LLM长上下文处理的难题。该架构的核心在于其创新的长期神经记忆模块，允许模型在测试时学习和记忆信息，从而将上下文窗口扩展到200万tokens。

Titans架构的核心：长期神经记忆模块

1. **挑战与思路:** 传统的注意力机制虽然能精确建模依赖关系，但计算成本高，限制了上下文窗口长度。Titans通过引入长期神经记忆模块来解决这个问题，该模块能够学习如何在测试时记忆或忘记信息，避免过拟合并提升泛化能力。

2. **学习过程与意外指标:** 模型通过在线学习的方式，学习记忆函数。它利用“意外指标”来衡量输入数据与过去数据的偏差，从而决定哪些信息需要记忆。该指标包含过去意外和瞬时意外两个方面，基于一个损失函数进行学习，目标是实现联想记忆（键值对存储）。

3. **遗忘机制:** 为了管理有限的记忆容量，Titans采用了自适应遗忘机制，允许模型忘记不再需要的信息。

4. **记忆架构与检索:** 长期记忆模块采用多层MLP架构，检索信息则通过简单的线性层投影和前向传递实现。

5. **并行化训练:** 为了高效训练，Titans利用了块式梯度下降、张量化以及参数作为块函数的简化方法，充分利用硬件加速器。

Titans架构的三个变体

1. **Memory as a Context (MAC):** 将记忆作为当前信息的上下文，结合注意力机制，选择性地利用历史信息。

2. **Memory as a Gate (MAG):** 使用滑动窗口注意力作为短期记忆，神经记忆模块作为衰减记忆，两者协同工作。

3. **Memory as a Layer (MAL):** 将循环模型和具有注意力机制的深度神经网络进行堆叠。

实验结果与结论

在语言建模、常识推理、基因组学和时序预测等任务中，Titans架构均表现优于Transformer和其他的现代线性循环模型，尤其在“大海捞针”任务中，其在超过200万tokens的上下文窗口下取得了更高的准确率。实验还表明，更深的记忆模块和更大的参数量可以提升模型性能。尽管训练速度略慢于部分基线模型，但Titans (MAL) 的训练速度仍然优于其他一些基线模型和神经记忆模块。

总而言之，Titans 架构提出了一种新颖的将记忆机制融入深度学习架构的方法，为处理超长上下文和提升LLM性能提供了新的思路。

联系作者

文章来源：Founder Park
作者微信：
作者简介：来自极客公园，专注与科技创业者聊「真问题」。