谷歌研究人员开发了一种新的神经网络架构——Titans,旨在解决大型语言模型(LLMs)在处理长序列时面临的一个重大挑战:如何在不大幅增加内存和计算成本的情况下扩展其推理时的记忆能力。
Titans架构结合了传统LLMs中的注意力模块与“神经记忆”层,使得模型能够高效地处理短期和长期记忆任务。据研究人员介绍,采用神经长期记忆的LLMs能够处理数百万个标记,且性能优于经典LLMs和如Mamba等其他替代方案,同时所需参数更少。
传统LLMs通常基于Transformer架构,使用自注意力机制来计算标记之间的关系,这是一种学习标记序列中复杂和细致模式的有效技术。然而,随着序列长度的增加,计算和存储注意力的成本呈二次方增长。
为解决这一问题,近期提出的一些替代架构具有线性复杂度,能够扩展而不导致内存和计算成本激增。但谷歌研究人员指出,线性模型在性能上无法与经典Transformer相媲美,因为它们会压缩上下文数据,容易遗漏重要细节。
研究人员认为,理想的架构应具备不同的记忆组件,这些组件能够协同工作,利用现有知识、记忆新事实并从上下文中学习抽象概念。他们提出,有效的学习范式应类似于人类大脑,包含不同但相互连接的模块,每个模块负责学习过程中的关键部分。
为了填补当前语言模型的空白,研究人员设计了一个“神经长期记忆”模块,该模块能够在推理时学习新信息,而无需使用完整的注意力机制带来的低效性。神经记忆模块学习一个函数,该函数能够在推理过程中记忆新事实,并根据遇到的数据动态调整记忆过程,从而解决了其他神经网络架构存在的泛化问题。
为了决定哪些信息值得存储,神经记忆模块采用“惊讶度”概念。标记序列与模型权重和现有记忆中存储的信息差异越大,就越令人惊讶,因此越值得记忆。这使得模块能够高效利用其有限内存,仅存储对模型已知知识有用的数据片段。
为了处理非常长的数据序列,神经记忆模块具有自适应遗忘机制,能够移除不再需要的信息,从而管理内存的有限容量。
Titans架构是一个模型系列,将现有的Transformer模块与神经记忆模块相结合。该模型包含三个关键组件:作为短期记忆的“核心”模块,使用经典注意力机制关注模型正在处理的输入标记的当前段落;使用神经记忆架构存储当前上下文之外信息的“长期记忆”模块;以及训练后保持固定的可学习参数“持久记忆”模块,用于存储与时间无关的知识。
研究人员提出了连接这三个组件的不同方法。总体而言,该架构的主要优势在于使注意力和记忆模块能够相互补充。例如,注意力层可以使用历史和当前上下文来确定当前上下文窗口中哪些部分应存储在长期记忆中。同时,长期记忆提供当前注意力上下文中不存在的历史知识。
研究人员对Titan模型进行了小规模测试,参数范围从1.7亿到7.6亿,测试任务包括语言建模和长序列语言任务。他们将Titans的性能与各种基于Transformer的模型、如Mamba等线性模型以及如Samba等混合模型进行了比较。结果显示,Titan在语言建模方面的表现优于其他模型,并且与大小相似的Transformer和线性模型相比,性能更佳。
在涉及长序列的任务中,如“大海捞针”(模型必须从非常长的序列中检索信息)和BABILong(模型必须在非常长的文档中跨事实推理)等,Titan的性能优势尤为明显。在这些任务中,Titan甚至超越了参数数量级更多的模型,包括GPT-4、GPT-4o-mini以及增强了检索增强生成(RAG)功能的Llama-3模型。
此外,研究人员能够将Titans的上下文窗口扩展到200万个标记,同时保持内存成本处于较低水平。尽管仍需对更大规模的模型进行测试,但论文中的结果表明,Titans的潜力尚未达到上限。
对于企业应用而言,随着谷歌在长上下文模型方面的领先地位,这种技术有望被应用于私有和开放模型,如Gemini和Gemma。LLMs支持更长的上下文窗口,这为创建能够通过提示融入新知识的应用提供了更多可能性,而无需使用如RAG等复杂技术。基于提示的应用的开发和迭代周期远快于复杂的RAG管道。同时,如Titans等架构有助于降低长序列的推理成本,使企业能够在更多用例中部署LLM应用。谷歌计划发布用于训练和评估Titans模型的PyTorch和JAX代码。