Analyze Transformer Self-Attention Complexity

mediumFREETransformersAlgorithms, ComplexityAdvanced Architectures

The self-attention mechanism is a critical component of the Transformer architecture. It processes a sequence of $N$ input vectors (tokens or patches), each with embedding dimension $D$ .

The computation requires calculating pairwise affinities between every element (query $q$ ) and every other element (key $k$ ) in the sequence, typically followed by scaling and multiplying by the value matrix $V$ .

Given:

$N$ : Sequence length
$D$ : Embedding dimension
$K$ : Convolution kernel size (for comparison)

Task: Determine the asymptotic time complexity for:

A single layer of self-attention: $O(N^2 \cdot D)$
A standard convolution layer: $O(N \cdot D^2 \cdot K^2)$

Return a tuple of strings representing both complexities.

Write each complexity as a string with no spaces, using ^ for powers and ***** for multiplication, exactly as shown in the example output.

Example

Input

N=100, D=512, K=3

Output

("O(N^2*D)", "O(N*D^2*K^2)")

Reasoning

Inputs: $N = 100$ tokens of dimension $D = 512$ , and a $3 \times 3$ kernel ( $K = 3$ ) for the comparison.
Self-attention: every query is compared with every key, giving an $N \times N = 10{,}000$ score matrix. Each score is a dot product of length $D$ , so that costs $100^2 \cdot 512 = 5{,}120{,}000$ multiply-adds; weighting the value vectors by those scores costs the same again, $10{,}240{,}000$ in total. The constant factor 2 drops, leaving $O(N^2 \cdot D)$ .
Convolution: at each of the $N$ positions, each of $D$ output channels sums over $D$ input channels times $K^2$ kernel taps: $100 \cdot 512^2 \cdot 9 = 235{,}929{,}600$ , i.e. $O(N \cdot D^2 \cdot K^2)$ .
Output: the two answers as strings, O(N^2*D) and O(N*D^2*K^2), attention first. They never change with the inputs: they describe growth, not a count.
Notice that for this short sequence the convolution is about 23 times costlier ( $235{,}929{,}600 / 10{,}240{,}000 \approx 23$ ), but attention grows quadratically in $N$ and overtakes it once $2N^2 D > N D^2 K^2$ , i.e. $N > D K^2 / 2 = 2304$ .

Constraints

$N$ (Sequence Length): $1 \leq N \leq 10^5$
$D$ (Embedding Dimension): $1 \leq D \leq 1024$
$K$ (Kernel Size): $1 \leq K \leq 10$

Learn this in

Advanced Topics & ApplicationsDiffusion Study Plan Introduction to LLMsLLM Study Plan Transformer ArchitectureLLM Study Plan Attention MechanismsNLP Study Plan

More from Advanced Architectures

Complete Mixer Block with Residualshard Channel Mixing with GELUmedium Token Mixing Layermedium Bidirectional Scan for Visionmedium Discretize Continuous State Space Modelhard

Concept primer

1. Background Knowledge

The Transformer architecture, introduced in "Attention is All You Need" (2017), relies on self-attention to process sequences without recurrence or convolution. For a sequence of $N$ tokens each with $D$ -dimensional embeddings $X \in \mathbb{R}^{N \times D}$ :

Self-attention computes:

Query ( $Q$ ), Key ( $K$ ), Value ( $V$ ) matrices: $Q = XW_Q$ , $K = XW_K$ , $V = XW_V$ where $W \in \mathbb{R}^{D \times D}$
Attention scores: $A =$ \text{softmax} $\left($ \frac{QK^T}{ $\sqrt{D}$ } $\right) \in \mathbb{R}^{N \times N}$
Output: $Z = AV \in \mathbb{R}^{N \times D}$

This captures all pairwise interactions between tokens, enabling long-range dependencies but at quadratic cost.

Standard convolution slides a $K \times K$ kernel over input, computing local interactions with complexity linear in sequence length.

2. Algorithm Approach

Direct matrix computation for self-attention:

Q, K, V = linear_projection(X) # O(N × D²)
attention = softmax(Q @ K.T / √D) # O(N² × D) [matrix mult] + O(N²) [softmax]
output = attention @ V # O(N² × D)

Total: $O(N^2 \cdot D + N \cdot D^2)$ , dominated by $O(N^2 \cdot D)$ when $N \gg D$ .

Convolution for 1D sequence:

For each position i in N:
 For each output channel cout in D:
 For each input channel cin in D:
 For each position in kernel K:
 output[i] += input[i:i+K] * kernel[cin, cout]

Total: $O(N \cdot D^2 \cdot K^2)$ .

3. Step-by-Step Strategy

Count Q, K, V projections: 3 matrix multiplies, each $O(N \cdot D^2)$
Analyze attention matrix: $QK^T$ is $N \times D$ × $D \times N = O(N^2 \cdot D)$ , softmax $O(N^2)$
Attention × V: $N \times N$ × $N \times D = O(N^2 \cdot D)$
Dominant term: $O(N^2 \cdot D)$ from quadratic $N^2$ terms
Convolution: Per position: $K$ positions × $D^2$ channels = $O(N \cdot D^2 \cdot K^2)$

Expected output: ("O(N² · D)", "O(N · D² · K²)")

4. Common Pitfalls

Forgetting projection costs: $O(N \cdot D^2)$ exists but is subdominant when $N \gg D$ (constraints: $N \leq 10^5$ , $D \leq 1024$ )
Multi-head confusion: $h$ heads → $O(h \cdot N^2 \cdot$ \frac{D}{h} $+ N \cdot h \cdot$ \frac{D}{h} $\cdot D) = O(N^2 \cdot D + N \cdot D^2)$ , same asymptotics
Ignoring softmax: $O(N^2)$ but negligible vs. $N^2 \cdot D$
Convolution dimensionality: Verify 1D vs. 2D; problem implies 1D sequence ( $N$ tokens)
Space complexity: Self-attention stores $O(N^2)$ attention matrix (65GB for $N=10^5$ !)

Operation	Self-Attention	Convolution
Time	O(N² · D)	O(N · D² · K²)
Space	O(N²)	O(N · D)
Scaling	Quadratic explosion	Linear scaling

5. Time & Space Complexity

Self-attention single layer:
└── Time: O(N² · D) [QKᵀ + AV dominate]
└── Space: O(N²) [attention matrix]

Standard convolution:
└── Time: O(N · D² · K²)
└── Space: O(N · D)

Key insight: Self-attention's global pairwise computation ( $N^2$ pairs × $D$ work/pair) vs. convolution's local sliding window ( $N$ positions × $K^2 \cdot D^2$ work/position).

solution.py

Python 3.13

Test results