Poly-attention: a general scheme for higher-order self-attention
Authors: Sayak Chakrabarti
Toniann Pitassi
Josh Alman
Venue: ICLR 2026 Poster (2026)
Keywords: computational complexity
polynomial method
fine-grained complexity
communication complexity
tensor generalizations
Primary area: learning theory
TL;DR: This article generalizes higher-order self attention, studies their representational strengths and computational complexities, and characterizes all self-attention mechanisms computable in quadratic time.

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 (Sanford et al., 2023) and Strassen attention (Kozachinskiy et al., 2025), 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 tradeoffs 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.

Published as a conference paper at ICLR 2026
POLY-ATTENTION" A GENERAL SCHEME FOR HIGHER-
ORDER SELF-ATTENTIO
Sayak Chakrabarti Computer Science, Columbia University,
Toniann Pitassi
Josh Alman
Computer Science, Columbia University,
Computer Science, Columbia University,
New York, NY 10027.
New York, NY 10027. ton (c.
New York, NY 10027.
sayaksc@gmail.. com
coni@ccs.oolubbi... www.mbll.aa.a.... edu
josh@cs www.coolloooooo. edu
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 result. Some higher- nigher-dimensional alternatives to self-attention have been proposed to address this, including higher-order attention (Sanford al., 2023) and Strassen attention (Kozachinskiy et al., 2025), which can perform some of these polyadic tasks in exchange for slower, superquadratic rumning times.
In this work, define 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-theorede lower bounds on the time complexity computing the attention matrix exactly well approximately, and tightly determining which polyadic tasks they can each perform. Our results give interesting trade-offs between different desiderata for these mechanisms, including 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.
INTRODUCTION
The transformer architecture, introduced by Vaswani et al. (2017), has the self-attention mechanism at its heart, which used to capture pair-wise wise correlations large language models. Since its inception. it has been used in variety of large language model (LLM) architectures, including BERT (Devlin
al. 2019), GPT series (Radford et al. 2018 Brown al. 20200; OpenAl 2023). Claude (Anthropic 2024). Llama (Grattafiori et al. 2024), and ol (OpenAl 2024). Its success has led to its prominent use in nearly every area of modern deep learning.
Transformers consist of three main components within each block: an input Multilayer Perceptron (MLP) layer, followed by a self-attention mechanism, then finally an output MLP layer Vaswani al. 2017). The self-ateention mechanism is function from Roxd - RR xxd which computes and combines weighted pairwise correlations between tokens in its input, and key to the success the Transformer model.

---

Published as a conference paper at ICLR 2026
Self-attention (Vaswani et al., 2017). For matrix M and index we write M, denote the ith row of M. Given a query matrix Q E gnxd key matrix K E R2 and value matrix V Raxd for specific input, the output of the self-attention function is given by the matrix Att E R=xd whose jth row is:
All; Ejel exp( 2 (Qi, K2)) V.
EzEn exp( Ky))
Despite the widespread use of self-attention in Transformers, there are limits to its expressive power, which is intuitively limited to capturing pairwise correlations between tokens. particular, researchers have defined number of basic tasks such as iterated function composition, Match3, Parity, Majority, and Dyck-1 which require higher order relationships than pairwise correlations and provably cannot be solved by simple self-attention networks (Sanford et al. 2024bb Peng et al.) 2024; Hahn 2020). Empirical studies have also confirmed this intuition, showing poor performance by simple Transformers
benchmark datasets like multiplication, logical puzzles and dynamic programming Dziri et al. (2023), memorized mappings (Zhang et al.] (2025) and other datasets like
SCAN (Lake & Baroni 2018), PCFG (Hupkes et al. 2020), CLUTRR (Sinha al. 2019). CoGS
(Kim & Linzen 2020), GFQ (Kyysers et al. 2020), and CREPE (Ma al. 2023).
In this paper, focus especially
type of task called function composition. As simple example, the language model may be given the query "If Sam lives in Toronto, Peter lives in Paris, Toronto is in Canada, and Paris is in France, which country does Sam live in?", and the model expected to
reply "Canada". This is composition of two functions: the first maps people to cities, and the second
maps cities countries. Several works including (Peng et al.) [0244] Dzzrr et al., [023 Lu al., 2023) have shown, both theoretically and experimentally, that simple language models are unable to perform these tasks. In order to overcome these representational limitations, several stronger attention mechanisms have been proposed, notably higher-order tensor atlention and Strassen attention which we define next.
Tensor-attention. Clift et al. (2020) came up with a tensor generalization of self-attention, called simplical attention, Shnchr Sanford al. (2024b) also studied the higher-order tensor attention (that we will call 3-tensor attention) for a query matrix Q(1) E Raxd, key matrices Q(2), Q(3) E R
and value matrices (22,, (33 E IRaxd The output is given by the matrix At((T) E Igadd whose ith row is given by:
AT Ennlee (2 080 Q ((2) V(3) (Q(") 0 (2 1
ls
Ea exp(
Here ④ denotes the element-wise product (also called Hadamard product), and for three vectors a. b, c E Rd, we define (a, b, c) Σ a[[]b[[]c[[].
Sanford et (2024b) showed that one 3-tensor attention head can solve more complicated tasks like Match3, which requires finding a triple of correlated tokens. They also defined a natural generalization to f-tensor attention, which can solve Match-t for t 3.
Strassen-attention. Later, Kozachinskiy al. (225 gave more efficient attention mechanism that can also perform Match3 and several other tasks difficult for self-attention. (As we will discuss shortly, 3-tensor attention can have prohibitive computational complexity, and Strassen-attention was defined as step toward addressing this.) This attention mechanism is again defined over query matrix Q()) E Raxd key matrices Q(2), Q(3) E R2xd and value matrices (2) V(3) E Rnnd The output matrix is (s) E paxdd where the ith row, for i E [n], is given by:
AT = Elshen exp( Q/ 10 2
1/ O v(
Eegbee exp() (( (2) + (0 Q (3) + (2 Q.
Quite recently, 3-tensor attention has been implemented and performances studied by Roy al. (2025). We refer the reader to Section B in which we survey other attention mechanisms and the landscape of results known about them more detail.
2

---

Published as a conference paper at ICLR 2026
1.1 RUNNING TIME CONSIDERATIONS
A natural trade-off arises in these proposed attention mechanisms: as the attention mechanism becomes more general to give more representational power, the required running time increases too. This can often be prohibitive: the quadratic running time of self-attention is already computational bottleneck which is mitigated in practice only by extensive hardware; a superquadratic running time may not be practical even with such hardware speedups.
We compare here the running times of various attention mechanisms as a function of n, the number of input tokens, where the embedding dimension is O(log n): see running times in Table [ below.
Exact Algorithms. The best algorithms for self elff-attenioon take time n22+((1), matching the straightfor- ward algorithm. For tensor attention, the best algorithm is also the straightforward algorithm, which for t-tensor attention (t 2 3) runs in superquadratic time (
The straightforward algorithm for Strassen attention, just following its definition, takes time n 3+++(()) However, (Kozachinskiy et al. (2025) give a faster algorithm for Strassen attention with running time O(n"), where w I 2.37 1371477 is the exponent of matrix multiplication (Alman et al. 2025), i.e., the constant such that still truly supercubic, and moreover, we note that the aforementioned bound on w comes from a highly theoretical algorithm, and typically either R 2.81 from Strassen's 's algorithm (Strassen 1969), or even w = 3 from the straightforward matrix multiplication algorithm, are used in practice. (Kozachinskiy et al. (2025) named it after Strassen's 's matrix multiplication algorithm emphasize this faster algorithm.)
matrices can be multiplied in time On This faster algorithm is
It is natural wonder whether even faster algorithms are possible, and particularly whether tensor attention or Strassen attention could be computed in quadratic time. In fact, these known running times are known be optimal under standard complexity-theoreti assumptions, so these algorithms cannot be improved. For self-attention and tensor attention, this was shown in prior work (Alman & Song. 2023 2024); for Strassen attention, we prove this here in Theorem 3.6 below.
Approximation Algorithms. In most cases, sufficiently accurate approximation of self-attention suffices, and this can sometimes be computed much faster. Alman & Song (2023) shows that as long as the entries of the query and key matrices are bounded (and all have magnitude at most B = (oooog n we can compute an entry-wise approximation of the self-attention matrix in almost linear time, n² +011) approximation of tensor attention !tt!T)) in n +o(() time, with a smaller bound on B. These prior works have also shown matching lower bounds, showing that these bounds B are tight: if the weights are even slightly larger, than the straightforward exact running times discussed above are unavoidable. (These lower bounds use standard assumptions from fine- ene-grained complexity theory; see Section for more details.) Many different lines of experimental work studied Transformers with reasonable
Alman & Song (0244) similarly showed how to compute an entry-wise
precision guarantees (Zafrir et al. 2019 Sun et al. 2019; Katharopoulos et al. 2020 Dettmers et al. 2022 Xiao et al. 2023 Dettmers et al.
2022 Perez et al. 2023 Roy et al. 2021 Han ct al. 2024).
In this paper, we build on this line of work and give the first fast approxi- mation algorithm for Strassen atten- tion. We show that, if all the weights bounded by B = o( o(vlog n),, then one can approximate Strassen atten- tion in almost linear time n tl)))), and if the weights are larger, then the ex- act running time of n01)) cannot be avoided (again using fine-grained complexity assumptions). This lower bound fits within a new, much more general lower bound on different gen- eralizations of attention which we will state in Theorem 3.6 In partic-
Mechanism
Exact cc
Apx cc
Bound
Self-attention
n++o(1)
n 2
o( n.)
1- T-Tensor
3
nl+o(1)
((o)g))1/))
Strassen
n nl+o(1)
o(vlog n
Tree (new)
2++0(1)
n t1+0(1))
o(viog
Poly (new)
n =1+((1)
n1+o(1))
o((log 1/kk)
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 max- imum 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 & Song 2003 2024) proved bounds for self-attention
and censor-attettoon while we prqve the rest. matrix where each entry most poly(n) far from the exact value.
An entry-wise approximation outputs
3

---

Published as a conference paper at ICLR 2026
ular, although the statement appears similar to prior work, proving this requires substantial new techniques, since prior techniques focused on proving cubic lower bounds, but Strassen attention actually has a subcubic (but superquadratic) time algorithm based matrix multiplication; Section 4 for more details.
1.2 POLY-ATTENTION IS ALL YOU NEED
In this work we introduce more general class of attention mechanisms called poly-attention that generalizes and improves upon these previous attention mechanisms. An instantiation of poly- attention is given by base polynomial, over variables, degree k and sparsity We will precisely define voly-attention shortly, and show that it includes self-attention, tensor attention, and Strassen attention as special eases.
Our main results include complete and exhaustive analyses of the running times one can achieve to compute or approximate different poly-attentions, as well as the expressive power of each one. Using these, we identify new, specific instantiations of poly-attention which are simultaneously more expressive than self-attention, and easier to compute than prior replacements to self-attention. One may also between expressiveness and computational complexity.
results to identify attention mechanisms of interest which achieve a desired trade-off'
Tree-attention. We particularly highlight subclass of our poly-attention mechanisms that we call treeaattenion, which loosely speaking is characterized by a subclass of degree-2 base polynomials h that possesses tree-like property. We find that all tree-attention mechanisms be computed in quadratic time, matching the running time of standard self-attention. Furthermore, we show that tree-attention solve r-fold function composition for any constant
This is a substantial improvement on prior attention mechanisms. Self-attention cannot even solve 2- 2old function composition. Meanwhile, 3-tensor attention and Strassen attention, which can solve 2-fold function composition, require superquadratic time, and furthermore, they cannot solve 3-fold function composition. Our new tree-attention can solve r-fold function composition for all r and can be computed in quadratic time (Theorem 3.4).
We give a more detailed analysis of tree-attention, including tight exact and approximation algorithms, in Section 3.2 (Theorem 3.5p. We posit tree-attention as the best of all worlds in terms of representa- tional strength and time complexity. self-attention mechanism, we will see that the runtime of tree-attention matches the best possible runtimes both the exact and approximate versions. We envision two types of users/applications.
addition to strictly improving the expressive power of the
if quadratic running time can be tolerated then the exact algorithm for tree-attention
if faster, almost linear rumning time is needed, then the user should find the largest bound on the weights which can be tolerated by their hardware and architecture, and then apply the most expressive tree-attention which can be approximated quickly for that B (we will explore the trade-off Section
We emphasize that exact and approximate algorithms for tree-attention only use straightforward matrix multiplication algorithms, and do not rely on bounds on w or other impractical fast matrix multiplication algorithms. See Section 5 for experimental validation.
Full characterization of poly-attention. Beyond tree-attention, we give a full characterization of the running time needed to compute poly-attention a function of the underlying properties of the base polynomial, ht We find that these mechanisms often require cubic or more time to compute exactly, but nonetheless have fast approximation algorithms when B (the bound on the weights) is small enough, and meanwhile can perform very complex tasks.
2 THE POLY- COLY-ATTENTIO MECHANISM
In this section, we define the general class of poly-attention mechanisms. They will be described by a special class of multi-lineal polynomials, which we will call attention polynomials.
4

---

Published as a conference paper at ICLR 2026
Definition 2.1 (Attention polynomial). We call polynomial h h(21...... 22) attention polynomial of degree k if it is multi-linear, it has coefficients only in {0, 1), and all its monomials have degree at least 2 and at most k.
Attention polynomials will be a central concept in this article. We will use them to concisely denote combinations of inner products of vectors. Given vectors Yi...., Yi R consider multi-linear monomial of an attention polynomial, m, of degree containing variables :'....... je where 1 S j1 .... < jk S t.. We denote m.y...... Y,) == (Yn:: -a.....; Yja). which an inner product of order k. Then, given an attention polynomial h(r1,--... It) containing s monomials m1,.... Mg. we define i(Yi,.... Y)) = Eie(s] (Y......... Yi).
Now, we describe our class of poly-attention mechanisms, of order I, using attention polyno- mial (x......, Iz) of degree k having 8 monomials (typically think of I, k, as small constants). Definition 2.2 (Poly-attention). For attention polynomial h(z)...., 244 having monomials of degree at most k, we define the poly-attention function from Ruxd to Raxd, which depends on h and has, as its paramelers, query-key weights W WQw)...., Woun Baxdd and value weights Wes...., Wy E Rdxdd
For an input
Raxd, the query-key matrices are denoted as Q XW S.W........... .0.))................ XW qte, and the value matrices as V(2) :- X w........ v(I) = XWwo.
The oulput of the poly-attention function will be given by the matrix
(Q Q(t), v(), ,..., v(o) TRaxd
where the 11-th row defined as:
......,!!) Etamnheen exp (2h(O)) ......
exp H ..... 0!!)
(1) 1/2 ( (33) O
((()
(1) Aut
(1)
We will often drop the ((0's and Vûl's from the notation (b) when it doesn't lead to ambiguity.
Here, Q(1) will be the query matrix as used in the usual self- mechanisms, and Q(2),.... Q( will be the key matrices, as the index of the row of Q(1) corresponds to the row of the output of poly-attention, and correlations variables (and hence, the matrices) in a symmetric sense, we denote both the query and the key matrices using Q(3) for ease of notation.
considered with respect that. However, since use all the
Lemma 2.3. Poly-attention captures all the previous higher-order self-attention techniques. In particular, (i) self seff-attention is poly-attention with the base polynomial h(21, x2) = 22 (ii) t- tensor attention is poly-attention with h(........ xx) 21 Ii and (iii) Strassen-attentte is
poly-attention with h(v1, D2, X3) = 3132 + D2223 + F3T1.
BEYOND SELF-ATTENTION. THE POWER OF POLY-ATTENTION
In this section, we study the strength and limitations of the poly-attention scheme. We begin in Section 3.I by studying an illustrative example. Thereafter, we will consider tree-attention and poly-attention full generality.
Mechanism
2-fold
3-fold
3.1 AN EXAMPLE: FUNCTION COMPOSITION
Self-attentin
No
No
3-Tesoor
Yes
No
To demonstrate the power of poly-attention, we analyze a special case when h(x1, 22, x3) x1 + T2 We show that this specific poly-attention efficiently solve important tasks faster than any other previous attention mechanisms.
Strassen
Yes
No
Tree (new)
Yes
Yes
Poly (new)
Yes
Yes
Table 2: Compositionality results show-
ing support for function composition. Peng et al. (0044) prove impossibil- ity bounds for self-attention, Kozachin skiy et al. (2025) simulate 2-fold with Strassen-attenten while we prove the rest.
To demonstrate the strength of this polynomial h, we de- fine the function composition problem demonstrated ear- lier. Mathematically, the 2-fold function composition prob- lem is: given two functions fi, fa [n] - [n] and x E [n].
5

---

Published as a conference paper at ICLR 2026
output f2))) To express this problem for attention mechanism, the input is X [[2++1)xd where X, for
i E [n] contains an encoding of fi(),, X, for j E [m+ + 1, 2n] contains an encoding f2(j - n)) and X2n+1 contains an encoding of z;; and our goal is output the value of (x)) in the (2n + 1)-th
entry of the output.
Peng et al. (024) proved that self- cannot simulate 2-fold function composition, and even that almost ni self-attention heads are needed in order to solve it. Since self-attention needs quadratic time compute, it would take cubic time compute n heads. All prior mechanisms that solve require superquadratic time. This leads to our punchline: poly-attention for this very simple polynomial h2 can simulate 2-fold function composition in just quadratic time!
this, including 3-tensor attention and
Theorem 3.1. Let h2(21, 12, xza) = 2112 + 1233. Poly-attention for h2 can simulate function composition using only one head. Furthermore, Att can be computed in O(r²) time.
We will tightly characterize what weights are needed for efficient approximation of all poly-attentions: in the case of Althaa,, we find:
Theorem 3.2. Given the polynomial (2), 22, va) = 2 22 + 22 where the entries of the query- key matrices are [-B, B):
1. If B ol ol) we can compute an entry-wise (V/ool((m) approximation of Att(ha) in time n (1))
2. If B = MyVlggn), n)), then every algorithm for computing an entry-wise (Vpoly(m))- approximation of (tt(hr)) requires time St(n²), unless SETH is false.
We consider next 3-fold function composition, in which the input is three functions, fı, f2, fs [n] + [n] and x E [n]. and we want to compute fa(f2(f To our knowledge, no prior attention mechanisms could perform 3-fold function composition. In particular, although Strassen-attention and 3-tensor attention were designed to solve problems like 2-fold function composition, we prove that they cannot compute 3-fold function composition when the precision is bounded:
Theorem 3.3. Strassen-attention and 3-tensor attention, require least H n1-o(11) heads to simulate 3-fold function composition when the precision is bounded.
However, we prove that poly-attention can indeed simulate 3-fold composition, and even more generally r-fold composition for any constant r, and still be evaluated in quadratic time!
Theorem 3.4. For any integer r 2 2, define the polynomial h, br, = x 22 + D2333 +...
I,Ir+1- Then, poly-attention for hy can simulate r-fold function composition, and Att can be computed exactly in time Or²3n? (input dimension here O(rn), not n).
In fact, we give a general characterization of which polynomials h can be used in (tt(h) to perform r-fold function composition. For example, we will also prove that poly-attention for h can not simulate r-fold function composition.
3.2 TREE-ATTENTION: POLYNOMIALS LEADING TO EFFICIENT POLY- ATTENTION
We saw in the previous section that instances of poly- attention which can be computed in quadratic time can have great representational strength. A natural question arises: what is the class of attention polyno- mials that can be exactly computed in only 2+) time? Could there be even stronger ones? We answer this by giving a complete characterization. We first define few notations to describe them.
322
33's
34
35
Z6
27
Figure I: Graphical representation for the tree polynomial h(FI, ... xT7 = 22 + T133 + 2124 + Z2 + D226 + 24
For an attention polynomial lzz1..... 22) of degree 2. say that simple graph G is the graphical rep- resentation of h, if G contains t vertices ....... Vt.. where vertex is corresponds to the variable I;, and

---

Published as a conference paper at ICLR 2026
there exists an edge between and v if and only if 242, is monomial present in h. If the graphical representation of h is a tree or a forest, we say that is a tree polynomial, and poly- attention for tree polynomial will be called rreeattention.
Our main result about tree-attention shows that it can be computed just as efficiently as self-attention, both for exact algorithms (where it can be computed in quadratic time) and approximate algorithms (which has the same bound B = of olyoog ni as self-attention, which is also the largest bound for any poly-attention):
Theorem 3.5. Given tree polynomial h, where the entries of the query-key matrices are in B-B, B}:
1. The output of tree-attention, Atth), can be exactly computed in n2+o(1) time.
2. 1f B =(ylog entry-wise approximation of Alt can be computed in n²++(1) time.
3. If =3Q(lo n). under standard complexity assumptions, entry-wise approximation of Att(h) requires (?n2)) time.
Tree polynomials include the polynomials h from Theorem 3.4 which can compute function composi- tion. More generally, the poly-attention for a tree polynomial, where the tree has depth can simulate (g - 1)-fold function composition, as well as a variety of tree generalizations. (Function composition can be naturally seen as corresponding to the path graph, which is the graphical representation of h,.)
We show next that, for any attention polynomial which is not tree polynomial (either because has degree more than 2, or because its graphical representation contains a cycle), its poly-attention requires superquadratic time to compute. Thus, as promised, tree-attentions form a complete characterization of quadratic-time poly-attentions.
3.3 COMPUTATIONAL COMPLEXITY OF NON-TREE POLY-ATTENTION
Next, we give a complete characterization of the computational complexity (both exact and approxi- mate) for poly-attention for all attention polynomials h.
Theorem 3.6. Given poly-attention for an attention polynomial M(z1...., ,Ie) of degree i and sparsity s which is not a tree polynomial, where the query-key matrices have entries [-B, B):
1. If B = (((logn)//5 entry-wise (nly approximation of Att(h) can be computed in almost-linear time.
2. If B 52((logn)) 1/k,,, time, assuming standard complexity assumptions.
entry-wise polyin) approximation of Att(h) requires superquadratic
Prior work gave this characterization for specific polynomials (Alman & Song 22033 for the usual self-attention (i.c., h(11, x2) I2)2), followed by Alman & Song (2024) for t-tensor attention i.c., h(2)...., xx) = I1 ...... 22). We discuss in Section below a number of technical hurdles which we 4
overcome generalize their results all attention polynomials and prove Theorem 3.6
Notably, for many polynomials such as h(21, F3) = F1 F223 + Ii (corresponding to prior
Strassen attention), there is subeubic algorithm which uses fast matrix multiplication, approaches, which can only prove cubic (or above) lower bounds, cannot apply. In fact, we generalize the Strassen attention algorithm (Kozachinskly al. 2025), and prove that for any degree-2 attention polynomial h whose graphical representation contains exactly one cycle, there is an exact algorithm for (h) running in subcubic time Onn)). and that this cannot be improved.
3.4 REPRESENTATIONAL STRENGTII POLY-ATTENTION
We have discussed function composition at length, but poly-attention is also able to perform variety of other basic expressive problems. As an example, Match3 has been highlighted by prior work (Sanford et al. (2024a| Kozachinskiy et al.) 2025) problem which requires detecting correlated triples of tokens. We define here a generalization called polynomial fondingnding which can be solved by poly-attention.
7

---

Published as a conference paper at ICLR 2026
The problem is defined in terms of fixed polynomial p(.1...... zn) (which, unlike attention polynomial, may have degree monomials, and may not be multi-linear). In the problem. given as input set S containing n integers, and the goal is to find y1.....: 1/2 E S such that p(y). ........... V/) 0.
Match3 is a special case of root- root-finding, corresponding to the simple polynomial p(æ1, 22, 2z3) =
T1 + 22 + 23. Circuit evaluation for constant sized circuits, and other related problems can also be captured by polynomial root-finding by using arithmetization. We prove that for any polynomial p. one can solve polynomial root-finding using poly-attention:
Theorem 3.7. For every polynomial p(æ1....,22) R2), there is attention polynomial h h( m2) such that a Transformer using two heads of poly-attention for can solve polynomial root-finding
Finding the attention polynomial h for a given polynomial p using our approach is straightforward but requires some details; it could be performed by a user who would like to answer query patterns corresponding to polynomial root-finding non-fin for a particular p.
3.5 IMPLICATIONS OF POLY-ATTENTION
As we have seen, tree-attention can solve many problems which self-attention cannot, and still it can be computed in quadratic time. We show that this quadratic time is indeed practicable by showing in Figure 5 that the time-complexity does not hide large constants.
We further show in experiments in Section H that tree-attention is indeed more expressive than self-attention. This seems to be a promising area of research, and it will be interesting to study the large scale deployment of tree-attention instead of self-attention in follow-up work. One can select an appropriate pree-polynomial to use depending on the relationships between the data that the model intends to process.
When we move to more general poly-attention, for any attention polynomial which is not a tree polynomial, we have shown in Theorem 3.6 that (without a small bound on the model weights) poly- attention provably requires superquadratie time. Thus, there is trade-off between expressiveness (most straightforwardly represented by the degree and order of the polynomial h, although it could also take into account which tasks like polynomial root-finding can be performed). and running time (depending on how bounded the entries must be). Model designers therefore have a choice, potentially depending on the hardware available to them, the desired running time, and the logical structures they expect see in their data and queries.
would be exciting, in future work, to further study the expressive power of tree-attentions other than the ones studied here, and find more examples of complicated tasks with tree-like logical structures that it can solve. As an example, Peng et al. (2024) proposed some more problems like relationship composition, spatial composition and temporal composition which current language models cannot solve; it would be interesting to see how well tree-attention performs on these problems.
4 TECHNIQUE OVERVIEW
Representational strength. Our representational strength results include both constructions (e.g.. showing that tree-attention can perform r-oldd function composition) and lower bounds (e.g., showing that Strassen-attention and 3-tensor attention cannot perform 3-fold function composition).
Our constructions use generalization of the "sum of squares" approach of Kozachinskiy et al. (025):: :
If one can design a simple polynomial e which checks possible outputs of function composition, so that it outputs 0 on correct outputs and large values on incorrect values, then the softmax underlying attention can detect Os and thus solve the problem. An interesting algebraic challenge arises of expressing in terms of the monomials available in an attention polynomial h.
Our lower bounds make use of communication complexity theory, similar to many other representa- tional lower bounds in the literature. We show that if mechanisms, then there is a resulting, very efficient communication protocol for a problem called myopic pointer jumping. Results from Chakrabarti (2007): Kozachinskiy al. 2025 showing that myopic pointer jumping cannot be solved with small communication can then be applied.
can be simulated by these
Fast approximation algorithms. For obtaining entry-wise approximation algorithms for poly- attention, use low-rank decomposition methods based the polynomial method, which were first applied in the context of Gaussian kernel density estimation (see Aggarwal & Alman (2022); Alman & Guan 2.24 In this approach, one critically approximates the exponential function (part
8

---

Published as a conference paper at ICLR 2026
of softmax) with low-degree single-variable variable polynomial. The bound B on the weights then naturally comes into play: the smaller the interval one must approximate the exponential on, the lower degree polynomial one may use.
similar approach has been used to design approximation algorithms for other variants on attention Allan & Song (2023) 2024 2025), although a number of intricacies arise in this general setting. For instance (recalling that is the number of variables in the attention polynomial h, and is the degree), directly applying the approach of Alman & Song (024) would yield an approximation algorithm whenever B of (logg n) V*), but our algorithm works even for the much larger bound o (log n)'"), This is significant improvement for > k- in tree-attention, one could choose / = 20 but A: = 2.
Lower bounds. Our running time lower bounds. where show that different poly- poly-attention mechanisms cannot be computed in quadratic time (for big enough bounds B on the weights), make of tools from fine-gre one-grained complexity theory. In particular, as in the previous works of Almmn & Song (2023; 2024 2025) on the fine-grained complexity of attention mechanisms, we use a popular conjecture called the Strong Exponential Time Hypothesis (SETH) to obtain conditional hardness results. First introduced in Impagliazzo & Paturi (2001), SETH is strengthening of the P * NP conjecture (so, proving SETH would imply P + NP), and perhaps the most widely used conjecture in fine-grained complexity.
Notably, the way SETH has been used prior work results in cubic (or higher) lower bounds, and makes it difficult to prove lower bounds for running time 2((=') from the matrix multiplication exponent w < Indeed, for such lower bound, our starting assumption must itself use matrix multiplication in some way!
In order to prove our lower bound against Strassen attention and other poly poly-attention mechanisms with Onn²) running times, we therefore use different conjecture, the Max-2SAT conjecture (see Alman & Vassilevska Williams (2020) and its uses in El Halaby 6016);Jansee & Włodarczyk (2024);
Bringmann & Slusallek (2021) Lincoln et al. (2018), which roughly asserts that our current best algorithm for the Max-2SAI problem cannot be substantially improved. We ultimately show that faster algorithm for Strassen attention could be used to design faster algorithm for Max-2SAT,
refuting the conjecture. Our proof of this makes use of the distributed PCP framework (Abboud et al. 2017) for reducing variants of SAT other problems through multi-party interactive communication protocols (Aaronson & Wigderson 2009; Rubinstein 2018
EXPERIMENTAL VALIDATION
We have proved that tree-attention can be computed in the same ( time as self-attention, and can simulate function composition (whereas self-attention cannot). We complement this with simple experiment to demonstrate empirical learnability and efficiency. We compare the following models: (i) model with one head and one layer of tree-attention; (ii) a model with one head and two layers of self-attention; and (iii) model with one head and one layer of self-attention. We train all three in the same way to solve function composition. As expected (proved by Peng et al. (2024)). one head and one layer of self self-attention is able learn function composition, but find that the other two are. Furthermore, we find that our tree-attention model learns function composition in many fewer training epochs. Lastly, our empirical evaluation of inference time validates that tree-attention takes roughly similar time as Self-aattentonnnJ]Ss See Figure 2 for summary, and Section for further details and quantitative results.
My Mayay My
Figure 2: Accuracy per epoch for learn- ing fi ff(()))) for sequence length 51, on single layer of tree-attention, one layer self-attention and two layer self- attention.
We also perform experiments comparing simple networks with self-attention and tree-attention on the COGS NI.P dataset Kim & Linzen (2020) This is a dataset which tests whether model can perform simple compositional tasks when processing language. We find that networks with tree-attention learn to higher accuracy in the same number of epochs. See Section H for more details and quantitative results.
As shown in Figure trce-attention takes around 1.3x time that of self-attention.

---

Published as a conference paper at ICLR 2026
6 ACKNOWLEDGEMENTS
We especially thank Todd Morrill for valuable discussions that greatly improved the experiments in this paper. We also thank Richard Zemel for helpful suggestions that strengthened both the paper and the experimental evaluation. SC thanks Hantao Yu for pointing out Kozachinskiy et al. (2025) and Roy al. (2025). Finally, we thank the anonymous reviewers for comments that improved the presentation.
7 ETHICS STATEMENT
We affirm that all aspects of this research comply with the ICLR Code of Ethics. This paper does not involve human subjects, personally identifiable data, sensitive applications, and we do not foresee direct ethical risks.
REPRODUCIBILTTY STATEMENT
The paper contains theoretical results to categorize higher-order self-attention mechanism, and provide a fundamental framework for future work. All these results, including theorems and algorithms, have complete proofs, presented in the appendix. roadmap to the proofs has been provided Section A. for the reader.
The code which produces the experimental results described in Sections and H can be found in the supplementary materials.
REFERENCES
Scott Aaronson and Avi Wigderson. Algebrization: A new barrier in complexity theory, ACM Transactions on Computation Theory (TOCT), 1(1):1-54, 2009.
Amir Abboud, Ryan Williams, and Huacheng Yu. More applications of the polynomial method to algorithm design. In Proceedings of the twenty-sixth annual ACM-SIAM symposium on Discrete algorithms, pp. 218-230. SIAM, 20145 4a.
Amir Abboud, Virginia Vassilevska Williams, and Oren Weimann. Consequences of faster alignment of sequences. In Automata, Languages, and Programming: 41st International Colloquium, ICALP 2014, Copenhagen, Denmark, July 8-11, 2014, Proceedings, Part 41, pp. 39-51. Springer, 2014b.
Amir Abboud, Arturs Backurs, and Virginia Vassilevska Williams. Tight hardness results for les and other sequence similarity measures. 2015 IEEE 56th Annual Symposium on Foundations of Computer Science, pp. 59-78. IEEE, 2015.
Amir Abboud, Aviad Rubinstein, and Ryan Williams. Distributed pep theorems for hardness of approximation p. 2017 IEEE 58th Annual Symposium on Foundations of Computer Science (FOCS), pp. 25-36. IEEE, 2017.
Elie Abboud and Noga Ron-Zewi. Finer-grained reductions in fine-grained hardness of approximation. Theoretical Computer Science, 1026: 114976, 2025.
Amol Aggarwal and Josh Alman. Optimal-degee polynomial approximations for exponentials and gaussian kernel density estimation. In 37th Computational Complexity Conference (CCC 2022). für Informatik, 2022.
Schloss Dagstuhl-LLeibniz-ae
Albert Alcalde, Giovanni Fantuzzi, and Enrique Zuazua. Clustering in pure-attention hardmax transformers and its role in sentiment analysis. arXiv preprint arXiv:2407.01602, 2024.
Josh Alman and Yunfeng Guan. Finer-grained hardness of kernel density estimation. In 39th Computational Complexity Conference, CCC 2024, July 22-25, 2024, Ann Arbor, MI. USA, volume 300 of LIPIcs, pp. 35:1-35:21, Schloss Dagstuhl Leibniz-Zentruu für Informatik, 2024.
Josh Alman and Zhao Song. Fast attention requires bounded entries. Advances in Neural Information Processing Systems, 17-631117-33555. 2023.
10

---

Published as a conference paper at ICLR 2026
Josh Alman and Zhao Song. How to capture higher-order correlations? generalizing matrix soft- max attention to kronecker computation. In The Twelfth International Conference on Learning Representations. ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024.
Josh Alman and Zhao Song. Fast rope attention: Combining the polynomial method and fast fourier transform. arXiv preprini arXiv:2505. 1892, 2025.
Josh Alman and Virginia Vassilevska Williams. OV graphs are (probably) hard instances. In Thomas Vidick (ed.), 11th Innovations in Theoretical Computer Science Conference, ITCS 2020, January 12-14, 2020, Seattle, Washington, USA, volume 151 of LIPIcs, pp. 83:1-83:18. Schloss Dagstuhl- Leibniz-Zentrum Zentrum für Informatik, 2020.
Josh Alman and Virginia Vassilevska Williams, refined laser method and faster matrix multiplica- tion. TheoretiCS, 3, 2024.
Josh Alman and Hantao Yu. Fundamental limitations subquadratic alternatives to transformers. In The Thirteenth International Conference on Learning Representations, 2025. URI. https
/openreview net/ wttp://wwww.et/teuteeemmeaa d T2dCgeb6y0
Josh Alman, Timothy M Chan, and Ryan Williams. Polynomial representations of threshold functions and algorithmic applications. In 2016 IEEE 57th Annual Symposium on Foundations of Computer Science (FOCS), pp. 467-476. IEEE, 2016.
Josh Alman, Timothy Chu, Aaron Schild, and Zhao Song. Algorithms and hardness for linear algebra on geometric graphs. In 2020 IEEE 61 Annual Symposium on Foundations of Computer Science (FOCS), pp. 541-552. IEEE, 2020.
Josh Alman, Ran Duan, Virginia Vassilevska Williams, Yinzhan Xu, Zixuan Xu, and Renfei Zhou. More asymmetry yields faster matrix multiplication. In Proceedings of the 2025 Annual ACM-SIAM Symposium on Discrete Algorithms (SODA). pp. 2005-20235 2039. SIAM, 2025.
AI Anthropic. The claude model family: Opus, sonnet, haiku. Claude-3 Model Card, 1:1, 2024.
Arturs Backurs and Piotr Indyk. Edit distance cannot be computed in strongly subquadratic time (unless seth is false). In Proceedings of the forty-seventh annual ACM symposium on Theory of computing, pp. 51 11-5, 2015.
Arturs Backurs, Moses Charikar, Piotr Indyk, and Paris Siminelakis. Efficient density evaluation for smooth kernels. In 2018 IEEE 59th Annual Symposium on Foundations of Computer Science (FOCS), pp. 615-626. IEEE, 2018.
Leon Bergen, Timothy O' Donnell, and Dzmitry Bahdanau. Systematic generalization with edge transformers. Advances in Neural Information Processing Systems, 1011111900033 2021.
Karl Bringmann and Jasper Slusallek. Current algorithms for detecting subgraphs of bounded treewidth are probably optimal. In Nikhil Bansal, Emanuela Merelli, and James Worrell (eds.), 48th International Colloquiun on Automata, Languages, and Programming, ICALP 2021, July 12-16, 2021, Glasgow, Scotland ( (Virtual Conference), volume 198 of LIPlcs, pp. 40:1-40:16. Zentrum für Informatik, 2021.
Schloss Dagstuhl-
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-19011 2020.
Amit Chakrabarti. Lower bounds for multi-player pointer jumping. In twenty-Second Annual IEEE Conference Computational Complexity (CCC'07), pp. 33-45. IEEE, 2007.
Timothy M Chan and Ryan Williams. Deterministic apsp, orthogonal vectors, and more: Quickly Proceedings of the twenty-seventh anmual ACM-SIAM
derandomizing razborov-smolenske symposium Discrete algorithmss, pp. 1246-1255. SIAM, 2016.
Beidi Chen, Tri Dao. Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. Advances in Neural Information Processing Systems, 34: 17413-17426, 2021.
11

---

Published as a conference paper at ICLR 2026
Beidi Chen, Tri Dao, Kaizhao Liang, Jiaming Yang. Zhao Song, Atri Rudra, and Christopher Re. Pixelated butterfly: Simple and efficient sparse training for neural network models. In International Conference Learning Representations (ICLR), 2022.
Bo Chen, Xiaoyu Li, Yingyu Liang, Jiangxuan Long, Zhenmei Shi, and Zhao Song. Circuit complexity bounds for rope- based transformer architecture. arXiv preprint arXiv:2411.07602 2024.
Lijie Chen and Ryan Williams. An equivalence class for orthogonal vectors. In Proceedings of the Thirtieth Anmual ACM-SIAM Symposium on Discrete Algorithms, pp. 21-40. SIAM, 2019.
David Chiang. Transformers uniform TC arXiv preprint arXiv:2409. 13009, 2024.
David Chiang, Peter Cholak, and Anand Pillay. Tighter bounds on the expressivity of transformer encoders. In International Conference Machine Learning, pp. 5544-5562. PMLR, 2023.
Rewon Child, Scott Gray, Alec Radford, and Iya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In ICLR. arXiv preprint arXiv:2009.1475 14794, 2021.
James Clift, Dmitry Doryn, Daniel Murfet, and James Wallbridge. Logic and the 2-simplicial transformer. In International Conference on Learning Representations, 2020. URL https
/ooeerrview n fo formm22d=rkecc66FVrV
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances neural information processing systems, 35: 16344-16359, 2022.
Giannis Daras, Nikita Kitaev, Augustus Odena, and Alexandros G Dimakis. Smyrf-efficient attention using asymmetric clustering. Advances in Neural Information Processing Systems, 33:6476-6489, 2020.
Jyotikrishna Dass, Shang Wu, Huihong Shi, Chaojian Li, Zhifan Ye, Zhongfeng Wang. and Yingyan Lin. Vitality: Unifying low-rank and sparse approximation for vision transformer acceleration with a linear taylor attention. In 2023 IEEE International Symposium on Highh-Perfommnce Computer Architecture (HPCA), pp. 415 428. IEEE, 2023.
Ronald B Dekker, Fabian Otto, and Christopher Summerfield. Curriculum learning for human compo- sitional generalization. Proceedings of the National Academy of Sciences, 119441)e22055821 19, 2022.
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Gpt3. int8 (): 8-bit matrix multiplication for transformers at scale. Advances in neural information processing systems, 35: 30318-30332, 2022.
Jacob Devlin, Ming-Wes Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume (long and short papers). pp. 4171-4186, 2019.
Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, Nanning Zheng, and Furu Wei. Longnet: Scaling transformers 1,000, 1,000,000,000 tokens. arXiv preprint 0rXiv.2307.02486. 2023.
Ran Duan, Hongxun Wu, and Renfei Zhou. Faster matrix multiplication via asymmetric hashing. In 2023 IEEE 64th annual symposium on Foundations of Computer Science (FOCS), pp. 2129-2138. IEEE, 2023.
12

---

Published as a conference paper at ICLR 2026
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Sean Welleck, Peter West, Chandra Bhagavatula, Ronan Le Bras, et al. Faith and fate: Limits of transformers on compositionality. Advances in Neural Information Processing Systems, 36:70293- 70332, 2023.
Mohamed EI Halaby. On the computational complexity of maxsat. In Electronic Colloquium on Computational Complexity (ECCC), volume 23, pp. 34, 2016.
Alhussein Fawzi, Matej Balog, Aja Huang, Thomas Hubert, Bernardino Romera-Paredes, Moham- madamin Barekatain, Alexander Novikov, Francisco R. Ruiz, Julian Schrittwieser, Grzegorz Swirszez, et al. Discovering faster matrix multiplication algorithms with reinforcement learning. Nature, 610(7930):47-53, 2022.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.11783, 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling, 2024. http:::tttpp:seeeeeeeee net wwt.:/uuteetttmmmmeeemmm forum?id
LEYskw1VY2
Michael Hahn. Theoretical limitations of self-attention neural sequence models. Transactions of the Association for Computational Linguistics, 8:156-171, 2020.
Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Simple on-the-fly length generalization for large language models. 2023.
Insu Han, Rajesh Jayaram, Amin Karbasi, Vahab Mirrokni, David Woodruff, and Amir Zandich. Hyperattention: Long-context attention in near-linear time. In The Twelfth International Confer- ence on Learning Representations, 2024. URL pettps://openrev/e net. / forum? iid=
EhCOd2BJIM
Yiding Hao, Dana Angluin, and Robert Frank. Formal language recognition by hard attention trans- formers: Perspectives from eircuit complexity. Transactions of the Association for Computational Linguistics, 10:800-810, 2022.
Pengeheng He, Xiaodong Liu, Jianfeng Gao, and Weizhu Chen. Deberta: Decoding- bert with disentangled attention. In International Conference on Learning Representations, 2021. URI
net www.nnee//oortieoee
Dieuwke Hupkes, Verna Dankers, Mathijs Mul, and Elia Bruni. Compositionality decomposed: How do neural networks generalise? Journal of Artificial Intelligence Research, 67:757-795, 2020.
Russell Impagliazzo and Ramamohan Paturi. On the complexity of k-sat. Journal of Computer and System Sciences, 02(2):367-372 2001.
Bart MP Jansen and Michal Włodarczyk. Optimal polynomial-time compression for boolean max csp. ACM Transactions Computation Theory, 16(1):1-20, 2024.
Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer: Fast transformers via sketches for polynomial kernels, 2024. URL
net./ http:/wwweeteeee/ee/// ?idd
YkCjojDG31
Tokio Kajitsuka and Issei Sato. Are transformers with layer self-attention using low-rank weight matrices universal approximators? The Twelfth International Conference on Learning Representations, 2024. URL htt.tppse//peanreveeeeeeeeeeeeee
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. International conference on machine learning, pp. 5156-5165. PMLR, 2020.
13

---

Published as a conference paper at ICLR 2026
Daniel Keysers, Nathanael Schärli, Nathan Scales, Hylke Buisman, Daniel Furrer, Sergii Kashubin, Nikola Momchev, Danila Sinopalnikov, Lukasz Stafiniak, Tibor Tihon, Dmitry Tsarkov, Xiao Wang. Mare van Zee, and Olivier Bousquet. Measuring compositional generalization: A comprehensive method on realistic data. In International Conference on Learning Representations, 2020. URL htttps://openreviie net htt..t/ooooomm//leescccccccceee
Najoung Kim and Tal Linzen. COGS: compositional generalization challenge based semantic interpretation. 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020. /
Alexander Kozachinskiy. Lower bounds on transformers with infinite precision. arXiv preprint arXiv:2412.20195, 2024.
Alexander Kozachinskiy, Felipe Urrutia, Hector Jimene7, Tomasz Steifer, Germán Pizarro, Matías Fuentes, Francisco Meza, Cristian Buc, and Cristóbal Rojas. Strassen attention: Unlocking compositional abilities in transformers based on a new lower bound method. arXiv preprint 19X1::2501.19215,1 2025.
Brenden Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of seqquence-t--sequeee recurrent networks. In International conference on machine learning, pp. 2873-2882. PMLR, 2018.
Andrea Lincoln, Virginia Vassilevska Williams, and Ryan Williams. Tight hardness for shortest cycles and paths in sparse graphs. In Proceedings of the Twenty-Ninth Annual ACM-SIAM Symposium on Discrete Algorithms, pp. 1236-1252. SIAM, 2018.
Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. In The Eleventh International Conference on Learning Representations, 2023a. URL htR.httss: http::::///peeeeeeeeeeeeeeee
net/ weeee//orummmddDeFFFPaom.
Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Re, et al. Deja vu: Contextual sparsity for efficient Ilms at inference time. In International Conference on Machine Learming, pp. 22137-22176. PMLR, 2023b.
Pan Lu, Baolin Peng, Hao Cheng. Michel Galley, Kai-Wei Chang, Ying Nian Wu, Song- Kong-Chun Zhu, and Jianfeng Gao. Chameleon: Plug-and-pla compositional reasoning with large language models. Advances in Neural Information Processing Systems,
2023.
Zixian Ma, Jerry Hong, Mustafa Omer Gul, Mona Gandhi, Irena Gao, and Ranjay Krishna. Crepe:
Can vasion-language foundation models reason compositionally? In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10910-10921, 2023.
Eran Malach. Auto-regressive next-token predictors are universal learners.
arXiv preprint
arXiv:2309.06977 06979, 2023.
Gary Mareus. Deep learning: A critical appraisal. arXiv preprint arXiv: 2018.
William Merrill and Ashish Sabharwal. A logic for expressing log- log-precision transformers. Advances in neural information processing systems, 36:52453 52463, 2023a.
William Merrill and Ashish Sabharwal. The parallelism tradeoff: Limitations of log-precision transformers. Transactions of the Association for Computational Linguistics, 11:5: 11:531-545, 2023b.
William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought. In The Twelfth International Conference on Learning Representations, 2024. URL https
/ooenreview n worumiiddoooooo..o. N
William Merrill, Ashish Sabharwal, and Noah A Smith. Saturated transformers constant-depth threshold circuits. Transactions of the Association for Computational Linguistics, 10:843-856. 2022a.
14

---

Published as a conference paper at ICLR 2026
William Merrill, Ashish Sabharwal, and Noah A. Smith. Saturated transformers constant-depth threshold circuits. Transactions of the Association for Computational Linguistics, 10:843-855 856, 2022b.
OpenAl. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
OpenAl. Introducing openai ol ol-preview, 2024.
URL https: openai.com/indee.......e
int aatrodtciig--opeai-ooo ol-peviewww
Matteo Pagliardini, Daniele Paliotta, Martin Jaggi, and François Fleuret. Faster causal attention over large sequences through sparse flash attention. arXiv preprint 11X.0:2306.01166, 2023.
Binghui Peng, Srini Narayanan, and Christos Papadimitriou. On limitations of the transformer architecture. In First Conference on Language Modeling, 2024.
Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing-complete. Journal of Machine Learning Research, 22(75):1-355 2021.
Sergio Perez, Yan Zhang, James Briggs, Charlie Blake, Josh Levy-Kramen Paul Balanca, Carlo Luschi, Stephen Barlow, and Andrew Fitzgibbon. Training and inference of large language models using 8-bit floating point. In Workshop on Advancing Neural Network Training: Computational Efficiency, Scalability, and Resource Optimization
2023), 2023. URL https
//openreview net. forum? d=nnrrrDDerrrrr
Zhen Qin, Weixuan Sun. Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. In International Confer- ence Learning Representations, 2022. URL httptttt//:stt https:://ppereviii B1 B18Cgrx2Up4.
net. fwrmm?ormmmmmmm
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53-68, 2021.
Aurko Roy, Timothy Chou, Sai Surya Duvvuri, Sijia Chen, Jiecao Yu, Xiaodong Wang, Manzil Zaheer, and Rohan Anil. Fast and simplex: 2-simplicial attention in triton. arXiv preprint arXiv:2507.02754. 2025.
Aviad Rubinstein. Hardness of approximate nearest neighbor search. In Proceedings of the 50th annual ACM SIGACT symposium on theory of computing, pp. 1260-1268, 2018.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. One-layer transformers fail to solve the induction heads task. arXiv preprini arXiv:2408. 14332, 2024a.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. Advances in Neural Information Processing Systems, 36, 2024b.
Koustuv Sinha, Shagun Sodhani, Jin Dong, Joelle Pineau, and William L. Hamilton. CLUTRR: A diagnostic benchmark for inductive reasoning from text. November 2019.
Volker Strassen. Gaussian elimination is not optimal. Numerische mathemarik, 13(4):354-355 1969.
Xiao Sun, Jungwook Choi, Chia-Yu Yu Chen, Naigang Wang, Swagath Venkataramani, Vijayalak- shmi Viji Srinivasan, Xiaodong Cui, Wei Zhang, and Kailash Gopalakrishnan. Hybrid 8-bit floating point (hfp8) training and inference for deep neural networks. Advances in neural information processing systems, 32, 2019.
Zhiqing Sun, Yiming Yang, and Shinjae Yoo. Sparse attention with learning to hash. In International Conference Learning Representations, 2021.
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention for transformer models. In International conference on machine learning. pp. 10183-10192. PMLR, 2021.
15

---

Published as a conference paper at ICLR 2026
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez. Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems. 30, 2017.
Sinong Wang, Belinda Li,, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprins arXiv:2006.04770 2020.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems, 35:1 15:12071-, 2022a.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, al. Chain-of-thought prompting elicits reasoning large language models. Advances in neural information processing systems, 35:24824-24837 2022b.
R Ryan Williams. Algorithms and resource requirements for fundamental problems. Carnegie Mellon University, 2007.
Ryan Williams. A new algorithm for optimal 2-constraint satisfaction and its implications. Theoretical Computer Science, 348(2-3):357-365, 2005.
Virginia Vassilevska Williams, Yinzhan Xu, Zixuan Xu, and Renfei Zhou. New bounds for matrix multiplication: from alpha omega. Proceedings of the 2024 Annual ACM-SIAM Symposium Discrete Algorithms (SODA), pp. 3792-3835. SIAM, 2024.
Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. International Conference Machine Learning, pp. 38087-38099. PMLR, 2023.
Ofir Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. Q8bert: Quantized 8bit bert. In 2019 Fifth Workshop on Energy Efficient Machine Learning and Cognitive Computing-NeurlPy Edition (EMC2-NIPS), pp. 36-39. IEEE, 2019.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283-17297, 2020.
Amir Zandieh, Insu Han, Majid Daliri, and Amin Karbasi. Kdeformer: Accelerating transformers via kernel density estimation. In International Conference on Machine Learning, pp- PMLR. 2023.
Aimen Zerroug, Mohit Vaishnav, Julien Colin, Sebastian Musslick, and Thomas Serre. benchmark for compositional visual reasoning. Advances in neural information processing systems, 35: 29778, 2022.
Zhongwang Zhang, Pengxiao Lin, Zhiwei Wang, Yaoyu Zhang, and Zhi-Qin John Xu. Complexity control facilitates reasoning-based compositional generalization in transformers. arXiv preprint arXiv:2501.08537, 2025.
CONTENTS
1 Introduction
Running time considerations
Poly-attention is all you need
2 The poly-attention mechanism
16

---

Published as a conference paper at ICLR 2026
3
Beyond self-attention: the power of poly-attention
An example: function composition
3.2
Tree-attention: polynomials leading to efficient poly-attention
3.3
Computational complexity of non-tree poly poly-attention
3.4
Representational strength of poly-attention)
3.5
Implications of poly-attention
Technique overview
5
Experimental validation
6
Acknowledgements
10
Ethics statement
10
8
Reproducibility statement
10
Preliminaries
18
Roadmap
18
A.2 Notation and background
18
A.3
Conjectured hard problems
20
Related works
20
Warm-up: Strassen-attention upper and lower bounds
21
Algorithm for Strassen-attention
................................
21
Hardness of Strassen-attention
26
................................
C.2.1
Conditional hardness Gap-
................................
26
2,2
Hardness of approximating btrassen-ateentin
29
D Proof's of Section 3.2) tree-attention
31
Proofs of Section 3.3; computational complexities of poly-attention
34
Polynomial method for poly-attention
35
Time lower bounds for poly-attention
37
E.2.1
Time lower bounds based on degree of polynomial using SETH
38
E.2.2
Time lower bounds based on substructure of polynomial using Max-ASAT conjecture
E.2.3
Time lower bounds for degree 2 polynomials using conjecture
48
Proofs of Section 3.1;
function composition
52
Proofs of Section 3.4
polynomial root-finding
57
17

---

Published as a conference paper at ICLR 2026
H Experimental details
59
Function composition
59
H.2 COGS Dataset
62
The use of Large Language Models (LLMs)
63
A PRELIMINARIES
A.I ROADMAP
In the rest this paper, we prove all the results that we have stated in the main version. After describing some relevant notations and conjectures that will use, prove the results in two parts. First we prove the computational complexities of the poly-attention mechanism, followed by proofs of representational strengths. The proofs of computational complexities use subdivisions- an upper bound where we show that if the entries of the query-key matrices are bounded, then we can compute an entry-wise approximation in near-linear time, and lower bound where we show that if the entries of the query-key matrices are large, then assuming certain fine-grained complexity conjectures, computing entry-wise approximations are difficult.
warm-up, we start with upper and lower bounds for Strassen-attention (Section C based which, we proceed to prove the same for general poly-attention Section
order to completely characterize time complexities for poly-attention, we also give quadratic time algorithms for tree-attentions in Section D
The proofs of the results stated in the main paper are given as follows:
For Theorem 3.1
poly-attention can simulate function-composition has been proved in 2, and the time complexity of (nn²) has been proved in Theorem
Theorem F.6 or
D.2
Theorem
3.2
Part has been proved in Theorem E.2 and Theorem
3.2 Part 2 has been
proved in
heorem
E.3
Part 1.
Theorem
3.3
has been proved in Theorem F.3 and Corollary
E4
Theorem
3.4
has been proved in Theorem F.6
Theorem
3.5
Part has been proved in Theorem D.2
Theorem 3.5
Part 2 has been proved in
Theorem
E.
and Theorem 3.5 Part 3 has been proved in Theorem E.3
Theorem
3.6
Part has been proved in Theorem and Theorem 3.6 Part has been proved
in Theorem
Theorem
3.7 has been proved in Theorem G.1
A.2 NOTATION AND BACKGROUND
Throughout this article, for natural number n we denote [n] as the set 21,2,..., n2), [i : as the set of integers {i, + .,....) ..} for i < J, and [i,j] as the set of real numbers between and Given a matrix M E ma for E [n] E [m], we denote [M],.j and more loosely My,j. as the (i. j-th entry of the matrix, M; as the i-th transpose of nz-dimensional vector. M(sssh will also denote the submatrix M having rows [i : fi] and columns [i2 j.].
as m-dimensional vector, and M., as the j-th column the
For two matrices A, B E R=xm, we define B as the entry-wise division, i.e., [1 = Art. Given B
vector E Rax1, by diag(X), we denote the n diagonal matrix such that [diag( X)) X
for all i E [n). Some other operators matrices are defined as follows.
Definition A.1 (Hadamard product (), Given to matrices A, B E Raxm we denote the Hadamard product, denoted by A & paxm as the entrywise product
[A B)b Ai. Bi.j
for E [n],
18

---

Published as a conference paper at ICLR 2026
Definition A.2 (iiw Kronecker product 2). For matrices A E Raxd B Rm we denote the row-wise Kronecker product as A 2 B E Romxd where
[A 2 B (i(--))+ = A, Bj,
for E [n],j [m]]
Definition A.3 (Entry-wise approximation). Given matrix M Raxdd say that M is eniry-wise praproximaioooooonoon of M for all E [n]:3 E [d], we have
MM. - Masl <.
Throughout this paper, we will choose 1/poly(n).
Definition A.4 (Entry-wise function). Given a function f IR - and a matrix M E Raxm, we define the matrix [M] as the n x matrix such that the the (i, element is
pM = f(Ma.)
We will use [M] as the entrywise exponentiation function, i.e., [M] = Mis, For a real number c, M/c will also refer to the matrix obtained by dividing each entry of M by
The coefficient of matrix multiplication, W, roughly refers to the exponent of such that two series of works trying matrices can be multiplied in time O(n²) for large enough n. There
to improve this coefficient (flfaiiientAAmaann [Alman & Williams (2024): Duan et al. (2023); Williams et al. (0024); Fawzi et al. (2022); Alman et al. (2025), with the fastest being Alman et al. (2025) that achieves However, these matrix multiplications require be quite large and the hidden constants are enormous, which does not make implementations feasible. There is an algorithm by Strassen Strassen (9699) which is more practicable and achieves w = 2.8, but in most cases, only the naive matrix multplication algorithm is used as GPUs work better on them.
We will use some more concepts define the ideas in this article. Given an integer t, we define the symmetric group of order r, ((.)), as the set of tuples:
( (
= f2 1011 S < S..
Note that ((9)) variables is defined as
())
Based this, elementary symmetric polynomial of degree having /
Σ
Ln Z...j.
Definition A.5 (Variable separability). We say that a polynomial h(z1....,, ..,) ) is variable separable if there exists maximum integer r and non-zero attention polynomials gı ,,.,. (2 x²)) where x²1 "2....,2" xt are disjoint subsets of the variables, such that h( 22...., = g1 x1) ... gr (x),
We denote each of the polynomials a+t) branches.
Note that this definition of variable separability differs slightly from the folklore usage as here allow f and to share at most one variable, T1.
In this paper, for given polynomial h, interested in computing the entry-wise approximation of Ao) For this, we define the following version of computing poly-attention approximately.
Definition A.6 (Entry-wise Approximate Poly-Attention Computation APAC (h) (n, v)), Let be an attention polynomial in variables of degree k having sparsity Given query-key matrices Qua),...,Q QQt)) E [-- and value matrices (22 ....y (()) E Raxd we want to output matrix
Att(h) E B such that for all i E [n],j E [d].
[[A ( ,.... Q(1), v(2) ...., v(0))a S7
19

---

Published as a conference paper at ICLR 2026
A.3 CONJECTURED HARD PROBLEMS
We define some commonly known problems in fine- tne-grained complexity and state conjectures which will be used to show conditional hardness of generalized attention computations. First, we start by defining few central problems in fine-grained complexity.
Definition A.7 (kIP problem). Given sets of vectors A ,t...., A S (0, 134, each of size and target inner product m [d], the problem of klP asks if there exists a E A a² E A A ,.....a* a A
such that (a a² ................................ a²)
For k = and m 0, it is the famous orthogonal vectors problem, which we will abbreviate 20V or just the OV problem, and for k 2 and arbitrary m, we will abbreviate the problem as IP.
Definition A.8 (T: In the kSAT problem for Z 2, given as input a k-CNF formula Φ, determine whether not & has satisfying assignment.
Definition A.9 (Max-kSAT). In the Max-kSATn.w problem for 2 2, given input k-CNF formula & in n variables and m clauses, determine the maximum number of clauses I that can be simultaneossyy satisfied by Boolean assignment to the underlying variables.
Based on these definitions, complexity that we will use to prove our (conditional) hardness results.
ready describe popular conjectures in fine- fine-grained
Hypothesis (SETH Impagliazzo Paturi (2001).. For every > O, there exists & 2 such that kSAT can not solved in time O
The current fastest known algorithm for 1SAT uses the reduction to OV with dimension log n..
The best known time complexity of OV is 2 (0 given by Abboud al. (2014a); Chan & Williams (2016).
Since ASAT special case of Max-kSAT, SETH implies that Max-kSAT also cannot be solved in time ((2(1-6)m) for every 8 O. The next hypothesis Alman Vassilevska Williams (2020) strengthens this further to sparse instances of Max-kSAT.
Hypothesis 2 (Sparse Max-kSAT Hypothesis). For every k 2 3 and every ō > 0, there exists e >> 0 such that Max-kSAT, n.en cannot be solved in time (22(1-))))
The fastest known algorithm for sparse instances of Max-xSAT n.on for h 2 3 takes time 2m(1-1/0("//)) Alman et al. (2016); therefore the above hypothesis is consistent with the state- of-the-art algorithms. In contrast to the special case of Max-kSAT for = 2, the hypothesis is false. The best algorithm for Max-2SAT Williams 22005 2007) runs in time 2ool"p) where w the matrix multiplication exponent. The tollowing Max-2SA hypothesis states that William's 's algorithm Williams 2005) is essentially optimal when h: 2.
Hypothesis 3 (Max2SAT hypothesis). For every > O, there exists a a 0 such that Max-2SAT n.c camnot be solved in time 0(2m(w/3-6)), where us is the matrix multiplication exponent.
The following theorem gives a reduction from &SAT to kIP, thus proving the hardness kPP under SETH.
Theorem A.10 (Williams (005); Abboud al. (2014b); Backurs & Indyk (015); Abboud et al. (2015)) Assuming SETH, for every k and d >> 0, there exists e >> O such that KIP n. cannot be solved in time (1-66))
B RELATED WORKS
The most similar prior works on attention mechanisms which are more expressive than self-attention are Sanford et al. (2024b) and Kozachinskiy et al, (2025), which we have already discussed in detail. There another attention mechanism, triangular attention, introduced by Bergen et al. (202T), whose design was inspired by logic programming, and which shown to perform better than self-attention on certain compositional tasks. However, Kozachinskiy et al. (2025) proved that cannot perform function composition.
20

---

Published as a conference paper at ICLR 2026
As we have discussed, the self-attention mechanism (Vaswani et al., 2017) is at the center of all large language models because of its expressivity in real-life applications, but the quadratic time complexity for computing its output is sometimes already prohibitively expensive. One extensive line of work has introduced faster heuristic algorithms, which work well on many inputs. These have used different approximation techniques, including sparsity assumptions, norm bounds, and kernel
density estimation (Zandich et al.] 2023 Han et al., (0044) Kitaev et al. (2020) Choromanski et al.
202T
Pagliardini et al.
20233 Child al. 2019; /Wan al. 2020; Daras el al. 2020] Katharopoulos
et al.
3020:/henn et al.
2021
2022 Qin et al.
2022
Liu et al. 2023b He et al. [2021] Kacham et al.
2024
Dao et al., 2022
Dao
2024,
Roy et
202 Sun al. 2021 Ding et al. 2023 Han et al.
2023
Zaheer et al. 2020
Dass et al. 2023))
Other alternatives have been considered which completely replace attention with different mechanisms. A simple example is Hardmax attention, in which the softmax is replaced by a (hard) max, but training Hardmax attention Transformer models appears difficult as we do not know an efficient way to perform gradient descent. The power of hardmax has been explored in Alcalde et al. (2024); Pérez al. (202T): Kajitsuka & Sato (2024). Instead of computing the output of self-attention faster, some other alternatives to Transformers have been proposed that completely replace attention with other mechanisms; examples include Synthesizer (Tay et al.) (021), routing Transformers (Roy et al.) 2021), and Mamba (Gu & Dao 2024). These alternatives can typically be computed much faster than attention (often in almost linear time by definition), but in exchange appear to have weaker expressive power (Alman & Yu 2025). This paper continues a long line of work on understanding the power and limitations of Transformers, and finding more expressive alternative models.
Some papers have studied the circuit complexity of Trransformers (Chiang 2024 Merrill & Sabharwal
2023a Merrill et al. 20226; Chen et al., 2024; Merrill & Sabharwal| 2023b Merrill et al.) 2022a
Chiang et al. 2023). Other works on the representational strength of Transformers focus on their relationship with other models of computation. For example, line of work has studied the ability of Transformers to approximate other models of computation (Pérez et al.) 202T Wei et al. 2022a
Malach, 2023 Liu et al. 2023a Hao et al.) (022). On the other hand, there are many more tasks, beyond those discussed here, which are difficult to solve by a Transformer, including compositional
reasoning (Dekker et al., 2022 Zerroug et al.) 2022 Marcus 2018 Kozachinskiy 2024 Sanford
et al.
2024a Peng et al. 2024).
Another approach to overcoming the limitations of Transformers is augment them in other ways.
An important example is cf (Wei et al.] 2022b). |Mrrill & Sabharwal (2024) studied the space and time complexity of chain-of-thought, and Peng et al. (2024) studied how this relates to function composition.
C WARM-UP: STRASSEN-ATTETIOON UPPER AND LOWER BOUNDS
As warm-up, we describe the polynomial method and show Max-2SAT-basee hardness results on Strassen-attention Since Strassen-attention is only special case of poly-attention, we will later move on to show similar algorithms and lower bounds on poly-attention in Section
C. ALGGORITHM FOR STRASSEN-ATTENTIE
In this section, give near-linear algorithm for computing entry-wise approximation of the output matrix of Strassen-attention when the entries of the query-key matrices are bounded. We will use the polynomial method, which has been used in entry-wise approximations of other attention mechanisms as well, like in self-attention Alman & Song (2023). tensor-attention Alman & Song
(0024), RoPE based attention Alman & Song (2025).
Our goal is to compute the x d matrix (tt(S), the output of Strassen-attention, for query- key matrices Q()), ((2) ((3) E [-r, r axd n and value matrices (LL), ((2) E R=xd Using the expression of Strassen-attention Kozachinskiy et al. (2025), it can also be written as F D (1 (T (tt = (Q 1m))
(2)
(imms) (Q ( QQe)) ( (()m)) la
for all i E [m], l e [d], where D ((()) ()) and D diag(v (Lm.e))
21

---

Published as a conference paper at ICLR 2026
We will compute the entry-wise approximations of the numerator and the denominator terms of Equation 2 separately. The main idea is to use low rank entry-wise approximations for each of Q1 (2))) and multiply the low rank matrices together-
(Q .I (2 something that can be done more efficiently. In order to obtain the low rank approximations, we will use the following lemma from Aggarwal & Alman (022)).
Lemma C.I Aggarwal & Alman (002 Let r > 1, 8 (0,0.1). There exists polynomial P(x) R[x]] of degree == θ max log(log(4.)/T)
r}) such that for all E [-r,r], we have IP(a) e < se". Furthermore, P can be computed in poly(t) time and its coefficients are rational numbers.
Using the previous lemma, we obtain the low rank matrix approximations as a corollary.
Lemma C.2 (Low rank approximation Alman & Song (2023 2024 ). Let 1/poly(n),
(log n), no(1), and B o(log n). Given matrices P. Q p-r.rjnx. compute
matrices U. W E Rexr in time O(n +)()) such that UWT entry-wise e-approximates PQT, that
is: [UU k5 [P I[
This is an instance of the Gaussian KDE which has widely been used in LLMs and machine learning algorithms Zandieh et al. (2023); Backurs et al. (2018); Katharopoulos et al. (2020): Alman et al. (2020): Aggarwal & Alman 2022 Alman & Song 2203 (2023|2024)
We will show that we can compute E [n], 7-approximations of denominator in Equation 2 in time (n)+)(()) and fixing any L E [d]. we can compute 7-approximations of the numerator in time O (oot)) Vi [n]. where ; = 1/poly(n). Once we find the values of the denominator and numerator, we perform division, to compute the apappoxiiatioon Att(S) 3,é, which takes a total time of (o11+ +oo1) + nd) = On)))))) Using this as the central idea, we prove the following result. Since Strassen-attention is special case of poly-attention with the polynomial (x1, x2, z3)= = 21 + T2Zy + Ty31, we state the following result:
(hs) Theorem C.3. There is an algorithm that solves APAC (n, d = 0(log), oVVogg 1). 1/poly(n)) with query-key matrices Q()))((2) ((2))((())) E - --r. Pjxd, and value matrices (22), v(t) Iaxd in time O (oo1))
The algorithm is summed up follows.
22

---

Published as a conference paper at ICLR 2026
Algorithm 1 Algorithm to compute entry-wise approximation of Att(S)
Input: A number I Vlooo n), query-key matrices Q(1,, Q(2) Q() E e[-r, r] ces Va), ((2) Rexd an approximation parameter 1/poly(n).
value matri-
Output: Entry-wise 7-approximation Att(S) of Att(s) I: Initialize Att(S) := Onxd-
2: Compute the low-rank 7-approximations (W of [2"")))))
(W
of (Q and (W of (QU)) using Lemma
C.2 where
U W U2 W U3 W E RRaxr for =n(())
O(n +o Dr) time.
3: D == diag(V (Im.0)) D == diag(V (1m.e))
b O(n) time.
4: Compute ú²2 D W2 == D2 E gaxr.
O(rr) time.
5: Compute A UW (W3)T U and B = (2 (W2)T
b O times.
rxy
rxr
6: for E [n], E [d] do
7: Compute the approxxaioiii of the denominator (Equation 2 as
R = :== A(W (iiss)) E IR.
> 0(r3) time.
8: Compute the ((7) approxxtttina of the numerator (Equation 2 as
pf = U(...) B(W()..))) E IR.
> 0(r2) time.
9: Compute the L-th row of the entry-wise θ(7), approximation of (tt
p
Att(S) i 8
Q
D 0(1) time.
10:
end for
11:
Return (t)(s)
Before proving the correctness this algorithm, we first show that the entries of the exponentiated matrices are bounded, which is necessary for applying Lemma C.2
Lemma C.4 (Bounded entries). The entries of ( TT are bounded as
r I QQ 10 1 1 Se
for all i.j [n].
Proof. Without loss of generality, we prove the upper bound only for X and the rest follows similarly. Since each entry of Q()), Q(2), ((3) are [-r, r). the value of (Q(2))) is
(02)) 10.")000
Σ (12)
qif
EE(d]
r s (0 2000) / //d r (Since - s (11) 0 S))
Therefore, - QQ s r² for all i,j E [n]] and we similarly bound (Q())) (
Now, we prove Theorem C
which is also the correctness of Algorithm
23

---

Published as a conference paper at ICLR 2026
Proof of Theorem
C.3
First,
compute the low-rank approximations of using Lemma
(Q
H (2) (
C.2
(Step of Algorithm
D.. However, in order for Lemma C.2 to succeed in Step of Algorithm
we need the entries of the
exponentiated matrices to be bounded, which is true due to Lemma
We compute the Strassen-attention matrix in two steps, first computing the denominator, and then the numerator in Equation 2 to compute the entire self-attention matrix.
Computing the denominator. This has been described in Step 7 of Algorithm I and we prove its correctness. Since the entries of [1Q4) (Q()))J) [ [a are bounded, we can apply Lemma C.2 find their low rank approximations. approximations of (Q))) ( [20(3) be U UWW1)T, (2 and (W respectively, with entry-wise error for s 1/poly(n), where each of U W E Rexr Namely, for all i, [n]]
( the low-rank
(OOK, U0" (W!)
< 0 </
0 - 102(12) bI
<%
(3)
1 -
(W3
< < 0
<
where y = se.". When we choose E as the inverse a large enough polynomial such that ee!
poly(n): we have 7 1/poly(a) (note that O(vlog Now, claim that
an approximation of (ttp:iiniiiiooloo000009000110010 |1"""""0""11110000"""""""""""""" For ease of notation,
let us denote X = === (Q) Now,
XXYZ), (U,,-UUUV((U(00
-IIX22)a- + -
+ (UU
(4)
< x - (w)) + -wwp)//""Y..-V""/"VVV"""
+ [ 21, wwwp:///wwwwwaw/ww""""/"//"""""
where the last inequality follows from triangle inequality.
Now, using Equation 3 in each of the three terms, we can show that this is bounded above by O(7).
The computation of ((n))))) time for = ()) (which is true for the choice of d = O(log n). B = 7 1/poly(n) using the parameters of Lemma C.2
(
for i E [n] from Algorithm | takes
Computing the numerator. An entry-wise aaapprxximation of the numerator of the Strassen- attention matrix (tt(5) e Ruxd (Equation 2 has been computed in Step 8 of Algorithm Here, essentially, we compute each entry D & Y 2 Z),,. for all E [r]] by fixing E [d] at a time.
We again make use of the low rank decompositions of X, Y, Z as above (Equation 3p. Note that the value of each element of (S) is given as
Au! = [XD>
24

---

Published as a conference paper at ICLR 2026
We claim that IU' www"yyy--tPPWyy---00wwy is appapoimmatiioononn of
IXDI.ED' k) Indeed, we have D - (W
+ w|wp://wwwuuuuooaa...aaao//////0 (5)
<
+ (v))))
+
which again follows from the triangle inequality, and each term can be shown to upper bounded by O()) using Equation 3
(S) is obtained by approximating the value D2 b) and then dividing by the approximate value of [XYZ]i,, Using
Wrapping up.
An approximation of the (i, f)-th element, Au
P U (w D D
R4 U
U. (WW))?2W3))' (W3)T,
have
[XD 1,. P. 00)),
and,
[X Ri S 50()),
for E [n], E [d].
Therefore, the error is given by
(XD - $
[XX 7), [X
+ ||X - -R (Triangle inequality)
< <(()),
which follows from Equations
5
repeated applications triangle inequalities, and the fact that & is
an inverse polynomial in n, and,
[XD Zaxl
Σ
X V.
| 11111///
j.kE(m]
1
[XY ΣjAeEn) Y,KYY,.KKK..
for all [n),e E [d] since the entries of Q(1,, ((2))2Q())))) are in [-r,r] (Lemma C.4 For
d Ologgn,, I' n() and [|) || poly(n), we can choose 70 1 poly for a large enough polynomial such that
P
Au(
< 70
R
where = 0(7) 1/poly(n), which is our required approximation parameter.
As described Algorithm 1 we compute this appproxxiattion for all E [n] in time O(n O(n!-o(1)) l-(()) and repeating this over all e [d] requires O(n' +o(1)) = O(n (t(()) time since = O(log n). This proves Theorem
25

---

Published as a conference paper at ICLR 2026
C.2 HARDNESS OF STRASSEN- ATTEENION
Now, we introduce the techniques that will be used to prove lower bounds in this paper. We establish the hardness of Strassen-attention in the high weight case, assuming the Max-2SAT conjecture (Hypothesis 3). Our reduction will proceed in three steps. First, we use a reduction from Alman
Vassilevska Williams (2200) that establishes the hardness of IP (Definition C.5) assuming Hy- pothesis 3 (hardness of Max-2SAT). Second, we prove the hardness of e- e-Gap-IPA (an approximate version of IPA defined in Deñiiiton (.6) from the hardness of I,A, in Section C.2.1 Lastly, we
prove the hardness of Strassen- attention from the hardness of e- G-Gap-IP A in Section
C.2.2
We begin by defining the problems IPA and G-a--IPP
Definition C.5 (IPA). Given three sets of vectors A A?, A C {0.1]4 |A 1A21 = [A
TI, and target inner products, m12, 1223, m331 E [o....,d). d}, the problem iPA.,a A2, A
A9, mm2, 1223, m311 asks whether there exist vectors 01 E A 02 A A2203 A such that, simultaneously, (a1, (2)) = M12, (a2, 0z) = m22, (a3, (1),01) = 1331-
Definition C.6 (s-Gap-IPA). Let E > O. Given three sets of vectors A A2,A2, c 134, with |A4] LA n. target inner product m {0....,dl d], and the promise that for every 03 E A
|A'|
a1 E A Q2 E A2.03
either (01, 02) S (1 =)m or (a1, 02) = m.
and, either (az, 0s) S (1 - e)m or (02, as) m.
and, either (az, a1) S (1 - m)m or (0z, a1) = m,
the problem Gap- ad (11 A2, A3, m) is to decide there exist vectors a1 A ,02 E A A?,ag ag A3 such that: (01, 02) (02, as) (as, a1) m.
For IPA and Gap-IPA, will drop the parameters m, d when they are clear from context. Note that even though IPA might have different inner products for all three pairs, for E-Gap-IPA, the three inner products being equal suffices as the reduction for proving its hardness accommoaates this property, and for proving hardness of approximating the output of Strassen-attention, we need them be equal.
As mentioned above, the first step uses a result due to Alman & Vassilevska Williams 2020) which proved that IPA is at least as hard as Max-2SAT:
Lemma C.7 (Alman & Vassilevska Williams 2020)). Assuming the Max-2SAT conjecture (Hy- pothesis 3 for every ô > O there exists c such that clon, cannot be solved in time O
C.2.1 CONDITIONAL HARDNESS OF GAPP-IPP
In this subsection we prove the following theorem, establishing hardness of e Gap-IIP assuming hardness of IPA.
Theorem C.8. For every 6, e > O, there exists c, > such that if G-aa--PPAPPAAA n,e log can be solved in time (n( then IP IPAn,e' log can be solved in time ((n(w-6/2),
Building on Aaronson & Wigderson 2009), Rubinstein Rubinstein 2018) gave reduction from the IP problem to the gap version, 8- Gaa--IP. That is, they proved similar reduction to what we want, but where IP and e- e-Gap-IP take as input two sets A². A instead of three sets. Chen & Williams (2019); Abboud & Ron-Zew, (2025) further improved their reduction; for our reductions, we will use and build upon the proof given by Abboud & Ron-Zewi 2025).
The following lemma was proven Abboud & Ron-Zewi (025) (see the proofs of Lemma 4.1 and Claim 4.3 in their paper).
Lemma C.9 Abboud & Ron-Zewi 2025)). For all n, d = O(logn) n), there exists d' = O(d), (1),, m' O(logn) n), such that for every instance of IP n.d given by sets of vectors A, B and target inner product
{O,...,d) dd, there set of q instances {(A', B mn) | [4]} of -Gap-IP n. computable in ((1++0(1)) time, where e E (0, 1) is constant such that:
26

---

Published as a conference paper at ICLR 2026
1. (Yes case) there exists (a. A B such that (a. m, then there exists i [g] such that (A', B', m') is yes instance of Gap-IP n.d'
2. (No case) for every pair (a,b) b) A x B, (a, + m, then for all i [g]. (A', B*, m') is no instance of e- Gap-apPP,,,. n.d'
Proof of Theorem C.8 We start with an instance of IP m, and matrices A. B, C each of dimension d, where the rows of A correspond to set of vectors, and similarly for B and C.
n). given by a target inner product
For the pair (A, B), apply Lemma C.9 create set of q instances of Gap--I p. n.d each with target inner product m':
(AAB. BAB) {(A' B 1 [q]}.
Similarly we apply the Lemma to the pair (B, C) to get G-Ga- instances
(BBc, CBC) (B CBC) [d]}
and to the pair () to get instances
(AAC, CAc) = ((A'Ac Cac) i [d]}
By Lemma C.9 the following properties are satisfied by (AAB, BAB):
(I) For all i E [9]. the instance (A'A B'B m') satisfies the gap property. That is, for every
dAA E A' BAB BB (A'AB. bAB) is either equal to m' or is at most e)m'.
(2) Correctness of the reduction:
(2h) there exists ba E A x B such that (a, b) m. then there exists i E [g]. and vectors a'An E A' bAn E Ban such that (a'an Aa = m'.
(2c) for every (a, b) E A x B, (a, b) + m, then for all E [d]. and for all vectors a'nn E b'An E Ban we have (a' b'A) S e)=')m².
By the same argument, the above two properties also satisfied by (BEc, Čεc) and (Ārc, CAC). Equipped with the above pairs of 3-dimensional tensors, we are now ready to describe our reduction from the instance (A, B, C) of IPAn,d to set of q² instances of E-Gap-IPAL IPA, nlo by:
denoted
(A, ((At.i Biihk , jí [q]}
For each i,j, E [g]. we define A to consist of the following set of length 3d' vectors:
A = {a'o 0f a'ic l''D E A'sn anc e Ahcc
Similarly we define Bisk and Cij.k:
Biiik = (b'A UgC 0 BAB E B' UPO e Block
Chi.k = 10 cgc cic cec E Ebcc. dio E Öhic
Gap Property First, prove that every instance (Abiik Bisck Chiik) satisfies the gap property. Consider generic triple (a'jik, biick c.o.k) E Aith x Biish x Chick, where
a' d d'c.
27

---

Published as a conference paper at ICLR 2026
blin bs B'AB
ddi = 0 CBC
Since (atJik bl..kk = (a'A b' AAB) we can apply property (1) to (AAB, BAB) to infer that this inner product is either m' or at most (1 e)m' By a similar argument we can show that (bAS cf...i) and (arsah ch..k) are either m' or at most (1 e)mmm This completes the proof of the gap property.
Proof of Correctness. We first consider the yes case, when there exists (a, b, c) A B x C such that (a, b) (b, c) (a, c) m. Applying property (2a) above, we have:
1. There exists i [a], a'AB ABA VAB E BAB such that (a' BAB) ml.
2. There exists j E [gl. bBB BB ee E ÖoCc such that (BBO 2Bc) m'.
3. There exists is E [4). alcc E Alc d'ec Et such that (a cac) m
Now consider the corresponding vectors aisik E Atiik, Biikk Biikk, and cfdck in Chick defined as:
atj.k = d'AB0" a'cc
b/.h An B
d = 0 égc
By inspection together with the above three properties (1, 2, 3), we have (a hiih c.i.tyy = m'.
(at Bl..-e (br.ikk, dudlky
thus completing the "yes" case of correctness.
In the no case, suppose for all (a, b, c) A B C, either (a, # m or (b,c) + or (a, c) of m. We want to show that for all i, & E [g] and for all (a-s.k Biakk chJ.it) E A'iik x Biskk x Chick at least one of the following holds: (i) (atd.k, birth I (1 E)m' or (ii (bisik, chi.k) (1 s)m' or (iii) (a (ar chi.k) < (1 - e)m'.
Fix j, k [q] and consider generic triple (ats.k, b...kk c.o.k) in Aisik Biikk x Chick, where
ahi.k a' d'AD0" ate
bhih 1/
dsh 20
Consider first the where (a, b) * m. Then by applying property (2b) to (ĀAB, BAB) have (a'AB, b'AB) S =)m', and therefore (a-iikk i.kk bio.ky & (1 =)mm,, so case (i) above holds.
Similarly in the second case where (b, c) # m. applying property (2b) (Bnc, Čnc) it follows that (bi-i.k cl...k S (1 E)m', so case (ii) holds. For the last case where (a, e) # m we can similarly use (2b) to show that (iii) holds.
This completes the proof of correctness of the reduction.
Time complexity. Assume we are able to solve E-Gap-IPA in time n= for constant å > 0. Then we can solve all g² instances (AAii. Bi.j.k Ch.j.k) of e-Gap-IPA in time g'n0-6, Since q n(()), q² < 8/2 for n sufficiently large, and thus the runtime of the (Turing) reduction from IPA to Gap-IPA is at most n666 This completes the proof of Theorem C.8
28

---

Published as a conference paper at ICLR 2026
C.2.2 HARDNESS OF APPROXIMATING STRASSEN-ATTENTIE
In this subsection, we prove the following theorem which is the last step of our reduction for proving the lower bound. The following theorem gives an efficient reduction from E E-Gap-IPA to Strassen- attention when the weights are large. We again use the fact that Strassen-attention is poly-attention for the polynomial ( 22, 22, 4142 + 22233 13311
Theorem C.10 (Hardness of Strassen-attention)) For every constart
O, every 8 E
(0, 0.01), every C, M. > O, there exist constants Cn > and C. > 0 such that if APA hss (2n, 2clogn, n. = Covlogn, = =n-Ca) (Definition A.6j with query-key matrices Qt) Q(1) E [Tr,Ijanx2clnnn, value matrices V(2). V(A) E R2xx2e 22l2ol log
can be solved time Onn then e- G.p-l (Definition C.6 with target inner product m M log n can also be solved in Onn==) time.
Proof. We start with an instance of e- Gap-IPA p.d=c defined by A, B,C target inner product m = .M log n for a constant M, satistying the promise given by the definition of e-Gap-IPA (e.g., for every pair of vectors from different sets, their inner product is either equal m or at most (1 E m) From this instance we now want to create an instance of Strassen attention, given by matrices Q(2), Q(2) ((3), v(), ((22.
10, ( and
Now, for positive real number
(11) that we will fix later, similar to Alman & Song 2023
2024.. we construct the matrices Q (1), Q(2,, Q(3) E Raxd for ń = 2n, d = 2d as:
Taj
1
b1
0₂7
C1
0g7
:
B
On
ld
Q12) B
bn
Oa
Q B
Gn
Od
0a
la
0g
la
0a
1d
:
00
La
Loa
Laj
10a
LaJ
2nx2d
2nx2d
2nx2d
We also define V(), (2) Raxd whose first columns are
[127
1 (1:2n,1) = 10
(1:2m.1) = 0
and the remaining entries are zeros.
Correctness of the construction. We have defined the matrices Q(1,, Q(2), Strassen-attention so that, for any i,j,k [n]] we will have (Q 2 10,
_Q(3) underlying
(Q( ( (aaa,, by) (bj, (k) + (ck, ax)), and the bottom half of the matrices, Q (22) (20) (1+1:))) will act as a normalizing terms when we compute the softmax.
(n+1:2n)' (n+1:2n)*
As before, computing the output of the Strassen-attention works in two steps: for all [M]. first calculate the value of the denominator XXY)jis where X (Q21))))))) (Q())T] and Z = upper and lower bounds on this. Next, we will compute the numerator, [X DXD 2 Z), for all EE [d], where 1 diag(v (122,,)) and 2 diag d2ag(V((2) 1)2m,.). Our approach is to show that if there exists some i [n] such that for some j, E [n). have (at, = M log n, (by, cx) M log and (Ck, a.) = M log n, then we will be able to find such an using the entry-wise approximation of one Strassen-attention head. Thus, further improvements the entry-wise approximation algorithm would imply an algorithm for solving Gap-IPA A in time n=-(((1 time.
The normalizing term will allow us to give similar
Bounds the denominator. We analyze the denominator term and give upper and lower bounds on [XYZ]ua For computing this value, we find the value of EyAet exp (H(( ) (0)a) Q1) (Q( Q, We only care about the first
29

---

Published as a conference paper at ICLR 2026
of the attention matrix this is where the existence of an IPA will be noticed. For i [n], this is equivalent to computing
IX =
Σ
Σ
j.klj]]
1:2 Σ
(6)
Σ
clatd=d)B3
kefn],kfmm+1:22))
j.ke (
Using the gap property that the inner products of any pairs of ai, hj. Ck are either less than (1
e)M log exactly equal to M logn, and denoting x == previous equation, we get
M Mlogn where d 2clog from the
XYZZ14
3(1-))AB²
E
(1+(1=))))B
J.kE(n
Je|n41:2n].kEMl
((1==))+1))=-
Σ
eB /2
j.kEn+ 1:2n|
n² 2e 3(1-))AB + 2n (1+11=))
n
e 3B* /2
n² e 3B* 12
We also have A < 1/2 since M < c. Now, an upper bound of [XYZ]i, can also be computed using (as, by), (bj, (ck, 04) S M log n and Equation as, as,
[XY 2lu S
Σ 3AB
Σ
e(1+1)B²
Σ AE|m)AE
(1+2))
3B"
jE AE|m)
S 3AB2
+ 2n
2 (1 + n
3
< 2n² 3
for large enough B when is constant.
Bounds on the numerator. We analyze bounds on [XD1. 22.1 Z}, when positive certificate of IPA contains a, versus when it does not.
Case 1: IPA present at i.
In this case, we have
IXDILyD2.1 1/,1
Σ
el 2 +000 2f " y
j.ke(m]
Σ
((Q! 21000 2 20 y// (Using values of ((1), (2)),
BAB + (n² 3)e3(1-)) > 3AB
j.kE[n]
since we have some j, [n]] such that (a,, (bi, cx) (ck,, ai) m.
Therefore,
X2LIYD2.1Z R.
GAB2
(7)
2n² 2
Case 2: IPA not present in i. Here, we have all (a., + (bj. ck) (ck, ar) S (2M + c)M) log n for all j. & since otherwise will contain a IPA. Therefore,
[DD.ID221 Z),
Σ
H10 of 2e /d (1)0))))
1.1
Σ j.AE/n]
(10 0000 + o (Q)) 2
Σ
(3-=)AB²
S Snz (3-))3)
=
j.ke[n]
j.ke(n]
30

---

Published as a conference paper at ICLR 2026
which implies
(2D. 7/
el3-s)AB²
(8)
e
Wrapping up. Let u, be the value of the approximation of the i-th entry of the first row of the Strassen-attention matrix, i.e.,
5.
XY
We will show that u, is distinguisher between the yes and no instances of IPA; in particular for appropriate settings of the parameters will see that the value of u, in Case (the yes case) is always greater than the value of u, in Case 2 (the no case).
In Case 1, using Equations
7 have
3AB²
Uj
[XY Z)i,i
7
2BBy2 7
and in Case 2, using Equation
8 we have
44 XD Zles
c(3-))AB
+
3B2 - 7..
Thus it suffices to verify the following inequality:
(33-))BB2
3AB2
e3B2/2
2n2e30 2n223B212
Y
which is indeed satisfied for y 2- and A
> n². Therefore, B = [)loogn suffices.
Therefore, we have reduced Gap-IPA to APAC(ks) where B 12(VlollVl) completing the proof of the lemma.
Therefore, if APAC (ks) could be solved in O(n"~A) time, then that would imply that IPA could be solved in O( ((8)) time (Theorem C.&). which in turn would imply Max-2SAT could be solved in 2(*/s-2(8))nn time (Lemma C.7 which can not be true for an absolute constant 8 >> 0 (Hypothesis 3
D PROOFS OF SECTION 3.2, TREE-ATTENTION
In this section, prove the first part of Theorem 3.5 by giving algorithm exactly compute the output of tree-attention. The second and third parts are computational complexities of special subcases of poly-attention, which has been proved in Section
Before giving an algorithm for the exact computation complexity of tree-attention, we show property of branchings in the graphical representation. This happens when the underlying polynomial for the poly-attention is a variable separable polynomial.
Lemma D.1 (Variable separability). If hi(zi...., some [t - 1] and some polynomials f.g of minimum possible degrees, i.e., h is variable separable (Definition A.S then we have Att(h) Attlf) Attla) and also the entrywise- approximation (h) Att(f) @ Att(a), the attention, !ft and Att(a), can be computed in time (m) and T (n) respectively, then computing the wwtpwwiss-pppoximmate.m of) output of poly-attentonn for h, (h)( can be performed time O(max (n:), T (n)) =nd).
f(z1...., 24) + I(I1,, It) for
of) outputs of poly-
31

---

Published as a conference paper at ICLR 2026
Proof. For all [n]. [d], we have,
Att Att
- Elsmme exp( ( ......,())) 6sk likk
Elamin exp(2
(2) ....,,,
Eltriint exp C9
((t1)
(t)
((tt)
V(()
look
Etermmn exp H9
Q(1))
(9)
Ela exp (2(((())) o
+ +( ((+)) 2
Eegiini exp (2) L ...., 2 + t Q
Elsmall exp( ) 00 ( (2)
Eink
Au j.k
Elgmidy exp (QQ) 0
This implies Att(f) @ Attls) Att(h) and we obtain entrywise approximations Att and Attle) respectively with enorr 7 paly(n)' then Att(f) o Att(a) will be an entrywise approximation of (l(()) with error To 0(7) poly(nj as well.
Note that the polynomials might not even contain the variable 2ı, in which case we all the rows of the output of the corresponding poly-attention matrix will be the same.
Now, we prove that (tt where h is a tree polynomial, can be computed in (n?)) time. Theorem D.2. If h is a tree polynomial (graphical representation of h is a tree or a forest), then we can compute (t((h) exactly in Ó(n²) time.
Proof. Algorithm Sarrtuhm@givee gives procedure for computing the output of tree-attention given query-key and value matrices as inputs. Indeed, if there were multiple forests, we could have computed the output of tree-attention for each of them separately, and composed them together using Lemma D.I
We start with a tree rooted at U1, and compute poly-attention on each of the subtrees (polynomials corresponding to the subtrees) where the query variable is the root of the subtree. The main idea to compute this is, whenever we have branching. we compute each of the subtrees separately, and compose them together using Hadamard product of Lemma D.I
Overview.
In Algorithm we fix each of the columns E E [d] (Step 2), and compute (tt (1 one at a time.
The computation of proceeds as computing the numerator and the denominator terms separately, from the graphical representation G (as in Equation l.. In this recursive formulation, we employ compute the values in DFS fashion, first, we fix the root of the tree given by variable z1 (vertex in the graph), having the corresponding query-key matrix (11), and proceed to computing the output of the poly-attention mechanism for its subtree polynomial.
Each branch. Without loss of generality, consider the root variable un, and for each branch from V1, consider an edge given by (Vi, vj.), i.e., 01- for i E [p], where p is the number of branches. When Uj is a leaf, we compute the poly-attention Att(+1=5,), and recursively pass it up the tree. The denominator and numerator of A(f(=115)) are defined in Step 12 of Algorithm 2 two vectors in Rax1 which can be computed in (n?) time and then their ratio is the poly-attention output for this branch (Step 13).
Next, when 2, is not a leaf, i.e., the tree proceeds as U1ーU, - let us assume the polynomial whose subtree rooted at Uj is given by gi ( .Ft) and that we have already computed ( 29)
'Query variable refers to variable the highest priority the polynomial (priority monomials and variables has been defined in Definition[E.) It is usually the variable æi. and the indices of the corresponding query-key matrix in the softmax computation correspond the rows of (tg(h) (see Equation
32

---

Published as a conference paper at ICLR 2026
Algorithm 2 Algorithm to compute tree attention A
Input: A polynomial h( whose graphical representation is tree, query-key matrices
Q(),, ..... Q(+)) Raxd and value matrices V(2),.... ,.... V(t) Raxd
Output: At((h) E IRaxd
1: Construct G as the graphical representation of h, with vertices ,1,..,, Ut.
2: for E [d] do
3: Let p be the number of children of V1-
4: for all child node v/ of U1, i [p] do
5:
if Uy. is not leaf then
6:
Let go ( xf) be the polynomial of the subtree rooted at Vj.-
7:
(oc(s, recursively, where U/ is the query variable, by computing the numerator term and the denominator term separately. Let the numerator term be
Compute All (1:n,8)
((s("1)))) E Rax1 and the denominator term be R(a(())))) E R=x!
8:
Define the numerator
(Q
p(ai(1g
and the denominator,
Rimen +o:0s, "= (()
where DV diag(V Rux=.
(1m,6)
9:
Compute
ta
Att (217 to((ss (1m.e)
ta (a
10:
else
11:
Here, , ) 0 since there is rooted at ey.
12:
Define the numerator
p(a=))) = (Q (1m.6)
and the denominator,
;( (QU)(( 1nx1
Compute
Att (115)) p(( (sm.)) p(
14:
end
15:
end for
16:
For composing the branches together, compute the final numerator
p(h) = p(aya, +91 x2)) )) の plays ()),
and the final denominator,
R() = p(a,0 to1(2s) o C p + (xsp())))))
where h = Jy + 91 (1 + + IJJJj + 9p (Tjj 2²) (by definition).
Define
p(h)
(h) (h)
(1m.0)
== R
18: end for return (tt(h)
33

---

Published as a conference paper at ICLR 2026
(the numerator and the denominator separately given to p(a(e,,)))) R(o((1), R respectively). By x', we simply denote the subset of variables other than rj, that the subtree consists of. The output of tree-attention of the subtree rooted at Vi essentially Att(#1n +o:(s; For this, the numerator and the denominator can be computed as in Step 8 both of these computations take Onn²) time. The final value of All (21 (m (²)) is given by Step 9, and we pass the numerator and denominator vectors up the tree recursively.
(1:n.)
Along a branching. For conglomerating the branches, let us say that the children nodes of on "/ar-.r2 "/2- where the polynomials corresponding their subtrees 91 (4"), :.--p- 9p (Ij,, xP) (r',...,22* x² are disjoint subsets of variables which are precisely the ones present in each of the subtrees, respectively). We also assume that we have recursively computed the é-th x)
columns of the poly-attention outputs All (1m,))
..... ....,,A(((... (1y(()))"))) in terms of the numera- tors (la ,²))), p(sp(*,,:"")), and denominators R (()))) ..... R(sn(ss, respectively. Now, the poly-attention output for the polynomial having the subtree rooted at vn, which is
, ..... 2²)) 21 Ij (1),, +2) +... 21 9jp (Fip,
computed in Steps 16-17, and the correctness of this computation follows from Lemma D.I
Time complexity. We show quadratic time-complexity for Algorithm 2 Let assume that recursively in branch, the numerator and the denominator of Attls:(=s.=)) can be computed in Ö(n²) time (Step 7). From this, extending the output matrix of poly-attention the current vertex (Steps 8-9, 12-13, followed by 16-17) each require Ó(n²) time. The number of these sub-tree attention computations required is at most the size of the tree, O(s), which a constant. Therefore, this gives a DFS-style procedure to compute the ( in time Ö(n²) since the graph is of constant size, and repeating for all é E [d]. we will be able to find the entire matrix (M)
E PROOFS OF SECTION 3.3 POLY-ATTENTION
COMPUTATIONAL COMPLEXITIES OF
Throughout this paper, we will compute the numerator and the denominator in Equation separately, where the numerator term is Eigmnile exp (f 2!!)) (2) + v(()
.... vi( and
the denominator term is Eeg... exp ( ... 2
(k))
We also define monomial ordering, which will help us proceed with the proofs of these theorems.
Definition E.1 (Monomial ordering). A monomial my is said to be higher preference than another monomial my if either of the following holds:
deg(m1) deg(m2),
deg(ma) = dog(m22 and my comes lexicographically before m2, i.e., if is the smallest index such that x; is present in exactly one of the monomials, then the monomial in which x; is present has higher preference.
We will order the monomials of h according to this order, and m, will denote the i-th monomial. Note that this definition can also be used with variables, where variable x; has higher preference than æy if and only if i <
The polynomial method Alman & Song 2023 2024 2025 can again applied to poly-attention, by reducing APAC (h) to the computation the output a larger t-tensor attention, where the query-key vectors in tensor attention are of dimension n x (sd). However, the bound the variables in this case of computing poly-attention will be o((og 1/k) in contrast that of tensor attention being o(loog v)) in) Alman & Song 2024).
For proving Theorem 3.6 we show the two parts, upper and lower bounds, separately. For upper bounds, we give polynomial method algorithm if the entries of the query-key matrices are bounded
34

---

Published as a conference paper at ICLR 2026
(Theorem E.2p. and if the entries are large, give hardness results for entry-wise approximation conditioned on fine-graing men-ggrined complexity conjectures (Theorem E.3
Theorem E.2 (Polynomial method poly-attention). Given attention polynomial f(z1...... It) of degree having monomials, where ki 8 are constants, there an algorithm that solves APAC (n, = Ollogn,, = o((logn) )), = 2/poly(n)) with query-key matrices Q E|-r, and value matrices V ,...,(() E Ruud in time ).
Theorem E.3 (Lower bound for poly-attention). Given attention polynomial h(rr,---.,3t 3ie) of degree k having & monomials, where t, ki are constants, we are interested in computing an entry- wise a-aproximaaion (ht having query-key matrices Q(1), "i.,,...QQ [-r,,]]x and value matrices V(2) y.... V(O E Buud for Ollog n), 1/Dol((n). Then, depending on the structure of h,
1. If 2 2. then assuming SETH (Hypothesis) Z.. entry-wise approximation of Att!h) can not be computed in time On h (..(()(()))))) 2(1) when I= ()(logn) l/k).
2. If contains an elementary symmetric polynomial (1te)) for some to S then assuming the MAx conjecture (Hypothesis 2 an entry-wise approximation of Att(h) can not be computed in time (n-ta-((() when [ 52((log ///
3. Ifk 6 and h is not free polynomial, then assuming the Max-2SAT conjecture (Hypothesis an entry-wise approximation of Att can not be computed in time O(n--1(1) when
2((lgn))13)
E.1 POLYNOMIAL METHOD FOR POLY-ATTENTION
In this section, prove Theorem We start with the polynomial h defined in Theorem E.2 and reduce the problem of computing entry-wise approximation of Att!h) E Rgdd to that of Al(F) (ed) by constructing query-key matrices K(11, ..... K(O) K(1),..., matrices W(1) ..... W(t) E R (sd),, such that the row-softmax matrix of
[Rxx((d) and value
k (K K((3) ..... K(O)²
same the softmax matrix of Att()), and (tg(h) exactly equal to Att (1:1,1:d) using these inputs, and the remaining entries of Att(T) are zeros.
Defining K(), We define K() E (ad) for all j E [t], by dividing the columns into s blocks, each having columns. These blocks are defined as, for j E [t]:
the i-th block, for i E [s), contains the matrix Q if the i-th monomial of h contains the variable xy.
otherwise, the i-th block, for i [s), contains the all ones matrix 1 1xxd-
Roughly, the query-key matrices can be seen as:
1,d
K
(
(3)
n.d
3. bot 178
2y not in ma
2, is in m2
(x()
Using these definitions, can be verified that for this choice of K'))'s 's we have
(K ((2) ..... k!!))
Σ
(K
(2)
E2. 1)44 1iid (((()
Ex 1)4 mad (10)
61 1a4
sE|s]
Σ
( (i)
Qenu
(11)
sE[s
35

---

Published as a conference paper at ICLR 2026
where the monomials of h are defined before (Definition 2.2
Defining W The value matrices for the t-tensor attention operation will be the same as that of poly-attention. order to match the embedding dimensions the query-key matrices and the value matrices of the t-tensor attention operation (as used in Alman & Song (2024)), consider the new x (sd) dimensional value matrices, WO's to contain the corresponding x d dimensional value matrices V( in the first d-columns, and all the remaining entries of WU) contain zero. More specifically,
simply
W [VG)
Onxd Onxd)nx(sd)
(12)
Now, in Equation note that the output be written as
D- --AW(2) @... e...0 w(o,
where Ruxn' is defined as
A k()) @... (..
and D is the n x n diagonal matrix
(K... k
D = diag
1 nxl .... O 1 nxl
times
This is precisely the form of a f-tensor attention mechanism. Next, in order to use the polynomial method on this matrix, we need the entries to be bounded.
Lemma E.4 (Bounded entries). Given Q(i) E [-r,I]mx and h defined above, we have
e- sexp (EM(OM) .....(())) s T
for all l 4 [n]. For r exp ..... ...., (.. are sub-polynomial in n.
of? (logn) o((logn)/M)) the entries
Proof. Since h is degree polynomial with constant coefficients, for each monomial my of h }m,(Q!)) 2 in in the range of [-r*, There are s monomials and the total value is bounded inside the interval [-s sT*), which gives the required result after exponentiation.
For completing the algorithm, we use results which follow from the proofs in Alman & Song 2024 Apx. E).
Theorem E.5 (Alman & Song (0244). Given matrices KQ),..., K E [-r, and value
matrices W(2).... W()) E Raxd we can compute an entry-wise haaproiiaamttiitiioonnno for = 1/poly(n),
of the following:
matrix All E Iand which the entry-wise apaproximmiono of the numerator matrix of tensor attention output
Att = k o...
",. 20..i//www.....o....o.. w(0)
that is, for all i [n].j E [d].
Allla Alli.a <7
36

---

Published as a conference paper at ICLR 2026
2. A diagonal matrix D R= which is an entry-wise approximation of the diagonal matrix E Raxn given by
(1) 0.. O.. 1 (..00)00000.0))))1111))) C... @1xx).
D = diag
that is, for all i E [n]]
IDis Di,il <.
A} ollog n) is met (where A = 10 @ 2 Kl the time complexity for finding the matrices Att, D, and hence entry-wise 27-approximation of D-1 All, n
Here, when the condition max log(1/7)
[og(log(1/1//A)
Using Lemma E.4. the value A in Theorem E.5 O(T*), and for the choice of r o((log n)²),
the quantity max Ing(1/9 [1]( A] is indeed o(log n), which gives our required almost-linear
complexity for computing (tt())
Summing up, the algorithm for computing entry-wise approximation of Att given as the following algorithm.
Algorithm 3 Algorithm compute entry-wise approximation of Al
Input: An attention Q Q() V(2) VlO E Raxd , poly(n)
polynomial
h. of degree
matrices
Output: Entry-wise 7-approximation Al((h) E Raxd of (l E Raad
1: Using Q!)), ..... ((t) and h, compute ( ..... ((t) (sx((s) (Equation 10
O(nd) time. O(nd) time.
2: Compute W12) ..... ((!) E [ax((d) from V22), ..... ((() (Equation 12
3: Compute entry-wise aappproximation Att E (=d) of
Att k()) (K o.. 2 K(0)) a. o W(O)
using Theorem E.2 Step 1.
> O(n 1 +o(1)d) time.
4: Compute entry-wise 7-approximation D E Raxn of
D = diag (K..
O... 0...0 0...01xx1),
which is a diagonal matrix, using Theorem E.2 Step 2.
(n()-)())d) time. O(nd) time.
5: Return Đ D-AAll(1m,1dd)
This proves Theorem E.2
E.2 TIME LOWER BOUNDS FOR POLY-ATTENTION
We complete the main complexity result of this paper, either we can compute an entry-wise ap- proximation of poly-attention in near-linear time, when the entries of the query-key matrices bounded; or we require at least superquadratic time, unless the polynomial for poly-attention a tree polynomial.
Our proofs for showing the hardness of entry approximation of (tt(h) consists of two reductions: (1) first we reduce from each of kIP, HypergraphIP, and IPA (which have popularly known hardness conjectures of SETH, Max-2SAT, MSx respectively) to ( instances of their respective gap versions, and (2) secondly, reduce each of those gap versions to entry-wise approximation of poly-attention. These subcases and the starting complexity assumptions will based on the structure of h provided, as categorized in Theorem E.3
37

---

Published as a conference paper at ICLR 2026
For proving Step 1. when prove the first case, get hard instances of E-Gap-kIP assuming SETH (Theorem E.7 For the second case, we assume Max-kSAT is true, reduce Max-kSAT using a known reduction (Lemma 10 to no(1) instances of HypergraphIP, and further reduce each of those instances to n2)(1) instances of =Gap-Hypergraphhl (Corollary E.12). 1. For the third case, we start with Max-2SAT and reduce that to n instances of IPA (Lemma C.7 and then to n((1)
instances E-Gap-IPA (Theorem C.8).
We complete the reductions for Step 2 in each of the following subsections.
E.2.1 TIME LOWER BOUNDS BASED ON DEGREE OF POLYNOMIAL USING SETH
In this section, we prove the first part of Theorem We first start with an instance of kiP, which is SETH-hard, reduce it Gap-kIP (Definition E.6 using some previous works Rubinstein (2018); Alman & Song 2024). and then using the instances of s-Gap-kIP, create query-key matrices for All such that an entry-wise a-approximation of Au would solve the instance of e-Gap-kIP.
Definition E.6 (c- Gap-klP), For every E E (0.1) and positive integers k 2 2, given sets of vectors A S {0, 1]d with [A' = = | = target inner product m E {o,...,,d dd. and the promise that for any 01 A ,..., O0 E A#,
A A!....,A
* either (01,.... ax) m.
or, (a1,...,,ak) ak) (1 E)
the problem of s-Gap-kIP, n.d is to decide if there exist vectors aj E A Og E A such that (ai..... ax) m.
Using Rubinstein 2018 like techniques, conditional hardness of e- Gapp--PP can be obtained.
Theorem E.7 (Alman & Song (2024): [Rubinstein (2018)). For every 5 >> 0 and every constant & E (0, 1), there exists a constant c O, such that e Gap-kllpPPPpee n,elog n for any target inner product m foo..., n) cannot be solved in time ())(1-)))), unless SETH is false.
Due to this result, we start with an instance of e-Gap-kIP and reduce that to an entry-wise approxi- mation of (u((h) If the entry-wise approximation of (t can be computed in n (1-6)) time for a constant 5 > 0, then e-Gap-kIP can be solved in ((1(1-6)k) time, which would refute SETH.
Lemma E.8 (e-Gap-kIP (l)) For every constant & O, every E (0,0.01), every of degree k 2 2 having monomials, where k, 8 are constants, there exist constants C. > O and C. O such that if APAC (2n, (8 1)c - Co (logg)) ,7= nCCa) (Definition A.6 with query-key matrices Q(1), ..... Q(t) [-r. (nxx Jjolog and value matrices (²) V 22*x 1)clog can be solved in time
M, > 0, given an attention polynomial h(zı,..., It)
qusty
Onnk-6), ), then 8- -Gap-klPP..ct n.clogn. (Definition E.6 with target inner product i. = M log n can also be solved in (m) time for any constant M.
Proof. Let us start with an instance of e-Gap-kIP n.d=clogn that we want to solve, with & sets of vectors A....., A C {0, 1]d, consisting of vectors each. The vectors are {af...., an == A' and the target inner product is m = M log n, for a constant M, with the promise of the gap condition for approximation factor We also assume that there does not exist an all one's 's vector in A for each i E [k], as that would violate the gap-property (as m needs to be smaller than d for hardness).
Using this instance of deciding G-Gap-kIP, reduce it to computing entry-wise approximation of (lth), with query-key matrices Q(1), _(.),..., Q(1) E [-r, r] axd and value matrices V(2) ,,,),, V(t)
päxd for ñ = 2n, = (s + 1)d = (s 1)c log and I' that we will choose later.
Let us assume that the highest preference monomial of h, a monomial of degree k. is given by Tri ... Tre, where 71 has the index of the highest preference that may or may not be 1.
We will construct the query-key matrices such that each matrix Qiry) will contain vectors from A for j < k. zeros otherwise. Having the monomials ordered according to descending order of the monomial ordering (Definition E. each of these ((ry)'s will consist of blocks of columns which
38

---

Published as a conference paper at ICLR 2026
correspond monomials- the i-th column block, containing d columns from (i 1)d + 1 to i.d, for i [s], will correspond to the monomial m,, and the last column block will be a normalizing block. The idea of the reduction is that only the degree & term Σy, ... ...Er, Ir, of h will contribute to computing the final inner product, the terms which are subsets this degree k term will cancel each other out,
and all the other terms will be zero, thereby not contributing anything to h(Q), Q . Q More specifically, we want,
MOQ (1)
2
Aaat
for ....... fra E [n], and some scaling factor A which we will see later.
Construction of matrices. Let us now define each block of Q [t], which will have 2n rows and (s 1)d columns. We will define them by defining each of the column-blocks using scaling factor B (1). Considering the set = (ri...., r we define:
1. For Qii''s jf $ T, we just make the entire matrix zero U2+x(*+1)d 1jd-
2. We now fix [k] and define Q(1s) (i.e., value of ry E T). We define first column block of ((1)) as:
2(1)
0
B
(1:2,.1d)
0g
0j
Od
2nxd
For column blocks E [s), if the monomial m, does not divide Tr. Tre, we just make that block all zeros
(ry) (12n,((-1)d+1:... 02nxd-
3. If monomial E [s] does indeed divide Ir, Tra consider ji as the index of the highest preference variable present in m, Trn Fr. for J....., Jk. E [K]] k2 < k. Let s; be the negation of the integer which the number of occurrences of this monomial m, along with coefficients, in each of the monomials ordered higher than and that divides Ir, ... Trk (these are the only non-zero monomials).
More specifically, s; is the sum defined by adding:
from the monomial mi.
- S5 whenever < f < i, the monomial my divides m,, the monomial m, divides my and the highest preference variable of m also present in m,-
-1 whenever 1 highest preference variable of me is not present in my-
the monomial my divides my and m. divides my, but the
O in all other cases.
If Zr, is not present in monomial i, we simply set
(1-2n.(-Ddd4iid Ozaxd:
39

---

Published as a conference paper at ICLR 2026
otherwise:
8,01 21
8;02
8,0
= B
Oa
Og
Oa
2nxd
where Trz is the highest preference variable m, and
022
Q (
a
(iin,
B
0a
Og
Od
for all other j's such that 2r, is present in monomial
4. The last column block for ((r1) 1 is the all ones matrix 1nxd with scaling factor, i.c..
(1))
B 12 xds
and for j [2 k), it is the matrix
Q (ry)
Onxa Inxd 2nxd
Roughly, the query-key matrices can be seen as:
di
02
83 0
1,1
((ri) B
On
83 On
0g
Oa
0a
ld
Oa
Oa
ma does not divide mi
Ley not in mi
ma divides mi
nx nx(++1)d))
and for all other [2: k).
0
83 a
0,
Qls) B
a
af
0
0
0g2
1
0.
0,
0a
1
my does not divide mes
Zeg not my
mg divides my
(x(++)))
40

---

Published as a conference paper at ICLR 2026
For the value matrices Vo) [2(n))(s+ E
T (1),
define the first column as,
v (1:2n.1) Lo
1
and for j E [ define the first column as,
v 1
(1:2n.1) [
All the other columns are completely zero O
Correctness of construction. We now show that for l,....,4 4 E [n], h(Q Q ..... 0 mi (Q )- By definition, my
In ....... Fra and it is easy note that = B² (ak.,.... ... ......,,,,,,,,, ....,,,,!" )- For all the other degree h terms, the inner products from their corresponding blocks are all zeros as we had defined Q as all zeros matrix for all j & T.
B* (a)
We want to show that for all other i's. mz ( m²(Q ................................ e the i-th column blocks for product m, if and only if m, divides m; (otherwise m((O 1 P 134+ zero), and no i has contribution for mial ordering. Now, from the choice of 8, as above, follows that mi(Q!)) ... Ei m, m):(()) ( (,
(1) ..... !!!!) 0. When we compute
have some contributions to the inner
(())
due the correctness of the mono-
( (1 =0
For bounding the values of s we use induction prove |&s| < sf. The base case obviously true. For the induction step, assuming |s;l < 8², for the (i +- 1)-th monomial, 81+1 needs to cancel the contribution to the inner product corresponding to m4+1 from each monomial my which is divisible by m441. The contribution is at most |sy] s² (from the induction hypothesis), and hence |si+1| < Σ; mm.eim;; [sg < Ei: m.+1 liz s² i.s' < s Therefore, we have |s;] < s", which implies I = O(s* and from the definitions, obviously have I > B as well. Since, O()), we have = ((B).
Further, these query-key and value matrices can be computed in O(n (o()) time.
Approximation yields gap property. We assume an entry-wise approximation of the self-attention matrix, and the goal is to compute two values, the numerator and the denominator, for computing the softmax. The numerator, for é [2m], is given by
Pr Σ
o...... exp (=MO " ) V
and the denominator by
Re, Σ exp L
4E(2n)
The 4tth row of the (h)(h) will be Psg Re, and we want to find an entry-wise approximation. Since in our choice of the value matrices, the first coordinate of Pr is the only non-zero one, and its summation is only upto the top half of the value matrices, E [n), for kt. The only non-zero part of the numerator, that we care about, is therefore given by
Pe. Σ
(-h(0 Q ......QQQ)) 9.!) exp d
GEM GET E JET
we have an eitry-wise apppooximttionn of (tt(h), let xy, -th be the approximation for the (f1, 1) entry of (h) By definition, we have
Pt. <
(13)
Rr,
41

---

Published as a conference paper at ICLR 2026
Bounds on denominator. Consider the summation Eig... E(2n] exp (1)(Q) 0 ...., )) Define the factor such that when only the ",,,..,; 7k coordinates are B. and the remaining are zeros, i.e.,
(1)
<d
h(O,. B.ld, B.1, O....,, Ad
is easy to see that = 1 + o(1), since the evaluation these values will give a B from the first monomial, and the other - 1 monomials will give at most (s 1)Bk-I oB....
For the choice of QU)'s, we have
Re
Σ
exp( Me(-hhQ@ 0 ..... 2!
1
(2)
& E(n+41221 GET E Jer
exp(AdB* /d) = 1B"
since all the Ofri'', s, for j E [k], have the zeros in the last column baluunn-block and 1a along with the scaling factor B, which makes all the monomials of h give inner product dB*
For the upper bound Re have to use the maximum possible value of h(Q)) Q!22 (
(1)
irrespective of whether Ey are in [n]] or [n + 1: 2n.. Let us consider a choice of lg..... E [2m]. If all the E,, ss for E (1)), are in [n + 1 2n]. then the value of h( (1) 0
Q obtained will be be Ad Otherwise, there are some (but not all) E in [n]] for E [2 kJ. where the monomial of degree < k containing only those variables will be at most 8 ddBBB and the maximum
value will be obtained from the first term, which can be at most B (a! ..... of = (d - 1)
Thus, in this case, the maximum value of ((()))) 269) ..... ..!) would be (d- 1 B* o which still less than A
Therefore,
Σ exp (EMOO!) " Q
R
(((m))
AddB* 21-1 t (t*
exp (
d )
(
Therefore,
n' 1 (B 221) Re, <2 (1)
(14)
Bounds on numerator. Now, we will show that if a vector tuple exists with the proper target inner product (a positive certificate for Gaap---lP) then P:; is so large that 2e, (Equation 13 is larger than fixed threshold. Here, we first show lower bound on Pr.. Otherwise, we will show that æz, is small since every inner product will be scaled down by gap due to the approximation promise.
Consider & [n] when there exists c ...., eo E [n] such that the inner product (aty ..... a
M log (it is quite possible that r in which
will only consider t 19. Then,
42

---

Published as a conference paper at ICLR 2026
have
Pti
E
exp
1
M Q!!)
d
A :ET
6 JST
Σ
Eh 2 , 2!!
exp
G=
[ 12T
(2n) k-k- exp
= Bk (a
ar
deg
M
(2n) exp
B* (s+ 1)c)
where the second equality follows from the construction of the Q63's Using the upper bound of Re Equation 14 we get,
(2n)* c (B* Git)
" (0 ((--))(ott))
Ph
Re
(2n) 1e Bt x1))
(2n)k-1
Pa. Rr.
Using 2t
? (Equation
we get
(2n)a-1
7
(15)
Now, for finding upper bound on 24, when exact inner product tuple does not exist, we
exp (EMOH ......, e!!))
Pa
Σ
E iET E j
Σ exp BB oi...... at. o
= (2n)* (22m)=Ae -- k-1
4E(N)
KE
using the construction of ((r))'s. Now, using the gap property of inner products in our instance of e-Gap-kIP, we have
Pr = (2m)
Σ exp
(BB*(at,
(at, oz,..... ................................ .......))
d
::(n))iie iET
Bh (1-
(2n))-k-1
Σ exp
((1-s)ae s)M logn
4 : iET
- < 2 k- 1n 1-1 ( (11-))B* (1
Finally, using the lower bound of Rs (Equation 14
we get
21 (1- 1)e
Ps
2 k-1 ( (Bx( ))om)))
R
n ( Bhh (B)
and the bound on Ie, from Equation 13 implies,
Pt Re
+ 2 (o-)e 2)(G1)) +7.
De.
(16)
Wrapping up. In order to differentiate between the cases, we must have the lower bound of zz, when a positive instance for E-Gap-kIP tuple exists, Equation 15 must be greater than the upper bound when such an instance does not exist, Equation T6
21-k-1
(eB el (a*(A ))/o)))
+
(2n) (et 3)(o)))) 7:
43

---

Published as a conference paper at ICLR 2026
which is true for the choice of y I CC and B Co(logn) for large enough constants Ca, C. O. This would make e 8 4) large enough and y small enough, such that the inequality will be valid.
Now, since is constant, the maximum absolute value of the entries of the query-key matrices are ((B) = (2((log ([((log))'/4 n)) which proves our result. Therefore, if we can find an algorithm for finding an entry-wise 7-approximation of Au(A) for APAC (h) with these parameters, that runs in time n² --0((1) then SETH will be refuted (Theorem E.7
E.2.2 TIME LOWER BOUNDS BASED ON SUBSTRUCTURE OF POLYNOMIAL USING MAX-KSAT CONJECTURE
In the second part of Theorem E.3 we prove a stronger lower bound when the monomials of h contains an elementary symmetric polynomial of degree k in variables where k to I t. The underlying conjecture for this lower bound is the Max-&SAT. We first start with a problem called HypergraphIP (Definition E.9p. which is least as hard Max-kSAT, show that its gap version, 8- Sa-Gaa--Hyaa Hypargrrrrrr (Definition E. 11. is also at least as hard as HypergraphIP using Rubinstein 2018); Abboud & Ron-Zewi (2025), and finally show that computing an entry-wise 7-approximation of Att!)) efficiently would solve H-Gar--ypergraphIP faster, thereby refuting Max-kSAT conjecture.
Definition E.9 (HypergraphIPP).) For positive integers t, k, given t sets of vectors , A' E t,k {0, 13dd with |A'l = .... = |A4| n, and target inner products (a......(())) the problem HypergraphIP"..
t,k is to decide if there exist vectors a1 E -A'...,a, as E A' such that for all sub- sets S E (12). we have (as(1))....,0 0.(s)) = my, where ms is the target inner product corresponding to the given k-sized subset among the (i) choices.
We will drop n, d from the superscript and not include the target inner products as the parameters to make the problem definitions less cumbersome. This problem again has a hardness result, as follows.
Lemma E.10 (Alman & Vassilevska Williams, 2020 Theorem 23)). Assuming the Max-kSAT conjecture (Hypothesis 2]. for every 8 >> O and every positive integer k, there exists constant log. cannot 0 and larget inner products m1....; (!) E (O....,.) d} such that HypergraphIP t.k be solved in time (1)))
We can again reduce HypergraphIP to its gap version Hypegaapraap to show that this problem hard as well.
t,k For every E (0, 1) and positive integers I, k, given t sets of vectors A A!...., A {0, 1]a with |A'l ... |A'| n, and target inner product m E (o....d)), .d), along with the promise that for every dn E A ,.,...,0 0g E A and YS (12).
Definition E.11 H=Gar--yperrraphhlP:
* either. (as[i)..... ... asss) = m,
or, (ast)): ...,, aspa) S (1 E)m,
of E A such the problem g- -Gap-HypergraphiPPt is to decide if there exist vectors an E A At,...,aa that 4S E (!)), we have
= m.
Again, similar to Gap-IPA, for Gap-HyperrraphIP, we consider the target inner products to be the same for all the subsets of inner products, since the Rubiisteii (2018)-like reduction accommodates this, and we need this condition for reducing Gap-HypergaaphIP to entry-wise approximation of (k)
The hardness of &- Hypergaprr follows from a proof very similar to Theorem C.8. given by the following corollary.
Corollary E. For positive integers t, k with l 2 3, and every 8 > assuming the Max-kSAT conjecture, there exists constant and target inner product m E ..l.....oooo n). the problem -Gaa--yyperrapphhph cannot be solved in time (1)6))))
44

---

Published as a conference paper at ICLR 2026
Proof. We can again the reductions of Lemma C.9 We start with instance of HypergraphIP having sets of vectors A ...., A S {0, 1} containing n vectors each, and reduce that to n no(1) in- stances of --ra---yprrraaphlPPr having sets of vectors B!. .....; B C {0, 10, for d (log n). where each B' contains vectors.
t.k
The proof goes as- for each k-tuple (i,,.. (J..... (kk (12). we reduce A ....... Adh, an instance of kIP, to (1) instances of E-Gap-kIP of dimension dg (using methods of Alman & Song (2024);
Rubinstein 2018): Abbbbud & Ron-Zewi (2025). Then, we combine each of the e- Ga instances
for all (1 (..)).)))))) E (1) by creating (+) column blocks, each of dimension do, done in the proof of Theorem C.8 where the block corresponding to (ik will contain vectors obtained from
the above reduction, and the rest will be zero. The hardness result also holds true when the target
.....,." inner product for every subset of B' ......,B. B are equal.
Now, to show hardness of computing an entry-wise 7-approximation of Att(A) where h satisfies the conditions of Part 2 of Theorem E.3 we reduce e H-Gap-HyperrrpphlP... forr (which we know is at least hard Max-&SAT), to an entry-wise approximation of Att(A), Armed with Corollary E.12 we are now ready to prove the following lemma which completes the second part of Theorem
to ())C ). For every constant E >> O, every 6 E (0, 0.01), every c, M > O. given an attention polynomial h(z1,..., of degree 2 hav- ing- s monomials, such that the set of monomials of h contains as a subset all the monomi- als of the elementary symmetric polynomial in to < variables of degree ki where 8, are constants, there exist constants C. > O and C. > O such that if APAC (2n, (s 1)clogn, Ca(logn) ,y -CCaaa Cu) (Definition A.6) with query-key matrices [-r, ()nx((4l+))) and value matrices V(2) via),...,V V(t) E R22 can be solved in time O(n'o=²), then E- Gaap-yyperrapphlPP... to.k (Definition E.11 with target inner product m M log n can also be solved in ((n'0-8) time for any constant M.
Lemma E.13 E-Gaa--yyergraphll
Q(t) e
Proof. First, we consider that the subset of the monomials of h, which constitute symmetric polynomial in to variables of degree k, is given by the set of subset variables Zr,...... Tre. Let us denote T = fri..... 7 S [t].
Let us start instance of e -Gaa--ypergrpphhin fo.k
with to sets of vectors be A C {0, 13d, having n vectors each, and the target inner product being m M log with promise of gap given with a constant approximation factor of e. More specifically, we want to check if there exists ,,,,..,,
A
brog E [n] such that for all (JI, (Jz...., jk) ((te)). we have (aZ)
m, i.e.
(10 m =: 770.
Σ
jk agr
)
where ma My log n. We also have the promise that for every other tuple er,...., trio E [n] where HypergraphIP... tor property is not satisfied,
((18 10 -1 m + (1- =)m =: (1 Eg) mo)
Σ
lat
ale,
<
31 ((Pa))
for another constant Eg =0= c/(2)
Constructing the matrices. Now, we define the matrices Qi)'s, such that
1( 2 = Σ
for a scaling factor A, in a construction quite similar the proof of Lemma E.8 The query-key matrices will be Q(),..., Q(t))) E [-r, rjxd for ñ = 2n, d (s + 1)d, defined as follows using a scaling value B w1)) which will choose later:
45

---

Published as a conference paper at ICLR 2026
1. For Q if & T, we just make the entire matrix zero (2+x(s+))d
2. For some E [s), if m, is equal to some 2ry Tr. for j1, .......... . E ((22) we define that block as:
42
(rx)) (12nn(i-11d4111..
0g
2nxd
for all e [K], and
( (i- 02nxds
for all other j E (4Mr........... J.
3. However, if for E [s], monomial has degree S k - 1, let this be equal to 2r, Tras where Jk. [fo], k; < k is the degree (note that the variables are anything outside T, we have defined the corresponding query-key matrices to be zeros anyway). Let 8; be the integer which is the negation of the number of occurrences of this monomials in each of the monomials ordered higher preference than i.
As before, 8; sum defined by adding:
whenever < i and my of degree
- Sg whenever highest preference variable of m, is also present in m,.
i, the monomial my is of degree I k, m, divides my, and the
-1 whenever L < i, the monomial my is of degree < k and m, divides me, but the highest preference variable me also present in m,. in all other cases.
If Fr, is not present in monomial i, we just set Q (r)) (1:2n. (1-nn(i-1)attiii.
O2nxd:
otherwise:
s,a
8,02
Q (122((1-1))+1114) s:ae)
=
0g
0a
2nxx
where 2r, is the highest preference variable of m, and
a
(
a
B
0g
0a
0a 2nxd
for all other j's such that mry is present in monomial i.
46

---

Published as a conference paper at ICLR 2026
4. The last column block for Q is the all ones matrix 1xd with scaling factor, i...
(ri)
B 12nxds
(1
and for j (2,.... . ,4), the zeros matrix
(n))
0 Onxa [nxd] 2nxd
For the value matrices V E [R(2n)x(++1)4 j (1).. we define the first column as,
v (1:2n,1)
0
and for j E [2: QVT, we define the first column
- 1
(1:2n,1)
with every other columns 07
Correctness construction. construction does indeed give
Again, similar to the proof Lemma E.8
we can prove that this
0.!)) B
Σ
)
ji..J S
and the entries of the query-key matrices are in [-r, for B
< (s B).
Also, these query-key and value matrices can be computed in OOn time.
Approximation yields gap property. As before, let us assume there exists entry-wise approxi- mation iit of the (f1, 1)-th element of A(( such that
Pa RE.
< 7:
where
1 [1 00 2000
Σ
Pe.
exp
4 iET f E.22n] : jēT
Re Σ
(1) 0000 1e........ .000
exp
4E
and the (e1, 1)-th element of (tt(h) is
Pri
Re,
Bounds on Pe. RG,
Similar to before, we can prove
Aak Re, 2/-1 1- (141)
n 1
where X = 1 o(1). For the numerator, we show that when positive certificate of --yeryyprrrapphii does exist (if ri * 1, then this holds for all éss otherwise, there will be a fixed & such that af is included in the positive certificate).
P (2n)* k- 1 (=)))))
which implies
(n Mo ())
Ph RE,
(2n)a-1 7.
(17)
Otherwise, no positive certificate of the positive certificate, if exists, does not contain the vector at,, then
exists when ri o 1, or when ri = 1,
Ph < 2 k (11-0)) (c+1)))
My
and therefore,
Ps
(nt(2 cg) 1o) -))c01)) +
It <
+ <
(18)
Qi
47

---

Published as a conference paper at ICLR 2026
Wrapping up. In order maintain gap between the cases of an HypergraphIP existing, require the lower bound (Equation 18) must be less than the upper bound (Equation
2
+
Mo
Mo Mosts)
(5
Mo MotD))
(2n)
Now, there exist large enough constants Cay Co such that this inequality is satisfied for y S n C. and B 2 Co (log n///k.
This proves that any algorithm for entry-wise 7-approximation of Att(h) having maximum value of the entries r (2((lgg m)/ requires time ((n" assuming the Max-ASAT conjecture, since APAC (h) could be solved in ((nfo=6) time, then that would imply Max could solved
in n) time (Corollary E. 12 something that can not true for an absolute constant (Hypothesis 2
Remark Lemma E.13 for computing APAC when is variables, of degree and
contains subpolynomial elementary symmetric polynomial in to variables and degree k. the time-com complexxyyy is lower bounded by S2(n'). This is the strongest time complexity lower bound we can achieve, as the trivial algorithm for summing over the indices of all the query-key matrix also requires (n²) time and we say that this is the best we can for!
E.2.3 TIME LOWER BOUNDS FOR DEGREE POLYNOMIALS USING MAX-2SAT CONJECTURE
In this section, we prove the final part of Theorem E.3 where we show a lower bound for a certain subcase of h when the degree is 2. For the remaining degree 2 cases, we have already shown in Sections 3.2 and that they can be computed in O 2) time, which is essentially tight from Part I of Theorem
Unlike using SETH which proves lower bounds which are integer powers of n, in order to prove lower bounds of the form n² use the Max-2SAT conjecture (Hypothesis) 3) by giving a reduction from e-Gap-IPP (Theorem C.8) to entry-wise approximation of (t)(
The reductions work as, we first use the reduction of Max-2SAT to IPA, then reduction IPA to new problem IP-DIR--CYC using Alman & Vassilevska Williams (2020), which then is re- duced to its gap version containing (1) instances of e-Gap-I--IIR-CCYC Finally, we reduce C-Gap-IP-DIR--CCC to computing an entry-wise approximation of Att!h),
For these sets of reductions, we first define the new problem of IP-DIR-rCYC, which introduced in Alman & Vassilevska Williams/ (2020).
Definition E.14 (IP-DIR-rCYC). For a positive integer v, given r sets of vectors A Ar C (0.1)* with [A'| = ... |AAi n. and target inner products mi..... My E {O..... d} the problem IP-DIR- CYC,,a is to decide if there exist vectors 01 E A A'....,a ar E A" such that my-11 (0,,a) (1)= my.
simultaneously (a1, 02) m., (a2, a.) m2...., (0,-1, (0.2--.,,0.-1000
Naturally, to prove hardness of entry-wise approximation of poly-attention, we will again require the hardness of the gap version of this problem, E G-Ga--IP-DR--CCY
Definition E.15 (a- E--Ga--IPDDI--CCOO vectors A A..... A E 13d with |A' = = [A"l = n, and a target inner product m e
For every 0 and positive integer r. given sets of
{0. .......0 d} along with the promise that for all [r]. for all vectors as A and 0:41 mod A 1+11 mod
either (a,, ((++)) mod r) = m,
* or (ai, ((5-1) mod + < =)m
the problem of e -aa-----DR--CCCC. is to decide if there exist vectors A ...., a A such that simultaneously (a1, (2) (02, = (u,-1, ar) (ur,a1) = m.
At,...,
Now, we know that I--DIR--CYC at least as hard as IPA, which in turn is at least as hard as Max-2SAT (Lemma C.7 given by the following lemma. An OV version of this lemma proved
48

---

Published as a conference paper at ICLR 2026
Lemma 21). i.e. when the target inner products zero, by reducing OVA to OV-DIR- but all the proofs work similarly for reducing IPA to IP-DIR-rCYG well.
in (Alman & Vassilevska Williams 2020
Lemma E.16 (IPA IP-DIR-rCYO Alman & Vassilevska Williams (2020)). For every 6 > 0 and positive integer 2 3. if IP-DIR- CYCIR- can be computed in O( time, then IPAn,d can
also be computed in time (n.-6)).
Again, the 8 Gaa-----PIR--CYY is at least as hard as IP-DIR-rCYC using proofs very similar to Theorem C8
Corollary E.17. For every ò > O, positive integer 2 3 and every constant 0, assuming the Max2SAT conjecture, there exists a constant e > 0 and larget inner product m E (0....,d) ... such that 8- D-Gap-IP-DIR+CYCC n,. cannol be solved in time O(n=-6).
Proof. We prove the hardness of G-Ga--I-DIR-ACCC by starting with hard instance of IP-DIR-rCYC containing sets vectors A ,..., A" C {0, 13d, where n = |A4' and d clog
Following the technique of the proof of Theorem C.8 we consider A', mod for each i E [r]
as a 2IP instance, and reduce it to (1) many instances of E-GapIP having two sets n vectors of dimension dg. For the final instance of S-Gap-IP-DI---CYC. we create vectors having blocks. each block having the dimension dg- The (( 1) mod r)-th block and the i-th block in the final instances the reduction will contain vectors from each of the instances of E-GapIP obtained from the instances of 2IP from ( mod r, A' and A ((++)) mod respectively, while the other blocks
will be zero, exactly similar to the proof of Theorem C,8 This hardness result also true when all the target inner products the same.
Therefore, for proving the hardness the entry-wise approximation of poly-attention based on Max-2SAT conjecture, it is sufficient to start with hard instance of e-Gap-IP-DIR-FCYY we prove the lower bound for poly-attention for all polynomials that are not tree polynomials (since we already know that tree polynomials have exact computational complexity O(n²)). If polynomial is not a tree polynomial, the graphical representation must contain at least one cycle.
Further,
Lemma E.18 (e-Gap-IP-IR--CYC
(M)) For every constant E O, every å
(0, 0.01), every M > O, given an attention polynomial h(z1...., of degree 2 having monomials, such that its graphical representation contains a cycle of size r. where t,s.r are constants, there exist constants C. O and C. > O such that if APAC (2n, (r + 1)clog n, r Co =n -C=)) (Definition A.6 with query-key matrices Q!")...., Q(i) E [-r, 1)2xx(++1)cllgg and value matrices V( v(t) E (enx(s+t)p 1)+ can be solved in time O(n then G-Gap-IP-DD Dia-I---DR--CCC,,..o m,oog (Definition E. with target inner product m = M log ni can also be solved in 01)-55) time for any constant M.
Proof. In our final part of Theorem E.3 we reduce Max-2SAT to entry-wise approximate com- putation of poly-attention. We start with an instance of G-Gap-P-DIR--CCY, ns log since we know that this is at least as hard as Max-2SAT (Theorem E. Lemma 16 consisting of sets of vectors A A" S {0, 13 where A for all i E [ has vectors [af..... a each. The target inner product is M log n, and the constant approximation factor is 8 for the gap condition. This is equivalent to checking if there exists a) E At ,a?, E A?,...,,a" aa. E A"
such that (a), ) + (a?,,03) + + ( ag,)) + (a 1..00 a/ M2 log n, or, due to the promise, if (a)) ) ) + (az, a ....4 (a)_) af,)) + (a (a S (1 eo log where Mo O(M), ((e).
For the graph G of the polynomial, we consider a vertex Uoa where the cycle of length r starts. If there are multiple cycles. we consider any one.
Let the cycle be of length r be given by (Vt., vt)+1), (Uto+1s U2o+2).... (to++--1 v..), without loss of generality. When we construct the matrices Qi"'s, 's the idea is to construct the instance of
e- --Gap-IP-IR--CCYV from Uh, (i.c., from ((to)), and make every other query-key matrix corre- sponding variables outside the cycle be zero.
49

---

Published as a conference paper at ICLR 2026
Similar to before, construct query-key matrices such that for all l1 le [n], 2 h2 "...,......"""" - A((ah,g aign +...- (a (ootttr- obio ..) + (a In abo 1), (19)
for scaling factor A.
Constructing the matrices. We form the matrices , (r + 1)d as follows, using scaling factor B w(1):
Q( E e--r, nt 2n, d
1. For Q s, j
to or
just make the entire matrix zero (1+
2. For defining ((oo), we define the first column block (starting of the cycle)
02
(te)
= B
a
(1:2m.1:d)
0a
Oa
2nxd
the r-th column block (end of the cycle) as,
02
(fo) (1:2n,(r-1dd+te
an
B
0a
0 2nxd
the final block that balances the inner product as
(to) Q (1 14
Bl2nxds
and all the other remaining blocks as 02nxd-
3. Now, for the matrices inside the cycle, i.e., j E 1lo + 1, +r- 1], we define Q as follows. For = 1 (which is the traversal inside the cycle from 05-1 to By, and ny リ respectively), we define that block as,
(ta 1)
(to
02
e (1:2n(0--1)4411..))
= B
a
- (to 1)
0g
0a
2nxd
the final block as,
Onxd
lnxd 2nxd
and all other blocks Ozaxd-
50

---

Published as a conference paper at ICLR 2026
For the value matrices Vo) [ -E [tu, tu 7 1]. we define the first column as,
v (1:2n.1)
1
L
and for all other j's, we define the first column as,
v (1:2n.1) 1
1 7
with every other columns 07
Correctness of construction. We prove that indeed Equation 19 is satisfied when l,.... ,,. [m]] When we consider h, all the monomials containing variables 2; for j < to or tu + vanish since QU''s are zero. Whenever we have monomial of the form 1z0g+1 J E Mlo lo + r 1/. survives and gives (at f-1o+1 (j-fo+2) mod ((gooo)))) ((s-(o mod .)?
These query-key and value matrices can be computed in O(n time.
Approximation yields gap property. as Ft,, and we have
We again consider the entry-wise approximation of (tr f1.1
Pei Rei
< 7:
for
Σ exp L T1 1.........? 0!0) ).)
Pe
Σ exp ( 00 ........ .(.....(!.)
Re
when the (61, 1)-th element of (tt(h) is
Pe.
R
Bounds on Pri+ Re,- For the lower bound on Ra using calculation exactly similar to that of the proof Lemma E.8 gives us
Re 2t0-1 1o-t--111
When positive certificate for the given e-Gap- DIR-- exists, we will have some r-1 a
C2 10 fot11...., 60 totr-1 E [n]] for which (a a
+ (af
) + (ap p
Mo log n. This would give
Pe. > (2m) (2n) -o--o-r-a ( B²)
which implies
c (Ba Me )
Pt.
(20)
Ré
(2n)*
Otherwise, no positive certificate for IP-Dir- exists, then
Pt < 2 2 1, 10 (1 (Bb))) be))
and therefore,
r- It, < + < 2 (5 ( Mo )) 7.
Pr
(21)
Re
Note that if a positive instance of E- G----IP-DIIR-CCY exists. then øe, is the greater than the lower bound (it greater for all & if to # 1, otherwise we choose only that & for which the C-Gap-IP-DIR--CCC instance contains at ) otherwise always lesser than the lower bound.
51

---

Published as a conference paper at ICLR 2026
Wrapping up. In order to maintain gap between the cases of positive instance of G-Gap-IP-DR-+CCC existing, we require the lower bound (Equation 21 must be less than the upper bound (Equation 20)
240 20--r11
+7
( (r+l)e)
102 Me Wer+))
(2n)re (=²(-- My M/++))))
Now, there exist large enough constants Ca, Co > 0 such that this inequality is satisfied for S n -C. and B Covlog n.
This proves that any algorithm for an entry-wise 7-approximation of Att(h) having maximum value of the entries I' = B S2( v /log n)) requires time 52(n"), assuming the Max-2SAT conjecture, since if APAC could be solved in ((n=~)) time, then that would imply that would imply Max-2SAT could be solved in 2(~)a-0()))n time (Corollary E.12 which can not be true for an absolute constant 5 0 (Hypothesis
F PROOFS OF SECTION FUNCTION COMPOSITION
In this section, describe poly-attention mechanism whose attention head can simulate t-fold function composition. In order to study the representational powers, it is important to also consider the number of bits stored for each entry for the matrices, denoted precision, p. Since the entries usually considered be polynomial in n, it is safe to assume p (1))), Furthermore, usual, we consider the embedding dimension (log
Before showing the representational strength of poly-attention, first show that Strassen-attention and 3-tensor attention cannot simulate 3-fold function composition. For this limitation result, we require communication lower bound proved in previous work of Chakrabarti (2007) myopic pointer jumping.
Definition F.1 (Myopic pointer jumping). similar to function composition, where we are interested in computing t-fold function composition, for imputs as functions fi,..., J.. [m] + [n] and a value x E [n]]. There are t players and a coordinator C, such that:
every 2 2, myopic pointer jumping can seen as
Plaayer 1 has as inputs x and fz.
Player i for E [2: 2t -1] have imputs x and fi...., fi41.
* Player t has inputs x and fi.....f--J
The Players E [1] can only send messages to C, and the goul of the protocol is for C to compute the value of fe(fe-1... (x)).
Now, the lower bound due to Chakrabarti 2007 for myopic pointer jumping is given as below.
Lemma F2 (Chakrabarti 2007 Theorem 1)). To solve the myopic pointer jumping problem, the players need to send at least Q(n/t) bits to C in order for C to compute f. (x)))
We want to study the representational strengths and limitations in terms of function composition. We say that an attention mechanism simulates t-fold function composition, if, given the input X R [tn+1)xxd containing descriptions of f1,...,f and an E [n]. the attention mechanism is able to output the value fe(f-11... ... fi (x)). As before, the input function f. will given as the i-th block of X, in X(G X((G-1nx+1:.n) for all E [t]. and will be given in X 1. and we want the attention mechanism to output the value of fe(ft-1... fi (x) in the (tn + 1)-th entry of the output.
The first limitation result, Strassen- attention can not simulate 3-fold function composition is given by: Theorem F.3. One layer of Strassen-attentinn attention requires least H > nl-((1) heads simulate 3-fold function composition.
Proof. Let us consider an instance of 3-fold function composition where, given fi. [m]] [n]. and x. [m), we want to compute (z))). As usual, the input X contains N 3n rows
52

---

Published as a conference paper at ICLR 2026
of embedding dimension d OOlog n.) where ((mm) corresponds to the values of f1(1). (n)). X (x+1:22m) corresponds to the values of f2(1), f2(n), X (2n+1:3m) corresponds to the values of fa( and finally X3n-1 corresponds to 2.
The main idea for proving this lower bound is by assuming that Strassen-attention can simulate 3-fold function composition using H heads. We are given the query-key and value matrices for H Strassen-attentisn heads such that the output of mechanism contains the value of (s(ff(fi Using these, we define a communication problem which will use computations required for outputting the matrix Att(S), that gives the value of fa(f2(fi (x))). Next, we will use existing lower bounds (Lemma| F22) to contradict this statement, which would give lower bound on the minimum number of heads of Strassen-
required to compute fa(fa(i (x))).
We now define the communication problem capture this setting. Consider 3 players with inputs,
Player has x; /2,
Player 2 has x; f1, f3.
* Player 3 has v. f1, f2.
and coordinator C. The communication channel is such that only the 3 players can send messages to the coordinator. The communication complexity is the total number of bits sent by the players to the coordinator such that the coordinator can compute the value of fa (f)()())))
As defined before, this communication setting is an instance of myopic pointer jumping for = 3, and the lower bound from Lemma F.2 implies that at least (()) bits need to be communicated.
Now, let us assume that there exists a composition using H heads, where will denote the index of the head superscript The weight matrices for query-key are WQux* W Wore)" W WQUs."
mechanism that computes 3-fold function
[H]. E Rdxd and the value weights are Wve", Wvo) E Rddd for the attention head u E [H]. Let the precision of the values be These matrices and the functions computed by the first and last MLP layers are known to the players and the coordinator. Assuming that Strassen-attention can simulate 3-fold function composition, we devise a communication protocol for the above problem using the value of Att(S) to obtain lower bounds on H using a proof inspired by works of Peng et al. (2024); Sanford et (2024b).
The output matrix of the w-th head of Strassen-attention, (S) for [M], is given as
EjRe(M
N (X, Wy Jr @ (X& Wyor") Ja
(5r (Sh
(22)
Nu "..""
where we have N = 3n + 1, which is the row of Att(S) where we want the value of fs(f2(f(())). and
N
exp (Xant (Woo"""W""""" + X, Woa" "Wao"""""""""" j.k
+ X W Wo ( (Won)"""X"
for all heads E [H]. The players have parts of X, i.e., for Si they have X (1::), for f2 they have (11+1:2n) for fa they have ((2n+1:3m) and for x they have X3
The communication protocol proceeds as follows, where the player sends the values for each Strassen- attention head E [H]:
Player sends L." 2 and L where Li" is an O(plog log ni- approximation of the binary expression of LP", and L is an O(p log log n)-bit approximation of the binary expression of Li", where
L
Σ
jES R
S S
((nn
53

---

Published as a conference paper at ICLR 2026
and
1
Σ JES ((5+t))),,,
(X, o (xWWvss")
TT"
for all E [H], to C.
2. Player 2 sends L3" and L2 where L2" is an O(plog log n) approximation of the binary expression of L2", and L2 an O(p log log ni-bit approximation of the binary expression of L where
L =
j,k
JES KES; St 1:3mn (3x (8m+11)
and
L
Σ
N j.k
L
JESt.KESt
S1 (3n (m(m)] (2 1:3n (2) (3n+1) 1}}
for all u E [H), to C.
3. Player 3 sends L3"" and L where Ls" is an O(plog log n)-bit approximation of the binary expression of L3", and T an Ogp log log nb-bit approximation of the binary expression of L3", where
L
Σ
Nu 3ik
jES, KES, S1 [n):2n+) 31732
and
LA
Σ JESI
Nu (X, (XXWVV" (XxWye)))
Lid'
S1.82 E ([m).
$1782
for all [H], to C.
4. C computes
Σie
Tho T.P. TRo.T.p
E Rdd
(23)
Evety/ L2"
as the N-th row of the (8) matrix.
Note that Equation 23 is the correct value of the approximation of Att (9),, (tt for all u E [H) since the values of L," L'"" are simply the partial sums, all of which amount Equation 22 with the given bounds on each of the summations. Sanford et al. (202bb) showed that using (p log log n) bits of precision is sufficient in this approximation, and this gives the correct value of [2( (r))) upto p bits of precision. The number of bits communicated is equal to O(dpH log log n), and using the lower bound from Lemma
F.2 we must have dpli 1(n/(Io log n)). Since we usually choose Oloog n), p (11), must have, the number of heads, H n--o(1)
Corollary E.4. One layer of 3-tensor attention requires at least n - heads to simulate 3-fold function composition.
Proof. The proof is very similar to that of Theorem where again we have players and a coordinator in myopic pointer jumping instance. Using the construction of 3-tensor attention, can again infer that the communication complexity will be Odpp log log N), which needs to be greater than ((n) from Lemma F2 This gives result.
54

---

Published as a conference paper at ICLR 2026
In fact, can show stronger result.
Theorem F.5. If h can be written as a variable separable polynomial, where each branch (see Definition A.5) has I in variables, then one layer of poly-altention for requires at least H n² ol1) heads to solve lo-fold function composition.
by constructing communication protocol for tu-fold function composition if poly poly-attention for can solve it, and using the lower bound result of Lemma F22 The input X contains N = lon + tokens, and we want the output to be in the last row of Att ( for each head E [
Proof. We use the same proof as of Theorem F.3
We define a communication problem again that of myopic pointer jumping, with to players and a coordinator C' who wants to compute fto (fto-1 -1.0--.........1.. ... fi (x)) (Definition F.I). Since tu is constant, Lemma F.2 states that this requires (2(n) bits of communication.
Now, we develop a communication protocol for function composition using the Att matrices, Yu E [H] which will have a communication complexity of O(Hdp log log N).. In computing the output of the poly-attention mechanism at the last row of (M we have the numerator term as
Σ
exp(h(Q
(he Q
(take )) (2) 3 ..... G. (yttay
V
f E(N)
and the denominator term as
Σ
exp(h(Q
(1he ohe
Que ((al))).
t [[N]
If the polynomial h is variable separable and has branches, where each branch is given by the polynomial g: 2x) having S to variables each, i.e., hh(11....,.2 v2) Eierr] g=(#1, then players devise a protocol to separately compute the (ton + 1) row of Attle:) for all i E [r]. Similar to the proof of Theorem F33 the summation of f2...., lo E [A] will be broken down to partial summations, which correspond to computations performed from the inputs of each player.
In computing the poly- poly-attention output of each branch (both numerator and denominator as in the proof of Theorem F.3). let the corresponding variables of that branch be I Fre. Now, Player would send the summations of fr,....,hy ,,, E [n +1 2n] U {ton 1}. Player would send the summations over l .a.....,eyy .rr, E [m] U [2n + 3n] U (fon + +1} except the tuples that have already been sent, and so on until Player i would send the summations over ......., br., E
I)n + 1] [ī.n +1: 1:::::0 (i + 1)n] U fton + 1} except the tuples that have already been sent. Since there are to 1 variables that are not fixed (& is fixed to N) and all the players with their given inputs completely cover the summation required in the softmax computation of Att(A)
In this way, the players can communicate (H log log N) bits before to compute the value of Attlady, y for all i E [r] and u E (H), and given the poly-attention outputs for all these branching polynomials, the coordinator can compute the value of (h) using Lemma D.1
Therefore, with total of dp ppp log log N) bits (since the number of branches, of the polynomial is constant), the coordinator will able solve folf function composition. By Lemma F2 Il dp log log N M(n), and considering O(logn), n), (1)), we require H n1-o(11)
Next we prove that certain class of tree-attention, given by polynomials of the form hy(m1,-.., x(41) = D2 + T2T3 +... + TyT141 can simulate i-fold function composition. This proves Theorem 3.4 which also the generalization of Theorem 3.1
Theorem F.6. For every integer 2 2, poly-attention for the polynomial
he(z,... .x) x142 + 3233 +.. + T:
can simulate t-fold function composition using one poly-attention head.
Proof. For solving the problem of t-fold function composition, we consider the functions fi..... f.. [n] + (n). The input (before the first MLP layer) is sequence of numbers ((1)...., (bnn+ [m].
fa:
55

---

Published as a conference paper at ICLR 2026
such that for & [n]. j [t), we have oll + 1)t) f (C), and finally Φ(3 + 1) Our task to compute the value of fe(fe-1... f1(x)), and we give a construction of the MLPs, the query-key weights and the value weights of poly-attention for he, such that this Transformer layer compute the same using only one head. We adopt the construction of Kozachinskiy et al. (2025) due to its simplicity, and use it define the parameters of poly-attention.
We define the first MLP layer such that its output, i.e., the positional encoding the i-th entry of the input to poly-attention, is given by:
X; =
(())
(o(1))²
Oak- 5| 1x33
for i E [tn + 1]. Here, a precision of (log n) can be used. Next, we construct the weight matrices W :ow...., W Qou.
Our goal is to create a them such that
he( ) -A² log n (Ó(()) - -82)² + (43 - 2(22))
+ (e1 2n rb(x)) + + (l++1 - (1- 1-m b(()))))
(24)
for constant A > 1. For L = tn + 1, this is maximized when
le (((1) @tnn +) 1-=1)=a =.
l = n + (((2) = n =n+ fi(z) f2(f((x)).
4+11 (t- n x(((-1) = (t fi 16 (())).
which is precisely our required value.
For constructing Q E RRxr3r for j E [t + 1). with such properties, we can define each row as:
1. for i = 1:
200 T
Q A Vlogn
6(E)
LOT 31-33 3kx1
2. for j 2 2:
03(j-2) 2)
2( (j 2)n)
(E ( 2)m)
00 = Av A log
(
((E)
O (2(0-))
3kx1
for all l E [n].
Note that, for any [t].
QQ) Eg = -A" - -o
which is consistent with Equation 24
While computing the softmax entries for é = in + 1, the Qu for all f2,..., (.11 that do not maximize this value, will be factor of n A less than the maximum value. Since while computing softmax, we take a sum over all f2...., l+ E [tn + 1]1 as long as we choose > l(VT). the maximum value will be obtained in the correct setting of ł
value of hr (Qt++1 (1)) ((2)
56

---

zeros. The error in the final output will be n²-A², and as long as this is less than the number of bits Published as a conference paper at ICLR 2026
For outputting the value, set the first column of all the V ss [2 t). ones, and the rest as zeros; and for v(++1) we define the first column as (-(++1) = L, for all E [tn + 1]. and the rest as ',1
of precision, we have the correct output.
As we will see, even though poly-attention for h2 will be able solve t-fold function composition. the previous theorem, Theorem F.5, shows that not only poly-attention for f4-1 can not simulate t-fold function composition, but neither can the poly-attention for the polynomial ..,11,,..,,,.)).. = 22 + T2 + I1 + 1114-11e42, which is polynomial in + variables!
Remark 2. From Theorem F.6 saw that poly-attention for ( 22, 22. 22 T23as can simulate 3-fold function composition just as Strassen-attention. Again, Strassen-attention is poly- attention for the polynomial h(21, D2, D2g 21 + I2Ts + Iz which is just one monomial different from h2g However, even though they might seem similar, the cost of this one monomial is huge- Att can be computed in (n²) time, while computing Att(S) requires at least 22("*) time.
G PROOFS OF SECTION 3.4
POLYNOMIAL ROOT-FINDING
In this final section of the proofs, we prove the strong characterization of representational strength of poly-attention introduced in Section 3.4] We show this by giving a construction of the weight matrices of poly-attention mechanism which solves polynomial root-finding (Theorem 3.7p.
In this problem of polynomial root-finding, for fixed polynomial p(æ1....,, .a) and given as input a set C 12", we are interested in finding if there are elements y1..... Wt E S such that p(y1,.... M/h) = 0. For the output, if vi. ...,, Me is root of p and Slj] = y then in the row of the output, we want to output an encoding of that root.
Theorem G.1 (Polynomial root- foot-finding,, For a polynomial p(æ1...., 24) of degree ko and given an input S R", for any integers k; 8 if polynomial h(x1,-.., v4) of degree l and sparsity is such that all the monomials of the polynomial p² divide at least some degree k monomial of h, then poly-attention for h with attention heads can perform polynomial root-finding for p with the imput.
Proof. We give a construction of the MLP layers, query-key weights and the value weights such that the Transformer can find root of the polynomial from S", and output it. First, given S, considering so the sparsity of p², we set the embedding dimension as d = 80.8. For the input X (ax((e))), let the embedding of X, after the first MLP layer be
X = [1
U1
2ND Di
0 2ho -1h0-111 (x(50.9)
where we require s0.s
> 2ko + 1.
Construction of first head. Now, our goal is to define the weight matrices such that after computing the query-key matrices Q!),..., where y; = -(],, and f,,..., [m]]
((), the value of (.......... will yield - -n2 -n?p(y,,.... ye,))
Choose a h (21....,, 224 of degree i (where i is number greater than the maximum number of variables in each monomial of p²), and is of any sparsity (satisfying su.s > 2kg), where each monomial of p² divides at least some degree & monomial of h. We assign each of these monomials of p² to exactly one degree k monomial m, of h for i E [s]. and we associate a set T; which stores all the monomials of p² that are assigned to this monomial m, of h.
Now, define Q())...,
...., (it) (=o.() where each column block is of size so, as:
1. For the i-th column block, where for each column j [sp] of the block, we consider the exponents of the variables p such that ((Q(1) Q..,))! will give evaluations of the j-th monomial of -p² at .y....... yr.), for all f,...., lG E [n]. For these values of i,j, we will simply denote these terms the monomial corresponding to this column (i 1)80 + j.
57

---

Published as a conference paper at ICLR 2026
(a) If the j-th monomial of p², for j E [su], consists of k, variables, and is in T; for some dry;
i E [s], let Cyæri" dn ..., be this monomial where 2r, is the highest preference variable. Then, we define the j-th column the i-th column block of Q(1i) as
Qia (im,(i--1)ott = 7
des
Cyyn Ur
and for < S kj.
drg
Q (lra) (im.(i-1)00-3 1= n
yn
For all r E [t] such that z, is a variable of m, and the j-th monomial of p² does not contain xy but is present in T, define
@ 1mm == n.1ns
and otherwise, if v, is not present in m
O lmm On
(b) If the j-th monomial of p², for j E [so], is not in T, then we define
(r) (1. (i=,,(i-1)sot))) On:
for all r [t].
2. Fixing an i such that m, is of degree < k we define the query-key matrices as before, to cancel out the terms which were defined in the degree k. Each degree term had so terms which could lead to non-zero values, and now for the block i, corresponding to the monomial i, the r-th column in that block will cancel out the j-th columns of each block obtained from the degree k-terrms, for j E [su].
Let s/ be the integer which is the number of occurrences of j-th monomial of p2 while com-
puting the monomial containing variables 21.... 2e corresponding to m, (Q(1) ..... (0) when consider each of the monomials m, ordered higher preference than (i.e., 8 << i), and divisible by m,
As before, sf is the sum defined by adding:
-C, whenever l < i, my divides my, degree of me exactly kt, and the highest priority variable of my is present in mi;
-s whenever < i, m, divides me, degree of m is less than k, and the highest priority variable m4 also present m,-
otherwise when the above conditions are not met but my divides m. in all other cases.
For every j E [so]. if x, is not present in monomial j of p² for r E [t], we just set
( (1m. On:
otherwise, for the highest preference Ir, variable of the j-th monomial of p², we define:
d., s V1
d.,
Q (nn)
s U/2
(in,,(i-1)oott) =
dri
8,yn
n. no.8
58

---

Published as a conference paper at ICLR 2026
and for all other such that 2, divides this monomial,
81
Q (1:n,(1-1)0005)
//2
Notice that in these constructions, we have only used linear combinations yę for E [1] and q E [2ku]. Therefore, weight matrices Woo E (*6.s)) exist for every fixed polynomial p such that
1/
2kg
kg
1/2
2/2
W Q++))
9n
y ko
0
0
yield the required (rr)ss For defining the value matrices, for the first coordinates, the r-th coordinate of (y),, E [2 t] stores the corresponding value of Σy, and all the other entries are of the coordinates in [2 {r] are one, and the first coordinate zero. More specifically, we define
0
1 V1
0
0
122
V
:
yn
1 0
where the r-th column has the values of the ye
since the degree h monomials of h are what contribute to (ye, ye,) from the corresponding column blocks. Inside each of these column blocks corresponding degree l monomials of h, the j-th column for j E [so] gives the value of the j-th monomial of -p² at (ye,.... )r.) Due to our construction, all the values of m. m²(Q(() ......QQ Q(!)) are zeros when m, 's are of degree k, which finally gives us the required result.
Using the construction defined above, we have h(Q (1)
Q!!)) = -n
(yc..... yc., -
Now, for each fixed é,, the value of h(Q ..... ) -p ,t.) which maximized for some indices ......, E0, is at least e"" factor larger than all the other values in the summation Ees....eein MQ ... With the given construction of ((r's, the values of veg.--.. Veo for which -y? +) is maximized, will be present in the first coordinates the output Atr. (N)
Construction of second head. Finally, we need to verify that if there exists some el such that the values of ,,, encoded in Att(e) indeed is root of the polynomial. For this, need the value of yr, 's for each of the 4i coordinate, and we incorporate this by using a second attention-head whose output matrix contains the vector [y1
ym.)" in the first column and all zeros elsewhere. Therefore, when we add the two attention heads, the fi row will contain values of (yE,,...,, Me. which maximizes the value of -p²(ye, +). Finally, can check using the output MLP layer if indeed the value is a root of the polynomial.
H EXPERIMENTAL DETAILS
H.1 FUNCTION COMPOSITION
In this section, explain the experimental setup behind Figure 2 We train transformers that uses self-attention for one layer, a transformer that uses self-attention for two layers, as well as transformer that uses tree-attention for one layer, for the attention polynomial h(21, (2,1,,2,23))
59

---

Published as a conference paper at ICLR 2026
1142 + 2223- (This is the polynomial from Theorem 3.1 findings that tree-attention better: is faster, more learnable, and uses less space compared to its representational counterpart, the two layer self-attention.
above.) We infer from the experimental
In the remainder this subsection, we first explain the details behind Figure 2 which shows that despite having less trainable parameters than two layer self-attention tree-attention is more learnable. Note that two layer self-attention requires two query matrices, two key matrices, and two MLP layers, while tree-attention requires only three query-key matrices and one MLP layer. Second, show that the time to compute tree-attention is comparable to the runtime to compute two-layer self-attention.
Problem set-up. We solve the task of function composition, where, for an integer n., given two functions f1, [n]] [n] and value E [ni]. interested in computing the value of f((2(z)).
We know that two layer transformer using self-attention can solve function composition but one layer can not Peng et al. (2024), and we further proved in Theorem 3.1 that tree attention can solve as well. We show that these theoretical results are in line with practice, where transformers with two layer self-attention as well as transformers with one layer tree-attention can both solve O-function composition for
(which means the number of tokens is + 51).
Input generation. As described above, we train the transformers to learn f1((((z)) where fi, [n] + [n], and x E [n], for 25. The functions f1, f2, and the value x, are generated uniformly at random from the set [n] for each batch in each epoch.
With these f1, f2 and x. the input tokens to the transformers are given by a sequence of 2n + 1 tuples, (i, (i)) for E [n]. followed by (i. f2(i n)) for [n 2n] and finally (2n + 1,
We define separate embedding tables for each of the index (first element of the tuple) and the value of the function (second element of the tuple), given by embedding matrices E, E IR (2)++1)xx and Eg E Raxd where is the embedding dimension. The input the self-attention for the token (i, j)
is then given as E, (i) + E2(j).
Architecture details. We choose sequence length of 51. The transformer has an embedding dimension 32, number of heads H = 4, followed by MLP layer which uses ReLU activation with one hidden layer of size 128. We also use the standard sinusoidal positional encoding from Vaswani et al. (2017), given by
PE122 sin
10000004
PE cos
100000%/
for i [n]. (0.....dyy (O....,d/2), which is added the i-th token.
Training details. For learning, model using an Adam optimizer. The model is trained for 100, 000 epochs on 2024 Apple Macbook Air with M3 Chip, and the evaluations have been shown in Figure and Figure
batch size of 64, learning rate of 0.001 and train the
60

---

Published as a conference paper at ICLR 2026
Frit Lna Compu
24
50
35
Tweatherier
kaya
Figure 3: Training loss per epoch, averaged over 10 seeds, for learning f1(f(())) for sequence length 51, on single layer of tree-attention, one layer self-attention and two layer self-attention. Tree-attention learns faster and has less fluctuations.
Test Accuracy Comparion toon
02-
04
00
0.0
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.
Observed inference running time. We plot the running time of computing various attention schemes for sequence lengths in {20, 50, 100]. We use vocabulary size = 32, embedding dimension = 64, number of heads 4, and hidden layer width 256 for the transformers, and evaluate on batch of size B 64.
With this architecture, we randomly choose query, key and value weights in Rdxd and random weights and biases for the MLP layer. Then randomly generate 1000 sets of inputs X E R Bxnx and compute the running time of the attention mechanisms. The average running time has been depicted in the following table.
Seq len
I-layer SA (ms)
2-layer SA (ms)
1-layer tree (ms)
1-layer 3-tensor (ms)
1-layer Strassen (ms)
20
1.076 + 0.057
1.775 = 0.085
1.367 = 0.057
1.442 + 0.062
1.593 +'2 0.086
50
1.079 + 0.048
1.757 = 0.060
1.363 - 0.055
2.911 + 0.044
1.594 + 0.088
100
1.080 + 0.048
1.781 = 0.097
1.374 = 0.060
13.813 + 0.051
3.395 + 0.081
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 the time complexity
computations are small.
Discussion.
We obtain the following conclusion about tree-attention from these experiments.
One layer tree-attention successfully learn function composition, despite having only three query-key matrices and only one MLP layer (compared to two-layer self-attention that has two query matrices, two key matrices and two MI. layers).
61

---

Published as a conference paper at ICLR 2026
One layer tree-attention exhibits better learnability for function composition than two layer self-attention as Figure 2 since accuracy increases faster for tree-attention.
aree-attention has an efficient inference time. From Table 5 we infer that it has running time comparable to self-attention, and our cases, even outperforms two layer self self-antention.
H.2 COGS DATASET
To give further evidence of the practical advantage of tree-attention, evaluate two simple models (one with self-attention and one with tree-attention) on a benchmark NLP task, the COGS dataset Kim & Linzen| (2020). We compare the models, which differ only in which attention mechanism they use, and evaluate the difference.
COGS dataset which challenges the model perform composition based task, in which must parse sentences into fragments, and understand the relationships of the different fragments throughout the sentence. In our experiment, the words the sentences, along with special characters, are input to the transformer after encoding by pre-trained tokenizer, and the error is computed on the output which is expected be a semantically parsed sentence.
Example 1. Input: melon was given to a girl by the guard.
Target: guard x x_ melon
AND give. theme ( 3, x_ AND give. recipient (x AND girl (x
6) _ AND give agent (x (x_3, 3,
From the above example, the transformer is supposed to figure out that 'guard' is the subject, and is present at position 9 of the sentence, where indexing starts from 0. The other nouns are 'melon' and girl', in positions I and 6 respectively. The verb "give', present position has logical forms given by a theme 'melon' (position 1), recipient 'girl' (position 6), and agent 'guard' (position 9).
We trained both tree-attention and self-attention on the training set of COGS, and tested them on in-distribution test set as well a generalization test Including generalization test set is important to make sure the model is not just memorizing the distribution. We compare the exact token match accuracy for both, and infer that tree-attention performs better the generalization set than self-attention. This gives strong evidence of the inherent ability of tree-attention to solve composition-relaan relatte tasks.
Figure 6 shows the table for the final accuracy results. The learning plots have been described in Figure 7 where tree-attention almost always out-performs self-attention. The training was performed over 10 randomly chosen seeds, and we see in Figure 8 that tree-attention achieves considerable performance of around 30-50% accuracy on almost half of the seeds.
Implementation details. We use simple 3 layer encoder-only transformers, having embedding dimension 64 having 4 heads, and an MLP with a hidden layer of size 256. Both transformer models (using tree-attention and self-attention), were trained for 200 epochs with a batch size of 32 (755 batches per epoch) with learning rate of 0.001, and tested on the in-distribution test set and the generalization test set. The results for inddisriibttinn test token accuracy were similar, both giving an exact match of around 97.5%. The generalization set accuracies have been plotted as follows.
Tree-attention
Self-attention
Generalization token accuracy Generalization exact match
0.727691 = 0.013486
0.723993 + 0.008649
0. 0.264919 = 0.127609
0. 0.239024 + 0.087351
Figure 6: Table for mean accuracies and standard deviation over 10 random seeds.
62

---

Published as a conference paper at ICLR 2026
Generalization Test Taken Accuracy
Genaralization Test Exact Match Accuracy
0.74
0.72
0.3
Exrit
0.88
Setharth tokten
Sallattn exart
0.66
23
100
125
150
175
200
25
1:0
130
175
Figure 7: Plots for mean + one standard deviation over 10 random seeds for token accuracy and exact match accuracy the generalization set. Tree-attention has higher accuracy than self-attention.
Ties sttemtonn
Sclf Exact Match
Sauri
05-
SAn
Sand 10
Sead
Seed
0.4
Netch
0.
10B
125
12.0
200
120
Sporne
Aports
Figure Exact match accuracies with each seed on the generalization for tree-attention and self-attention. Tree-altention reaches ~ 40% accuracy for 4 out of 10 random seeds.
Conclusion. From the learning experiments, we infer that tree-attention is more expressive when it is used to solving composition based task. As can be inferred in Figures and g tree-attention noticeable performance benefits for several seeds, which calls for future work to explore learning heuristics for further strengthening the results.
I
THE USE OF LARGE LANGUAGE MODELS (LLMs)
Large language models have been used to find related works, and to polish the code for experiments.
63


Reviews:
[{"review_id": "TOAwUsGtE8", "reviewer": "Reviewer_Pxmp", "summary": "This paper introduces poly-attention, a generalized class of self-attention mechanisms that incorporate higher-order (tensor) computations and arbitrary relational structures between input tokens. The authors show that existing higher-order attention mechanisms, such as tensor attention and Strassen attention, are special cases of poly-attention. They provide a systematic analysis of the computational complexity and expressive power of these mechanisms, including new algorithms and matching lower bounds. A key contribution is the introduction of tree-attention, a subclass of poly-attention that can be computed in quadratic time and can perform multi-fold function composition—a task beyond the reach of standard self-attention and prior higher-order variants under similar time constraints.", "questions": "1. Could the authors evaluate tree-attention on more complex or realistic tasks, such as logical reasoning, mathematical problem-solving, or long-range dependency tasks, to better demonstrate its practical advantage?\n\n2. How does tree-attention scale with sequence length in practice, especially when n is large (e.g., >10K tokens)? Are there any memory or numerical stability issues?\n\n3. Could the authors provide more intuition or a high-level overview of the lower-bound proofs, especially for readers less familiar with fine-grained complexity?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality: The paper introduces a unified framework (poly-attention) that generalizes several recent higher-order attention mechanisms. The introduction of tree-attention and its theoretical and empirical validation is novel and meaningful.\n\nQuality: The technical contributions are substantial. The authors provide rigorous complexity analyses, including both upper and lower bounds, and support their claims with proofs and experiments. The connection to fine-grained complexity (e.g., SETH, Max-2SAT) is well-executed.\n\nClarity: The paper is generally well-written, with clear definitions and a structured presentation. The use of graphical representations for tree-attention is intuitive.\n\nSignificance: The work addresses a key limitation of transformers—their inability to capture higher-order dependencies efficiently—and offers a practical path toward more expressive and scalable attention mechanisms. The results could influence both theoretical and applied research in efficient transformer architectures.", "weaknesses": "Experimental Scope: While the paper includes an experimental validation of tree-attention, the evaluation is limited to a single task (function composition). More diverse benchmarks (e.g., on standard NLP or reasoning tasks) would strengthen the practical relevance of the proposed mechanisms.\n\nPresentation of Lower Bounds: The lower-bound proofs, especially those based on fine-grained complexity, are highly technical and may be difficult to follow for a general audience. A more intuitive explanation or summary of the proof techniques would improve accessibility.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TrnxDKJIqj", "reviewer": "Reviewer_8kRi", "summary": "The authors unified some of the existing attention variants into a new attention variant (poly-attention). In addition they showed that (with methematical theorems and proofs) under some conditions poly-attention can \"efficiently\" represent some functional forms that self attention, higher-order attention and strassen attention cannot represent \"efficiently\". One such example is the function composition problem. They showed that 2-fold function composition can be represented by poly-attention with $O(n^2)$ computational complexity while exiting attention variants require \"superquadratic\" time complexity. At the end they also provide an experiment which shows that poly attention learns 2-fold composition problem faster (in the sense of epochs) compared to 2-layer self attention, in the same experiment thay showed that 1 layer self attention cannot learn 2-fold composition problem.", "questions": "I feel that every poly-attention head defined in the paper can be written as a (single) t-order attention head, if we allow the tensor head to use a feature dimension that is $s$ times larger, where s is the number of monomials in the attention polynomial h. Here is a rough proof sketch\n\nLet\n$$\nh(x_1,\\dots,x_t) = \\sum_{i=1}^s m_i(x_1,\\dots,x_t)\n$$\nbe the attention polynomial used by the poly-attention, where each monomial\n$m_i$ uses some subset of the variables, e.g.\n$m_i(x_1,\\dots,x_t) = x_{j^i_1}\\cdots x_{j^i_{k_i}}$.\nFor the input $X$ the poly-attention first forms the usual projections\n$Q^{(j)} = X W_Q^{(j)}$ and $V^{(j)} = X W_V^{(j)}$.\n\nNow consider a $t$-tensor attention head whose key vectors have dimension\nsd (that is: s blocks, each of size d).\nFor each slot $j=1,\\dots,t$ we build its key as follows:\nfor monomial $m_i$, if $x_j$ appears in $m_i$ we put the row\n$Q^{(j)}_{\\ell,:}$ in block $i$; if $x_j$ does not appear in $m_i$ we put\nthe all-ones vector in block $i$ (which can be acquired by adding biasses to the Q, K projections).\nSo every key is just $s$ copies of either “the real query” or “all ones.”\n\nWhen the tensor attention takes the order-$t$ inner product of these keys,\nthe product over block $i$ multiplies exactly the variables that $m_i$\nwanted, and multiplies $1$ for the variables that $m_i$ does not use.\nTherefore the tensor score is\n\n$$\\sum_{i=1}^s m_i\\bigl(Q^{(1)}_{\\ell_1},\\dots,Q^{(t)}_{\\ell_t}\\bigr)\n= h\\bigl(Q^{(1)}_{\\ell_1},\\dots,Q^{(t)}_{\\ell_t}\\bigr),$$\ni.e. the same score as the poly-attention.\n\nFor the values do the same trick. Hence the tensor attention output, restricted to the first $d$ coordinates, coincides with the poly-attention output, for every input X. This shows that, up to a factor s increase in feature dimension, poly-attention does not define a strictly larger class of functions than standard t-tensor (higher-order) attention. At first sight this increase in feature dimension might look like an artificial blow-up. However, throughout the paper the authors do not argue that d must stay small; d is simply the embedding size chosen by the model designer. In that regime, allowing a constant (or even moderate) factor increase in d in order to simulate poly-attention with a standard $t$-tensor attention is not a strong penalty. On top of that, in practice we almost never need exact equality of attention scores. Since the model is trained end-to-end, it can approximate the desired score function well enough with fewer parameters than the exact tensor construction. So the sd-dimensional tensor above should be read as a worst-case. A trained model would likely get close with a smaller dimension. That makes the “poly-attention is strictly more expressive” claim weaker in a practical sense. \n\nShortly, for any given poly-attention head, there are parameter settings of a higher-order/tensor attention head that can reproduce it. So poly-attention is not a strict superset of all higher-order variants; it’s another way to specify the same class of interactions. Thus my question is did you carefully consider this possibility while writing the paper?\n\nAnother caveat is practical: in real Transformer setups we usually don’t know the right interaction structure in advance, so fixing a specific polynomial might be brittle. From that viewpoint the current architecture looks more theoretical, and it would be good to follow up with experiments to show when/if poly-attention actually helps in some general situations, which brings us to the weakness W2 I mentioned.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper proposes a \"unified poly-attentin\" framework that organizes several high order attention machanisms. That seems useful especially for follow up theory work. I like the the represantation ability comparisons without ignoring the computational complexities. I didnt check every proof and the proofs in detail, but the upper/lower bounds story is believable and it explains why other attention variants blow up to superquadratic. Overall exposition is quite clear, mapping the special cases early helps both theory minded and experimental readers.", "weaknesses": "W1)About the motivation (I feel this is crucial): \nAlthough it can possibly achieve other functions in the paper the main motivation/achievement of poly-attention is stated around the functional composition problem (even your \"punchline\" is around it). However, you should explain more clearly why this particular functional form matters. What makes it a good choice? In principle, we can always cook up an unusual functional form and then design a neural network that works well for it. Because of that, we need some justification that this is not just an arbitrary example but a form that’s actually important or likely to appear in practice.\n\nW2)About the experiments:\nThe experimental \"validation\" seems strange. All of the paper is about representational abilities of some models but the experiments are about the learning dynamics. Thus the experiments are not specifically about \"validating\" the theories. Seeing that the main thing about the experiment is comparing 2 layer self attention and poly-attention learning speed, they are just experiments themselves. However, that does not mean that these experiments are not relevant. Seeing that the paper does not have any theory about the learning dynamics that makes the provided experiments even more important, so I feel authors should improve them: \n1)The learning curves in the figure 2 and the appendix figures seem unusual. The accuracy slightly increases at the begining stays constant for some time then abruptly increases again and the epoch number it abruptly increases the second time seems \"random\" at first look. Therefore I strongly suggest repeating the same experiment with multiple random seeds (this experiment should not even require much computation the networks are very small so I am sure it can be repeated very easily), and eventually provide final plot with confidence bounds. \n2)Also, seeing that the main discussion in the paper is around the computational efficiency, the figure 2 should be replaced with another version in which the horizontal axis is the number of flops used instead of the epochs. \n\nas small note, in the experimental validations instead of using words like \"roughly\" they can easily provide the ratio of the time it takes (even a short footnote may be okay, instead of referring to appendix). \n\nW3) About the literature review: the hyper(feature)attention paper available at “https://arxiv.org/abs/2506.06179” appears to be closely related and should probably be discussed. That’s the first paper that comes to my mind, but there are probably other, closely related ones (with very similar self-attention variants). I just wanted to check whether you’ve already reviewed the literature and how confident you are in your coverage.\n\nIn addition to these, I discuss a fundamental possible weakness in the questions section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ym1A2zsjlq", "reviewer": "Reviewer_Henj", "summary": "This paper presents a significant and elegant theoretical contribution to the understanding of attention mechanisms. It introduces \"poly-attention,\" a general framework parameterized by an \"attention polynomial\" $h$ that generalizes standard self-attention and other recent higher-order proposals. The authors provide a comprehensive theoretical analysis of this framework, charting the trade-offs between computational complexity and representational power. While the theoretical work seems solid, the paper is held back by a critical lack of clarity in its core definition and insufficient experimental validation to support the breadth of its claims.", "questions": "See above.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "1. **Generalization of attention mechanism:** The poly-attention framework itself is an elegant generalization that unifies standard self-attention, tensor attention, and Strassen attention under a single, intuitive parameter: the attention polynomial $h$.\n\n2. **Rigorous theoretical analysis:** The paper's main strength is its rigorous analysis. The authors use tools from computational complexity theory to provide a tight characterization of the running time. The finding that tree polynomials are the only class of poly-attention computable in quadratic time is an excellent and complete theoretical result.", "weaknesses": "1. **Clarification on core definitions:** The paper's central contribution, the poly-attention mechanism, is presented in Definition 2.2. This definition is inconsistent, making the paper's subsequent claims impossible to verify. Please correct me if I am wrong. The attention polynomial $h$ is defined over **$t$** variables ($x_1, ..., x_t$). However, the exponent in Definition 2.2 is written as $h(Q_{l_{1}}^{(1)},...,Q_{l_{k}}^{(k)})$, incorrectly using the polynomial's *degree* **$k$** as the number of arguments.\n\n2. **Overstatement of representational power:** The claim that poly-attention can \"solve polynomial root-finding\" (Theorem 3.7) is a severe overstatement. However, the proof sketch in Appendix G does not describe any learnability result. Instead, it details a complex, static encoding where the polynomial $p^2$ is explicitly hard-coded into the query-key weights. This is just construction of a specific circuit, not evidence that a model can *learn* to solve this task.\n\n3. **Meaningless Runtime Baselines:** The paper claims tree-attention's $O(n^2)$ complexity is a practical advantage over superquadratic methods (Strassen, 3-tensor). However, the runtime table (Figure 4) omits these mechanisms entirely, comparing only against 1- and 2-layer self-attention. Without these baselines, the practical benefit of the $O(n^2)$ complexity is completely unsubstantiated.\n\n4. **Insufficient experiments on real world datasets :** A theoretical paper can be valuable. But when a paper makes explicit claims about practicality, efficiency, and solving tasks, it has a duty to validate them empirically. This paper's experiments are so weak, and they cannot demonstrate whether the proposed method is really useful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0BoB3zkDk9", "reviewer": "Reviewer_9YpC", "summary": "This paper introduces **Poly-Attention**, a unifying mathematical framework that generalizes standard self-attention into a family of *higher-order attention mechanisms*. Each variant is parameterized by an “attention polynomial” defining the interaction structure among tokens. The framework subsumes existing mechanisms such as **self-attention**, **tensor-attention**, and **Strassen-attention** as special cases. The authors analyze both **expressive power** and **computational complexity**。 Overall, the work bridges representational theory and fine-grained complexity, positioning Poly-Attention as a fundamental generalization of Transformer attention mechanisms.", "questions": "See weakness", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* Introduces a **novel formal framework** — “attention polynomials” — that unifies multiple strands of higher-order attention research under a single algebraic abstraction.\n* The theoretical results are carefully stated with proofs in the appendix, covering both upper- and lower-bound analyses.\n* Experimental validation, though limited, correctly demonstrates the empirical viability of Tree-Attention.", "weaknesses": "* Experiments are confined to small synthetic tasks (function composition on sequences of length 25). While appropriate for illustrating theory, they do not demonstrate real-world scalability or integration into large models, since transformers are always stacked in a lot of layers. So I think a multi-layer experiments in real-world scenerio (like perplexity in NLP tasks) will make great contribution to the paper.\n\n* The paper claims tree-attention has similar runtime to self-attention, but only reports per-epoch times qualitatively. Quantitative benchmarks on memory, throughput, or gradient stability would strengthen the practical argument.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)