Approaching I/O-optimality for Approximate Attention¶

会议: NeurIPS 2026
arXiv: 2605.23751
OpenReview: https://openreview.net/forum?id=EKuz0q2V2L
论文: https://arxiv.org/abs/2605.23751
领域: 优化/理论
关键词: 近似注意力, I/O复杂度, FlashAttention, 红蓝卵石博弈, 多项式近似

一句话总结¶

针对基于多项式近似的近线性时间注意力机制，本文建立了其在两级存储层次下的紧致与近紧致 I/O 复杂度上下界，证明其数据移动量在序列长度维度上达到近线性，大幅突破了精确注意力 FlashAttention 的二次方 I/O 瓶颈。

研究背景与动机¶

标准自注意力机制的计算需要先计算 Query 与 Key 矩阵的乘积 $\mathbf{Q}\mathbf{K}^\top$，再做逐元素指数化与归一化得到注意力矩阵，最后乘以 Value 矩阵 $\mathbf{V}$。当序列长度 $n$ 极大时，标准矩阵乘法不仅带来 $O(n^2 d)$ 的算术复杂度，更在两级存储层次（容量为 $M$ 的高速缓存与无限容量的慢速内存）之间引发严重的数据搬运瓶颈。虽然 FlashAttention 及其后续改进通过算子融合与分块重排将 I/O 复杂度优化至 $\Theta(n^2 d^2 / M)$，并在 $M > d^2$ 时达到了精确注意力计算的渐进 I/O 下界，但其数据传输量对序列长度 $n$ 依然具有无法回避的二次方依赖。

另一方面，理论界为了打破长上下文下的二次方算术瓶颈，提出了基于多项式近似的近似注意力算法（如 Alman & Song, NeurIPS 2023）。该方法利用度数为 $g$ 的多项式消除 Softmax 中的非线性指数项，将注意力计算转化为低秩特征展开矩阵的两次连续矩阵乘法，从而在算术复杂度上实现了关于 $n$ 的近线性时间 $n^{1+o(1)}$。然而，现代计算硬件的实际运行时间不仅受限于算术浮点数（FLOPs），更加受制于存储层次之间的数据移动开销。此前学术界对该类近似注意力算法在存储层次中的 I/O 开销完全缺乏理论刻画，究竟能否在保持近线性算术复杂度的同时实现近线性的 I/O 数据搬运量，在理论上一直悬而未决。

本文的核心切入点正是将 Alman & Song 的近似注意力计算图与经典红蓝卵石博弈（Red-Blue Pebble Game）及 $S$-partition 技术相结合，系统探究多项式展开矩阵在生成与乘法过程中的数据复用机制。核心 idea：利用多项式展开特征项生成与输入变量载入的非线性放大效应，针对高速缓存容量 $M$ 与多项式度数 $g$、展开维度 $r$ 的相对比例，设计出突破传统正方形分块范式的高维矩形瓦片（Tile）分块调度算法，并证明其在多数参数区间内达到关于序列长度 $n$ 的近线性 I/O 复杂度，同时通过 $S$-划分建立匹配的紧致理论下界。

方法详解¶

整体框架¶

本文针对的问题是加法误差近似注意力（Additive-error Approximate Attention）。给定误差容忍度 $\epsilon > 0$ 以及矩阵 $\mathbf{Q}, \mathbf{K}, \mathbf{V} \in \mathbb{R}^{n \times d}$，目标是计算 $\widetilde{\mathbf{A}} = \widetilde{\mathbf{D}}^{-1} (\mathbf{U}_1 (\mathbf{U}_2^\top \mathbf{V}))$，其中 $\mathbf{U}_1, \mathbf{U}_2 \in \mathbb{R}^{n \times r}$（$r = \binom{d+g}{g}$）是输入向量在度数为 $g$ 的多项式基下的展开矩阵，$\widetilde{\mathbf{D}} = \text{diag}(\mathbf{U}_1(\mathbf{U}_2^\top \mathbf{1}))$ 为归一化对角阵。整体流程将注意力计算转化为两步矩阵乘法：首先计算中间低秩乘积 $\mathbf{H} := \mathbf{U}_2^\top \mathbf{V} \in \mathbb{R}^{r \times d}$，随后计算 $\mathbf{U}_1 \mathbf{H} \in \mathbb{R}^{n \times d}$。

因为输入特征维度 $d$ 与多项式项数 $r$ 通常满足 $r \cdot d \ll n \cdot d$，中间矩阵 $\mathbf{H}$ 的规模远小于输入和输出矩阵，因此整个 I/O 瓶颈主要集中在第二阶段矩阵乘法 $\mathbf{U}_1 \mathbf{H}$（第一阶段乘法对称同理）。与传统标准矩阵乘法不同的是，矩阵 $\mathbf{U}_1$ 的每一行并非直接从慢速内存中读取，而是通过从 $\mathbf{Q}$ 的对应行载入 $w$ 个变量后在高速缓存中就地展开计算出多达 $\tau(w) = \binom{w+g}{g}$ 个多项式项。这一特性彻底改变了最优数据流分块的几何形态。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入矩阵 Q, K, V ∈ R^(n×d) 及多项式度数 g"] --> Param["缓存与多项式维度判定<br/>四分类参数区间 Case I - IV"]
    Param --> Tile["非线性瓦片分块调度<br/>解方程 τ(w) = Θ(M)/w 确定生成集尺寸 w*"]
    Tile --> Gen["单行生成集重叠消除<br/>d·g/w 组超图划分与字典序归属规则"]
    Gen --> Multiply["两阶段矩阵乘法执行<br/>H = U2^T V 与 U1 H 的分块加载与原位展开"]
    Multiply --> Bound["红蓝卵石博弈下界分析<br/>2M-划分与内部乘法节点容积极值刻画"]
    Bound --> Out["近线性 I/O 复杂度的近似注意力输出 A~"]

关键设计¶

1. 非线性瓦片分块调度：基于变量生成项平衡的高维矩形分块 在标准矩阵乘法中，为了最大化缓存中输入与输出的数据复用率，经典策略是构造尺寸为 $\Theta(\sqrt{M}) \times \Theta(\sqrt{M})$ 的正方形瓦片。然而在近似注意力场景下，$\mathbf{U}_1$ 的元素不是直接载入的输入，而是由 $\mathbf{Q}$ 的某一行中的 $w$ 个元素在缓存中组合生成。载入 $w$ 个变量可以生成 $\tau(w) = \binom{w+g}{g}$ 个特征项。因此，为了使得单次分块所涉及的慢速内存读取量控制在 $\Theta(M)$，不能机械地划分子块，而是需要求解平衡方程： $$\tau(w^*) = \frac{\Theta(M)}{w^*}$$ 根据解出的最优生成变量数 $w^*$，算法从 $\mathbf{Q}$ 中挑选 $\frac{\Theta(M)}{w^*}$ 行、每行载入 $w^*$ 个变量，从而在高速缓存中无需任何额外 I/O 即可动态恢复出 $\mathbf{U}_1$ 中局部块的 $\tau(w^*)$ 列。这构成了高度不均匀的矩形分块，从而在输入载入代价极低的前提下，最大化每个瓦片内执行的标量乘法数量。

2. 单行生成集重叠消除：字典序超图组合覆盖与原位项生成 在实际分块中，若要覆盖行向量展开的所有多项式单项式，必须对 $\mathbf{Q}$ 的同一行载入不同的变量子集。然而，不同变量子集展开出的多项式单项项存在大量重叠（例如包含变量数少于 $g$ 的交叉项）。如果无序重叠计算，会导致有效瓦片宽度骤降，瓦片重复读取带来高昂代价。为此，算法将每行的 $d$ 个变量划分为 $\frac{d \cdot g}{w}$ 个组，每组包含 $\frac{w}{g}$ 个变量，并通过枚举 $\binom{\frac{d \cdot g}{w}}{g}$ 种组组合来构建聚合瓦片。为了严格消除冗余，算法制定了全局字典序规则：对于能由多个组组合生成的项，强制将其唯一归属于字典序最小的那个组合瓦片。每个瓦片在高速缓存内部枚举 $\tau(w)$ 个候选单项时，通过极低开销的原位检查过滤出所属项，确保每个多项式特征项恰好被计算且参与聚合一次，保证了总体 I/O 复杂度严格受限于理论极值。

3. 红蓝卵石博弈下界证明：基于 2M-划分的未覆盖内部节点极值分析 为了证明算法的渐进最优性，本文基于 Hong & Kung 的红蓝卵石博弈模型构建计算有向无环图（DAG），利用 $2M$-划分（$2M$-partition）理论推导 I/O 转移次数的下界。计算图中总共包含 $n \cdot r \cdot d$ 个内部乘法节点。对于任意一个大小受限的划分块 $D_i$，其支配集（Dominator Set）和极小集（Minimum Set）大小至关重要。论文将 $\mathbf{U}_1$ 的行依据在支配集中的节点数是否达到阈值划分为两类子矩阵：对于高密度行，直接受限于支配集容量；对于低密度行，其能派生出参与有效乘法的节点数严格被 $\tau(w) \leq \frac{M}{4w}$ 压制。通过极值代数放缩证明，每个划分块中最多只能包含 $O(\frac{M^2}{w})$ 个未覆盖的内部运算节点，从而导出完成整个 DAG 计算所需的最小划分数 $\Omega(\frac{n \cdot r \cdot d \cdot w}{M^2})$，最终确立了 $\Omega(\frac{n \cdot r \cdot d \cdot w}{M})$ 的 I/O 复杂度下界。

损失函数 / 训练策略¶

本工作为纯理论算法与计算复杂度分析，不涉及经验训练损失函数或微调训练策略。多项式逼近函数 $P(x) = \sum_{l=0}^g c_l x^l$ 针对 $[-B, B]$ 闭区间上的指数函数 $\exp(x)$ 进行一致逼近，多项式系数可在 $\text{poly}(g)$ 时间内精确计算为多项式位长的有理数，且多项式度数满足 $g := \Theta\left(\max\left\{\frac{\log(1/\epsilon)}{\log(\log(1/\epsilon)/B^2)}, B^2\right\}\right)$。

实验关键数据¶

主实验¶

本文通过理论渐进界分析，在四种具有明确硬件与算法参数物理意义的区间（Case I - IV）中，给出了近似注意力（Approximate Attention, AA）与标准注意力矩阵乘法、FlashAttention（FA）在 I/O 复杂度上的全景对比（其中 $r = \binom{d+g}{g}$，当 $g=O(1), d=O(\log n)$ 时 $r = n^{o(1)}$）。

参数区间	约束条件	本文算法 I/O 上界	匹配的 I/O 下界	对比 FlashAttention
Case I	$M = \Omega(d \cdot r)$	$\Theta(n \cdot d)$	$\Omega(n \cdot d)$（显然下界）	$\Theta(\frac{n^2 d^2}{M})$，本文节省 $\frac{n \cdot d}{M}$ 倍
Case II	$M = o(d \cdot r), g = o(\log M)$	$O\left(\frac{n \cdot r \cdot d \cdot (4e^2)^g}{M^{g/(g+1)}}\right)$	$\Omega\left(\frac{n \cdot r \cdot d \cdot g}{M^{g/(g+1)}}\right)$	本文近线性，FA 为二次方 $O(n^2)$
Case III	$M = o(d \cdot r), \Omega(\log M) \le g \le o(\sqrt{M})$	$O\left(\min\left(\frac{n r d^2}{M}, \frac{n r d}{\sqrt{M}}\right)\right)$	$\Omega\left(\frac{n \cdot r \cdot d \cdot g}{M}\right)$ ($d \ge 5g$)	本文近线性，FA 为二次方 $O(n^2)$
Case IV	$M = o(d \cdot r), g = \Omega(\sqrt{M})$	$\Theta\left(\frac{n \cdot r \cdot d}{\sqrt{M}}\right)$	$\Omega\left(\frac{n \cdot r \cdot d}{\sqrt{M}}\right)$ ($d \ge 5g$)	行为退化为标准矩阵乘法

消融实验¶

为分析不同生成集变量数 $w$ 及参数假定对 I/O 复杂度的影响，下表对比了 Case II 与 Case III 中选择不同 $w$ 时的理论分析结果与紧致性：

理论配置	变量 $w$ 取值设定	导出的 I/O 上界	导出的 I/O 下界	理论边界状态
Case II (最优选择)	$w_0 = \frac{1}{4e} g M^{\frac{1}{g+1}}$	$O\left(\frac{n r d (4e^2)^g}{M^{g/(g+1)}}\right)$	$\Omega\left(\frac{n r d g}{M^{g/(g+1)}}\right)$	当 $g=O(1)$ 时常数级紧致
Case II (通用瓦片基线)	固定正方形瓦片	$O\left(\frac{n \cdot r \cdot d}{\sqrt{M}}\right)$	$\Omega\left(\frac{n r d g}{M^{g/(g+1)}}\right)$	上界次优，存在 $M^{\frac{g-1}{2(g+1)}}$ 差距
Case III (强假设扩展)	$w = g, d = O(g^{2-\delta_1}), M=O(g^{\delta_2})$	$O\left(\frac{n \cdot r \cdot d \cdot g}{M}\right)$	$\Omega\left(\frac{n \cdot r \cdot d \cdot g}{M}\right)$	达到完全紧致 $\Theta(\frac{n r d g}{M})$
Case IV (极小缓存极限)	元素级就地缓存	$\Theta\left(\frac{n \cdot r \cdot d}{\sqrt{M}}\right)$	$\Theta\left(\frac{n \cdot r \cdot d}{\sqrt{M}}\right)$	完全紧致，多项式特性失效

关键发现¶

在长上下文主流配置下（$M = \Omega(d \cdot r)$，例如 $M = \sqrt{n}$ 或常数级较大缓存且 $g=O(1)$），近似注意力的 I/O 复杂度严格达到平凡下界 $\Theta(n \cdot d)$，即只需要完整读入输入和写出结果即可，内部完全无需重复刷写慢速内存。
相比于 FlashAttention 的 $\Theta(\frac{n^2 d^2}{M})$，本文算法在 Case I - III 下将 I/O 对序列长度 $n$ 的依赖从二次方 $n^2$ 压制至近线性 $n^{1+o(1)}$，在长序列极限下带来的 I/O 削减比率高达 $O(\frac{n \cdot d}{M})$。
当多项式度数相对缓存较小（$g=o(\log M)$）且为常数时，本文算法的上界与下界在渐进意义上完全闭合，证明了基于生成集的高维矩形分块调度策略的最优性。

亮点与洞察¶

首次揭示了近似算法在存储层次数据搬运上的本质优越性：近似注意力不仅降低了浮点运算 FLOPs，而且通过代数解耦打破了非线性依赖，从根本上释放了分块优化的自由度，使得 I/O 复杂度同样实现从 $O(n^2)$ 到近线性 $O(n)$ 的质变。
提出了基于代数展开的非平凡瓦片分块方法：突破了传统矩阵乘法基于输入直接切块的正方形瓦片教条，针对输入通过多项式组合产生高维展开的特点，通过反解组合数方程定制高维非对称矩形瓦片，极大提升了缓存内部的数据生成与利用密度。
在计算理论层面拓展了红蓝卵石博弈在动态展开计算图中的下界证明技术：通过将 DAG 节点按照支配集覆盖容量与多项式项数极值分类，建立了一套分析非平庸数据依赖图 I/O 复杂度的完整理论框架。

局限与展望¶

展开维度 $r = \binom{d+g}{g}$ 具有强烈的组合爆炸特性。当特征维度 $d$ 较大（如目前主流大模型的 $d=128$）且为了追求极高精度导致多项式度数 $g > 3$ 时，中间维度 $r$ 会极其庞大，在工程实现中可能抵消理论上的渐进优势。
本文完全聚焦于两级存储模型（单级 Cache + RAM）的红蓝卵石博弈理论分析，尚未在具体的现代 GPU 架构（多级存储如 HBM-SRAM-Register、Tensor Core 矩阵乘单元、Warp 并行通信）中给出工程实现与实测速度评测。
未考虑带因果掩码（Causal Mask）的自回归解码场景下下三角注意力的特殊分块调度与状态复用问题。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次系统研究近似注意力的 I/O 复杂度，建立了多项式展开下的非正方形分块理论。
实验充分度: ⭐⭐⭐⭐☆ 理论推导完备严谨，四种区间分类详实，若有实际 GPU 硬件核函数性能对比会更为完满。
写作质量: ⭐⭐⭐⭐⭐ 逻辑严密，定理叙述与直觉解释层次分明，技术引理与红蓝卵石博弈分析非常清晰。
价值: ⭐⭐⭐⭐⭐ 为长上下文高效注意力机制的软硬件协同设计与底层算子开发提供了坚实的理论基石。

参数区间	约束条件	本文算法 I/O 上界	匹配的 I/O 下界	对比 FlashAttention
Case I	\(M = \Omega(d \cdot r)\)	\(\Theta(n \cdot d)\)	\(\Omega(n \cdot d)\)（显然下界）	\(\Theta(\frac{n^2 d^2}{M})\)，本文节省 \(\frac{n \cdot d}{M}\) 倍
Case II	\(M = o(d \cdot r), g = o(\log M)\)	\(O\left(\frac{n \cdot r \cdot d \cdot (4e^2)^g}{M^{g/(g+1)}}\right)\)	\(\Omega\left(\frac{n \cdot r \cdot d \cdot g}{M^{g/(g+1)}}\right)\)	本文近线性，FA 为二次方 \(O(n^2)\)
Case III	\(M = o(d \cdot r), \Omega(\log M) \le g \le o(\sqrt{M})\)	\(O\left(\min\left(\frac{n r d^2}{M}, \frac{n r d}{\sqrt{M}}\right)\right)\)	\(\Omega\left(\frac{n \cdot r \cdot d \cdot g}{M}\right)\) (\(d \ge 5g\))	本文近线性，FA 为二次方 \(O(n^2)\)
Case IV	\(M = o(d \cdot r), g = \Omega(\sqrt{M})\)	\(\Theta\left(\frac{n \cdot r \cdot d}{\sqrt{M}}\right)\)	\(\Omega\left(\frac{n \cdot r \cdot d}{\sqrt{M}}\right)\) (\(d \ge 5g\))	行为退化为标准矩阵乘法

理论配置	变量 \(w\) 取值设定	导出的 I/O 上界	导出的 I/O 下界	理论边界状态
Case II (最优选择)	\(w_0 = \frac{1}{4e} g M^{\frac{1}{g+1}}\)	\(O\left(\frac{n r d (4e^2)^g}{M^{g/(g+1)}}\right)\)	\(\Omega\left(\frac{n r d g}{M^{g/(g+1)}}\right)\)	当 \(g=O(1)\) 时常数级紧致
Case II (通用瓦片基线)	固定正方形瓦片	\(O\left(\frac{n \cdot r \cdot d}{\sqrt{M}}\right)\)	\(\Omega\left(\frac{n r d g}{M^{g/(g+1)}}\right)\)	上界次优，存在 \(M^{\frac{g-1}{2(g+1)}}\) 差距
Case III (强假设扩展)	\(w = g, d = O(g^{2-\delta_1}), M=O(g^{\delta_2})\)	\(O\left(\frac{n \cdot r \cdot d \cdot g}{M}\right)\)	\(\Omega\left(\frac{n \cdot r \cdot d \cdot g}{M}\right)\)	达到完全紧致 \(\Theta(\frac{n r d g}{M})\)
Case IV (极小缓存极限)	元素级就地缓存	\(\Theta\left(\frac{n \cdot r \cdot d}{\sqrt{M}}\right)\)	\(\Theta\left(\frac{n \cdot r \cdot d}{\sqrt{M}}\right)\)	完全紧致，多项式特性失效