Fast attention mechanisms: a tale of parallelism

Jingwen Liu, Hantao Yu, Clayton Sanford, Alexandr Andoni, Daniel Hsu

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.

Primary Area: Theory (e.g., control theory, learning theory, algorithmic game theory)

Submission Number: 22189