The transformer model, introduced by Vaswani et al., has become a cornerstone of modern natural language processing due to its capability to capture long-range dependencies effectively through the self-attention mechanism. However, this capability comes with a computational cost, particularly concerning time and space complexity. Understanding the computational complexity of self-attention within transformers is crucial for developing efficient models, especially for large-scale applications.

Self-Attention Mechanism

The self-attention mechanism allows the model to weigh the relevance of different words in a sentence while encoding a particular word representation. In a typical transformer self-attention block, given an input sequence of length $n$ and dimensional size $d$, attention is computed as follows:

1. Linear Projections: Each input token is linearly projected onto three different linear subspaces to generate queries ($Q$), keys ($K$), and values ($V$), using learnable parameter matrices $W^Q$, $W^K$, and $W^V$: $Q = XW^Q, \quad K = XW^K, \quad V = XW^V$ The dimensions for these matrices are typically such that $W^Q, W^K, W^V \in \mathbb{R}^{d \times d_k}$.
2. Attention Scores: The query and key matrices are then used to compute attention scores, essentially dot products that signify the compatibility or attention between various token representations: $\text{Attention Scores} = QK^T$ This results in an $n \times n$ matrix.
3. Scaling: The attention scores are scaled by the inverse square root of the dimension $d_k$: $\text{Scaled Scores} = \frac{QK^T}{\sqrt{d_k}}$
4. Softmax: The scaled scores are passed through a softmax function to obtain attention weights: $\text{Attention Weights} = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)$
5. Weighted Sum: Finally, attention weights are used to compute a weighted sum of the value vectors: $\text{Self-Attention Output} = \text{Attention Weights} \cdot V$

Computational Complexity

Time Complexity

The major computational cost of self-attention comes from the matrix multiplication involved, primarily from the $QK^T$ and subsequent operations:

• Query-Key Matrix Multiplication: The computation of $QK^T$ incurs a time complexity of $O(n^2d_k)$. As $d_k$ is often proportional to $d$, this can be approximated as $O(n^2d)$.
• Attention Output Calculation: Multiplying the attention weights with the value matrix $V$ also requires $O(n^2d)$ operations.

Thus, the overall time complexity per layer of the self-attention mechanism is dominated by the $O(n^2d)$ factor. This quadratic dependency on the sequence length $n$ becomes a bottleneck for long sequences.

Space Complexity

Self-attention primarily requires memory for storing three main components:

• Projection Matrices ($Q$, $K$, $V$): The storage for these projections and the matrices themselves is $O(nd_k)$ each. With typical configurations, this is $O(nd)$ combined.
• Attention Scores and Weights: Storing the intermediate attention scores and weights adds another $O(n^2)$ memory requirement.

Overall, the space complexity predominantly arises from $O(n^2)$ because all $n \times n$ scores and weights need to be stored in memory.

Complexity Summary

Optimizing Self-Attention Complexity

Several modifications and advancements have been proposed to mitigate the self-attention overhead:

1. Efficient Attention Mechanisms: Techniques like Linformer and Reformer propose approximate methods to reduce complexity by projecting attention calculations into lower-dimensional spaces or using locality-sensitive hashing.
2. Sparse Attention Patterns: Limiting attention to only specific parts of a sequence, as seen in Longformer and BigBird, effectively reduces computational overhead while retaining the model's attention capabilities where most needed.
3. Pruned Models: Techniques like quantization and pruning remove redundant components, thereby reducing the memory footprint and improving inference speed.

These advancements aim to balance the power of transformers with computational efficiency, making them applicable in scenarios where resource constraints are a concern. By managing complexity, transformers can achieve performance gains necessary for their deployment in both small device contexts and extensive computational environments.