计算MHA的时间复杂度

计算MHA的时间复杂度

1 矩阵乘的时间复杂度公式

　　L*M   *  M  * N = out L*N； 时间复杂度O(L*M*N);

2 MHA需要计算的部分

　　1 3个W； 2 QK；3sfmx；4score*V；5结果*Wout；

3 全部时间复杂度

　　1 Step 1：生成 Q、K、V 投影

X (L×d) · W_Q (d×d_k) → Q (L×d_k) 复杂度：O(L·d·d_k)

X (L×d) · W_K (d×d_k) → K (L×d_k) 复杂度：O(L·d·d_k)

X (L×d) · W_V (d×d_v) → V (L×d_v) 复杂度：O(L·d·d_v)

若 d_k = d_v = d（简化单头情况）：三者均为 O(L·d²)

　　2 Step 2：QK^T 注意力分数

Q (L×d_k) · K^T (d_k×L) → Score (L×L)

复杂度：O(L · d_k · L) = O(L²·d_k)

你写的 O(L²) 是把 d_k 省略了，严格来说是 O(L²·d_k)

　　Step 3：Softmax

Softmax(L×L) → (L×L)

逐行做 softmax，复杂度：O(L²)（通常被上下项吸收，可忽略）

　　Step 4：Attention × V（你这步算错了）

Attn (L×L) · V (L×d_v) → Output (L×d_v)

根据矩阵乘法公式 O(m·n·p)：

• m = L, n = L, p = d_v
• 复杂度 = O(L · L · d_v) = O(L²·d_v)

你写的 O(L·dim) 是错误的，正确是 O(L²·d_v)

　　Step 5：输出投影 W_out

Output (L×d_v) · W_out (d_v×d) → Final (L×d)

复杂度：O(L · d_v · d) = O(L·d²)（当 d_v = d 时）

4 全部相加

　　3O(L*d*d) + O(L²·d) + O(L²) + O(L²·d) + O(L·d²) = 2 O(L²·d) + 4 O(L·d²) + O(L²) 

5 缩减

          

 6 优化

   当 L >> d（超长文本）：O(L²·d) 主导 → Attention 是瓶颈

   当 d >> L（短文本大模型）：O(L·d²) 主导 → 线性变换是瓶颈

    这也正是 IndexCache 论文针对长上下文（L 很大时）优化 indexer 的 O(L²) 项如此有效的根本原因。