提出ATLAS模块,让模型在测试时动态记忆上下文,提升长序列理解能力。
ATLAS: Learning to Optimally Memorize the Context at Test Time
- 设计可优化历史与当前输入的高容量记忆模块,突破传统在线更新限制。
- 在1000万上下文长度上,比Transformer高出80%准确率,显著超越现有模型。
- 适合需要长程依赖建模的任务,如超长文本推理、复杂问答等场景。
Transformers虽因上下文检索能力强而成为序列建模主流,但其二次方时间与内存复杂度限制了长序列应用。尽管现代循环神经网络(即长时记忆模块)在诸多下游任务中表现优异,但在需要长上下文理解与外推的任务中仍显不足。我们发现其局限源于三点:(1)记忆容量受限于架构和输入特征映射;(2)更新为在线模式,仅基于最新输入优化记忆;(3)对固定大小记忆的表达力弱。为此,我们提出ATLAS——一种高容量长时记忆模块,能基于当前及过往标记优化记忆,克服在线学习缺陷。在此基础上,我们构建一类严格推广原始Transformer的新架构,称为DeepTransformers。在语言建模、常识推理、回忆密集型任务和长上下文理解任务上的实验表明,ATLAS性能优于Transformer及近期线性递归模型。在BABILong基准上,其1000万上下文长度下的准确率提升达80%。
原文摘要 · Abstract (English)
Transformers have been established as the most popular backbones in sequence modeling, mainly due to their effectiveness in in-context retrieval tasks and the ability to learn at scale. Their quadratic memory and time complexity, however, bound their applicability in longer sequences and so has motivated researchers to explore effective alternative architectures such as modern recurrent neural networks (a.k.a long-term recurrent memory module). Despite their recent success in diverse downstream tasks, they struggle in tasks that requires long context understanding and extrapolation to longer sequences. We observe that these shortcomings come from three disjoint aspects in their design: (1) limited memory capacity that is bounded by the architecture of memory and feature mapping of the input; (2) online nature of update, i.e., optimizing the memory only with respect to the last input; and (3) less expressive management of their fixed-size memory. To enhance all these three aspects, we present ATLAS, a long-term memory module with high capacity that learns to memorize the context by optimizing the memory based on the current and past tokens, overcoming the online nature of long-term memory models. Building on this insight, we present a new family of Transformer-like architectures, called DeepTransformers, that are strict generalizations of the original Transformer architecture. Our experimental results on language modeling, common-sense reasoning, recall-intensive, and long-context understanding tasks show that ATLAS surpasses the performance of Transformers and recent linear recurrent models. ATLAS further improves the long context performance of Titans, achieving +80\% accuracy in 10M context length of BABILong benchmark.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。