How To Reduce The Memory Usage Of The Self-Attention

With a bit of magic, we take a very inefficient computation like the Self-Attention and make it super memory-optimized for the specific hardware we use for training and inference. And we all need a bit of magic!

• Self-attention Does Not Need O(N²) Memory

• The GPU Architecture

• The FlashAttention-1

• The FlashAttention-2

• The FlashAttention-3

---

So far, we have mainly explored how to reduce the complexity of the attention mechanisms by approximating the vanilla attention. The vanilla attention has a strict O(N²) time complexity, but the space complexity doesn't need to be O(N²)! Computing Q^TK requires ~O(N²) operations, but the full N x N alignment scores and attention matrices do not need to be fully materialized all at once in memory. As models and sequence lengths scale, it becomes essential to minimize the memory requirements at training and inference time to better utilize the underlying hardware.

Self-attention Does Not Need O(N²) Memory

Let's consider again the computation of the context vectors:

\(\begin{align} \mathbf{c}_i &= \frac{\sum_{j=1}^N \exp\left(\frac{\mathbf{q}_i^\top \mathbf{k}_j}{\sqrt{d_\text{model}}}\right) \mathbf{v}_j}{\sum_{j=1}^N \exp\left(\frac{\mathbf{q}_i^\top \mathbf{k}_j}{\sqrt{d_\text{model}}}\right)}. \end{align}\)

Here is how we arrive at the typical O(N²) space complexity:

• The typical assumption is that we first compute the dot product between the query q_i and all the keys [k₁, …,k_N]:

  \( \mathbf{e}_i = \left[\frac{\mathbf{q}_i^\top \mathbf{k}_1}{\sqrt{d_\text{model}}}, \ldots, \frac{\mathbf{q}_i^\top \mathbf{k}_N}{\sqrt{d_\text{model}}}\right] \)

  Where e_i is the alignment score vector of size N for the query q_i, which leads to the N x N matrix for the N queries.

• We then perform the softmax transformation:

  \( \mathbf{a}_i = \frac{1}{\sum_{j=1}^N\exp\left(\frac{\mathbf{q}_i^\top \mathbf{k}_j}{\sqrt{d_\text{model}}}\right)}\left[\exp\left(\frac{\mathbf{q}_i^\top \mathbf{k}_1}{\sqrt{d_\text{model}}}\right), \ldots, \exp\left(\frac{\mathbf{q}_i^\top \mathbf{k}_N}{\sqrt{d_\text{model}}}\right)\right] \)

  Here, a_i is the attention vector of size N for the query q_i. Again, for N queries, it leads to the typical N x N attention matrix.

• And finally, we project a_i onto the different values [v₁, …,v_N], which leads to c_i:

  \(\begin{align} \mathbf{c}_i &= \mathbf{a}_i^\top V\nonumber\\ &=\sum_{j=1}^Na_{ij}\mathbf{v}_j \end{align}\)

Therefore, naively computing the alignment scores and the attention matrices first forces the materialization of those matrices in memory, which leads to the O(N²) space complexity.

However, we do not need to order the computations in this manner! In 2021, Rabe and Staats realized that by reordering the operations, we can greatly reduce the requirements on the memory. The idea is to consider the unnormalized context vector c̃_i and the normalization constant ∑exp(q_i^Tk_j) of the softmax transformation separately:

\(\begin{align} \tilde{\mathbf{c}}_i &= \sum_{j=1}^N \exp\left(\frac{\mathbf{q}_i^\top \mathbf{k}_j}{\sqrt{d_\text{model}}}\right) \mathbf{v}_j \nonumber\\ s_i&=\sum_{j=1}^N\exp\left(\frac{\mathbf{q}_i^\top \mathbf{k}_j}{\sqrt{d_\text{model}}}\right)\nonumber\\ \mathbf{c}_i &= \frac{\tilde{\mathbf{c}}_i}{s_i} \end{align}\)

Because c̃_i and s_i are just sums, we can easily loop through the key-value pairs to compute the context vector:

At any point during the for-loop, we only need to store the intermediary values of c̃_i and s_i. c̃_i is a vector of size d_model (ignoring heads for simplicity), and s_i is a scalar. Therefore, for one query, we need constant space complexity O(1) to compute one context vector. Even iterating through all the queries, we never need to capture more than the intermediary values of c̃_i and s_i, so we can compute the full attention mechanism in O(1) space complexity:

In reality, this is not a practical solution because sequential operations are not adapted to the parallelization capability of the CPU, GPU, or TPU hardware that is commonly used for neural network computations. In practice, the queries, keys, and values are partitioned into chunks to allow for a high degree of parallelization while keeping the memory requirement low. Let's assume that we partition the queries into n_q chunks and the keys and values into n_k chunks:

\(\begin{align} Q&=\left[Q_1, Q_2, \ldots, Q_{n_q}\right] \nonumber\\ K&=\left[K_1, K_2, \ldots, K_{n_k}\right] \nonumber\\ V&=\left[V_1, V_2, \ldots, V_{n_k}\right] \end{align}\)

where each Q_i is a N / n_q x d_model matrix and K_i, V_i are N / n_k x d_model matrices. Let's call N_q = N / n_q, the number of queries per chunk, and N_k = N / n_k, the number of key-value pairs per chunk. We can now iterate through the chunks exactly in the same way:

As before, we need to store intermediary values of C̃_i and S_i. In this context, Q_iK^T_j is a matrix of size N_q x N_k, and so is A_ij. C̃_i is a matrix of size N_q x d_model and S_i is a vector of size N_q. Therefore the space complexity is O(N_q x N_k + N_q x d_model). To balance the number of chunks and the number of key-value pairs per chunk, they chose N_k = n_k = √N and fixed N_q = 1024. This results in a space complexity:

\(\mathcal{O}(1024\sqrt{N} + 1024 d_\text{model})=\mathcal{O}(\sqrt{N})\)

This approach allows for efficient tensor operations within each chunk while dramatically reducing the peak memory requirements. Note that no approximation has been made, and it is mathematically equivalent to the vanilla attention mechanism. However, this approach is slower (8-13% slower during the forward pass and 30-35% slower during the backward pass) due to the sequential computations, but it enables the processing of much longer sequences that would otherwise be impossible due to memory constraints. Along with the FlashAttention, it is one of the memory optimization strategies used in the xFormers package developed by Meta and used in the development of the Llama models.

Stabilizing The Computations

Until now, we have been ignoring the numerical stability of the softmax computation, but most implementations (PyTorch, TensorFlow, ...) are using a couple of tricks to ensure its stability. Let's remind ourselves of the softmax function:

\( \text{Softmax}(x_i) = \frac{e^{x_i}}{\sum_{j=1}^N e^{x_j}}\)

Computing e^xi can be tricky because if x_i ≥ 89, then e^xi ≅ 4.4e38, which exceeds the floating-point limit of 3.4e38 for a 32-bit float number, potentially leading to float overflow errors. To prevent this from happening, we typically modify the exponent by finding the maximum x_i value: