Approaching I/O-optimality for Approximate Attention
- URL: http://arxiv.org/abs/2605.23751v1
- Date: Fri, 22 May 2026 15:23:26 GMT
- Title: Approaching I/O-optimality for Approximate Attention
- Abstract summary: We revisit the I/O complexity of attention in large language models.<n>We present a technique for computing attention where the I/O cost only depends almost-linearly on $n$ in most parameter regimes.
- Score: 45.67330863443465
- License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
- Abstract: 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.
Related papers
- Stronger Memory-Query Tradeoffs for Convex Optimization: The Limitations of Subquadratic Memory [65.64123585249297]
We prove two lower bounds for the first order complexity of minimizing a $d$-dimensional $1$-Lipschitz convex function over the unit ball with $m$ bits of memory.<n>We first show that any such (possibly randomized) algorithm must make $tilde(fracd2sqrtm)$ oracle queries.<n>For deterministic optimization algorithms, we show that $tilde(mind1.6,fracd8/3m2/3)$ queries are required.
arXiv Detail & Related papers (2026-07-21T00:40:59Z) - HSR-Enhanced Sparse Attention Acceleration [19.776342074253435]
We introduce a novel approach to accelerate attention computation in Large Language Models (LLMs)<n>We leverage the inherent sparsity within attention mechanisms, both in conventional Softmax attention and ReLU attention.<n>Our method only introduces provably negligible error for Softmax attention.
arXiv Detail & Related papers (2024-10-14T05:18:02Z) - How Sparse Attention Approximates Exact Attention? Your Attention is Naturally $n^C$-Sparse [9.552839922307587]
Sparse Attention is a technique that approximates standard attention computation with sub-quadratic complexity.<n> Variations of this technique, such as pruning KV cache, sparsity-based fast attention, and Sparse Transformer, have been extensively utilized for efficient Large Language Models (LLMs) deployment.
arXiv Detail & Related papers (2024-04-03T12:37:34Z) - The I/O Complexity of Attention, or How Optimal is Flash Attention? [4.126403496388015]
Self-attention is at the heart of the popular Transformer architecture, yet suffers from quadratic time and memory complexity.
The breakthrough FlashAttention algorithm revealed I/O complexity as the true bottleneck in scaling Transformers.
We show an I/O complexity lower bound that matches the upper bound provided by FlashAttention for any values of $M geq d2$ within any constant factors.
arXiv Detail & Related papers (2024-02-12T06:50:45Z) - One Pass Streaming Algorithm for Super Long Token Attention
Approximation in Sublinear Space [11.735802740426294]
Attention computation takes both the time complexity of $O(n2)$ and the space complexity of $O(n2)$ simultaneously.
We introduce a new algorithm that only reads one pass of data in a streaming fashion.
Notably, our algorithm exhibits exceptional memory-efficient performance with super-long tokens.
arXiv Detail & Related papers (2023-11-24T18:35:00Z) - Oblivious Stochastic Composite Optimization [47.48197617884748]
We show that our algorithms converge without any prior knowledge on the parameters of the problem.<n>All three algorithms work without prior knowledge of the diameter of the feasible set, the Lipschitz constant or smoothness of the objective function.<n>We extend our framework to relative scale and demonstrate the efficiency and robustness of our methods on large-scale semidefinite programs.
arXiv Detail & Related papers (2023-06-30T08:34:29Z) - Provably Efficient Reinforcement Learning via Surprise Bound [66.15308700413814]
We propose a provably efficient reinforcement learning algorithm (both computationally and statistically) with general value function approximations.
Our algorithm achieves reasonable regret bounds when applied to both the linear setting and the sparse high-dimensional linear setting.
arXiv Detail & Related papers (2023-02-22T20:21:25Z) - Low Rank Matrix Completion via Robust Alternating Minimization in Nearly Linear Time [8.808780937701522]
We take a major step towards a more efficient and error-robust alternating minimization framework.
Our algorithm runs in time $widetilde O(|Omega| k)$, which is nearly linear in the time to verify the solution.
arXiv Detail & Related papers (2023-02-21T23:49:36Z) - Gaussian Process Bandit Optimization with Few Batches [49.896920704012395]
We introduce a batch algorithm inspired by finite-arm bandit algorithms.
We show that it achieves the cumulative regret upper bound $Oast(sqrtTgamma_T)$ using $O(loglog T)$ batches within time horizon $T$.
In addition, we propose a modified version of our algorithm, and characterize how the regret is impacted by the number of batches.
arXiv Detail & Related papers (2021-10-15T00:54:04Z) - List-Decodable Mean Estimation in Nearly-PCA Time [50.79691056481693]
We study the fundamental task of list-decodable mean estimation in high dimensions.
Our algorithm runs in time $widetildeO(ndk)$ for all $k = O(sqrtd) cup Omega(d)$, where $n$ is the size of the dataset.
A variant of our algorithm has runtime $widetildeO(ndk)$ for all $k$, at the expense of an $O(sqrtlog k)$ factor in the recovery guarantee
arXiv Detail & Related papers (2020-11-19T17:21:37Z) - Streaming Complexity of SVMs [110.63976030971106]
We study the space complexity of solving the bias-regularized SVM problem in the streaming model.
We show that for both problems, for dimensions of $frac1lambdaepsilon$, one can obtain streaming algorithms with spacely smaller than $frac1lambdaepsilon$.
arXiv Detail & Related papers (2020-07-07T17:10:00Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.