The I/O Complexity of Attention, or How Optimal is Flash Attention?

• URL: http://arxiv.org/abs/2402.07443v1
• Date: Mon, 12 Feb 2024 06:50:45 GMT
• Title: The I/O Complexity of Attention, or How Optimal is Flash Attention?
• Abstract summary: 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.
• Score: 4.126403496388015
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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. Given two levels of memory hierarchy, a fast cache (e.g. GPU on-chip SRAM) and a slow memory (e.g. GPU high-bandwidth memory), the I/O complexity measures the number of accesses to memory. FlashAttention computes attention using $\frac{N^2d^2}{M}$ I/O operations where $N$ is the dimension of the attention matrix, $d$ the head-dimension and $M$ the cache size. However, is this I/O complexity optimal? The known lower bound only rules out an I/O complexity of $o(Nd)$ when $M=\Theta(Nd)$, since the output that needs to be written to slow memory is $\Omega(Nd)$. This leads to the main question of our work: Is FlashAttention I/O optimal for all values of $M$? We resolve the above question in its full generality by showing an I/O complexity lower bound that matches the upper bound provided by FlashAttention for any values of $M \geq d^2$ within any constant factors. Further, we give a better algorithm with lower I/O complexity for $M < d^2$, and show that it is optimal as well. Moreover, our lower bounds do not rely on using combinatorial matrix multiplication for computing the attention matrix. We show even if one uses fast matrix multiplication, the above I/O complexity bounds cannot be improved. We do so by introducing a new communication complexity protocol for matrix compression, and connecting communication complexity to I/O complexity. To the best of our knowledge, this is the first work to establish a connection between communication complexity and I/O complexity, and we believe this connection could be of independent interest and will find many more applications in proving I/O complexity lower bounds in the future.

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)
• Approaching I/O-optimality for Approximate Attention [45.67330863443465]
  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.
  arXiv  Detail & Related papers  (2026-05-22T15:23:26Z)
• 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)
• Fast Multipole Attention: A Divide-and-Conquer Attention Mechanism for Long Sequences [1.5484595752241124]
  We present Fast Multipole Attention, a new attention mechanism that uses a divide-and-conquer strategy to reduce the time and memory complexity of attention for sequences of length $n$. The hierarchical approach groups queries, keys, and values into $mathcalO( log n)$ levels of resolution, where groups at greater distances are larger in size and the weights to compute group quantities are learned. We find empirically that the Fast Multipole Transformer performs much better than other efficient transformers in terms of memory size and accuracy.
  arXiv  Detail & Related papers  (2023-10-18T13:40:41Z)
• How to Capture Higher-order Correlations? Generalizing Matrix Softmax Attention to Kronecker Computation [12.853829771559916]
  We study a generalization of attention which captures triple-wise correlations. This generalization is able to solve problems about detecting triple-wise connections that were shown to be impossible for transformers. We show that our construction, algorithms, and lower bounds naturally generalize to higher-order tensors and correlations.
  arXiv  Detail & Related papers  (2023-10-06T07:42:39Z)
• Quadratic Memory is Necessary for Optimal Query Complexity in Convex Optimization: Center-of-Mass is Pareto-Optimal [23.94542304111204]
  We show that quadratic memory is necessary to achieve the optimal oracle complexity for first-order convex optimization. We prove that to minimize $1$-Lipschitz convex functions over the unit ball to $1/d4$ accuracy, any deterministic first-order algorithms using at most $d2-delta$ bits of memory must make $tildeOmega(d1+delta/3)$ queries.
  arXiv  Detail & Related papers  (2023-02-09T22:37:27Z)
• Cumulative Memory Lower Bounds for Randomized and Quantum Computation [1.52292571922932]
  Cumulative memory is a measure of time-space complexity. We prove the first lower bounds on cumulative memory complexity for both sequential classical computation and quantum circuits.
  arXiv  Detail & Related papers  (2023-01-13T17:57:02Z)
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness [80.3586155104237]
  FlashAttention is an IO-aware exact attention algorithm for Transformers. It reduces the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip. FlashAttention and block-sparse FlashAttention enable longer context in Transformers.
  arXiv  Detail & Related papers  (2022-05-27T17:53:09Z)
• The First Optimal Acceleration of High-Order Methods in Smooth Convex Optimization [88.91190483500932]
  We study the fundamental open question of finding the optimal high-order algorithm for solving smooth convex minimization problems. The reason for this is that these algorithms require performing a complex binary procedure, which makes them neither optimal nor practical. We fix this fundamental issue by providing the first algorithm with $mathcalOleft(epsilon-2/(p+1)right) $pth order oracle complexity.
  arXiv  Detail & Related papers  (2022-05-19T16:04:40Z)
• Exponential Separation between Quantum and Classical Ordered Binary Decision Diagrams, Reordering Method and Hierarchies [68.93512627479197]
  We study quantum Ordered Binary Decision Diagrams($OBDD$) model. We prove lower bounds and upper bounds for OBDD with arbitrary order of input variables. We extend hierarchy for read$k$-times Ordered Binary Decision Diagrams ($k$-OBDD$) of width.
  arXiv  Detail & Related papers  (2022-04-22T12:37:56Z)
• Linear Bandit Algorithms with Sublinear Time Complexity [67.21046514005029]
  We propose to accelerate existing linear bandit algorithms to achieve per-step time complexity sublinear in the number of arms $K$. We show that our proposed algorithms can achieve $O(K1-alpha(T))$ per-step complexity for some $alpha(T) > 0$ and $widetilde O(stT)$ regret, where $T$ is the time horizon.
  arXiv  Detail & Related papers  (2021-03-03T22:42:15Z)
• Hybrid Stochastic-Deterministic Minibatch Proximal Gradient: Less-Than-Single-Pass Optimization with Nearly Optimal Generalization [83.80460802169999]
  We show that HSDMPG can attain an $mathcalObig (1/sttnbig)$ which is at the order of excess error on a learning model. For loss factors, we prove that HSDMPG can attain an $mathcalObig (1/sttnbig)$ which is at the order of excess error on a learning model.
  arXiv  Detail & Related papers  (2020-09-18T02:18:44Z)
• Approximate Multiplication of Sparse Matrices with Limited Space [24.517908972536432]
  We develop sparse co-occuring directions, which reduces the time complexity to $widetildeOleft((nnz(X)+nnz(Y))ell+nell2right)$ in expectation. Theoretical analysis reveals that the approximation error of our algorithm is almost the same as that of COD.
  arXiv  Detail & Related papers  (2020-09-08T05:39:19Z)
• An Optimal Separation of Randomized and Quantum Query Complexity [67.19751155411075]
  We prove that for every decision tree, the absolute values of the Fourier coefficients of a given order $ellsqrtbinomdell (1+log n)ell-1,$ sum to at most $cellsqrtbinomdell (1+log n)ell-1,$ where $n$ is the number of variables, $d$ is the tree depth, and $c>0$ is an absolute constant.
  arXiv  Detail & Related papers  (2020-08-24T06:50:57Z)
• 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.