arXiv:2501.19215cs.LGcs.AI2025-01NeurIPS被引 9

提出新型注意力机制,突破单层Transformer的推理能力极限。

Strassen Attention, Split VC Dimension and Compositionality in Transformers

  • 引入斯特拉森注意力,解决单层Transformer在复杂推理任务中的理论瓶颈。
  • 实验显示其在三项任务上显著优于标准注意力,且运行时间低于立方级。
  • 适合关注Transformer推理能力与高效注意力设计的研究者。

本文首次证明了任意精度(包括无穷精度)的单层softmax Transformer在三个需要高级推理的任务上存在理论局限性:匹配三元组(Match 3)、函数复合与二元关系复合。我们正式证明这些任务无法被单层Transformer解决。为克服此限制,提出斯特拉森注意力机制,证明其可使单层Transformer理论上解决所有上述任务,并具有亚立方时间复杂度,优于此前的高阶注意力等机制。实验对比了斯特拉森注意力、标准注意力(Vaswani et al., 2017)、高阶注意力(Sanford et al., 2023)和三角注意力(Bergen et al., 2021),结果揭示各机制的优劣。斯特拉森注意力在所有任务中显著超越标准注意力,验证了其有效性与可扩展性。理解这些理论边界有助于推动更高效、更强推理能力的Transformer设计。

原文摘要 · Abstract (English)

We propose the first method to show theoretical limitations for one-layer softmax transformers with arbitrarily many precision bits (even infinite). We establish those limitations for three tasks that require advanced reasoning. The first task, Match 3 (Sanford et al., 2023), requires looking at all possible token triplets in an input sequence. The second and third tasks address compositionality-based reasoning: function composition (Peng et al., 2024) and binary relations composition, respectively. We formally prove the inability of one-layer softmax Transformers to solve any of these tasks. To overcome these limitations, we introduce Strassen attention and prove that, equipped with this mechanism, a one-layer transformer can in principle solve all these tasks. Importantly, we show that it enjoys sub-cubic running-time complexity, making it more scalable than similar previously proposed mechanisms, such as higher-order attention (Sanford et al., 2023). To complement our theoretical findings, we experimentally studied Strassen attention and compared it against standard (Vaswani et al, 2017), higher-order attention (Sanford et al., 2023), and triangular attention (Bergen et al. 2021). Our results help to disentangle all these attention mechanisms, highlighting their strengths and limitations. In particular, Strassen attention outperforms standard attention significantly on all the tasks. Altogether, understanding the theoretical limitations can guide research towards scalable attention mechanisms that improve the reasoning abilities of Transformers.

Transformer注意力机制推理能力理论分析

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。