Poly-Attention — A Better Attention Mechanism That's Still Fast
Most improvements to the Transformer's self-attention mechanism work at the level of approximation — finding faster ways to compute the same pairwise dot products. Sparse attention, linear attention, flash attention: all of these keep the fundamental computation the same and just make it cheaper. They share an assumption: pairwise token interactions are the right thing to compute, and the goal is to do it faster.
A recent paper from Columbia University — Poly-Attention: A General Scheme for Higher-Order Self-Attention (Chakrabarti, Pitassi, Alman, 2026) — makes a different move. It asks whether pairwise interactions are expressive enough, proves they often aren't, and then builds a unified mathematical framework for richer attention mechanisms. The punchline is a new mechanism called tree-attention that is simultaneously more expressive than standard self-attention and runs in the same quadratic time.
1. The Problem — What Self-Attention Cannot Do
Self-attention is built around one operation: the dot product between two token vectors. For token \( i \) attending to token \( j \), the score is \( \langle Q_i, K_j \rangle \) — a single number measuring pairwise similarity. The full attention output for token \( i \) is:
This is powerful for tasks where pairwise relationships matter — coreference resolution, subject-verb agreement, word-to-word translation. But many real reasoning tasks require looking at triples or chains of tokens simultaneously.
The canonical example is function composition. Given the facts "Sam lives in Toronto" and "Toronto is in Canada", answering "which country does Sam live in?" requires a two-hop lookup: first find Sam's city, then find that city's country. Formally, given two functions \( f_1, f_2 : [n] \to [n] \) and an input \( x \), the task is to compute \( f_2(f_1(x)) \).
Peng et al. (2024) proved that standard self-attention cannot do this in one layer. You would need almost \( n \) heads of self-attention — making the total cost cubic — just to solve this single, simple reasoning task. Empirically, LLMs fail at multiplication tables, logical puzzles, dynamic programming, and spatial composition for exactly this reason: the pairwise attention mechanism lacks the structural vocabulary for chained lookups.
2. Prior Attempts — Tensor Attention and Strassen Attention
Two prior mechanisms tried to fix this by generalizing the dot product to involve three tokens at once instead of two.
3-Tensor Attention replaces the pairwise dot product with a three-way inner product. For vectors \( a, b, c \in \mathbb{R}^d \) define \( \langle a, b, c \rangle = \sum_{\ell=1}^d a[\ell]\,b[\ell]\,c[\ell] \). The output becomes:
The double sum over \( \ell_2, \ell_3 \) makes this O(n³) — already a showstopper at scale. And it still cannot solve 3-fold function composition.
Strassen Attention uses a cleverer score: instead of one three-way product, it sums three pairwise products forming a triangle between tokens \( i, \ell_2, \ell_3 \):
The triangular structure allows a matrix-multiplication speedup, bringing the cost down to \( O(n^\omega) \) where \( \omega \approx 2.37 \) — subcubic but still superquadratic. And like tensor attention, it provably cannot solve 3-fold function composition (shown in Theorem 3.3 of the paper).
| Mechanism | Score involves | Exact time | 2-fold composition | 3-fold composition |
|---|---|---|---|---|
| Self-attention | pair (i, j) | \(O(n^2)\) | ✗ | ✗ |
| 3-Tensor attention | triple, 3-way product | \(O(n^3)\) | ✓ | ✗ |
| Strassen attention | triple, sum of pairs | \(O(n^\omega)\) | ✓ | ✗ |
| Tree-attention (new) | chain of pairs | \(\mathbf{O(n^2)}\) | ✓ | ✓ |
3. The Framework — Poly-Attention
The paper introduces a single unified framework that captures all the mechanisms above and many more. The key object is an attention polynomial.
Definition (Attention Polynomial): A polynomial \( h(x_1, \ldots, x_t) \) is an attention polynomial of degree \( k \) if it is multi-linear (no variable appears with power > 1), has coefficients only in \( \{0, 1\} \), and every monomial has degree between 2 and \( k \).
Each monomial encodes a specific combination of token interactions. For a monomial \( m \) containing variables \( x_{j_1}, \ldots, x_{j_k} \) and vectors \( Y_1, \ldots, Y_t \in \mathbb{R}^d \), define:
Then for a full polynomial with \( s \) monomials \( m_1, \ldots, m_s \):
The poly-attention output is then defined for each query token \( \ell_1 \) as:
The polynomial \( h \) is a recipe: each monomial votes on which token combinations should increase the attention score. The exponential + softmax then selects the highest-scoring combinations, just as in standard attention.
Every prior mechanism is a special case. Self-attention uses \( h(x_1, x_2) = x_1 x_2 \). Tensor attention uses \( h(x_1, \ldots, x_t) = x_1 x_2 \cdots x_t \). Strassen attention uses \( h(x_1, x_2, x_3) = x_1 x_2 + x_2 x_3 + x_3 x_1 \). The framework doesn't just unify them — it opens up the entire space of polynomials as candidate mechanisms.
4. The Key Insight — Tree Polynomials
For a degree-2 attention polynomial, define its graphical representation: a graph \( G \) with \( t \) vertices (one per variable), with an edge between \( v_i \) and \( v_j \) whenever \( x_i x_j \) is a monomial in \( h \).
If this graph is a tree or forest, \( h \) is called a tree polynomial, and poly-attention for it is called tree-attention.
The graphical structure tells you everything about the computational cost. A tree can always be split into two independent sub-polynomials that share exactly one variable (the root). This means the attention softmax factorizes — you compute each subtree's attention separately and combine them with an element-wise product:
where \( h = f + g \) and \( f, g \) share only \( x_1 \). Each of \( f \) and \( g \) involves a sum over independent sets of key tokens — so the double summation over all token pairs never appears. This keeps the cost at \( O(n^2) \).
A cycle breaks this factorization. In the Strassen triangle \( x_1 x_2 + x_2 x_3 + x_3 x_1 \), no matter how you try to split the polynomial, the two pieces always share more than one variable — creating a dependency that forces you to sum over all pairs \( (\ell_2, \ell_3) \) jointly, pushing the cost to \( O(n^3) \) or at best \( O(n^\omega) \) via matrix multiplication.
5. Function Composition — The Proof of Power
The specific tree polynomial the paper uses to demonstrate expressiveness is the path polynomial:
This is a chain: each monomial connects consecutive variables. Variable \( x_2 \) appears in both \( x_1 x_2 \) and \( x_2 x_3 \) — it's the bridge. This chain structure mirrors the chain of lookups in function composition exactly.
Theorem 3.4: For any integer \( r \geq 2 \), poly-attention for \( h_r \) can simulate \( r \)-fold function composition using one head, computable exactly in time \( O(r^3 n^2) \).
The construction encodes the functions so that \( h_r \) evaluated at the query-key vectors gives:
where \( \phi(i) \) encodes function values and \( A > 1 \) is a constant. This expression is maximized (equals 0) exactly when:
The softmax concentrates all weight on these correct indices, effectively reading out the composed function value. Since \( r \) is a fixed constant, \( O(r^3 n^2) = O(n^2) \) — quadratic time.
For comparison: self-attention needs \( \Omega(n) \) heads to solve even 2-fold composition (proved by Peng et al. 2024), costing \( O(n^3) \) total. One head of tree-attention does it in \( O(n^2) \).
6. The Complete Complexity Picture
The paper proves a tight characterization of every poly-attention mechanism, both for exact computation and for approximation. The approximation results hinge on a bound \( B \) — the maximum absolute value of entries in the query-key matrices. When \( B \) is small, the softmax is smooth and can be approximated by low-degree polynomials. When \( B \) is large, the softmax is sharply peaked and approximation becomes hard.
Theorem 3.5 (Tree-Attention, complete): For a tree polynomial \( h \) with query-key entries in \( [-B, B] \):
- Exact: \( \text{Att}^{(h)} \) computable in \( n^{2+o(1)} \) time.
- Fast approx: If \( B = o(\sqrt{\log n}) \), approximation runs in \( n^{1+o(1)} \) time.
- Lower bound: If \( B = \Omega(\sqrt{\log n}) \), approximation requires \( \Omega(n^2) \) time (under SETH).
Theorem 3.6 (Non-tree poly-attention): For any attention polynomial \( h \) that is not a tree polynomial, with degree \( k \):
- Fast approx: If \( B = o\!\left((\log n)^{1/k}\right) \), approximation runs in \( n^{1+o(1)} \) time.
- Lower bound: If \( B = \Omega\!\left((\log n)^{1/k}\right) \), approximation requires superquadratic time (under complexity assumptions).
| Mechanism | Exact time | Approx time | Bound \(B\) for fast approx |
|---|---|---|---|
| Self-attention | \(n^{2+o(1)}\) | \(n^{1+o(1)}\) | \(o(\sqrt{\log n})\) |
| t-Tensor attention | \(n^{t+o(1)}\) | \(n^{1+o(1)}\) | \(o((\log n)^{1/t})\) |
| Strassen attention | \(n^{\omega+o(1)}\) | \(n^{1+o(1)}\) | \(o(\sqrt{\log n})\) |
| Tree-attention (new) | \(\mathbf{n^{2+o(1)}}\) | \(\mathbf{n^{1+o(1)}}\) | \(\mathbf{o(\sqrt{\log n})}\) |
| General poly-attention | \(n^{t+o(1)}\) | \(n^{1+o(1)}\) | \(o((\log n)^{1/k})\) |
7. Two Forward Passes — Same Weights, Different Polynomial Structure
It helps to see concretely how the same set of weight matrices produces different behaviors depending on which polynomial \( h \) you choose. The diagram below shows the attention graph for two different polynomials on the same five tokens.
In self-attention, token \( x \) can find \( f_1(x) \) in one step — but to then look up \( f_2(f_1(x)) \), the result of the first step would need to be the input to a second attention operation, requiring a second layer. Tree-attention's polynomial \( x_1 x_2 + x_2 x_3 \) does both hops simultaneously in one softmax: \( x_2 \) acts as the bridge variable, attending both to the input token \( x_1 \) and to the \( f_2 \) table via \( x_3 \).
8. Representational Breadth — Polynomial Root-Finding
Function composition is just one application. The paper proves a much broader representational result: poly-attention can solve polynomial root-finding for any polynomial.
The problem: Given a polynomial \( p(x_1, \ldots, x_t) \) and a set \( S \) of \( n \) integers, find \( y_1, \ldots, y_t \in S \) such that \( p(y_1, \ldots, y_t) = 0 \).
This captures: Match3 (find three tokens summing to zero, i.e. \( p = x_1 + x_2 + x_3 \)), circuit evaluation (via arithmetization of logic gates), and many other reasoning tasks.
Theorem 3.7: For every polynomial \( p \), there exists an attention polynomial \( h \) such that a Transformer with two heads of poly-attention for \( h \) can solve polynomial root-finding.
The construction uses \( p^2 \) as the building block. Since \( p^2 \geq 0 \) always, and \( p^2 = 0 \) exactly at roots of \( p \), we construct \( h \) such that:
The softmax then concentrates entirely on root tuples (score = 0) versus non-roots (score = \( -n^2 \cdot \text{something positive} \)). Head 1 finds the root, Head 2 retrieves the actual values for verification by the output MLP.
9. Proof Techniques — How the Results Are Established
The paper uses three distinct technical ingredients.
Fast approximation algorithms use the polynomial method: approximate \( e^x \) with a low-degree polynomial \( P(x) \), then replace the exponential in the softmax. When query-key entries are bounded by \( B \), the exponent lies in \( [-sB^k, sB^k] \). On this interval, \( e^x \) can be approximated to precision \( \varepsilon \) by a polynomial of degree \( O(\max\{\log(1/\varepsilon)/\log(\log(1/\varepsilon)/B^k), B^k\}) \). For \( B = o((\log n)^{1/k}) \), this degree is \( o(\log n) \), and the resulting low-rank factorization enables near-linear computation. The key improvement over prior work: the bound depends on degree \( k \), not number of variables \( t \) — so for tree-attention with \( t = 20 \) variables but degree \( k = 2 \), the bound is \( o(\sqrt{\log n}) \) rather than the much smaller \( o((\log n)^{1/20}) \).
Running time lower bounds use fine-grained complexity. For degree-\( k \) polynomials, a reduction from \( k \)-IP under SETH gives an \( \Omega(n^k) \) lower bound when \( B \) is large. For degree-2 non-trees (like Strassen), SETH only gives cubic lower bounds — useless since Strassen already has a subcubic algorithm. Instead the paper uses the Max-2SAT conjecture, reducing from ε-Gap-IP△ (a triangular inner product problem) via the distributed PCP framework. Faster Strassen attention would imply faster Max-2SAT, contradicting the conjecture.
Representational lower bounds use communication complexity. The proof that Strassen and tensor attention cannot solve 3-fold composition reduces to myopic pointer jumping: a multi-player communication task where players hold different parts of the input and must collectively compute a composed function. Chakrabarti (2007) proved this requires \( \Omega(n) \) bits of communication. If an attention mechanism with \( H \) heads could solve 3-fold composition, a communication protocol using \( O(d \cdot p \cdot H \cdot \log\log n) \) bits would follow. Setting this \( \geq \Omega(n) \) with \( d = O(\log n) \) and \( p = n^{o(1)} \) gives \( H > n^{1-o(1)} \) — nearly \( n \) heads required.
10. Experimental Results
The theoretical results are backed by two experiments. Both use a simple architecture: embedding dimension 32–64, 4 heads, one MLP layer with hidden size 128–256.
Experiment 1 — Function Composition. Three models are compared on learning \( f_1(f_2(x)) \) for functions \( f_1, f_2 : [n] \to [n] \) with \( n = 25 \) (sequence length 51):
| Model | Learns composition? | Speed of learning | Parameters |
|---|---|---|---|
| 1-layer self-attention | ✗ (flat at ~0%) | — | fewest |
| 2-layer self-attention | ✓ (slowly) | slower, noisier | most |
| 1-layer tree-attention | ✓ (fast) | fastest, most stable | fewer than 2-layer SA |
One-layer self-attention never learns — exactly as the theory predicts. Tree-attention not only learns, but does so faster and more consistently than two-layer self-attention, despite having fewer parameters. The inductive bias of the polynomial \( h = x_1 x_2 + x_2 x_3 \) essentially gifts the model the chain structure of the task.
Inference timing on an NVIDIA A100 GPU confirms the O(n²) theory has small hidden constants:
| Seq len | 1-layer SA | 2-layer SA | 1-layer Tree | 1-layer 3-Tensor | 1-layer Strassen |
|---|---|---|---|---|---|
| 20 | 1.076 ms | 1.775 ms | 1.367 ms | 1.442 ms | 1.593 ms |
| 50 | 1.079 ms | 1.757 ms | 1.363 ms | 2.911 ms | 1.594 ms |
| 100 | 1.080 ms | 1.781 ms | 1.374 ms | 13.813 ms | 3.395 ms |
Tree-attention is ~1.3× self-attention at all sequence lengths — the overhead is small and constant. 3-tensor attention blows up from 1.4 ms to 13.8 ms as length doubles from 50 to 100 — the O(n³) theory made visible. Tree-attention is actually faster than two-layer self-attention while being more expressive.
Experiment 2 — COGS NLP Benchmark. Both models are evaluated on compositional sentence parsing. In-distribution accuracy is similar (~97.5% for both). On the harder generalization test set (unseen sentence structures):
| Model | Gen. token accuracy | Gen. exact match |
|---|---|---|
| Self-attention | 0.724 ± 0.009 | 0.239 ± 0.087 |
| Tree-attention | 0.728 ± 0.013 | 0.265 ± 0.128 |
Tree-attention achieves ~40% exact match accuracy on 4 out of 10 random seeds — a threshold self-attention never reaches on any seed. The higher variance suggests the optimization landscape for tree-attention is richer; finding better training heuristics is left as future work.
11. What This Means in Practice
The paper gives model designers a concrete decision framework. The choice of attention polynomial \( h \) is now a design variable — not a fixed architectural decision.
If you need almost-linear time: use approximate tree-attention. The fast approximation is available whenever \( B = o(\sqrt{\log n}) \) — the same condition as standard self-attention. So if your model already supports fast approximate attention, tree-attention inherits that speedup automatically.
For highly complex tasks: use general poly-attention with a non-tree polynomial of higher degree. You pay a superquadratic exact cost but gain more expressiveness. The tradeoff is now precisely quantified by Theorem 3.6.
The broader implication is architectural. The reason LLMs struggle with multi-hop reasoning, compositional generalization, and function composition may not be a training data problem or a scale problem — it may be a structural problem with pairwise attention itself. Poly-attention suggests a path forward that doesn't require making the model bigger or training it longer, just making the attention polynomial richer in a principled way.
Whether tree-attention scales to GPT-sized models remains to be seen. The experiments here use small models on controlled tasks. But the theory is clean, the empirical evidence is consistent, and the computational cost is genuinely manageable. That combination doesn't come along often.
Comments
Post a Comment