arXiv preprint ·
Softmax Attention with Any Temperature in Fixed Dimension: Matching the Best Known Maximum-Inner-Product Bounds
- Published
- Source
- arXiv preprint, posted before peer review
- Original
- Read it at arXiv
- Topics
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 authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.
Read it on ReadSearch
Each opens after you create an account, and brings you back to this paper.
More in Data Structures & Algorithms
Get the few papers worth your time in Data Structures & Algorithms every morning.
ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.