谷歌推出基于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