Computational Complexity of Self-Attention

As introduced, the standard self-attention mechanism, while powerful, carries a significant computational burden, especially as input sequences grow longer. Understanding this computational cost is fundamental to appreciating the motivation behind the more efficient architectures we will explore later in this chapter.

Let's analyze the core operations within the scaled dot-product attention formula, which forms the basis of self-attention in Transformers:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

Here, $Q$ (Query), $K$ (Key), and $V$ (Value) are matrices derived from the input sequence embeddings. Let $N$ be the sequence length and $d_k$ be the dimension of the keys and queries, and $d_v$ be the dimension of the values.

The primary computational steps are:

1. Query-Key Similarity Calculation: Computing the matrix product $QK^T$.

   • $Q$ has dimensions $(N \times d_k)$.
   • $K^T$ has dimensions $(d_k \times N)$.
   • The resulting attention score matrix $S = QK^T$ has dimensions $(N \times N)$.
   • The computational cost of this matrix multiplication is approximately $O(N \cdot d_k \cdot N) = O(N^2 d_k)$ floating-point operations (FLOPs).

2. Scaling and Softmax: Scaling the scores by $1/\sqrt{d_k}$ and applying the softmax function row-wise.

   • Scaling involves $N^2$ element-wise multiplications.
   • Softmax involves exponentiation and normalization for each of the $N$ rows, each requiring operations proportional to the row length $N$. The total cost is roughly $O(N^2)$.
   • Compared to the matrix multiplications, this step is typically less computationally dominant for large $N$ and $d_k$.

3. Value Aggregation: Multiplying the softmax output (the attention weights matrix $A$, dimensions $(N \times N)$) by the Value matrix $V$.

   • $A$ has dimensions $(N \times N)$.
   • $V$ has dimensions $(N \times d_v)$.
   • The resulting output matrix $O = AV$ has dimensions $(N \times d_v)$.
   • The computational cost of this matrix multiplication is approximately $O(N \cdot N \cdot d_v) = O(N^2 d_v)$ FLOPs.

Dominant Factors and Overall Complexity

Combining these steps, the total computational complexity is dominated by the two large matrix multiplications: $O(N^2 d_k + N^2 d_v)$.

In many standard Transformer configurations, the dimensions $d_k$ and $d_v$ are proportional to the overall model embedding dimension $d_{model}$ (often $d_k = d_v = d_{model} / h$, where $h$ is the number of attention heads). Therefore, the complexity is commonly summarized as:

$$O(N^2 \cdot d_{model})$$

This quadratic dependence on the sequence length $N$ is the critical bottleneck. While the operations are highly parallelizable across the sequence dimension (unlike recurrent models), the total number of operations grows rapidly with $N$.

Memory Complexity

Beyond computation, there's also a significant memory requirement. The intermediate attention score matrix $QK^T$ (before or after softmax) has dimensions $(N \times N)$. Storing this matrix requires:

$$O(N^2)$$

memory. For long sequences (e.g., $N > 4096$), storing an $(N \times N)$ matrix of floating-point numbers can exceed the memory capacity of typical accelerators like GPUs, even before considering the memory needed for activations, gradients, and model parameters.

The graph illustrates the rapid divergence in computational cost. As sequence length $N$ increases, the $O(N^2)$ cost of standard self-attention quickly surpasses linear $O(N)$ growth, making it impractical for very long sequences. The Y-axis uses a logarithmic scale to accommodate the large range of values.

Implications for Long Sequences

This quadratic computational and memory complexity severely limits the application of standard Transformers to tasks involving very long sequences, such as:

• Processing entire high-resolution images treated as sequences of patches.
• Analyzing full lengthy documents or books.
• Modeling long-form time series data (e.g., audio, sensor readings).
• Working with genomic sequences.

Even with powerful hardware, sequence lengths beyond a few thousand tokens become challenging from both a runtime and memory perspective during training and inference. This limitation directly motivates the development of the alternative attention mechanisms and architectural variants discussed in the following sections, which aim to reduce this quadratic dependency.