Near-Linear Attention in Three Dimensions

Abstract

Given n query vectors, n key vectors, and scalar values, softmax attention returns one exponentially weighted average per query. We give an n 2^{O(√(log n))}-time randomised algorithm in three dimensions on a RAM with O(log n)-bit words. Coordinates are logarithmic-bit rationals bounded by n¹⁰, values lie in [−1, 1], and the simultaneous additive error is at most n⁻¹⁰ with probability at least 2/3. Thus the optimal word-RAM exponent is α(3) = 1. The algorithm builds small sampled hulls. Within a simplicial normal cone, keys satisfying its three facet inequalities have a nonnegative score-deficit formula; dyadic coefficient bins let many queries share their moment sums. The other keys recurse. A bounded-degree triangulation makes each sampled facet occur in at most nine frames, and the expected total facet-conflict size is linear. Queries follow single paths, whose recorded local maxima also recover the global maxima. All moment sums and polynomial reconstruction use exact integers. Without a time cutoff, every execution meets the error bound and the stated running time holds in expectation.