通过稀疏注意力加速长文本处理，速度提升显著。
HSR-Enhanced Sparse Attention Acceleration
- 利用半空间报告结构识别注意力矩阵中的关键非零项。
- 生成解码时复杂度降至O(mn^4/5)，比原始方法快得多。
- 适合需要高效处理长上下文的LLM应用，如文档摘要。
大语言模型在各类应用中表现出色，但在长上下文任务上的性能常受注意力机制计算复杂度限制。本文提出一种新方法，加速大语言模型在长上下文场景下的注意力计算。该方法利用注意力机制中固有的稀疏性，涵盖传统Softmax注意力与带有$ℓℓ^α$激活的ReLU注意力（$α∈ ℕ_+$）。通过引入半空间报告（HSR）数据结构，高效识别注意力矩阵中的非零或“大幅激活”项。本文对生成解码和提示预填充两种关键场景进行了理论分析。在生成解码中，运行时间从朴素方法的$O(mn)$降低至$O(mn^{4/5})$，其中$n$为上下文长度，$m$为查询长度，$d$为隐藏维度。在提示预填充中，时间复杂度从$O(mn)$降至$O(mn^{1 - 1 / ⌊ d/2⌋} + mn^{4/5})$。该方法对Softmax注意力引入的误差可忽略不计。本工作为实现大语言模型高效长上下文处理迈出了重要一步。
原文摘要 · Abstract (English)
Large Language Models (LLMs) have demonstrated remarkable capabilities across various applications, but their performance on long-context tasks is often limited by the computational complexity of attention mechanisms. We introduce a novel approach to accelerate attention computation in LLMs, particularly for long-context scenarios. We leverage the inherent sparsity within attention mechanisms, both in conventional Softmax attention and ReLU attention (with $\mathsf{ReLU}^α$ activation, $α\in \mathbb{N}_+$), to significantly reduce the running time complexity. Our method employs a Half-Space Reporting (HSR) data structure to identify non-zero or ``massively activated'' entries in the attention matrix. We present theoretical analyses for two key scenarios: generation decoding and prompt prefilling. Our approach achieves a running time of $O(mn^{4/5})$ significantly faster than the naive approach $O(mn)$ for generation decoding, where $n$ is the context length, $m$ is the query length, and $d$ is the hidden dimension. We can also reduce the running time for prompt prefilling from $O(mn)$ to $O(mn^{1 - 1 / \lfloor d/2\rfloor} + mn^{4/5})$. Our method introduces only provably negligible error for Softmax attention. This work represents a significant step towards enabling efficient long-context processing in LLMs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。