接近 I/O 最优性以获得近似注意力

我们重新审视 大语言模型 中注意力的 I/O 复杂性。给定查询键值矩阵 $Q,K,V\in\mathbb{R}^{n\times d}$ 和具有快速内存大小 $M$ 的机器，目标是使用快速内存​​和慢速内存之间的最少数据传输来计算“注意力矩阵”$A=\text{softmax}(Q K ^{\top}/\sqrt{d}) V$。文献中的现有方法，尤其是 FlashAttention 及其变体，会产生与 $n$ 成二次方关系的 I/O 成本，而一个简单的下限仅需要 $Ω(nd)$ I/O 来读取输入和写入输出。在这项工作中，我们提出了一种计算注意力的技术，其中 I/O 成本在大多数参数范围内几乎线性地依赖于 $n$。这是通过开发 I/O 高效算法来实现的，该算法受到 Alman 和 Song 最近的近似注意力框架的启发。我们还证明了每个参数范围中相应的下限，以表明我们的算法确实接近 I/O 最优。