新架构 Titans 引入神经长期记忆:突破 200 万上下文且兼顾速度
研究人员在 arXiv 发布新架构 Titans。它把注意力机制当短期记忆,配合可在测试时学习的神经长期记忆模块,在保持快速并行训练与推理的同时,将上下文扩展至 200 万 token 以上,表现超越传统 Transformer。
让大模型读完一整套长篇巨著还能随时回答细节,一直是业界面临的头疼难题。要么算力成本呈二次方暴涨,要么像循环神经网络那样把所有内容压缩成固定状态,读到后面把前面的细节忘个精光。
研究人员在 arXiv 上发布了一篇重磅论文,提出了名为 Titans 的新模型架构家族。它给大模型换了一种记忆机制,把上下文窗口一举推到了 200 万 token 以上,而且在长文本检索中表现依然扎实。
把注意力当短期记忆,配上测试时学习的长期大脑
过去十年里,模型架构主要在循环模型和注意力机制之间寻找平衡。注意力机制建模极为精准,但计算代价巨大,往往只能固定在有限窗口;传统循环模型虽然快,信息压缩后却很容易丢失关键细节。
Titans 的做法很巧妙,它从人类认知结构中汲取灵感,把注意力机制降级成了「短期记忆」。在 Titans 看来,注意力机制既然擅长精准捕捉依赖,就让它专注处理当前的局部上下文即可。
与此同时,研究团队引入了一个全新的神经长期记忆模块,充当模型的「长期记忆」。这个模块能在测试推理阶段持续学习并记住历史上下文,再回过头来辅助注意力机制提取过去的远期信息。
更难得的是,这个神经记忆系统不仅保持了极快的推理速度,还能像现代模型一样进行高度并行的快速训练。论文围绕如何将两者结合,一共给出了三种 Titans 变体架构。
突破 200 万上下文,在大海里轻松捞针
这套双记忆组合的效果非常直接。实验结果显示,在常规的语言建模、常识推理任务上,Titans 的表现全面超越了传统 Transformer 以及近期涌现的各类现代线性循环模型。
除了自然语言,Titans 在基因组学和时间序列等对超长序列极其敏感的领域,也展现出了更强的建模能力。
在业界检验超长文本能力的「大海捞针」测试中,Titans 成功将上下文窗口扩展到了 200 万 token 以上。在如此庞大的数据中查找隐蔽信息,它的准确率依然明显优于现有的各类基线模型。
为什么重要
对开发者和研究人员来说,Titans 指出了一条很有吸引力的新路径。如果后续能够开源并成熟落地,处理百万行代码库、全套基因组序列或者数千页技术文档的硬件门槛,都有望大幅降低,不再需要为了顾及显存而牺牲模型对关键细节的记忆能力。
从行业竞争来看,过去两年大模型在上下文长度上的比拼,大多依靠稀疏注意力技巧或纯线性状态空间模型。但纯线性模型在大海捞针等高精度检索任务上常常乏力,Transformer 又太消耗算力。Titans 证明了把短期注意力与测试时神经记忆分工协作,完全可以在保留精准度的同时解决计算瓶颈。
我们认为,这项研究最聪明的地方在于重新划定了注意力机制的职责边界,没有死磕单一结构,而是用模块化思路让长短记忆各司其职。不过,在实际生产环境中,测试时动态记忆的状态管理、多轮长会话的容错与安全机制究竟如何,还需要真正的工程落地来进一步验证。
来源 (1)
- arxiv.org(在新标签页打开)原文arxiv.org