Title: Fast attention mechanisms: a tale of parallelism
Venue: NeurIPS 2025
Accepted: yes
Keywords: Transformer theory

representational strength

nearest neighbor search

massively parallel computation
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nI 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.\n\nQUESTIONS:\nI listed my questions and remarks in the strengths and weaknesses section.\n\nLIMITATIONS:\nMy main concern about this paper is the absence of a limitations section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nThe 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.\n\nQUESTIONS:\nLine 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\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\nQ1 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-07"}