HyperAttention: Long-context Attention in Near-Linear Time

📄 arXiv: 2310.05869v3 📥 PDF

作者: Insu Han, Rajesh Jayaram, Amin Karbasi, Vahab Mirrokni, David P. Woodruff, Amir Zandieh

分类: cs.LG, cs.AI

发布日期: 2023-10-09 (更新: 2023-12-01)

💡 一句话要点

提出HyperAttention以解决长上下文注意力计算效率问题

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

关键词: 长上下文 注意力机制 大型语言模型 近似算法 计算效率 局部敏感哈希 模块化设计

📋 核心要点

  1. 现有方法在处理长上下文时面临计算复杂度高的问题，尤其是在注意力矩阵条目不受限的情况下。
  2. HyperAttention通过引入两个参数来捕捉问题的复杂性，设计了一种线性时间的近似注意力机制。
  3. 实验结果表明，HyperAttention在32k上下文长度下使ChatGLM2的推理时间提高了50%，在131k上下文长度下实现了5倍的速度提升。

📝 摘要（中文）

本文提出了一种名为HyperAttention的近似注意力机制，以应对大型语言模型（LLMs）中长上下文带来的计算挑战。研究表明，在最坏情况下，除非注意力矩阵的条目被限制或矩阵具有低稳定秩，否则需要二次时间复杂度。我们引入两个参数来衡量问题的复杂性，并在这些参数较小的情况下，即使在矩阵条目不受限或具有大稳定秩的情况下，也能实现线性时间的采样算法。HyperAttention的模块化设计便于与其他快速低级实现（如FlashAttention）集成。通过局部敏感哈希（LSH）识别大条目，HyperAttention在多个长上下文数据集上表现优越，显著提高了速度。

🔬 方法详解

问题定义：本文旨在解决大型语言模型中长上下文注意力计算的高复杂度问题。现有方法在处理不受限的注意力矩阵时，通常需要二次时间复杂度，限制了其应用。

核心思路：HyperAttention通过引入最大列范数和行范数比率两个参数，精细化地捕捉问题的复杂性，从而实现线性时间的采样算法，即使在矩阵条目不受限的情况下也能有效处理。

技术框架：HyperAttention的整体架构包括参数设置、局部敏感哈希（LSH）用于识别大条目，以及与FlashAttention等快速实现的模块化集成。

关键创新：HyperAttention的主要创新在于其能够在不限制注意力矩阵条目的情况下，依然实现线性时间复杂度的注意力计算，突破了传统方法的限制。

关键设计：论文中设置的两个参数（最大列范数和行范数比率）是关键设计，能够有效捕捉问题的复杂性。此外，采用LSH技术来识别和处理大条目，进一步提升了计算效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，HyperAttention在32k上下文长度下使ChatGLM2的推理时间提高了50%，而在131k上下文长度下实现了5倍的速度提升。这些结果表明HyperAttention在处理长上下文时的显著性能优势，超越了现有的最先进解决方案如FlashAttention。

🎯 应用场景

HyperAttention的研究成果在自然语言处理、对话系统和其他需要处理长上下文的任务中具有广泛的应用潜力。其高效的注意力机制能够显著提升大型语言模型的推理速度，降低计算资源消耗，推动相关技术的实际应用和发展。

📄 摘要（原文）

We present an approximate attention mechanism named HyperAttention to address the computational challenges posed by the growing complexity of long contexts used in Large Language Models (LLMs). Recent work suggests that in the worst-case scenario, quadratic time is necessary unless the entries of the attention matrix are bounded or the matrix has low stable rank. We introduce two parameters which measure: (1) the max column norm in the normalized attention matrix, and (2) the ratio of row norms in the unnormalized attention matrix after detecting and removing large entries. We use these fine-grained parameters to capture the hardness of the problem. Despite previous lower bounds, we are able to achieve a linear time sampling algorithm even when the matrix has unbounded entries or a large stable rank, provided the above parameters are small. HyperAttention features a modular design that easily accommodates integration of other fast low-level implementations, particularly FlashAttention. Empirically, employing Locality Sensitive Hashing (LSH) to identify large entries, HyperAttention outperforms existing methods, giving significant speed improvements compared to state-of-the-art solutions like FlashAttention. We validate the empirical performance of HyperAttention on a variety of different long-context length datasets. For example, HyperAttention makes the inference time of ChatGLM2 50\% faster on 32k context length while perplexity increases from 5.6 to 6.3. On larger context length, e.g., 131k, with causal masking, HyperAttention offers 5-fold speedup on a single attention layer.