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

Abstract: The most widely used 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.

• The paper demonstrates that a symmetry-aware Taylor expansion of the attention kernel achieves fixed per-token compute cost regardless of sequence length.
• By exploiting tensor symmetries, the method significantly reduces the dimensionality of monomial features, leading to orders-of-magnitude memory and runtime improvements.
• Empirical tests with four Taylor terms recover performance comparable to Float16 quantization, paving the way for scalable and energy-efficient Transformer architectures.

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

Motivation and Background

Transformer models relying on self-attention dominate contemporary AI workloads. Canonical self-attention exhibits $\bigO(n^2)$ complexity in sequence length $n$, a bottleneck for scaling context windows and accommodating inference workloads with persistent histories. Existing mitigation strategies encompass low-rank/sparse approximations, structured context windows, and alternative architectures such as linear/state-space models. This work bypasses these heuristic modifications, presenting a mathematically rigorous approach that renders self-attention computable to arbitrary precision with fixed per-token cost, independent of sequence length.

Theoretical Construction: Symmetry-Aware Taylor Expansion

The cornerstone of the approach is a Taylor expansion of the exponential kernel in attention, $\exp( q^\top k / c )$, systematically decomposed into sums over symmetric chains of tensor products. For a Taylor truncation order $P$, the exponentiation can be approximated as a weighted sum of monomial inner products of increasing degree:

$$\exp \left( \frac{q^\top k}{c} \right) \approx \sum_{p=0}^{P-1} \alpha_p \left< \Phi_p(q), \Phi_p(k) \right>_{C_p}$$

Here, $\Phi_p(\cdot)$ extracts the minimal monomial basis (upper hyper-triangular region of the symmetric tensor), while $C_p$ applies combinatorial weights reflecting monomial multiplicities. This approach exploits the symmetric structure of the tensor products to substantially reduce feature dimensionality: instead of $d_K^p$ monomials for degree $p$, only $\binom{d_K + p - 1}{p}$ unique basis elements are computed and stored, minimizing both computational and storage overhead.

Figure 1: Illustration of the symmetric tensors of increasing order in Eq. (2), with the blue region denoting the minimal monomial basis relevant for the Taylor expansion.

Computational Complexity and Efficiency Analyses

For fixed $d_K$ (key dimension), $d_V$ (value dimension), and Taylor order $P$, the hidden state per token is:

$(d_V + 1) \binom{d_K + P - 1}{P - 1}$

FLOPs per token for the forward-pass similarly simplify to:

$$\left(4 d_V + \frac{2 (P d_K + 1)}{d_K + 1} + 2\right) \binom{d_K + P - 1}{P - 1}$$

These are sequence-length-independent, in contrast to conventional attention's $\bigO(n)$ per-token storage and $\bigO(n^2)$ global cost. Empirically, four Taylor terms suffice for recovery with errors comparable to Float16 quantization, enabling orders-of-magnitude memory and compute reductions as context grows.

Empirical Validation and Implementation

A PyTorch-based proof-of-concept implementation demonstrates exact recovery of conventional attention outputs as $P$ increases, with robust reduction in both peak memory (per token) and runtime as context length increases. For example, for context lengths up to $10^8$ tokens, both peak memory allocation and runtime decrease by approximately three orders of magnitude relative to canonical self-attention.

Optimizations not yet explored in the implementation (e.g., bypassing unnecessary data copies and exploiting hierarchical index structure) suggest even greater practical gains in low-level, hardware-optimized deployments.

Architectural Implications

A salient architectural consequence is the decoupling of per-token compute/memory from context length and, specifically, the reduction of cost by increasing the number of attention heads with smaller dimension: per-head costs become inversely proportional to $d_K$ and $d_V$, rather than linear in head count. This enables attention to be distributed over many narrow heads—an infeasibility for standard implementations due to linear scaling of per-token state.

Trade-offs and Tunable Precision

Rather than imposing sparsity, locality, or changing the attention pattern, the method trades precision for compute explicitly via the Taylor order $P$. Experiments indicate diminishing returns above $P = 4$ for realistic head dimensions, with higher-order terms subsumed by quantization noise.

Contrasts with Prior Taylor/Linear Attention Work

Previous Taylor-based attention approximations restricted truncation order, limited by the combinatorial growth in monomial dimensionality. By exploiting tensor symmetries, the present formulation realizes tractable computation at high truncation order, maintaining constant cost. Moreover, unlike single-feature linear attention, this method aggregates multiple independently accumulated components corresponding to each Taylor degree, all amenable to parallel scan.

Limitations and Future Directions

Limitations include lack of end-to-end downstream evaluation—current benchmarks isolate the attention mechanism from full-model training and inference pipelines. The proof-of-concept suffers in runtime due to generic library calls and sequential evaluation of Taylor terms, but these are amenable to kernel-level fusion and parallelization optimizations.

Open research directions include: training models with the proposed mechanism at scale, assessing downstream generalization and convergence behavior, compressing feature spaces for higher-order Taylor expansion, and extending symmetry-aware approximations to further analytic kernels beyond the exponential.

Conclusion

By decomposing the self-attention kernel as a symmetry-aware Taylor expansion and compactly representing polynomial interactions via minimal monomial bases, this work enables self-attention with fixed per-token compute and storage, irrespective of context length. This lays the groundwork for scalable, energy-efficient Transformer architectures capable of persistent, long-horizon inference without recourse to context truncation or heuristic approximations, with broad practical and theoretical implications for large-scale AI deployment.