提出近似注意力新算法，显著降低大模型推理时的内存传输开销。
Approaching I/O-optimality for Approximate Attention
- 基于近似注意力框架设计高效内存访问策略
- 在多数参数场景下实现接近线性的内存传输量
- 适合大规模语言模型部署与高效推理场景
我们重新审视大语言模型中注意力机制的输入输出（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$ 二次增长，而理论下界仅需 $Ω(nd)$ 次 I/O 读写输入输出。本文提出一种计算注意力的新技术，使 I/O 成本在多数参数范围内几乎线性依赖于 $n$。该方法受 Alman 与 Song 最近提出的近似注意力框架启发，设计了高效的 I/O 算法，并在各参数区间证明了相应下界，表明所提算法确实接近 I/O 最优。
原文摘要 · Abstract (English)
We revisit the I/O complexity of attention in large language models. Given query-key-value matrices $Q,K,V\in\mathbb{R}^{n\times d}$, and a machine with fast memory size $M$, the goal is to compute the "attention matrix" $A=\text{softmax}(Q K ^{\top}/\sqrt{d}) V$ with the minimal number of data transfers between fast and slow memory. Existing methods in the literature, most notably FlashAttention and its variants, incur an I/O cost that depends quadratically on $n$, while a trivial lower bound only requires $Ω(nd)$ I/O's to read the inputs and write the output. In this work, we present a technique for computing attention where the I/O cost only depends almost-linearly on $n$ in most parameter regimes. This is achieved by developing I/O-efficient algorithms inspired by the recent approximate attention framework of Alman and Song. We also prove corresponding lower bounds in each parameter regime to show that our algorithms are indeed close to I/O-optimal.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。