How to Capture Higher-order Correlations? Generalizing Matrix Softmax Attention to Kronecker Computation
Authors: Josh Alman
Zhao Song
Venue: ICLR 2024 spotlight (2024)
Keywords: Attention computation
kronecker computation
Primary area: learning theory

Abstract: In the classical transformer attention scheme, we are given three $n \times d$ size matrices $Q, K, V$ (the query, key, and value tokens), and the goal is to compute a new $n \times d$ size matrix $D^{-1} \exp(QK^\top) V$ where $D = \mathrm{diag}( \exp(QK^\top) {\bf 1}_n )$. Here, $\exp()$ is applied entry-wise and ${\bf 1}_n$ denotes a length-$n$ vector whose entries are all ones.

Intuitively, attention computation captures pairwise information between words in a sentence, but not higher-order information. Indeed, recent work \cite{sht23} has shown that attention units cannot solve simple problems about detecting triples of connected words.

In this work, we study a generalization of attention which captures triple-wise  correlations. The generalization is based on computations involving tensors defined by tuples of words. More formally, given five $n \times d$ size matrices $Q, K_1, K_2, V_1$ and $V_2$ (generalized query, key, and value tokens), our new goal is to compute an $n \times d$ size matrix $D^{-1} \exp( Q ( K_1 \oslash K_2)^\top ) (V_1 \oslash V_2) $ where $D = \mathrm{diag}( \exp( Q ( K_1 \oslash K_2)^\top ) {\bf 1}_{n^2} )$ and $K_1 \oslash K_2 \in \mathbb{R}^{n^2 \times d}$ denotes the column-wise Kronecker product of $K_1$ and $K_2$. This generalization is indeed able to solve problems about detecting triple-wise connections that were shown to be impossible for transformers.

The potential downside of this generalization is that it appears as though computations are even more difficult, since the straightforward algorithm requires cubic time in $n$. However, we show that in the bounded-entry setting (which arises in practice, and which is well-studied in both theory and practice), there is actually a near-linear time algorithm. More precisely, we show that bounded entries are both necessary and sufficient for quickly performing generalized computations:

$\bullet$ On the positive side, if all entries of the input matrices are bounded above by $o(\sqrt[3]{\log n})$ then we show how to approximate the ``tensor-type'' attention matrix in $n^{1+o(1)}$ time.

$\bullet$ On the negative side, we show that if the entries of the input matrices may be as large as $\Omega(\sqrt[3]{\log n})$, then there is no algorithm that runs faster than $n^{3-o(1)}$ (assuming the Strong Exponential 
Time Hypothesis from fine-grained complexity theory).


We also show that our construction, algorithms, and lower bounds naturally generalize to higher-order tensors and correlations. Interestingly, the higher the order of the tensors, the lower the bound on the entries needs to be for an efficient algorithm. Our results thus yield a natural tradeoff between the boundedness of the entries, and order of the tensor one may use for more expressive, efficient attention computation.

Our constructions make use of a novel connection with a higher-order variant on the kernel density estimation problem. They combine a number of technical tools, including the polynomial method, algebraic geometry codes, and multiparty Merlin-Arthur communication protocols.

Published as a conference paper at ICLR 2024
How TO CAPTURE HIGHER-ORDER CORRELATIONS? GENERALIZING MATRIX SOFTMAX ATTENTION TO
KRONECKER COMPUTATION
Josh Alman
Zhao Song
Columbia University
Adobe Research Seattle, WA, USA
New York, NY, USA
joshêcs col umbi edu
zsongeadobe..omm com
ABSTRACT
In the classical transformer attention scheme,
given three d size matrices Q, K, V (the query, key, and value tokens), and the goal is to compute a new n x d size matrix D exp(QK' V) where D diag(exp(QK' )I.). Here, exp() is applied entry-wise and 1, denotes a length-n vector whose entries are all ones.
Intuitively, attention computation captures pairwise information between words in sentence, but not higher-order order information. Indeed, recent work Sanford et al. (2023) has shown that attention units cannot solve simple problems about detecting triples of connected words.
In this work, we study a generalization of attention which captures triple-wise correlations. The generalization is based on computations involving tensors defined by tuples of words. More formally, given five n d size matrices Q, K1, K2, V1 and V22 (generalized query, key, and value tokens), our new goal is to compute an n x d size matrix D- exp((
K2) (Vi @ 1/2) where D diag diag(exp( Q(K T1)T) and Ki Ø K2 E R=xx denotes the column- wise Kronecker product of K₃ and Kg. This generalization is indeed able to solve problems about detecting triple-wise connections that were shown to impossible for transformers.
The potential downside of this generalization is that appears as though compu- tations are even more difficult, since the straightforward algorithm requires cubic time in n. However, we show that in the bounded-ent couneedeentry setting (which arises practice, and which is well-studied in both theory and practice), there is actually near-linear time algorithm. More precisely, we show that bounded entries are both necessary and sufficient for quickly performing generalized computations:
On the positive side, if all entries of the input matrices are bounded above by of VTo n) then we show how to approximate the "tensor-type" attention matrix in n (+((l) time.
On the negative side, we show that if the entries the input matrices may be as large Qi Vlogn), then there is algorithm that runs faster than 3--o(1) (assuming the Strong Exponential Time Hypothesis from fine-grained complexity theory).
We also show that our construction, algorithms, and lower bounds naturally gener- alize to higher-order tensors and correlations. Interestingly, the higher the order the tensors, the lower the bound on the entries needs to be for an efficient algorithm. Our results thus yield natural tradeoff between the boundedness of the entries, and order of the tensor one may use for more expressive, efficient attention computation.
Our constructions make of novel connection with higher-order variant on the kernel density estimation problem. They combine a number of technical tools, including the polynomial method, algebraic geometry codes, and multiparty Merlin-Arthur communication protocols.

---

Published as a conference paper at ICLR 2024
INTRODUCTION
Large language models, such as Transformer Vaswani et al. (2017), BERT Devlin et al. (2018), GPT-1 Radford et al. (2018), GPT-2 Radford al. (2019), GPT-3 Brown et al. (2020), PaLM Chowdhery et al. (2022), OPT Zhang et al. (2022), GPT-3.5, Bard, GPT-4 OpenAl (2023), Llama Touvron al. (2023a); Rozière al. (2023), Llama 2 Touvron al. (2023b) and its successors, have gained immense importance and found a wide range of applications due to their ability to understand and generate human-like text. These models are trained on massive amounts text data, enabling them to learn patterns, structures, and nuances of human language. They have applications in many areas, including understanding natural language, content generation, improved human-computer interaction, translation and multilingual communication, and rapid prototyping.
The fundamental computational structure at the core of LLMs is called an attention unit. When length-n input is given to the attention unit (like a sentence or paragraph of n words), we embed into three matrices Q, K, (the query, key, and value token matrices) where each has n rows and d columns. Here d is the feature dimension; one has d < in the long sentence regime. Mathematically, the attention unit computes D exp(QK VV, where D = diag(exp(((i )1,) is diagonal matrix, 1n denotes the length-n vector with all entries equal to 1, and exp is applied entry-wise.
Intuitively, the attention unit is finding pairwise correlations between tokens in the input since it computes inner products between pairs of tokens when computing QK However, if the input data has correlated triples of tokens, is not clear attention unit can detect this.
recent and exciting work Sanford et al. (2023) formalized this intuition. They defined a simple task about learning correlations between triples of words, and showed that attention units unable solve it. By contrast, they are able to solve the analogous problem of learning correlations between pairs of words. Toward resolving this, Sanford al. (2023) proposed generalization of attention computation:
Definition 1.1 (Tensor generalization of attention scheme). Given as input n d matrices K,KI, V1, V2, the goal is to construct another X d matrix
D- @ V2).
Here
V1 O V2 Rn denotes the column-wise Kronecker product of V1 and Vz. Similarly for K1 @ K2 E Rn below. (The column-wise Kronecker product of matrices K₁ Rand K2 Ruxd matrix K == K1 @ K2 E R=nxx defined as Ki 11p (K1) (K2)aag Vi1,22 E [m].] j E [d].)
A Ruxn² is the n² matrix exp @ K2)T/a) where exp is applied entry-wise.
Ruxx is the n diagonal matrix diag(exp(Q(U) @ (2) /d)1..)
1,, here denotes length-n? vector whose entries are all ones.
One may naturally view A as n X tensor, which is why we call this a 'tensor generalization': this view will be important in our proofs below.
In this generalization, entries of the matrix A now correspond to triples of tokens, so one may hope that this generalization can detect triple-wise correlations. And indeed, Sanford et al. (2023) show that this is the case: the tensor generalization gets around their expressivity barrier and is able to detect correlations among triples of tokens.
A fundamental question arises naturally: how quickly can generalized attention computations be performed? The running time of attention computations critically important, since forms the time bottleneck of L.M training and inference. By generalizing attention make it more expressive, have we also made intractably slow?
To answer this question, focus approximate version of the tensor attention computation problem. In practical applications, it is sufficient to approximately perform these computations Child al. (2019); Kitaev et al. (2020); Wang et al. (2020); Choromanski et al. (2021); Daras al. (2020);
2

---

Published as a conference paper at ICLR 2024
Katharopoulos et al. (2020); Chen et al. (2021; 2022); Qin et al. (2022); Zandich al. (2023): Liu et al. (2023); Zhang et al. (2023); Kacham et al. (2023); Dao et al. (2022); Dao (2023), and this often helps lead to faster algorithms.
Definition 1.2 (Approximate Tensor Attention Computation ATAttC d, B, Let Eu O, B > be parameters. Given five matrices Q, K1, K2, V1, V/2 E Raxd that satisfy the following bounded constraints,
. |111 S B, |K 2 B, 1|2 S B. 1V B. 2V2 S
we want to generate a matrix T E Budd which is able to entry-wisely approximate D-1 AV, iee,
|[T A-V @ Ve) co
Here,
the too norm for matrix N E Bgxd is writlen as ||N maxiel [Ntg], and
the other matrices are defined as in Definition 1.1 above.
We focus here the natural setting with d O(log n) (so that modeling long sequences) and 6a = 1/ poly(n) (so that one can combine the errors from attention computations over an entire network).
In the case of (non-tensor) attention, the computational complexity of exact and approximate attention computation is very well-understood. Keles et al. (2023) showed that the trivial (n²) time algorithm is essentially optimal for exact computation, assuming the Strong Exponential Time Hypothesis (SETH). SETH Impagliazzo & Paturi (2001) is popular conjecture from fine-grained complexity which posits that one cannot substantially improve our current best algorithms for k-SAT; see the survey Williams (2018) for more details. Since k-SAT algorithms are very well-studied, it is not commonly believed that major improvements are possible, and so much of fine- tieeggrined complexity theory based on this assumption.
Alman & Song (2023) studied the approximate (non-tensor) attention problem and showed that its complexity depends on the magnitude of the entries the matrices Q, K: If they are smaller than (looog n.), then there is a fast algorithm running time n +++o();; this near-linear time algorithm is essentially as fast as one could hope for. On the other hand, if they are at least D log n). then there no algorithm substantially faster than the trivial (n²) assuming SETH. This theoretical result mirrors practical observations that bounded entries are essential for fast attention Zafrir et al. (2019); Sun al. (2019); Katharopoulos al. (2020); Dettmers al. (2022b); Xiao al. (2023); Dettmers et al. (2022a); Perez et al. (2023); Shen et al. (2023).
1.1
OUR RESULTS
Our main results tightly resolve the computational complexity of the tensor generalization of attention. Generalizing the situation for (non-tensor) attention, we show that whether or not there is a fast algorithm for AAttC depends the parameter B, the magnitudes of the entries in the query, key, and value matrices.
We first show lower bound, that when B Ω Vlogn). it is impossible design truly subeubic- time algorithm (assuming SETH). Note that the straigtforward algorithm for this problem runs in cubic time, our result shows that one cannot substantially improve on the straightforward algorithm when the entries have magnitude at least O Vlog n).
Theorem 1.3 (Lower bound, informal version of Theorem B.2). Assuming SETH, for every > 0, there are constants C, Ca. C. > such that: there no algorithm running in time O for the problem ATAttC(n, d C' log n. B Cs Vlogn, n, = n Ca).
Our second result is a new algorithm, showing that when B < of V)ogn, then there is an almost linear time algorithm for solving the problem.
Theorem 1.4 (Upper bound, informal version of Theorem E.3). There is an algorithm (Algorithm 1) that solves ATAttC(n, d O(log n), , B of VTg n), 1/ poly(n)) in time n nl+o(1),
3

---

Published as a conference paper at ICLR 2024
Our Theorems 1.3 and 1.4 together show that the complexity of ATAttC has very tight transition at B = e Vlo When B < of ylgg n.) is smaller than the threshold, the problem can be solved essentially as quickly as one could hope for, in time +lo0 Meanwhile, when 2 Q Vlog is greater than the threshold, it is impossible achieve subeubic running time, matter what algorithmic techniques are used (assuming SETH).
It is exciting that, even for the more expressive tensor generalization of attention, there is linear time algorithm in the bounded entry regime. Interestingly, though, the bound must be smaller than for regular attention: for regular attention to have near-linear time algorithm, it is necessary and sufficient that B < Vlog n, whereas for tensor- based attention, we show it is necessary and sufficient that Vlogn.
More generally, for any positive integer k 2 2, we study a higher-order tensor generalization of attention which can detect k-wise correlations. (Regular attention corresponds k 2 and ATAttC corresponds to k = 3.) For this problem, we further generalize our results to show that there is a near- linear time algorithm when the entries satisfy B Vlog n. and that the trivial (nn) time essentially cannot be beaten otherwise. This suggests an intriguing tradeoff between the boundedness of the entries, and the expressiveness of attention we can perform quickly: Given vectors corresponding to tokens for LLM training inference, we let B be the largest magnitude of entry, then select the largest for which < ş/log n, and we can quickly perform k-th order attention computations for our tokens, but not higher-order attention.
Definition 1.5 (k-th order generalization of Definition 1.1). Suppose K, K2, .. Kg-1 and V1, V2,, ... VV our target is to construct another n. matrix
given X d matrices
D A(V 2 V2 2 2.... ...
Vk-1)
Here
1..0....000 E RRnn"xxd is the column-wise tensor product V1 ". Va-1
E Ruxnn= is the n. x n kn1 size matrix exp(Q(Ki 2 K2 Ø ... ..0 Kk-1) 1) 1/ /dd)
D E Rux is the x diagonal mairix diag(exp( Ø K2 Ø ... K 1))))/))mn--1
1, is the length-n vector whose entries all ones.
Roadmap.
In Section 2, we provide number basic notations and definitions. Section 3, we give a technique overview, summarizing our proofs for both our upper bound result and our lower bound result. In Section 4, we prove the key intermediate results for our lower bound result. Our upper bound result, and the remainder of our lower bound result, are proved in the Appendix.
PRELIMINARY
Hadamard Product
Definition 2.1 (o Hadamard product). Given A, B. E Raxd we use C == A e to denote their entry-wise product, i.e., the matrix C Raxd given by Cij Ai Bt.j We similarly define to denote the entry-wise product of vectors or tensors. This is often called the Hadamard product in the literature.
Tensor Operations
Many of proofs will involve manipulating tensors. Here introduce three different tensor operations we will frequently use.
Definition 2.2 ( tensor computation). Given matrices A E Rud B e Raxd C E Raxd we use T A B C denote an n n x n tensor whose entries are given by Tej
Ed= Aia Bj.a Cr.y, V [n],j E [m), E [p].
We note that a tensor can be written the form A O B O like this and only if its tensor rank is most d.
4

---

Published as a conference paper at ICLR 2024
Definition 2.3 (@ Kronecker product). Given two matrices K1 Rnxd and K2 Ruxd, we define K K1 @ K2 E 1R xa² as follows (iit(tt--))m::-22--0 (Ki) (K2)kadz [n], E J1 E [d]. [d].
Vi,
In this work, we will primarily use the following column-wise version of the Kronecker product. Definition 2.4 (2 column-wise Kronecker product), Given matrices K1 E IRuxd K2 E Iaxdd we define matrix K K1 Ø K2 R= xd follows (ix+(i (Ki) (K2) Vi1, [n]. E [d).
Matrix Multiplication Finally, our algorithm will make use matrix multiplications. For positive integers n. m, d, we write Tmaa (n, d, m) denote the time multiply given d matrix A and a d x m matrix B. The straightforward algorithm shows that Tmat (n, d, m) = O(ndm), and this will suffice for our algorithms here; we will typically apply this when of n. m, are very small compared to the third, and in this situation, more clever matrix multiplication algorithm do not yield substantial speedups.
TECHNIQUE OVERVIEW
Generalizing prior work on the computational complexity of the attention problem to our tensor generalization requires overcoming number of technical challenges. Here summarize our approach, with an emphasis on differences with the prior work on (non-tensor) attention that we build on Rubinstein (2018); Katharopoulos et al. (2020); Alman & Song (2023); Sanford et al. (2023).
3.1 ALGORITHM
Tool for the column-wise Kronecker product We begin by introducing basic tool for manipulat- ing computations involving the column-wise Kronecker product Ø (see details in Lemma C.3 below). Define the following matrices.
Given A E Ruxd, A2 E Raxd, we define == (A1 Ø A2) E R=ixd,
Given B1 Raxds, B2 Runds, define (Bi @ B2) E R
We define C Rdixda C == A B, and similarly define C₁ A B₁ and C2 A. B2.
Then, we prove that we have C1 o C2 C. Using this identity, C can be computed in time O(Tmar(d, n, dz)) given the matrices A,, A2, B1, Bz.
Approximating D In order perform generalized attention, we aim compute the matrix /d)l,²) [d))//)),> Notice that the intermediate matrix (Xp(Q(K) Ø K2)T/d
D = diag(exp(((e) @ K2)T has n² entries. We thus cannot compute in subcubic time. We instead aim use an implicis representation of an approximation of this matrix which can be quickly manipulated.
Toward this goal, we find appropriate matrices U1, U2 U3 (which we discuss in more detail shortly) and formulate = diag(U (U2 @ U,)),,²) such that As D. Given the matrices U1, U2 U3 and using the above tool for Ø we can compute D quickly in O(nd) time.
Approximating A We can similarly approximate the attention matrix A exp(Q(K, @ K2) via À = U, Ø TT) such that ~ A (in Eoo norm). Again, in contrast to D, we cannot compute the entries of A since it has n² entries. We instead directly approximate A(Vi O V2) by computing A(V, Vz). This again be done in O(Tmat (d, n, d)) ++(1) time by using the above tool.
Finding approximating matrices U,, U2, U3 Thus, it remains to find matrices U1, U2, Ua which appropriately approximate and A as above. We show how to efficiently find such matrices as long as the inputs Q. K1. K1, V1, V22 have bounded entries. The key idea is the polynomial method, a key tool from prior work Aggarwal Alman (2022); Alman Song (2023) which allows one to find low-rank representations of matrices.
5

---

Published as a conference paper at ICLR 2024
The method generally says that if M is low-rank matrix, and p is low-degree polynomial, then p(M) (where p is applied entry-wise) also has relatively low rank. Furthermore, its low-rank decomposition can be found efficiently given the decomposition of M. By applying this method where p is an appropriate polynomial approximation of the exp function (see Aggarwal & Alman (2022)), get low-rank approximation of exp(M).
This polynomial method approach was also taken in the prior work on (non-tensor) attention Alman & Song (2023). Here we generalize it, showing that the same line of attack can be applied to low-rank tensors. Viewing A interchangeably as both an n X n x tensor and an n X n² matrix allows us to take advantage of this low-rank tensor approximation as well as the aforementioned matrix multiplication algorithms, and U1. U2 U3 are ultimately the low-rank approximation expression for this tensor. Notably, as the bound B on the entries increases, the degree of the polynomial to approximate exp also increases, but the degree needs to be small enough to give nontrivial algorithm. See details in Lemma E.
3.2 HARDNESS
Gap-MaxIP Our hardness proof proceeds by introducing and considering a new intermediate problem we call Gap-MaxIP (Definition 4.6), promise version of the more common 3-Max IP problem. In this problem, one is given as input 3n vectors di,... Bny b bi,..., bp 11, Cl,..., {0, 13dd as well as threshold t, and the goal is distinguish between the cases
* (a;, by, cx) I t for all i,j,k k E [n]. or
* (a; (y.b),a)) 2 2t for some i,j, k E [n].
(If neither is the case, we may give any output.) Here, (as, by, cu) denotes the 3-way inner product [1] by[e] cx((].
We first prove that Gap- MaxIP cannot be solved in truly subcubic time assuming SETH. We then show that truly subeubic time algorithm for our generalized ATAttC (Definition 1.2) problem with large entries would yield one for MaxIM as well.
Previous work (non-tensor) attention Alman & Song (2023) used its intermediate problem the approximate Hamming Nearest Neighbor problem. However, it not obvious how to directly generalize this to the tensor setting, since there is no way to define 'distance' function for triples of vectors which satisfies the needed properties to generalize the original proof. We instead investigate the Sap-MaxiP problem, which can itself be seen generalization of an intermediate step in the proof of hardness for approximate Hamming Nearest Neighbor Rubinstein (2018).
Hardness of MaxIP Fine-grained complexity results for approximation problems like Gap-MaxIP have previously been shown using a distributed probabilistically checkable proof framework Abboud et al. (2017); Rubinstein (2018), which we also use here.
We begin by generalizing the approach of Rubinstein (2018) using Merlin-Arthur (MA) communica- tion protocols (Babai (1985); Goldwasser & Sipser (1986); Arora & Barak (2009)). We construct four party communication protocol for the disjointness problem: Alice, Bob and Charlic are each given subsets of a universe, and want to determine whether there is an element in all three of their sets. In an MA protocol, Merlin first sends an advice string to the three players to convince them their sets are disjoint. Alice, Bob and Charlie may then flip private random coins and communicate to come to an answer. (See details in Theorem 4.5).
Generalizing known three-party protocols for disjointness Aaronson & Wigderson (2009); Rubinstein (2018), our protocol algebraic in nature, and critically makes use of algebraic geometry codes from coding theory Shum (2000); Shum et al. (2001).
We then use this protocol to reduce from SAT to G A standard reduction Williams (2005) shows that SAT reduces to the 30V problem, which is a computational version of the three player disjointness problem. We can convert inputs to this problem into vectors by corresponding entries of the vectors to possible transcripts of the communication protocol. The gap in inner products will arise naturally from the correctness guarantees of the protocol. See reduction details in Theorem 4.7 and its proofs.

---

Published as a conference paper at ICLR 2024
Reducing from Gap-MaxIP to ATAttC Finally, reduce the Gap-MaxIP (Definition 4.6) problem to our ATAttC (Definition 1.2) problem. The key idea is that. by defining the matrices Q, K,K,,K V1, V/ of generalized attention in terms of the inputs to Gap-MaxIP, large entries of the attention matrix A correspond to the triples with largest inner product. (See Lemma B.1 below for illustration.) Some manipulation similar to prior work Alman & Song (2023) allows us to detect large entries from the output of ATAttC. This approach has been used for the fine-grained hardness of many attention and kernel density estimation problems Backurs al. (2018); Katharopoulos et al. (2020); Alman et al. (2020); Aggarwal & Alman (2022); Alman & Song (2023). See details Lemma B.1 and its proofs.
can make
HARDNESS
In this section, begin the formal proof of our hardness result. We begin by introducing the fine-grained hypotheses we will use.
Hypothesis 4.1 (Strong Exponential Time Hypothesis (SETH), Impagliazzo & Paturi (2001)). For every > O there exists integer k 2 3 such that CNF - SAT on formulas with clauses size at most k (the so called k-SAT problem) and variables cannot be solved ((2(1-e)a) time even by randomized algorithm.
Definition 4.2 (3OV). Given three sets A, A,B,C {0, where A| B = |C| = n, the goal is find tuple (11, i2, iz) E [mx] x [n] [n) such that (aty biyy 0.
Conjecture 4.3 (Orthogonal Vectors Conjecture (3OVC) Williams (2005); Abboud et al. (2014)).
For every > 0, there is c 2 such that 30V cannot be solved 3- time on instances with d = clog
is known that SETH implies 30VC; see, e.g., Williams (2005).
4.1 ALGEBRAIC GEOMETRY CODES FROM PREVIOUS WORK
We state important tool from the field of algebraic geometry codes. For more background on algebraic geometry codes, we refer the reader to Goppa (1981); Tsfasman et al. (1982); Shum (2000); Shum et al. (2001); Sudan (2013).
Theorem 4.4 (Shum al. (2001); see also Rubinstein (2018)). There is constant qo N such that,
for every prime q 2 Yu there are two systematic code families C == {Ca} and C' : { [C*] whose
codewords are given by functions Rn E q² for some appropriate subset R., C F 01lg n) The
codes C,C' C' satisfy four key properties:
Systematicity. There exists a subset Sn C RRn of cardinality |S,1 ((), such that for any assignment x Sn -> Fqz, there exists a codeword w E C such that w|s..
3-way Polynomial Closure. C and are linear codes. For each wi, W2, Wa E C, there exists w' C' such that for each Rn- w w'(i) wi(i) w2(1) wy(1)
Efficiency. Both codes be encoded in poly(n) time and checked in poly(n) time.
Parameters. Both codes have relative rate at least 0.01 and relative distance least 0.01.
4.2 A FOUR PARTY MA COMMUNICATION PROTOCO1.
Prior work (Rubinstein (2018)) constructed protocol for three party communication, which includes Merlin, Alice and Bob. Here we modify this protocol for four parties.
Theorem 4.5. For any T [2, m). There is MA-communication protocol for Set Disjointness over universe (m]. This protocol computationally efficient.
In particular, the details of protocol are
Merlin sends Alice O moloo"? (mmloTT) bits
Alice, Bob, Charlie toss Olog m) coins
7

---

Published as a conference paper at ICLR 2024
Charlie sends Alice O(T log T) bits
Bob sends Alice ( log T) bits
Alice returns Accept or Reject
If the three sets do not have any element in common, Alice always accepts. Otherwise, she accepts with probability al most 1 1/2
Proof. We assume that T divides m. i.e., there is some positive integer r such that m Tr. Otherwise, increase m to the nest multiple of T: this at most doubles m. We partition the universe into T disjoint sets of size r: [m/] = U² U..... ... U U Let 3, S [m] denote the inputs of Alice, Bob, and Charlie. Our goal is to determine whether there is an element in the intersection a n n y.
ynt We will next encode these parts using an algebraic geometry code. Let q be prime greater than T, and let C' be algebraic geometry code over the field IF 222, and let C" be its associated code for the polynomial closure property. Let pc, Sc be the rate and distance of the code; recall these are at least positive constant. Let nc
For each E (7), we define the t-th parts of the three sets: a' == anU', == BnU*,
(T) be the length of the codewords of C.
T-pc
For each E [T]. we write C(a"), C(3*), C(7*) to denote the encodings of a², 3 and 7², Thus, their entry-wise product μ² (i.e., u1 == C(a'): C(B)), C(r²), ) is a codeword in the second code C". Furthermore, since C' is a linear code, the entry-wise sum of the yt (p: Et= 1C) is also codeword of C
C is a systematic code, so we may assume that for each i E [n/ /T, the entries C(a'), C(8);, C(r"); are from {0. } and represent membership in the set. Similarly, μ E {0. and the sets are disjoint and only if u 0 for all i E [m/T] and [T), or equivalently, μu; = for E [m/T].
Now the protocol proceeds as follows:
Step 1.
Merlin sends Alice fi, which supposed be the encoding of y
Step Charlie, Bob and Alice pick random i* E [nc]
Step 3.
Charlie sends Alice C(+*);- for E (T)
Step 4. Bob sends Alice C(8*): for all E [Z]
Step 5. Alice accepts iff all of the following hold:
is codeword in C" Thp- Σt= C(at) CC()): C(7));- and fu: 0 for all
i E [m/T]
First, observe that Merlin's message length is n. log T = T Charlie's message lengths are T Olog 7) as desired. To see correctness, note that if Alice ever accepts given Merlin's 's message ő. then z must in particular be a codeword of C'. If Alice accepts with probability greater than 1 - &c. (where &c is a positive constant) then ğ is also equal to the true ş by definition of &g. This means μ4 0, Vi e [m/T], so the sets are disjoint.
m/T). T and both Bob and
4.3 SHOWING 3-MAX-IP IS HARD
We define the appropriate gap 3-MAX-IP problem, which we use as our intermediate hard problem. Definition 4.6 (Gap approximate maximum inner product search (Gap- Maapp-axlPPn,,,. d, (,)))). Suppose the following conditions hold
We use t >> 0 represent threshold parameter.
We to represent an accuracy parameter.
Suppose n, d denote two positive integers.
8

---

Published as a conference paper at ICLR 2024
Given three sets of points.A {a1,... .an), B {bi, .b.. C for.-. ,ca. c []
For every index E [n], we need to distinguish the following two cases
.
Case 1. There exists pair (J1, J2) E [r]] [n] such that (0s bj2 C/2) 2 L.
Case 2. For all pairs (f1,22 E [n]] [n] we have (at, (ai,bj,: : Oj2) S - c) -t.
Implicit in previous work (Rubinstein (2018)) is proof that the analogue of Gap-Mapt with two sets of points is hard. Here we generalize this to three sets.
Theorem 4.7. Unless SETH and OVC are false, the following holds: for every 5 0 there are constants 04 > Q2 > 0 such that for integer solving Gap-MaxlP(nn 01 log 02 logn, n, = 1/2) requires time (2n18-6).
Proof. We reduce from 3 OV to Max-M. Let gov 8/2. Our reduction takes as input an instance (Aov, Bov. Cov) of orthogonal vectors over {0, 1mm. These sets have sizes |Aov| = |Bovl= Cov| 2m/e for a constant e depending on dov from Definition 4.2 and Conjecture 4.3, and 3OVC posits there no algorithm solving this problem in time (((2m/o)3-6ox)
For constant k 0 to be determined, pick > 0 to be constant such that kclog log(1/e) < 8/2. log(1/e)
log(1/e) We the protocol of Theorem 4.5, instantiated with parameter T T(e) Ol Tog l/g
Suppose that 7" ( T)-T) is representing the number of different possible messages sent by Bob and Charlie in the protocol. Let us choose T so that T = (1/cc. For each vector y E Cov, -> we construct a new vector こ {0, 1). ([(T'(X'm xm by setting iig.c.. := iff Charlie send message ic E [T'] on input 7' and randomness E [m]] (The value is independent of in.)
For each vector B E Bov, construct Bob sends message ig E [T"] on input 3 and randomness [m]. (The value is independent ic.)
vector T {0, (T by setting 0 in 1 iff
For each Merlin-message μa E 10, (0(00og))m//T and vector Aov, construct new vector quia E {0, (T"))mmm as follows: 7. "igec = iff Alice accepts on
input a;
message y from Merlin,
message ig from Bob, message ic from Charlie, and randomness
Notice also that the inner product of three vectors (7 D 27) is exactly proportional to the probability that Alice, Bob and Charlie accept on inputs a, 8, 7 and message μ from Merlin.
In particular, if a, B and y are orthogonal (i.e., (a, 3,7) >> 0), then the inner product is most (āria, (,,,), 2 < m/2. Otherwise, there exists μ that Merlin could send make the players accept, meaning that (a 68, 27) m.
In partieular, these can be distinguished by an algorithm for
Gap-MaxIP(n= = 2m/ 20 log" log 1/c/ log 1/c) = 2(T') m, m, 1/2),
which must therefore be hard as solving the original instance of 30V. By 30VC, this means requires time
(|Aovl + |ovl |Cov/) (2m/eg3-bur Bov = n²² n22/2m(fow/e log log(1/e)) log(//) < n2 3-5
where the last step follows from choosing & large enough in the definition of
At the end, we notice that the vectors we construct have dimension 2(T')² m = 0(m) = O desired.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
The authors would like to thank Yichuan Deng, Yeqi Gao, Junze Yin, Lichen Zhang, Ruizhe Zhang, Tianyi Zhou for helpful discussions of attention literature.
REFERENCES
Scott Aaronson and Avi Wigderson. Algebrization: new barrier in complexity theory. ACM Transactions Computation Theory (TOCT), 1(1):1-54, 2009.
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, 2014.
Amir Abboud, Aviad Rubinstein, and Ryan Williams. Distributed pep theorems for hardness of approximation in p. In 2017 IEEE 58th Anmual Symposium Foundations of Computer Science (FOCS), pp. 25-36. IEEE, 2017.
Amol Aggarwal and Josh Alman. Optimal-degres polynomial approximations for exponentials and gaussian kernel density estimation. In 37th Computational Complexity Conference (CCC 2022). Schloss Pngstuhl-Leibniiz-eenttu für Informatik, 2022.
Thomas D Ahle, Michael Kapralov, Jakob BT Knudsen, Rasmus Pagh, Ameya Velingker, David Woodruff, and Amir Zandieh. Oblivious sketching of high-degree polynomial kernels. In Pro- ceedings of the Fourteenth Annual ACM-SIAM Symposium on Discrete Algorithms (SODA), pp. 141-160. SIAM, 2020.
Josh Alman and Zhao Song. Fast attention requires bounded entries. In NeurlPS. arXiv preprint arXiv:2302. 13214, 2023.
Josh Alman, Timothy Chu, Aaron Schild, and Zhao Song. Algorithms and hardness for linear algebra on geometric graphs. In 2020 IEEE 61st Annual Symposium on Foundations of Computer Science (FOCS), pp. 541-552. IEEE, 2020.
Sanjeev Arora and Boaz Barak. Computational complexity: modern approach. Cambridge University Press, 2009.
László Babai. Trading group theory for randomness. In Proceedings of the seventeenth annual ACM symposium on Theory of computing, pp. 421-429, 1985.
Arturs Backurs, Moses Charikar, Piotr Indyk, and Paris Siminelakis. Efficient density evaluation for smooth kernels. In 2018 IEFE 59th Annual Symposium on Foundations of Computer Science (FOCS), pp. 615-626. IEEE, 2018.
Aditya Bhaskara, Moses Charikar, Ankur Moitra, and Aravindan Vijayaraghavan. Smoothed analysis of tensor decompositions. In Proceedings of the forty-sixth annual ACM symposium on Theory of computing (STOC), pp. 594-603, 2014.
Aditya Bhaskara, Aidao Chen, Aidan Perreault, and Aravindan Vijayaraghavan. Smoothed analysis for tensor methods in unsupervised learning. Mathematical Programming, pp. 1-51, 2020.
Jan van den Brand, Zhao Song, and Tianyi Zhou. Algorithm and hardness for dynamic attention maintenance in large language models. arXiv e-prints, pp. arXiv-2304, 2023.
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 (NeurIPS), 13:1177-1901. 2020.
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 (NeurIPS), 34: 17413-17426, 2021.
10

---

Published as a conference paper at ICLR 2024
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.
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: /904. /044/// 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. ICLR. arXiv preprint arXiv:2009.1475 147944 2021.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311 2022.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXivv2307.08691, 2023.
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 (NeurIPS), 33: 1476-6489) 2020.
Yichuan Deng, Zhao Song, and Junze Yin. Faster robust tensor power method for arbitrary order. arXiv preprint arXiv:2306.00406 2023.
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer, Llm. int8 (): 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.07339, 2022a.
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Gpt3, int8 O: 8-bit matrix multiplication for transformers at scale, Advances in Neural Information Processing Systems, 35: 30318 -70322, 2022b.
Jacob Devlin, Ming-We Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint /rXiv:/81000044 2018.
Huaian Diao, Zhao Song, Wen Sun, and David Woodruff. Sketching for kronecker product regression and p-splines. In International Conference Artificial Intelligence and Statistics, pp. 1299-1308. PMLR, 2018.
Huaian Diao, Rajesh Jayaram. Zhao Song, Wen Sun, and David Woodruff. Optimal sketching for kronecker product regression and low rank approximation. Advances neural information processing systems, 32, 2019.
Yeqi Gao, Zhao Song, Weixin Wang, and Junze Yin. A fast optimization view: Reformulating single layer attention in I1m based on tensor and svm trick, and solving it in matrix multiplication time. arXiv preprint arXiv:2309.07411 07448, 2023a.
Yeqi Gao, Zhao Song, and Xin Yang. Differentially private attention computation. arXiv preprint arXi::2305.0470. 2023b.
Yeqi Gao, Zhao Song, and Junze Yin. Gradientcoin: A peer-to-peer decentralized large language models. arXiv preprint arXiv:2308. 10502, 2023c.
Shafi Goldwasser and Michael Sipser. Private coins versus public coins in interactive proof systems. In Proceedings of the eighteenth annual ACM symposium on Theory of computing, pp. 59-68, 1986.
Valerii Denisovich Goppa. Codes on algebraic curves. Doklady Akademii Nauk, volume 259:6, pp. 1299- Russian Academy of Sciences, 1981.

---

Published as a conference paper at ICLR 2024
Russell Impagliazzo and Ramamohan Paturi. On the complexity of k-sat. Journal of Computer and System Sciences, 62(2):367-375, 2001.
Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer: Fast transformers via sketches for polynomial kernels. arXiv preprint arXiv:2310.01655, 2023.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Flcuret. Transformers mns::
Fast autoregressive transformers with linear attention. In International Conference on Machine Learning (ICLR), pp. 5156-5165, PMLR, 2020.
Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self celf-attention. In International Conference on Algorithmic Learning Theory, pp. 597-619. PMLR, 2023.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04455 2020.
Zichang Liu, Jue Wang, Tri Dao. Tianyi Zhou, Binhuang Yuan, Zhao Song, Anshumali Shrivastava, Yuandong Tian Ce Zhang, Christopher Re, and Beidi Chen. Deja Contextual sparsity for efficient ilms at inference time. In ICML, 2023.
OpenAl. Gpt-4 technical report. https//aariiooo//////00...7777 2023.
Sergio P. Perez, Yan Zhang, James Briggs, Charlie Blake, Josh Levy-Kramen Paul Balanca, Carlo Lushi, Stephen Barlow, and Andrew Fitzgibbon. Training and inference of large language models using 8-bit floating point.
17.2449........fff..... 2023.
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 ICLR. arXiv preprint arXiv:2202.08791, 2022.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre- mee-trining... .> 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jérémy Rapin, et al. Code Ilama: Open foundation models for code. arXiv preprint arXiv:2308. 12950, 2023.
Aviad Rubinstein. Hardness of approximate nearest neighbor search. In Proceedings of the 50th annual ACM SIGACT symposium theory of computing, pp. 1260-1268. 2018
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. arXiv preprint arXiv:2306.022 02896, 2023.
Haihao Shen, Naveen Mellempudi, Xin He, Qun Gao, Chang Wang, and Mengni Wang. Efficient post-training quantization with fp8 formats. arxiv preprint hupp::/aartioorhhhp/2.000555 2023.
Kenneth W Shum, Ilia Aleshnikov, Vijay Kumar, Henning Stichtenoth, and Vinay Deolalikar. A low-complexity algorithm for the construction of -eommeic- codes better than the filbert-varshamor bound. IEEE Transactions on Information Theory, 47(6):2225-22 2241, 2001.
Kenneth Wing-Ki Shum. A low-complexity construction of algebraic geometric codes better than the Gilbert-Varshamvv
Gilbert- Varshamov bound. University of Southern California, 2000.
Zhao Song, David P Woodruff, and Peilin Zhong. Relative enror tensor low rank approximation. In Proceedings of the Thirtieth Annual ACM-SIAM Symposium on Discrete Algorithms, pp. 2772- 2789. SIAM, 2019.
Zhao Song, David Woodruff, Zheng Yu, and Lichen Zhang. Fast sketching of polynomial kernels of polynomial degree. International Conference on Machine Learning, pp. 9812-9823. PMLR, 2021a.
12

---

Published as a conference paper at ICLR 2024
Zhao Song, Lichen Zhang, and Ruizhe Zhang. Training multi-layer wer-paaamettrizd neural network in subquadratic time. arXiv preprint arXiv:2111207628, 2021b.
Zhao Song, Xin Yang, Yuanyuan Yang, and Lichen Zhang. Sketching meets differential privacy: fast algorithm for dynamic kronecker projection maintenance. In International Conference on Machine Learning (ICML), pp. 32418 -24662. PMLR, 2023a.
Zhao Song, Mingquan Ye, Junze Yin, and Lichen Zhang. A nearly-optimal bound for fast regression with fo guarantee. In ICML, 2023b.
Zhao Song, Mingquan Ye, and Lichen Zhang. Streaming semidefinite programs: O O(yn) passes, small space and fast runtime. arXiv preprint arXiv:2309.05133 2023c.
Madhu Sudan. Lecture notes 8 of madhu sudan's class, and scribed by josh alman. In Essential coding theory. http::/peoolll..ca...oo waraaa..aaaard..ed madhusudan/MIT/ST //sc////b////c/tt...pppp 2013.
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.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.1397/ 2023,
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023b.
Michael A Tsfasman, SG Vlädutx, and Th Zink. Modular curves, shimura curves, and goppa codes, better than varshamov-gilbert bound. Mathematische Nachrichten, 109(1):21-28, 1982.
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 (NeurIPS), 30, 2017.
Sinong Wang. Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprins MrXiv:2006.04765 2020.
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. On some fine-grained questions in algorithms and complexity. In Proceedings of the international congress of mathematicians: Rio de janeiro 2018, pp. 3447-3487. World Scientific, 2018
Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pp. 38087-38099. PMLR, 2023.
Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes. Advances in Neural Information Processing Systems, 32, 2019.
Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006.1455 2020a.
Greg Yang. Tensor programs ii:: Neural matrix laws, arXiv preprint arXiv:2009. 2020b.
Greg Yang and Edward Hu. Tensor programs Feature learning in infinite-width neural networks. In International Conference on Machine Learning, pp. 11727-11777 PMLR, 2021.
Ofir Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. Q8bert: Quantized 8bit bert. In 2019 Fifth Workshop on Energy Efficient Machine Learning and Cognitive Cempuuting-euurlP Edition (EMC2-NIPS), pp. 36-39. IEEF, 2019.
13

---

Published as a conference paper at ICLR 2024
Amir Zandieh, Insu Han, Majid Daliri, and Amin Karbasi. Kdeformer: Accelerating transformers via kernel density estimation. In ICML. arXiv preprint 1rXivv2302.02451, 2023.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint M.Xiv:2205.01660 2022.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, and Beidi Chen. H₂O: Heavy- hitter oracle for efficient generative inference of large language models. In NewrIPS. arXiv preprint arXiv:2306. 4048, 2023.
14

---

Published as a conference paper at ICLR 2024
APPENDIX
Roadmap.
In Section A, provide the definitions of several notations. In Section C, provide the running time proofs for our upper bound result. In Section D, we provide the error analysis for our upper bound result. In Section E, we combine everything together, and also present algorithm. In Section B. we show how to reduce our problem to Max-M.
A PRELIMINARY
We write R to denote the real numbers, and write Raxd to denote the set x matrices whose entries real numbers.
For any positive integer n, we write [n]] to denote {1, ... . .
For matrix A E Rnd and indices i [m]. j [d], we write A., to denote the entry of A in the i-th row and j j-th column.
We 1, to denote length-n vector whose entries are all ones.
For vector w E IR,, we use diag(w) E _axx" denote a diagonal matrix with (diag(w))s, = wy; all the other (off-diagonal) entries of the matrix are zero.
D E R is diagonal matrix, we write D- E R for its inverse, which is the diagonal matrix whose i-th entry the diagonal is 1/D,,i, and whose off-diagonal entries are all zero.
For matrix A Raxd, use AT Ødx= to denote its transpose.
For a vector x E R", we use exp(x) E R to denote the entry-wise exponential of xr i.e., the length-n vector with exp (2); exp(x;) for all E [n].
For matrix E R=*",, we similarly use exp(X) E Rax" to denote the matrix with exp(X)), = exp(X.,)).
For any matrix A E Raxd, we define ||A||r = EE Ej=1 222 1/2 to be the Frobenius norm of A.
For vector a, E R", we write (a, b) to denote their inner product, Σ a,b;
For matrix A e Raxd we write || to denote its 6oo norm, i.c., ||A|] == max inje id) |Ar For tensor 7 E Raxxxn, we similarly write || to denote the &oo norm of that tensor T, i.e. ||T|| 00 maxj [Tgg.]-
Definition A.1 (k-wise inner product). For le vectors aj...., ag E Rd we define
(aI...., 0x) = EI a.c
1=
Definition A.2 (3-MAX-IP). Given three sets A, B,C.C. C (0, 13 of vectors where |A| = =B] |C| n1 the goal is to compute
max CEA.
(a, b,,b,c)
B HARDNESS: FROM MaxIP TO OUR PROBLEM
In Section B. we show how to reduce our problem to Gap -aapM. In Section B.2, we present our main lower bound (hardness) result.
B. REDUCTION
We now generalize the hardness proof of Alman & Song (2023) the tensor attention case.
15

---

Published as a conference paper at ICLR 2024
and every C >> Co 0, there exist constants Ca > and C. > and such that, if ATAttC (Definition D.1) for parameters (2n, d = 2C'log n, B C. Vlgg ca Ca) be solved in time T, then Map--MaxP((n,, d= C'logn. n. Co log n. <) (Definition 4.6) can be solved in time O(T + n²=Cx).
Lemma B.1. For every constant C. 0.,0.1), every
Proof. We give algorithm for Gap-Maxx (axIP d Cog n, Co log n,,) (Definition 4.6). Let d1.-.", On b..... bn, c1. E {0, 1] denote the inputs to this problem. Using them, we will construct uppropriate inputs the ATAttC problem that its output will help us detect triples with large inner product.
Let 0 and 2 d be parameters to be determined (in Eq. (5) and Eq. (2) below). Define >> 0 by
T := exp(3/2).
(1)
We pick these parameters so that will be upper bound entries of the attention matrix, namely exp(B(at, by. cd2)
max Je(l
We will use an algorithm for the ATAttC d, ea) problem with parameters:
n := 2n,
d:= 2d
(2)
B = Voogon, where C. /40C/(Cocc.
(3)
Ea == n= where C. == -L(
(4)
Furthermore, set
= B3.
(5)
We define the query and key matrices, Q Ruxd and K1, K2 E Rüxd
b T
cI
0
1
by 0
C
0,
Q 3/B
a
I
K1== VB-
TT
O 1
and K2 = vB.
C
0f
l
02
1
di
di
of
0
l
0
I
0g
I
0
I
0
I
0
1
Since each entry of Q and K1. K2 is either 3/3 or 0, it follows that
|| I V 3 B
| 0o I v/3 3 B
|K2 1o0 S VB B
3.d = =3= = B3.
||Q / <
In terms of these matrices Q E IR R and K Kı 2 K₂ Ra xd
the attention matrix
= exp(QKT /d) E R
is naturally partitioned into eight submatrices
A
A
A
A4
A
A
As
A7
As
where each A; (Vi E [8) is a matrix of size n x n2, [n]..ja ja [n).
defined as follows. For each jo E [n].h e
16

---

Published as a conference paper at ICLR 2024
The (Ju.31 + (f2 - 1)n)-thh entry of A₁ is exx((B((u)o 2/2= S22 + (Id: O2, 0z))/d) exp(3 2/ (22)/d)
The (Jo.Ji + (j2 1 entry of A2 is
exxp(B((a)o رزز 0a) (ajo Oa, la))/d) exp(0)
The (Ju.1 + (j2 1) ) entry of A3 is
-exp(B((0)o- Oz Cj.) (ajo: 1a, 0a))/ = exp(0)= =1
The (jo.J1 + (j2 - -h entry of A is
- -exp(3((ajo Og, 0z) + (Ia, 1a la))/d) exp(Bd/d) =
The (Jo.11 + (j2 - Ith entry of Ag is
- -exp(B((02) by Sta) + (Id, 0g, 0a))/d) = cxp(0) =
The (Jo,h + (j2 - 1) 1t entry of A is
-exp(B((0z: bi. 0a) + (aju. : Oa, 1a)) d)) = exp(0)
The (jo. + (j2 1)m entry of A, is
-exp(3((02 Oa, (32) + (aju 1a. 0a))/d) exp(0)
The (Jo, + (j2 1 entry of Ag is
eexp(B( (0a, Oa, 0a) (1d, 1a, Ix)) expp///))
For each (1,11.f2) [n] x [n] [n], we know that
A + 13m = exp(8 (a;, by cd2
s exp(3 ||ailloo 110 Ileia llo d/d)
< ex)(3/2)
r.
(6)
Here we used the fact that d (see Eq. (2)), and the last step uses the definition (see Eq. (1)).
We also know that for each (2) E [n] x [n] [n].
A A 20 0
(7)
since it is the exponential of an entry of QK /d. By combining our expression for A with Eq. (6) and Eq. (7), we see that
n²r S (A1 S 4n²r, Vi e [m].
Since D,,- = (A1;); it follows that
n²T < Di, < 4n²r,
vi e [?]
Choose the vector e E R (recall that 772 4n²) defined as
0 ,
e
0 0
0,2
17

---

Ca. Published as a conference paper at ICLR 2024
We define as
7:= exp(0.2531/0)/((nm)
8
3
We that 8 ca follows:
exp(0.2531/d-3377 3/2)
12m
1
exp( -x56 + 0.25/1/d)
12m
exp( -05,3 0.258Cc/CC
12m
+ 0.253Co/C - 10
12
1
exp(-0.5CG logn + (.250)(Co/C)) log 2 -2log
12
2 n
Here, the last two steps follow from Eq. (4).
We now use an algorithm for ATAttC(m,, d, B, ca), with the value matrix V with one row o and the rest O. Since f 2 the result is vector u E R such that, for all i E [7]]
144 - (D- Ab)
Recall that for each [n] we need to distinguish between two cases: either there is pair (f1,22) E [n]] such that (a;, , 2 f, else for all pairs (2, E [n] the inner product (as, C12 S (1 et).. We will distinguish between these cases by checking whether u, is greater than threshold value to = 2Ł. We next consider the two cases to see why this is.
Case 1.
For a given E [n], if there are (f1,22 E [n] x [n]] such that (at, bJ 0j2) 2 t, then
Bla, b )2) c/2)/d
0.5 B(a,, by Gjz)/d
2 0.25 Bt/d,
where the 1st step follows from 2d d (see Eq. (2)). This means desired that
u; (xpp0.225tt/d))(4rr -
= -
27
To-
Case 2.
For given i E [n], if for all (f1,12) E [n] [m] we have (a;, bj, Cja) < t(1 c), then
B(a,, cz2) 0.253 (1 c)t/d.
Hence, as desired,
ui (n² (xp(0025.011 -000002000110///0////11++
exp((0.5531/)) (4n² (412/ex((0.255344)) +ī
37. (An2/ exx(().55.eel/))) +7
by definiton of T, see Eq. (8)
s 3F. 4 +ē
18

---

Published as a conference paper at ICLR 2024
2t
tu.
Here, the 4th step follows because, by our choice of B and t, we have
exp(0.25,ee//d d) = exp( log n)/d)
exp (0.25..eBe
= exp(10
16n²;
(9)
where we used that / Co log (by Lemma statement), that d = C' C'logr n. that B = B (Eq. (5)) and the choice of B (Eq. (3)).
B.2 MAIN HARDNESS RESULT
We can finally conclude our main lower bound.
Theorem B.2 (Lower bound, formal version of Theorem 1.3). Assuming SETH, for every @ > there are constants C, Ca. Co > such that: there is no algorithm running in time O(n3-9) for the problem AAttC(n, d C' log n, C. A..loog Eo. =n Co).
Proof. Follows from combining Theorem 4.5, Theorem 4.7, and Lemma B.I.
C UPPER BOUND: RUNNING TIME
In Section C.I. review the standard "matrix" attention computation problem. In Section C.2, we define the "tensor" attention computation problem. In Section C.3, we provide an efficient tool for implementing tensor related computations. In Section C.4, we provide several tools for rearranging tensor computations that we will use in our algorithm.
C.1 CLASSICAL ATTENTION COMPUTATION
We first review the attention computation definition in Vaswani al. (2017); Devlin et (2018); Radford et al. (2018; 2019); Brown et al. (2020); OpenAI (2023); Zandich et al. (2023); Alman & Song (2023); Brand et al. (2023); Gao al. (2023b;c;a).
Definition C.1. Suppose there are three n. x d size matrices Q.K, V E Raxd our plan is to generate the n x d matrix Att(Q, K, V) defined by
Att Q K
D-
V
nxd axd axd
RXn uXn axd
nxd
axd
where A E Buxn and diagonal matrix D E Buxn are defined
exp Q K /d), and D diag( A 1n
aXB
dxn
axn
nxl
C.2 TENSOR ATTENTION COMPUTATION
Given two d matrices, there are two standard variants on their Kronecker product one may consider: The standard Kronecker product (denoted (9) is a new n² x d² matrix, whereas the column-wise Kronecker product (denoted ( is n² d matrix. For more literature on tensor computations and their applications learning algorithms, we refer the readers to Bhaskara et al. (2014); Song al. (2019); Diao et al. (2018; 2019); Yang (2019; 2020a;b); Bhaskara et al. (2020); Ahle et al. (2020); Song et al. (2021a;b); Yang Hu (2021); Song et al. (2023a;b;c); Deng et al. (2023).
Next, we generalize the matrix attention computation (in Alman Song (2023)) into tensor attention computation follows:
19

---

Published as a conference paper at ICLR 2024
Definition C.2 (TensorAtt). Given matrices Q, K1, K2 E Raxd and matrices V1, V2 Ruxd the goal of tensor attention computation is to compute
TensorAtt(Q. K1, K2, V1, V2) = D
nxd
xd
where
A E Ruxn" is defined as A =
Q
K1
K2
nxd
nxd
D E Raxn is defined as D = diag(
A
n²
nxn2
n²
V E IR=2xd xd is defined as V =
V2
nxd
nxd
C.3 EFFICIENT COLUMN-WISE KRONECKER COMPUTATION
We prove an important tool which will be used in analyze the running time of our algorithm. Lemma C.3. If the following condition holds
Let @ be defined Definition 2.4.
Given A1 E Raxdi, A2 TRnxdi we define A = (A1 @ A2) E Rdh'xh,
Given B₁ E Ruxxd, B2 E Rundd, we define B = (B₁ @ B2) E RRn2xd2
We define C E Rd11x2 as := B
We define C= B1, C2 == B2 B₂
Then, we have
Part C/002=
Part Given input A1, A2, B1, B2, we can compute C in Tinat (d1, n, d2) time.
Proof. For each E [n), let II, denote the i-th row of A E Raxd,
For each i [ni], let aT,; denote the i-th row of Ag E Baxd
For each i [n:]. let bT, denote the i-th row of By.
For each i [n:], let bJ, denote the i-th row of By.
For each i [d]. let A.,; E T denote the i-th column of matrix A x
Recall that C1 E dR xda and C2 Rdiixds
C A C2:= := A. B2
Thus, we see that
(Ci)k
(C2)
Σ a2 b2
20

---

Published as a conference paper at ICLR 2024
Then, write C' Rdxx as
=
di xd2
di xxe n xd2
n2²
A.s B.
1xd2
A 1)n B B
1xd2
(11s o (2.5)
(bi o b2.j)
(10)
x1
1xd2
where the first step follows from definition of C E Rddd the second step follows from the matrix can written as the summation n² rank-1 matrices, the third step follows from changing the index, the forth step follows from Ax,++0-1jn
01,5 02.
x1
d,x1
dyx1
From the above, we can calculate that the entry of C' in location kı, ki is
.k
(a1 02.)k; (b1,s be De
01 02.j. b1 b2.j.
k2 (2.j.k))
(C1) k2 (C2)
where the first step follows from Eq. (10), the second step follows from simple algebra, the third step follows from separating the summation over i and the summation over j, and the last step follows from definition of matrices C1 and C2.
Thus, we can conclude
C C= o Cz.
The algorithm will first compute C1 and C2, whic takes Tiat (di, d2) time. Then it calculates C Cz, which takes (d2d2) time.
C.4 SIMPLE EQUIVALENT TOOLS FOR TENSOR NOTATIONS
We define standard tensor notation, for example see Song al. (2019).
Definition C.4 ((.,..) tensor operator), Given tensor T E R XM2Xn, let X E Rmixd Y E Xd, Z E R xds
We define T(X, Y, Y, E pdd xd2xds as follows
(XX,Y2)i,j
YojzKo Va [die) [da], c E [da].
Next, we present several equivalence results for tensors.
Lemma C.5. the following conditions hold
Let @ be defined Definition 2.4.
21

---

Published as a conference paper at ICLR 2024
Let @ be defined Definition 2.2.
Let ((yi)) operator be defined as Definition C.4.
Let be defined Definition 2.1.
Ler Q, K2, V1, V2 E R
Let AERANN' be Q(1i 2 K2)
Le be Q @ K @ K2.
Then, we have
Part 1. A = Aiji.i for E [m], E [n],ja E [m] (This means A can be viewed the lensor version of A)
Part 2. ,1,, A(I, I,; 1a) Q o (1, (2)
Part 3. A(V 2 V2) Q(K 2 (22T(V( 2 122 Q((K V1) (K. V2))
Proof. Proof of Part 1.
Directly follows from definition of A and A.
Proof of Part 2.
Follows from tensor notations in Definition 2.2 and Definition C.4.
Proof of Part 3.
Directly follows from applying Part 1 of Lemma C.3 here.
D UPPER BOUND: ERROR ANALYSIS
In Section D.1, we provide the definition of approximate tensor attention computation. In Section D.2, we state a polynomial approximation tool from previous work. In Section D.3, we show bound on the entries of the attention matrix. In Section D.4, we provide a low-rank decomposition for the tensor version of the attention matrix. Finally, in Section D.5, we compute the error propagation from A to D, then in Section D.6, analyze the propagation from A and D to the attention matrix.
D.1 APPROXIMATE TENSOR ATTENTION COMPUTATION
Definition D.1 (A tensor generalization of standard attention computation, restatement of Defini-
tion 1.2). Let ea B 0 be parameters. Given five matrices Q, K1, K2, V1, V2 E Raxd such that
101 B. |111 B, 122 S B. 1V B. V22 B
Our goal is to find a matrix T E Raxd which can entry-wisely approximate D-1 AV, in particular, means the following Eo norm guarantee,
- DD DV'IV cc
Here,
A: exp(Q(Ki 2 T) flattening of an n. X n X n tensor)
= (We remark that can also view matrix A as the
V Vi Ø V2 R
D diag( A1 RRxx is n. size positive diagonal matrix.
Notice that the straightforward algorithm for this problem will spend at least 2(x²2) time to write the matrix A (we also think of A as tensor that has size n).
22

---

Published as a conference paper at ICLR 2024
D.2 AN ERROR CONTROL TOOL FROM PREVIOUS WORK
We state tool from previous work.
Corollary D.2 (Corollary 2.2 in Alman & Song (2023)). Suppose the following conditions hold
Let B 1.
Lete E000..1)
Let := O(max( (1/cl/c)
logtlog(/ee//B)> B}).
There is polynomial P R + R of degree-g such that for all æ E [-B, B), we have
c) exp(x) P(x) (1+ exp(x)).
D.3
TENSOR Q @ K1 K2 HAS BOUNDED ENTRIES
Lemma D.3 (Bounded entry). Suppose the following conditions hold
Suppose B 2
Assume matrices Q, K1, E Ruxd have 11 I B, |Klloo & B. ||22llxc S B.
Let @ operation be defined as Definition 2.2.
Then, we have
IIQ @ K @ K2/d|loo S B
Proof. For every index triple (i,j1: [m] [n] x [n], we are able to prove
[QK Jinh
Σ (K))
1=1
< ||0 |K, 11,0 ||
< dd B²,
where the 2nd step is because triangle inequality, the 3rd step is using fox norm bound on Q, K1, K2. Now, complete the proofs.
D.4 TENSOR LOW-RANK APPROXIMATION
In the following definition, we view the n X n² size matrix an X X n size attention matrix. Definition D.4. Assume the following parameters regime,
We 2 1 to denote positive integer.
We use e E (0,0.1 to represent an accuracy parameter.
Suppose there is 3rd order tensor A R B
We say 3rd order lensor À E mxxaxa is an (e, rapproximmaiinon of tensor A if
20
Ā = UUQQU U2 2 Us for some matrices U1, U2, U3 Buxr (i.e.. À has rank most r), and
A Aiguual S Aijiis for all (i, (i, [n] [n] x [n).
23

---

Published as a conference paper at ICLR 2024
D.5 FROM A TO D
In this section and the next, we generalize the proof of Alman & Song (2023) for error propagation from the matrix setting to the tensor setting. The proofs are nearly identical.
Lemma D.5. Let A E TRaxn' be matrix with positive entries, and 6 (0,0.1) be any parameter.
Let A Raxn" be an approximation to A, meaning for all (i,1) E [n] [n²], we have
A A,,l S CA A.j-
We consider two diagonal matrices D, D paxx which can be formally written as D == diag( mA, and D diag( (Al,)) Then, for every index i [n], the following bound holds
1De. Ds.el S cA Des-
Proof. We calculate that
n²
Di.. - Di,ll
Σ
Ac, Σ
Ard
1=1
j=1
n²
A - Ari
1=1
n²
6
Aij
1=1
E
Di.i-
where the second step follows from triangle inequality. This completes the proof.
D.6 FROM A, D TO TENSOR ATTENTION
The goal of this section is to prove Lemma D.6. Lemma D.6. Suppose the following conditions are true
Let CA E (0.0.1)
*
Let B be bounded parameter,
We use V = -V E g²² xd to represent matrix with |[ < B².
Let A E R be positive matrix,
and let let Br R be matrix such that, for every tuple (i, E [r]] x [n²]] have
A 16i-Apa A,,l S EA A
Suppose D, D E TRun are diagonal matrices with positive diagonal entries, and such that for every index E [n], we have
DD. Di.el S ED Dr.i-
Then, we have
15-1V D-- S (EA+ B²
24

---

Published as a conference paper at ICLR 2024
Proof. By the triangle inequality, we know
-D-IAV D' D- & I | - D ÄVIl + +|ID-1 ÄV - DD D-'AVllo. (11)
We bound each of these two terms to get our desired result.
First of all, for every index pair (i.j) E [r]] [d]
n² AD- D- D- =
Σ
(D -DT.)
Au Vasl
1=1
n²
10 i,
D Aud |V
a
Dis Dest A., ||V|
Di,
D Ar,l- ||
S <D
D A.gl V
<p
1=1
(D |V| S ED B
(12)
Here the 2nd step uses the triangle inequality, the 4th step follows from the assumption that [(Pi, Dr.))/D Da,al S Ep. the 5th step follows because Dj11 and As, are positive numbers, and the final step follows by foo norm of V is bounded by B²
Second, for every (i,j) [m] x [d].
a² D - (A,) Visl
1DD 1 ĀV - D- D--AVV- AV) =
1=1
D I(A,,, - A.,)) -IV
1=1
Σ
D 1(Ar,
Σ Djd -EAAi. B²
= EA B2.
(13)
Here, again, the 2nd step the triangle incquality, the 3rd step follows because D is positive, the 4th step follows from the assumption that [A As,l S < Aj,1 and the final step follows by definition of Dr,--
The Lemma conclusion then becomes true by substituting Eq. (12) and Eq. (13) into Eq. (11).
E UPPER BOUND: PUTTING IT ALL TOGETHER
In Section E. we provide a low-rank decomposition for approximating the original attention tensor. In Section E.2, we calculate the running time of constructing that low-rank decomposition. In Section E.3, we put everything together, and prove main upper bound theorem.
25

---

Published as a conference paper at ICLR 2024
E.1 DECOMPOSING INTO U1, U2 AND U3
The goal of this section is to prove Lemma E.I. Lemma E.1. If the following conditions hold
We use M := X Xoooo Ruxax to represent tensor that is constructed by X,Y, Z R
Let P(x) denote a degree-g single-variable polynomial. We apply P(M) entry-wisely, i.e, P(M)isa = P(Ma,,)
Let be rank parameter that is (S(g+d)).
3g
There is an algorithm running in time ( which, given as imput X, Y, Z, constructs matrices U1, U1,U2,U3 gaxr such that P(M) U1 @ U2 U3.
Proof. Expand P as sum of monomials as
P(x) Σ x².
Consider the function Rd Rd Rd defined by, for u, u, Rd
K(u, w) P(uu, w)).
We define set V and provide names for variables in set V in the following sense,
-u..... Ud U1. Va, W1.-.. wa)
Thus, function K can be viewed degree-3g polynomial in the 3d entries in V of the vectors u, 8, count the number of monomials.
We define set F as
F= ff: f0.1. 10,1,2, .3.3)
f(v) 3g
oVV
Let us count the size of set F
(3d + 3g
F
3g
There exists coefficients tother E R such that
K(u,u,w)
Σ
c
II (B)) BEV
IEF
We define partitions of V:
Vg = {ui,"" ud):
V. = fur,..., val, Vu == fwi, w
We define du . + F by. for any t E F,
Ou(u,, (dd) = Ct II
((as)
u
N; N.EV.
Similarly, define d. Rdd RR by, for any F.F.
D0 ... Ud)) = II
((%.))
weV.
26

---

Published as a conference paper at ICLR 2024
and we define ow IR + R by, for any E F,
Φ (w1, wa)t
f(ws)
EV.
Here, can view K function as
K(u, w, = (du (o( (m)))
For every index i E [n)] suppose X; Rd is the i-th row of X, assume Y; R is the i-th row of Y. and let Z, E Rdd denote the i-th row of Z.
Therefore, we should construct three matrices U1, U2 and U3 as the following way, for each E [n]
the i-th row of the matrix U1 E [ is the vector (h. (3:),
the i-th row of the matrix U2 E R is the vector @i (y:),
the i-th row of the matrix U3 E I is the vector @i (z:).
These n x r matrices can be constructed in time (nnrg) the straightforward way, since each entry depends on variables.
Algorithm 1 Our Polynomial Method Tensor Attention Algorithm
procedure POLYTENSOOATTENTIONNQ E Raxd K₁ E Ruxd K2 E Baxd x Vi E Raxd V/2 Raxd E N4, d E N4, B > 0, (0, 00))
D. Theorem 1.4
2:
n can viewed as the length the sentence D d can be viewed as the feature of dimension é is the accuracy output B
3:
4:
5:
ma |1 K2ll
Vol [/ V/2
+ (max log(1/c) log log( B*)
6:
r F +=(3(++d)
7:
8:
*SStep 1$/
9:
Construct U1, U2, U3 E Ruxr via Lemma E.2
b O time
10:
StSpp 2
11:
U1
(U2 U3)
1 1.) n
O(rr) time
nxr
rxn2
12:
Stepp 3+/
13:
D-1 = diag(w 11-
D O(n) time
*Step 4%/
15:
Compute (U2 2 Uz) ((V Ø V2) Rrdd
Takes Tiaa (r. n, d) time
Step 5*1
17:
Compute Un ((U2 @ Uz) [Vi 2 V2))
Tmua (n, r, d) time
StSepp 6*/
19:
T D (U1 ((U2 @ U3) T (V1 0 V2)))
O(nd) time D T E Raxd
20:
return T
21:
end procedure
E.2
TIME FOR CONSTRUCTING U1, U2 U3
Lemma E.2. Suppose five matrices Q, K1, K2, Vi, V2 Raxd satisfy
1Qlloo S B,
|Kilo B, /22 I B,
|V B. and ||V & B.
27

---

Published as a conference paper at ICLR 2024
We define tensor A = exp(Q @ K1 C K2/d) E R
For bounded number >> and accuracy parameter E (0, 1), there are positive integers and the condition for
9 0 max
log(L/e)
Ba
(og(log(1/c)/
the condition for T:
3(g + d)
3g
such that:
There is third order tensor E R that
is (c, approxxmmatnn (Definition D.4) of A RaxAX.
Let U1, U2 and U3 E Bxx be the matrices defining A, ie., À UU-00
it takes (nr) time to construct U1, U2 and U3.
Proof. We define M = Q @ Ki C K2/de R=xax., Using Lemma D.3, we can show that || B².
Recall that the definition (e, (e,r))-aprrone apppoopoxiitaionn can be found Definition D.4.
Next, we will apply Corollary D.2 (with replacing B by B*), there is degree-g polynomial P R R such that the tensor À ((M) is (e, approximatiiii to tensor A. Here apply P to M entrywisely.
We then compute U1, U2, and U3 using Lemma E.I, which gives the bound
[3(g + d)
<
3g
Therefore, finish the proof.
E.3 MAIN ALGORITHMIC RESULT
We present our main algorithmic result as follows: Theorem E.3 (Upper bound, formal version of Theorem 1.4), There is an algorithm (Algorithm 1)
that solves ATAttC(n, d = O(log n), B = of Vlog n), ca = 1/ poly(n)) in time n (+o(1)
Proof. Proof of Running Time.
Using Lemma E.2, know that Step (in Algorithm 1) can be implemented in O(nrg)
time
Using Lemma C.3, we know that Step (in Algorithm 1) can be implemented in O(nr) time
Step can implemented in O(n) in a straightforward way.
To compute Step 4 efficiently, need to use Lemma C.3 again.
Computing Step 5 is just standard matrix multiplication
Step 6 is just rescaling the d matrix
Proof of Correctness.
We combine Corollary D.2, Lemma D.5, Lemma D.6, and simple algebra.
28


Reviews:
[{"review_id": "75riXveNic", "reviewer": "Reviewer_dthX", "summary": "This work explores the computational complexity of generalized matrix attention scheme, which is used to capture high-order correlation among the input sequence. \nTo capture the correlation among triplets of tokens in an input sequence of length $n$, the generalized attention scheme outputs an attention matrix by computing the column-wise Kronecker products based on one query matrix, two key matrices and two value matrices. \nIn such a case, this work shows if the max entries of the query, key and value matrices are bounded by $o(\\sqrt[3]{\\log n})$, one can compute an approximation to the generalized attention matrix in almost linear time. On the other hand, if the max entries of the input matrices are at least $\\Omega(\\sqrt[3]{\\log n}$), one cannot hope to efficiently approximate  the attention matrix in less than cubic time, assuming SETH. The latter hardness result is shown by reducing from the Gap Max IP problem, whose hardness is then shown through a combination and generalization of previous techniques.  Furthermore, the work shows the techniques developed above can be extended to characterize the gap in computational complexity in $k$-th order generalization of the attention scheme.", "questions": "I am not a complexity expert. I have no comments on the proof techniques presented in this work. However, I do have a few questions for the authors.\n\n- In Definition 1.2, why is approximating the higher order attention matrix in the $\\ell_{\\infty}$ norm considered a good metric to evaluate the approximation of a matrix that captures higher order correlation? Why not the other norms?\n\n- In Section 3.1, does “$\\widetilde{D} \\approx D$” mean $\\tilde{D}$ and $D$ close in the $\\ell_{\\infty}$ norm? (and so does “$\\widetilde{A} \\approx A$”?)\n\n- In Section 3.1, why can $\\widetilde{A}(V_1 / V_2)$ be computed in $O(n^{1+o(1)})$ time, while $\\widetilde{D}$ needs to be computed in $O(nd)$ time?\n\n- Why does the construction of the algorithm in Section 3.1 fail when there are large entries $\\Omega(\\sqrt[3]{\\log n})$ in the input matrices?\n\n- In Section 4, what is the major challenge of extending the three-party communication protocol to a four-party communication protocol in Section 4.2? Why does one need to use the algebraic geometry code?\n\n- In Section 4, where does $B = O(\\sqrt[3]{\\log n})$ pop up in the LB proof?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- This work considers an interesting problem of computing generalized matrix attention scheme. Since the generalized schemes involves computing the Kronecker products between a set of matrices, this is apparently a computationally expensive operation. It is hence natural to explore the computational complexity of this problem.\n\n- Both the upper bound and the lower bound results (especially the latter one) presented in the work are interesting. \n\n- A high-level summarization and intuition behind the techniques used to derive the upper and the lower bound helps the reader.", "weaknesses": "The presentation needs to be improved. \n\n- In Section 3.2 hardness, “3-Max IP” and “Gap-MaxIP” seem to refer to the same problem. It is confusing to give two names to the same problem.\n\n- It might be clearer to present the upper bound (UB) and the lower bound (LB) in two separate sections, instead of giving an overview of the UB + LB, and then elaborate on the LB.\n\n- Is it possible to give a few sentences of description of the mysterious $U_1, U_2, U_3$ matrices and how to compute them in Section 3.1?\n\n- In Section 4 “hardness” which elaborates on key steps in showing the LB, presenting Hypothesis 4.1, Definition 4.2, Conjecture 4.3 and Theorem 4.4, all of which are from prior works, in the main paper do not help much on understanding and appreciating the novelty / challenges addressed in extending and generalizing current proof techniques to show the LB on computational complexity for approximating the generalized attention matrix. Some of them can indeed be moved to the Appendix. It would be better to give more intuition on the (technical) difference between the three-party and four-party communication protocol that computes set disjointness, how algebraic geometry code is applied in extending the proof from three-party to four-party communication and how the new protocol is used in showing the LB on computation time of Gap Max-IP.\n\n- Minor issue: the last paragraph in page 5 states “showing that the same line of attack”? “attack” here means “techniques”, I assume?", "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": "JJq04YVtq8", "reviewer": "Reviewer_FrM6", "summary": "The paper studies the computational aspects of a so-called \"third-order tensor attention variant\" recently defined by Sanford et al (2023). While standard attention captures pairwise interactions between tokens, the third-order tensor variant was suggested in the context of capturing triplet interactions. \n\nThe current paper shows that if the input entries are bounded, and if it suffices to compute each entry of the attention output approximately rather than exactly, then there is an almost-linear time algorithm for computing this operation. The bound on the entry magnitude is asymptotically (conditionally) tight, as the paper also shows that without this bound, computing this attention output in time significantly better than the trivial n^3 time computation is SETH-hard, in the complexity theoretic sense. Both the algorithm and the hardness result are generalizations (in terms of both the results themselves, and the techniques used to prove them) of a recent work of Alman and Song (2023) that proved them for standard attention. The results also extend to tensors of higher orders (than 3).", "questions": "I'd be interested to hear from the authors what do they consider to be the importance and implication of their algorithm, and whether they deem it implementable within a neural network architecture?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is generally well-written and the mathematical content seems interesting.", "weaknesses": "The paper is purely theoretical and seems quite removed from application. It is entirely about a form of tensor attention that has been suggested as a bit of an afterthought in a recent work (Sanford 2023), that deemed it likely impractical and anyway did not implement it. Thus it is not about an architecture that is actually used or presently considered usable. This raises the question of what is the import of showing the existence of an almost-linear time approximation algorithm for it, and whether this algorithm makes sense as part of a neural network (I believe the paper does not touch on this point). A negative outlook on this paper would be that the interesting results and the conceptual message in this line of research were already given in Alman and Song (2023), and the extension to higher-order tensors in this manuscript might be an elegant intellectual and mathematical exercise, albeit without much consequence to the ML community. Nonetheless, since the content does seem elegant, and there is always the issue of benefit-of-doubt about whether a piece of theoretical research would have implications down the road, I prefer to vote for acceptance.", "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": "ZXzkWJ7NVc", "reviewer": "Reviewer_7Xxi", "summary": "This paper solves the transformer attention scheme over triplets under some mild complexity theoretical assumptions.", "questions": "Page 2:\nCould you please explain the column-wise Kronecker product of V_1 and V_2 here as well?\n\nPage 3:\nPlease elaborate on SETH.\nI know it is standard, but perhaps too strong? :)\n\nPage 4:\nPlease provide some intuition for Definition 1.5.\n\nPage 5:\nI do not understand \"Approximating A.\"\n\nPage 6:\nCould you please sketch an short example for the reduction from GapMaxIP to ATAttc?\n\nPage 8:\nCould you please add some discussion about Theorem 4.7?\n\nPage 9:\nLine 12 from bottom:\nWhy is there such a \\mu?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Originality:\nThis work is original to the best of my knowledge.\n\nQuality:\nQuality is high.\n\nClarity:\nWriting is clear.\n\nSignificance:\nThe results of this paper are important due to their connections to LLMs and other AI applications.", "weaknesses": "None.", "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 (spotlight)