How to Capture Higher-order Correlations? Generalizing Matrix Softmax Attention to Kronecker Computation

Josh Alman, Zhao Song

Introduction

Large language models, such as Transformer , BERT , GPT-1 , GPT-2 , GPT-3 , PaLM , OPT , GPT-3.5, Bard, GPT-4 , Llama , Llama 2 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 of 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 a length- $n$ input is given to the attention unit (like a sentence or paragraph of $n$ words), we embed it into three matrices $Q,K,V$ (the query, key, and value token matrices) where each has $n$ rows and $d$ columns. Here $d$ is the feature dimension; one has $d\ll n$ in the long sentence regime. Mathematically, the attention unit computes $D^{-1}\exp(QK^{\top})V$ , where $D=\diag(\exp(QK^{\top}){\bf 1}_{n})$ is a diagonal matrix, ${\bf 1}_{n}$ 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^{\top}$ . However, if the input data has correlated triples of tokens, it is not clear an attention unit can detect this.

A recent and exciting work formalized this intuition. They defined a simple task about learning correlations between triples of words, and showed that attention units are unable to solve it. By contrast, they are able to solve the analogous problem of learning correlations between pairs of words. Toward resolving this, proposed a generalization of attention computation:

Given as input $n\times d$ matrices $Q,K_{1},K_{2},V_{1},V_{2}$ , the goal is to construct another $n\times d$ matrix

${\bf 1}_{n^{2}}$ here denotes a length- $n^{2}$ vector whose entries are all ones.

One may naturally view $A$ as an $n\times n\times n$ 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, 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 is critically important, since it forms the time bottleneck of LLM training and inference. By generalizing attention to make it more expressive, have we also made it intractably slow?

To answer this question, we focus on an approximate version of the tensor attention computation problem. In practical applications, it is sufficient to approximately perform these computations , and this often helps lead to faster algorithms.

$\|Q\|_{\infty}\leq B$ , $\|K_{1}\|_{\infty}\leq B$ , $\|K_{2}\|_{\infty}\leq B$ , $\|V_{1}\|_{\infty}\leq B$ , $\|V_{2}\|_{\infty}\leq B$

the other matrices are defined as in Definition 1.1 above.

We focus here on the natural setting with $d=O(\log n)$ (so that we are modeling long sequences) and $\epsilon_{a}=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. showed that the trivial $O(n^{2})$ time algorithm is essentially optimal for exact computation, assuming the Strong Exponential Time Hypothesis ( $\SETH$ ). $\mathsf{SETH}$ is a popular conjecture from fine-grained complexity which posits that one cannot substantially improve our current best algorithms for $k$ -SAT; see the survey for more details.

studied the approximate (non-tensor) attention problem and showed that its complexity depends on the magnitude of the entries of the matrices $Q,K$ : If they are smaller than $o(\sqrt{\log n})$ , then there is a fast algorithm running in time $n^{1+o(1)}$ ; this near-linear time algorithm is essentially as fast as one could hope for. On the other hand, if they are at least $\Omega(\sqrt{\log n})$ , then there is no algorithm substantially faster than the trivial $O(n^{2})$ assuming $\SETH$ . This theoretical result mirrors practical observations that bounded entries are essential for fast attention .

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 $\mathsf{AAttC}$ depends on the parameter $B$ , the magnitudes of the entries in the query, key, and value matrices.

We first show a lower bound, that when $B\geq\Omega(\sqrt{\log n})$ , it is impossible to design a truly subcubic-time algorithm (assuming $\SETH$ ). Note that the straigtforward algorithm for this problem runs in cubic time, so our result shows that one cannot substantially improve on the straightforward algorithm when the entries have magnitude at least $\Omega(\sqrt{\log n})$ .

Assuming $\mathsf{SETH}$ , for every $q>0$ , there are constants $C,C_{a},C_{b}>0$ such that: there is no algorithm running in time $O(n^{3-q})$ for the problem $\mathsf{ATAttC}(n,d=C\log n,B=C_{b}\sqrt{\log n},\epsilon_{a}=n^{-C_{a}})$ .

Our second result is a new algorithm, showing that when $B<o(\sqrt{\log n})$ , then there is an almost linear time algorithm for solving the problem.

There is an algorithm (Algorithm 1) that solves $\mathsf{ATAttC}(n,d=O(\log n),B=o(\sqrt{\log n}),\epsilon_{a}=1/\poly(n))$ in time $n^{1+o(1)}$ .

Our Theorems 1.3 and 1.4 together show that the complexity of $\mathsf{ATAttC}$ has a very tight transition at $B=\Theta(\sqrt{\log n})$ . When $B<o(\sqrt{\log n})$ is smaller than the threshold, the problem can be solved essentially as quickly as one could hope for, in time $n^{1+o(1)}$ . Meanwhile, when $B\geq\Omega(\sqrt{\log n})$ is greater than the threshold, it is impossible to achieve a subcubic running time, no matter what algorithmic techniques are used (assuming $\mathsf{SETH}$ ).

It is exciting that, even for the more expressive tensor generalization of attention, there is a near-linear time algorithm in the bounded entry regime. Interestingly, though, the bound must be smaller than for regular attention: for regular attention to have a near-linear time algorithm, it is necessary and sufficient that $B<\sqrt{\log n}$ , whereas for tensor-based attention, we show it is necessary and sufficient that $B<\sqrt{\log n}$ .

More generally, for any positive integer $k\geq 2$ , we study a higher-order tensor generalization of attention which can detect $k$ -wise correlations. (Regular attention corresponds to $k=2$ and $\mathsf{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<\sqrt[k]{\log n}$ , and that the trivial $O(n^{k})$ 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 or inference, we let $B$ be the largest magnitude of an entry, then we select the largest $k$ for which $B<\sqrt[k]{\log n}$ , and we can quickly perform $k$ -th order attention computations for our tokens, but not higher-order attention.

Suppose we are given $n\times d$ matrices $Q,K_{1},K_{2},\cdots,K_{k-1}$ and $V_{1},V_{2},\cdots,V_{k-1}$ , our target is to construct another $n\times d$ matrix

${\bf 1}_{n^{k-1}}$ is the length- $n^{k-1}$ vector whose entries are all ones.

In Section 2, we provide a number of basic notations and definitions. In 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

Tensor Operations

Many of our proofs will involve manipulating tensors. Here we introduce three different tensor operations we will frequently use.

We note that a tensor $T$ can be written in the form $A\odot B\odot C$ like this if and only if its tensor rank is at most $d$ .

In this work, we will primarily use the following column-wise version of the Kronecker product.

Matrix Multiplication

Technique Overview

Generalizing prior work on the computational complexity of the attention problem to our tensor generalization requires overcoming a number of technical challenges. Here we summarize our approach, with an emphasis on differences with the prior work on (non-tensor) attention that we build on .

We begin by introducing a basic tool for manipulating computations involving the column-wise Kronecker product $\oslash$ (see details in Lemma C.3 below). Define the following matrices.

Approximating DD

In order to perform generalized attention, we aim to compute the matrix $D=\diag(\exp(Q(K_{1}\oslash K_{2})^{\top}/d){\bf 1}_{n^{2}}).$ Notice that the intermediate matrix $\exp(Q(K_{1}\oslash K_{2})^{\top}/d)$ has $n^{3}$ entries. We thus cannot compute it in subcubic time. We instead aim to use an implicit representation of an approximation of this matrix which can be quickly manipulated.

Toward this goal, we find appropriate matrices $U_{1},U_{2},U_{3}$ (which we discuss in more detail shortly) and formulate $\widetilde{D}=\diag(U_{1}(U_{2}\oslash U_{3})^{\top}{\bf 1}_{n^{2}})$ such that $\widetilde{D}\approx D$ . Given the matrices $U_{1},U_{2},U_{3}$ , and using the above tool for $\oslash$ , we can compute $\widetilde{D}$ quickly in $O(nd)$ time.

Approximating AA

Finding approximating matrices U1,U2,U3U_{1},U_{2},U_{3}

Thus, it remains to find matrices $U_{1},U_{2},U_{3}$ which appropriately approximate $D$ and $A$ as above. We show how to efficiently find such matrices as long as the inputs $Q,K_{1},K_{1},V_{1},V_{2}$ have bounded entries. The key idea is to use the polynomial method, a key tool from prior work which allows one to find low-rank representations of matrices.

The method generally says that if $M$ is a low-rank matrix, and $p$ is a 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 ), we get a low-rank approximation of $\exp(M)$ .

This polynomial method approach was also taken in the prior work on (non-tensor) attention . 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\times n\times n$ tensor and an $n\times n^{2}$ matrix allows us to take advantage of this low-rank tensor approximation as well as the aforementioned matrix multiplication algorithms. See details in Lemma E.1.

2 Hardness

Our hardness proof proceeds by introducing and considering a new intermediate problem we call ${\sf Gap}{\rm-}{\sf MaxIP}$ (Definition 4.6). In this problem, one is given as input $3n$ vectors $a_{1},\ldots,a_{n},b_{1},\ldots,b_{n},c_{1},\ldots,c_{n}\in\{0,1\}^{d}$ as well as a threshold $t$ , and the goal is to distinguish between the cases

$\langle a_{i},b_{j},c_{k}\rangle\leq t$ for all $i,j,k\in[n]$ , or

$\langle a_{i},b_{j},c_{k}\rangle\geq 2t$ for some $i,j,k\in[n]$ .

We first prove that ${\sf Gap}{\rm-}{\sf MaxIP}$ cannot be solved in truly subcubic time assuming $\SETH$ . We then show that a truly subcubic time algorithm for our generalized $\mathsf{ATAttC}$ (Definition 1.2) problem with large entries would yield one for ${\sf Gap}{\rm-}{\sf MaxIP}$ as well.

Previous work on (non-tensor) attention used as its intermediate problem the approximate Hamming Nearest Neighbor problem. However, it is not obvious how to directly generalize this to the tensor setting, since there is no way to define a ‘distance’ function for triples of vectors which satisfies the needed properties to generalize the original proof. We instead investigate the ${\sf Gap}{\rm-}{\sf MaxIP}$ problem, which can itself be seen as a generalization of an intermediate step in the proof of hardness for approximate Hamming Nearest Neighbor .

Hardness of 𝖦𝖺𝗉−𝖬𝖺𝗑𝖨𝖯{\sf Gap}{\rm-}{\sf MaxIP}

Fine-grained complexity results for approximation problems like ${\sf Gap}{\rm-}{\sf MaxIP}$ have previously been shown using a distributed probabilistically checkable proof framework , which we also use here.

We begin by generalizing the approach of using Merlin-Arthur (MA) communication protocols (). We construct a four party communication protocol for the disjointness problem: Alice, Bob and Charlie 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 , our protocol is algebraic in nature, and critically makes use of algebraic geometry codes from coding theory .

We then use this protocol to reduce from $\mathsf{SAT}$ to ${\sf Gap}{\rm-}{\sf MaxIP}$ . A standard reduction shows that $\mathsf{SAT}$ reduces to the $\mathsf{3OV}$ 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.

Reducing from 𝖦𝖺𝗉−𝖬𝖺𝗑𝖨𝖯{\sf Gap}{\rm-}{\sf MaxIP} to 𝖠𝖳𝖠𝗍𝗍𝖢\mathsf{ATAttC}

Finally, we reduce the ${\sf Gap}{\rm-}{\sf MaxIP}$ (Definition 4.6) problem to our $\mathsf{ATAttC}$ (Definition 1.2) problem. The key idea is that, by defining the matrices $Q,K_{1},K_{2},V_{1},V_{2}$ of generalized attention in terms of the inputs to ${\sf Gap}{\rm-}{\sf MaxIP}$ , we can make large entries of the attention matrix $A$ correspond to the triples with largest inner product. Some manipulation similar to prior work allows us to detect large entries from the output of $\mathsf{ATAttC}$ . This approach has been used for the fine-grained hardness of many attention and kernel density estimation problems . See details in Lemma B.1 and its proofs.

Hardness

In this section, we begin the formal proof of our hardness result. We begin by introducing the fine-grained hypotheses we will use.

For every $\epsilon>0$ there exists an integer $k\geq 3$ such that $\mathsf{CNF}-\mathsf{SAT}$ on formulas with clauses size at most $k$ (the so called $k$ - $\mathsf{SAT}$ problem) and $n$ variables cannot be solved in $O(2^{(1-\epsilon)n})$ time even by a randomized algorithm.

Given three sets $A,B,C\subset\{0,1\}^{d}$ where $|A|=|B|=|C|=n$ , the goal is to find a tuple $(i_{1},i_{2},i_{3})\in[n]\times[n]\times[n]$ such that $\langle a_{i_{1}},b_{i_{2}},c_{i_{3}}\rangle=0$ .

For every $\epsilon>0$ , there is a $c\geq 1$ such that $\mathsf{3OV}$ cannot be solved in $n^{3-\epsilon}$ time on instances with $d=c\log n$ .

It is known that $\SETH$ implies $\mathsf{3OVC}$ ; see, e.g., .

We state a important tool from the field of algebraic geometry codes. For more background on algebraic geometry codes, we refer the reader to .

3-way Polynomial Closure. ${\cal C}$ and ${\cal C}^{\prime}$ are linear codes. For each $w_{1},w_{2},w_{3}\in{\cal C}$ , there exists $w^{\prime}\in{\cal C}^{\prime}$ such that for each $i\in{\cal R}_{n}$ , $w^{\prime}(i)=w_{1}(i)\cdot w_{2}(i)\cdot w_{3}(i)$

Efficiency. Both codes can 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 at least $0.01$ .

2 A Four Party MA Communication Protocol

Prior work () constructed a protocol for three party communication, which includes Merlin, Alice and Bob. Here we modify this protocol for four parties.

For any $T\in[2,m]$ . There is a $\mathsf{MA}$ -communication protocol for Set Disjointness over universe $[m]$ . This protocol is computationally efficient.

In particular, the details of protocol are

Merlin sends Alice $O(\frac{m\log T}{T})$ bits

Alice, Bob, Charlie toss $O(\log m)$ coins

If the three sets do not have any element in common, Alice always accepts. Otherwise, she accepts with probability at most $1/2$ .

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^{1}\cup\cdots\cup U^{T}.$

Let $\alpha,\beta,\gamma\subseteq[m]$ denote the inputs of Alice, Bob, and Charlie. Our goal is to determine whether there is an element in the intersection $\alpha\cap\beta\cap\gamma$ .

For each $t\in[T]$ , we define the $t$ -th parts of the three sets:

Let $\rho_{C},\delta_{C}$ be the rate and distance of the code; recall these are at least a positive constant. Let $n_{C}=\frac{m}{T\cdot\rho_{C}}=O(m/T)$ be the length of the codewords of $C$ .

For each $t\in[T]$ , we write $C(\alpha^{t})$ , $C(\beta^{t})$ , $C(\gamma^{t})$ to denote the encodings of $\alpha^{t}$ , $\beta^{t}$ and $\gamma^{t}$ . Thus, their entry-wise product $\mu^{t}$ ( i.e., $\mu_{i}^{t}:=C(\alpha^{t})_{i}\cdot C(\beta^{t})_{i}\cdot C(\gamma^{t})_{i}$ ) is a codeword in the second code $C^{\prime}$ . Furthermore, since $C^{\prime}$ is a linear code, the entry-wise sum of the $\mu^{t}$ ’s ( $\mu_{i}=\sum_{t=1}^{T}\mu_{i}^{t}$ ) is also a codeword of $C$ ’.

$C$ is a systematic code, so we may assume that for each $i\in[n/T]$ , the entries $C(\alpha^{t})_{i},C(\beta^{t})_{i},C(\gamma^{t})_{i}$ are from $\{0,1\}$ and represent membership in the set. Similarly, $\mu_{i}^{t}\in\{0,1\}$ , and the sets are disjoint if and only if $\mu_{i}^{t}=0$ for all $i\in[m/T]$ and $t\in[T]$ , or equivalently, $\mu_{i}=0$ for all $i\in[m/T]$ .

Step 1. Merlin sends Alice $\widehat{\mu}$ , which is supposed to be the encoding of $\mu$

Step 2. Charlie, Bob and Alice pick a random $i^{*}\in[n_{C}]$

Step 3. Charlie sends Alice $C(\gamma^{t})_{i^{*}}$ for all $t\in[T]$

Step 4. Bob sends Alice $C(\beta^{t})_{i^{*}}$ for all $t\in[T]$

Step 5. Alice accepts iff all of the following hold:

$\widehat{\mu}$ is a codeword in $C^{\prime}$

$\widehat{\mu}_{i^{*}}=\sum_{t=1}^{T}C(\alpha^{t})_{i^{*}}\cdot C(\beta^{t})_{i^{*}}\cdot C(\gamma^{t})_{i^{*}}$

$\widehat{\mu}_{i}=0$ for all $i\in[m/T]$

First, we observe that Merlin’s message length is $n_{c}\cdot\log T=O((\log T)\cdot m/T)$ , and both Bob and Charlie’s message lengths are $T\cdot O(\log T)$ , as desired. To see correctness, note that if Alice ever accepts given Merlin’s message $\widehat{\mu}$ , then $\widehat{\mu}$ must in particular be a codeword of $C^{\prime}$ . If Alice accepts with probability greater than $1-\delta_{C^{\prime}}$ (where $\delta_{C^{\prime}}$ is a positive constant) then $\widehat{\mu}$ is also equal to the true $\mu$ by definition of $\delta_{C^{\prime}}$ . This means $\mu_{i}=0,\forall i\in[m/T],$ so the sets are disjoint.

3 Showing 3-𝖬𝖠𝖷\mathsf{MAX}-𝖨𝖯\mathsf{IP} is hard

We now define the appropriate gap $3$ - $\mathsf{MAX}$ - $\mathsf{IP}$ problem, which we use as our intermediate hard problem.

We use $t>0$ to represent a threshold parameter.

We use $\epsilon$ to represent an accuracy parameter.

Suppose $n,d$ denote two positive integers.

$A=\{a_{1},\cdots,a_{n}\}\subset\{0,1\}^{d}$

$B=\{b_{1},\cdots,b_{n}\}\subset\{0,1\}^{d}$

$C=\{c_{1},\cdots,c_{n}\}\subset\{0,1\}^{d}$

For every index $i\in[n]$ , we need to distinguish the following two cases

Case 1. There exists a pair $(j_{1},j_{2})\in[n]\times[n]$ such that $\langle a_{i},b_{j_{1}},c_{j_{2}}\rangle\geq t$ .

Case 2. For all pairs $(j_{1},j_{2})\in[n]\times[n]$ we have $\langle a_{i},b_{j_{1}},c_{j_{2}}\rangle\leq(1-\epsilon)\cdot t$ .

Implicit in previous work () is a proof that the analogue of ${\sf Gap}{\rm-}{\sf MaxIP}$ with two sets of points is hard. Here we generalize this to three sets.

Unless $\mathsf{SETH}$ and $\mathsf{OVC}$ are false, the following holds: for every $\delta>0$ there are constants $\alpha_{1}>\alpha_{2}>0$ such that for integer $n$ , solving ${\sf Gap}{\rm-}{\sf MaxIP}(n,d=\alpha_{1}\log n,t=\alpha_{2}\log n,\epsilon=1/2)$ requires time $\Omega(n^{3-\delta})$ .

We reduce from $3\OV$ to ${\sf Gap}{\rm-}{\sf MaxIP}$ . Let $\delta_{\OV}=\delta/2$ . Our reduction takes as input an instance $(A_{\OV},B_{\OV},C_{\OV})$ of orthogonal vectors over $\{0,1\}^{m}$ . These sets have sizes $|A_{\OV}|=|B_{\OV}|=|C_{\OV}|=2^{m/c}$ for a constant $c$ depending on $\delta_{\OV}$ from Definition 4.2 and Conjecture 4.3, and $\mathsf{3OVC}$ posits there is no algorithm solving this problem in time $O((2^{m/c})^{3-\delta_{\OV}})$ .

For a constant $k>0$ to be determined, pick $\epsilon>0$ to be a constant such that

We use the protocol of Theorem 4.5, instantiated with parameter

Suppose that $T^{\prime}=2^{O((\log T)\cdot T)}$ is representing the number of different possible messages sent by Bob and Charlie in the protocol.

Let us choose $T$ so that $T^{\prime}=O(1/\epsilon)$ .

For each vector $\gamma\in C_{\OV}$ , we construct a new vector $\widetilde{c}^{\gamma}\in\{0,1\}^{(T^{\prime})^{2}\times m}$ by setting $\widetilde{c}^{\gamma}_{i_{B},i_{C},j}:=1$ iff Charlie send message $i_{C}\in[T^{\prime}]$ on input $\gamma^{\prime}$ and randomness $j\in[m]$ . (The value is independent of $i_{B}$ .)

For each vector $\beta\in B_{\OV}$ , we construct a new vector $\widetilde{b}^{\beta}\in\{0,1\}^{(T^{\prime})^{2}\times m}$ by setting $\widetilde{b}^{\beta}_{i_{B},i_{C},j}:=1$ iff Bob sends message $i_{B}\in[T^{\prime}]$ on input $\beta^{\prime}$ and randomness $j\in[m]$ . (The value is independent of $i_{C}$ .)

For each Merlin-message $\mu\in\{0,1\}^{O((\log T)\cdot m/T)}$ and vector $\alpha\in A_{\OV}$ , we construct a new vector $\widetilde{a}^{\mu,\alpha}\in\{0,1\}^{(T^{\prime})^{2}\times m}$ as follows: $\widetilde{a}^{\mu,\alpha}_{i_{B},i_{C},j}:=1$ iff Alice accepts on

message $i_{B}$ from Bob, message $i_{C}$ from Charlie, and randomness $j$ .

Notice also that the inner product of three vectors

is exactly proportional to the probability that Alice, Bob and Charlie accept on inputs $\alpha,\beta$ , $\gamma$ and message $\mu$ from Merlin.

In particular, if $\alpha$ , $\beta$ and $\gamma$ are orthogonal (i.e., $\langle\alpha,\beta,\gamma\rangle=0$ ), the inner product is at most

Otherwise, there exists a $\mu$ such that

In particular, these can be distinguished by an algorithm for

which must therefore be as hard as solving the original instance of $3\mathsf{OV}$ . By $\mathsf{3OVC}$ , this means it requires time

where the last step follows from choosing $k$ large enough in the definition of $\epsilon$ .

At the end, we notice that the vectors we construct have dimension $2(T^{\prime})^{2}\cdot m=O(m)=O(\log n)$ as desired. ∎

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.

Appendix

In Section A, we provide the definitions of several notations. In Section C, we 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 our algorithm. In Section B. we show how to reduce our problem to ${\sf Gap}{\rm-}{\sf MaxIP}$ .

Appendix A Preliminary

For any positive integer $n$ , we write $[n]$ to denote $\{1,2,\cdots,n\}$ .

We use ${\bf 1}_{n}$ to denote a length- $n$ vector whose entries are all ones.

Given three sets $A,B,C\subseteq\{0,1\}^{d}$ of vectors where $|A|=|B|=|C|=n$ , the goal is to compute

Appendix B Hardness: From 𝖬𝖺𝗑𝖨𝖯\mathsf{MaxIP} to Our Problem

In Section B.1, we show how to reduce our problem to ${\sf Gap}{\rm-}{\sf MaxIP}$ . In Section B.2, we present our main lower bound (hardness) result.

We now generalize the hardness proof of to the tensor attention case.

For every constant $C_{\gamma}\in(0,0.1)$ , every $\epsilon>0$ , and every $C>C_{0}>0$ , there exist constants $C_{a}>0$ and $C_{b}>0$ and such that, if $\mathsf{ATAttC}$ (Definition D.1) for parameters $(2n,d=2C\log n,B=C_{b}\sqrt{\log n},\epsilon_{a}=n^{-C_{a}})$ can be solved in time $T$ , then ${\sf Gap}{\rm-}{\sf MaxIP}(n,d=C\log n,t=C_{0}\log n,\epsilon)$ (Definition 4.6) can be solved in time $O(T+n^{3-C_{\gamma}})$ .

We give an algorithm for ${\sf Gap}{\rm-}{\sf MaxIP}(n,d=C\log n,t=C_{0}\log n,\epsilon)$ (Definition 4.6). Let $a_{1},\cdots,a_{n},b_{1},\cdots,b_{n},c_{1},\cdots,c_{n}\in\{0,1\}^{d}$ denote the inputs to this problem. Using them, we will construct appropriate inputs to the $\mathsf{ATAttC}$ problem so that its output will help us to detect triples with large inner product.

Let $\beta>0$ and $\widetilde{d}\geq d$ be parameters to be determined (in Eq. (5) and Eq. (2) below). Define $\tau>0$ by

We pick these parameters so that $\tau$ will be an upper bound on entries of the attention matrix, namely

We will use an algorithm for the $\mathsf{ATAttC}(\widetilde{n},\widetilde{d},B,\epsilon_{a})$ problem with parameters:

Since each entry of $Q$ and $K_{1}$ , $K_{2}$ is either $\sqrt{\beta}$ or $0$ , it follows that

is naturally partitioned into eight submatrices

where each $A_{i}$ ( $\forall i\in$ ) is a matrix of size $n\times n^{2}$ , defined as follows. For each $j_{0}\in[n],j_{1}\in[n],j_{2}\in[n]$ ,

The $(j_{0},j_{1}+(j_{2}-1)n)$ -th entry of $A_{1}$ is

$\exp(\beta(\langle a_{j_{0}},b_{j_{1}},c_{j_{2}}\rangle+\langle{\bf 1}_{d},{\bf 0}_{d},{\bf 0}_{d}\rangle)/\widetilde{d})=\exp(\beta\langle a_{j_{0}},b_{j_{1}},c_{j_{2}}\rangle/\widetilde{d})$

The $(j_{0},j_{1}+(j_{2}-1)n)$ -th entry of $A_{2}$ is

$\exp(\beta(\langle a_{j_{0}},b_{j_{1}},{\bf 0}_{d}\rangle+\langle a_{j_{0}},{\bf 0}_{d},{\bf 1}_{d}\rangle)/\widetilde{d})=\exp(0)=1$

The $(j_{0},j_{1}+(j_{2}-1)n)$ -th entry of $A_{3}$ is

$\exp(\beta(\langle a_{j_{0}},{\bf 0}_{d},c_{j_{2}}\rangle+\langle a_{j_{0}},{\bf 1}_{d},{\bf 0}_{d}\rangle)/\widetilde{d})=\exp(0)=1$

The $(j_{0},j_{1}+(j_{2}-1)n)$ -th entry of $A_{4}$ is

$\exp(\beta(\langle a_{j_{0}},{\bf 0}_{d},{\bf 0}_{d}\rangle+\langle{\bf 1}_{d},{\bf 1}_{d},{\bf 1}_{d}\rangle)/\widetilde{d})=\exp(\beta d/\widetilde{d})=\tau$

The $(j_{0},j_{1}+(j_{2}-1)n)$ -th entry of $A_{5}$ is $1$

$\exp(\beta(\langle{\bf 0}_{d},b_{j_{1}},c_{j_{2}}\rangle+\langle{\bf 1}_{d},{\bf 0}_{d},{\bf 0}_{d}\rangle)/\widetilde{d})=\exp(0)=1$

The $(j_{0},j_{1}+(j_{2}-1)n)$ -th entry of $A_{6}$ is $1$

$\exp(\beta(\langle{\bf 0}_{d},b_{j_{1}},{\bf 0}_{d}\rangle+\langle a_{j_{0}},{\bf 0}_{d},{\bf 1}_{d}\rangle)/\widetilde{d})=\exp(0)=1$

The $(j_{0},j_{1}+(j_{2}-1)n)$ -th entry of $A_{7}$ is $1$

$\exp(\beta(\langle{\bf 0}_{d},{\bf 0}_{d},c_{j_{2}}\rangle+\langle a_{j_{0}},{\bf 1}_{d},{\bf 0}_{d}\rangle)/\widetilde{d})=\exp(0)=1$

The $(j_{0},j_{1}+(j_{2}-1)n)$ -th entry of $A_{8}$ is

$\exp(\beta(\langle{\bf 0}_{d},{\bf 0}_{d},{\bf 0}_{d}\rangle+\langle{\bf 1}_{d},{\bf 1}_{d},{\bf 1}_{d}\rangle)/\widetilde{d})=\exp(\beta d/\widetilde{d})=\tau$

For each $(i,j_{1},j_{2})\in[n]\times[n]\times[n]$ , we know that

Here we used the fact that $d<\widetilde{d}$ (see Eq. (2)), and the last step uses the definition of $\tau$ (see Eq. (1)).

We also know that for each $(i,j_{1},j_{2})\in[n]\times[n]\times[n]$ ,

since it is the exponential of an entry of $QK^{\top}/\widetilde{d}$ .

By combining our expression for $A$ with Eq. (B.1) and Eq. (7), we see that

Since $D_{i,i}=(A{\bf 1}_{\widetilde{n}})_{i}$ , it follows that

We can see that $\widetilde{t}\geq\epsilon_{a}$ as follows:

Here, the last two steps follow from Eq. (4).

Recall that for each $i\in[n]$ we need to distinguish between two cases: either there is a pair $(j_{1},j_{2})\in[n]$ such that $\langle a_{i},b_{j_{1}},c_{j_{2}}\rangle\geq t$ , or else for all pairs $(j_{1},j_{2})\in[n]$ the inner product $\langle a_{i},b_{j_{1}},c_{j_{2}}\rangle\leq(1-\epsilon_{a})t$ . We will distinguish between these cases by checking whether $u_{i}$ is greater than a threshold value $\widetilde{t}_{0}:=2\widetilde{t}$ . We next consider the two cases to see why this is.

For a given $i\in[n]$ , if there are $(j_{1},j_{2})\in[n]\times[n]$ such that $\langle a_{i},b_{j_{1}},b_{j_{2}}\rangle\geq t$ , then

where the 1st step follows from $2d=\widetilde{d}$ (see Eq. (2)). This means as desired that

For a given $i\in[n]$ , if for all $(j_{1},j_{2})\in[n]\times[n]$ we have $\langle a_{i},b_{j_{1}},c_{j_{2}}\rangle<t(1-\epsilon)$ , then

Here, the 4th step follows because, by our choice of $\beta$ and $t$ , we have

where we used that $t=C_{0}\log n$ (by Lemma statement), that $d=C\log n$ , that $\beta=B^{3}$ (Eq. (5)) and the choice of $B$ (Eq. (3)). ∎

B.2 Main Hardness Result

We can finally conclude our main lower bound.

Assuming $\mathsf{SETH}$ , for every $q>0$ , there are constants $C,C_{a},C_{b}>0$ such that: there is no algorithm running in time $O(n^{3-q})$ for the problem $\mathsf{AAttC}(n,d=C\log n,B=C_{b}\sqrt{\log n},\epsilon_{a}=n^{-C_{a}})$ .

Follows from combining Theorem 4.5, Theorem 4.7, and Lemma B.1. ∎

Appendix C Upper Bound: Running Time

In Section C.1, we 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.

We first review the attention computation definition in ,

C.2 Tensor Attention Computation

Given two $n\times d$ matrices, there are two standard variants on their Kronecker product one may consider: The standard Kronecker product (denoted $\otimes$ ) is a new $n^{2}\times d^{2}$ matrix, whereas the column-wise Kronecker product (denoted $\oslash$ ) is a new $n^{2}\times d$ matrix. For more literature on tensor computations and their applications in learning algorithms, we refer the readers to .

Next, we generalize the matrix attention computation (in ) into tensor attention computation as follows:

C.3 Efficient Column-wise Kronecker Computation

We prove an important tool which will be used in analyze the running time of our algorithm.

Let $\oslash$ be defined as Definition 2.4.

We define $C_{1}:=A_{1}^{\top}B_{1},C_{2}:=A_{2}^{\top}B_{2}$

For each $i\in[n]$ , let $b_{1,i}^{\top}$ denote the $i$ -th row of $B_{1}$ .

For each $i\in[n]$ , let $b_{2,i}^{\top}$ denote the $i$ -th row of $B_{2}$ .

From the above, we can calculate that the entry of $C$ in location $k_{1},k_{2}$ is

where the first step follows from Eq. (C.3), 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 $C_{1}$ and $C_{2}$ .

C.4 Simple Equivalent Tools for Tensor Notations

We define a standard tensor notation, for example see .

Next, we present several equivalence results for tensors.

Let $\oslash$ be defined as Definition 2.4.

Let $\odot$ be defined as Definition 2.2.

Let $(\cdot,\cdot,\cdot)$ operator be defined as Definition C.4.

Let $\circ$ be defined as Definition 2.1.

Part 1. $A_{i,j_{1}+(j_{2}-1)n}=\mathsf{A}_{i,j_{1},j_{2}}$ for $i\in[n],j_{1}\in[n],j_{2}\in[n]$ (This means $\mathsf{A}$ can be viewed as the tensor version of $A$ )

Part 2. $A{\bf 1}_{n^{2}}=\mathsf{A}(I,{\bf 1}_{n},{\bf 1}_{n})=Q\odot({\bf 1}_{n}^{\top}K_{1})\odot({\bf 1}_{n}^{\top}K_{2})$

Part 3. $A(V_{1}\oslash V_{2})=Q(K_{1}\oslash K_{2})^{\top}(V_{1}\oslash V_{2})=Q((K_{1}^{\top}V_{1})\circ(K_{2}^{\top}V_{2}))$

Directly follows from definition of $A$ and $\mathsf{A}$ .

Follows from tensor notations in Definition 2.2 and Definition C.4.

Directly follows from applying Part 1 of Lemma C.3 here. ∎

Appendix 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 a 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, we analyze the error propagation from $A$ and $D$ to the attention matrix.

$\|Q\|_{\infty}\leq B$ , $\|K_{1}\|_{\infty}\leq B$ , $\|K_{2}\|_{\infty}\leq B$ , $\|V_{1}\|_{\infty}\leq B$ , $\|V_{2}\|_{\infty}\leq B$

Notice that the straightforward algorithm for this problem will spend at least $\Omega(n^{3})$ time to write the matrix $A$ (we can also think of $A$ as an tensor that has size $n\times n\times n$ ).

D.2 An Error Control Tool From Previous Work

Let $g:=\Theta(\max\{\frac{\log(1/\epsilon)}{\log(\log(1/\epsilon)/B)},B\})$ .

D.3 Tensor Q⊙K1⊙K2Q\odot K_{1}\odot K_{2} Has Bounded Entries

Let $\odot$ operation be defined as Definition 2.2.

For every index triple $(i,j_{1},j_{2})\in[n]\times[n]\times[n]$ , we are able to prove

D.4 Tensor Low-Rank Approximation

In the following definition, we view the $n\times n^{2}$ size matrix as an $n\times n\times n$ size attention matrix.

We use $r\geq 1$ to denote a positive integer.

We use $\epsilon\in(0,0.1)$ to represent an accuracy parameter.

$|\widetilde{\mathsf{A}}_{i,j_{1},j_{2}}-A_{i,j_{1},j_{2}}|\leq\epsilon\cdot\mathsf{A}_{i,j_{1},j_{2}}$ for all $(i,j_{1},j_{2})\in[n]\times[n]\times[n]$ .

D.5 From AA to DD

In this section and the next, we generalize the proof of for error propagation from the matrix setting to the tensor setting. The proofs are nearly identical.

Then, for every index $i\in[n]$ , the following bound holds

where the second step follows from triangle inequality.

D.6 From A,DA,D to Tensor Attention

The goal of this section is to prove Lemma D.6.

Suppose the following conditions are true

Let $\epsilon_{A},\epsilon_{D}\in(0,0.1)$

We bound each of these two terms to get our desired result.

First of all, for every index pair $(i,j)\in[n]\times[d]$ ,

Second, for every $(i,j)\in[n]\times[d]$ ,

Here, again, the 2nd step uses the triangle inequality, the 3rd step follows because $D_{i,i}^{-1}$ is positive, the 4th step follows from the assumption that $|\widetilde{A}_{i,l}-A_{i,l}|\leq\epsilon_{A}\cdot A_{i,l}$ and the final step follows by definition of $D_{i,i}$ .

The Lemma conclusion then becomes true by substituting Eq. (D.6) and Eq. (D.6) into Eq. (11). ∎

Appendix E Upper Bound: Putting It All Together

In Section E.1, 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 our main upper bound theorem.

The goal of this section is to prove Lemma E.1.

Let $P(x)$ denote a degree- $g$ single-variable polynomial. We apply $P(M)$ entry-wisely, i.e, $P(M)_{i,j,l}=P(M_{i,j,l})$ .

Let $r$ be rank parameter that is $r:={3(g+d)\choose 3g}.$

We define set $V$ and provide names for variables in set $V$ in the following sense,

Thus, function $\mathsf{K}$ can be viewed a degree- $3g$ polynomial in the $3d$ entries in $V$ of the vectors $u,v,w$ .

Here, we can view $\mathsf{K}$ function as

Therefore, we should construct three matrices $U_{1},U_{2}$ and $U_{3}$ as the following way, for each $i\in[n]$

These $n\times r$ matrices can be constructed in time $O(nrg)$ in the straightforward way, since each entry depends on $g$ variables. ∎

E.2 Time for Constructing U1,U2,U3U_{1},U_{2},U_{3}

$\|K_{1}\|_{\infty}\leq B$ , $\|K_{2}\|_{\infty}\leq B$ ,

$\|V_{1}\|_{\infty}\leq B$ , and $\|V_{2}\|_{\infty}\leq B$ .

For bounded number $B>0$ and accuracy parameter $\epsilon\in(0,1)$ , there are positive integers $g$ and $r$

it takes $O(nr)$ time to construct $U_{1},U_{2}$ and $U_{3}$ .

Recall that the definition of $(\epsilon,r)$ -approximation can be found in Definition D.4.

We can then compute $U_{1},U_{2}$ , and $U_{3}$ using Lemma E.1, which gives the bound

E.3 Main Algorithmic Result

We present our main algorithmic result as follows: