还在卷长文本?谷歌最新论文直接把文本干到…无限长了
作者|苗正
邮箱|miaozheng@pingwest.com
当人们还在比拼上下文窗口的时候,谷歌发布了这样一篇论文《Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention》。论文中写到,团队发明了一种新的注意力技术,叫做“无限注意力”(Infini-attention),通过这项技术,能使transformer大模型在有限的计算资源条件下,处理无限长度的输入。
在一个transformer模型中,注意力的作用是允许模型根据当前位置的输入元素(如词元或token)来分配权重给序列中其他所有位置的元素。而上下文窗口则限制了注意力机制的实际操作范围,即模型在计算注意力时仅考虑当前元素周围特定范围内(前后若干位置)的其他元素。
无限注意力允许允许模型在处理无限长输入序列的时候,仍能保持对上下文信息的访问能力,再也不需要在处理新输入时丢弃前一输入段的注意力状态了。那么也就是说,它的上下文窗口可以是……无限。
无限注意力机制背后的关键技术叫做压缩记忆系统,这是一种能够以紧凑形式存储和检索大量信息的结构,通过改变自身参数来捕获新信息,确保信息可以在之后被恢复。单从运行的逻辑上来讲,压缩记忆系统和咱们日常生活里压缩文件是一模一样的。
压缩记忆系统最大的作用是克服transformer标准注意力机制在处理长序列时存在的内存足迹和计算时间的二次复杂度问题,只需要通过使用固定数量的参数存储和召回信息,确保存储和计算成本保持在可控范围内。因为其参数数量不随输入序列的增长而变化,那也就是说,无论输入序列长度有多长,也不会对影响模型的复杂度。
接下来无限注意力机制会将输入序列划分为一系列娇小的、连续的子序列,每个段具有固定的长度,使得模型在处理这些较短的段时,能够保持较低的内存需求和计算复杂度。这种分段方法避免了一次性加载和处理整个无限长序列的挑战,允许模型以流式(streaming)方式逐步处理输入,即每次仅处理一个或几个段,而非一次性加载全部数据。
在每个分段内部,无限注意力模型采用局部注意力机制来处理该段内的上下文信息。局部注意力限制了模型对当前段内token之间的注意力计算范围,通常采用因果(causal)或自回归(autoregressive)的形式,确保模型在处理当前token时,只能看到该令牌之前的所有token,而不能看到未来(即当前token之后)的任何token。
在输出结果时,无限注意力模型为了生成最终的上下文输出,要从压缩记忆中检索到的长期记忆信息与当前局部注意力计算出的上下文结合起来。这种融合确保模型既考虑了当前输入段的局部依赖,又充分利用了历史输入的长期上下文。
当你理解了无限注意力机制后再回到标题,无限注意力模型能够以流式方式处理极端长的输入序列,无需一次性加载整个无限长的输入,而是会根据历史记录进行分批次处理。那对于模型来说,就能够在有限的内存和计算资源约束下,适应并处理无限长度的上下文。
论文首先在长上下文语言建模基准上评估了无限注意力模型的表现,与包括transformer-XL在内的多种模型进行了对比。
采用无限注意力的模型在PG19(长文档数据集)以Arxiv-math(数学数据集)上都取得了远超于transformer-XL的结果,同时实现了114倍的内存压缩率,在保持低困惑度的同时提高了模型效率。
为了进一步验证无限注意力机制的性能,论文将一个10亿参数的大语言模型进行改造,把这个模型的多头注意力(MHA)模块换成了无限注意力,并继续对其进行预训练。验证过程是,团队要求模型在长达100万tokens的输入中定位并检索隐藏的密钥信息。
预训练阶段,模型使用的输入序列长度仅为4K个tokens,以适应无限注意力的处理模式。经过3万步的预训练后,对密钥检索任务进行微调。在微调阶段,为了模拟实际应用中可能遇到的更长上下文环境,模型在包含5K个token的长度输入上进行微调。
点个“在看”,再走吧👀