Analyze Transformer Self-Attention Complexity

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:

1. Query ($Q$), Key ($K$), Value ($V$) matrices: $Q = XW_Q$, $K = XW_K$, $V = XW_V$ where $W \in \mathbb{R}^{D \times D}$
2. Attention scores: $A =$\text{softmax}\left(\frac{QK^T}{$\sqrt{D}$}\right) \in \mathbb{R}^{N \times N}
3. 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

1. Count Q, K, V projections: 3 matrix multiplies, each $O(N \cdot D^2)$
2. Analyze attention matrix: $QK^T$ is $N \times D$ × $D \times N = O(N^2 \cdot D)$, softmax $O(N^2)$
3. Attention × V: $N \times N$ × $N \times D = O(N^2 \cdot D)$
4. Dominant term: $O(N^2 \cdot D)$ from quadratic $N^2$ terms
5. 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$!)

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).