谷歌推出基于Transformer的创新技术——无限注意力

添加书签

专注AIGC领域的专业社区,关注微软&OpenAI、百度文心一言、讯飞星火等大语言模型(LLM)的发展和应用落地,聚焦LLM的市场研究和AIGC开发者生态,欢迎关注!

随着ChatGPT、Sora、Midjourney、Suno等生成式AI产品的火爆出圈,Transformer几乎成为文生图、视频、音频、文本等大模型的标配基础架构。

但Transformer也有明显的缺点,就是一次性处理超长序列数据时会遭遇内存、计算的瓶颈,成本和复杂度将呈指数级增长并导致大模型的性能下降。

所以,谷歌基于Transformer提出了一种创新技术Infini-Attention(无限注意力),通过将压缩记忆、局部遮蔽注意力等模块融合到传统的自注意力机制中,可轻松处理无限序列数据,同时保持有限的内存和计算资源

论文地址:https://arxiv.org/abs/2404.07143

压缩记忆

压缩记忆是Infini-Attention的核心模块之一,允许大模型以固定数量的参数存储和回忆信息,从而实现对长文本的有效处理。

这与传统的注意力机制不同,压缩记忆不会随着输入序列长度的增加而增长,而是通过改变其参数来添加新信息,从而提升推理、计算效率并降低内存使用。

在处理新的输入序列时,大模型会压缩记忆中的key和value对。该过程通过一个简单的关联绑定操作来完成,将新的key和value对与记忆中现有的信息结合起来。这种更新机制确保了即使在处理极长的序列时,大模型也能以较低的内存来完成并保持稳定。

此外,为了提升大模型的检索效率,在处理每个新的输入序列时,Infini-attention会使用当前的查询从压缩记忆中检索相关的信息。

这个过程主要通过一个线性注意力机制来实现,将查询与记忆中的键进行匹配并返回相应的值。这种检索机制使得大模型能够在处理长文本时,有效地利用之前序列中的信息。

记忆更新

记忆更新主要负责在处理新的输入序列时更新压缩记忆模块中的key和value对,既能保持大模型的性能,又能减少内存占用和节省计算资源。

记忆更新是基于神经科学非常知名的概念“Hebbian学习原则”。即当一个神经元A足够接近地参与到另一个神经元B的激活时,会发生某种过程或一系列的过程,这使得A的效率作为B的一个激活源之一得到加强。

换句话说,如果两个神经元同时激活,那么它们之间的连接会变得更加强大。这种机制支持了学习和记忆的过程,使得神经网络能够根据经验调整其内部连接。

而在Infini-Attention中,新的输入信息通过注意力查询与现有的记忆键进行匹配,并将匹配的结果直接更新到记忆中。

在每次更新过程中,记忆参数(通常是关联矩阵)会根据新的输入信息进行调整,通过计算新的key和value对与当前记忆状态的差值,然后将这个差值应用到关联矩阵上来实现。

策略方面,记忆更新模块采用的是“增量式更新策略”,每次只更新与当前输入相关的部分记忆。这种策略不仅提高了更新过程的效率,还减少了因大规模参数更新带来的不稳定性。

局部遮蔽注意力

Transformer的自注意力机制,允许大模型在序列中的任何位置查找信息,这在处理短文本时效果很好。但当输入序列变得非常长时,会导致计算成本呈指数级增长。

Infini-Attention的遮蔽机制模块主要用于限制大模型的注意力范围,使其只能关注当前处理的局部上下文,而不是无限制的全序列检索白白浪费算力

对于每个输入序列,模型会创建一个遮蔽矩阵,该矩阵与序列的维度相同。遮蔽矩阵中的元素通常是一个二值矩阵,其中1表示允许注意力聚焦的标记,而0表示需要遮蔽的标记。

在计算自注意力分数时,大模型会将遮蔽矩阵与注意力分数矩阵相乘。这样,被遮蔽的标记对应的注意力分数会被设置为一个很小的值。

通过遮蔽注意力分数,大模型可以计算加权和,即每个标记的上下文表示。这个加权和只包括未被遮蔽的标记,确保模型的输出集中在当前的局部上下文上,从而更好地捕捉和理解序列的连贯性。这对于生成长文本和处理复杂的语言结构非常重要。

开发人员将Infini-Attention在长上下文语言建模、密钥检索任务和书籍摘要任务中进行了综合测试。

结果显示,Infini-Attention模型在这些基准测试中的表现优于Transformer-XL和Memorizing Transformers等基础模型。尤其是在使用100K序列长度进行训练时,Infini-Attention的困惑度得分大幅度降低。

本文素材来源Infini-Attention论文,如有侵权请联系删除

END

本篇文章来源于微信公众号: AIGC开放社区