⏶23
ATLAS：学习在测试时最优地记忆上下文
发表
由 
Ali Behrouz 提交
Ali Behrouz 提交作者: Ali Behrouz, Zeman Li, Praneeth Kacham, Majid Daliri, Yuan Deng, Peilin Zhong, Meisam Razaviyayn, Vahab Mirrokni
摘要
Transformer模型已成为序列建模中最流行的骨干网络，这主要得益于它们在上下文检索任务中的有效性和规模化学习的能力。然而，它们的二次方的内存和时间复杂度限制了其在更长序列中的适用性，因此促使研究人员探索有效的替代架构，如现代循环神经网络（又称长期循环记忆模块）。尽管它们最近在各种下游任务中取得了成功，但在需要长上下文理解和外推到更长序列的任务中，它们却难以胜任。我们观察到这些不足源于其设计的三个独立方面：（1）内存容量有限，受限于内存架构和输入特征映射；（2）在线更新的特性，即仅针对最后一个输入优化内存；以及（3）对固定大小内存的管理表达能力较弱。为了增强所有这三个方面，我们提出了ATLAS，一个高容量的长期记忆模块，它通过基于当前和过去标记优化内存来学习记忆上下文，克服了长期记忆模型的在线特性。基于这一见解，我们提出了一类新的类Transformer架构，称为DeepTransformers，它们是原始Transformer架构的严格泛化。我们在语言建模、常识推理、需要大量回忆的任务和长上下文理解任务上的实验结果表明，ATLAS超越了Transformer模型和近期线性循环模型的性能。ATLAS进一步提高了Titans的长上下文性能，在BABILong基准的10M上下文长度上实现了+80%的准确率提升。
Transformer已成为序列建模中最流行的骨干模型，主要归功于其在上下文内检索任务中的有效性以及大规模学习的能力。然而，其二次方的内存和时间复杂度限制了它们在长序列中的应用，因此促使研究人员探索有效的替代架构，例如现代循环神经网络（也称为长时循环记忆模块）。尽管它们在各种下游任务中取得了成功，但在需要理解长上下文和外推到更长序列的任务中却表现不佳。我们观察到这些缺点源于其设计的三个不同方面：（1）受限于记忆架构和输入特征映射的有限记忆容量；（2）在线更新的性质，即仅针对最后一个输入优化记忆；以及（3）对固定大小记忆的管理表达能力较弱。为了增强所有这三个方面，我们提出了ATLAS，一个具有高容量的长时记忆模块，它通过基于当前和过去的token优化记忆来学习记忆上下文，克服了长时记忆模型的在线性质。基于这一见解，我们提出了一系列新的类Transformer架构，称为DeepTransformers，它们是原始Transformer架构的严格泛化。我们在语言建模、常识推理、密集回忆和长上下文理解任务上的实验结果表明，ATLAS超越了Transformer和近期线性循环模型的性能。ATLAS进一步提升了Titans的长上下文性能，在BABILong基准测试的10M上下文长度上实现了+80%的准确率提升。