Self-Attention at Constant Cost per Token via Symmetry-Aware Taylor Approximation

AI summary

Overview

Research area: Efficient Transformer architectures and attention mechanisms (machine learning systems, linear attention, kernel methods).

Technical level: Advanced. The paper builds on tensor algebra, reproducing kernel Hilbert spaces (RKHS), symmetric tensors, and Taylor expansions of the exponential kernel, though the core intuition is accessible.

Scope: The paper derives, implements, and empirically validates a formulation of self-attention whose memory and compute costs per token are constant rather than growing with context length, by approximating the exponentiated query–key dot product with a symmetry-aware truncated Taylor expansion.

What This Paper Is About

Standard self-attention costs O(n) space and time per token and O(n²) time for the sequence as a whole, where n is the number of tokens in context, because each token must attend to all previous tokens. This drives demand for storage, compute, and energy that the authors say is outstripping society's ability to provide it. The authors' goal is not to replace or modify attention, but to show that the same attention can be computed efficiently to arbitrary precision at constant cost per token — O(1) — with a fixed-size hidden state and a constant number of FLOPs per token.

Key Contributions

1. A decomposition of the Taylor expansion into symmetric tensor chains. The authors show that each power (qᵀk)^p in the expansion of exp(qᵀk/c) can be written as a sum over order-p symmetric tensors q^⊗p and k^⊗p, whose unique elements live in an upper hyper-triangular region.

2. Identification and tight packing of the minimal polynomial-kernel basis. They show the upper hyper-triangular region contains m_p = binom(d_K + p − 1, p) unique monomials, and that grouping identical monomials and weighting each by its multiplicity reduces stored features and FLOPs per token by orders of magnitude versus naively evaluating all d_K^p permutations.

3. A succinct, embarrassingly parallel feed-forward feature transformation. The feature map reduces to def Phi(x): return x[..., M].prod(dim=-1), where M is a constant matrix of precomputed hierarchical indices satisfying i₁ ≤ i₂ ≤ … ≤ i_p.

4. An implemented, verified proof of concept. The authors implement the formulation in PyTorch, verify it recovers conventional attention with increasing accuracy as Taylor terms are added, and confirm order-of-magnitude reductions in memory and run time up to 100M tokens of context.

Main Findings

• Constant per-token cost. For a fixed approximation order P, the hidden state has (d_V + 1) · binom(d_K + P − 1, P − 1) elements and the forward pass executes (4d_V + 2(P·d_K + 1)/(d_K + 1) + 2) · binom(d_K + P − 1, P − 1) FLOPs per token — both fixed, independent of context length n.

• Four Taylor terms suffice. The authors find that P = 4 recovers conventional attention with elementwise errors of approximately the same magnitude as Float16 resolution, which they note is not too surprising because the scaling constant α_p declines to that resolution by the fourth Taylor term.

• Recovery improves with approximation order. Over causal sequences of 100K tokens, with queries, keys, and values sampled from N(0,1)^d, each additional Taylor term decreases reconstruction error. Outputs are compared to a Float64 implementation treated as ground truth, at four head sizes and four truncation numbers P, with differences measured in decimal orders of magnitude.

• Memory per token falls by three orders of magnitude. In an Nvidia GPU benchmark from 1K to 100M tokens of context across four head sizes, peak memory allocated per token declined three orders of magnitude below conventional attention at 100M tokens.

• Run time per token falls by almost three orders of magnitude. Using the mean of seven executions, run time per token declined almost three orders of magnitude below conventional attention at 100M tokens.

• The results are conservative. Both memory and run time are described as unfairly penalized by implementation issues, so the authors state that the figures likely understate achievable savings in optimized implementations.

• Costs scale inversely with head size. Both space and time complexity are inversely proportional to d_K and d_V, so smaller heads with more heads per token reduce fixed per-token costs — the opposite of conventional attention, where per-token costs grow with the number of heads.

• Comparison to conventional attention. The conventional formulation has a KV cache of n(d_K + d_V) elements and executes n(2d_K + 2d_V + 3) FLOPs per token, and cannot be computed via parallel scan because its stepwise transformation is not associative over tokens.

• Prior Taylor-based work stopped at second order. The authors attribute this to the perceived complexity of evaluating higher-degree polynomial interactions without exploiting symmetry.

Methodology in Plain English

The authors start from the observation that attention's core operation is exponential of a scaled dot product between a query and a key. Because the exponential function can be written as an infinite power series, they can approximate it with a finite number of terms, where the number of terms P controls precision.

The problem is that writing out each power of the dot product looks like it would require an enormous number of polynomial terms — d_K^p of them — which grows combinatorially and makes higher-order terms impractical. The authors' insight is that these products form symmetric tensors: a product like q₁q₂ is mathematically the same as q₂q₁, so many terms are redundant. Only the unique, non-repeated monomials matter, and these live in a triangular region of the tensor.

By collecting just the unique monomials, packing them tightly into a vector, and multiplying each by a precomputed count of how many ways it can be permuted, the authors obtain a small set of features. These are computed with a single parallelizable operation: index into the input vector using a precomputed index matrix, then take a product along the last dimension.

The resulting formulation is a form of linear attention. Instead of one inner product and one scan, it needs P inner products and P independent parallel scans, one per Taylor term, whose outputs are aggregated. Each term maintains its own accumulated state that is updated by a simple linear recurrence as tokens arrive.

For validation, the researchers implement the method in PyTorch, compare its outputs to a Float64 conventional implementation on synthetic sequences, and benchmark peak memory and run time per token against a conventional implementation over contexts ranging from 1K to 100M tokens.

Why This Matters

Impact on research: The paper challenges the assumption that Taylor-based linear attention must stop at the quadratic term. It shows the obstruction was a matter of not exploiting symmetry, and provides a general mathematical technique — symmetric tensor decomposition plus minimal-basis identification — that the authors describe as of independent interest. It also opens what they call a new region of the architectural design space, where head dimension becomes a tunable lever for trading off cost against configuration.

Real-world applications:

• Long-context AI services that maintain extended personalized histories or step-by-step reasoning chains without per-token costs growing with history length.
• Unbounded token generation, where a model continues generating at modest fixed cost per token provided its positional encoding scheme accommodates indefinite sequence lengths.
• Infrastructure planning for large-scale Transformer deployments, where memory capacity, processor utilization, communications bandwidth, and energy consumption are the binding constraints.
• Multi-head configurations that would otherwise be infeasible because key–value caching dominates resource usage.

Industry relevance: The authors frame the work as a response to demand for data centers and energy sources outstripping society's current ability to provide them. Their primary claimed benefit is enabling models to generate tokens in perpetuity at modest fixed cost, substantially reducing the infrastructure and energy demands of large-scale Transformer models.

Future Directions

• Training and evaluating real models. The next step the authors identify is training Transformer models end-to-end and via transfer learning using this attention mechanism, and evaluating downstream performance and convergence behavior. Doing so at scale would require fused hardware kernels.

• Hardware-optimized implementations. The current code is a proof of concept built from high-level PyTorch calls. Identified targets include avoiding m_p × p temporary copies of query and key elements, exploiting the hierarchical index structure of M_p, evaluating Taylor terms in parallel rather than sequentially on a single stream, and writing low-level on-device code such as a CUDA kernel.

• Compressing higher-order feature spaces. The authors suggest reducing dimensionality for P ≥ 4 by dropping basis components with low-magnitude coefficients in C_p, finding best-fit lower-rank basis approximations, and using random sampling or random projections for fast approximations.

• Extending beyond the exponential kernel. The approach could be applied to other analytic kernels, and the authors note that unbounded-context generation also requires positional encoding schemes that remain well-behaved at extreme sequence lengths — a complementary problem they do not solve.

• Choosing the approximation order. The appropriate P may depend on model architecture, head size, and representation statistics, and the authors recommend treating it as a tunable hyperparameter.

Target Audience

Researchers and engineers working on efficient Transformer architectures, linear attention, state space models, and long-context inference will benefit most, along with practitioners concerned with the memory, compute, and energy costs of serving large models. Readers need comfort with tensor notation, polynomial feature maps, and kernel methods to follow the derivation, though the high-level argument and empirical results are accessible to a broader machine learning audience. The paper explicitly notes that its benchmarks isolate the attention mechanism itself rather than measuring end-to-end throughput, so readers seeking production-ready deployment guidance should treat the runtime results as conservative and preliminary.

Authors’ abstract

The most widely used artificial intelligence (AI) models today are Transformers employing self-attention. In its standard form, self-attention incurs costs that increase with context length, driving demand for storage, compute, and energy that is now outstripping society's ability to provide them. To help address this issue, we show that self-attention is efficiently computable to arbitrary precision with constant cost per token, achieving orders-of-magnitude reductions in memory use and computation. We derive our formulation by decomposing the conventional formulation's Taylor expansion into expressions over symmetric chains of tensor products. We exploit their symmetry to obtain feed-forward transformations that efficiently map queries and keys to coordinates in a minimal polynomial-kernel feature basis. Notably, cost is fixed inversely in proportion to head size, enabling application over a greater number of heads per token than otherwise feasible. We implement our formulation and empirically validate its correctness. Our work enables unbounded token generation at modest fixed cost, substantially reducing the infrastructure and energy demands of large-scale Transformer models. The mathematical techniques we introduce are of independent interest.