面向近似注意力的 I/O 最优性

我们重新审视大型语言模型中注意力的 I/O 复杂度。给定查询-键-值矩阵 $Q,K,V\in\mathbb{R}^{n\times d}$，以及内存大小为 $M$ 的机器，目标是计算注意力矩阵 $A=\text{softmax}(Q K ^{\top}/\sqrt{d}) V$，而实现这一目标所需的在快速存储器和慢速存储器之间的数据传输次数最小。文献中现有方法，最著名的是 FlashAttention 及其变体，所需的 I/O 成本随 $n$ 平方增长，而一个平凡的下界只要求 $\Omega(nd)$ 的 I/O 才能读取输入并写入输出。在本文中，我们提出一种计算注意力的方法，其 I/O 成本在大多数参数区间几乎线性地依赖于 $n$。这通过发展受 Alman 和 Song 最近近似注意力框架启发的 I/O 高效算法来实现。我们也证明了相应的下界，以显示我们的算法在每个参数区间都接近 I/O 最优。