证明了Transformer计算无法比独立计算每头更快,揭示其理论极限。
On the Computational Hardness of Transformers
- 从复杂性角度分析多头多层Transformer的计算下界。
- 小嵌入时需至少 $LHN^{2+o(1)}$ 时间,大嵌入时需 $LHN^{ω−o(1)}$ 次运算。
- 首次建立Transformer的非平凡计算下界,适合理论学习者与算法研究者。
Transformer已在语言、视觉等领域彻底革新现代AI。其由 $L$ 层构成,每层并行运行 $H$ 个注意力头,每个头处理 $N$ 个长度为 $m$ 的向量。注意力机制涉及三个 $N \times m$ 矩阵相乘,并对中间结果应用Softmax。现有算法独立计算每个注意力头。我们提出核心问题:能否比 $LH$ 次独立计算更高效?本文在负面上回答此问题,在小嵌入情形($m = N^{o(1)}$)下证明所需时间 $LHN^{2+o(1)}$ 已近乎最优(基于SETH)。在大嵌入情形($m = N$)下,独立计算需 $LHN^{ω+o(1)}$ 次算术运算(含指数),我们证明该代价不可再降,即至少需要 $LHN^{ω−o(1)}$ 次运算(当 $ω > 2$ 时)。该下界依赖于对Baur-Strassen定理的新颖应用,这一工具是反向传播算法的核心基础。
原文摘要 · Abstract (English)
The transformer has revolutionized modern AI across language, vision, and beyond. It consists of $L$ layers, each running $H$ attention heads in parallel and feeding the combined output to the subsequent layer. In attention, the input consists of $N$ tokens, each a vector of dimension $m$. The attention mechanism involves multiplying three $N \times m$ matrices, applying softmax to an intermediate product. Several recent works have advanced our understanding of the complexity of attention. Known algorithms for transformers compute each attention head independently. This raises a fundamental question that has recurred throughout TCS under the guise of ``direct sum'' problems: can multiple instances of the same problem be solved more efficiently than solving each instance separately? Many answers to this question, both positive and negative, have arisen in fields spanning communication complexity and algorithm design. Thus, we ask whether transformers can be computed more efficiently than $LH$ independent evaluations of attention. In this paper, we resolve this question in the negative, and give the first non-trivial computational lower bounds for multi-head multi-layer transformers. In the small embedding regime ($m = N^{o(1)}$), computing $LH$ attention heads separately takes $LHN^{2 + o(1)}$ time. We establish that this is essentially optimal under SETH. In the large embedding regime ($m = N$), one can compute $LH$ attention heads separately using $LHN^{ω+ o(1)}$ arithmetic operations (plus exponents), where $ω$ is the matrix multiplication exponent. We establish that this is optimal, by showing that $LHN^{ω- o(1)}$ arithmetic operations are necessary when $ω> 2$. Our lower bound in the large embedding regime relies on a novel application of the Baur-Strassen theorem, a powerful algorithmic tool underpinning the famous backpropagation algorithm.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。