Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Poly-attention: a general scheme for higher-order self-attention

Sayak Chakrabarti, Toniann Pitassi, Josh Alman

TL;DR

This work presents poly-attention, a unifying framework that generalizes self-attention to higher-order tensor computations and arbitrary token-relationship structures. It formalizes Att^{(h)} via a polynomial h over multiple inputs, showing that standard self-attention, tensor-attention, and Strassen-attention are special cases, and introduces tree-attention as a quadratic-time, highly expressive subclass. The paper establishes a near-complete complexity landscape: tree-attention achieves exact quadratic-time computation and can realize r-fold function composition, while non-tree polynomials typically incur superquadratic costs unless weight-bounded approximations are used, with subcubic exact algorithms available in some degree-2-cycle cases. It also provides representational results such as polynomial root-finding with two heads, and experimental validation that tree-attention improves compositional tasks while maintaining practical runtimes. Overall, poly-attention offers a principled, tunable trade-off between expressiveness and computational efficiency for higher-order relational reasoning in Transformers.

Abstract

The self-attention mechanism, at the heart of the Transformer model, is able to effectively model pairwise interactions between tokens. However, numerous recent works have shown that it is unable to perform basic tasks involving detecting triples of correlated tokens, or compositional tasks where multiple input tokens need to be referenced to generate a result. Some higher-dimensional alternatives to self-attention have been proposed to address this, including higher-order attention and Strassen attention, which can perform some of these polyadic tasks in exchange for slower, superquadratic running times. In this work, we define a vast class of generalizations of self-attention, which we call poly-attention mechanisms. Our mechanisms can incorporate arbitrary higher-order (tensor) computations as well as arbitrary relationship structures between the input tokens, and they include the aforementioned alternatives as special cases. We then systematically study their computational complexity and representational strength, including giving new algorithms and matching complexity-theoretic lower bounds on the time complexity of computing the attention matrix exactly as well as approximately, and tightly determining which polyadic tasks they can each perform. Our results give interesting trade-offs between different desiderata for these mechanisms, including a tight relationship between how expressive a mechanism is, and how large the coefficients in the model may be so that the mechanism can be approximated in almost-linear time. Notably, we give a new attention mechanism which can be computed exactly in quadratic time, and which can perform function composition for any fixed number of functions. Prior mechanisms, even for just composing two functions, could only be computed in superquadratic time, and our new lower bounds show that faster algorithms for them are not possible.

Poly-attention: a general scheme for higher-order self-attention

TL;DR

This work presents poly-attention, a unifying framework that generalizes self-attention to higher-order tensor computations and arbitrary token-relationship structures. It formalizes Att^{(h)} via a polynomial h over multiple inputs, showing that standard self-attention, tensor-attention, and Strassen-attention are special cases, and introduces tree-attention as a quadratic-time, highly expressive subclass. The paper establishes a near-complete complexity landscape: tree-attention achieves exact quadratic-time computation and can realize r-fold function composition, while non-tree polynomials typically incur superquadratic costs unless weight-bounded approximations are used, with subcubic exact algorithms available in some degree-2-cycle cases. It also provides representational results such as polynomial root-finding with two heads, and experimental validation that tree-attention improves compositional tasks while maintaining practical runtimes. Overall, poly-attention offers a principled, tunable trade-off between expressiveness and computational efficiency for higher-order relational reasoning in Transformers.

Abstract

The self-attention mechanism, at the heart of the Transformer model, is able to effectively model pairwise interactions between tokens. However, numerous recent works have shown that it is unable to perform basic tasks involving detecting triples of correlated tokens, or compositional tasks where multiple input tokens need to be referenced to generate a result. Some higher-dimensional alternatives to self-attention have been proposed to address this, including higher-order attention and Strassen attention, which can perform some of these polyadic tasks in exchange for slower, superquadratic running times. In this work, we define a vast class of generalizations of self-attention, which we call poly-attention mechanisms. Our mechanisms can incorporate arbitrary higher-order (tensor) computations as well as arbitrary relationship structures between the input tokens, and they include the aforementioned alternatives as special cases. We then systematically study their computational complexity and representational strength, including giving new algorithms and matching complexity-theoretic lower bounds on the time complexity of computing the attention matrix exactly as well as approximately, and tightly determining which polyadic tasks they can each perform. Our results give interesting trade-offs between different desiderata for these mechanisms, including a tight relationship between how expressive a mechanism is, and how large the coefficients in the model may be so that the mechanism can be approximated in almost-linear time. Notably, we give a new attention mechanism which can be computed exactly in quadratic time, and which can perform function composition for any fixed number of functions. Prior mechanisms, even for just composing two functions, could only be computed in superquadratic time, and our new lower bounds show that faster algorithms for them are not possible.
Paper Structure (31 sections, 15 theorems, 26 equations, 4 figures, 2 tables, 1 algorithm)

This paper contains 31 sections, 15 theorems, 26 equations, 4 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Self-attention vaswani2017attention.
  3. Tensor-attention.
  4. Strassen-attention.
  5. Running time considerations
  6. Poly-attention is all you need
  7. The poly-attention mechanism
  8. Beyond self-attention: the power of poly-attention
  9. An example: function composition
  10. Tree-attention: polynomials leading to efficient poly-attention
  11. Computational complexity of non-tree poly-attention
  12. Representational strength of poly-attention
  13. Implications of poly-attention
  14. Technique overview
  15. Representational strength.
  16. ...and 16 more sections

Key Result

Lemma 2.3

Poly-attention captures all the previous higher-order self-attention techniques. In particular, (i) self-attention is poly-attention with the base polynomial $h(x_1, x_2) = x_1x_2$; (ii) $t$-tensor attention is poly-attention with $h(x_1, \dots, x_t) = x_1\dots x_t$; and (iii) Strassen-attention is

Figures (4)

  • Figure 1: This summarizes the running times of both exact and approximate algorithms for these attention variants. For entry-wise approximation (Apx cc), the bound $B$ is the maximum absolute value of the matrix entries such that we can entry-wise approximate the output matrix in near-linear time; the attention polynomial is in $t$ variables and has degree $k$. alman2023fastalman2023capture proved bounds for self-attention and tensor-attention, while we prove the rest.
  • Figure 1: Graphical representation for the tree polynomial $h(x_1, \dots, x_7) = x_1x_2+x_1x_3+x_1x_4+x_2x_5+x_2x_6+x_4x_7$
  • Figure 2: Compositionality results showing support for function composition. peng2024limitations prove impossibility bounds for self-attention, kozachinskiy2025strassen simulate 2-fold with Strassen-attention, while we prove the rest.
  • Figure 2: Accuracy per epoch for learning $f_1(f_2(x))$ for sequence length $51$, on a single layer of tree-attention, one layer self-attention and two layer self-attention.

Theorems & Definitions (30)

  • Definition 2.1: Attention polynomial
  • Definition 2.2: Poly-attention
  • Lemma 2.3
  • Theorem 3.1
  • Theorem 3.2
  • Theorem 3.3
  • Theorem 3.4
  • Theorem 3.5
  • Theorem 3.6
  • Theorem 3.7
  • ...and 20 more