Fast attention mechanisms: a tale of parallelism
Authors: Jingwen Liu
Hantao Yu
Clayton Sanford
Alexandr Andoni
Daniel Hsu
Venue: NeurIPS 2025 poster ()
Keywords: Transformer theory
representational strength
nearest neighbor search
massively parallel computation
Primary area: theory

Abstract: Transformers have the representational capacity to simulate Massively Parallel Computation (MPC) algorithms, but they suffer from quadratic time complexity, which severely limits their scalability. We introduce an efficient attention mechanism called Approximate Nearest Neighbor Attention (ANNA) with sub-quadratic time complexity. We prove that ANNA-transformers (1) retain the expressive power previously established for standard attention in terms of matching the capabilities of MPC algorithms, and (2) can solve key reasoning tasks such as Match2 and $k$-hop with near-optimal depth. Using the MPC framework, we further prove that constant-depth ANNA-transformers can simulate constant-depth low-rank transformers, thereby providing a unified way to reason about a broad class of efficient attention approximations.

Fast attention mechanisms: a tale of parallelism
Jingwen Liu Columbia University New York, NY jingwenliucss www.woowl.m.....ooo
Hantao Yu
Clayton Sanford Google Research San Francisco, CA chsanf woonoorrdgggge.coo
Columbia University New York, NY
muntao.yu@columiaa. edu
Alexandr Andoni Columbia University New York, NY Celooi@cs.colummete edu
Daniel Hsu
Columbia University New York, NY
djhsuêcs. wow.uucc.ccllmbia..o... edu
Abstract
Transformers have the representational capacity to simulate Massively Parallel Computation (MPC) algorithms, but they suffer from quadratic time complexity, which severely limits their scalability. We introduce efficient attention mecha- nism called Approximate Nearest Neighbor Attention (ANNA) with sub-quadratic time complexity. We prove that ANNA-transformer (I) retain the expressive power previously established for standard attention in terms of matching the capabilities of MPC algorithms, and (2) can solve key reasoning tasks such as Match2 and k-hop with near- near-optimal depth. Using the MPC framework, we further prove that constant-depth
can simulate constant-depth low-rank trans- formers. thereby providing a unified way to reason about broad class of efficient attention approximations.
Introduction
The transformer [57] has become the dominant neural architecture deep learning due its ability
to select and compose complex information structures from large inputs [9 24]. which in turn enables capabilities such as "in-context learning" At the core of transformers is the attention mechanism, which leverages parallelism to gain important advantages over non-paran non-parallel architectures (e.g., recurrent nets), including training stability [46)
that are crucial for downstream applications.
55 60). One recently highlighted advantage comes the Massively Parallel Computation (MPC) model [33] that captures the power of large-scale distributed computing frameworks like MapReduce [17 efficient MPC algorithms can be simulated by small-size transformers, and vice versa [53 55). This correspondence suggests that broad class of computational tasks can be efficiently solved by transformers.
and representational power [1] 31 from a coarse theoretical relationship between transformers
41
45
52
Despite these advantages, transformers suffer from a key limitation: the quadratic time complexity of (standard) attention with respect input size, which is likely be unavoidable in the worst- To address this limitation, a growing body of work proposes alternatives to the attention mechanism that are more computationally efficient (e.g., sub-quadratic time). These
case 2 3 34
alternatives employ variety of techniques, ranging from low-rank approximation [8] [4
[6 32 58
to efficient nearest neighbor search [27 35 ST 64 Although many of these techniques show
promising empirical performance, it is unclear whether they preserve the representational advantages of attention.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

In this work, we study a particular efficient (sub-quadratic time) attention mechanism, called Approx- imate Nearest Neighbor Attention (ANNA). and we prove that ANNA retains key representational advantages of standard attention over non-parallel architectures. We additionally prove that ANNA- transformers can simulate class of attention mechanisms based on low-rank approximation. In doing so, we provide a unified way to reason about low-rank and nearest neighbor approaches to efficient attention.
1.1 Standard transformers and MPC
Prior work 53 55 established a coarse relationship between standard transformers and MPC by proving the following (for any constant (0, 1)):
For any MPC algorithm with inputs of size N that uses R rounds (of computation and commu- nication), O(N) machines, and O(Ne) words of local memory per machine (for some constant 5 E (0, 1)), there is an equivalent transformer of width O(N++t) and depth O(H). (By width, we mean number of heads per layer times the embedding dimension.)
2. For any transformer operating on inputs of size N that has L layers and width O(Ns) (for some constant (0, 1)), there an equivalent MPC algorithm that uses (L)) rounds, O(N2) machines, and O(N++6) words of local memory per machine.
Notice that the "loss" in the number of rounds depth is only constant factor, and the "loss" in the local memory size or width is only a ((N²) factor (where 5 > O) can be arbitrarily small). Consequently, these results were sufficient to give new results about transformer representational power (e.g., for various graph reasoning tasks (53 (55),, and also give a (conditional) logarithmic- depth lower bound for transformers that solve a multi-hop reasoning problem called k-hop (55)
However, there a important gap in the MPC algorithm for simulating a transformer (Item 2 above): it may up to 22 machines. This gap suggests the possibility that transformers are strictly more powerful than MPC algorithms, so the relationship established in prior work is very coarse. Moreover, this gap is likely to be inevitable. Indeed, when the local computation in an MPC algorithm is fast (say, polynomial-time) and the local memory size is N for small e E (0, 1), then a single round of MPC on M machines can be simulated on a sequential machine in roughly ( time. However, evaluation of attention over N inputs is believed to require least 2 time (for every constant
> 0) on a sequential machine 2 3 34]. So any such MPC algorithm that simulates attention in (1) rounds should use N2-6-0(e) machines.
This gap naturally leads us to the following question: Is there an alternative attention mechanism that more tightly captures the computational power of MPC?
1.2 Our contributions
We prove that ANNA-transformen more sharply capture the power of MPC algorithms than standard transformers do, and in particular avoid the aforementioned gap in prior works' characterization of transformers using MPC [53 55
The core ANNA mechanism is designed to perform c-approximate
nearest neighbor search 30 38 for a given set of N "queries" against a database of N "key-value" pairs. Here >> 1 is a parameter of the ANNA mechanism that, for the purpose of this present informal description, should be thought as being at least a large positive constant (e.g., c 2 10). Throughout, N denotes the input size.
Theorem 1.1 (Informal version of Theorems 4.1 and 4.2). Fix any constants E, ô E (0, and ANNA parameter e >> 1. For any R-round MPC algorithm with N words of local memory, there is an MWNA-transformer of width O(Ne+)). For any L-layer tran-traasformer with width N there is an equivalent (T.))-round MPC algorithm that uses No+6 words of local memory and N1-5+011/2) machines.
equivalent O(R)-layer
Observe that Theorem 1.1
except for the resources needed by the MPC algorithm in Item [ the local memory stays the same, but the number of machines required can be strongly sub-quadratic (and in fact, close to linear) in N. This gives positive answer to the question at the end of Section
essentially mirrors Items and 2 in Section I.]

---

We also study efficient attention mechanisms based on low-rank approximation, and by leveraging Theorem I.I. we prove that transforrrss have at least the same representational power.
Theorem 1.2 (Informal version of Theorem 4. 4.4. there is an equivalent OLy)-aayer
For any L-layer low-rank attention Transformer, (of comparable width).
Finally, we illustrate the power of ANNAA-ransformers on two concrete reasoning tasks (Match2 52 and k-hop induction heads [12 55 We give theoretical constructions of ANNA transformers for these tasks that nearly match the efficiency achievable by standard transformers (Theorems 5.2 and [.6). and we also show empirically that trnA--rrnsformers can be trained to approximately solve these tasks (Section 5.3
1.3 Other related works
The representational power of standard attention is relatively well understood from variety of perspectives. This includes: universal approximation properties in the large depth/width limit [22 44] (59 631, ability to recognize formal languages [10] [5] 41 55 62), computational bounds in terms of circuit classes [28 41 43 and other parallel computation models (discussed below) [53] 55), and bounds for specific compositional tasks 12
54] In general, these and other prior works do not consider the representational power of transformers based on efficient (sub-quadratic time) alternatives to attention. The exceptions works that give lower bounds for low-rank attention
15
36
41
49
and sparse attention [55] 61], sequential architectures [1 31 45 52
60], more generally, any What is missing in these prior works, however, is a characterization (and, in particular, upper bounds) for an efficient architecture.
mechanism that be evaluated in sub-quadratic time 2
Massively Parallel Computation (MPC) [5 9 intended to capture the power of MapReduce [17 and other large-scale distributed computing frameworks. Many works, including those that originally helped to define the MPC model, gave efficient algorithms for a variety of basic data and graph processing tasks. A connection between MPC and circuit models was given by (50) and used to formalize barrier on proving lower bounds for certain graph problems (e.g., connectivity) in MPC. Nevertheless, certain conjectured lower bounds in MPC are widely believed (e.g., the 1-vs-2 cycle conjecture
20
23
29
33
is model of parallel computing
(29]) and have been used to establish conditional lower bounds for other problems (21]. The same conjectured lower bounds have also been used to establish depth lower bounds for transformers via the aforementioned relationship between MPC and transformers
53 55
Our main result for ANNA-transfomme neural networks (GNNs). The computational power of such GNNs is characterized by the CONGEST model of distributed computing [48] (a refinement of LOCAL 7 40 47).. This fact was established by 42 and in turn used to give upper- and lower-bounds on the size of GNNs needed to solve various graph problems (e.g., subgraph detection). The GNN/CONGEST equivalence is almost immediate, since the communication network is fixed by the input graph in both models, and the model definitions semantically and syntactically similar. In contrast, while the layered representations of transformers are reminiscent of the alternation between communication and computation rounds in MPC, the "communication patterns" themselves are dynamic, and this dynamism greatly complicates the simulation of MPC algorithms by either standard transformers or trNA--tansformers.
has an analogue in the context of message- passing graph
The ANNA mechanism that we study is inspired by many prior efficient attention mechanisms [e.g.. | 64 based on locality-sensitive hashing (30, which one of the key techniques for nearest neighbor search. Some of these mechanisms have guarantees about the quality of approximation under structural assumptions on the attention matrix they are meant to approximate. The motivation of our analysis is largely orthogonal: we instead seek to characterize the representational power of ANNA-transformerr in terms of other well-understood models of parallel computation. We further compare the capabilities of LSH-hased efficient attention other sub-quadratic alternatives, including
those based on low-rank approximations of self-attention matrices 16 32 58
2 Preliminaries
2.1 Standard attention and transformers
We first define the (standard) attention mechanism and transformers.
3

---

Definition 2.1 (Attention). A (standard) attention head Attnq.KV K.V is specified by query, key, value embedding fumctions Q, K, Va R2 + IR. On input X E RN RNNxs Xxd it computes
K. (X) = softmax (Q(X)
Nxm.
where Q, K, V, and softmax are applied row-wise. We say N is the context length, and m is the embedding dimension. The rows Q (resp., ( V V(X)) ) are the queries (rexp., keys, values).
q: Q(X), k2 (X), and 0j = ( then the i-th row of Attnq.x.v is
expl k2))
Al E Wi where Wij Ey==1 exp (gp, ky)))
An H -headed attention layer f: RNxx RN Nxd consists of attention heads
and
x d matrices (Wi) it computes f(X) E Attnqu. Kn, K,,V,, (X)Wn.
Definition 2.2 (Transformer). An L-layer transformer TT is specified by attention layers fi....,fu R - R. xxd and output function &p R - R.. Given input X RNxx define
( = and x := fe(X(((1)) for forl- = I,....,,
The output T on input X is (x(())) with v/ being applied row-wise.
In this paper, we consider the context length N (the number of input tokens) as the principal scaling parameter. This reflects the modern paradigm of long-comm ong-context LLMs, where the context-length can exceed 106 [9), enabling book-length textual inputs. Consequently, we typically want the size parameters m, H. and L to be sub-linear in N (and L ideally constant). The O(N2) runtime of attention therefore remains the main bottleneck of the transformer architecture.
Following 155 we allow the element-wise operations (Qn, K₁, Vn, v)) to be arbitrary functions !T (limited only by bit-precision; [55] Appendix A.1J). Therefore, can view transformer computational model that alternates between arbitrary per-token computation, and communication between tokens. This motivates connection massively parallel computation, defined next.
2.2 Massively Parallel Computation
The Massively Parallel Computation (MPC) framework (29] models computation on large inputs by distributed computing system that alternates between rounds of local computation and rounds of restricted all-to-all communication. We formally state the definition of MPC protocols (with sublinear memory) follows.
Definition 2.3 (MPC). For constants 7,5 > O. an R-round (7.E) -MPC protocol specifies the following computation on inputs of N words (where word p = O(log N) bits, represented by an element of (xr) by P ((N1-1 machines, each with local memory s = O(N)) words:
Initially, the input is arbitrarily distributed across the first [N]] machines.
In each round, each machine prepares, as an arbitrary function of its local memory, messages to send other machines. The total size of messages prepared by any machine is at most words.
3. At the end of each round, the messages are placed in the local memory of the intended recipients. The protocol ensures that the messages received by any machine has total size at most s words.
After the R-th round, the output is stored in the memory of the first [N] machines.
We say an MPC protocol π computes a function f 22 + Eny, for any E EN ((X) f(X). where π (X) is the output of π with the input X.
The primary measure of complexity considered in this paper is the number of rounds The round complexity of numerous classical algorithmic problems is well understood. For example, there are simple MPC protocols for graph connectivity (Oog N) rounds) and sorting (0(1) rounds) (20]
Such assumptions are necessary for establishing the equivalence with the MPC model, since the MPC model allows arbitrary computation on the local memory for each machine. That said, many concrete MPC algorithms do have simple local algorithm which can be simulated by small MLP.

---

3 Approximate Nearest Neighbor Attention
In this section, we introduce Approximate Nearest Neighbor Attention (ANNA), an attention mech- anism inspired by the approximate nearest neighbor (ANN) search problem. We first outline the approximate nearest neighbor search problem and present locality-sensitive hashing (LSH), a core technique for ANN (Section 3.T). We then formally define ANNA and provide sub-quadratic time algorithm based on LSH for computing ANNA with theoretical guarantees (Section 3.2
3.1 Approximate nearest neighhor and locality sensitive hashing
We first define the Approximate Nearest Neighbor (ANN) search problem.
Definition 3.1 (ANN search problem R0 (. Given dataset D of N points lying in metric space Y and parameters c, r > O, build a data structure that, given query q Y within distance at most r from D, returns any point in D that is within distance er from q-
In the modern machine learning setting, we want to develop fast algorithms for ANN search in the high-ding high-dimensional metric space with sub-linear query time. A well-known tool that achieves this runtime and provable approximation guarantees is locality sensitive hashing (LSH), For simplicity, we assume the metric space is n-ddiensional Euclidean space.
Definition 3.2 (Locality Sensitive Hashing (30). Fix a parameter >> O, an approximation factor > I and a set U. Then family H of hash functions h. Rmm - U is (r, cr,p), Pt.) if the following holds for any I. y E m
if || al S r, then PLiSи [h(x)= h(y)] 2 PL, and
if || vl cr, then PrneH(h((F) = h(w) S p2-
The family H. called LSH family with quality p
(1/1) 11p2)"
typical LSH-based algorithm can solve the ANN search problem with space O(N++)) and query time ONN)), and can be as small as c/c²
3.2 Transformer based Approximate Nearest Neighhor Attention
We first define family of models where only tokens with sufficiently "neighborly" queries and keys attend one another and then provide an efficient implementation of subset of this class. ANNA attention units treat attention query vectors queries to approximate nearest neighbor search and key vectors as data points. ANNA retrieves and weights value vectors according to approximate nearest neighbor thresholds. The following definition formalizes this family of models.
Definition 3.3 (ANN Attention). An Approximate Nearest Neighbor Attention (ANNA) mechanism ANNAQ, v with query, key, and value embedding functions Q,K,V Bd - RTM and (non- negative) parameters r, C, E, TJ, is (possibly randomized) mechanism that performs the following computation on an imput X E &R. Xdd
ANNAQ.Kv(X); = Σ WisUi for all i (N),
j=1
for some non-negative weights waj 2 0 with E3 waj that satisfy the following. With probability least for all i e [N].
Wi. >0= k2 E NNig,
k, EN((..)) = Wt.j 2 (Wi(g.er)-11)4)
where qi = Q(X), := K(X) oi V(X), and Nd(,N) = (k (ky) 1u-kli S 1).
We define ANNA layers and ANNA transformers in a completely analogous fashion as (standard) attention layers and transformers defined (Definitions) 2.7 and 2.2
5

---

Algorithm ANNA implementation with LSH family H. hash tables, and z hash functions/table
Input: Input X
Nxd
Output: ANNA output for each of the query.
I: Let q )t, K2 (X),,, and 24 ((X), for all i [N].
2: for u = 1 to & do
Preprocessing phase
3: Sample z hash functions h.. hu.2.... ..... .. ha,s i.i.d. from H.
4: Create empty hash table Tu indexed by hash codes (below).
5: for each key-value pair (Ry, 1j) do
6: Compute hash code 9u (kz) = (bv,, (k;g (l, (kz)( N.. (k)))
7:
if T. [9. (kj) empty, then Tu (ga (kj)) == (Dj, 1) else Tulgu (ki)) += (vz:
8: Initialize a dictionary attn t {(1 0), (92. O)..... 0)}
9:
for each query g: do
D. Query phase
10: Usum <- 0; count +-
11:
for u = to l do
12:
Compute hash code gu (q:) (hu,, (q:).... (q:))
13: if T. (9)())) = (v, a) is not empty, then Usum += v and count += a
14: attn [qs] < Usum/ /count
15:
return attn
The parameters r, have the same semantics as in ANN search. The parameter E captures how much "attention weight" is spread over keys that are not r-near neighbors of a query. The failure probability 7) allows for randomization, which is typical of ANN search algorithms like LSH.
The above definition represents set of constraints that all ANNA units must satisfy rather than specific algorithmic implementation. As a result, wide variety of models satisfy the definition, including softmax attention with bounded query and key vectors, and Exact-Match Attention (EMA) where W,jj > O if and only if kj. Not all such models admit computationally efficient implementations. To identify sub-quadratic ANNA models, we present an LSH- implementation of ANNA that computes satisfying weight vectors Wi, for specific choices of parameters r, r,c,e,
We define a hash function family G = : + (h1 (p). H. Vi [=]] and sample l hash functions, 9l, -t. from G independently and uniformly at random, giving é hash tables. Each hash code corresponds to a hash bucket in the hash table and each hash bucket maintains sum of e's and count of k's that falls into this bucket. We preprocess all the key, value pairs by storing them in the hash tables. For each (k;, uz). E [N], compute the hash codes of ki, gi ((()).((k)))))) .... (k(). and update the sum and count for the buckets corresponding to 91 ( (2(k)) ).................. ................................ ge k2) respectively, For each query qi, i E [N], search and retrieve all the values and counts from 91 (4)), 92 (q:)..... ge
E U h2 E
(q:). Then compute the averaged value by summing up all the values in the& e buckets, divided by the sum of counts. See Algorithm| [ for the details.
Theorem 3.4 (LSH algorithm guarantee for ANNA), Fix c > v3. LSH family H. that is (r.cr, cr, Py p
sensitive with quality p < 1/3, l = (NN3p log N), and @log1/pp N). Then Algorithm I (with H. E, and 2) implements an ANNA mechanism with parameters c, and O(1/N1
We leave the full proof of Theorem 3.4 to Appendix A The total runtime of Algorithm I is O(mN1-3p logi/ps N), assuming sampling from the LSH family and evaluating a hash function requires ((m) time and the numerical inputs Algorithm are specified with p = O(lo N) bits of precision. The total space used is O(m bits.
Remark 3.5. The memory complexity of Algorithm Dean be further improved to Ó(mN) bits with the same time complexity by storing only one hash table with each entry keeping track of the value for each query. See Algorithm for the detailed implementation in Appendix
Remark 3.6. The weight Wi.j depends the number of hash collisions between qi and kej, and is typically a function of the distance A == ||q: - k2ll- For example, if we use the random hyperplane LSH frof (6]. then Wa,j x exp( - 2 log(m) (mgmmm/(4 ())
In the remainder of this paper, our ANNA- constructions are meant to refer to their efficient implementation by Algorithm with suitable choice of r and arbitrarily large c.

---

Efficient transformers and MPC
We prove a sharp equivalence between memory and sub-quadratic number of machines (Sections) 4.1 and 4.2). We also show that ANNA subsumes alternative low-rank sub-quadratic attention mechanisms. (Section 4.3)
and MPC in the regime of sub-linear local
4.1 ANnA-transfrrmer can simulate MPC
The following theorem shows that any R-round MPC protocol with sub-linear local memory be simulated by an ANNA-transformer with O(f)) layers and sub-linear number of heads and embedding dimension. The full proof of Theorem 4.1 is in Appendix B
Theorem 4.1 (ANNA simulates MPC). Fix constants O < S'<E'' <. For any deterministic R-round (E, EEPE)-MPC protocol л, there exists an ANNA-transfommer T with = O(R) layers, H- = ((N(c')))/AA heads per layer, and embedding dimension m = O(Ne)), such that T(input) = (input) for all input Z.y.
Proof sketch. In fact, we show that the special case of ANNA-transformet whose approximation factor - and 0 is already suffice to simulate MPC. In such case, ANNA for each query is equivalent to finding only the keys that exactly match the query; we call this Exact-Match Attention (EMA), and formally define it in Appendix B
In our simulation, we treat each input token as the local machine, and all the local computation is handled by the element-wise functions Q, K, V. The bulk of the proof is to handle the message delivery between machines using EMA. By Proposition 24 of [53] we can assume that each machine only sends messages to at most O(N) machines for some < S. We assign a unique positional encoding or identifier each machine, and this encoding serves as unique key to retrieve the message in each machine. The high level idea is to create query for each machine and key for each destination machine and the associated value is the embedding of the message sent to the destination machine in the protocol. Since each machine can send at most messages to other machines, create a EMA heads and each head is responsible for one outgoing message for all the N machines. Each machine retrieves the message sent to them by having a query in each head. Since the messages are averaged together, we use the same embedding mechanism from Lemma 3.2 of (55] to allow error correction in the element-wise operations.
This gives us sub-quadratic time reduction from MPC tion process can be implemented in near linear time, whereas it is quadratic for standard attention. In addition, this ties ANNA-transfomen by an R-round, O )(N=)- m-emory MPC protocol can be solved by O(R)-layer with H = ( for some 8 >> O. For example, following Theorem 3.1 of 26] O(1)-layer ANNA-transformer can solve 3-SUM with mH = O[N12245)
i.e., the communica-
the existing MPC hierarchy [53]: any problem solvable
4.2
MPC can simulate UIN-trrasformer
The following theorem (proved in Appendix C) shows that any I.-layer ANNA-transformer (as implemented by Algorithm 1 can be simulated by a O(.)--ound MPC protocol. Since Algorithm randomized, it uses a random seed to sample the hash functions from the LSH family. The simulation assumes access to the random seeds needed for all layers in the ANNA-transformer.
Theorem 4.2 (MPC simulates ANNA). Fix constants 0 < é' 1. For any L-layer ANNA- transformer T (as implemented by Algorithm 7 with mH = O(N),, there exists a O(L((e'-e)) round MPC protocol with local memory =O(Ne)) and P O(NI+s-E 43332/2 machines such that ग (input) = T(input), for all input E Z
Observe that the number of machines used the simulation the ANNA-trrasfoommr can be strongly sub-quadratic (and in fact, close to linear when e is large). In contrast, the simulation of a standard transformer from (55] requires N2 machines. As previously discussed (Section 1.2). this shows that MINA-transformer more sharply characterizes efficient MPC protocols than standard transformers do. On the other hand, by Theorem 4.2 round-complexity lower bounds for MPC directly imply depth lower bounds for ANNA- ArnA-transformers. This argument was used in [55] to establish (conditional)
7

---

depth lower bounds for standard transformers on problems such as graph connectivity and k-hop induction heads; these lower bounds also hold for ANNA-transformet
4.3 ANNA-transformer can simulate low-rank transformers
As mentioned in Section | there are many proposals for efficient attention alternatives. In this section, we focus on the sub-quadratic alternatives based on low-rank approximations of the attention matrix. Specifically, we ask the following: what problems are intrinsically easy for ANNA but hard for low-rank approximation attention, and vice versa?
Definition 4.3 (Low-rank attention). A low-rank attention is specified by two feature maps Y.K" K' IRd + IR" for some < N (with the intention of approximating (xximax((X))((C))) by ("XXK"(XX K" (X) E "x and then left-nultip by (X).
On input X @Nxd, it computes ( by first computing
Note that (55) gives lower bound of any low-rank attention for the k-hopp problem. Later in Section 5 we give a construction of Oloog -deephh ANNA-transformer solving k-hop, and this directly gives a type of problem that is easy for ANNA but hard for low-rank attention. However, is there any problem that is easy for low-rank attention but hard for ANNA?
The following theorem answers the question by showing any L-layer low-rank attention-trann transformer can be simulated by -layer MANA-transformer. So, under the time and parameter-efficien regime (sub-linear rank and embedding dimension), low-rank tratiion-trnnfrmmer no stronger than ANNA-transformer
Theorem 4.4 (ANNA simulates low-rank attention). For constants O < e' < 1, any low-rank attention based transformer with depth L, rank r, embedding dimension m and rm = O(N)) ) can
be simulated by an ANNA-transformer with depth O(L), number of heads H O(N(c' and embedding dimension m O(N'))
We prove Theorem 4.4 by first using OLL/(s' - -round MPC to simulate L-layer low-rank transformer, and then Theorem 4.1 give us the simulation of L-layer low-rank transformer with ANNA-transformet through MPC. The full proof is given Appendix D
Other efficient attention mechanisms based on nearest neighbor search. Reformer 35 is another efficient attention based on LSH. In Reformer, the input tokens are sorted by their (scalar) hash values. Then, this sorted list is split into equal-sized chunks, each containing only O(1)-many tokens. Standard attention is applied within each chunk. We show that the expressive power of Reformer must come from the sorting operation: without sorting, the restriction of attention within each constant-size chunk prevents Reformer from even computing basic functions like "average" with ((1) layers (regardless of the embedding dimension); details are given in Appendix
KDEformer [64] and HyperAttention (27] approximate softmax attention matrices as a sum of sparse matrix and a low-rank matrix. They use LSH techniques to find sparse elements (i.e., heavy elements in the attention matrix) and low-rank attention for the remaining components. Theorem 4.4 indicates this low-rank part does not substantially increase the representational power.
5 ArNA-transfrrmer for reasoning tasks
In this section, study ANNA-transformer on two concrete reasoning tasks: Match2 [52 and k-hop (55) These tasks are benchmarks for evaluating the reasoning capabilities of transformers, and they separate different neural architectures in terms of their representational strengths.
5.1
ANNA-transformer solves Match2
The Match2 task (52) measures the ability of model to associate paired elements with another. We show that a single ANNA mechanism can solve Match2.
Definition 5.1 (Match2). Given input sequence X (II...., .N) .ZN) E [M] for some M S (N), the i-th output of Match2(X) is 1(3j.0; 24 + Ij = mod M) for all E [N]. Theorem 5.2. For any N, M = NO(1), there exists an IrNA--ransformer T with one layer, one attention head, and embedding dimension 1 such that T(X) Match2(X) for all X [M].
8

---

5.2 ANNA-transformer solves k-hop
The induction heads (a.k.a, associative recall) task [8]] a reasoning task that predicts the next token by completing the most recent bigram. It has been identified as an important mechanism for the emergent "in-context learning" ability of LLMs.
Definition 5.3 (Induction heads). Let Σ be afnite alphabet and w E Σ>, For each E [N]. define
a(w,1) = max {0} U {j E EENNN i i, Wj-1 = w.}}
The induction head task is to compute, for each S S N, the value of Wo(w,,)-
For example, let Σ = {a, b, c} and w = aabcbabca. Then (z(w = b because the 9th token is and the last occurrence of before position 9 (which is in position 6) is followed by b.
The following theorem shows that our ANNA-transformer can solve induction heads problem using constant number of layers and sub-linear embedding dimension and number of heads.
Theorem 5.4. Fix constants 0 < < É' < 1. There exists an ANNA trnn-rransformer T with L = O())
layers, H = O(N(''(()) heads per layer, and embedding dimension m O(Ne') such that T(w), (oww,)) Vi E for all Σ.N
We prove Theorem 5.4 (in Appendix F.2) by constructing a constant-round MPC algorithm for induction heads, and then applying Theorem 4. to convert it into an ANNA-transformer.
The induction heads task was generalized by 55 to k-step variant called "k-hop".
Definition 5.5 (k-hop induction heads). Let Σ be a finite alphabet and w E ΣN Let ( i) denote k-fold composition of a (w,-) from the previous definition. The k-hop induction head task is to compute Gi (w,s) for each S S N.
Using the same example where Σ fa. b. c}, aabcbabca and i: we have (a(a((99) because the last occurrence of b before position 7 is followed by a.
As was done in [55], construct O(log -r-uu MPC algorithm for k-hop using function composition, thus yielding a logarithmic depth scaling for ANNA-transformerrs on this task.
Theorem 5.6. Fix constants exists an ANNA-transformer T with L = (log k) layers, = O(N(c'-6)44) heads per layer, and embedding dimension m O(Ne') such that Tww); W (w.,)) Vi E [or all E....
E' < 1, any k N and alphabet E with Σ = = There
The full construction given in Appendix F.2 We note that while prior works have given transformer (and Match2 (52)), these results do not directly imply ANNA- transformers constructions, given the differences in the architectures.
constructions for k-hop 12 55
Prior work [55] showed that multi-layer recurrent nets and low-rank sub-quadratic attention [6 32
unable to solve k-hopp unless the depth is ((k) their memory size/embedding dimension is (2N/kk On contrast, ANNA-transformer achieves both Ologg de-thphh and sublinear (in N) width. In this sense, the h-hop task separates ANNA-transformer from these other efficient neural architectures.
5.3 Experiments on Match2 and induction heads
We empirically test the performance of ANNA-transformer the Match2 and induction heads tasks. Experimental details are given in Appendix G Since Algorithm
not differentiable,
we train softmax version of attention as surrogate, and then distill from the trained model to an normalizes all the queries and keys in (XX) and K (X)
(based on Algorithm with angular LSH (6). Our softmax attention X to have unit norm, and computes softmax(8 with a tunable temperature parameter 3 >
The Match2 dataset is generated the same way as 37 with context length N = 32 and upper bound are able achieve zero error with { 8 hash tables and 2 = 1 hash function per table. See Figure Ta lalfor for the detailed performance. For induction heads, we
M 37. One-layer
use the dataset from (55) with number of hops = 1, context length N 100 and alphabet size
| = 4. A two-layer ANNA-transforme achieved highly nontrivial enror with e = 32 hash tables
and z hash functions per table; with more hash tables, the error rate was as low as 0.1. See
Figure lb for the detailed results.

---

Broy vs. tur, tash for dfeertt ues
Meiimum Erer everz
7.230
0.202
0.238
0.233
0:198
5.201
0.355
0:001
5.15
0.144
0.100
0.230
0:340
8.007
0.120
0.324
0.222
0:233
6:224
5.123
0.16
2.118
0.125
0:218
0:214
0.123
1:100
325
0:335
8.007
5.007
3.105
0.12
7,075
0.655
055
0.032
0.332
0:793
5,793
7,009
20
26
20
24
28
32
(a) Match2
(b) Induction heads
Figure I: All errors are averaged over runs. (a) Error rate on Match2: x-axis denotes the number of hash tables f, and different colors correspond to different numbers of hash functions per hash table. (b) Error rate on induction heads: Rows correspond to the number of hash tables in the first layer, columns correspond to the number of hash tables in the second layer. The reported error rate is the best achieved over the choice of & E 2,
Conclusion and future work
In this work, propose more efficient class of neural architecture, the ANNA trnn-ftrasformmer, which not only preserve the representational power of standard transformer characterized by the MPC framework but also yield tighter equívalence with the MPC model. Furthermore, show that constant layers A -NNA-transformers can simulate constant layers of low-rank transformer, and can solve reasoning tasks such Match2 and k-hop tasks in near-optimal depth.
There are some interesting directions we leave for future work. While our ad hoc training method was effective as proof-of-concep,, it is desirable to develop principled training method that directly optimizes the performance of an ANNA-transforme (or a differentiable variant thereof), rather than that of surrogate model. Also, empirical validation was limited to small synthetic datasets; extending these experiments to targe-scale, real-world benchmarks is a promising next step.
Acknowledgements
Part of this work was done at the "Modern Paradigms in Generalization" and "Special Year Large Language Models and Transformers, Part 1" programs at the Simons Institute for the Theory of Computing, Berkeley in 2024. We acknowledge support from the ONR under grants N00014-24-1 2700, N00014-22- 1-2713, from the NSF under grant CCF2008733, and an award from the Columbia Center of AI Technology in collaboration with Amazon.
References
[1] Josh Achiam, Steven Adler, Sandhini Agarwal. Lama Ahmad, 1lge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. GPT-4 technical report. arXiv preprint arXiv:2303.08877, 2023.
[2] Josh Alman and Zhao Song. Fast attention requires bounded entries. Advances in Neural Information Processing Systems, 36:631 066:63117-333 17-63135, 2023.
[3] Josh Alman and Hantao Yu. Fundamental limitations subquadratic alternatives to trans- formers. The Thirteenth International Conference on Learning Representations, 2025, URL https://openreviie htt.://www.....eeoooododddaaaaod
[4] Alexandr Andoni and Piotr Indyk. Near-optimal hashing algorithms for approximate nearest neighbor high dimensions. In 2006 47th Annual IEEE Symposium on Foundations of Computer Science (FOCS 06), pages 459-468, 2006. doi: 10.1109/FOCS.2006..
10

---

[5] Alexandr Andoni, Aleksandar Nikolov, Krzysztof Onak, and Grigory Yaroslavtsev. Parallel algorithms for geometric graph problems. Proceedings of the forty- annual ACM symposium on Theory of computing, pages 774-583, 2014.
[6] Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, Ilya Razenshteyn, and Ludwig Schmidt. Prac- tical and optimal Ish for angular distance. In Proceedings of the 29th International Conference on Neural Information Processing Systems- Volume 1, NIPS' 15, page 1225-1233, Cambridge. MA, USA, 2015, MIT Press.
[7] Dana Angluin. Local and global properties in networks of processors. In STOC, pages 82-93, 1980.
[8] Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, James Zou, Atri Rudra, and Christopher Re. Simple linear attention language models balance the recall- throughput tradeoff. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 41st Inter- national Conference Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 1763-1840. PMLR, 21-27 Jul 2024.
[9] Paul Beame, Paraschos Koutris, and Dan Suciu. Communication steps for parallel query processing. Journal of the ACM (JACM), 64(6):1-58, 2017.
[10] Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the Ability and Limitations of Transformers to Recognize Formal Languages. In Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu, editors, Proceedings of the 2020 Conference on Empirical Methods in Natural
Language Processing (EMNLP), pages 7096-7116, Online, November 2020. Association
for Computational Linguistics, doi: 10.18653/v 122000.emmll-mminn.5
URL https://
org.llogyoor////// acl aclanthology.oog wwo.orr//2000.. 2020. e01220..mmlp.mmmmmm emnlp-main heappmmai..577
[11] Satwik Bhattamishra, Michael Hahn, Phil Blunsom, and Varun Kanade. Separations in the representational capabilities of transformers and recurrent architectures. In The l'hirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[12] Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: A memory viewpoint. Advances in Neural Information Processing Systems, 36: 1560-1588, 2023,
[13] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Vone Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems 33, 2020.
[14] Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatter- brain:
Unifying sparse and low-rank attention. A. Beygelzimer, Y. Dauphin, Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL htttps://pperreiiee
[15] Lijie Chen, Binghui Peng, and Hongxun Wu. Theoretical limitations of multi-layer transformer. arXiv preprini urXiv:24/2.00995. 2024.
[16] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos. Peter Hawkins, Jared Quincy Davis. Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference on Learning Representations, 2021. URL https //openrevvew.net/// httt:/oormmmmtt-aaauuuuoooo
[17] Jeffrey Dean and Sanjay Ghemawat. Mapreduce: simplified data processing large clusters. Commun. ACM, 51(0):177-113, January 2008. ISSN 0001-0782, doi: 10.1 10.1442/1327452.
1327492. URL https: //oii org/10. 1145/1327452.1382 1327492
11

---

[18] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Hemighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, 7ac Hatfield-Dodee Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, htt.://mmmrriiiiiiiiuuuu///22//
Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt,
2021.
[19] GeminiTeam, Rohan Anil, Sebastian Borgeaud, Jean- Bean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, Katie Millican, et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312.11880 2023.
[20] Mohsen Ghaffari. Massively parallel algorithms, June 2019, Lecture notes by Davin Choo, Computer Science, ETH Zurich.
[21] Mohsen Ghaffari. Fabian Kuhn, and Jara Uitto. Conditional hardness results for massively parallel computation from distributed lower bounds. 2019 IEEE 60th Annual Symposium on Foundations of Computer Science (FOCS), pages 1650-1663, 2019. URL https:/:::/ap semanticscholar: :ttp:oorpuII::2111100
[22] Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. In International Conference on Machine Learning, 2023.
[23] Michael T. Goodrich, Nodari Sitchinava, and Qin Zhang. Sorting, searching, and simulation in the mapreduce framework. In Proceedings of the 22nd International Conference on Algorithms and Computation, ISAAC" 11, page 373-383, Berlin, Heidelberg, 2011. Springer- Verlag. ISBN 9783642255908. doi: 10. 0111799978---64222551---- 178-3-642-25591--1
URL https:
[24] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501.12945 2025.
[25] Michael Hahn. Theoretical limitations of self- in neural sequence models. Transactions of the Association for Computational Linguistiics, 8: 8:156-171, December 2020. ISSN 2307-387X. doi: 10.1 121//6cccclaa 0012000cl__000066....__ /Tt./appp///br....ooor////0.--. 1ttp//00..112/ta///.
Decebbr
doi
00306
[26] MohammadTaghe Hajiaghayi, Silvio Lattanzi, Saeed Seddighin, and Cliff Stein. Mapreduce meets fine-grained complexity: Mapreduce algorithms for apsp. matrix multiplication, 3-sum, and beyond. arXiv preprint 19Xiv: 2019
[27] Insu Han, Rajesh Jayaram, Amin Karbasi, Vahab Mirrokni, David Woodruff, and Amir Zandich. Hyperattention: Long-context attention in near-linear time. The Twelfth International Conference on Learning Representations, 2024. URL htttts://opporeeeaieeee httpriiiwwweneforoeoeeee net/forum? id EddEho0d2BJIM
[28] Yiding Hao, Dana Angluin, and Robert Frank. Formal language recognition by hard attention transformers: Perspectives from circuit complexity. Transactions of the Association for Compu- tational Linguistics, 10:800-8100 2022. ISSN 2307-387X. doi: 162/ta// a 00490. URL http://dx doi wttp:////......//aaaaaaaaaaaa//a
[29] Sungjin Im, Ravi Kumar, Silvio Lattanzi, Benjamin Moseley, Sergei Vassilvitskii, et al. Mas- sively parallel computation: Algorithms and applications. Foundations and Trends@ in Opti- mization, 5(4):340-417, 2023.
[30] Piotr Indyk and Rajeev Motwani. Approximate nearest neighbors: towards removing the curse of dimensionality. In Proceedings of the Thirtieth Annual ACM Symposium on Theory of Computing, STOC '98, page 604-613, New York, NY, USA, 1998. Association for Computing Machinery. ISBN 0897919629. doi: 10.1 14/1276668822777666 URL https:: :/do:i:///..... //do:.org/oo.0o org/10
1 1145/ 1145/276693 276698 276876
12

---

[31] Samy Jelassi, David Brandfonbrener, Sham M. Kakade, and Eran Malach. Repeat after me: Transformers are better than state space models at copying. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the s1st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 21502-21521. PMLR, 21-27 Jul 2024.
[32] Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. TolySketcchorrme:: Fast transformers via sketching polynomial kernels. Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 41 International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 12748-22711 PMLR, 21-27 Jul 2024.
[33] Howard Karloff, Siddharth Suri, and Sergei Vassilvitskii. A model of computation for mapre- duce. In Proceedings of the Twenty- Anmual Annual ACM-SIAM Symposium on Discrete Algorithms, SODA '10 page 238-948, USA, 2010. Society for Industrial and Applied Mathematics. ISBN 9780898716986.
[34] Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the com- putational complexity of self-attention. In International Conference on Algorithmic Learning Theory, 2023.
[35] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020. URL httpss://ppenrevier net/fforun?ddrrkgWWktttt
[36] Alexander Kozachinskiy. Lower bounds on transformers with infinite precision. arXiv preprint
arXiv:2412.20195, 2024.
[37] Alexander Kozachinskiy, Felipe Urrutia, Hector Jimenez, Tomasz Steifer, Germán Pizarro, Matías Fuentes, Francisco Meza, Cristian B. Calderon, and Cristóhal Rojas. Strassen attention: Unlocking compositional abilities in transformers based on new lower bound method, 2025.
URL
29211
[38] Eyal Kushilevitz, Rafail Ostrovsky, and Yuval Rabani. Efficient search for approximate nearest neighbor in high dimensional spaces. In STOC, 1998.
[39] Zhiyuan Li. Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems. In The Twelfth International Conference on Learning Representations, 2024. URL httts::/ wwwprrrwwwie...eeeeeooeeeeeeeeee net/forumm iideendddddii 3EWTTEY9N
[40] Nathan Linial. Locality in distributed graph algorithms. SIAM Journal on computing, 21(1): 201, 1992.
[41] Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Trans- formers learn shortcuts to automata. In The Eleventh International Conference on Learning Representations, 2023. URL http::ttt/// httt:::://ppnnrevieee
[42] Andreas Loukas. What graph neural networks cannot learn: depth vs width. In International Conference on Learning Representations, 2020. URL hattps://openreviees wtw.:eww.ee/foorrrieee
1d=B112bp4YwS
[43] William Merrill and Ashish Sabharwal. The parallelism tradeoff: Limitations of oog-prrcision transformers. Transactions of the Association for Computational Linguistics, 1 11:531-545, 2023.
[44] William Merrill, Ashish Sabharwal, and Noah A Smith. Saturated transformers are constant- depth threshold circuits. Transactions of the Association for Computational Linguistics, 10: 843-856, 2022.
[45] William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. In Forty-first International Conference on Machine Learning, 2024. URL https //openreview. net/ peet/foorum2idd=/g3333//pp
[46] John Miller and Moritz Hardt. Stable recurrent models. International Conference on Learning Representations, 2019. URL
13

---

[47] Moni Naor and Larry Stockmeyer. What can be computed locally? In STOC, pages 184-193, 1993.
[48] David Peleg. Distributed computing: a locality--enssitvv approach. SIAM, 2000.
[49] Binghui Peng, Srini Narayanan, and Christos Papadimitriou. On limitations of the transformer architecture. In First Conference on Language Modeling, 2024. Utt.htttps:/oppenreee/
net/forum?id=Kidymlt
[50] Tim Roughgarden, Sergei Vassilvitskii, and Joshua Wang. Shuffles and circuits (on lower bounds for modern parallel computation). Journal of the ACM, 65(6))1-24, 2018.
[51] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 68, 02 2021. ISSN 2307-387X. doi: 10.1 10.1162/aac1. a 00353. URL https a03533
//doi.org/10.111
[52] Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. In Thirty-sevent seventh Conference on Neural Information Processing Systems, 2023.
URR.itttpp://operrevvewe net:/oorumidd-33000000
[53] Clayton Sanford, Bahare Fatemi, Ethan Hall, Anton Tsitsulin, Mehran Kazemi, Jonathan Halerow, Bryan Perozzi, and Vahab Mirrokni. Understanding transformer reasoning capabiliitiss via graph algorithms. Advances in Neural Information Processing Systems, 37:78320-78370, 2024.
[54] Clayton Sanford, Daniel Hsu, and Matus Telgarsky. One-layer transformers fail to solve the induction heads task. arXiv preprint arXiv:2408. 14332, 2024.
[55] Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and logarithmic depth. Forty- First International Conference on Machine Learning, 2024.
[56] Lena Strobl, William Merrill, Gail Weiss, David Chiang, and Dana Angluin. What formal lan- guages can transformers express? a survey, Transactions of the Association for Computational Linguistics, 12:543-561, 2023. URL https://app semanticscholar. Irg/CorpusID: 264833196
[57] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, ed- itors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, htwp://wooooediinnn...o.eee..pp. neurips. cc/ htcc//pperfffle/peeeeeeeeee www.fflllspppppee ffilee//paperr201//iiia Paper pdf
U. Von Luxburg,
Inc., 2017. URL https: htpp::/t////oooeeednggg 41fbe2446677ee911tt003311ccl3663
[58) Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04766 047668 2020.
[59] Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. In Advances Neural Information Processing Systems 35, 2022.
[60] Kaiyue Wen, Xingyu Dang, and Kaifeng Lyu. RNNs are not transformers (yet): The key bottleneck in-context retrieval. In The Thirteenth International Conference Learning Representations, 2025. URL
[61] Kai Yang, Jan Ackermann, Zhenyu He, Guhao Feng, Bohang Zhang, Yunzhen Feng, Qiwei Ye, Di He, and Liwei Wang. Do efficient transformers really save computation? In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 55928-55947. PMLR, 21-27 Jul 2024.
14

---

[62] Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages. Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics, pages 3770-3785, 2021.
[63] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequeeee functions? International Conference Learning Representations, 2020. URL https://ppenrevier. www.eewwweeeoooeee net/forum?
BddByxRMONtvv
[64] Amir Zandieh, Insu Han, Majid Daliri, and Amin Karbasi. KDEformer: Accelerating trans- formers via kernel density estimation. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th Inter- national Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 40605-40623. PMLR, 23-29 Jul 2023.
15

---

NeurIPS Paper Checklist
Claims
Do the main claims made the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: We summarized all the results and contribution of the main paper in the our results section in the introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contribu- tions made in the paper and important assumptions and limitations. No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We have all the assumptions in the preliminaries part and briefly discuss the limitation about training in the last paragraph.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiscless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., the approach was only tested on few datasets with a few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each thcoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
16

---

Justification:
We have all the assumptions applicable to all the theorems in the preliminaries part and theorem- -perei--pecific assumptions in the theorem statement. We have all the proofs in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experi- mental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide all the experimental details the Appendix G Our main claims only focus theoretical properties and experiments do not serve as the main contribution.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution a dataset and/or model, the authors should describe the steps taken to make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility results, access to a hosted model (c.g., in the case a large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed instructions for how replicate the
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution a new model (e.g., large language model), then there should either be way to access this model for reproducing the results a way to reproduce the model (c.g.. with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
Open access data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, described in supplemental material?
17

---

Answer: [No]
Justification: We will try to clean up the code and provide open access for the camera-ready version. Our main claims only focus on theoretical properties and experiments do not serve the main contribution.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines http://///iis.ccccpppppe guides/Cooeeuumiis het.bbmmssioooooooog.. for more details.
While we encourage the release code and data, we understand that this might not possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-souree benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips. cc/ Bttee/guides//ooeSuue wttpaaummissiiooooooooooooooooee oliicy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLS to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide these design choices in the Appendix G
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Question: Does the paper report eπтоr bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All the error reported are averaged over 10 runs and the variance is very small.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main
claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar the standard deviation or the standard error of the mean.
18

---

It is OK to report 1 I-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of emrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the ex ments?
Answer: [Yes]
Justification:
We used GPUs and provide the type of the GPU in Appendix G
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experi- mental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http://p////oopp..cooooo/ooottta
Answer: [Yes]
Justification: All the experiments are in the toy theoretical setting, and the data are all synthetic data which don't involve any human.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there is special consideration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work focus understanding the theoretical properties of the computational model we proposed. We see the societal impart to be minimal.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address socictal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
19

---

The conference expects that many papers will be foundational research and not tied to par- ticular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: Our experiments are on synthetic data, at most layer small transformers, and do have any pre-trained model.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We used the dataset from 555. We cite this paper in the experiment section and follow the license.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the package should be provided. For popular datasets, wwwppswwwtttt.. thcode. com/ com/dataset has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
20

---

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: we do not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their sub- missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, applicable, as well as details about compensation (if any)?
Question:
Answer: [NA]
Justification: This a theory paper and do not involve any human subject.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included main paper.
According the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This is a pure learning theory paper and do not involve these.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non- standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
21

---

Justification:
LLMs are only used to assist writing and editing the paper.
Guidelines:
The answer NA means that the core method development in this research does not involve
LLMs as any important, original, or non-standard components.
Please refer our LLM policy oettee///eeerrrgeecc//ooooeeeeer for what should or should not be described.
22

---

A Proof of Theorem 3.4
We restate Theorem 3.4 here, which gives the theoretical guarantee for Algorithml D
Theorem A.I (LSH algorithm guarantee for ANNA; Theorem 3.4 Fix c > v3, ISH family H that is (r, cr, Pi, P2) sensitive with quality p 1/3, l = @(N log N), and O(logi/vg N).
Then Algorithm 7 (with H. f, and 2) implements an ANNA mechanism with parameters r, and = O(1NN'
Proof. Our algorithm applies for the regime with large approximation factor, i.e., c > v3. Since only want the nearest neighbors within distance with the query point, we want to bound the probability of two points with distance greater than cr to fall into the same bucket. Consider the family G with Prgea(g(x) = s(v)] S 0.1 if | yll > cr. Then for each bucket, the expected number of collision (z. y fall into the same bucket and || 1--y/ y/| > cr) is less than N. Prgge [g(x) = [ S 0.1 Therefore, by Markov's s inequality, for each bucket, with probability greater than 1 0.1 N7, there is no collision within the bucket. Then, by union bound over all the non-empty bucket (there are at most N of them), with probablity greater than 1 0.1 there is no collision in one hash table. By 30). Ollogi/p2 N), i.e., each hash function g & is composed of O N) hash functions sampled from the LSH family H. which suffices to achieve Proec(g(r) = g(v)l S 0.1 whenever || y/ > cr.
N
NZ-
N
On the other hand, since the probability of collision is very small, the success probability (when || - y/| I r). namely p = Pr = -vv]] = - (recall that log 1/p1 ). is also somewhat small. However, we can boost the success probability by using multiple hash tables. Let & denote the number of hash tables. Then for each qs. the probability its r-nearest neighbor & (A NV(g,))) r)) falls into different bucket with qi for all e tables is upper bounded by (1 p)² By union bound over all possible nearest neighbors and all q 's, the failure probability is bounded by 2 p) Assume
log 1/p2
we want the failure probability to be less than some 5 >> 0, then we want 22 y) < 8. Taking logarithm of both sides, and using Taylor expansion of log(1 - x)) for sufficiently small æ, we find that & f= O(N3p(log N + log 1 suffices for success probability 1 - 8.
Therefore, by union bound over all é hash tables, with probability 1 3 there is no collision in all the hash tables, which implies wij = || qill > cr. By setting ó N we get l = O0N3 log N... Hence, the total failure probability 7 is bounded by ö + ạ O(1/N' -30).
0.1
which is
If [A; 2:1 S r, from the guarantee above, we know that ky collides with qi at least once in the hash bucket. This implies Wig 2 1/count, where count is the number of all the collisions in the l hash buckets that qi retrieves. In the worst case, all the k E (V(g)) collides with qa in all & hash tables except for hy only colliding once. Therefore, count I (N(g)) - 1) l, and this gives us Wiy Z (1
Runtime and memory usage. One can see that for each query, we need to evaluate O(N3p logi/ps N) hash functions and compute sum of m-dimensional vectors, so the total run- time is OmmN1-30 logi/p; N). During the preprocessing, we need to store N hash tables and the sum of values, each with at most N buckets, so the total memory is O(mN 1+3p log N) bits. In fact, the space used can be further improved to O(mN log N) bits. Instead of maintaining & hash tables, one can just store hash table of size O(mN log N) with each entry responsible for tracking the values for each query. For each round of hashing (E rounds total), hash all queries using the hash functions and creates empty buckets for them. Then, hash each key, and if the key hashes to an existing query bucket, its value is added (along with a count). After processing keys, each query accumulates the values and counts from its corresponding bucket. We give the memory- implementation in Algorithm 2
OOtimall (data-oblivious) LSH schemes achieve p = 1/e² + o(1) 4.. Since we assume e > V3, the failure probability 1/poly(AV) decreases zero with N.
23

---

Algorithm 2 Linear memory ANNA implementation with LSH family H. hash tables, and 2 hash functions/table
Input: Input X -RNx
Output: ANNA output for each of the query.
1: Let qi Q()); K = ((X),, and v2 V(X): for all E [N].
2: Initialize an array of tuples A, and A[i] + Vi [N].
3: for u
e do
4: Sample z hash functions 1, 2......... ha,s i..d. from H.
5: Create empty hash table T, indexed by hash codes of queries (below).
6: for each query do
7:
Compute hash code gu ga(4)) = (Nu., (q:). Au.=(
8:
Create an entry in Tu indexed by gu();) and (9))) - (0.
9:
for each key-value pair (kij, 03)) do
10:
Compute hash code (a(k) (hv.1 (kiz), 2 (ks),...,,uu (kz)).
if T. 19 (ky)] exists in T., then Tul (k;)] + (vj, 1).
12:
for each query do
13:
A[i] += [g:]
14: Initialize a dictionary attn + {(q1, 0), (02, .)..... (qx. 0)}.
15: for each query 9: do
16: attn < A[]]0///A
17:
return attn
B ANNA-transformee can simulate MPC
Our simulation of MPC using Exact Match Attention (EMA). In EMA, we require the key to be exactly the same as the query for it to be considered in the attention matrix. We show that this special case already suffices to simulate MPC.
uses only special case of ANNA, which we call
Definition B.1 (EM Attention). Let X E Nxd be the input embedding, Q, K. V. R - Rxd be query/keyvvaue embedding functions. For any query q, let (g)) [kg K ky = q). For each query qu the Exact Match attention computes
Σ JEN((.)
by i(
V
0
otherwise.
EMA layer and E.MA-rannformmr are defined analogously. To see that EMA is special case of ANNA, notice that ANNA, we can set r 0. > oo and Wa, = (N((:)) such that it becomes exactly the same as EM attention. EMA also admits a near linear-tima algorithm: sort all the keys first (using lexicographic ordering) in time (dN (dg N) and space OddN); ); at query time, perform binary search in time (d (dg N) per query.
We first give simulation that directly simulates the R-round (c,E)-MPC using L. R + 1 layers but large embedding dimensions to showease the core idea of the proof.
Theorem B.2 (EMA simulates MPC). For constant 0
L, any deterministic R-round MPC pro- tocol with N machines with = O(N)) words local memory, there exists an EMMA-tranffrmmr T with depth R+1, number of heads H O(N=), and embedding dimension m O(N log N), such that T Tinputt) (innutt) for all input E Zř.
Proof. For any R-round MPC protocol N with machines that maps the input to output, we define the intermediate steps for local computation phase and message transimission phase. We denote the in- put all the machines before the local computation as MachineIn,, MachineIn.......... ..aehin...MMMMieIIiI... and denote the information after deterministic local computations (Local))
MachineIng.
as MachineOutg. where MachineDut! = Mocal((Machine
MachineOut;, MachineOut,. ... ... In the communication (message transimission) phase, we need to route the messages the correct. machines ie from MachineOut, to MachineIn,:1.
24

---

In our simulation, each token input to the EMA-transformer plays the role of machine in the MPC protocol. We simulate the local computation functions (Local))re(M,see)) by the element-wise functions Q(-,, ((), (-) in the architecture. Therefore, the simulation process can be partitioned into 3 different parts:
1. Initialization. The input feeded into EMA-transformer is distributed in the N tokens, and we need to transfer than into the first N tokens/machines to match MachineIng.
Routing (message transmission). After the local computation in each round r, we need to communicate the messages from MachineOut, to MachineIng+1-
3. Final output. The MPC output is distributed in the first N tokens/machines, and we need to distributed them back to the tokens.
The following 3 lemmas construct the elements for each of these 3 parts.
We first show the message transmission part of MPC be simulate by the EMAAtransformerr Recall that after r rounds of local computation, each machine i has a set of messages it wants to send to other machines, denoted by MachineDut, the set of machine indices that machine i will sent the message to and Maggest is the message machine i send to machine dest. After the message communication phase, each machine has the set of messages it receives from other machines, denoted by MachineIn; = {Msg, Src) (Msg, i) E MachineDut, Sre) Since each machine can only send/receive s words, we have Eeettgeett [Msg] S and Σ (iteg..)/Mecchueeut [Mag] for machine We call this process the routing process of MPC. The following lemma shows that each routing round MPC can be simulated by one layer of EMA-transformer.
{Msga dest): dest sent'), ), where sent' is
Lemma B.3 (Routing). For any R-round MPC protocol ा having q machines each with local memory and any r R 11, there exists EMA- EMA-trrasformer route, with H O heads and m = (log q) for Q and K, m = O( log q) for V that takes input X Machine In, and produces oufput route, (X) MachineIn.+11
Proof. Follow the assumption in (55,, we encode the local computation into the element operations Q(-), K(.), V-) of transformer. The main part of the proof will focus using EMA to route MachineDut, to MachineIn,+1-
We assign unique positional encoding identifier to each machine i, denoted by ps. This can be done with Ologg q) bits. This encoding serves as a unique key to retrieve the message in each machine. The high level idea is to create query for each machine i and a key for each dest E sent, and the associated value is the message Magdess sent to dest in the protocol. Since each machine can send at most s messages other machines, create EMA heads and each head is responsible for one message for all the q machines. Each machine retrieves the message sent to them by having a query in each head. Because each query attend only perfectly matching keys, each distinct outbound message must be passed by a different attention head, but multiple inbound messages may be received by the same attention head.
Specifically, let Qh, Kh,, yh be the query, key, value embedding after the machine local computation for each head h (s). Set q p: for all h, so
P2
Q' Q Q
Let ksf Pawst; for destj, E sent = {desti, dests,.... dest where dest is the destination machine index for the jth word message that machine i sends. The key matrices are constructed as follows:
Pdest] Paeat
Paest; Poest?
Paeat? Poset
KL
K²
K
Paestz/
Paeat?/
Paeaty
25

---

Let oe be some embedding of (Ms dest 1), denoted by 2. = emb; (Msg (aeet) dest i) for some emb, defined later, and embj (Msg desti, 1 1) emby (Ms gaett? dest
emb dest 1)
emb MMg Magaat dest? 2)
V
V
emb (Mag /dest?
dest?, q)
emb (Mag dest q)
By such construction of Q, K, V. in our EMA, each query will retrieve the average value of the messages whose key exactly matches the query. However, by setting the value matrix this way, we might corrupt the message when there are more than one k. E K that are equal to the same query. To solve this problem, we can apply the same multiple hashing-bass -aased encoding in Lemma 3.2 from [55], which encodes each message in multiple fixed locations generated by a sparse binary matrix and have extra "validity bit" indicating whether the message is corrupted not. We restate an adapted version of their Lemma 3.2 here.
Lemma B.4 (Lemma 3.2 of [55; message encoding in sparse averaging)- For any message size N, message count bound E N, there exist an encoding function such that & takes in (Msgaesc;: dest i) defined above where the size of it is bounded by A for all i [g] and h [a] and encodes it into (Msgaestttt dest i) E IR2 with m = O(a2 log q). and a decoder function such that ip takes in the output of the EMA with Q, K, V defined above and decodes it into (Msgaest) desti; i).
Let rcvd = fsrci, sreh...,,scch where src is the jth source machine index that machine i receives message from. Because |sent'| < and [revd*| I s, in each head of EMA, there are at most values get retrieved and averaged for each query. Thus, here we can just apply Lemma B.4 with = A which gives us embedding dimension bound = Os5 log 4).
We then show with one layer EMA-transformer,, we can properly initialize the setup of MPC by ... .....
inputa) to Machine the input before the first round of MPC computation, ie the input is distributed evenly on the first [ machines.
converting Input = (input,, inputt..........................
Lemma B.5 (Initialization). For any R-round MPC protocol having q machines each with local memory s and n-word input, there exists an EMA- AilA--tranffrmer init with for Q, K and ((s) for V that takes input and outputs init(input) = MachineIn..
head m O 4)
Proof. The input should be distributed accross each machine 1 < [ with MachineIn; = {iinputiazs idx) idx E s(i- ((i-1))=e +1 +.... min ( n, si}}. }. Let qsn = be the number of machines to store the initial input. Since the input given to init are n tokens (here we treat each token as a machine), we need to rearange the memory that the input distributed the first 4in tokens.
Same as before, we use the positional encoding p; to be the unique identifier for each machine. We create key value pair for each input token and the key corresponds to the identifier of the machine that inputiax goes to and the value be (inputiax: idx). Also, create a query for each machine i [dmm].
For each machine [g.n]. define the query embedding qi = Pt+
P
P2
For each token input idx [n], let destiax the key embedding kidx Pdestiac
idx
be the machine storing the token, define
Poeat, Pleatz
K
P Plest./
26

---

Let = idx mod For each token inputiax idx: idx [n]. define the value embedding Viax E R be (input idx) in the 21' 2i'-th entry and 0 in all other entry.
input 0
0
input
2
2
0
0
input,
V
imput
s+1
+
0
0 0
By setting the value matrix like this, we can avoid corrupting the messages.
Last, we show that with an additional one layer EMA-transformer,, we can map the final round MachineIng to the output of MPC protocol where the output is store in the first = machines.
Lemma B.6 (Final Output). For any R-round MPC protocol having q machines each with local memory and n-word input, there exists an EMA-transformer out with H = head, m O 9)
for Q,K K and m O(s) for V that takes MachineIng and outputs out (MtachinnIIn)) T (input)) output.
Proof. First, the element-wise operations can compute MachineOuta from MachineIng. The output distributed accross each machine < iis = - qous. with memory of machine be output {(outputien iax: idx) idx E s(i 1) +....... min (n, si}}. Then, we wejuust need to retrieve the output tokens from the Yout machines and distribute them back to tokens. This step does the inverse job of init. We create query for each token output,ar for all idx [n). Let sr
idz be the
machine token outputiaz is in.
For each token output jax idx E [n]. define the query embedding Pidx = Percue
Parc, Pare,
Q =
Parc
For each machine [Rour], create key k; Pi.
D.
P2
K =
Paunt
The value associates with each key is the memory MagOut, stored in each machine i. Define the value embedding = MagOut,
MsgOut, MsgOut
=
MsgOut 9oa
By choosing proper element-wise function yp, out(MachineIIm)a output,.
The theorem follows from stacking the elements from these three lemmas. Each lemma gives single layer of the final EMA-transforrmer 7' with embedding dimension m (SN55 log N):
T out o route o ... ... route; o init
Remark B.7 (General (y,,)-MPC) The above simulation works the same for (7: MPCC by padding max(0,
N) empty chain-of ohouhhthtttht tokens in the input.
Remark B.8 (Number of heads). The standard transformer can simulate MPC using the same embedding dimension but only attention head [53 (5). Here the EMA needs O(Ne) heads, and we leave how improve the number of heads for future work.
27

---

Since EM attention is a special case of ANN attention with r= 0 and e - 0o, the simulation of MPC with ANNA-transformer naturally follows from Theorem B.2
Corollary B.9 (ANNA simulates MPC). For constant 0 < any deterministic R-round MPC protocol π with N machines with O(N)) words local memory, there exists an ANNA- transformer T with depth L R 1, number of heads
ONN)), and embedding dimension m = 0( log N), such that T (iimtut) = π (input) for all input E Ezz.
Theorem B.2 B.2only gives us an MPC simulation the sublinear local memory regime when & ((N1/5-0) for any > However, a lot of MPC protocol algorithms require s = (2(N1/2) local, such as MPC algorithm for 3-SUM [26] and algorithms for graphs [53]. The above simulation using EMA-transformer does not yield sublinear embedding dimension. (53] further gives a simulation of MPC with sublinear local memory using transformer with sublinear embedding dimension by simulating one round of MPC protocol with ((1) layers of transformer, instead of just one layer. Their improvement also applies here.
Theorem B.10 (EMA simulates MPC with improved embedding dimension). For constant 0 < E < E' < any deterministic R-round MPC protocol ा with N machines with s = (NN) words local memory, there exists an EMA-transformer of depth I, = (R), number of heads H ONN and embedding dimension m = O(Ne)) such that T(input) = (inputt) for all input E Z
Proof. The proof relies simulating any MPC protocol by a restricted version of MPC protocol [53] which limits the number of machines each machine can send message to. Then, use modified version of Theorem B.2 to simulate this restricted version of MPC.
Definition B.11 (Definition 3 of [53]). For constants E.E.P. > 0, a (7,e, MMPC protocol is a (Y, MPC protocol with an additional constraint: in each round, each machine can only send/receive messages from k = ((n)) machines, while the total size of messages it can send and receive is still 8 = =ONN). We refer to k as the communication capacity.
gives a construction that simulates a R-round (7, eM-PC protocol with O(R)-round p) MPC protocol. We restate their proposition here.
53
Lemma B.12 (Proposition 24 of [53]: (y,E, MMCC simulates (7,5)-MPCE -MMCC) For constants ", s >> and p E (0,e/2) function f can be computed by an R-round (7,E)-MPC, it can also be computed by a ( R(++)* )-round (7, e, -MPC protocool
Therefore, we just need to simulate (y, -MPC protocol using EMA-transformert The simu- lation follows the same recipe as Thereom B.2 where we have the initialization, message passing and final output phase. Since the initialization and output of (7, p)-MPC follow the same rule as (7, cMMPC, we only need to modify the message passing part of the simulation which corresponds to the routing Lemma B.3
Lemma B.13 (EMA simulates (t.s,p)-MPCC For constant O < < < 1, any deterministic R- round MPC protocol π with machines with s = O(N)) words local memory and communication capacity k = O(Nº), there exists an E.M-traasformer of depth T, = R + 1, number of heads H = O(N)) and embedding dimension m ONN++A0 log N) such that T(input) = o (inputt) for all input Zgy.
Proof. For the initialization and final output part, we just use the same init and out constructed in Lemma B.5 and Lemma B.6 In the routing part (Lemma B.3), because |sent'| S k, |reva'] S k. we only need l heads and in each head of EMA, there are at most i keys matching each query and thus at most values get averaged for a single query. Therefore, we can apply Lemma B.4 with O( log N). This gives us new
= k and A = s, leading to an embedding dimension route', with reduced number of heads and embedding dimension for each round r.
Likewise, we stack the 3 building blocks of one-layer EMA transformer
and have an (R + 1) layer
T out route ..... route' init
and this finishes the construction for the lemma
28

---

Let p= min(s/2, c)/44) -e) (7.E)-MPC by an = 0 min( (c)-))a) simulate this Rrruo (y, p) -MCC by an R² + 1-layer EMA transformer with ((NP) heads and
In this setting, c. By Lemma B.12 we can simulate the R-round
R R(1+))
) )-round EY)E, ().:P) Then, by Lemma B.13 we can
embedding dimension O(N log N) O(Ne).
Again, the improved simulation result of ANNA-transformer follows from Theorem B.10
Corollary B.14 (ANNA simulates MPC with improved embedding dimension). For constant
e' < L, any deterministic R-round MPC protocol π with N machines with s = =[N=) words local memory, there exists an tNNA-transformer of depth L = O(R), number of heads H = O ) and embedding dimension m. = O(N')) such that T(input) = ग (inputt) all input E Zy.
C MPC can Simulate ANNA-traansormn
As warm-up, we first simulate EMA-transformers using MPC, and then generalize it to the simulation of ANNA-transfommers.
Theorem C.1 (MPC simulates EMA). Fix constants O < <E< << < 1. For any L-layer EMA- transformer ' with mH O(N)), there exists a ((o-n)-round MPC protocol with local memory =O(Ne) and P = O(N-=z-e') machines such that (iinuut) = T(input) for all input E .zy..
Proof. We first show how to use MPC to simulate one layer of EMMA-tranffommrr In the high level, for each token I,, we have token Machine i which is responsible for computing the key, query and value embedding for æ; and other sement-wisse computation on ø4. The main bulk of the proof is to search for the exact matching keys for each query and send the averaged values associated with the matching keys to the token machines. In order to do this, we sort all the key and value pairs (As, v;) in the order defined by the key. We divide the sorted key and value pairs into buckets such that each bucket contains the same keys. For each bucket, we have a "meta-info" machine to store the indices of the machines that contains the keys in the bucket. We then compute the averaged values within each bucket and store the averaged value into the "meta-info" machine and propagate the value to all the queries that match with the key.
To begin with, let denote the space of query and key, and we define a comparator < over in order to sort. Without loss of generality, just define it to be the lexicographical ordering comparator. Based on this comparator, we define a query ranking permutation of [N] by o = (Ø), o....., σx) and key ranking permutation of [N] by d'
o'y) such that
go; < 9az <... qNN
and log < Kog < ... < kio'
For the" mnte-met- machine, we use a uniform hash function to their corresponding "meta-info" machine. Recall that for uniform hash function h, P(h(a) l(bb) N. for any a, b E X' and a o b. Therefore,
- [N] to map queries and keys
PP3i such that the size of bucket h(q)) 2 s) IPs3s different elements fall into one bucket)
<
1 N
NeT
With high probability, each "meta-info" machine is responsible for at most keys queries. We divide the machines into different types and summarize the role of each type of machine here:
For E [N), Machine i is the token machine for This machine performs all the element-wise computation for token Specifically, it computes the query, key, value embeddings 4, h; vs and element-wise operations after the attention layer.
29

---

For i E [[mm//s],, Machine (i. Q) is data structure machine for sorting queries and storing the ith chunk of the sorted list of queries after sorting. In other words, let ny [s/m] be the number of queries each machine can store and, at the end of sorting, machine (1, Q) stores (10,-1)nntt (0:n..
For E [2mN/s], Machine (i, is data structure machine for sorted list of key and value pairs, In other words, let ny = [s/2m] be the number of key and value pairs each machine can store and, the end of sorting, machine (i, KV) stores {(Koje Ding12 fate 1) 12 (Kot., "ol )
For E [N], Machine (i, h2) is the "meta-info" machine for the queries whose hash value is i. i}. This machine stores the location information of q E hg in the
Let h = {gilā E [N]. sorted list. Specifically,, for all q E hig- this machine stores the start machine index, i.e. (start, Q) where start = arg m Machine (j. Q) and the end machine index, i.c. (end, Q) where end arg maxi(q E Machine (j, Q)).
For i E [N]. Machine (i, hb) is the "meta-info" machine for the keys whose hash value is i. Let hig = [kylā E [N], k((k)) = 1}. This machine stores the location information of & E hig in the sorted list. Specifically, for all k E Kg this machine stores the start machine index, i.e. (start, KV) where start = arg min, (k E Machine (j, KV),, and the end machine index, i.e. (end, Q) where end = arg max, (k E Machine (j, KV)).
The auxiliary machines needed for message propagation.
We proceed to discuss the MPC protocol for computing the output of one layer single head EM- attention transformer. In the first round, (same as the token dispersion stage of (55]), route each token I; to its corresponding token machine i.
In the second round, each token machine computes the query, key value embedding q: = =((x) K K(x)), U, = ((x)) and sends (g;,i) to the sorting query data structure machine ([mi/s], Q) and (ki, U2 i) to the sorting key data structure machine ([2mi// s/ KV).
Then, sorting query data structure machines ((i, Q) for all E [mN/s]] sorts the queries. Sorting in MPC has been well studied, and this can be done in constant number of rounds 223.
Lemma C.2 (MPC Sorting). There exists an MPC protocol with local memory 8 = O(N')) that can sort N items and each item has size O(N5), < in O rounds with O(N machines.
After sorting, for each E [N]. we need to send the location information of g. ie which data structure machines contains q. to its "meta-info" machine (h(q)), hy). The idea is to build an -aary tree structure to aggregate the information and each query data structure machine is a leaf node this tree. Recall that each machine (i, () stores the queries S {9at--)) lowng }. If S contains the start and end of a particular query vector q. then (i, Q) sends a message (g,, (i, Q)) to machine (h(q)) hg)) Machine (i, Q) also sends the first and last query to its parent machine in the tree, i.e. sends the messages (fo (1)1))) (i, Q), first) and (qo.... (i, Q), last). After the parent node collects all the messages from the leaf node, it then does the same as its child: it contains the start and end of certain query q. it sends to the location information (the first and last machine that store it) of the query to its corresponding "meta-info" machine (h(q). hgy), and sends the first and last query and their location information to its parent machine. This is done recursively, and since there are [mN/s
query data structure machines in total, the depth of this A ary tree is O(loss/ mN/s) = 0 ( which means ((zLz) rounds and Om/N/s) machines suffice.
We do the same for (ks v) pairs. The sorting data structure machines (i, KV) sort the (k, v) pairs based on the order of k. As before, we build 2m -ary tree to send the location information to the "meta-info" machine of each key. The different part from query is that we combine the values that have the same key. For each machine in the 2m, -ayy tree, it computes the averaged value associated with each key it contains and sends the averaged value to the corresponding "meta-info" machine. In particular, for each k; the "meta-info" machine for ki (h(k;)) hx), contains the information (k;, 6) where T is the average of vy 's such that k3- ki.
30

---

Next, the "meta-info" machines of query and key need to exchange information to retrieve the corresponding value for each query. Each (i, hy.) sends the (k, ū) pairs it has to the machine (i, hg)) Then, each (i, hig) machine matches the y and kt and sends the associated value i to the q. Note that this step can be done by back propagating the m -ary tree constructed for sending the location information of q to (h(q), hag). In other words, we can just reverse the message sending direction in this tree. Therefore, each query in the query data structure machine receives the value retrieves and from each query data structure machine (i, Q), we can send the retrieved value for each query to its corresponding token machine, which is the inverse of the second round.
To summarize, the total rounds needed is (-²ze) and the number of machines needed is O(mN/s)
O(NI+e-e'). To make this work for H heads, we can create H copies of this and each copy runs in parallel. Since mHH O(N*), the bounds for number of rounds and machines still hold. By creating this MPC simulation for each of the L-layers, we stack them in the order of layers yielding
the complete simulation for L-layer
Next, we generalize the above algorithm and proceed to simulate the ANN attention that can be computed by Algorithm 1 Since Algorithm iss randomized algorithm, we assume that the MPC protocol shares all the random seeds needed for all the layers of ANNA-transformer.
Theorem C.3 (MPC simulates ANNA). Fix constants < & < E' <. For any L-layer ANNA- transformer' T (as implemented by Algorithm [ with mH ON+),, there exists a O(L/(E - -))- round MPC protocol with local memory s =O(Ne) and P +(N +2/2) machines such that T (input) = T(input), for all input E Z
is very similar to simulating EMA. We have the same kinds of machines as before. The biggest difference is that, instead of having one hash table for queries and keys, we now have hash tables, one for each round of hashing, and we sort the queries and keys based on the hash values of queries and keys. Again, we first outline different types of machines we will use.
Proof. The high level idea of simulating
For E [NV machine is the token machine for 24. This machine performs all the element-wise computation for token i. Specifically, it computes the query, key, value embeddings qs h;s U2 and element-wise operations after the attention layer.
For E [[mN/s]], t E [C], machine (i, Q, A²) is a data structure machine for sorted queries for the t-th hash table and the i-th chunk of the sorted list of queries, where the ordering of sorting is based on 9r (g) from Algorithm| [
For i E [2mN/s], s1, Machine (i, KV, hh") is a data structure machine for sorted list of key and value pairs for the t-th hash table and the i-th chunk of the sorted list of key and value pairs, where the ordering of sorting is based on (k)).
For i E [N], E [8]. Machine (Ir(),) hy, () is the "meta-info" machine for the queries whose 1-th hash value is 9r (.:). Let hn = [gglū [N), ge(g,) = 9e (4.)). This machine stores the location information of q E Mr. in the 1-th hash table. Specifically, for all E h2 this machine stores the start machine index, i.e. (start, Q,11') where start = arg min, {g E Machine (j, Q, h()), and the end machine index, i.e. (end, Q21²*) where end = arg max E Machine (j,Q, (.)).
For i E [N], E [8]. Machine (Ax), hg. 1) is the "meta-info" machine for the keys whose L-th hash value is gi (k:). Let ing (kylū 1 [N], 9,,)), = ge (k:)). This machine stores the location information of E hi, in the f-th hash table. Specifically, for all E hig this machine stores the start machine index, i.e. (start, KV, 12)) where start = arg min, (k E Machine (j, KV, h))), and the end machine index, i.c. (end, Q, h²t) where end = arg max, {k E Machine (j, KVKK
The auxiliary machines needed for message propagation.
Like before, we still use each token machine to compute the embeddings qi, kis 02 E R Then, each token machine need to send (q.,) i) and (kss 0y i) to the data structure machines, machine ([mi/s], Q, h') and machine ([2mi/s KV, h²). for all e L. Because L = 3r,,, we use the ary
31

---

tree to propagate the queries and keys to the corresponding data structure machines. This takes
O O(---)) rounds and O(N1+3p-6--') machines.
Then, for each query hash table E [e], the data structure machines sort the queries based on the hash value of the queries. Same as Theorem C.I we use the & --ary tree to send the location information of each hash bucket to its corresponding "meta-info" machine. For each key, value pair hash table t E [f], the data structure machines sort the key, value pairs based on the hash value of the keys. After that, use the 2m -ary tree to propagate the information to the corresponding "meta-info" machine. The difference from the EMA simulation is that each machine in this 2y:: -ary tree maintains the sum of values whose key has the same hash values instead of the averaged value, and also maintains count of the number of keys. These can be done in Ol ((--e) rounds and O O(NI+Spts==)) number of machines.
m
Next, the key "meta-info" machine send the sum of values and count to the corresponding query "meta-info" machines, i.e. machine (9,(,,), hg.t) sends to machine (ge (g:)) hyg t). Each query "meta-info" machine then follows the - -ary tree, broadcasting the sum of values and counts the queries in the hash table. finally, each query in the hash table needs to propagate the information back to its original token machine. Since each token machine will receive message from é N machines, we again reverse the m -aary tree that send the query to each data structure machine. During the aggregation, each machine in the counts it receives. After receiving the sum of values and counts, each token machine then calculates ANNA( = sum of the values divided by the counts.
m -ary tree still maintains the sum of values and the sum of
The above simulates one layer of ANNAA-trnsffrmert in O( rounds and using ONN 1(3+++p
machines, where = 1/c².. Therefore, by stacking the simulation for L layers, this gives O = rounds in total. To extend to H heads, we just need to instantiate the above simulation for H parallel copies and because mHH (e), the total number of rounds and machines still remains the same. D
D ANN/EM Attention can simulate low-rank Attention via MPC
We simulate the low-rank attention using ANN attention by first giving MPC algorithm for comput- ing low-rank attention and then convert it to MENA-transformer.
Theorem D.1 (ANNA/EMA simulates low-rank Attention). For constants O < < g' < 1, any low- rank attention based transformer with depth T., rank r, embedding dimension m and rm O(N))
can be simulated by an
with depth O( (() number of heads H O
and embedding dimension m = O log N).
Proof. We prove this theorem by first proving that any one-layer of low-rank attention can simulated by constant number of rounds of MPC.
Lemma D.2 (MPC simulates low-rank Attention). For constants 0 < E < < 1, any one-layer low-rank attention with rank r, embedding dimension m. and rm O(N*) can be simulated by O 2 -rrond MPC protocol with local memory s = O(Ne)) and O(N) machines.
Proof. Assume rm = O(N) and local memory of MPC S = O(Ne') where E < e'. Same as what do in MPC simulating EMA, for each token 2;, E [N], have token machine to compute the embedding of x; but we need to compute it in the kernel space, i.e. 9i = Q' (2x1,, k, K' K* and v; Vira). To compute K" (X(X)V(X), recall that
K'X) = Σ kad
We just need to compute the sum of N matrices of size r x m. Each token machine i computes the matrix kacy and we construct - J-rry tree of machines to compute the sum. The leaves of the tree are all the token machines and each node is responsible for computing the sum of _ number of matrices. We know from the previous simulation that the depth of the tree is O ). After we obtain the matrix M (XX Rx in order to compute Q(X) (X) (X), we just
32

---

need to propagate the matrix M to all the token machines. And each token machine computes q. M. By reversing the direction of message propagation the computing sum tree, we can propagate M to all the token machines in (g---) rounds. Therefore, we can simulate kernel attention with O rounds in total.
For L layers of low-rank attention transformer, construct the MPC for each layer using Lemma [D.2] and again we use the local computation of each token machine to simulate the element- computation. We stack the MPCs together, which has O (g/ rounds. The theorem follows from applying Theorem B.2 and Corollary B.9
Since the core of the proof is through MPC simulating low-rank Attention, we can also apply Theorem B.10 and Corollary [B. which simulate MPC with better embedding dimension to get improved embedding dimension for simulating low-rank attention transformer.
Corollary D.3 (ANNA/EMA simulates low-rank Attention with improved embedding dimension). For constants 0 < & < e' < 1, any low-rank attention based transformer with depth L, rank r, embedding dimension m and rm O(Ne) can be simulated by an EMMAANNA-rranffoomr
with depth O (TTre)) (ee-e)minle2 number of heads H O(N "") and embedding dimension
m O(N=).
E Discussion on Reformer
We formally define Reformer as a computational model here. Definition E.1 (Reformer attention). Given query; key, value embeddings Q(X),K(X) v(x) R N xm such that qi == k; Q(x)[i, = (x)),,, = V(X)(i, Reformer attention proceeds as follows:
1. Apply hash function h. m - U on (91.....9N' :x);
2. Sort all qi is's (and thus k; 's) by h(gg) and partition all q, 's into chunks of size B S 0(1), and let h (q)) be the label of the chunk that is in (the queries in each chunk can have different hash values);
For each qu only attend to ký such that they are in the same chunk.
The output embedding for is therefore
exp((g;, k))
Σ
Σ
Uy
J(M((k))=hh()))
We define fe : [N] + [ as the function that specifies the set keys each query should compute inner product with in the E-th layer. From the Reformer constraints, we have Vi [N]:
1. fe( = {aI, 02, .g....,ag./ ag} E [ is set (no repetition).
2. E fe()).
3. For any j E fe(1), fe(i) fe()).
In the f-th layer attention computation for each query 4i, Reformer computes
expliq,, k.))
Vz.
exp((gr, ky))
jEfe(f)
F'EJ((
We first study restricted version Reformer that fix the communication pattern beforehand i.e. fe input-independent for all e [L], and show that it can not compute the sum of all the input tokens.
33

---

Definition E.2 (SUM), Given input X (21, E [M.. and M = NNO0, the SUM task is defined SUM(()) Σ ii= Tz. We say Reformer T computes SUM if for all X,
x = SUM(X).
Here T(X) the N-th output of T given the input X. One can choose any position to be the final output position, and here WLOG choose the last token to follow the autoregressive generation model convention.
Proposition E.3. Fix = and (feléer Any Reformer T with L layers and each layer the attention pattern is specified by {fe] can not compute SUM(X): there exists an X, NT(X)N SUM(X)) 2 e, for any <<<<> M/2.
We denote each layer's element- wise computation by { Let (X), denote the i-th output of T after layers of computation. We prove this proposition by induction.
Proof.
Inductive hypothesis: T (X); is function of most B different 2; X.
Base case: t=1
exp((gi, k;))
T'));
E
Vjj
Σ
expl(gi, ky))
jEfi()
T'EA(()
exp((Q(<:), K(z)))
Σ JEN())
V(zy)
E exp((Q(x:)) ((x))))
f'Eh(()
$1 (T Toz..... Fan) wherefi (t) fa...... ag)
which is function of most B x, 's in X. Inductive step: consider
Tet (x) = Σ Σ
(X)) v(T)()), j'Elet1(0)
jEfe+i(()
= (T) ( (X)an) where ( = ={a {ai..... aB}
Since each (X)),, )a, is a function of at most Be variables from X, 7 (X); is function of at most B B B variables from X.
Therefore, if (X); is function of all {x1,.... IN), need B > N and thus L = 2(log N). In the case B = (11) and L= 0(1, (L)))x)))) ; is a function of B < variables. WLOG, consider (LX)N is function of (B), Then, TB441 can be any number in [M] that makes (*)(N far from SUM(X).
Therefore, if Reformer has any power, it must come from the sorting part, because the sorting algorithm have access to the information of all the token inputs.
Although constant-layer Reformer can not compute SUM, one can easily show that one layer of ANNA-ttrasformer can compute SUM by setting v Nz, and kı kg KN gN for all i [N], thereby retrieving all the v, 's and averaging them.
ArrA-trranerrrmer Solves k-hop and Match2
F1 Aranforetra Solves Match2
Theorem F.I. For any N, M (O(1), there exists an EMAA-ransformmr T with one layer, one attention head, and embedding dimension 1 such that T(X) Match2(X) for all X [M]
34

---

Proof. Given input X E [0, M Mx1,
MMM Let Q(X) o(X)Q, - o(x),, V(X) O(X)V,
where Q, K, are matrices in 2x1 Define by ((x) = (z. 1) and
= ()
-
M ) V=
K
such that
I1
M 21
I2
M 22
x(X)==
(XX)K-
xXXVV =
JN
M NN
As result, for each 1 isis S N, if there exists j S N such that x; + 2. = M, then
((6(x))Q)((((X)))))))}
K [N] + MH Otherwise, ((o(xx)))(((x))))]]i. for all S j S N. Finally, we can calculate that Hi E [N] 24 +2 = MY * 0, then
EMA(Ø(X)Q (X)), <())v|
11 [N] :ntty MY = 1,
[ xi + My
and if i(j E [N] x; + カæ = M] 0, then
(MA(6(X)Q. (XX)
That gives the same result for ANNA-transformers:
Corollary F.2. For any N, M NOO), there exists an ANNA--trasffommer T with one layer, one attention head, and embedding dimension 1 such that T (X) Match2(X) for all X E [M]N.
F.2 ANN transformer Solves k-hop
We first show that ANNA transformers can solve induction head (I-hop).
s' < and |>] S N. There exists an ArNA-tranrfrrmer with = ((Te--Ep)) layers, H = O(N(''=a)/) heads per layer, and embedding dimension m ONN') such that T(w): = (o(w,)) if (w, i) # 0: T(w): = (f(( E [or all w E N
Lemma F.3. Fix constants
Proof. We prove this lemma by designing constant-round MPC algorithm with local memory = O(N) and N/s machines to solve 1-hop. Since Σ < N, each token can be embedded with Ologg N) bits (O(1) words). Denote the input woN = (I1 .2,..., .N). The MPC algorithm works as following:
1. For each vi retrieve the next token x1+1 and each token the machine is stored the embedding of (I,, i, 2441, 1).
2. Define comparator < for the object (z,, i, 1,41, 1). For two tuples (1,,4, + 1) and
(uj,d,=f+1,) + 1), if v; * xz, then 24 < x, 7 (,,,,, 1) < (1 + 1); z; Tj. then < j = (14,1,
< ( +1). Sort (xx, 2, 1) by the
comparator <.
3. Each token (a, 1,411 + 1) in the sorted list retrieves the token before it the sorted list, denoted by 2. + 1). Update the embedding of token: if z, Fy the embedding of the token z; becomes (1, (... + i.. (i, Wn(m,)) or (w. i)); if wy # w;, then the embedding of the token 2, becomes (i, L.).
4. Send each (i, (w.))) o(w,)) i)) to the correct output machine [ and output (o(w,)) for token i.
35

---

For step 1. each machine only needs to send message to its neighbor machine: machine sends message machine 1, and this only takes 1 round. In step 2, each tuple is only Ollog N) bits. so by Lemma C.2 the sorting takes O rounds. In step 3, again each machine only needs to send message to its neighbor machine: machine sends message to machine i i+ 1, and this only takes I round. In step 4, each machine for the sorted list sends at most 8 tuples stored in it to the correct output machine which takes round. Thus, the MPC algorithm has O(2) rounds total.
Then, we convert this MPC algorithm to an ANNA-transformer. By Corollary B.14 this gives us an ANNA-transforma with number of heads H = O(N((' -)44. embedding dimension m = O(N2) and number of layers L O 0(<-6'--3>>
Now we show that Arnn-transformers can solve k-hop with Ollog k) layers.
Theorem F.4. Fix constants 0 < < < < 1,
| S N and any k E N. There exists an ANNA- layers, H = O(N(e' -6)/1) heads per layer, and embedding dimension m O(Ne)) such that T(w); Wi (wr,)) Vi E for, all w E NNN
transformer T with L = (5 +
log (6'-e)2
Proof. We prove this theorem by constructing Ollog -rouuss MPC with O(Ne) local memory and (N/ss) machines. We show that this MPC algorithm can compute k-hop by induction. log k] k22 and k f-11 - ky2', where ky 10.
Let k
Inductive hypothesis: after ) + 2 rounds of MPC computation, the token embedding for each token i encodes the information of this tuple
(i, o2 (w,z): (w. i), W.g o (uc,i)
Base case: l = =0, is implied by Lemma F3 Now consider k 8 1. For each the machine (Machine [i/s]) that contains (1, w Wi (w.s), 0²6 (w, i), Wo Ait (w, i)) sends the message (i, (w, i)) to machine that con-
After step 3, we have (i, (w.)) (w, i)).
tains o² (w, i) as the first entry of the tuple which is machine send the following tuple to machine [j/s]:
p (wrS))
Machine (an.4) then
(02 (w. i), a (w. (w.a)) oz (w, (w, i)), o o (w, (w,,o"(w,,)) (w,)))
= (o² (w, i), a (m,,)) 02
(w, (w, w oe (w,s) o i))
= (w, i), h1, 12, 13. 14)
Since each machine has at most tuples and the function (w, i) is one-to-one except for J., the number of messages each machine sends and receive is bounded by s. After machine [i/s] receiving the above message, it update the tuple for the token i:
1. if kg token i updated as: (i, t1, 12, W te (ww,)) k (w, i))
2. if kg 1, token i is updated as: (1, tx.
By definition, the embedding of token i now is:
(i, g (m.e)) 2 (w, i), w wrreaaan (wa)))) kets (w, i))
The above inductive step only takes 2 rounds of MPC. Therefore, the total round is 0 + 2(E + 1). When (log k] + 1, this algorithm compute the output for k-hop.
Again, we can convert this MPC algorithm to an ArNA-transformer. By Corollary B.14 this gives us an ANNA-transformer with number of heads H = O(N(C-))/44 embedding dimension m ONN')) and number of layers I. = O( -ej
log k
(e (e
36

---

G Experimental details
Here are the details of the experimental setup, All the experiments are launched on GPUs: NIVIDIA Titan RTX and NVIDIA Titan Xp.
We train a modified version of the attention matrix and then distill from the trained model us- ing ANNA implemented by the angular distance LSH family from [6]. Our softmax atten- tion normalizes all the queries and keys in Q(X) and ((X) to have unit norm, and computes softmax(3 (X))) with hyperparameter 8 >> 0.
G.1 Match2 experiments
Dataset generation.
Inspired by the way [37 generating data for Match3 task, triple-wise version
of Match2, we generate the data for Match2 using the same algorithm but change to pair-wise relation when computing the label. Each sample a tuple (X, Y). where X = (#1, 22 , FNN),
). and each x; an integer sampled from f1,2. ..... 36); Y (y1. U2...... YN and each y/ 1{3j 2; 2, O mod 37). The sequence length N set to 32. When sampling the data, we ensure that each batch is balanced by having the distribution of one's in Y the same: each batch has 4 bins and each bin corresponds to each percentage [0, 25%), [25%, 50%), (50%, 75%), [75%, 100%] of one's in Y; each bin size is 1 1/4 of the batch size. See Algorithm 3 for details.
Algorithm 3 Match2 Dataset Generation
Input: N 32, M 37, dataset size
Output: Dataset of (X, Y) pairs
1: Initialize 4 empty bins for oneepperce ranges: [0, 25), [25, 50), (50, 75), [75, 100]
2: N₂ D/4
3: while total examples in bins < D/40 do
4: Sample X E {1...., M) uniformly random
5: Calculate the percentage p of one's in Y
6: if size of the bin p is in N₂ then
7: Add (X, Y) to the correct bin
8: for each bin do
9: while size of bin < N₁ do
10:
Randomly sample (X, Y) from bin Sample permutation и over [O...., N 1] (()]
11:
12:
X
13:
Add (X", Y*) to bin
14: Combine and shuffle all bins into final dataset
15:
return Dataset
Training details. We trained 3 models with 3 E {0.1, 1, 10) respectively, with Adam optimizer cross-entropy loss and learning rate 0.01. Each model has one layer, one attention head, embedding dimension 64 and MLP with width 4m and GeL.U activation. The dataset size, batch size, training steps are 10000, 32, 20000 respectively.
We apply ANNA with number of hash tables & (1, 16} and number of hash functions for each table z {1, .,..., can be 0). Since the implementation of ANNA is randomized, for each combination of (e, 2), 10 times and report the averaged error over the runs. See Figure Ta for plotted performance when 0.1. In this setting, 2 1 can achieve error the test set with 256 test samples.
, 6) on the trained models, and 3 = 0.1 has the best performance (error
G.2 Induction heads experiments
Dataset generation. We use the data generation algorithm from (55). Each sample (X, Y) of the form X = (k, X') for h E {0, 1} for training samples and k for test samples, X' E N1 (0. Y"), where Y is the k-hop label for X". Here the sequence length N 100 and alphabet size Σ 4.
37

---

Training details. We trained models with 3 E {0.1, 1, respectively, with Adam optimizer on cross-entropy loss and learning rate 0.01. Each model has layers, each layer with one attention head, embedding dimension m 128 and an MLP with width 4m and GeLU activation. We use online training: at each training step, sample fresh new data train. The batch size and training steps are 32, 400000 respectively.
We apply ANNA on all the 3 trained models. For the first layer, the number of hash tables is chosen from {32, 40, 48, ..... 96} and is chosen from {1, 2, 3, 4J. For the second layer, the number of hash tables is chosen from {4, 8, 12, ....... 32} and is chosen from {1, 2, 3, 4). When evaluating the test error, we compute the eπror on all the tokens. Note that this is different from [55], where they only compute the eπror on the tokens whose induction head exists avoid overestimating the performance when k is large and has a large fraction of null outputs. In our setting, k = 1 which doesn't 't have many null outputs, and important for the model to learn when to output the null token.
We found B has the hest performance, so we report it in Figure Tb Again, the errors are averaged over 10 runs for each combinations and taken the minimum over s's. One can see that 32 hash tables in the first layer and 4 hash tables in the second layer already gives highly non-trivial performance: the error rate 0.2 over 100 samples and each sample has 100 token predictions, while random guess would give 0.75 error rate. With more hash tables in the first layer, the error rate can go below 0.1.
38


Reviews:
[{"review_id": "bI5PP4zpBh", "reviewer": "Reviewer_XSKf", "summary": "In this paper, the authors introduce Approximate Nearest Neighbor\nAttention (ANNA), a sub-quadratic attention mechanism for transformers\ndesigned to improve scalability without sacrificing expressive\npower. In particular, they first establish a tight connection between\nANNA-transformers and the Massively Parallel Computation (MPC)\nmodel. More precisely, every MPC protocol with R rounds and sub-linear\nlocal memory can be simulated by an ANNA-transformer with O(R) layers\nand a sub-linear number of heads and embedding dimension; and every\nANNA-transformer with L layers can be simulated by an MPC protocol\nwith O(L) rounds. Then the authors prove that each low-rank\nattention-transformer with L layers can be simulated by an\nANNA-transformer with O(L) layers.\n\nTo show the usefulness of ANNA-transformers, the authors use this\narchitecture in two concrete reasoning tasks. More precisely, the\nauthors provide constructions of ANNA-transformers for these tasks\nthat nearly match the efficiency achievable by standard transformers,\nand they provide empirical evidence that ANNA-transformers can be\ntrained to approximately solve thes e tasks.", "questions": "Q1 Could you discuss the practical implications of key architectural\nparameters in ANNA, such as the memory overhead of locality sensitive\nhashing tables, the failure probability, and the trade-offs involved\nin selecting the approximation factor c?\n\nQ2 Do you have any empirical results on real-world benchmarks or\nlarge-scale datasets to support the practical effectiveness of ANNA\nbeyond the current proof-of-concept experiments?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "**Strong Points**\n\nS1 The paper establishes a tight connection between ANNA-transformers\nand MPC, improving the previous connections that had larger gaps in\nmachine usage for simulating attention.\n\nS2 The paper provides an algorithm to implement ANNA in sub-quadratic\ntime, which is based on using locality sensitive hashing.\n\nS3 The paper provides a unified framework to reason about low-rank and\nnearest neighbor approaches to efficient attention.\n\n**Weak Points**\n\nW1 The empirical evaluation is limited, and it mostly serves as\nproof-of-concept. Performance evaluation on real-world benchmarks or\nlarge-scale datasets is not provided.\n\nW2 The paper does not discuss in practical terms the implications of\nimportant factors for the architecture such as memory overheads of\nlocality sensitive hashing tables, failure probabilities, and\ntrade-offs involving the approximation factor c for ANNA.", "was_revised": true, "final_justification": "After the authors’ response, I am convinced that this paper makes a significant contribution, so I keep my acceptance score.", "rebuttal": "Thank you for the review. We appreciate all the feedback and questions provided. \n\n> W2 The paper does not discuss in practical terms the implications of important factors for the architecture such as memory overheads of locality sensitive hashing tables, failure probabilities, and trade-offs involving the approximation factor c for ANNA.\n\n> Q1 Could you discuss the practical implications of key architectural parameters in ANNA, such as the memory overhead of locality sensitive hashing tables, the failure probability, and the trade-offs involved in selecting the approximation factor c?\n\nMemory overhead: In fact, Algorithm 1 can be implemented using linear memory ($O(N)$) with the same time complexity. Instead of maintaining $\\ell$ hash tables, one can just store 1 hash table of size $O(N)$ with each entry responsible for tracking the values for each query. For each round of hashing ($\\ell$ rounds in total), hash all queries using the hash functions and creates empty buckets for them. Then, hash each key, and if the key hashes to an existing query bucket, its value is added (along with a count). After processing keys, each query accumulates the values and counts from its corresponding bucket. We will add a remark about the memory complexity and this linear memory algorithm to the final version of the paper. \n\nFailure probability: The failure probability decays at a polynomial rate with respect to the context length $N$. Context length $N$ is our main scaling parameter for our asymptotic analysis, so the failure probability tends to zero rapidly with $N$. (Note that $N=128000$ for GPT-4 and $N=1000000$ for Gemini.)\n\nTrade-offs for selecting approximation factor: whether there are any trade-offs for selecting $c$ is highly problem-dependent. For example, for the Match2 and k-hop tasks in this paper, even a very large $c$ (e.g., $c=1000$) works based on the theoretical results in this paper. However, for problems that are conjectured to require quadratic time, such as Orthogonal Vectors, even $c=1.1$ is probably bad.\n\n> W1 The empirical evaluation is limited, and it mostly serves as proof-of-concept. Performance evaluation on real-world benchmarks or large-scale datasets is not provided.\n\n> Q2 Do you have any empirical results on real-world benchmarks or large-scale datasets to support the practical effectiveness of ANNA beyond the current proof-of-concept experiments?\n\nThe primary goal of our paper is to make theoretical contributions to the formal understanding of transformers and related architectures, and, as a result, we did not perform empirical studies on real-world benchmarks or large-scale datasets.\n\nWe'd love to discuss more if you have any further questions."}, {"review_id": "vd0nQidgQn", "reviewer": "Reviewer_1KAS", "summary": "This paper studies the ability of the transformer architecture to simulate the MPC model when the standard attention mechanism in the transformer is replaced with sparse attention based on Approximate Nearest Neighbor. The authors show that ANNA transformers (i.e. transformers that use Approximate Nearest Neighbor Attention) have the same expressive power as standard transformers in terms of simulating the MPC. In addition, the authors show that ANNA transformers may be strictly more powerful than transformers that use low-rank attention, in the sense that there are problems that are easy for ANNA but hard for low-rank attention, and ANNA transformers can simulate low-rank attention based transformer. Finally, the authors show that ANNA transformers can solve some canonical reasoning tasks such as Match2 and k-hop with near-optimal depth. Overall, this paper extends prior study on the transformer's representational capacity (with respect to simulating MPC algorithms) to an interesting case where the attention matrix is sparse.", "questions": "Line 355-356, I think there is an error. In this example, sigma(w, 9) = 7, so sigma(w, sigma(w,9)) = sigma(w, 7) = 6, and hence w_{sigma(w, sigma(w,9))} = a, not b.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The paper is very well-written and easy to follow. The theoretical results are a nice addition to our understanding of transformer's representational capacity in terms of simulating MPC algorithms, when the attention matrix is restricted to be sparse and computable in sub-quadratic time.\n\nThe accompanying empirical results in the appendix are short, but nonetheless they validate some of the theoretical results in the paper.\n\nOn the weak side, as in prior work [51,53], this paper assumes that the element-wise operations (Q, K, V, ψ) can be arbitrary functions. This deviates from the typical transformer architectures that are used in practice.", "was_revised": true, "final_justification": "After reading the author response as well as the feedback from other reviewers, I would like to keep my score and recommend acceptance. I do agree (with other reviewers) that additional practical experiments would be extremely valuable and make this work stronger.", "rebuttal": "Thank you for your review, and we appreciate the suggested change. \n\n> On the weak side, as in prior work [51,53], this paper assumes that the element-wise operations (Q, K, V, ψ) can be arbitrary functions. This deviates from the typical transformer architectures that are used in practice.\n\nSuch assumptions are necessary to establish an equivalence to MPC, since the MPC model also allows arbitrary computation by each machine on their local memory content. This is a reasonable abstraction because our goal is to understand the capabilities/limitations due to the attention (or attention-like) mechanism. (That said, many concrete MPC algorithms do have a simple local algorithm.)\n\n> Line 355-356, I think there is an error. In this example, sigma(w, 9) = 7, so sigma(w, sigma(w,9)) = sigma(w, 7) = 6, and hence w_{sigma(w, sigma(w,9))} = a, not b.\n\nThank you for catching this error! We will fix the example in the final version of the paper.\n\nWe'd love to clarify more you have any further questions."}, {"review_id": "tFT4VE6iFZ", "reviewer": "Reviewer_sfjv", "summary": "This paper studies Approximate Nearest Neighbor Attention (ANNA), a subquadratic transformer mechanism that matches the expressiveness of vanilla attention. More precisely, they show an equivalence between ANNA and MPC algorithms (Theorem 1.1) and an equivalence between ANNA and low-rank transformers (Theorem 1.2).  \n\nThe motivation for their work is that prior work showed a coarse correspondence between transformers and MPC algorithms. However, this correspondence required $N^2$ machines to simulate 1 layer of attention.  A key selling point of the paper is that ANNA narrows this gap to an almost linear dependency on the number of machines.\n\nIn Section 2, the authors define the main objects: standard attention, transformers, and recall what MPC algorithms are. In Section 3, the authors introduce ANNA attention mechanism. In Section 4, the authors state their main results: they first show a correspondence between ANNA-transformers and MPC (Theorems 4.1 and 4.2)  and then use this link to compare ANNA with other sub-quadratic attention schemes like low-rank transformers (Theorem 4.4). Lastly, in Section 5, the authors show that ANNA-transformers can solve reasoning tasks such as Match-2, induction heads and k-hops.", "questions": "I listed my questions and remarks in the strengths and weaknesses section.", "limitations": "My main concern about this paper is the absence of a limitations section.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "I am not an expert at all in the field related to this paper. Therefore, I apologize in advance for my low-confidence assessment.\n\n**Strengths**:\n\nEven though I am not a specialist in this area, the core contribution is clear and significant: ANNA attention tightens the transformer–MPC correspondence from Θ(N²) to almost-linear (in N) number of MPC machines.\nThe motivation for ANNA is natural: Definition 3.3 abstracts the essential locality property (“attend only to r-near keys”) without tying it to one concrete algorithm.\nUnification of the two main “fast-attention” families: Theorem 4.4 shows that every low-rank transformer can be simulated by an ANNA transformer. Researchers can therefore analyse a single model (nearest-neighbour attention) instead of a zoo of unrelated approximations.\n\n\n\n\n**Weaknesses**:\n\nBelow I list the main questions or concerns that, in my view, the authors should still address. But again, I have very little understanding in the field. \n\n- **Novelty of the paper?**: the ingredients ­– LSH-based attention and the correspondence transformers MPC – are known. The paper’s contribution is to tighten this correspondence from $O(N^2)$ to $O(N^{1+o(1)}$ machines and, through this, to prove that ANNA strictly subsumes low-rank attention. It would help if the authors highlighted more explicitly which lemmas are new, e.g. the almost-linear-machine simulation (Theorem 4.2) and the reduction from low-rank to ANNA (Theorem 4.4). \n- **The authors should add a limitation section**: the experiments distill from a soft-max model; the paper does not yet demonstrate an end-to-end differentiable training scheme for ANNA, although such schemes exist for related LSH attentions (Reformer, Routing Transformer). \n- **Memory & constant factors**: in line 230, the authors say that the runtime is sub-quadratic in $O( N^{1+3\\rho}) $. With c=10 (which is already pretty aggressive), we get $\\rho=0.01$, so the exponent is 1.03 and $\\ell = N^{0.03} log N$. For N=16 000 that is 64 hash tables per layer – still a noticeable overhead, though far from the “900 tables” suggested by a literal reading of the bound. A discussion of how these constants affect wall-clock speed and memory would be valuable.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We appreciate all the feedback provided. \n\n> Novelty of the paper?: the ingredients – LSH-based attention and the correspondence transformers MPC – are known. The paper’s contribution is to tighten this correspondence from $O(N^2)$ to $O(N^{1+o(1)})$ machines and, through this, to prove that ANNA strictly subsumes low-rank attention. It would help if the authors highlighted more explicitly which lemmas are new, e.g. the almost-linear-machine simulation (Theorem 4.2) and the reduction from low-rank to ANNA (Theorem 4.4).\n\nThanks for your suggestion! We will highlight these connections more in the main results part of the paper. The ANNA simulation of MPC (Theorem 4.1) is new, because ANNA is a more restrictive and different model from transformer, and thus the correspondence between transformer and MPC does not apply to ANNA. Put differently, we propose ANNA not as an approximator to standard transformers (implemented via, say, LSH), but rather as an alternative, more efficient model in itself (and hence the model to be trained to start with).\n\n> The authors should add a limitation section: the experiments distill from a soft-max model; the paper does not yet demonstrate an end-to-end differentiable training scheme for ANNA, although such schemes exist for related LSH attentions (Reformer, Routing Transformer).\n\nThank you for the suggestion; we will add a limitation section discussing these issues. Note that Reformer and Routing Transformer are also not differentiable, so they also do not have end-to-end differentiable training schemes. (Cf. \"Differentiable Approximations for Distance Queries\" by Abdelkader and Mount, which gives a scheme for approximate differentiable distance query computations that is efficient in low-dimensional spaces.) But it is a beautiful open question!\n\n> Memory & constant factors: in line 230, the authors say that the runtime is sub-quadratic in $O(N^{1+3\\rho})$. With c=10 (which is already pretty aggressive), we get $\\rho=0.01$, so the exponent is 1.03 and $\\ell=N^{0.03}\\log N$. For N=16 000 that is 64 hash tables per layer – still a noticeable overhead, though far from the “900 tables” suggested by a literal reading of the bound. A discussion of how these constants affect wall-clock speed and memory would be valuable.\n\nOur theoretical analysis is not fine-grained enough to pin down the constants suppressed by the big-O notation, so wall-clock time and the precise number of hash tables is beyond the scope of our paper. Please also see our reply to Reviewer XSKf regarding (lack of) memory overhead.\n\nWe'd love to discuss more if you have any further questions."}]

Author rebuttal:

Meta-review: This paper introduces Approximate Nearest Neighbor Attention (ANNA), a subquadratic transformer mechanism that preserves the expressiveness of standard attention, which usually requires Quadratic time. Specifically, the authors establish an equivalence between ANNA and Massively parallel computation. The paper demonstrates the empirical validity on two reasoning tasks. Overall this paper should be of interest to  Neurips community.
Decision: Accept (poster)