Transformer多头注意力的计算量和单头注意力计算量比较

多头注意力机制（Multi-Head Attention）在Transformer中引入了多个并行的注意力头，每个注意力头可以学习到不同的特征表示。尽管这种机制增强了模型的表达能力，但也增加了一些计算量。下面详细比较一下多头注意力和单头注意力的计算量。

单头注意力计算量

假设输入序列的长度为 \( n \)，词嵌入维度为 \( d_{\text{model}} \)。

1. **线性变换**：
   将输入矩阵 \( \mathbf{X} \) 转换为查询 \( \mathbf{Q} \)、键 \( \mathbf{K} \)、值 \( \mathbf{V} \) 矩阵，每个变换的计算量为：
   \[
   \mathcal{O}(n \times d_{\text{model}} \times d_{\text{model}})
   \]
   由于有三个这样的变换，总计算量为：
   \[
   3 \times \mathcal{O}(n \times d_{\text{model}}^2)
   \]

2. **注意力得分计算**：
   计算查询和键的点积，得到注意力得分矩阵 \( \mathbf{Q} \mathbf{K}^T \)，计算量为：
   \[
   \mathcal{O}(n \times d_{\text{model}} \times n)
   \]

3. **softmax和加权求和**：
   计算注意力权重并对值进行加权求和，计算量为：
   \[
   \mathcal{O}(n^2 \times d_{\text{model}})
   \]

多头注意力计算量

假设有 \( h \) 个注意力头，每个头的维度为 \( d_k = \frac{d_{\text{model}}}{h} \)。

1. **线性变换**：
   将输入矩阵 \( \mathbf{X} \) 转换为每个头的查询 \( \mathbf{Q} \)、键 \( \mathbf{K} \)、值 \( \mathbf{V} \) 矩阵，每个变换的计算量为：
   \[
   \mathcal{O}(n \times d_{\text{model}} \times d_k)
   \]
   总计算量为：
   \[
   3 \times h \times \mathcal{O}(n \times d_{\text{model}} \times d_k) = 3 \times \mathcal{O}(n \times d_{\text{model}}^2)
   \]

2. **注意力得分计算**：
   计算每个头的查询和键的点积，得到注意力得分矩阵 \( \mathbf{Q}_i \mathbf{K}_i^T \)，每个头的计算量为：
   \[
   \mathcal{O}(n \times d_k \times n)
   \]
   总计算量为：
   \[
   h \times \mathcal{O}(n \times d_k \times n) = \mathcal{O}(n^2 \times d_{\text{model}})
   \]

3. **softmax和加权求和**：
   计算每个头的注意力权重并对值进行加权求和，计算量为：
   \[
   h \times \mathcal{O}(n^2 \times d_k) = \mathcal{O}(n^2 \times d_{\text{model}})
   \]

4. **拼接和线性变换**：
   将所有头的输出拼接并进行一次线性变换，总计算量为：
   \[
   \mathcal{O}(n \times d_{\text{model}} \times d_{\text{model}})
   \]

总计算量比较

- **单头注意力**的总计算量为：
  \[
  3 \times \mathcal{O}(n \times d_{\text{model}}^2) + \mathcal{O}(n^2 \times d_{\text{model}})
  \]

- **多头注意力**的总计算量为：
  \[
  3 \times \mathcal{O}(n \times d_{\text{model}}^2) + \mathcal{O}(n^2 \times d_{\text{model}}) + \mathcal{O}(n \times d_{\text{model}}^2)
  \]
可以简化为：
  \[
  4 \times \mathcal{O}(n \times d_{\text{model}}^2) + \mathcal{O}(n^2 \times d_{\text{model}})
  \]

从上述计算可以看出，多头注意力机制在计算量上增加了一倍的线性变换计算量，但主要的计算瓶颈仍然在 \(\mathcal{O}(n^2 \times d_{\text{model}})\) 这一项上，即注意力得分的计算和softmax的加权求和。这表明多头注意力虽然计算量更大，但并不会成倍增加复杂度，而是通过增加计算资源来增强模型的表达能力和学习不同特征的能力。