提出可测试时记忆的神经模块,突破注意力机制长文本处理瓶颈。
Titans: Learning to Memorize at Test Time

- 设计新型神经记忆模块,实现测试时长期上下文记忆
- 在200万词长上下文下完成针堆任务,准确率超越基线
- 适合需要超长依赖建模的复杂序列任务,如基因组分析
十余年来,研究者持续探索如何有效利用循环模型与注意力机制。尽管循环模型将数据压缩为固定大小的隐藏状态(隐状态),而注意力机制能捕捉所有标记间的直接依赖关系,但其二次计算开销限制了上下文长度。本文提出一种新型神经长期记忆模块,学习在测试时记忆历史上下文,帮助注意力关注当前上下文并利用远期信息。该模块支持快速并行训练,同时保持高效推理。从记忆视角看,注意力因上下文有限但依赖建模精确,表现为短期记忆;而神经记忆因具备数据记忆能力,起到长期、更持久的记忆作用。基于此,我们提出新架构家族 Titans,并展示三种变体以有效整合记忆机制。实验表明,Titan 在语言建模、常识推理、基因组学及时间序列任务中优于 Transformer 及近期线性循环模型。其可有效扩展至超过 200 万词长上下文窗口,在针堆任务中精度更高。
原文摘要 · Abstract (English)
Over more than a decade there has been an extensive research effort on how to effectively utilize recurrent models and attention. While recurrent models aim to compress the data into a fixed-size memory (called hidden state), attention allows attending to the entire context window, capturing the direct dependencies of all tokens. This more accurate modeling of dependencies, however, comes with a quadratic cost, limiting the model to a fixed-length context. We present a new neural long-term memory module that learns to memorize historical context and helps attention to attend to the current context while utilizing long past information. We show that this neural memory has the advantage of fast parallelizable training while maintaining a fast inference. From a memory perspective, we argue that attention due to its limited context but accurate dependency modeling performs as a short-term memory, while neural memory due to its ability to memorize the data, acts as a long-term, more persistent, memory. Based on these two modules, we introduce a new family of architectures, called Titans, and present three variants to address how one can effectively incorporate memory into this architecture. Our experimental results on language modeling, common-sense reasoning, genomics, and time series tasks show that Titans are more effective than Transformers and recent modern linear recurrent models. They further can effectively scale to larger than 2M context window size with higher accuracy in needle-in-haystack tasks compared to baselines.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。