提出无需低温假设的快速注意力算法,突破传统复杂度瓶颈。
Subquadratic Algorithms and Hardness for Attention with Any Temperature
- 针对常数头维度设计亚二次时间算法,依赖输入值大小的多对数因子
- 在低秩情况下实现与训练全程兼容的渐近加速,且运行时间接近理论下界
- 证明现有方法难以进一步优化,揭示注意力计算的根本极限
尽管Transformer架构广受欢迎,但标准注意力计算在上下文长度n上仍存在二次时间复杂度。Alman和Song(NeurIPS 2023)表明,在头维度d=Θ(log n)时,仅当输入绝对值有界于B=o(√log n)时,注意力才可在强指数时间假设(SETH)下实现亚二次计算;等价于高温度软最大化情形。此类算法运行时间随B呈指数增长,无法推广至大输入范围。本文解决此核心问题:在任意温度下能否高效计算注意力?我们首次为常数头维度d=O(1)构建了亚二次时间算法,复杂度为Õ(n^{2-1/d} · polylog(B)),即使输入值较大也有效。若输入矩阵低秩,该算法同样适用,并可用于梯度计算及全模型训练。此外,我们证明任何显著改进均极不可能:当d=2^{Θ(log* n)}时,注意力仍需n^{2-o(1)}时间。最后,在头维度为poly(n)时,标准算法在细粒度复杂性假设下已是最优。
原文摘要 · Abstract (English)
Despite the popularity of the Transformer architecture, the standard algorithm for computing Attention suffers from quadratic time complexity in context length $n$. Alman and Song [NeurIPS 2023] showed that when the head dimension $d = Θ(\log n)$, subquadratic Attention is possible if and only if the inputs have small entries bounded by $B = o(\sqrt{\log n})$ in absolute values, under the Strong Exponential Time Hypothesis ($\mathsf{SETH}$). Equivalently, subquadratic Attention is possible if and only if the softmax is applied with high temperature for $d=Θ(\log n)$. Running times of these algorithms depend exponentially on $B$ and thus they do not lead to even a polynomial-time algorithm outside the specific range of $B$. This naturally leads to the question: when can Attention be computed efficiently without strong assumptions on temperature? Are there fast attention algorithms that scale polylogarithmically with entry size $B$? In this work, we resolve this question and characterize when fast Attention for arbitrary temperatures is possible. First, for all constant $d = O(1)$, we give the first subquadratic $\tilde{O}(n^{2 - 1/d} \cdot \mathrm{polylog}(B))$ time algorithm for Attention with large $B$. Our result holds even for matrices with large head dimension if they have low rank. In this regime, we also give a similar running time for Attention gradient computation, and therefore for the full LLM training process. Furthermore, we show that any substantial improvement on our algorithm is unlikely. In particular, we show that even when $d = 2^{Θ(\log^* n)}$, Attention requires $n^{2 - o(1)}$ time under $\mathsf{SETH}$. Finally, in the regime where $d = \mathrm{poly}(n)$, we show that the standard algorithm is optimal under popular fine-grained complexity assumptions.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。