Strassen Attention, Split VC Dimension and Compositionality in Transformers

TL;DR

This paper establishes theoretical limitations of one-layer softmax transformers on complex reasoning tasks and introduces Strassen attention, which overcomes these limitations with improved scalability and performance.

Contribution

The paper introduces Strassen attention, a novel mechanism that enables one-layer transformers to solve advanced reasoning tasks, overcoming previous theoretical limitations.

Findings

Strassen attention enables solving complex reasoning tasks.

It has sub-cubic time complexity, improving scalability.

It outperforms standard attention on all tested tasks.

Abstract

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.