Squeezed Attention: Accelerating Long Context Length LLM Inference

📄 arXiv: 2411.09688v3 📥 PDF

作者: Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Monishwaran Maheswaran, Sebastian Zhao, June Paik, Michael W. Mahoney, Kurt Keutzer, Amir Gholami

分类: cs.CL

发布日期: 2024-11-14 (更新: 2025-06-11)

备注: 22 Pages

期刊: ACL 2025

🔗 代码/项目: GITHUB

💡 一句话要点

Squeezed Attention：通过离线聚类和稀疏注意力加速长文本LLM推理。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长文本推理 注意力机制 模型加速 K-means聚类 稀疏注意力 大型语言模型 上下文压缩

📋 核心要点

  1. 长文本LLM应用面临推理效率挑战，因为推理成本随序列长度线性增长。
  2. Squeezed Attention通过离线聚类固定上下文的键，并在推理时仅关注语义相关的键，从而减少计算量。
  3. 实验表明，该方法在长文本基准测试中显著降低了KV预算，并实现了推理加速，精度损失可忽略不计。

📝 摘要（中文）

本文提出Squeezed Attention，旨在加速长输入上下文的大型语言模型（LLM）应用，尤其针对那些大部分输入上下文在不同用户输入中保持不变的场景。该方法首先离线使用K-means聚类，根据语义相似性对固定上下文的键（key）进行分组，并用单个质心值代表每个簇。在推理时，将用户输入中的查询（query）token与质心进行比较，预测固定上下文中哪些键在语义上相关，然后仅使用这些重要的键计算精确的注意力，从而降低带宽和计算成本。此外，还提出了一种分层版本的算法，可以将注意力复杂度从线性降低到对数级别（相对于固定上下文长度）。在LongBench等长上下文基准测试中，该方法在LLaMA-2-7B-32K、LWM-Text-Chat-1M和Longchat-7B-v1.5-32K模型上实现了3.1倍的KV预算减少，且没有明显的精度损失，甚至在精度仅下降0.5个百分点的情况下，实现了高达8倍的减少。通过实现质心比较和带有重要键的稀疏FlashAttention的内核，在长上下文推理的预填充和生成阶段实现了超过4倍的加速。

🔬 方法详解

问题定义：论文旨在解决长文本LLM推理中，由于输入上下文长度过长导致的计算效率低下问题。现有方法在处理长文本时，需要对所有上下文进行注意力计算，计算量随文本长度线性增长，效率较低。尤其是在许多应用场景中，输入文本的大部分上下文是固定的，对所有上下文都进行计算造成了资源浪费。

核心思路：论文的核心思路是利用固定上下文中键（key）的语义相似性，通过离线聚类的方式对键进行压缩，只保留每个簇的代表性质心。在推理时，只计算查询（query）与这些质心的相似度，从而快速筛选出相关的键，并仅对这些相关的键进行精确的注意力计算。这样可以显著减少需要计算的键的数量，从而降低计算复杂度。

技术框架：Squeezed Attention的整体框架包含离线聚类和在线推理两个阶段。离线聚类阶段，使用K-means算法对固定上下文的键进行聚类，得到每个簇的质心。在线推理阶段，首先计算查询与质心的相似度，选择相似度最高的几个簇，然后只对这些簇中的键进行精确的注意力计算。此外，论文还提出了分层版本的算法，通过构建层次化的聚类结构，进一步降低计算复杂度。

关键创新：最重要的技术创新点在于利用离线聚类来压缩固定上下文的键，从而在推理时只需要计算少量键的注意力。与现有方法相比，Squeezed Attention避免了对所有上下文进行注意力计算，显著降低了计算复杂度，尤其是在固定上下文较长的情况下。分层聚类的引入进一步将复杂度降低到对数级别。

关键设计：K-means聚类的簇的数量是一个关键参数，需要根据实际情况进行调整。论文中使用了FlashAttention的稀疏版本，只对选择出的键进行注意力计算，进一步提高了计算效率。分层聚类中，每一层的簇的数量也需要进行调整，以平衡计算复杂度和精度。

🖼️ 关键图片

img_0

📊 实验亮点

实验结果表明，Squeezed Attention在LongBench等长文本基准测试中，能够在LLaMA-2-7B-32K、LWM-Text-Chat-1M和Longchat-7B-v1.5-32K模型上实现3.1倍的KV预算减少，且没有明显的精度损失。在精度仅下降0.5个百分点的情况下，实现了高达8倍的减少。此外，通过优化内核，在长上下文推理的预填充和生成阶段实现了超过4倍的加速。

🎯 应用场景

Squeezed Attention适用于需要处理长文本输入，且输入文本中包含大量固定上下文的LLM应用，例如文档分析、代码生成、法律文本处理等。该方法可以显著提高这些应用的推理效率，降低计算成本，使得在资源受限的设备上部署长文本LLM成为可能。未来，该方法可以进一步扩展到处理动态变化的上下文，从而应用于更广泛的场景。

📄 摘要（原文）

Emerging Large Language Model (LLM) applications require long input context in order to perform complex tasks like document analysis and code generation. For these long context length applications, the length of the input prompt poses a significant challenge in terms of inference efficiency since the inference costs increase linearly with sequence length. However, for many of these applications, much of the context in the prompt is fixed across different user inputs, thereby providing the opportunity to perform offline optimizations in order to process user inputs quickly, as they are received. We propose Squeezed Attention to accelerate LLM applications where a large portion of the input context is fixed. We first leverage K-means clustering offline to group the keys for the fixed context based on semantic similarity and represent each cluster with a single centroid value. During inference, we compare query tokens from the user input with the centroids to predict which keys from the fixed context are semantically relevant, and then compute exact attention using only the important keys, thereby reducing bandwidth and computational costs. We also present a hierarchical version of our algorithm which can reduce the complexity of attention from linear to logarithmic with respect to the fixed context length. We evaluate our method on long-context benchmarks including LongBench, where it achieves a 3.1$\times$ reduction in KV budget with no noticeable accuracy loss and up to an 8$\times$ reduction with only a 0.5 point accuracy gap for the LLaMA-2-7B-32K, LWM-Text-Chat-1M, and Longchat-7B-v1.5-32K models. Futhermore, we implement kernels for centroid comparison and sparse FlashAttention with important keys, achieving more than 4$\times$ speedups during both the prefill and generation phases for long-context inference. Our code is available at https://github.com/SqueezeAILab/SqueezedAttention.