Simple linear attention language models balance the recall-throughput tradeoff

📄 arXiv: 2402.18668v2 📥 PDF

作者: Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, Christopher Ré

分类: cs.CL, cs.LG

发布日期: 2024-02-28 (更新: 2025-03-07)

🔗 代码/项目: GITHUB

💡 一句话要点

提出BASED架构以平衡语言模型的召回与吞吐量

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 语言模型 注意力机制 召回能力 内存效率 自然语言处理 生成模型 深度学习

📋 核心要点

  1. 现有基于注意力的语言模型在推理时面临KV缓存的高内存消耗，影响了效率。
  2. 本文提出BASED架构，结合线性和滑动窗口注意力，通过调整参数实现召回与内存消耗的平衡。
  3. BASED在真实任务中比现有模型提高了6.22的准确率，并在生成速度上实现了24倍的提升。

📝 摘要（中文）

近期研究表明，基于注意力的语言模型在召回能力上表现优异，但在推理过程中KV缓存的内存消耗限制了其效率。本文探讨如何在不妥协召回能力的情况下，提高语言模型的效率。通过实验和理论分析，我们发现模型状态大小与召回能力之间存在关键权衡。我们提出了一种简单的架构BASED，结合线性和滑动窗口注意力，通过调整窗口大小和线性注意力特征维度，能够在召回与内存消耗之间找到最佳平衡。实验表明，BASED在真实世界的召回密集任务中超越了现有模型，并在生成速度上实现了显著提升。

🔬 方法详解

问题定义：本文旨在解决基于注意力的语言模型在推理过程中由于KV缓存导致的高内存消耗问题。现有方法在召回能力与效率之间存在显著的权衡，影响了实际应用。

核心思路：论文提出BASED架构，通过结合线性注意力和滑动窗口注意力，允许在不同的窗口大小和特征维度之间调整，从而优化模型的状态大小与召回能力的平衡。

技术框架：BASED架构主要包括线性注意力模块和滑动窗口注意力模块，用户可以根据需求调整窗口大小和特征维度，以实现最佳的召回与内存消耗平衡。

关键创新：BASED架构的创新在于其能够在保持召回能力的同时，显著降低内存消耗，突破了传统注意力机制的限制。与现有模型相比，BASED在性能与效率上实现了双重提升。

关键设计：在BASED中，窗口大小和线性注意力特征维度是关键参数，通过实验确定最佳设置。此外，开发了IO感知算法，使得在生成1024个token时，BASED的吞吐量比FlashAttention-2高出24倍。

📊 实验亮点

BASED在真实世界的召回密集任务中超越了现有的最强子二次模型（如Mamba），准确率提高了6.22个百分点。此外，在生成1024个token时，BASED的吞吐量比FlashAttention-2高出24倍，显示出显著的性能优势。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和文本生成等。通过提高语言模型的效率和召回能力，BASED可以在实时应用中提供更好的用户体验，推动智能助手和自动化内容生成的进步。

📄 摘要（原文）

Recent work has shown that attention-based language models excel at recall, the ability to ground generations in tokens previously seen in context. However, the efficiency of attention-based models is bottle-necked during inference by the KV-cache's aggressive memory consumption. In this work, we explore whether we can improve language model efficiency (e.g. by reducing memory consumption) without compromising on recall. By applying experiments and theory to a broad set of architectures, we identify a key tradeoff between a model's state size and recall ability. We show that efficient alternatives to attention (e.g. H3, Mamba, RWKV) maintain a fixed-size recurrent state, but struggle at recall. We propose BASED a simple architecture combining linear and sliding window attention. By varying BASED window size and linear attention feature dimension, we can dial the state size and traverse the pareto frontier of the recall-memory tradeoff curve, recovering the full quality of attention on one end and the small state size of attention-alternatives on the other. We train language models up to 1.3b parameters and show that BASED matches the strongest sub-quadratic models (e.g. Mamba) in perplexity and outperforms them on real-world recall-intensive tasks by 6.22 accuracy points. Implementations of linear attention are often less efficient than optimized standard attention implementations. To make BASED competitive, we develop IO-aware algorithms that enable 24x higher throughput on language generation than FlashAttention-2, when generating 1024 tokens using 1.3b parameter models. Code for this work is provided at: https://github.com/HazyResearch/based.