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

원문 초록

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$.