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

Sayak Chakrabarti; Toniann Pitassi; Josh Alman

cs.LGFeb 2, 2026

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

Authors: Sayak Chakrabarti, Toniann Pitassi, Josh Alman

Organizations: Computer Science, Columbia University, New York, NY 10027.

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.

Figures & tables

Mechanism	Exact cc	Apx cc	Bound
Self-attention	$n^{2+o(1)}$	$n^{1+o(1)}$	$o(\sqrt{\log n})$
$t$ -Tensor	$n^{3+o(1)}$	$n^{1+o(1)}$	$o((\log n)^{1/t})$
Strassen	$n^{\omega+o(1)}$	$n^{1+o(1)}$	$o(\sqrt{\log n})$
Tree (new)	$n^{2+o(1)}$	$n^{1+o(1)}$	$o(\sqrt{\log n})$
Poly (new)	$n^{t+o(1)}$	$n^{1+o(1)}$	$o((\log n)^{1/k})$

Table 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

. Alman and Song (2023) ; Alman and Song (2024) proved bounds for self-attention and tensor-attention, while we prove the rest.

Mechanism	2-fold	3-fold
Self-attention	No	No
3-Tensor	Yes	No
Strassen	Yes	No
Tree (new)	Yes	Yes
Poly (new)	Yes	Yes

Table 2: Compositionality results showing support for function composition. Peng et al. (2024) prove impossibility bounds for self-attention, Kozachinskiy et al. (2025) simulate 2-fold with Strassen-attention, while we prove the rest.

Panel of Figure 1 — Figure 1: Graphical representation for the tree polynomial $h(x_{1},\dots,x_{7})=x_{1}x_{2}+x_{1}x_{3}+x_{1}x_{4}+x_{2}x_{5}+x_{2}x_{6}+x_{4}x_{7}$

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.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

R_{i}:=U^{1}_{(i,1:r)}A(W^{3}_{(i,1:r)})^{T}\in\mathbb{R}.

Appendix

Algorithm 1 Algorithm to compute entry-wise approximation of

Att^{(S)}

Att^{(x_{1}x_{j_{i}}+g_{i}(x_{j_{i}},\bar{x^{i}}))}_{(1:n,\ell)}:=\frac{P^{(x_{1}x_{j_{i}}+g_{i}(x_{j_{i}},\bar{x^{i}}))}}{R^{(x_{1}x_{j_{i}}+g_{i}(x_{j_{i}},\bar{x^{i}}))}}.

Appendix

Algorithm 2 Algorithm to compute tree attention

Att^{(h)}

Att=[\frac{1}{d}K^{(1)}(K^{(2)}\oslash\ldots\oslash K^{(t)})^{T}]^{e}W^{(2)}\oslash\ldots\oslash W^{(t)},

Appendix

Algorithm 3 Algorithm to compute an entry-wise approximation of

Att^{(h)}

Figure 3: Training loss per epoch, averaged over 10 seeds, 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. Tree-attention learns faster and has less fluctuations.

Figure 4: Accuracy per FLOP, averaged over 10 seeds, for tree-attention, 1-layer self-attention and 2-layer self-attention for learning function composition. Notice that tree-attention learns more efficiently and the learning is stable.

Seq len	1-layer SA (ms)	2-layer SA (ms)	1-layer tree (ms)	1-layer 3-tensor (ms)	1-layer Strassen (ms)
20	$1.076\pm 0.057$	$1.775\pm 0.085$	$1.367\pm 0.057$	$1.442\pm 0.062$	$1.593\pm 0.086$
50	$1.079\pm 0.048$	$1.757\pm 0.060$	$1.363\pm 0.055$	$2.911\pm 0.044$	$1.594\pm 0.088$
100	$1.080\pm 0.048$	$1.781\pm 0.097$	$1.374\pm 0.060$	$13.813\pm 0.051$	$3.395\pm 0.081$

Appendix

Figure 5: Average running time of various attention schemes implemented on NVIDIA A100 GPU. Tree-attention performs as fast as self-attention, implying that hidden constants in the time complexity computations are small.

	Tree-attention	Self-attention
Generalization token accuracy	$\mathbf{0.727691\pm 0.013486}$	$0.723993\pm 0.008649$
Generalization exact match	$\mathbf{0.264919\pm 0.127609}$	$0.239024\pm 0.087350$

Appendix

Figure 6: Table for mean accuracies and standard deviation over 10 random seeds.

Figure 7: Plots for mean $\pm$ one standard deviation over 10 random seeds for token accuracy and exact match accuracy on the generalization set. Tree-attention has higher accuracy than self-attention.

Figure 8: Exact match accuracies with each seed on the generalization set for tree-attention and self-attention. Tree-attention reaches $\sim 40\%$ accuracy for 4 out of 10 random seeds.

Explore similar work

Beyond Pairwise Attention: Higher-Order Modular Attention for Efficient Sequence Learning
Mar 11, 2026Shirin Amiraslani, Xin GaoSequence Modeling Higher-Order Structures
Department of Mathematics and Statistics York University Toronto, ON M3J 1P3, Canada
The Head Complexity of Boolean Functions in Single-Layer Attention
Sep 3, 2026Rajmohan Rajaraman, Ravi Sundaram, Amanuel TesfayeMulti-Head Attention Structural Complexity
Northeastern University.
Lower bounds for one-layer transformers that compute parity
May 12, 2026Daniel HsuLinear Attention Rectified Linear Unit Networks

Mar 11, 2026, cs.LGJ/K move · Enter open · S save

Beyond Pairwise Attention: Higher-Order Modular Attention for Efficient Sequence Learning

Shirin Amiraslani, Xin Gao

Department of Mathematics and Statistics York University Toronto, ON M3J 1P3, Canada

Sequence Modeling Higher-Order Structures

Open paper AskPDF arXiv

Sequence modeling tasks can involve intrinsic higher-order dependencies, while standard self-attention assigns scores to token pairs and does not explicitly parameterize such interactions. We introduce Higher-Order Modular Attention (HOMA), which fuses pairwise attention with an explicit triadic attention pathway made tractable through overlapping blocks, local windows, and a low-rank projection. We compare HOMA with matched pairwise and purely triadic baselines on controlled PARITY and MATCH3 tasks, as well as TAPE benchmarks. HOMA is competitive with or outperforms the baselines, with its clearest advantages when the underlying dependencies extend beyond the explicitly modeled triadic order. These advantages are accompanied in several settings by faster convergence and improved parameter efficiency, with the learned nonlinear fusion providing an effective mechanism for combining the pairwise and triadic representations. Overall, our results provide empirical evidence that HOMA is an effective attention design when task structure extends beyond pairwise interactions.

Abstract

Explore similar work

Beyond Pairwise Attention: Higher-Order Modular Attention for Efficient Sequence Learning

The Head Complexity of Boolean Functions in Single-Layer Attention

Lower bounds for one-layer transformers that compute parity