Transformer 的复杂度怎么算？
time、memory、parameter、FLOPs
先弄清楚：每个概念是什么
- FLOPs 与 2mkn 规则
- 一次乘加算 2 个 FLOPs。\((m\times k)(k\times n)\) 的结果有 \(mn\) 个元素，每个要 \(k\) 次乘加，所以是 \(2mkn\)。例：\((2\times3)(3\times4)\) 是 \(2\cdot2\cdot3\cdot4 = 48\) FLOPs。
- n² 和 d² 从哪来
- 每个 token 要和所有 token 各打一次分，所以有 n²。每个 token 还要乘几个 \(d\times d\) 的矩阵（projection、FFN），所以有 d²。
30 秒回答
I count matrix multiplies: an (m×k)(k×n) product costs 2mkn FLOPs. Per layer, the Q, K, V and output projections cost 8nd², the FFN 16nd², and the two attention matmuls 4n²d. So a layer is O(n²d + nd²), attention dominates once n exceeds 6d, and the score matrix needs O(n²) memory unless we use FlashAttention. Parameters are about 12d² per layer, so forward costs about 2N FLOPs per token and training 6N. With a KV cache, decoding costs O(nd) per token per layer.
怎么算：数一层里的矩阵乘
把一层里的每个矩阵乘按 \(2mkn\) 加起来。\(n\) 是序列长度，\(d\) 是 d_model，\(h\) 个 head，\(d_k = d/h\)，forward、batch = 1。
| 步骤 | 矩阵乘的形状 | FLOPs |
|---|---|---|
| Q、K、V projection | \((n\times d)(d\times d)\)，3 次 | \(6nd^2\) |
| \(QK^\top\) | \((n\times d_k)(d_k\times n)\)，\(h\) 个 head | \(2n^2d\) |
| softmax 后乘 \(V\) | \((n\times n)(n\times d_k)\)，\(h\) 个 head | \(2n^2d\) |
| output projection | \((n\times d)(d\times d)\) | \(2nd^2\) |
| FFN（hidden 4d） | \((n\times d)(d\times 4d)\) 和 \((n\times 4d)(4d\times d)\) | \(16nd^2\) |
| 合计 | \(24nd^2 + 4n^2d\) |
每层 time 是 \(O(n^2d + nd^2)\)，\(L\) 层再乘 \(L\)。两项相等时：
\[ 4n^2d = 24nd^2 \;\Longrightarrow\; n = 6d \]\(n > 6d\) 时 attention 的 \(n^2\) 项才主导。
parameter：QKVO \(4d^2\) + FFN \(8d^2\) = 每层 \(12d^2\)，全模型 \(\approx 12Ld^2\) 加 embedding。
per token FLOPs：每个参数对每个 token 做一次乘加，forward \(\approx 2N\)。backward 要对 activation 和 weight 各求一次 gradient，约是 forward 的 2 倍，所以 training \(\approx 6N\)（Kaplan et al. 的 \(C \approx 6NBS\)）。
memory：标准实现每层要存 \(h\) 个 \(n\times n\) 的 score 矩阵，\(O(n^2)\)。FlashAttention 不存整个矩阵，额外 memory 降到 \(O(n)\)，FLOPs 不变。KV cache 存 \(2\cdot L\cdot n\cdot d\) 个数，算例见 KV cache。
inference：prefill 一次处理整个 prompt，每层 \(O(n^2d + nd^2)\)。decode 每步只算 1 个 token，新 query 对 cache 里的 \(n\) 个 key 做 attention，每层 \(O(nd + d^2)\)。没有 cache 就要重算前缀，第 \(t\) 步每层 \(O(t^2d + td^2)\)。
手算例子：LLaMA-7B
32 层，\(d = 4096\)，32 个 head，context 2048（LLaMA Table 2）：
- 参数：\(12 \times 32 \times 4096^2 \approx 6.44\)B，加上 embedding，论文写 6.7B。SwiGLU 的 3 个矩阵合计仍是 \(8d^2\)。
- 每层 forward，\(n = 2048\)：\(24nd^2 \approx 825\) GFLOPs，\(4n^2d \approx 69\) GFLOPs，attention 只占约 8%，要到 \(n = 6d = 24576\) 才相等。
- 训练 1T token：\(6 \times 6.7\text{B} \times 10^{12} \approx 4\times10^{22}\) FLOPs。
必答点
- 方法：数矩阵乘，\((m\times k)(k\times n) = 2mkn\)
- 每层 \(24nd^2 + 4n^2d\)，即 \(O(n^2d + nd^2)\)；\(n > 6d\) 时 attention 主导
- memory：score 矩阵 \(O(n^2)\)，FlashAttention 降到 \(O(n)\)、FLOPs 不变
- 参数 \(12d^2\)/层；forward \(2N\)、training \(6N\) per token
- decode 有 KV cache：每 token 每层 \(O(nd + d^2)\)
追问
序列很长时怎么办？
restricted attention 只看邻近 \(r\) 个 token，time 降到 \(O(rnd)\)，最长路径变成 \(O(n/r)\)。KV cache 太大就用 MQA / GQA。
和 RNN 比呢？
RNN 每层 \(O(nd^2)\)，但串行步数是 \(O(n)\)。见 RNN vs Transformer。
causal mask 能省一半计算吗？
按块计算的 kernel 可以跳过整块被 mask 的上三角。朴素实现先算满 \(n\times n\) 再加 mask，没有省。
为什么 decode 慢、是 memory-bound？
batch 小时每个 token 都要把全部权重从 GPU memory 读一遍，每个参数只做约 2 FLOPs。arithmetic intensity 见 06。
来源Vaswani et al. 2017 Table 1（p.6）· Kaplan et al. 2020 §2.1 eq. 2.2（p.6）· Dao et al. 2022, FlashAttention Theorem 1（p.5）· Touvron et al. 2023, LLaMA §2.2、Table 2（p.3）