Fast attention mechanisms: a tale of parallelism
Jingwen Liu, Hantao Yu, Clayton Sanford, Alexandr Andoni, Daniel Hsu
cs.LG
2025-09-11
ANNA is a sub-quadratic nearest-neighbor attention that still simulates MPC; it solves k-hop in O(log k) layers, where low-rank attention needs linear depth.
Standard attention is quadratic in length N, so long context hits a wall. Low-rank maps, sparsity, and hash-based neighbor attention all cut that cost. Few of them prove that the cheaper operator still has the representational power of full attention.
Columbia University and Google Research tie this question to Massively Parallel Computation (MPC), the MapReduce-style model with many machines, sublinear local memory, and rounds of local compute plus communication. Prior work showed that standard transformers can simulate MPC algorithms, while the reverse simulation may need about N² machines. That gap is likely inherent: attention evaluation is believed to need near-quadratic sequential time, so an MPC round with fast local work has to pay in machine count.
The question is whether a sub-quadratic attention can characterize efficient MPC more tightly.
Approximate Nearest Neighbor Attention (ANNA) treats queries as ANN queries and keys as database points. Each query may put weight only on nearby keys: r-neighbors must receive enough mass, and keys farther than c·r must get zero weight, with high probability. The implementation uses locality-sensitive hashing: ℓ tables accumulate key-value sums into buckets, and a query averages the colliding buckets. For approximation factor c larger than √3 and hash quality ρ<1/3, runtime is O(m N^{1+3ρ} log N) and space can be Õ(mN). ρ can be as small as about 1/c², so the cost is close to linear.
In the limit c→∞ and r=0, ANNA becomes Exact-Match Attention, aggregating only tokens whose key equals the query. That special case already suffices to simulate MPC: each token is a machine, elementwise Q/K/V do local work, and exact match delivers messages. In the other direction, MPC simulates the LSH implementation of ANNA with N^{1-δ+O(1/c²)} machines, near-linear when c is large, not N².
The same MPC bridge shows that constant-depth ANNA can simulate constant-depth low-rank attention transformers. For Reformer-style "hash then attend inside fixed chunks," the power comes from the sort before chunking; without that sort, constant depth cannot even compute a global average.
The main theorems are two-way simulations with constant-factor depth overhead and an arbitrarily small N^δ width overhead. On concrete tasks:
| Task | ANNA depth | Comparison |
| Match2 | 1 layer, 1 head, embedding dim 1 | Standard transformers also solve it at constant depth |
| induction heads | constant depth, sublinear width | Transcribed from a constant-round MPC protocol |
| k-hop | O(log k) layers, sublinear width | Low-rank attention or RNNs need Ω(k) depth or width Ω(N/k⁶) |
Match2 asks, for each xi, whether some xj satisfies xi+xj ≡ 0 (mod M). k-hop composes the induction-head step "find the previous copy of this token and take the next one" k times. Low-rank attention is known to need linear depth or near-linear width on k-hop; ANNA gets logarithmic depth.
Because the LSH operator is not differentiable, training uses a unit-norm softmax surrogate and distills into ANNA. On Match2 with N=32 and M=37, a 1-layer model with 8 hash tables and 1 hash per table reaches zero error. On induction heads with N=100, alphabet size 4, and k=1, a 2-layer model reaches error around 0.1 with enough tables.
Sub-quadratic attention is often sold as an engineering approximation. This paper draws a provable boundary: ANNA still captures what MPC can do, and matches MPC machine counts more tightly than softmax attention. For people building efficient attention, k-hop-style composition is a filter. Low-rank routes are structurally weak here; neighbor/hash routes are at least theoretically intact. Constant-depth ANNA can also simulate low-rank layers, so anything those layers represent, ANNA can represent.
This is a theory paper. It does not ship a trained drop-in for FlashAttention.
ANNA is a family of weight constraints, not one differentiable kernel. Training goes through a softmax surrogate; the authors ask for a more principled direct optimizer. Experiments stay on synthetic tasks with N=32 and N=100, with no language-modeling or long-context benchmark. The LSH implementation is randomized, and the MPC simulation assumes access to the hash seeds. Q/K/V may be arbitrary elementwise maps (bit-precision limited) to match MPC's arbitrary local computation, which is still a step away from real MLPs. Conditional MPC lower bounds transfer, so logarithmic-depth barriers for connectivity and k-hop still apply to ANNA.