Research map / Architectures and efficiency / Attention and transformers

Fast attention mechanisms: a tale of parallelism

NeurIPS 2025 · Architectures and efficiency › Attention and transformers · approach: Attention head mechanisms · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Introduces Approximate Nearest Neighbor Attention (ANNA), a sub-quadratic LSH-based attention, and proves it matches MPC simulation power of standard transformers with far fewer machines, simulates low-rank transformers, and solves Match2 and k-hop with near-optimal depth.

Problem

Standard attention has quadratic cost, and prior MPC-transformer equivalences require up to N^2 machines, leaving a gap; it is unclear whether efficient sub-quadratic attention alternatives preserve transformers' representational advantages.

Approach

Define ANNA, a class of attention mechanisms whose weights are supported on approximate nearest neighbors, with an LSH-based sub-quadratic implementation; use MPC to show ANNA-transformers simulate MPC with near-linear machine count and can simulate low-rank attention transformers.

Key results

• ANNA-transformers simulate R-round MPC with O(R) layers and width O(N^{eps+delta}), while needing N^{1-delta+O(1/c^2)} machines versus N^2 for standard transformers.
• One-layer ANNA-transformer achieves zero error on Match2 with N=32, M=37, using 8 hash tables and 1 hash function per table.
• A two-layer ANNA-transformer on induction heads (N=100, |Sigma|=4) reaches an error rate as low as about 0.1 with more hash tables.

Measured effects reported in the paper

Surprising findings

• The Reformer's expressive power comes entirely from its sorting step; without sorting, constant-size chunked attention cannot compute averages in O(1) layers.

Limitations

• Training uses a softmax surrogate distilled to ANNA rather than directly optimizing ANNA, which is ad hoc.
• Experiments limited to small synthetic datasets (Match2, induction heads).
• EMA simulation needs O(N^eps) heads, a number not yet optimized.

Experimental setup

Datasets: Match2 synthetic dataset induction heads dataset

Benchmarks: Match2 induction heads

Metrics: error rate

Baselines: standard softmax attention low-rank attention Reformer KDEformer HyperAttention

Builds on: MPC transformer equivalence (Sanford et al.) locality-sensitive hashing Reformer low-rank attention

Closest papers by meaning

• Mixture-of-Top-$k$ Attention: Efficient Attention as Scalable Fast Weights (NeurIPS 2026 · similarity 0.89)
• kNN Attention Demystified: A Theoretical Exploration for Scalable Transformers (ICLR 2025 · similarity 0.88)
• Transformers, parallel computation, and logarithmic depth (ICML 2024 · similarity 0.87)
• Positional Attention: Expressivity and Learnability of Algorithmic Computation (ICML 2025 · similarity 0.87)
• Theory, Analysis, and Best Practices for Sigmoid Self-Attention (ICLR 2025 · similarity 0.87)
• Towards Interpretable and Efficient Attention: Compressing All by Contracting a Few (NeurIPS 2025 · similarity 0.87)