量子算法实现注意力计算的亚线性时间近似,大幅提速大模型推理。
Sublinear Time Quantum Algorithm for Attention Approximation
- 利用量子奈氏斯特姆近似与采样技术,仅通过行查询即可快速估计注意力输出。
- 预处理时间复杂度为 $\widetilde{O}(ε^{-1} n^{0.5} (s_λ^{2.5} + s_λ^{1.5} d + α^{0.5} d))$,单次查询仅需 $\widetilde{O}(s_λ^2 + s_λd)$。
- 适用于大规模语言模型中需高效注意力计算的场景,尤其适合高维稀疏数据。
给定查询、键和值矩阵 $Q, K, V\in \mathbb{R}^{n\times d}$,注意力模块定义为 $\mathrm{Att}(Q, K, V)=D^{-1}AV$,其中 $A=\exp(QK^\top/\sqrt{d})$,$\exp(\cdot)$ 逐元素作用,$D=\mathrm{diag}(A{\bf 1}_n)$。注意力是现代 Transformer 与大语言模型的核心,但显式构造软最大矩阵 $D^{-1}A$ 需 $Ω(n^2)$ 时间,促使众多近似方法将时间降至 $\widetilde O(nd)$。本文提出一种量子数据结构,仅通过 $Q, K, V$ 的行查询即可近似 $\mathrm{Att}(Q, K, V)$ 的任意行。预处理耗时 $\widetilde{O}\left( ε^{-1} n^{0.5} \left( s_λ^{2.5} + s_λ^{1.5} d + α^{0.5} d \right) \right)$,其中 $ε$ 为精度目标,$s_λ$ 为由 $Q$ 与 $K$ 定义的指数核的 $λ$-统计维度,$α$ 衡量 $V$ 的行失真,上限为 $d/{\rm srank}(V)$(稳定秩)。每次行查询可在 $\widetilde{O}(s_λ^2 + s_λd)$ 时间内完成。据我们所知,这是首个在 $n$ 上实现亚线性时间的注意力行近似量子数据结构。方法依赖于指数核的量子奈氏斯特姆近似、量子多变量均值估计以计算 $D$,以及量子杠杆得分采样进行与 $V$ 的乘法。
原文摘要 · Abstract (English)
Given the query, key and value matrices $Q, K, V\in \mathbb{R}^{n\times d}$, the attention module is defined as $\mathrm{Att}(Q, K, V)=D^{-1}AV$ where $A=\exp(QK^\top/\sqrt{d})$ with $\exp(\cdot)$ applied entrywise, $D=\mathrm{diag}(A{\bf 1}_n)$. The attention module is the backbone of modern transformers and large language models, but explicitly forming the softmax matrix $D^{-1}A$ incurs $Ω(n^2)$ time, motivating numerous approximation schemes that reduce runtime to $\widetilde O(nd)$ via sparsity or low-rank factorization. We propose a quantum data structure that approximates any row of $\mathrm{Att}(Q, K, V)$ using only row queries to $Q, K, V$. Our algorithm preprocesses these matrices in $\widetilde{O}\left( ε^{-1} n^{0.5} \left( s_λ^{2.5} + s_λ^{1.5} d + α^{0.5} d \right) \right)$ time, where $ε$ is the target accuracy, $s_λ$ is the $λ$-statistical dimension of the exponential kernel defined by $Q$ and $K$, and $α$ measures the row distortion of $V$ that is at most $d/{\rm srank}(V)$, the stable rank of $V$. Each row query can be answered in $\widetilde{O}(s_λ^2 + s_λd)$ time. To our knowledge, this is the first quantum data structure that approximates rows of the attention matrix in sublinear time with respect to $n$. Our approach relies on a quantum Nyström approximation of the exponential kernel, quantum multivariate mean estimation for computing $D$, and quantum leverage score sampling for the multiplication with $V$.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。