Google发布下一代Transformer模型Infini-Transformer，高效处理无限上下文
1.360开源智脑7B系列大模型
2.埃隆·马斯克将在5月发布Grok2.0版本
3.谷歌功能最强大的生成式 AI 模型：Gemini 1.5 Pro
4.谷歌推出三个开源工具 MaxDiffusion、JetStream、MaxTest
5.苹果：推出多模态大模型Ferret-UI
6.蚂蚁集团：CodeFuse推出"图生代码"
7.aiXcoder：开源aiXcoder-7B模型
8.Uncharted Labs 推出音乐创作工具 Udio
9.HuggingFace开源模型：Parer-TTS
10.Mistral AI：1760亿参数MoE登开源榜首本研究提出了一种新型的注意力机制
Infini-attention，它允许基于Transformer的大型语言模型（LLMs）高效处理无限长的输入序列，同时保持内存和计算资源的有限性。Infini-attention在传统的注意力机制中加入了压缩记忆，并在单个Transformer模块中集成了隐性局部注意力和长期线性注意力机制。通过在长上下文
语言建模基准测试、1M序列长度的token上下文块检索任务以及500K的书籍摘要任务中使用 1B 和 8B LLMs，证明了机制有效性。这个机制引入了最小有界内存参数，并实现了LLMs的快速流式推理。
关键词：Google、下一代模型Infini-Transformer、创新的注意力机制、Infini-attention、最小化改进、处理极长输入的能力、超越基线的性能、强大的泛化能力
点关注不迷路
一、简介
抛出问题
记忆是AI智能的基石，它能够根据不同的场景进行高效的计算。然而，由于注意力机制的特性，Transformer及其衍生的大型语言模型（LLMs）的上下文依赖记忆受到了限制。
Transformer注意力机制的局限性，特别在处理极长序列时的内存和计算挑战
Transformer中的注意力机制在内存占用和计算时间上呈现出二次方的复杂度。例如，对于一个参数量为500亿的模型，批量大小为512，上下文长度为2048，其注意力的键值（KV）状态的内存占用高达3TB（Pope等人，2023年）。实际上，使用标准的Transformer架构将LLMs扩展到更长的序列（例如100万个标记）是具有挑战性的，而且随着上下文模型越来越长，其经济成本也随之上升。
压缩记忆系统在处理极长序列时，比注意力机制更具可扩展性和效率。与随输入序列长度增长的数组不同，压缩记忆主要通过维持一定数量的参数来存储和回忆信息，从而实现有界的存储和计算成本。在压缩记忆中，新信息通过改变其参数添加到记忆中，目标是之后能够恢复这些信息。然而，目前的LLMs尚未看到一种既有效又实用的压缩记忆技术，该技术能够在保持简单性的同时还能保持高质量的性能。
核心创新点：Infini-attention
- 能够有效地处理无限长的输入数据，同时保持内存占用和计算量在可控范围内。
- 将
压缩记忆整合到传统的注意力机制中，并在单个Transformer模块内同时构建了隐性的局部注意力和长期的线性注意力机制。 - 允许现有的LLMs通过
持续的预训练和微调，自然地扩展至无限长的上下文。
二、实现原理
一种新型的Transformer模型：Infini-Transformer，能够高效地处理极长输入序列，同时保持内存和计算资源的有限性。该模型的核心是一个名为Infini-attention的创新注意力机制。
2.1 Infini-attention
Infini-attention机制：通过将压缩记忆（compressive memory）整合到传统的注意力机制中，实现了对无限长输入的处理。与传统的Transformer注意力层不同，Infini-attention在单个Transformer块中同时构建了掩蔽局部注意力和长期线性注意力机制。不仅能够重用标准注意力计算中的所有键、值和查询状态，以实现长期记忆的整合和检索，而且还能够在处理后续序列时，通过注意力查询状态从记忆中检索值。最终，Infini-attention通过聚合长期记忆检索到的值和局部注意力上下文来计算最终的上下文输出。
2.1.1 缩放点积注意力
缩放点积注意力（scaled dot-product attention）是大型语言模型（LLMs）的主要构建块。它的自注意力（self-attention）变体在自回归生成模型中被广泛利用，因为它能够建模上下文依赖的动态计算，并且具有时间掩蔽的便利性。在标准的多头缩放点积注意力（multi-head scaled dot-product attention）中，每个注意力头通过计算输入序列的注意力查询、键和值状态来生成注意力上下文。然后，通过加权平均所有其他值来计算注意力上下文，其中权重是通过softmax函数计算得到的。
2.1.2 压缩记忆
不同于传统注意力机制的压缩记忆，不会随着输入序列长度的增加而增长。相反，它主要通过维护一定数量的参数来存储和回忆信息，从而实现有界存储和计算成本。在压缩记忆中，新信息的添加是通过改变其参数来实现的，目的是之后能够恢复这些信息。Infini-attention通过重用点积注意力计算中的查询、键和值状态，而不是像传统注意力机制那样丢弃。然后，通过使用注意力查询状态从记忆中检索值，以处理后续序列。此外，Infini-attention还采用了一种与先前工作类似的关联矩阵（associative matrix）来参数化记忆，这使得记忆更新和检索过程可以被视为线性注意力机制，从而利用相关方法中的稳定训练技术。
在Infini-attention机制中，实现
压缩记忆核心原理的三个关键步骤：记忆检索、记忆更新和长期上下文注入。
1.记忆检索（Memory retrieval）
记忆检索是指从压缩记忆中获取信息的过程。在Infini-attention中，这一过程使用当前的注意力查询（Q）状态来从之前的记忆（Ms-1）中检索与当前输入相关的信息。检索到的内容（Amem）通过一个非线性激活函数（σ）和一个归一化项（zs-1）来计算，这两者共同决定了从记忆中检索出的值。
2.记忆更新（Memory update）
记忆更新是指在处理完一个输入序列段后，如何将新的键值对（KV）状态整合到压缩记忆中的过程。在Infini-attention中，更新过程包括将新的键值对添加到现有记忆中，并更新归一化项。这样，新的记忆状态（Ms）和归一化项（zs）就可以传递到下一个序列段中，从而在每个注意力层中建立起递归。
记忆更新的数学表达式为：
3.长期上下文注入（Long-term context injection）
长期上下文注入是指将从压缩记忆中检索到的内容与局部注意力状态结合起来，以形成最终的注意力输出。这一步骤通过一个学习的门控标量（β）来实现，它决定了长期记忆内容和局部注意力内容在最终输出中的权重。
通过这三个步骤，Infini-attention机制能够在处理极长输入序列的同时，保持内存占用的有限性，并有效地利用长期和短期上下文信息。
2.2 记忆与有效上下文窗口
Infini-Transformer模型通过使用压缩记忆，实现了一个无界上下文窗口，同时保持内存占用的有限性。
表1列出了以前的基于段的记忆模型，并根据模型参数和输入段长度定义了它们的上下文记忆占用和有效上下文长度。Infini-Transformer在单层中为每个头存储压缩上下文的Ms和zs，其内存复杂度为dkey × dvalue + dkey，而其他模型的复杂度随着序列维度的增长而增长——无论是Transformer-XL的缓存大小，还是Compressive Transformer和Memorizing Transformers的压缩记忆复杂度，或者是RMT和AutoCompressors的软提示大小。Infini-Transformer通过这种方式，实现了在保持内存占用有限的同时，处理极长输入序列的能力。
三、评估相关
在涉及极长输入序列的基准测试中评估了Infini-Transformer模型，任务包括长上下文语言建模、100万长度的密码上下文块检索以及50万长度的书籍摘要任务。对于语言建模基准测试，从头开始训练模型；而对于密码检索和书籍摘要任务，持续地对现有的大型语言模型（LLMs）进行预训练，以此来突出具备的即插即用式长上下文适应能力。
实验结果显示，Infini-Transformer在长上下文语言建模基准测试中超越了基线模型，并且在内存大小方面比具有65K长度向量检索基线的Memorizing Transformer模型有114倍的压缩比率提升。此外，当训练序列长度增加到100K时，模型的困惑度进一步降低。
Infini-Transformer还展现了在长序列长度任务上的优越性能。在1M长度的密钥检索任务中，通过在5K长度输入上进行微调，模型成功解决了长达1M的上下文长度问题。
在500K长度的书籍摘要任务中，通过持续预训练8B规模的LLM模型，并在书籍摘要任务上进行微调，Infini-Transformer在BookSum数据集上达到了新的SOTA结果，展示了其在处理长文本时的强大性能。
一个具有10亿参数的LLM能够自然地扩展到1M序列长度，并在应用Infini-attention后解决了密钥检索任务。此外，一个具有80亿参数的模型在经过持续的预训练和任务微调之后，在500K长度的书籍摘要任务上取得了新的最先进成果。
四、Infini-attention 相关工作
本节主要探讨了与Infini-attention相关的研究领域，主要包括压缩记忆、长上下文持续预训练以及高效注意力机制。
压缩记忆
压缩记忆是一种受生物神经元可塑性启发的记忆系统，它通过参数化函数来存储和检索信息，与Transformer模型中的键值记忆数组不同，它通过维持固定数量的参数来提升计算效率。这些参数会根据特定的规则更新以保存信息，并通过记忆读取机制来检索。传统Transformer模型压缩输入序列以有效处理长上下文，但通常为了容纳新的记忆条目而丢弃旧的段级记忆，限制了上下文窗口的大小。Infini-attention通过递增式记忆更新，循环利用固定数量的记忆参数，解决了这一问题。
长上下文持续预训练
长上下文持续预训练的目标是让大型语言模型能够处理更长的上下文，包括引入稀疏性到注意力层和操作位置编码。虽然基于位置编码的方法在数据效率上有优势，但推理成本较高。Infini-attention通过分段处理长序列，并使用固定的局部注意力窗口来降低成本。这种方法使得Infini-Transformer能够在训练中处理长达1M的输入序列。
高效注意力
高效注意力技术致力于通过近似计算或系统级优化提高点积注意力的效率。研究者们尝试了多种方法，包括基于稀疏性和线性注意力的近似。线性注意力变体与关联记忆矩阵和元学习神经记忆相关，它们在快速权重中存储键值对，根据新上下文信息进行更新。此外，一些系统级优化技术通过特定硬件架构提高了精确注意力计算的效率。
通过这些相关工作的讨论，展示了Infini-attention机制在理论和实践上的进步，以及它在处理长上下文信息时的潜力和优势。
五、结论
本研究通过在传统的点积注意力层中集成压缩记忆模块，提出了一种新型的注意力机制-Infini-attention，它极大地提升了大型语言模型（LLMs）处理长上下文的能力。这种机制不仅能够有效地建模长距离和短距离的上下文依赖，还能够在有限的内存和计算资源下处理极长的输入序列，实现对无限长上下文的扩展。Infini-attention的设计支持即插即用的持续预训练和长上下文适应，使得LLMs能够自然地扩展到百万长度级别的输入序列，并在长上下文语言建模和书籍摘要任务上取得了优于基线模型的性能。此外，该方法还展现了出色的泛化能力，例如，通过在5K长度的密钥实例上进行微调，1B模型能够解决1M长度的问题。研究的主要贡献在于引入了Infini-attention这一实用且强大的注意力机制，它对标准的缩放点积注意力进行了最小化的改进，使得Transformer LLMs在有限资源下处理长文本成为可能。
引用
- https://arxiv.org/pdf/2404.07143.pdf
关于我
欢迎关注，一起进步一起成长～
点关注不迷路