Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention
Tsendsuren Munkhdalai, Manaal Faruqui and Siddharth Gopal Google
Abstract
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.
中文速览
大型语言模型（LLM）受限于注意力机制的二次方复杂度，处理超长文本时内存和计算代价会急剧膨胀，难以扩展到百万级别的上下文长度。为此，研究者提出了一种名为"Infini-attention"的新型注意力机制，将压缩记忆（compressive memory）嵌入标准注意力层，同时在同一个Transformer块中融合局部掩码注意力和长程线性注意力，并通过可学习的门控标量将两者的输出自适应地加权合并。实验表明，搭载Infini-attention的1B参数模型可自然扩展到100万token的上下文并成功完成密钥检索任务，8B模型在50万token的书籍摘要任务上达到新的最优水平，同时相比Memorizing Transformers实现了逾114倍的内存压缩率。这项工作的意义在于：它以极小的参数改动和有界的内存开销，让现有Transformer LLM通过持续预训练即可"即插即用"地具备无限上下文处理能力，为实际部署超长文本应用提供了高效可行的路径。
原文 arXiv:2404.07143;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2404.07143v2