Machine Learning PapersLandmark papers
← All papers
RESEARCH PAPER / 2026-11-04

Softmax Attention with Any Temperature in Fixed Dimension: Matching the Best Known Maximum-Inner-Product Bounds

Zixuan Huang, Guozhang Chen

Peking UniversityPeking University · arXiv

Affiliations: Peking University

Abstract

Attention is central to modern AI models for language, vision, and scientific modeling. However, direct evaluation compares every query with every key, requiring quadratic time in the context length, which limits its scalability to long contexts. Consequently, characterizing the complexity bounds of attention is of fundamental importance. For every fixed dimension $d$, we give an algorithm for $m$ queries, $n$ key-value pairs, and additive error $\varepsilon$, with expected time $ \widetilde{\mathrm{O}}\bigl((mn)^{k/(k+1)}+m+n\bigr),\qquad k=\lfloor d/2\rfloor. $ For fixed dimension and a fixed number of value channels, soft-O hides factors of the form $\operatorname{poly}_d\bigl(b,\log(m+n),\log B,\log(1/\varepsilon)\bigr),$ where $b$ bounds input encoding lengths and $B$ bounds coordinate and value magnitudes. In particular, the dependence on $B$ and $1/\varepsilon$ is polylogarithmic, which allows arbitrary positive temperature. When $m=n$ and $d\ge2$, the bound is $\widetilde{\mathrm{O}}(n^{2-2/(\lfloor d/2\rfloor+1)})$: near-linear in dimensions two and three, with exponent $4/3$ in dimensions four and five and $3/2$ in dimensions six and seven. It improves the previous $\widetilde{\mathrm{O}}(n^{2-1/d})$ upper bound and matches the best known batch maximum-inner-product bound in each fixed dimension. Additionally, we reduce integer batch maximum inner product computation to one attention call with the same dimension and input sizes $m,n$.

The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.

Research summary

The paper provides a tight upper bound for fixed-dimensional softmax attention that matches the best-known bounds for batch maximum inner product, utilizing advanced geometric decomposition and random sampling techniques. This work significantly advances the theoretical understanding of attention complexity, demonstrating that the quadratic bottleneck can be broken in low dimensions with near-linear complexity, thereby establishing fundamental limits and possibilities for efficient transformer architectures.

Methodology

The paper presents a rigorous theoretical framework for computing softmax attention in fixed dimensions. The core methodology involves reducing the attention computation to a series of geometric range-searching problems. Specifically, it utilizes a "common-anchor block" decomposition combined with random prefix chains to handle the exponential weights efficiently. The authors employ Lagrange interpolation to decouple row-dependent factors and use convex hull geometry (specifically triangulations and ray shooting) to organize near-maximal pairs. This approach allows them to match the best-known bounds for the Maximum Inner Product (Max-IP) problem, which is a significant theoretical breakthrough as it bridges the gap between attention complexity and classical computational geometry.

Results & evaluation

As a theoretical computer science paper, there are no empirical experiments, benchmarks, or dataset evaluations. The "results" are mathematical theorems and complexity bounds. The evaluation is based on the tightness of the asymptotic bounds compared to previous work (e.g., improving from $O(n^{2-1/d})$ to $O(n^{2-2/(\lfloor d/2\rfloor+1)})$).

Reproducibility

The paper is highly reproducible in the sense that it provides full mathematical proofs. However, it does not provide code implementations. The algorithms described are complex geometric data structures (persistent segment trees, triangulations) that would be difficult to implement efficiently for practical use, but the theoretical correctness is established via formal proof.

Limitations

The primary limitation is that the results are for *fixed* dimension $d$. In practical Transformers, the head dimension is often large (e.g., 64, 128, 256), and the complexity bounds may degrade or become irrelevant as $d$ grows. The paper acknowledges this by noting that for large $d$, the exponent approaches 2 (quadratic). Additionally, the algorithms are Las Vegas (randomized with guaranteed correctness) or deterministic but complex, making them unsuitable for immediate practical deployment in standard deep learning frameworks.

Broader impact

This paper has high impact on the theoretical understanding of Transformers. It establishes that, in low dimensions, attention is not fundamentally harder than maximum inner product search, resolving a key open question in the complexity of attention. This provides a theoretical baseline for what is possible and helps delineate the boundaries of subquadratic attention algorithms. It influences future research in efficient attention mechanisms by providing tight lower and upper bounds. The paper provides a tight upper bound for fixed-dimensional softmax attention that matches the best-known bounds for batch maximum inner product, utilizing advanced geometric decomposition and random sampling techniques. This work significantly advances the theoretical understanding of attention complexity, demonstrating that the quadratic bottleneck can be broken in low dimensions with near-linear complexity, thereby establishing fundamental limits and possibilities for efficient transformer architectures.

Author of this paper? Send a correction.