---

站长提醒：注意力机制中的计算复杂度是一个非常关键的知识点，是深入理解transformer的核心知识，希望大家认真掌握。

1、Q 的形状是 (n, d)，K^T 的形状是 (d, n)，两者相乘结果是 (n, n)，计算复杂度是 O(n² · d)。问：为什么计算复杂度是 n² · d 呢？

矩阵乘法 \(Q(n,d) \times K^T(d,n)\) 得到结果矩阵 \((n,n)\)。我们来数一下到底做了多少次乘法。

1.1、结果矩阵的每个元素

结果矩阵 \(S = QK^T\) 是 \(n \times n\)，其中第 \(i\) 行第 \(j\) 列的元素是：

$$ S_{ij} = \sum_{k=1}^{d} Q_{ik} \cdot K_{jk} $$

也就是 \(Q\) 的第 \(i\) 行（长度 \(d\)）和 \(K^T\) 的第 \(j\) 列（长度 \(d\)）做点积。

计算一个 \(S_{ij}\) 需要：\(d\) 次乘法 + \((d-1)\) 次加法，量级为 \(O(d)\)。

1.2、总共有多少个元素？

结果矩阵有 \(n \times n = n^2\) 个元素。

1.3、总量

$$ \underbrace{n^2}_{\text{元素个数}} \times \underbrace{O(d)}_{\text{每个元素的计算量}} = O(n^2 d) $$

1.4、直观理解

维度	含义	数量
行	结果的行数	\(n\)
列	结果的列数	\(n\)
每个元素内部	点积长度	\(d\)

• 外层循环：遍历 \(n\) 行
• 中层循环：遍历 \(n\) 列
• 内层循环：累加 \(d\) 次

三重循环嵌套：

for i in range(n):        # n 次
    for j in range(n):    # n 次
        s = 0
        for k in range(d):  # d 次
            s += Q[i][k] * K[j][k]
        S[i][j] = s

总乘法次数 = \(n \times n \times d = n^2 d\)。

1.5、一句话总结

结果矩阵有 \(n^2\) 个元素，每个元素都要做一次长度为 \(d\) 的点积，所以是 \(O(n^2 d)\)。

这也是为什么自注意力在序列长度 \(n\) 很大时开销急剧上升——它对 \(n\) 是平方级的。

2、对 (n, n) 矩阵做 softmax，计算复杂度是 O(n²)。问：计算复杂度为什么是n²？

Softmax 是对矩阵逐元素做的，每个元素独立计算，所以计算量就是元素个数。

2.1、Softmax 怎么算

对 \((n,n)\) 矩阵 \(S\)，softmax 通常是按行做的（每一行做一次 softmax），但无论按行还是逐元素，总计算量都是看有多少个元素。

对一行 \(s_1, s_2, \dots, s_n\)：

$$ \text{softmax}(s_i) = \frac{e^{s_i}}{\sum_{j=1}^{n} e^{s_j}} $$

分三步：

1. 求指数：对每个元素算 \(e^{s_i}\) → \(n\) 次
2. 求和：把这一行的 \(n\) 个指数加起来 → \(n-1\) 次加法
3. 归一化：每个元素除以这个和 → \(n\) 次除法

每行大约 \(O(n)\) 的计算量。

2.2、总共有多少行？

有 \(n\) 行。

$$ \underbrace{n}_{\text{行数}} \times \underbrace{O(n)}_{\text{每行计算量}} = O(n^2) $$

2.3、换个角度：逐元素看

更直接地：

• 矩阵一共有 \(n \times n = n^2\) 个元素
• 每个元素至少要做一次指数运算（\(e^{s_{ij}}\)）
• 再参与一次归一化（除以分母）

所以至少是 \(n^2\) 次基本运算。

2.4、三重循环视角

for i in range(n):          # 每一行
    # 第一步：算指数
    for j in range(n):      # n 次
        E[i][j] = exp(S[i][j])
    # 第二步：求和
    total = sum(E[i])       # n 次
    # 第三步：归一化
    for j in range(n):      # n 次
        P[i][j] = E[i][j] / total

外层 \(n\) 行，内层每行 \(O(n)\)，总共 \(O(n^2)\)。

2.5、和矩阵乘法的对比

操作	结果尺寸	每个元素计算量	总量
\(QK^T\)	\(n \times n\)	\(O(d)\)（点积）	\(O(n^2 d)\)
softmax	\(n \times n\)	\(O(1)\)（指数/除法）	\(O(n^2)\)

关键区别：矩阵乘法每个元素要做长度为 \(d\) 的点积，而 softmax 每个元素只做常数次运算，所以 softmax 少了一个 \(d\) 因子。

2.6、一句话总结

softmax 是逐元素的常数级运算，矩阵有 \(n^2\) 个元素，所以总量是 \(O(n^2)\)。

3、softmax(...)V 的计算结果为(n, n) × (n, d) → (n, d)，问：计算复杂度为什么是 O(n2 · d)？

这是矩阵乘法 \((n,n) \times (n,d) \to (n,d)\)，和第一次的 \(QK^T\) 结构一样，只是维度角色换了。

3.1、结果矩阵的每个元素

设 softmax 结果为 \(P\)（\(n \times n\)），\(V\) 为 \((n,d)\)，结果是 \(O = PV\)（\(n \times d\)），其中第 \(i\) 行第 \(j\) 列：

$$ O_{ij} = \sum_{k=1}^{n} P_{ik} \cdot V_{kj} $$

也就是 \(P\) 的第 \(i\) 行（长度 \(n\)）和 \(V\) 的第 \(j\) 列（长度 \(n\)）做点积。

计算一个 \(O_{ij}\) 需要：\(n\) 次乘法 + \((n-1)\) 次加法，量级为 \(O(n)\)。

3.2、总共有多少个元素？

结果矩阵 \(O\) 有 \(n \times d\) 个元素。

3.3、总量

$$ \underbrace{n \times d}_{\text{元素个数}} \times \underbrace{O(n)}_{\text{每个元素的计算量}} = O(n^2 d) $$

3.4、三重循环视角

for i in range(n):          # 结果的行数 n
    for j in range(d):      # 结果的列数 d
        s = 0
        for k in range(n):  # 点积长度 n
            s += P[i][k] * V[k][j]
        O[i][j] = s

总乘法次数 = \(n \times d \times n = n^2 d\)。

3.5、和第一次矩阵乘法对比

操作	左矩阵	右矩阵	结果	点积长度	计算量
\(QK^T\)	\((n,d)\)	\((d,n)\)	\((n,n)\)	\(d\)	\(n^2 d\)
\(PV\)	\((n,n)\)	\((n,d)\)	\((n,d)\)	\(n\)	\(n^2 d\)

两次都是 \(O(n^2 d)\)，但来源不同：

• \(QK^T\)：结果有 \(n^2\) 个元素，每个点积长度 \(d\)
• \(PV\)：结果有 \(n \cdot d\) 个元素，每个点积长度 \(n\)

一个 \(n^2 \times d\)，一个 \(nd \times n\)，乘积都是 \(n^2 d\)。

3.6、一句话总结

结果矩阵有 \(n \cdot d\) 个元素，每个元素要做一次长度为 \(n\) 的点积，所以是 \(n \cdot d \cdot n = n^2 d\)。

4、整个注意力机制汇总

步骤	操作	计算量
1	\(QK^T\)	\(O(n^2 d)\)
2	softmax	\(O(n^2)\)
3	\(PV\)	\(O(n^2 d)\)

总计算量 \(O(n^2 d)\)，其中 \(n\) 是序列长度，\(d\) 是每个头的维度。对 \(n\) 是平方级，这是自注意力在长序列上开销大的根本原因。

---

---