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

Mar 11, 2026cs.LG

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

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.

Shirin Amiraslani, Xin Gao

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

Sequence Modeling Higher-Order Structures

Ask PDF

Sep 3, 2026cs.CC

The Head Complexity of Boolean Functions in Single-Layer Attention

What can a single layer of self-attention compute? We study head complexity: the minimum number of attention heads required to compute a function in a one-layer attention-only model. We establish an exact hierarchy under this measure:

k

heads compute

k

-bit parity but cannot compute

(k+1)

-bit parity. The lower bound is unconditional in the two resources a transformer might otherwise exploit; it holds at unbounded embedding dimension and unbounded numerical precision. The proof rests on an alternating-sum obstruction: after clearing the softmax denominators, every monomial in the resulting decision polynomial omits at least one of the

k+1

input bits, forcing its correlation with parity to vanish. The same obstruction yields lower bounds for related tasks, including the well-studied multi-hop induction-head task. We also establish compactness bounds for embedding dimension and numerical precision. Specifically, a compactness theorem shows that any function computable at all can be computed with embedding dimension and precision bounded by the discrete data of the task, namely, head count, alphabet size, and length. Thus, potentially unbounded dimension or precision provably cannot substitute for heads. Finally, we derive nearly matching universal bounds for general binary functions:

2^n

heads suffice to compute every

n

-bit binary function, with one head per monomial in its multilinear expansion, while a counting argument shows almost all such functions require

Ω(2^n/n^2)

heads. This lower bound matches the upper bound to within a

\operatorname{poly}(n)

factor, even when dimension and precision are unbounded. Together, these results characterize head requirements for Boolean computation in this model.

Rajmohan Rajaraman, Ravi Sundaram, Amanuel Tesfaye

Northeastern University.

Multi-Head Attention Structural Complexity

Ask PDF

May 12, 2026cs.LG

Lower bounds for one-layer transformers that compute parity

This note shows that no self-attention layer post-processed by a rational function can sign-represent the parity function unless the product of the number of heads and the degree of the post-processing function grows linearly with the input length. Combining this lower bound with rational approximation of ReLU networks yields a margin-dependent extension for self-attention layers post-processed by ReLU networks.

Daniel Hsu

Linear Attention Rectified Linear Unit Networks Transformer Architectures Lower Bounds Approximation

Ask PDF