Approaching I/O-optimality for Approximate Attention

Pál András Papp, Aleksandros Sobczyk, and Anastasios Zouzias, at Huawei Technologies, analyze the I/O complexity of approximate attention rather than arithmetic complexity, using the polynomial-based framework of Alman and Song from NeurIPS 2023. For Q, K, V in R^{n×d}, they show that replacing exp(QK^T/d) with a degree-g polynomial expansion yields matrices U1 and U2 with r = binom(d+g, g) features, and this structural change lets attention be tiled far more efficiently than FlashAttention. Their main result is a four-regime characterization of optimal data movement in the red-blue pebble game model: when fast memory satisfies M = Ω(dr), the I/O cost is Θ(nd), matching the trivial input-output lower bound; when M is smaller but g = o(log M), the cost is O(nrd(4e^2)^g / M^{g+1}) with a matching lower bound up to constants; when g is larger, they prove nearly linear-in-n bounds of O(min(nrd^2/M, nrd/√M)) and a corresponding Ω(nrdg/M) lower bound under d ≥ 5g; and for tiny caches, M = O(g^2), the cost is Θ(nrd/√M). The key novelty is a nonstandard tiling strategy that loads only w entries from a row of Q or K to generate up to τ(w)=binom(w+g, g) terms, enabling tiles whose shape depends on the combinatorics of the polynomial approximation rather than on d alone. Compared with FlashAttention, whose I/O cost is Θ(n^2d^2/M) in the favorable regime, these bounds remove a quadratic dependence on sequence length n and can save an almost-linear factor in I/O.