“Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention”
通过引入Infini-Transformer，1B大模型上下文长度可扩展到1M（100万token，大约相当于10部小说），并能完成Passkey检索任务。
8B大模型在500K上下文长度的书籍摘要任务中，拿下最新SOTA。
论文地址：https://arxiv.org/pdf/2404.07143.pdf
摘要
本文介绍了Infini-attention，可以将基于Transformer的大型语言模型（LLMs）扩展到无限长的输入，同时控制内存和计算。Infini-attention将压缩内存引入到基本的注意力机制中，并在单个Transformer块中构建了掩码本地注意力和长期线性注意力机制。通过在长上下文语言建模基准测试、1M序列长度的密码块检索和500K长度的书籍摘要任务中，使用1B和8B LLMs展示了该方法的有效性。
简介
Transformers和基于Transformer的LLMs在处理长序列时存在内存和计算复杂度的问题。压缩内存系统可以更有效地处理无限长序列。本文提出了一种新的注意力机制Infini-attention，通过将压缩内存与传统的注意力机制结合起来，实现了对无限长输入的处理。
Infini-attention通过存储和检索旧的注意力键值状态来实现长期记忆的整合和检索。实验证明，使用Infini-attention的模型在长上下文语言建模任务上表现优于基线模型，并且具有更小的内存占用。Infini-attention还支持持续预训练和长上下文适应。Infini-attention使得Transformer LLMs能够以有限的内存和计算资源处理无限长的上下文。
方法
与Transformer-XL类似，Infini-Transformer在一个序列的片段上进行操作。每个片段内部计算标准的自回归点积注意力。与此不同的是，Infini-Transformer在处理下一个片段时不会丢弃前一个片段的注意力状态，而是重新利用它们来保持整个上下文历史。因此，Infini-Transformer的每个注意力层都具有全局压缩和局部细粒度状态。这种高效的注意力机制被称为Infini-attention。
Infini-attention
Infini-attention计算局部和全局上下文状态，并将它们合并为输出。类似于多头注意力（MHA），每个注意力层还维护了H个并行的压缩内存（H是注意力头的数量），除了点积注意力。
缩放点积注意力
多头缩放点积注意力(multi-head scaled dot-product attention， MHA)是语言模型中的主要构建模块之一。它具有建模上下文相关动态计算的强大能力，并且在自回归生成模型中广泛应用。在标准的多头缩放点积注意力中，每个头计算注意力上下文向量，并将它们连接起来，最后将连接向量投影到模型空间得到注意力输出。
压缩内存
在Infini-attention中，我们重用点积注意计算中的Q、K和V。点积注意力和压缩内存之间的状态共享和重用不仅可以实现高效的插件长上下文适应，而且可以加快训练和推理速度。为了简单和计算效率，我们使用关联矩阵参数化内存。这种方法进一步允许我们将内存更新和检索过程转换为线性注意力机制，并利用相关方法中的稳定训练技术。
内存检索。我们使用查询Q从内存M中检索新内容A，如下所示：
我们将所有键的总和记录为归一化项z s−1，并使用元素ELU + 1作为激活函数。
内存更新。一旦检索完成，我们用新的KV条目更新内存和归一化项，并获得下一个状态为：
新的内存状态M s和z s然后传递到下一个片段s + 1，在每个注意力层中建立递归。方程(4)中的右边项σ (K) T V被称为结合约束算子。
受到delta法则成功的启发，我们也将其纳入Infini-attention。增量规则尝试进行稍微改进的内存更新，首先检索现有值项，并在将关联绑定应用为新更新之前从新值中减去它们。
如果KV绑定已经存在于内存中，则该更新规则(线性+ δ)保留不修改的关联矩阵，同时仍然跟踪与前一个(线性)相同的归一化项，以保持数值稳定性。
长期上下文注入。我们将局部注意状态A点和内存检索内容A点通过学习门限标量β聚合在一起：
这只增加了单个标量值作为每个头的训练参数，同时允许模型中长期和局部信息流之间的可学习权衡。
与标准的MHA类似，对于多头无限注意，我们并行计算H个上下文状态，并将它们连接并投射到最终的注意输出
内存和有效上下文窗口
与其他模型相比，Infini-Transformer具有恒定的内存复杂度，并且能够在每个头部的单个层中存储压缩上下文的复杂度为d key × d value + d key。相比之下，其他模型的内存复杂度会随着序列维度的增加而增加。
Transformer-XL通过对上一段的KV状态进行缓存来扩展上下文窗口，其额外的内存占用为(d key + d value) × H × N × l。Compressive Transformer在Transformer-XL的基础上添加了第二个缓存，并存储了过去段落激活的压缩表示。Memorizing Transformers则选择将整个KV状态作为上下文存储，但由于存储成本过高，它们只在单个层中进行上下文计算。通过利用快速kNN检索器，Memorizing Transformers可以在增加存储成本的情况下构建覆盖整个序列历史的上下文窗口。
实验证明，Infini-Transformer LM在Memorizing Transformers的基础上可以实现超过100倍的压缩率，并进一步提高困惑度得分。
实验
长上下文语言建模
我们在PG19和Arxiv-math 基准上训练和评估了小型Infini-Transformer模型。实验表明，Infini-Transformer优于Transformer-xl和Memorizing Transformer，同时与Memorizing Transformer模型相比，使用基于矢量检索的KV存储器，其存储参数减少了114倍。
100K长度训练。我们进一步将训练序列长度从32K增加到100K，并在Arxiv-math数据集上训练模型。100K训练进一步降低了线性和线性+ Delta模型的困惑度得分，分别为2.21和2.20。
门控分数可视化。训练后的无限注意出现了两种类型的头：门控得分接近0或1的专业化头和得分接近0.5的混合型头。专用头通过局部注意力计算处理上下文信息或从压缩内存中检索，而混合头将当前上下文信息和长期记忆内容聚合到一个输出中。我们还观察到在整个前向计算中长期和短期内容检索的交错。
LLM持续预训练
我们对现有LLM的长上下文适应性进行了轻量级的连续预训练。实验中，片段长度N被设置为2K。
1M passkey retrieval基准。我们用Infini-attention替换了1B LLM中的普通MHA，并继续对长度为4K的输入进行预训练。
passkey任务将一个随机数隐藏到一个长文本中，并在模型输出时要求它返回。分散注意力的文本的长度可以通过多次重复文本块来改变。我们仅在5K长度输入上进行测试，以测试1M长度。
对于每个测试子集，我们控制passkey的位置，使其位于输入序列的开始、中间或结束附近。Infini-Transformers在对5K长度输入进行400步微调后，解决了高达1M上下文长度的任务。
500K长度的图书摘要(BookSum)。我们连续预训练一个8B LLM模型，输入长度为8K，步长为30K。然后，我们对图书摘要任务BookSum 进行了微调。我们将输入长度设置为32K进行微调，并将其增加到500K进行评估。
我们的模型优于之前的结果，并通过处理book中的整个文本在BookSum上实现了新的SOTA。有一个明显的趋势表明，随着更多的文本从书籍中输入，我们的Infini-Transformers提高了它的摘要性能指标。
总结
本文将压缩内存模块紧密集成到普通的点积注意力层中。这种对注意力层的微妙但关键的修改使LLM能够用有限的内存和计算资源处理无限长的上下文。实验表明，我们的方法可以自然地扩展到100万长度的输入序列，同时在长上下文语言建模基准和书籍摘要任务上优于基线。我们还证明了我们的方法具有很好的长度泛化能力。1B模型在5K序列长度的passkey实例上进行了微调，解决了1M长度的问题。
评论0