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

TL;DR

This work tackles the prohibitive resource demands of self-attention in long-context Transformers by reexpressing the exponential kernel of attention as a truncated Taylor expansion, then organizing the resulting polynomial interactions into symmetry-aware, minimal monomial bases. By mapping each term to a compact, precomputable feature space with diagonal weights and leveraging parallelizable inner products, the authors achieve a constant per-token cost with a fixed hidden state size that scales inversely with head dimensions. The approach aligns with linear-attention paradigms while enabling higher-order Taylor terms to be computed in parallel, yielding near-conventional attention accuracy with far smaller memory and compute footprints, and enabling scalable long-context inference. However, the current implementation is a proof of concept and requires optimized hardware kernels and end-to-end training studies to fully validate performance and task-specific effects. Overall, the method opens a new architectural design space where attention costs can be made independent of context length, potentially transforming the scalability of large Transformer deployments.

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.

Figures (6)

• Figure 1: Hidden state size and FLOPs per token (4 Taylor terms), compared to conventional formulation.
• Figure 2: Illustration of the symmetric tensors of increasing order in \ref{['eq:chains_of_tensor_products']}. The blue region in the rightmost plot contains the minimal monomial basis for representing $( q^\top k )^{p}$, as discussed in Section \ref{['sec:minimal_basis']}.
• Figure 3: Reduction in hidden state size and FLOPs per token achieved by our tight packing method.
• Figure 4: Size of hidden state and FLOPs per token in multi-head configurations.
• Figure 5: Scaling constant of each Taylor term, $\alpha_p$, compared to resolution of common floating-point formats.