Approaching I/O-optimality for Approximate Attention
Pál András Papp∗
Aleksandros Sobczyk
Anastasios Zouzias
arXiv:2605.23751v1 [cs.LG] 22 May 2026
Computing Systems Lab Huawei Technologies Zurich, Switzerland
Abstract We revisit the I/O complexity of attention in large language models. Given query– key–value matrices Q, K, V ∈ Rn×d , and a machine with fast memory size M , √ the goal is to compute the “attention matrix” A = softmax(QK⊤ / 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 [3, NeurIPS’23]. We also prove corresponding lower bounds in each parameter regime to show that our algorithms are indeed close to I/O-optimal.
1
Introduction
Since its discovery [24], the attention mechanism has reshaped the entire field of machine learning, and, in particular, Large Language Models (LLMs). Given a sequence of tokens of length n, and a feature dimension d, the attention mechanism takes as input three matrices Q ∈ Rn×d (query), K ∈ Rn×d (key), and V ∈ Rn×d (value), and returns the attention matrix: A = Att(Q, K, V) := D−1 exp(M)V,
(1)
where M = √1d QK⊤ , exp(M) is element-wise exponentiation, and D is a diagonal matrix with Pn Di,i = j=1 exp(Mi,j ) (in the literature D−1 exp(M) is often written as softmax( √1d QK⊤ )). The straightforward algorithm, which first computes the matrix QK⊤ using standard matrix multiplication, has an arithmetic complexity of O(n2 d). In the so-called “long-context” regime, where n can grow very large, the quadratic complexity often forms a computational bottleneck. Overcoming the O(n2 ) cost of attention is one of the most interesting computational challenges of the last decade. Indeed, the literature on fast and efficient attention algorithms is quite extensive, where common techniques to reduce the arithmetic intensity include approximation algorithms [10, 14, 7, 17, 25], sparsification [6, 26], or replacing the softmax with a kernel function [16]. A particularly important work on approximate attention is the recent algorithm by Alman and Song [3], which also comes with strong theoretical bounds on the approximation guarantee. The aforementioned works can considerably reduce the arithmetic complexity of attention. However, in modern systems, the computational runtimes also depend on several other aspects. One particularly important factor is the amount of data movements within the memory hierarchy. Due to this, there is a line of research analyzing computations with respect to their I/O complexity or I/O cost, i.e. the data movement they require when executed in a two-level memory hierarchy with a fast memory of ∗ Contacts: [email protected], [email protected], [email protected]
limited capacity M (e.g. cache) and a slow memory of unlimited capacity (e.g. RAM). For attention, if we simply use the I/O-optimal method for standard matrix multiplication, we obtain an I/O cost of 2 ·d ). However, the celebrated FlashAttention algorithm by Dao et al [9], and follow-up works O( n√M 2
2
[8, 22], allow to execute the same algorithm with an I/O cost of only O( n M·d ) by fusing the two multiplications and re-arranging the order of execution. This is indeed an improvement over the naive method when M > d2 , i.e. the cache size is sufficiently large, which is often the case in practice. The work of Saha and Ye [21] also proves tight (conditional) lower bounds for both of these cases up to constant factors, showing that these algorithms (naive matrix multiplication and FlashAttention) are indeed I/O-optimal in the corresponding cache size regimes. In this work, we combine the two lines of research above: we study the most I/O-efficient way to execute attention if we use the state-of-the-art algorithm from [3]. Our results show that this approximate attention algorithm can significantly outperform classical attention not only in terms of computational, but also in terms of I/O complexity. In particular, its asymptotic I/O cost is also significantly lower than that of the FlashAttention algorithm for classical attention. 1.1
Contributions
Our main results are tight and nearly-tight bounds for the I/O complexity of attention when computed using the algorithm of Alman and Song [3]. It turns out to that this I/O complexity can behave rather differently depending on the relations between several parameters: the feature dimension d of the input matrices, the polynomial degree g used to approximate the exponential function in [3], the size r = d+g of the approximation matrices in [3], and the fast memory capacity M . For clarity, the g role of these parameters is summarized in Table 1. We distinguish the following cases: • Case I: when M = Ω(d · r), i.e. the fast memory is so large that it can essentially fit all of intermediate matrix U⊤ 2 V from [3], up to constant factors; • Case II: when M = o(d · r), and also g = o(log M ), i.e. the fast memory is more than exponentially larger than the degree g of the approximating polynomial; √ • Case III: when M = o(d · r), and g = Ω(log M ) but g = o( M ), i.e. the fast memory is between exponential and quadratic in the degree g; √ • Case IV: when M = o(d · r) and g = Ω( M ), i.e. the fast memory is at most quadratic in g. We prove different I/O bounds for these cases. In Case I, when the fast memory essentially fits the intermediate matrix U⊤ 2 V in [3], a simpler strategy produces tight bounds up to a constant factor. Theorem 1.1 (Case I.). If M = Ω(d · r), then the optimal I/O complexity of Approximate Attention is both upper and lower bounded by Θ(n · d). In Case II, when the degree of the approximating polynomial is small (compared to the fast memory size), the upper bounds we obtain are somewhat more complex. Theorem 1.2 (Case II.). If M = o(d · r) and g = o(log M ), then the optimal I/O complexity of Approximate Attention is upper bounded by n · r · d · (4e2 )g O , g M g+1 and lower bounded by n·r·d·g Ω . g M g+1 When the polynomial degree is larger or fast memory capacity is small, the problem becomes even more technical. In these cases we need an additional assumption d ≥ 5g for our lower bounds. √ Theorem 1.3 (Case III.). If M = o(d · r) and Ω(log M ) ≤ g ≤ o( M ), then the optimal I/O complexity of Approximate Attention is upper bounded by n · r · d2 n · r · d O min , √ , M M and when d ≥ 5g, it is also lower bounded by Ω n·r·d·g . M 2
Table 1: Notation and summary of the main parameters used throughout our paper. Parameter n (sequence length) d (feature dimension) M (fast memory capacity) g (degree of approximating polynomial)
Origin
Description
Input problem
Number of rows in Q, K, V.
Input problem
Number of columns in Q, K, V. Often assumed to be O(log n) for long-context attention [3].
Architecture
The amount of available fast memory (i.e. cache size), which determines the optimal I/O cost.
Alman & Song’s algorithm
Degree of the polynomial used to approximate the exponentiation. Depends on the accepted approximation error and the magnitude of the entries. Assumed to be o(log n) in [3].
r (dimension of the approximating matrices)
Alman & Song’s algorithm
The maximal number of terms in a polynomial of degree g on d variables. Its value equals d+g . This determines the g number of columns in the approximation matrices U1 and U2 .
w (generating set size)
Our I/O analysis
The number of entries we decide to load from slow memory in a row of Q (or K), used in our proofs. If w ≤ d, this allows to compute at most w+g entries in a row of U1 (or U2 ). g
Finally, for tiny cache sizes M = O(g 2 ), the problem behaves like standard matrix multiplication. √ Theorem 1.4 (Case IV.). If M = o(d·r) and g = Ω( M ), the optimal I/O complexity of Approximate √ Attention is upper and lower bounded by Θ n·r·d , where the lower bound also requires d ≥ 5g. M These theorems provide a comprehensive overview of the optimal I/O complexity of Approximate Attention. From a broader perspective, our results can also be understood as an I/O analysis of matrix multiplication in a special setting where one of the input matrices is sparse, or behaves unusually in some way regarding I/O. We believe that out proof techniques could also inspire further results in sparse attention computation or in entirely different domains. 1.2
Parameters and discussion
We discuss the bounds in more detail with some example parameters for clarity. For the polynomial degree g, Alman and Song consider g = o(log n) [3]; together with d = O(log n), this is already enough to establish r = no(1) in their analysis. However, in practice, it is more realistic to have the degree g even smaller, e.g. a constant. This is indeed a valid choice of g in the approximation setting of [3] when the approximation accuracy ϵ and the magnitude B of the entries are constants, i.e., they do not grow as a function of n, which is reasonable for long–context attention. We point out that if g = O(1), then the upper and lower bounds in Theorems 1.2 and 1.3 are also tight up to a constant factor. This is easy to see for Theorem 1.2. In Theorem 1.3, the gap between √ the bounds is at most gM ; with g = Ω(log M ), we can upper bound this by 2O(g) , which is again a constant. Theorems 1.1 and 1.4 are less interesting from this angle: these bounds are always tight. In order to compare the bounds to FlashAttention, let us consider d2 < M < n · d; this is the range where FlashAttention improves upon naive matrix multiplication [21]. Recall that the parameter choices in [3] ensure that r = no(1) . The condition of Case I is M = Ω(d·r); √ hence we are in this setting if M is larger than no(1) , e.g. if M = n or M = n0.01 . Our upper bound 2 2 here matches the trivial lower bound of Ω(n · d), while FlashAttention incurs a cost of Θ( n M·d ) [21]. Thus our bounds indeed outperform FlashAttention for M = o(n · d). The factor of difference n·d M between our method and FlashAttention becomes even more significant as n grows. For the remaining cases, the bounds we obtain are again almost-linear in n: for Cases II, III and IV, respectively, we get upper bounds of 1+o(1) 1+o(1) 2 1+o(1) n ·d n · d n1+o(1) · d n ·d √ √ O , O min , and O . g g+1 M M M M 3
In particular, the first two of these expressions are at most O 2
1+o(1) n
M
·d2
in the given parameter
2
regimes. FlashAttention has a very similar I/O cost of Θ( n M·d ), which depends quadratically on n; hence compared to it, our algorithms save at least an almost-linear factor of I/O cost in n.
2
Preliminaries
2.1
Approximate attention
Many existing works on attention algorithms target the exact computation of the attention matrix. While this is not necessarily “unreasonable”, it raises some intricate questions. In particular, if attention can be computed exactly in T steps, then exp(x) can be also computed exactly in T + O(1) steps (see Appendix D.1). However, assuming that the entire sequence of digits of exp(x) can be returned in a single “time-step” might raise concerns regarding the machine model; see e.g. the discussion in [13]. To avoid such intricacies, we target approximate attention algorithms, rather than exact ones. We consider the following approximate attention definition from [3]. Problem 2.1 (Additive-error Approximate Attention [3]). Given ϵ > 0, and matrices Q, K, V ∈ e which satisfies: Rn×d , return a matrix A e i,j − Att(Q, K, V)i,j ≤ ϵ. max A i,j
This specific definition has proven useful to obtain sharp bounds on the arithmetic complexity of p attention. When d = O(log(n)) and Q, K, V have entries of magnitude o( log(n)), the authors of [3] describe an algorithm with nearly-linear n1+o(1) complexity in the algebraic model; we describe this in detail in Section p 2.2. On the other hand, the authors also prove that if the magnitude of entries in Q, K, V is Θ( log(n)), then no algorithm can solve Problem 2.1 up to 1/ poly(n) error in sub-quadratic time under the Strong Exponential Time Hypothesis (SETH). 2.2
The algorithm of Alman and Song
In brief, the algorithm of Alman and Song [3] that solves 2.1 works as follows: 1. Compute the coefficients of a degree-g polynomial P (x), which approximates exp(x) in the domain [−B, B], where B is a bound on the magnitudes of the elements of Q, K. ⊤ n×r 2. Construct two for some r = no(1) , such that U1 U⊤ 2 = P (QK ) ≈ √ matrices U1 , U2 ∈ thR ⊤ th exp(QK / d). Intuitively, the i row of U1 (i row of U2 , respectively) is obtained from the ith row of Q (ith row of K), via an expansion of the polynomial P (·).
e ←D e −1 (U1 (U⊤ V)), where D e = diag(U1 (U⊤ 1)). 3. Return A 2 2 The key idea that makes this algorithm efficient is that the polynomial approximation completely removes the non-linearity of the softmax. This allows to execute the matrix multiplication U⊤ 2 V first, ⊤ and then U1 (U2 V) afterwards, which altogether needs significantly fewer arithmetic operations. Going in more detail, in step 2 we have that exp(QK⊤ )i,j = exp(q⊤ k), where q⊤ is the ith row of Q and k is the j th column of K⊤ . When exp(x) is replaced by the polynomial P (x), we can write: P (q⊤ k) = P (q1 k1 + q2 k2 + . . . + qd kd ) =
g X
cl (q1 k1 + q2 k2 + . . . + qd kd )l .
l=0
If we expand (q1 k1 +q2 k2 +. . .+qd kd ) for any l ∈ {1, . . . , g}, we obtain a sum with d+l−1 additive l l1 l2 ld terms. Each additive term is of the form (q1 k1 ) · (q2 k2 ) · . . . · (qd kd ) , where l1 + l2 + . . . + ld = l. Pg Altogether, we get r = l=0 d+l−1 additive terms. Ultimately, since each term is the product l of some powers of the elements of q and k, the polynomial P (q⊤ k) can be written as a bilinear form P (q⊤ k) = u⊤ Cv, where u ∈ Rr consists of products of (powers of) the elements of q, while v ∈ Rr consists of products of (powers of) the elements of k, and C ∈ Rr×r is a diagonal matrix with scalar coefficients that depend on the polynomial P (·), but are independent from n and d. l
4
n o log(1/ϵ) 2 To ensure |P (x)−exp(x)| ≤ ϵ·exp(x), setting g := Θ max log(log(1/ϵ)/B is sufficient [1, 2) , B 3]. The coefficients of P (·) can be computed as poly(g)-bit rationals in poly(g) time. Below we introduce a specific notation for the number of possible terms in a polynomial of degree g on w variables; this will play a central role in our analysis. Definition 2.1. Let us define the function τ : Z+ → Z+ as a shorthand notation for g X w+l−1 w+g τ (w) := = . l g l=0
w+l−1
Intuitively, is a combination with repetition, i.e. the number of ways we can divide a total l degree of g among w variables. The number of columns in U1 and U2 is r = τ (d). The closed form expression of the sum in the definition can be shown through an induction; we defer the proof to 2(w+g) Appendix D. We note that the work of [3] uses a looser bound of to bound r. 2g Note that the algorithm above executes matrix multiplications with a standard inner product-based algorithm. Previous works on the I/O analysis of attention also focus on standard matrix multiplication [21]. It is well-known that “Strassen-like”, fast matrix multiplication algorithms (see e.g. [2, 4] and references therein) can achieve better computational and I/O complexity [11]. However, these algorithms are often known to be rather impractical, and their analysis is more complex than the standard algorithm, especially when rectangular matrices are involved. Throughout the paper, we analyze I/O-efficient algorithms based on the aforementioned Approximate Attention method. Note that our upper bounds naturally carry over to solving Problem 2.1 in general. On the other hand, lower bounds in standard I/O-complexity models are inherently coupled to a concrete algorithm, so our bounds here are specific to this Approximate Attention technique. 2.3
Model of computation
The most prominent model to analyze the I/O complexity of computations is the red-blue pebble game of Hong and Kung [15]. In this setting, a computation is captured as a directed acyclic graph (DAG), where the nodes represent operations, and the directed edges represent data dependencies, i.e. that the output of an operation is required as an input for another operation. The computational DAG corresponding to the approximate attention algorithm from [3] is illustrated in Figure 1. Note that the algorithm only uses basic algebraic operators {+, −, ×, /}, and square roots, hence avoiding the intricacies of exact exponentiation that we mentioned in Section 2.1. The red-blue pebble game models the execution of the computation in a two-level memory hierarchy: we have a fast memory with limited capacity M , and slow memory of unlimited capacity. If the output of an operation is currently stored in fast (slow) memory, then this is indicated by having a red (blue) pebble on the node. The number of red pebbles can not exceed the cache capacity M at any time. Loading a value from (respectively, saving a value to) slow memory then corresponds to placing a red (blue) pebble on a node that already has a blue (red) pebble. A new output value can only be computed (placing a red pebble on the node) if all the parents of the node have a red pebble. The execution of the computation is modeled by a sequence of steps, where each step can load, save, compute or delete a value. Initially, the source nodes of the DAG have a blue pebble, and the process finishes when all the sink nodes have a blue pebble. Since the goal of the model is to capture I/O complexity, the cost of the pebbling sequence is the total number of save and load operations executed; compute and delete steps are considered free. Red-blue pebbling has been thoroughly analyzed in several works, and widely used to establish upper and lower bounds on the I/O cost of specific computations [15, 19, 23, 12, 20, 18, 5].
3
The I/O Cost of Approximate Attention
In this section, we outline the main intuition behind the bounds in Theorems1.1-1.4, and the main ingredients of the proofs. The detailed proofs are deferred to Appendices A, B and C. Our computation U1 (U⊤ 2 V) consists of two consecutive matrix multiplications; we denote the intermediate matrix by H := U⊤ 2 V. Throughout the proofs, we focus on the second multiplication 5
d K⊤
U⊤ 2
r
V
H
r
Q
U1
t1
H
n
t 3 U1
t2
output
Figure 2: Illustration of matrix multiplication with tiles, for the example U1 H. The output matrix is split to tiles of size t1 × t2 , and each of these is aggregated in tiles of width t3 (respectively, height t3 ) in the corresponding row strip of U1 (column strip of H). Altogether, the output matrix is computed via tn1 · td2 · tr3 tiling steps.
Figure 1: Sketch of the computational DAG for the approximate attention algorithm [3]. Each box represents a set of nodes. In matrix multiplication, we first multiply pairs of entries from the input matrices, and then sum these up to from an entry in the output. In U1 (and U2 ), each node is computed from at most g nodes of Q (and K).
U1 H; note that U1 ∈ Rn×r and H ∈ Rr×d . The lower and upper bounds are then easy to extend to the entire computation U1 (U⊤ 2 V). Considering the sub-problem U1 H separately is in fact a restriction of I/O strategies: the entries in H here are all inputs that need to be loaded from slow memory. This rules out “fusion-based” strategies that never save the entries of H, but instead immediately go on to multiply them with U1 . Intuitively, these strategies can be ignored because H is smaller than Q, K and V, so saving/loading it does not affect the magnitude of the total cost. 3.1
Intuition: optimal tile design
In standard matrix multiplication, the key to optimal I/O strategies is to divide the matrix into rectangular tiles to maximize reusing the same inputs and outputs. In general, when multiplying m2 ×m3 matrices A ∈ Rm1 ×m2 and B √∈R √ , the optimal strategy uses square tiles. That √is, we read √ the inputs in a square-shaped Θ( M )×Θ( M ) tile in A, and another square-shaped Θ( M )×Θ( M ) tile in B, and we execute the part of the matrix multiplication. This gives us a partial √ corresponding √ sum for each entry in a Θ( M ) × Θ( M ) tile of the output matrix, which we can add (and save) to the already aggregated values for this specific output entry. Each such tiling √ step only requires a fast memory capacity of Θ(M ), it has an I/O cost of only Θ(M ), and covers Θ( M )3 = Θ(M 3/2 ) of the scalar multiplications that compose the entire operation. With a total of m multiplications 1 · m2 · m3 √ 2 ·m3 . This general in the entire operation, this allows for an algorithm with an I/O cost of Θ m1 ·m M tiling idea is illustrated in Figure 2 for the matrix multiplication U1 H. Intuitively speaking, this tiling is optimal for standard matrix multiplication because it allows to execute the highest number Θ(M 3/2 ) of the scalar multiplications while the tile sizes in the input/ output matrices is still Θ(M ). The best such tiling is naturally obtained when we use square-shaped tiles in the inputs and the output. However, with our special matrices U1 and U2 , the situation is significantly different. The entries in a given row of U1 are not inputs, but values computed from the corresponding row of Q; in fact, by loading only w entries from the ith row of Q (for some w ≤ d), we can compute up to τ (w) = g+w different entries in the ith row of U1 . w 6
This means that in order to form an essentially square-shaped tile in U1 with Θ(M ) load operations, we actually need to solve the following equation for w: τ (w) =
Θ(M ) . w
(2)
) ∗ Let us denote the solution of this by w∗ . We can then select Θ(M w∗ rows of U1 , and load w values in ∗ the corresponding rows of Q, thus generating up to τ (w ) entries in each row of U1 .
Forming tiles of this shape is the key idea behind our results. Intuitively, Case II marks the general case where this choice of w∗ is indeed possible. The other settings can be be understood as special cases. In Case I, we essentially have w∗ ≥ d; since we cannot load more than d values in a row of Q, we can only form ‘taller’ rectangular tiles. In Cases III and IV, we have w∗ ≤ g, but we still need to load g values from Q in many cases because many degree-g terms in the polynomial actually contain g different variables; hence we need to form ‘wider’ rectangular tiles. Note that obtaining a closed-form expression for w∗ from Equation (2) is non-trivial; instead, we rely b b on standard bounds on binomial coefficients. Specifically, we use the fact that ab ≤ ab ≤ e·a . b 3.2
Generic upper bounds √
√
Note that if we naively apply the standard matrix multiplication method with tiles of size 4M × 4M √ √ ). on our computation, we get an algorithm of I/O cost O( n·r·d This is indeed viable when M ≥ d, √M i.e. each of the matrices can fit the square tiles. When M < d, we can instead form tiles of size M M M 4d × d in both H and the output matrix, and tiles of size 4d × 4d in U1 (assuming that these tiles M M fit, i.e. we have 4d ≤ r). Each of these tiles only needs 4 I/O operations, because in U1 , we can obtain all the entries of a row with only d I/O steps. The number of tiling steps is altogether nrd2 n r nrd2 M · M = O( M 2 ). With an I/O cost of O(M ) per tile, we altogether get O( M ) I/O steps. 4d 4d √ 2 √ ≤ nrd Note that n·r·d M ≥ d, so the two bounds above can be easily combined. M exactly when M This gives a generic upper bound that holds over Cases II, III and IV. Lemma 3.1. If M = o(d · r), the optimal I/O cost of Approximate Attention is upper bounded by n · r · d2 n · r · d O min , √ . M M 3.3
Case I: proof of Theorem 1.1
In Case I, we have d·τ (d) = d·r = O(M ), i.e. there is a constant 0 < c0 < 1 such that c0 ·d·r ≤ 14 M . M This means that we can select 4·d rows from U1 , load all the d values in the corresponding rows of 1 Q (at an I/O cost of 4 M ), and generate all the r entries in each of these rows in U1 . M For the computation, we can then split the output matrix into tiles of height 4·d and width c0 · d. We M can compute each tile in a single iteration: we create the corresponding 4·d rows in U1 as discussed, we load the tile of shape r × (c0 · d) from H that needs to multiply it (at an input cost of at most 1 1 4 M ), and then output the final result on the tile (output cost of at most 4 M ). In each tile, all the loaded entries in Q and H and the aggregated values in the output matrix can be kept in fast memory simultaneously (since 34 M < M ) for the entire duration of computing the tile. On the other hand, the entries of U1 are always recomputed from the entries in Q, without requiring any new I/O steps. The n number of tiles needed to cover the output matrix is M/(4d) · c0d·d = O( n·d M ), altogether resulting in an I/O cost of O(n · d).
The first part of the computation, i.e. the other matrix multiplication H = U⊤ 2 V, can be executed in a similar way at the same I/O cost. Here d entries in a row of K allow us to generate the entire ⊤ M M column of U⊤ 2 , so we can form tiles of shape r × 4·d in U2 , 4·d × (c0 · d) in V and r × (c0 · d) in n H, aggregating each tile in the output matrix in M/(4d) steps. Note that the corresponding lower bound is straightforward: each entry of the input matrices Q, K and V must be loaded at least once, giving an I/O cost of at least 3 · n · d. 7
3.4
Case II: when g ≤ w∗ < d
In Case II, the solution to Equation (2) satisfies g ≤ w∗ < d. The best strategy here, as discussed ) above, is to tile U1 by loading w = w∗ values from Θ(M distinct rows of Q. Consider a specific w row q of Q, and let u be the corresponding row in U1 . Hypothetically, if every set of w loaded values from q could generate τ (w) distinct entries in u, then the strategy would provide matching upper and lower bounds. In this case, the corresponding tiles in H and U1 H could both have height Θ(M ) and w width Θ(w), so all inputs and outputs in a tile would fit into fast memory. This altogether gives Θ
n
M w
· wd tiles in the output matrix, each aggregated from Θ( Mr ) separate column strips in w
rw ndrw U1 and row strips in H. The total number of tiling steps would be Θ( nd M · M ) = Θ( M 2 ). With ndrw Θ(M ) I/O steps per tile, the total cost is Θ( M ). In Appendix A, we show that no I/O strategy can be more efficient than this. The proof analyzes the so-called S-partitions of the computational graph, which is a long-established tool to derive I/O lower bounds on different computations [15, 19].
In contrast to this hypothetical setting, the w loaded values cannot generate fully-disjoint sets of entries from u in each tile. To ensure that every entry of u is obtainable in at least one tile, the entries of q actually need to be loaded in multiple different tiles. This means that there will be many entries in u that can be generated in multiple tiles. We can assign each such entry of u to an arbitrary one of these tiles, but this still implies that the width of the average tile in U1 will be lower than τ (w). As such, we need to select different subsets of w values each from q, and assign each entry in u to one of these subsets that can generate it. For this, we divide the d entries of q into d·g w groups of size d·g w w possible combinations of g different groups as generator sets for a g each, and consider all the g tile. All entries in u can be obtained in at least one combination, since they are generated from at most g entries in q. For entries in u that can be generated in multiple tiles, we simply assign it to the first such tile, hence some of our tiles in U1 will be wider than others. d·g Altogether, each tile of the output matrix is aggregated through wg steps. With tiles of size Θ(M ) × Θ(w) in an output matrix of dimensions n × I/O w d, and an cost of O(M ) per tile, this gives d·g dg d n w w a total I/O cost of Θ(M ) · Θ(w) · g · O(M ) = O n · d · g . w
Formally, the observations above can be expressed through the following lemma. Key Lemma 3.2. Let us choose an integer w such that g ≤ w ≤ d and w · τ (w) ≤ 14 M . Then the optimal I/O cost of Approximate Attention in Case II is upper bounded by dg ! O n·d· w , g and lower bounded by Ω
n·r·d·w M
.
The proof is provided in Appendix A. Theorem 1.2 then follows from this more general statement. Proof of Theorem 1.2. Given Key Lemma 3.2, the bounds in Theorem 1.2 are obtained by choosing 1 1 · g · M g+1 4·e for w. Note that this w0 is only slightly different from the real w∗ , due to only approximating the binomial coefficient in τ (w). First note that w0 ≥ g; this is equivalent to M ≥ (4e)g+1 , which indeed holds if g = o(log M ). We then show that w0 satisfies w0 · τ (w0 ) ≤ 41 M . Indeed, we have g w0 + g 2 · e · w0 w0 · τ (w0 ) = w0 · ≤ w0 · g g
w0 =
using the upper bound on the binomial coefficient, and w0 ≥ g. Substituting w0 into this, we get g 1 1 g 4e · 2g · M . This is indeed smaller than 4 M , since g < 2 for any positive integer g. From this, it also follows that w0 ≤ d, since in Case II, we have d · τ (d) = ω(M ). 8
Using w = w0 in the lower bound of Key Lemma 3.2 directly gives our lower bound in Theorem 1.2: n·r·d n·r·d·g Ω · w0 = Ω . g M M g+1 Similarly, the upper bound of Key Lemma 3.2 provides the upper bound in Theorem 1.2 if we dg d+g g substitute w = w0 , and use wg ≤ ( edg ≥ ( dg )g : wg ) and r = g ! dg ! g g 2 g g e · d n · d · (4e ) · d n · r · d · (4e2 )g w 0 = O n·d· = O = O . O n·d· g g g w0g g g · M g+1 M g+1 3.5
Case III: when w∗ ≤ g
In Case III, the optimal w∗ for square tiles would be smaller than g. However, many of the entries in U1 still require at least g values to generate. The best upper bounds here are mostly obtained simply via the generic bounds in Lemma 3.1; this is also what we find in Theorem 1.3. Recall from the introduction that for g = O(1), this is tight to the lower bound up to a constant factor. Note that even 2 for g = ω(1), the bounds are also tight in some cases: for instance, if we have d = O(g), then n·r·d M again matches the lower bound up to a constant factor. On the other hand, when d = Ω(g 2 ), one can show that the gap between the bounds in Theorem 1.3 is at least g = ω(1). The regime between d = O(g) and d = Ω(g 2 ) is more challenging. With some further work, we can actually extend the tight upper bounds to d being almost quadratic in g. The proof of this is more technical; it essentially adapts the upper bound idea from Case II with a choice of w = g, and some further adjustments. It also requires the extra assumption that M is polynomial in g. Lemma 3.3. In Case III, if we also assume that d = O(g 2−δ1 ) and M = O(g δ2 ) for some constants δ1 , δ2 > 0, then the optimal I/O complexity of Approximate Attention is upper bounded by O( n·r·d·g M ). The lower bound in Theorem 1.3 uses the same tools (S-partitions) as in Case II, but requires a more complex proof. Intuitively, the terms that can be generated from less than g variables have a much larger role here. Due to this, we add an extra assumption here that d ≥ 5g, which is indeed realistic if e.g. g = O(1). With d ≥ 5g, a technical lemma shows that at least a constant fraction of the columns in U1 use at least g2 different variables. We can then restrict our analysis to this sub-matrix where any term requires loading g2 values in Q, and use a similar proof to Case II. The formal proofs of Lemma 3.3 and the lower bound in Theorem 1.3 are available in Appendix B. 3.6
Case IV
Finally, for tiny cache sizes M = O(g 2 ), we actually have another special case. Intuitively, here it becomes more beneficial to store the values of U1 than to recompute them; this provides tight bounds n·d·r of Θ √M , similarly to standard matrix multiplication. These proofs are discussed in Appendix C.
4
Discussion and Conclusion
In this work we studied the I/O complexity of Approximate Attention. We provided sharp upper and lower bounds for different parameter regimes, and showed that this approach can have significantly lower I/O cost than state-of-the-art methods like FlashAttention. Our proof techniques may also be used to derive I/O bounds in other settings with special kinds of matrix multiplication. For instance, consider a sparse attention approach where the attention matrix is guaranteed to have at most O(d) non-zero entries in each row: this can be addressed in a very similar way to Lemma 3.1. Other, more complex use cases may have matrices with dynamically generated values, where the I/O analysis may be conducted with tools similar to Cases II-IV. While the theoretical bounds are promising to improve the I/O cost of attention algorithms, it is wellknown that the theoretically-fastest algorithms are not always the best-performing ones in practice. For example, the combinatorial nature of the parameter r (which originates from the analysis of [3]) can be prohibitively large for some practical parameter regimes. It is an interesting direction for future work to also investigate the efficiency and applicability of our algorithms in practice. 9
References [1] Amol Aggarwal and Josh Alman. Optimal-degree polynomial approximations for exponentials and gaussian kernel density estimation. In 37th Computational Complexity Conference (CCC), page 1, 2022. [2] Josh Alman, Ran Duan, Virginia Vassilevska Williams, Yinzhan Xu, Zixuan Xu, and Renfei Zhou. More asymmetry yields faster matrix multiplication. In Proceedings of the 2025 Annual ACM-SIAM Symposium on Discrete Algorithms (SODA), pages 2005–2039. SIAM, 2025. [3] Josh Alman and Zhao Song. Fast attention requires bounded entries. Advances in Neural Information Processing Systems (NeurIPS), 36:63117–63135, 2023. [4] Josh Alman and Hantao Yu. Improving the leading constant of matrix multiplication. In Proceedings of the 2025 Annual ACM-SIAM Symposium on Discrete Algorithms (SODA), pages 1933–1971. SIAM, 2025. [5] Toni Böhnlein, Pál András Papp, and Albert-Jan N. Yzelman. Red-blue pebbling with multiple processors: Time, communication and memory trade-offs. In International Colloquium on Structural Information and Communication Complexity (SIROCCO), pages 109–126. Springer, 2025. [6] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019. [7] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference on Learning Representations (ICLR), 2021. [8] Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In International Conference on Learning Representations (ICLR), 2024. [9] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in neural information processing systems (NeurIPS), 35:16344–16359, 2022. [10] Giannis Daras, Nikita Kitaev, Augustus Odena, and Alexandros G Dimakis. Smyrf-efficient attention using asymmetric clustering. Advances in Neural Information Processing Systems (NeurIPS), 33:6476–6489, 2020. [11] Erik D Demaine, Andrea Lincoln, Quanquan C Liu, Jayson Lynch, and Virginia Vassilevska Williams. Fine-grained i/o complexity via reductions: New lower bounds, faster algorithms, and a time hierarchy. In 9th Innovations in Theoretical Computer Science Conference (ITCS), pages 34–1. Schloss Dagstuhl–Leibniz-Zentrum für Informatik, 2018. [12] Erik D Demaine and Quanquan C Liu. Red-blue pebble game: Complexity of computing the trade-off between cache size and memory transfers. In Proceedings of the 30th ACM Symposium on Parallelism in Algorithms and Architectures (SPAA), pages 195–204, 2018. [13] Jeff Erickson, Ivor Van Der Hoog, and Tillmann Miltzow. Smoothing the gap between np and er. SIAM Journal on Computing, 53(6):FOCS20–102, 2022. [14] Insu Han, Rajesh Jayaram, Amin Karbasi, Vahab Mirrokni, David Woodruff, and Amir Zandieh. Hyperattention: Long-context attention in near-linear time. In International Conference on Learning Representations (ICLR), 2024. [15] Jia-Wei Hong and Hsiang-Tsung Kung. I/O complexity: The red-blue pebble game. In Proceedings of the thirteenth annual ACM symposium on Theory of computing (STOC), pages 326–333, 1981. [16] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning (ICML), pages 5156–5165. PMLR, 2020. 10
[17] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations (ICLR), 2020. [18] Grzegorz Kwasniewski, Marko Kabić, Maciej Besta, Joost VandeVondele, Raffaele Solcà, and Torsten Hoefler. Red-blue pebbling revisited: near optimal parallel matrix-matrix multiplication. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC), pages 1–22, 2019. [19] Pál András Papp, Aleksandros Sobczyk, and Albert-Jan N Yzelman. The impact of partial computations on the red-blue pebble game. In Proceedings of the 37th ACM Symposium on Parallelism in Algorithms and Architectures (SPAA), pages 328–338, 2025. [20] Pál András Papp and Roger Wattenhofer. On the hardness of red-blue pebble games. In Proceedings of the 32nd ACM Symposium on Parallelism in Algorithms and Architectures (SPAA), pages 419–429, 2020. [21] Barna Saha and Christopher Ye. I/O complexity of attention, or how optimal is flashattention? In Proceedings of the 41st International Conference on Machine Learning (ICML), pages 43024–43042, 2024. [22] Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. Advances in Neural Information Processing Systems (NeurIPS), 37:68658–68685, 2024. [23] Aleksandros Sobczyk. I/o complexity and pebble games with partial computations. Information Processing Letters, page 106637, 2026. [24] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems (NeurIPS), 30, 2017. [25] Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström-based algorithm for approximating self-attention. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pages 14138–14148, 2021. [26] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems (NeurIPS), 33:17283– 17297, 2020.
A
Proof of Key Lemma 3.2
A.1
Upper bound proof
The upper bound proof follows the idea outlined before. We split the output matrix into tiles of height M M 4w and width w. This already specifies a row strip of size 4w × r in U1 and a column strip of size d·g r × w in H that are needed to compute this output tile. We aggregate the output tile in wg steps, splitting the row strip of U1 horizontally (and the column strip of H vertically) into aggregation tiles. The aggregation tiles are obtained as follows. We split the column indices {1, . . . , d} of Q into groups of size wg : let G1 = {1, . . . , wg }, G2 = { wg +1, . . . , 2· wg }, . . ., G dg = {d− wg +1, . . . , d}. w Each aggregation tile is formed by selecting g distinct groups, i.e. the number of aggregation tiles is dg w . In each aggregation tile, we consider the union of the g groups combined, and consider the g M resulting w distinct column indices; specifically, for each of the corresponding 4w rows, we load the 1 values from these w columns of Q. This corresponds to 4 M load operations for each aggregation tile. Note that some of the terms in the polynomial require an entry from g distinct groups; these columns of U1 can only be generated on the specific tile that combines these g groups. Since any term contains 11
at most g variables, each entry of U1 can be generated in at least one tile. However, there will also be numerous terms that can be generated in multiple different tiles: e.g. a term that only uses variable indices in G1 and G2 can be assigned to any of the tiles that combine G1 , G2 and a selection of (g − 2) arbitrary further groups. For simplicity, we will always assign each term to the first tile that can generate it. That is, we iterate through the group subsets via a lexicographic order of their sorted indices, and we assign each term to the first tile that can generate it in this order. Note that for any term, it is straightforward to determine the tile that generates it: we consider the groups that contain at least one variable of this term, and if the number of groups is smaller than g, we use the smallest non-included group indices for the rest of the groups. As such, for any tile and term, it is easy to decide whether the term is assigned to this given tile or not. This means that for any given tile, we can iterate over all the τ (w) terms that can be generated from the corresponding w inputs, and generate only the terms that belong to this specific tile, avoiding duplications. With this method, the width of each tile in U1 will be at most τ (w) (the maximum number of terms we can generate from w variables), and at least ( wg )g (the terms using g distinct groups that can only be obtained on this tile). We point out that in Case II, the aggregation tiles we form are always taller than wider. Indeed the M M height of the tiles is 4w , and their width is at most τ (w). We have 4w ≥ τ (w) due to w · τ (w) ≤ 14 M . This is crucial to ensure that the corresponding tiles formed in H indeed fit into fast memory: since M their height is at most τ (w) ≤ 4w , and their width is w, they contain at most 41 M entries. Altogether, each aggregation tile requires at most 14 M loaded inputs from both U1 and H, which are always kept in fast memory. The results are aggregated in a tile of size 14 M in the output matrix, which are written to slow memory in the end. As such, each sub-computation fits into the working n · wd , and memory of M , and has I/O cost of O(M ). The number of tiles in the output matrix is M 4w dg each are aggregated in wg steps. This results in a total I/O cost of dg ! dg 4nw d . · · w ·M =O n·d· w g g M w We note that in the red-blue pebbling model, generating the entries in U1 requires no further working memory; the process of iterating through the entries is incorporated into the pebbling strategy. However, in a practical implementation, this requires a few further fast memory entries, e.g. index variables to iterate through the τ (w) possible combinations. We also note that when aggregating the results in a tile of size 14 M above, this actually incurs an I/O cost of 12 M , i.e. two I/O steps for each entry of the matrix (in the appropriate red-blue pebbling variant; see Appendix A.3 for a discussion). Indeed, we first need to load the current aggregated value for the entry in the output matrix, add the newly computed terms to this value, and then save the new (partial) aggregated value to slow memory again. However, this is simply a clarification of a technical detail; since this cost is still in O(M ), it has no effect on the analysis of I/O costs above. Extending the algorithm to the entire computation U1 (U⊤ 2 V) is straightforward. We simply execute the two multiplications separately, considering H = U⊤ V as an output in the first and as an input in 2 the second. In the first computation, we use the same tiling strategy as discussed above, just with M the role of the matrices exchanged. That is, we form column strips of width 4w in U⊤ 2 and the dg M corresponding row strips of height 4w in V. We split each of these strips to wg tiles (of differing height) in U⊤ 2 . This results in tiles of width w and height at most τ (w) in H. Each such tile will be dg d n w aggregated in M steps from the different column strips in U⊤ 2 and row strips in V. With g · w 4w tiles in H, the total I/O cost is the same as in the second matrix multiplication. We point out that in almost all cases, the upper bound for Theorem 1.2 (obtained via Key Lemma 3.2) is tighter than the generic upper bound from Lemma 3.1. In particular, when M = O(d2 ) and √ hence the generic upper bound is O( n·r·d ), then our upper bound in Theorem 1.2 is smaller when M 1
1
1
M 2 − g+1 ≥ (4e2 )g holds. For any g ≥ 2, the left side here is lower bounded by M 6 , and then the claim indeed follows from g = o(log M ). On the other hand, when M = ω(d2 ), and hence 1 2 g+1 · (4e2 )g . the generic upper bound is O( n·r·d M ), then our upper bound is smaller when d ≥ M This also holds if we have any constant δ > 0 such that M ≤ dδ . Indeed, then (4e2 )g can be 12
1
δ
δ
upper bounded by M 2 (because g = o(log M )) and M g+1 can also be upper bounded by M 2 asymptotically (assuming that g = ω(1); otherwise, the upper and lower bounds are tight anyway). The generic upper bound may only be superior for very specific parameter combinations, e.g. if we g+1 . have M = dg A.2
Lower bound proof
The lower bound proof corresponds to the hypothetical case above where each set of w entries from q generates τ (w) distinct entries in u, and hence it is likely somewhat lower than the actual optimum. When proving lower bounds for I/O cost, one of the most widely used tools are so-called S-partitions. Defined by Hong and Kung in their original paper on red-blue pebbling and I/O complexity [15], an S-partition is a way to partition the node of the computational graph such that, intuitively, each class can be computed with only 2S I/O steps. Definition A.1. An S-partition of a computational DAG (for some integer parameter S) is a disjoint partitioning D1 , ..., Dk of the vertices of the DAG such that • the classes D1 , ..., Dk form a topological order of the computation, i.e. there are no edges from Dj to Di with i < j; • for each class Di , there is a so-called dominator set of size S in the DAG, i.e. a set of nodes such that every path from a source node to a node in Di contains a node in this set; • for each class Di , the so-called minimum set of Di , i.e. the set of nodes in Di without a child in Di , has size at most S. For more details on S-partitions, we refer the reader to [15]. The crucial property of S-partitions is that they directly allow to lower bound the optimal I/O cost when using S = 2M . Lemma A.1 (From Hong & Kung [15]). If MIN2M denotes the minimal number of classes in any 2M -partition of a computational DAG, then the optimal I/O cost of this computation is at least M · (MIN2M − 1). In the lower bound proof, we analyze the number of classes that each 2M -partition must have in our computation. Through the lemma above, this directly yields the lower bound in Key Lemma 3.2. Proof of Key Lemma 3.2, lower bound. Consider first the computational DAG restricted to the multiplication U1 H, i.e. where the entries of H are all source nodes. Let us refer to the nodes that correspond to (the output of) multiplying an entry of U1 and an entry of H as internal nodes. The DAG altogether has n · r · d internal nodes. Similarly to previous I/O lower bounds on matrix multiplication variants, the main idea of the proof is to upper bound the number of internal nodes that can be contained in each partition, thus obtaining a lower bound on MIN2M . Let Di be any partition in a 2M -partition. Consider the minimum set of Di and a dominator set of size at most 2M for Di . This is altogether at most 4M nodes; even if they are all internal, this gives at most 4M internal nodes. Let Γ be the set of all other internal nodes contained in Di ; let us call these uncovered internal nodes. For every uncovered node v ∈ Γ, we must have that (i) on each directed path from a source node to v, we have a node in the dominator set, and (ii) at least one successors of v is in the minimum set. Note that each internal node has exactly one parent in H that is a source node, so for uncovered nodes, these must all be in the dominator set. An uncovered node v ∈ Γ also has a parent in U1 ; however, here there are two options, either the parent in U1 is in the dominator set, or all the parents of the parent (located in Q) are in the dominator set. As for the successor of v in the minimum set, this may be the sink node in the output matrix that corresponds to v, or any node in the summation tree from v to the sink. For any i ∈ [n], let wi (Q) and wi (U1 ) , respectively, denote the number of entries in the ith row of Q and the ith row of U1 that are in the dominator set. Let wi = wi (Q) + wi (U1 ) . Let us consider the number of nodes in the ith row of U1 that can have a child in Γ. These nodes must either be in the dominator set, or have all their parents in the dominator set, so their number is at most τ (wi (Q) ) + wi (U1 ) . With wi = wi (Q) + wi (U1 ) , this is at most τ (wi ). 13
(a)
(b)
Let us sort U1 into two sub-matrices: let U1 contain those rows where wi ≥ w, and U1 contain those rows where wi < w. The matrix multiplication can similarly be split into two sub(a) (b) (a) multiplications U1 H and U1 H. In U1 , each row has at least w nodes in the dominator set, so 2M the number of rows is at most w . On the other hand, the number of nodes of H in the dominator set (a) is also at most 2M . Each row in U1 has only one value that is multiplied with any of the 2M nodes (a) 4M 2 in H, hence the number of uncovered internal nodes in U1 H is at most 2M w · 2M = w . (b)
M On the other hand, in U1 , each row can have at most τ (w) ≤ 4w nodes that have a child in Γ. Recall that each node in Γ has a successor in the minimum set; this means that there are at most 2M sink nodes in the output matrix such that each node in Γ is a predecessor of one of these sink nodes. M For such a sink node v0 , we can have at most 4w nodes in Γ that are predecessors of v0 , because at M most 4w nodes of U1 can have a child in Γ in the given row. This limits the number of uncovered 2 (b) M internal nodes in U1 H to 4w · 2M = M 2w . 2
2
2
M M Altogether, the number of nodes in Γ is at most 4M + 4M w + 2w = O( w ). With n · r · d internal n·r·d·w nodes, this means that the number of partitions is at least Ω( M 2 ). Using Lemma A.1, we get that the optimal I/O cost is at least Ω(2 · n·r·d·w − 1) · M = Ω( n·r·d·w ). This finishes the proof of the M2 M lower bound.
In order to extend the proof from the multiplication U1 H to the entire computation U1 (U⊤ 2 V), we need to consider the fact that the nodes in H = U⊤ V are in fact not sources. That is, instead of 2 having an entry of H in the dominator set, we could include some of its predecessors. However, each node in H has n parents in the same column of V, which would all need to be in the dominator set in this case. As such, if a dominator set in our previous proof contained x ∈ {0, ..., r} entries from a specific column of H, we could only replace these by n nodes in V (and further nodes in U2 or K, but these can be ignored now). Intuitively, with n > r, this only increases the size of the dominator set. As such, for the whole computation U1 (U⊤ 2 V), it still holds that there can be at most 2M nodes in H with a child that belongs to Γ, which allows us to bound the number of uncovered internal nodes (b) in U1 H exactly as before. A.3
Discussion on red-blue pebbling variants
We note that the standard version of the red-blue pebble game raises some modeling questions regarding operators with many inputs, such as the summation part in the matrix multiplication. If the summation is modeled through a single node, then this requires all inputs in fast memory at the same time, which means that there is no viable pebbling strategy at all if M ≤ r. Previous works have often avoided the discussion of this issue by forming a binary ‘summation tree’ in the computational DAG, and assuming that aggregation happens this way [21]. This is in fact a strong restriction on the execution strategies considered. However, for heavily symmetric operations like matrix multiplication, this had no effect on the optimal I/O complexity: since the optimal strategy was to form tiles of specific sizes, the tiles can naturally consist of neighboring rows/columns in the matrices, and hence the summation within a tile essentially corresponds to a sub-tree in the summation tree of the final aggregated value. The same summation tree approach becomes somewhat more problematic for our Approximate Attention problem. Our results indicate that for an I/O-efficient execution, we need to group together specific columns of the matrix U1 ; hence it becomes critical that the corresponding group of nodes also form a sub-tree in the summation tree of the given entry of the output matrix. As such, this requires a model that inherently entangles the DAG representation of the computation with the optimal I/O strategy to execute it. In general, a more appropriate approach to resolve the question of summation trees is to consider a generalization of the red-blue pebble game with partial computations [19, 23]. This generalized model allows to keep the clean representation of the aggregations as a single node with r incoming edges (as in Figure 1), requiring no summation trees at all. The aggregation process is then modeled by the pebble game itself, which allows to store partially computed results in fast or slow memory. This extended model is a much more accurate description of how the computation is executed in practice. Our algorithm descriptions in the upper bounds are also easiest to interpret in this model. 14
On the other hand, deriving I/O lower bounds for this generalized pebble game also requires a small adjustment to the S-partitioning concept, using so-called S-edge-partitions [19]. While our lower bound proofs are presented in terms of the standard red-blue pebble game and S-partitions, they also carry over to this generalized model and S-edge-partitions by simply shifting the focus from the internal nodes to their (single) outgoing edges, similarly to the proofs adaptations for classic computations [19].
B
Proof of Theorem 1.3
B.1
Upper bound discussion
The upper bound in Theorem 1.3 simply re-states the generic upper bounds from Lemma 3.1. Note that once again when g = O(1), the bounds√are tight up to a constant factor. For instance, the factor √ of difference between n·r·d and n·r·d·g is gM . Due to g = Ω(log M ), we have M = 2O(g) ; for M M √
M g = O(1).
g = O(1), this implies M = O(1), and thus
We point out that besides g = O(1), the bounds are also tight for some other parameter configurations. 2 That is, let us assume g = ω(1). In this case, for instance when we have d = O(g), then n·r·d is M again only a constant factor away from the lower √ bound. On the other hand, when d = Ω(g 2 ), the gap between the upper and lower bound is min(d,g M ) , which is at least g (since M = Ω(g 2 )). As such, this case exhibits a larger gap than a constant factor. At this point, it is a natural question how the upper bounds behave when d is between g and g 2 (with g = ω(1)). As a further contribution, we present another algorithm that is very similar to that of Key Lemma 3.2, but adapted to Case III, in order to answer this question. The upper bound obtained by this algorithm will be weaker than the generic bounds of Lemma 3.1 when d = Ω(g 2 ). However, for the case when there are constants δ1 , δ2 > 0 such that d = O(g 2−δ1 ) and M = O(g δ2 ), this algorithm will once again provide a tight upper bound of O( n·r·d·g M ); this is formally stated in Lemma 3.3. This result extends the tight upper bound from d = O(g) to d almost quadratic in g, assuming that M is at most polynomial in g. B.2
Detour: proof of Lemma 3.3
More specifically, our algorithm for Lemma 3.3 proves an upper bound of d 4·r·g O n·d· + g M
(3)
when d = O(g 2−δ1 ) and M = O(g δ2 ) for some constant δ > 0. We first show that in our parameter regime, we have dg ≤ 4·r·g M ; this implies that the formula in Equation (3) can be upper bounded by 8·n·r·d·g , which is indeed a constant factor away from the M lower bound of Theorem 1.3. The claim we need to prove is equivalent to M r ≤ d = 4g g
d+g g d g
=
(d + g) · . . . · (d + 1) . d · . . . · (d − g + 1)
g 2−δ1 The expression on the right-hand side can be lower bounded by ( d+g ), this can d ) . With d = O(g 2−δ′ g 1 +g g be further lower bounded by for some constant δ1′ such that δ1 > δ1′ > 0. This is then 2−δ ′ g
g
1−δ1′
+1 ′ g 1−δ1
!g
1
= 1+
1 ′ g 1−δ1
g
= 1+
1 ′ g 1−δ1
(g1−δ1′ )
(gδ1′ )
.
The outer parenthesis goes to e as g goes to infinity, hence from some point we can lower bound ′ it by 2e . We also have M = O(g δ2 ) ≤ g δ2 for some δ2′ > δ2 . Hence we only need to show 15
(gδ1′ ) ′ g δ2 ≤ 4 · g · 2e ; this indeed holds asymptotically for any constants δ1′ , δ2′ , since the left side is polynomial, while the right side is exponential in g. It remains to show the I/O strategy that proves the upper bound in Equation (3). This is very similar to the upper bound proof of Key Lemma 3.2, with a specific choice of w = g. That is, we split the d M output matrix into tiles of height M 4g and width g. The row strip of size 4g × r in U1 is split into g aggregation tiles horizontally, and the column strips of width g in H are similarly split vertically. For the aggregation tiles, now each subset of {1, . . . , d} of size g will form a separate tile; this corresponds to having |G1 | = . . . = |Gd | = 1 in the key lemma proof. In each aggregation tile, we load the corresponding g values; this gives one term that can only be generated on this tile, and other terms that can be generated in multiple tiles. We once again assign each of the latter terms to the first possible tile in the lexicographic order. Similarly to the analysis before, this approach results in n d M · g tiles in the output matrix. 4g
The key difference to Key Lemma 3.2 here is that the width of the tiles formed in U1 might be larger than their height: that is, we may have τ (g) > M 4g . In this case, the corresponding tile in H, which has size τ (g) × g, might not fit into fast memory. Due to this, we split each of these tiles further into M 1 sub-tiles of height at most M 4g ; with this, the corresponding tile in H has size 4g · g = 4 M . Executing this split only requires us to memorize the current position in the aggregation tile in U1 . For each sub-tile, the same M 4g × g values can be kept from Q, the tile aggregated in the output can also remain in fast memory. However, from H, we need to load up to 14 M new values for each sub-tile. Let the width of the tiles in U1 be t1 , t2 , . . .; note that their sum is r. Then the number oftimes we need to d t1 t1 start a new sub-tile within a tile is at most M + M + . . . = Mr = 4·r·g M . With g tiles, this means 4g 4g 4g that the number of sub-tiles is at most dg + 4·r·g M altogether. This results in a total I/O cost of d 4·r·g n d + ·M, · · M g g M 4g which is equivalent to Equation (3). The strategy extends to the multiplication U⊤ 2 V exactly as in the proof of Key Lemma 3.2: dwe split d d M U⊤ to column strips of width , split each of these to parts vertically, and form 2 4g g g · g tiles in n the output H, each of which is aggregated over M steps. Similarly to the second multiplication, we 4g
M might have tiles in U⊤ 2 that are taller than 4g ; these are then split vertically into sub-tiles of height M M M ⊤ 4g at most. For all of the sub-tiles, the tile of size g × 4g in K and the tile of size 4g × g in V can remain the same, whereas in H, we save different sub-tiles each time.
B.3
Lower bound proof
The lower bound proof of Theorem 1.3 is similar to that in Key Lemma 3.2. However, here we only focus on the columns of U1 that have at least g2 distinct parents in Q. One can show that if d ≥ 5g, this indeed amounts to a constant fraction of the columns of U1 . This follows from the lemma below, which can be understood as the reverse direction: the number of terms generated from at most ( g2 − 1) variables is at most δ · r. Lemma B.1. For d ≥ 5g, there exists a constant 0 < δ < 1 d g ·τ −1 ≤ δ ·r. g 2 2 −1 Since this is only a technical lemma that offers no further insight into the proof, we defer its proof to Appendix D. In our proof, we only consider a subset of the matrix multiplication: the at least δ · r columns in U1 that satisfy this condition, and the corresponding at least δ · r rows in H. We assume that all other I/O operations are free; the lower bound derived this way clearly carries over to the original matrix multiplication. 16
The proof begins similarly to Key Lemma 3.2. Let Di be a partition in a 2M -partition of this computation. Consider its minimum set and a dominator set of size at most 2M . Let Γ be the set of all the uncovered internal nodes in Di that are not among these 4M nodes. All uncovered nodes v ∈ Γ must have their parent in H in the dominator set, and also either their parent in U1 or all the parents of this parent in the dominator set. For any i ∈ [n], let wi (Q) and wi (U1 ) again be, respectively, the number of entries in the ith row of Q and the ith row of U1 that are in the dominator set. Let wi = wi (Q) + wi (U1 ) . Consider two (a) (b) sub-matrices of U1 : let U1 contain the rows with wi ≥ g2 , and U1 contain the rows with wi < g2 . (a) (b) (a) This splits the matrix multiplication into two sub-multiplications U1 H and U1 H. In U1 , each g row has at least 2 nodes in the dominator set, so the number of rows is at most 2M = 4M g g . The 2
(a)
number of nodes of H in the dominator set is at most 2M . Each row in U1 has only one value (a) multiplied with any entry in H, so the number of uncovered internal nodes in U1 H is at most 2 4M 8M g · 2M = g . (b)
In U1 , none of the nodes in (our restricted) U1 can have all their parents in the dominator set; (b) hence if a node in Γ has its parent in U1 , then this parent has to be in the dominator set. Thus for simplicity, we can assume that wi (Q) = 0, i.e. none of the source nodes are in the dominator set, (b) since these nodes could be removed from the set anyway. With wi (U1 ) < g2 , each of the rows in U1 g can have at most 2 nodes which have a child in Γ. Similarly to before, each node in Γ must also have a successor in the minimum set, so there are at most 2M sink nodes in the output matrix with a (b) predecessor in Γ. Each sink node has at most g2 predecessors in Γ, since number of nodes in U1 g with a child in Γ is at most 2 in the given row. This means that the number of uncovered internal (b) nodes in U1 H is at most g2 · 2M = g · M . 2
In Case III, we have M = ω(g 2 ). This means that g · M is also upper bounded by O( Mg ), so the 2
2
M total number of nodes in Γ is at most 4M + 8M g + g · M = O( g ). Having n · r · d internal nodes,
the number of partitions is at least Ω( n·r·d·g M 2 ). Using Lemma A.1, the optimal I/O cost in this case is n·r·d·g at least Ω( M ). The proof can be extended to the entire computation U1 (U⊤ 2 V) exactly as in Key Lemma 3.2.
C
Proof of Theorem 1.4
The upper bound in Theorem 1.4 follows easily from the generic upper bound in Lemma 3.1. For the lower bound, the proof starts identically to that of Theorem 1.3. We only consider the columns of U1 with at least g2 distinct parents in U1 . This once again amounts to a constant fraction of the columns of U1 . (a)
For any row i ∈ [n], let wi , wi (Q) and wi (U1 ) be as before. We split the matrix again into U1 for (b) (a) rows with wi ≥ g2 , and U1 for rows with wi < g2 . In the sub-multiplication U1 H the number of 2 (a) uncovered internal nodes in U1 H is at most 8M g as before. Otherwise, the proof is analogously to the lower bound √ for standard matrix multiplication [15], essentially repeating the same arguments with a split at M instead of g2 . Recall that for each row in (b) (b) (b ) U1 , we can assume wi (Q) = 0. Let us further divide the matrix U1 , with U1 1 containing the √ √ (b) (b ) (b) rows of U1 where wi (U1 ) ≥ M , and U1 2 containing the rows of U1 where wi (U1 ) < M . √ (b ) 2M = 2 M . Each of these rows has at most one value In U1 1 , the number of rows is at most √ M multiplied with any of the (at most 2M ) nodes of H in the dominator set, so the number of nodes √ √ (b ) (b ) of Γ in U1 1 H is at most 2M · 2 M = 4M · M . In U1 2 , we again consider the at most√2M sink nodes in the output matrix that have a predecessor in Γ. Each of these can have at most M 17
predecessors in Γ, since the number of nodes in Γ in any row is at most √ (b ) at most 2M · M nodes in U1 2 .
√
M . This means that Γ has
√ (b) Altogether, the number of nodes of Γ in U1 H in this case is at most 6M · M . Together with (a) the nodes in U1 H and the internal nodes directly in the dominator/minimum sets, this is at most √ √ 2 2 4M + 8M M nodes. Since Case IV has M = O(g 2 ), we also have Mg = O(M · M ). g + 6M · √ As such, this expression is altogether in O(M · M ). With n · r · d internal nodes, the number of √ ), and the optimal I/O cost according to Lemma A.1 is Ω( n·r·d √ partitions is Ω( Mn·r·d ). · M M
D
Technical lemmas
D.1
Using attention to compute the exponential function
Assume that attention can be computed exactly in T steps. To compute the exponential exp(x) for any x, simply set Q = 1, K⊤ = (x 2x), V = (1; 0), then compute A = Att(Q, K, V) = exp(x) exp(x)+exp(2x) , and finally return z = 1/A − 1 = exp(x). D.2
Closed form for τ (w)
For completeness, we establish the closed-form expression on τ (w). Lemma D.1. For any integer w ≥ 1, we have g X w+l−1 w+g = . l g l=0
Proof. We can show this by an induction on g. The claim clearly holds for small values. For g = 0, w we have w−1 on the left-hand side and = 1 0 = 0 on the right-hand side. For g = 1, we have 0 w w+1 1 + 1 = w + 1 on the left-hand side and 1 = w + 1 on the right-hand side. Now consider g ≥ 2, and assume the claim holds for (g − 1). This means that we only need to show w+g−1 w+g w + (g − 1) = − . g g (g − 1) Indeed, the right-hand side is identical to (w + g) · (w + g − 1) · . . . · (w + 1) (w + g − 1) · (w + g − 2) · . . . · (w + 1) − = g! (g − 1)! (w + g) ·
(w + g − 1) · . . . · (w + 1) (w + g − 1) · . . . · (w + 1) −g· = g! g! (w + g − 1) · . . . · (w + 1) · w = g!
D.3
w+g−1 . g
Proof of Lemma B.1
Below we prove Lemma B.1, i.e. that for d ≥ 5g, we have a constant 0 < δ < 1 with d g ·τ −1 ≤ δ ·r. g 2 2 −1 Proof. For simplicity, we instead prove the stronger statement d g · τ − 1 ≤ δ ·r. g 2 2 18
Using the definition of r and τ , this is equivalent to 3 d 2g − 1 ≤ δ · d + g . · g g g 2 2 −1 This can be further expanded to ( 32 g − 1)! d! (d + g)! · ≤δ· g g g ( 2 )! · (d − 2 )! ( 2 − 1)! · g! d! · g! and then
g
d · . . . · (d −
( 3 g − 1) · . . . 2 g + 1) · 2 ≤ δ · (d + g) · . . . · (d + 1) . 2 ( g2 )!
We can upper bound ( 23 g − 1) · . . . g2 by 12 · g g . We can lower bound ( g2 )! by ( g5 )g/2 . We also have g (d + g) · . . . · (d + 1) ≥ d2 . g d · . . . · (d − 2 + 1)
Hence it is enough to prove that g g gg 1 1 g · g g/2 = · 5 2 · g 2 ≤ δ · d 2 . 2 (5) 2
This is equivalent to g1 1 . 4 · δ2 By selecting a constant δ close to 1, we have 4·δ1 2 ≤ 31 . The left side is then upper bounded by 1 for any g ≥ 1, and hence it is sufficient to have d ≥ 5 · g, which holds due to our assumption. d ≥ 5g
We note that if we modify the restriction condition from g2 to a smaller linear factor of g, the above proof can be applied similarly; this would allow us to loosen the condition d ≥ 5 · g to obtain a slightly smaller constant factor between d and g, if desired.
19