ATLAS：学习在测试时最优地记忆上下文
计算与语言
2025-05-30 v1 人工智能
摘要
Transformer 已成为序列建模中最流行的骨干网络，主要归功于其在上下文检索任务中的有效性以及大规模学习的能力。然而，其二次方的内存和时间复杂度限制了其在更长序列中的适用性，促使研究人员探索有效的替代架构，如现代循环神经网络（又称长期循环记忆模块）。尽管它们近期在多样的下游任务中取得了成功，但在需要长上下文理解和外推到更长序列的任务中仍显不足。我们观察到这些缺陷源于其设计中三个独立的方面：(1) 受限于记忆架构和输入特征映射的有限记忆容量；(2) 更新的在线特性，即仅针对最后一个输入优化记忆；以及 (3) 对其固定大小记忆的表达能力较弱的管理。为了增强这三个方面，我们提出了 ATLAS，一种具有高容量的长期记忆模块，它通过基于当前和过去的 token 优化记忆来学习记忆上下文，克服了长期记忆模型的在线特性。基于这一见解，我们提出了一种类似 Transformer 的新架构家族，称为 DeepTransformers，它们是原始 Transformer 架构的严格泛化。我们在语言建模、常识推理、召回密集型和长上下文理解任务上的实验结果表明，ATLAS 超越了 Transformer 和近期的线性循环模型的性能。ATLAS 进一步提升了 Titans 的长上下文性能，在 BABILong 基准的 10M 上下文长度下实现了 +80% 的准确率。
引用
@article{arxiv.2505.23735,
title = {ATLAS: Learning to Optimally Memorize the Context at Test Time},
author = {Ali Behrouz and Zeman Li and Praneeth Kacham and Majid Daliri and Yuan Deng and Peilin Zhong and Meisam Razaviyayn and Vahab Mirrokni},
journal= {arXiv preprint arXiv:2505.23735},
year = {2025}
}