【论文阅读】谷歌Infini-Transformer，一种长上下文的注意力机制

bylander

已于 2024-04-16 22:29:12 修改

阅读量1.6k

点赞数 19

分类专栏： AI论文阅读文章标签：论文阅读 chatgpt 人工智能 transformer

于 2024-04-16 22:07:42 首次发布

本文链接：https://blog.csdn.net/bylander/article/details/137842245

版权

AI论文阅读专栏收录该内容

54 篇文章 2 订阅

订阅专栏

对于LLM，长上下文功能非常重要。最近谷歌的一篇论文提出了一种新的Infini-attention注意力机制，将长上下文的注意力KV状态压缩到记忆中，有效解决LLM在处理长序列时的内存消耗大和计算成本高的问题。
在这里插入图片描述

论文方法
通过下图可以看的比较直接，图的右侧是常规的多头注意力；左侧的处理，通过当前状态的Q值，去检索历史的KV状态，然后计算Amem（具体计算后续描述）；然后聚合，再经过线性投影计算。
在这里插入图片描述
从图中可以看到两个主要问题，一个是历史的KV如何更新，一个是如何根据当前Q值检索历史KV状态。

文章中提到的Memory retrieval（记忆检索）是Infini-attention机制中的一个关键步骤，它涉及到从压缩记忆中检索与当前输入序列相关的信息。具体的方法如下：

使用查询向量（Query Vectors）：在Infini-attention中，每个注意力头（attention head）都会生成一个查询向量（Q），这个查询向量用于与压缩记忆中存储的信息进行匹配。
压缩记忆（Compressive Memory）：压缩记忆由两个主要部分组成，即关联矩阵（associative matrix）和归一化项（normalization term）。关联矩阵存储了键（K）和值（V）状态之间的绑定关系，而归一化项则用于调整检索过程，确保检索的稳定性。
检索过程：给定当前的查询向量Q，Memory retrieval的目标是从压缩记忆中检索出与之相关的值向量（Amem）。这个过程通过下面的公式表示：

在这里插入图片描述
其中，σ(Q)是一个非线性激活函数，Ms-1 是上一次迭代中的关联矩阵，zs-1是归一化项，σ(Q)分别与关联矩阵和归一化项进行逐元素的点积。

非线性激活和归一化：为了确保训练的稳定性，选择适当的非线性激活函数和归一化方法非常关键。文章中提到，使用了元素级的ELU（Exponential Linear Unit）加1作为激活函数，并记录所有key值的和作为归一化项zs-1。
记忆更新（Memory Update）：在检索完成后，需要更新压缩记忆以包含新的键值对。更新过程可以通过下面的公式表示：

公式里分别是更新后的关联矩阵和归一化项。
增量更新（Delta Rule）：为了进一步提升记忆更新的效率，文章还引入了增量更新规则（Delta Rule）。这种方法首先检索现有的记忆内容，然后从新值中减去这些内容，再应用关联绑定操作作为新的更新。这有助于保持关联矩阵的稳定性，同时仍然能够跟踪相同的归一化项。

在图1中，通过门控机制，将常规的注意力计算和内容检索的计算进行聚合，就是下面的公式：

最后的投影和常规的MHA类似，将多头Infini-attention进行级联后通过投影矩阵。

方法对比
图2展示了Infini-Transformer 与Transformer-XL的区别。

对于Transformer-XL，除了当前状态外，Transformer-XL 还会计算从上一段缓存的 KV 状态的注意力，来扩展上下文窗口。如果针对每一层完成扩展了上下文窗口，就会具有很大的上下文记忆复杂性。更进一步，Memorizing Transformers 选择将整个 KV 状态存储为输入序列的上下文。由于在这种情况下存储变得非常昂贵，因此将上下文计算限制为仅单层。实验表明，Infini-Transformer可以在 Memorizing Transformer 的基础上实现超过 100 倍的压缩率，同时进一步提高困惑度分数。

对于Infini-Transformer，通过压缩记忆来维护整个上下文历史。Infini-Transformer能够保留整个上下文历史，而不仅仅是最近一个段的信息。这使得Infini-Transformer能够处理更长的序列，同时保持较低的内存占用。图2展示了Infini-Transformer如何实现一个无界（unbounded）的有效上下文窗口，同时保持有界的内存占用。
在这里插入图片描述
图3说明，训练后，Infini-attention中出现了两种类型的头部：门控得分接近0或1的专业头部和得分接近0.5的混合头部。专用头或者通过局部注意力机制处理上下文信息，或者从压缩存储器中检索，而混合头将当前上下文信息和长期存储器内容一起聚合成单个输出。

借助Kimi，对文章进行总览
在这里插入图片描述
这篇文章的标题是《Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention》，作者们提出了一种新的方法来扩展基于Transformer的大型语言模型（LLMs），使其能够高效处理无限长的输入序列，同时保持内存和计算资源的有界性。以下是对文章的主要内容进行的分析：

解决的问题：

长序列处理的内存和计算限制：传统的Transformer模型在处理长序列时面临着内存消耗大和计算成本高的问题。随着输入序列长度的增加，模型的内存占用和计算时间呈二次方增长。
长序列建模的效率问题：现有的长序列处理方法在扩展上下文窗口和维持计算效率之间难以平衡，导致在实际应用中受到限制。

主要方法：

Infini-attention技术：文章提出了一种名为Infini-attention的新型注意力机制，该机制通过将压缩记忆（compressive memory）集成到传统的注意力机制中，同时在一个Transformer块中构建了掩蔽局部注意力和长期线性注意力机制。
压缩记忆：Infini-attention通过重用标准注意力计算中的键（key）、值（value）和查询（query）状态，存储旧的KV状态到压缩记忆中，而不是像标准注意力机制那样丢弃它们。在处理后续序列时，通过使用注意力查询状态从记忆中检索值。
长短期信息融合：Infini-attention通过学习门控标量β来聚合局部注意力状态和从记忆中检索的内容，实现了长短期信息的融合。

贡献：

长短期上下文依赖的高效建模：提出了Infini-attention，这是一种实用且强大的注意力机制，能够有效地模拟长短期上下文依赖关系。
对标准的注意力计算的最小改变：Infini-attention对标准的缩放点积注意力进行了最小化的改动，并通过设计支持即插即用的持续预训练和长上下文适应。
有界资源下的无限上下文扩展：该方法使得Transformer LLMs能够通过流式处理极长输入来扩展到无限长的上下文，同时保持有界的内存和计算资源。

文章通过在长上下文语言建模基准、1M序列长度的密钥检索任务和500K长度的书籍摘要任务上的实验，证明了其方法的有效性。特别是在内存压缩比方面，与现有技术相比，Infini-Transformer在保持性能的同时，实现了114倍的压缩率。此外，文章还展示了1B和8B规模的LLMs在经过Infini-attention改进后，在长序列任务上的性能提升。

有效的记忆系统对于使用大型语言模型（LLMs）理解长上下文至关重要，它不仅有助于理解长文本，还对推理、规划、持续适应新知识，甚至学习如何学习都有着重要作用。本文介绍了将压缩记忆模块紧密集成到传统的点积注意力层中的方法。这种微妙但关键的修改使得LLMs能够处理无限长的上下文，同时保持内存和计算资源的有界性。我们展示了我们的方法可以自然地扩展到输入序列的百万长度级别，同时在长上下文语言建模基准和书籍摘要任务上超越了基线模型。我们还证明了我们的方法具有很好的长度泛化能力。在5K序列长度的密钥实例上进行微调的1B模型解决了1M长度的问题。总的来说，本文的研究为长序列建模提供了一种新的有效方法，不仅提高了模型处理长文本的能力，还为未来的研究工作提供了新的方向，特别是在如何提高大型语言模型处理长序列数据的效率和准确性方面。

(更多论文阅读，请关注微信公众号：老无谈藏书票)