On the Computational Hardness of Transformers

本文在 SETH 假设下证明了，无论是小嵌入维度还是大嵌入维度，多头多层 Transformer 的计算复杂度均无法显著优于将 $LH$ 个注意力头独立计算，从而确立了 Transformer 计算复杂度的首个非平凡下界。

要点

这篇论文探讨了一个非常核心的问题：Transformer（目前最火的 AI 模型架构）到底能不能算得更快？还是说，我们现在的“笨办法”其实已经是极限了？

为了让你轻松理解，我们可以把 Transformer 想象成一个超级庞大的“信息处理工厂”。

1. 工厂的运作模式：多头注意力机制

想象这个工厂有 $L$ 层（比如 12 层或 96 层），每一层都有 $H$ 个并行的“专家小组”（也就是论文里说的“多头”）。

• 输入：工厂接收一串信息（比如一句话，有 $N$ 个词）。
• 任务：每个“专家小组”都要去分析这 $N$ 个词之间的关系。比如，在句子“猫吃了鱼”里，第一个专家要分析“猫”和“吃”的关系，第二个专家分析“吃”和“鱼”的关系，以此类推。
• 计算量：每个专家小组都要把 $N$ 个词两两配对进行比对。如果词的数量是 $N$，那么两两比对就需要 $N^2$ 次操作。
• 现状：因为工厂有 $L$ 层，每层有 $H$ 个小组，所以总的工作量大约是 $L \times H \times N^2$。这就是目前大家公认的“笨办法”（Naive Algorithm），也就是论文里说的“分别计算每个头”。

2. 核心疑问：能不能“批量处理”来偷懒？

这就引出了论文要回答的终极问题：

能不能像“批发”一样，一次性把 $L \times H$ 个小组的任务合并起来算，从而比“一个一个单独算”快得多？

这就好比你有 100 个工人要分别做 100 道不同的数学题。

• 传统做法：让每个工人独立做题，做完为止。
• 偷懒想法：能不能发明一种“超级流水线”，让这 100 个工人在同一时间、用一种巧妙的方法，瞬间算出所有答案？

在计算机科学的其他领域，确实有些任务可以这样“批量偷懒”（比如多项式求值）。所以，大家自然希望 Transformer 也能这样。

3. 论文的结论：想多了，不行！

这篇论文给出了一个令人失望但非常确定的答案：不行。你无法通过“批量处理”来显著加速 Transformer。

作者证明了：

• 小模型情况（嵌入维度较小）：如果你试图用比“一个一个算”更快的方法，在数学上几乎是不可能的（除非某些极其复杂的数学猜想是错的）。现在的“笨办法”已经是最优解了。
• 大模型情况（嵌入维度很大）：即使你用了最厉害的矩阵乘法技巧，想要算得比“分别计算”快，也是不可能的。

通俗比喻：
想象你要把 100 个不同的快递（每个头）送到 100 个不同的地址（输入词）。

• 以前的想法：也许我们可以发明一种“超级快递车”，一次能同时送 100 个快递，比 100 辆小车分别送要快。
• 这篇论文的发现：经过严密的数学推导，他们证明了这种“超级快递车”根本不存在。无论你怎么设计路线，100 辆小车分别跑（分别计算）已经是物理极限了。任何试图合并计算的尝试，最终都会发现工作量并没有减少，甚至可能因为合并的开销而更慢。

4. 他们是怎么证明的？（两个关键工具）

为了证明“无法偷懒”，作者用了两把“数学手术刀”：

1. 小模型时的“正交向量”测试：
   他们把 Transformer 的计算任务伪装成了一个著名的数学难题（3-OV 问题，类似于在一大堆数字里找三个互不相关的数）。他们证明了：如果你能更快地算出 Transformer，那你就能瞬间解开这个数学难题。但大家都相信这个数学难题是极难解的（需要 $N^2$ 的时间）。所以，Transformer 也跑不快。

2. 大模型时的“反向传播”魔法（Baur-Strassen 定理）：
   这是论文最精彩的部分。他们利用了一个叫 Baur-Strassen 的定理（这也是神经网络训练时“反向传播”算法的数学基础）。

   • 比喻：假设你有一个黑盒子（Transformer），输入是矩阵，输出是结果。作者说，如果你能很快算出这个黑盒子的结果，那么利用这个定理，你就能很快算出黑盒子里所有的中间变化率（导数）。
   • 结果：他们发现，如果能快速算 Transformer，就能快速算出 $L \times H$ 个独立的矩阵乘法。但数学界早就证明了，算这么多独立的矩阵乘法，必须花那么多时间，没法偷懒。所以，Transformer 也没法偷懒。

5. 这对我们意味着什么？

• 对 AI 开发者：别指望通过算法层面的“魔法”来让 Transformer 的速度产生质的飞跃（比如从 $N^2$ 变成 $N^{1.5}$）。现在的瓶颈是硬性的数学限制。
• 未来的方向：既然算法上很难突破，未来的加速可能更多依赖于：
  • 硬件优化：比如更好的芯片、更快的内存（像 FlashAttention 那样优化数据读取）。
  • 近似算法：如果允许一点点误差，也许可以算快一点（但这会牺牲精度）。
  • 架构创新：彻底换一种不是 Transformer 的架构。

总结一句话：
这篇论文给所有想给 Transformer“提速”的人泼了一盆冷水，但也指明了方向：在数学本质上，Transformer 的“分别计算”已经是最高效的极限了，想再快，只能靠换硬件或换架构，靠算法“走捷径”是行不通的。

技术摘要

这篇论文题为《On the Computational Hardness of Transformers》（关于 Transformer 的计算难度），由加州大学圣地亚哥分校（UCSD）和哥伦比亚大学的研究人员共同完成。文章主要探讨了 Transformer 架构在计算复杂性理论层面的下界问题，特别是针对多头（Multi-head）多层（Multi-layer）Transformer 的计算效率。

以下是该论文的详细技术总结：

1. 研究背景与核心问题

背景：
Transformer 架构通过自注意力机制（Self-Attention）彻底改变了自然语言处理和计算机视觉等领域。一个标准的 Transformer 包含 $L$ 层，每层有 $H$ 个并行的注意力头。每个注意力头处理长度为 $N$ 的输入，嵌入维度为 $m$。
目前的通用算法是独立计算每个注意力头，其时间复杂度为 $O(LHN^2m)$（小嵌入维度下）或 $O(LHN^{\omega+o(1)})$（大嵌入维度下，$\omega$ 为矩阵乘法指数）。

核心问题：
在理论计算机科学中，存在一个经典的“直接求和”（Direct Sum）问题：同时解决多个相同问题的实例，是否能比单独解决每个实例更高效？
对于 Transformer 而言，这意味着：是否存在一种算法，能够比独立计算 $L \times H$ 个注意力头更高效地计算整个 Transformer？之前的研究主要集中在单个注意力头的下界（如基于 SETH 假设的 $N^{2-o(1)}$ 下界），但针对多层多头的 Transformer 整体下界尚不明确。

2. 主要贡献

论文给出了 Transformer 计算的首个非平凡下界，证明了在大多数情况下，朴素算法（即独立计算每个头）在理论上几乎是最优的。

主要定理与结果：

1. 小嵌入维度情形 ($m = N^{o(1)}$)：

   • 结果： 在 $m = \Omega(\log N)$ 且 $L, H = \text{poly}(N)$ 的情况下，计算 Transformer 需要 $\Omega(LHN^{2-o(1)})$ 时间。
   • 假设： 基于 3-OV 假设（3 正交向量假设）或 SETH（强指数时间假设）。
   • 意义： 这一结果将下界从之前仅针对单个头的 $N^{2-o(1)}$ 提升到了与朴素算法 $LHN^2$ 匹配的规模，证明了无法通过“批量处理”显著加速。

2. 大嵌入维度情形 ($m = N$)：

   • 结果： 在扩展算术电路（Extended Arithmetic Circuits, eAC）模型下，计算 $L$ 层 $H$ 头、嵌入维度为 $N$ 的 Transformer 需要 $\Omega(LHN^{\omega-o(1)})$ 的大小（Size）。
   • 假设： 无需额外假设，仅基于矩阵乘法的下界（当 $\omega > 2$ 时）。
   • 意义： 证明了即使利用快速矩阵乘法，Transformer 的计算难度本质上等同于计算 $\Theta(LH)$ 个独立的矩阵乘法实例。

3. 方法论与技术细节

论文针对两种不同的嵌入维度采用了不同的归约策略和证明工具：

A. 小嵌入维度 ($m = N^{o(1)}$)：基于 3-OV 的归约

• 构造思路： 作者构造了一个特定的 Transformer $T_C$，其输入包含两个向量集 $A, B$ 和一个向量集 $C$（大小分别为 $N, N, LH$）。
• 机制设计：
  • 利用 $L$ 层和 $H$ 个头，将 $C$ 中的每个向量 $c_{h,\ell}$ 映射到第 $\ell$ 层第 $h$ 个注意力头。
  • 通过精心设计的查询（Query）、键（Key）和值（Value）映射，使得注意力机制能够检测 $A, B, C$ 中是否存在正交三元组（即 $a_i \cdot b_j \cdot c_{h,\ell} = 0$）。
  • 利用 Hardmax 注意力（可被 Softmax 近似）的特性：如果存在正交三元组，注意力输出的特定位置值会发生显著变化（从 2 变为 $\le 1.5$）。
• 结论： 如果存在计算 Transformer 的超快算法，则可以在 $o(LHN^2)$ 时间内解决 3-OV 问题，这与 3-OV 假设矛盾。

B. 大嵌入维度 ($m = N$)：基于 Baur-Strassen 定理的扩展

• 挑战： 直接归约矩阵乘法到 Transformer 很困难，因为 Transformer 的输出是多个注意力头的和，而矩阵乘法的和（Sum of Products）有时可以通过比单独计算更快的算法完成（例如 $N$ 个矩阵乘积的和可以在 $O(N^{3.251})$ 时间内完成，而单独计算需 $O(N^{3.372})$）。因此，简单的归约无法证明下界。
• 核心工具： Baur-Strassen 定理。该定理指出，如果一个算术电路计算函数 $f$，那么计算 $f$ 的所有偏导数所需的电路大小与计算 $f$ 本身的大小同阶（$O(s)$）。
• 扩展模型： 由于注意力机制涉及指数运算（Softmax），作者引入了扩展算术电路（eAC），允许包含 $\exp$ 和 $\ln$ 门。
• 证明策略：
  1. 构造一个 Transformer，其输出包含 $LH$ 个独立矩阵乘积 $A_k B_k^\top$ 的指数和（通过 Denormalized Attention 实现）。
  2. 引入辅助输入变量 $C_{kij}$ 和 $D_{ki}$。
  3. 利用扩展的 Baur-Strassen 定理，从计算 Transformer 的电路中“提取”出所有偏导数。
  4. 通过计算特定偏导数（$\frac{\partial f}{\partial C_{kij}}$ 和 $\frac{\partial g}{\partial D_{ki}}$），可以恢复出 $A_k B_k^\top$ 的每一个元素。
  5. 由于恢复 $LH$ 个独立矩阵乘积需要 $\Omega(LHN^{\omega-o(1)})$ 的大小，因此计算 Transformer 的电路也必须至少这么大。
• 关键引理： 证明了对于低次函数（如二次函数），扩展算术电路（含 $\exp, \ln$）并不比标准算术电路更强，从而可以将 eAC 的下界转化为标准矩阵乘法的下界。

4. 关键结论总结

1. 直接求和定理（Direct Sum Theorem）： 对于 Transformer，不存在显著的“批量加速”效应。计算 $L \times H$ 个注意力头的总成本，在渐进意义上等于独立计算每个头的成本之和。
2. 小维度下界： 在 $m = \Theta(\log N)$ 时，下界为 $LHN^{2-o(1)}$，匹配朴素算法。
3. 大维度下界： 在 $m = N$ 时，下界为 $LHN^{\omega-o(1)}$，匹配利用快速矩阵乘法的朴素算法。
4. 模型无关性： 即使去除了 MLP 层，或者假设 MLP 和嵌入映射可以在线性时间内计算，这些下界依然成立。

5. 意义与影响

• 理论突破： 这是首次针对多层多头 Transformer 建立非平凡的计算下界，解决了该领域长期存在的“直接求和”问题。
• 算法设计指引： 结果暗示，试图寻找比 $O(LHN^2)$ 或 $O(LHN^\omega)$ 更优的精确 Transformer 计算算法（在通用输入下）可能是徒劳的。这解释了为何现有的亚二次方（Sub-quadratic）注意力近似算法（如 FlashAttention 的优化主要在 I/O，或近似算法如 Performer）往往以牺牲精度或特定假设（如低秩、稀疏）为代价。
• 硬件优化方向： 既然算法层面的渐进加速受限，未来的优化方向可能更多集中在硬件层面的 I/O 优化（如 FlashAttention）或针对特定数据分布的近似算法，而非通用的计算复杂度降低。
• 工具创新： 将 Baur-Strassen 定理应用于深度学习模型（Transformer）的复杂性分析，为未来研究其他深度网络结构的计算下界提供了新的技术路径。

6. 开放问题

论文最后提出了一些开放方向：

• 能否将大嵌入维度的下界扩展到 Word-RAM 模型（通用计算模型），而不仅仅是算术电路？
• 对于中间嵌入维度（$m = N^a, 0 < a < 1$），能否填补下界与上界之间的差距？
• 是否存在特定的输入分布或结构假设，使得 Transformer 能够被更高效地计算？

综上所述，该论文从理论计算机科学的角度严格证明了 Transformer 的计算难度，表明其核心瓶颈在于必须处理大量的独立注意力计算，目前的“暴力”计算方法在理论极限上已接近最优。