cs.LGJun 5, 2026

Attention at the Theoretical Minimum: A Mathematics of Arrays Framework for Memory-Optimal Transformer Kernels

Authors: Lenore MullinGaetan Hains

Abstract

The attention mechanism is the dominant computational bottleneck in modern transformer-based AI. Its standard implementation incurs quadratic memory traffic in the sequence length~nn, and DRAM accesses cost 100--1000×\times more energy than arithmetic operations on contemporary hardware, so any analysis focused solely on FLOP counts fundamentally mischaracterises the bottleneck. We present a Mathematics of Arrays (MoA) reformulation of scaled dot-product attention and its numerically stable softmax, deriving a Denotational Normal Form (DNF) that eliminates all intermediate arrays -- including the implicit transposed-key buffer and every softmax temporary -- by algebraic construction rather than empirical tuning. The DNF achieves O(ndk+ndv)O(n_{dk} + n{_{dv}}) data movement versus O(n2+ndk+ndv)O(n^2 + n_{dk} + n_{dv}) for the standard implementation, where nn is the sequence length, dkdk is the key dimensionality and dvdv the value dimensionality, and is verified numerically against PyTorch at full double-precision floating-point on concrete inputs. Unlike hardware-specific accelerators or empirical tiling schemes such as FlashAttention, MoA simultaneously provides array fusion, shape-transformation correctness, and predictive cost models from a single algebraic framework. Memory minimality is a theorem established before any code is written. A predictive performance model projects 22--100×100\times speedup and 22--50×50\times energy reduction, with the advantage widening at exascale. The derivation establishes a formally verified pipeline from Python specification through (ONF) Operational Normal Form, and dimension-lifted hardware mapping, providing performance-portable AI kernels of direct relevance to DARPA edge-deployment and DOE exascale priorities.

Explore similar work

CardsList
  1. Approaching I/O-optimality for Approximate Attention

    May 22, 2026Pál András Papp, Aleksandros Sobczyk, Anastasios ZouziasQuadratic AttentionMatrix Multiplication