固定维度下任意温度的 Softmax 注意力：匹配已知最佳最大内积界

arXiv:2611.06055v1 公告类型：交叉发布 摘要：注意力是现代语言、视觉和科学建模 AI 模型的核心。然而，直接计算需要将每个查询与每个键进行比较，所需时间随上下文长度呈二次增长，从而限制了其向长上下文扩展的能力。因此，刻画注意力的复杂度界具有根本性的重要意义。对于每个固定维度 $d$，我们给出了一种处理 $m$ 个查询、$n$ 个键值对且加性误差为 $\varepsilon$ 的算法，其期望运行时间为 $ \widetilde{\mathrm{O}}\bigl((mn)^{k/(k+1)}+m+n\bigr),\qquad k=\lfloor d/2\rfloor. $ 在维度和值通道数量均固定时，软 O 记号隐藏了形式为 $\operatorname{poly}_d\bigl(b,\log(m+n),\log B,\log(1/\varepsilon)\bigr)$ 的因子，其中 $b$ 是输入编码长度的上界，$B$ 是坐标和值的绝对值的上界。特别地，对 $B$ 和 $1/\varepsilon$ 的依赖为多对数级，因此允许任意正温度。当 $m=n$ 且 $d\ge2$ 时，该界为 $\widetilde{\mathrm{O}}(n^{2-2/(\lfloor d/2\rfloor+1)})$：在二维和三维时接近线性，在四维和五维时指数为 $4/3$，在六维和七维时指数为 $3/2$。这一结果改进了此前的 $\widetilde{\mathrm{O}}(n^{2-1/d})$ 上界，并在每个固定维度下均匹配已知最佳的批量最大内积界。此外，我们将整数批量最大内积计算归约为一次注意力调用，同时保持维度和输入规模 $m,n$ 不变。