Sublinear Time Quantum Algorithm for Attention Approximation
Zhao Song, Jianfei Xue, Jiahao Zhang, Lichen Zhang
TL;DR
This work addresses the quadratic bottleneck of attention in transformers by introducing a quantum data-structure that achieves sublinear time in the sequence length $n$ for row-wise attention queries. It integrates a quantum Nyström approximation for the exponential kernel, quantum mean estimation for normalization, and quantum leverage-score sampling for the value matrix, enabling efficient approximation of Att$(Q,K,V)$ via $ ilde D^{-1} ilde A ilde V$ with provable Frobenius-norm guarantees. The preprocessing cost scales as $ ilde O( ext{ε}^{-1}n^{0.5}s_λ^{0.5})$ (for $Q,K$) and $ ilde O( ext{ε}^{-1}n^{0.5}oldsymbol{ ext{α}}^{0.5})$ (for $V$), while per-row queries run in $ ilde O(s_λ^2+s_λ d)$ time, with $s_λ$ the kernel’s statistical dimension and $oldsymbol{ ext{α}}$ the row-distortion of $V$. The method yields the first quantum sublinear-time row-query data structure for attention, supported by a detailed treatment of bit complexity, error propagation, and empirical parameter regimes on large models. This approach could enable scalable attention computations for very long sequences in quantum-accelerated settings, potentially impacting inference and training of long-context transformers.
Abstract
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$.