Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention

📄 arXiv: 2404.07143v2 📥 PDF

作者: Tsendsuren Munkhdalai, Manaal Faruqui, Siddharth Gopal

分类: cs.CL, cs.AI, cs.LG, cs.NE

发布日期: 2024-04-10 (更新: 2024-08-09)

备注: 9 pages, 4 figures, 4 tables (v2 adds: background, implementation details, recent citations and acknowledgments)

💡 一句话要点

提出Infini-attention以解决长输入Transformer模型的效率问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长输入处理 Transformer模型 Infini-attention 语言模型 高效推理

📋 核心要点

  1. 现有的Transformer模型在处理长输入时面临内存和计算效率的挑战，限制了其应用范围。
  2. 本文提出的Infini-attention技术通过压缩记忆和结合多种注意力机制，解决了长输入处理的效率问题。
  3. 实验结果表明，该方法在长上下文语言建模和摘要任务中显著提升了性能，支持快速流式推理。

📝 摘要（中文）

本文提出了一种高效的方法，通过引入Infini-attention技术，将基于Transformer的大型语言模型（LLMs）扩展到无限长输入，同时保持内存和计算的界限。Infini-attention在传统注意力机制中融入了压缩记忆，并在单个Transformer块中结合了掩蔽局部注意力和长期线性注意力机制。我们在长上下文语言建模基准、1M序列长度的密码上下文块检索和500K长度的书籍摘要任务上验证了该方法的有效性，使用了1B和8B的LLMs。该方法引入了最小的有界内存参数，并实现了快速流式推理。

🔬 方法详解

问题定义：本文旨在解决现有Transformer模型在处理无限长输入时的内存和计算效率问题。传统注意力机制在长输入时会消耗大量资源，导致性能下降。

核心思路：提出的Infini-attention技术通过引入压缩记忆，结合掩蔽局部注意力和长期线性注意力机制，优化了注意力计算，提升了处理效率。

技术框架：整体架构包括一个单一的Transformer块，内部集成了Infini-attention机制，能够在保持计算效率的同时处理长输入。该框架通过动态调整注意力范围，确保了信息的有效传递。

关键创新：Infini-attention是本文的核心创新，它与传统注意力机制的本质区别在于引入了压缩记忆，允许模型在有限的内存中处理无限长的输入。

关键设计：在设计中，采用了最小的有界内存参数设置，并优化了损失函数和网络结构，以支持快速流式推理和高效的长上下文处理。具体参数设置和网络结构细节在实验部分进行了详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，使用Infini-attention的模型在长上下文语言建模基准上表现优异，能够处理1M序列长度的输入，并在500K长度的书籍摘要任务中显著提升了性能。与传统模型相比，推理速度得到了显著提高，且内存消耗保持在可接受范围内。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理中的长文本理解、实时信息检索和大规模文本生成等。通过提高长输入的处理效率，Infini-attention技术能够推动大型语言模型在实际应用中的广泛使用，尤其是在需要快速响应的场景中。

📄 摘要（原文）

This work introduces an efficient method to scale Transformer-based Large Language Models (LLMs) to infinitely long inputs with bounded memory and computation. A key component in our proposed approach is a new attention technique dubbed Infini-attention. The Infini-attention incorporates a compressive memory into the vanilla attention mechanism and builds in both masked local attention and long-term linear attention mechanisms in a single Transformer block. We demonstrate the effectiveness of our approach on long-context language modeling benchmarks, 1M sequence length passkey context block retrieval and 500K length book summarization tasks with 1B and 8B LLMs. Our approach introduces minimal bounded memory parameters and enables fast streaming inference for LLMs.