DL0033 Transformer Computation

Written by

In a Transformer architecture, which components are the primary contributors to computational cost, and why?

Answer

It depends on sequence length. For short sequences, the feed-forward network (FFN) usually dominates: its two wide GEMMs cost $O(n d^2)$ while attention’s quadratic term is still small. For long sequences, multi-head attention takes over: forming the $n \times n$ score matrix costs $O(n^2 d)$ and grows much faster than anything else in the block.

(1) Multi-Head Attention: Q/K/V projections cost $O(nd^2)$ , but the score matrix $QK^\top$ and its product with V cost $O(n^2 d)$ , the quadratic term that explodes on long sequences.
(2) Feed-Forward Network: Two dense layers with expansion factor 4 cost $O(nd^2)$ , dominant when $n$ is small but $d$ is large.
(3) Crossover Point: Attention overtakes the FFN roughly when $n \approx 2d$ ; near $n = 1024$ for the classic $d_{model} = 512$ design.

Line chart of MHA and FFN share of total FLOPs versus sequence length on a log scale, crossing at sequence length 1024 where each takes fifty percent.

Figure 1: FLOP share vs sequence length ( $d = 512$ ): FFN dominates below $n \approx 1024$ ; past the crossover, attention’s quadratic cost becomes the bottleneck.

Mathematical Formulation (per block):
$\text{Cost}_{\text{attn}} = \underbrace{4nd^2}_{\text{QKVO projections}} + \underbrace{2n^2 d}_{QK^\top \text{ and } AV}$
$\text{Cost}_{\text{FFN}} = 2 \cdot n \cdot d \cdot 4d = 8nd^2$

Where:

$n$ is the sequence length and $d = d_{model}$ ; constants count multiply–adds per GEMM.
The $2n^2 d$ term comes from multiplying $(n \times d) \cdot (d \times n)$ to form scores, and again to mix values.
Softmax itself is cheap elementwise work but also scales with $n^2$ entries.

Sequence Length n	MHA Share (%)	FFN Share (%)	Dominant Component
64	34.69	65.31	FFN
256	38.46	61.54	FFN
1024	50.00	50.00	Tie
4096	71.43	28.57	MHA

Table 1: FLOP breakdown at $d = 512$ : the tie at n = 1024 marks where quadratic attention catches up with the linear-in-n FFN.

Flowchart of one Transformer block annotating each operation with its FLOP complexity: projections linear in n, score matrix quadratic in n, FFN linear in n but quadratic in d.

Figure 2: Where the FLOPs live inside one block: only the $QK^\top$ and $AV$ pair grows quadratically with sequence length.

Practical Caveat: The FFN’s dominance at short $n$ assumes the standard 4x expansion; shrink the expansion to 1x and the Q/K/V projections become the largest term, while efficient attention variants (sliding-window, linear attention) specifically attack the $n^2$ term for long contexts.

Transformer

DL0033 Transformer Computation

Comments

Leave a Reply Cancel reply

More posts

MSD0064 Notification Ranking System

MSD0063 Image Copyright Detection

MSD0062 Ads Ranking System

MSD0061 Nearby Place Recommendation