Skip to main content

multi head attention

Aaron L9/10/25About 4 minMachine LearningMachine LearningAttention

多头注意力（Multi-Head Attention）数学详解

多头注意力是 Transformer 的核心机制，它通过并行执行多个独立的注意力头（attention heads）来捕捉输入的不同表示子空间。以下是完整的数学推导（所有公式用 $ 包裹）：

1. 输入表示

输入序列矩阵：

X∈Rn×dmodelX \in \mathbb{R}^{n \times d_{\text{model}}}

其中：

  • nn：序列长度
  • dmodeld_{\text{model}}：模型维度（如 512）

2. 多头投影

对每个头 i∈{1,2,…,h}i \in \{1, 2, \dots, h\}，使用独立的权重矩阵进行投影：

Qi=X⋅WiQ,WiQ∈Rdmodel×dkKi=X⋅WiK,WiK∈Rdmodel×dkVi=X⋅WiV,WiV∈Rdmodel×dv\begin{align*} Q_i &= X \cdot W_i^Q, \quad &W_i^Q \in \mathbb{R}^{d_{\text{model}} \times d_k} \\ K_i &= X \cdot W_i^K, \quad &W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k} \\ V_i &= X \cdot W_i^V, \quad &W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v} \end{align*}

参数说明：

  • hh：头数（如 8）
  • dk=dv=dmodelhd_k = d_v = \frac{d_{\text{model}}}{h}：每个头的维度（如 64）
  • 投影后维度：Qi,Ki∈Rn×dkQ_i, K_i \in \mathbb{R}^{n \times d_k}，Vi∈Rn×dvV_i \in \mathbb{R}^{n \times d_v}

3. 单头注意力计算

对每个头 ii 独立计算缩放点积注意力：

headi=Attention(Qi,Ki,Vi)=softmax(QiKiTdk)Vi\text{head}_i = \text{Attention}(Q_i, K_i, V_i) = \text{softmax}\left( \frac{Q_i K_i^T}{\sqrt{d_k}} \right) V_i

其中：

  • QiKiTdk∈Rn×n\frac{Q_i K_i^T}{\sqrt{d_k}} \in \mathbb{R}^{n \times n}：缩放后的注意力分数矩阵
  • 输出维度：headi∈Rn×dv\text{head}_i \in \mathbb{R}^{n \times d_v}

4. 多头输出拼接

将所有头的输出在特征维度拼接：

MultiHead(Q,K,V)=Concat(head1,head2,…,headh)\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, \dots, \text{head}_h)

拼接后维度：Concat(⋯ )∈Rn×(h⋅dv)=Rn×dmodel\text{Concat}(\cdots) \in \mathbb{R}^{n \times (h \cdot d_v)} = \mathbb{R}^{n \times d_{\text{model}}}

5. 输出线性变换

将拼接结果通过可学习权重矩阵 WOW^O 投影回模型维度：

Output=Concat(head1,…,headh)⋅WO,WO∈Rdmodel×dmodel\text{Output} = \text{Concat}(\text{head}_1, \dots, \text{head}_h) \cdot W^O, \quad W^O \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}

最终输出维度：Rn×dmodel\mathbb{R}^{n \times d_{\text{model}}}

完整数学公式汇总

MultiHead(X)=Concat(head1,…,headh)⋅WOwhereheadi=softmax((XWiQ)(XWiK)Tdk)(XWiV)with{dk=dv=dmodelhWiQ,WiK∈Rdmodel×dkWiV∈Rdmodel×dvWO∈Rdmodel×dmodel\boxed{ \begin{align*} & \text{MultiHead}(X) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) \cdot W^O \\ & \text{where} \quad \text{head}_i = \text{softmax}\left( \frac{(X W_i^Q)(X W_i^K)^T}{\sqrt{d_k}} \right) (X W_i^V) \\ & \text{with} \quad \begin{cases} d_k = d_v = \dfrac{d_{\text{model}}}{h} \\ W_i^Q, W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k} \\ W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v} \\ W^O \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}} \end{cases} \end{align*} }

维度变化可视化

输入 X: [n, d_model]
│
├─ Head 1 ──┐
│ Q1 = X·W1^Q [n, d_k]  │
│ K1 = X·W1^K [n, d_k]  → Attn1 [n, n] → head1 [n, d_v]
│ V1 = X·W1^V [n, d_v]  │               │
├─ Head 2 ──┤                              │
│ ...      │                              ├→ Concat [n, h·d_v] = [n, d_model]
├─ Head h ──┤                              │
│          │                              │
└──────────┘                              │
                                          │
Output = Concat·W^O [n, d_model] ←────────┘

数学优势分析

  1. 子空间分解
    每个头学习不同的投影：

    WiQ,WiK,WiV 相互独立W_i^Q, W_i^K, W_i^V \text{ 相互独立}

    使模型关注不同方面的信息（如语法/语义/位置）。

  2. 计算效率
    虽然头数 hh 增加，但单头维度 dkd_k 减小：

    单头计算量=O(n2dk)=O(n2dmodelh)\text{单头计算量} = O(n^2 d_k) = O\left(n^2 \frac{d_{\text{model}}}{h}\right)

    总计算量 O(n2dmodel)O(n^2 d_{\text{model}}) 与单头相同。

  3. 表达能力增强
    输出是多个子空间的非线性组合：

    Output=f(∑i=1hheadi⋅WiO)\text{Output} = f\left( \sum_{i=1}^h \text{head}_i \cdot W^O_i \right)

    （WOW^O 隐含分解为子矩阵 WiOW^O_i）

实例演算

设 dmodel=4d_{\text{model}} = 4, h=2h = 2, 则 dk=dv=2d_k = d_v = 2
输入 X=[10.5−12−0.510.3−2]X = \begin{bmatrix} 1 & 0.5 & -1 & 2 \\ -0.5 & 1 & 0.3 & -2 \end{bmatrix}

头1计算：

W1Q=[0.10.4−0.20.31.0−0.50.50.2], Q1=XW1Q=[0.85−0.05−0.290.74]W_1^Q = \begin{bmatrix} 0.1 & 0.4 \\ -0.2 & 0.3 \\ 1.0 & -0.5 \\ 0.5 & 0.2 \end{bmatrix}, \ Q_1 = X W_1^Q = \begin{bmatrix} 0.85 & -0.05 \\ -0.29 & 0.74 \end{bmatrix}

（类似计算 K1K_1, V1V_1 后求 head1\text{head}_1）

头2计算：

W2Q=[−0.30.21.10.6−0.40.80.7−0.1], Q2=XW2Q=[2.150.65−1.321.02]W_2^Q = \begin{bmatrix} -0.3 & 0.2 \\ 1.1 & 0.6 \\ -0.4 & 0.8 \\ 0.7 & -0.1 \end{bmatrix}, \ Q_2 = X W_2^Q = \begin{bmatrix} 2.15 & 0.65 \\ -1.32 & 1.02 \end{bmatrix}

拼接与输出：

Concat=[head1(1)head1(2)head2(1)head2(2)]\text{Concat} = \begin{bmatrix} \text{head}_1^{(1)} & \text{head}_1^{(2)} & \text{head}_2^{(1)} & \text{head}_2^{(2)} \end{bmatrix}

WO=[0.5−0.21.10.30.10.8−0.40.6−0.30.70.21.00.9−0.50.3−0.8], Output=Concat⋅WOW^O = \begin{bmatrix} 0.5 & -0.2 & 1.1 & 0.3 \\ 0.1 & 0.8 & -0.4 & 0.6 \\ -0.3 & 0.7 & 0.2 & 1.0 \\ 0.9 & -0.5 & 0.3 & -0.8 \end{bmatrix}, \ \text{Output} = \text{Concat} \cdot W^O

与传统单头对比

特性单头注意力多头注意力
参数量3dmodel23d_{\text{model}}^23dmodel23d_{\text{model}}^2 (相同)
计算复杂度O(n2dmodel)O(n^2 d_{\text{model}})O(n2dmodel)O(n^2 d_{\text{model}})
表达能力单一表示空间hh 个正交子空间
并行性低高（头间完全并行）

为什么有效？

  1. 多样性：不同头关注不同模式（如局部依赖 vs 全局依赖）

    headi 学习独立的关系表示\text{head}_i \text{ 学习独立的关系表示}

  2. 维度分解：将高维空间分解为低维子空间，降低学习难度

    Rdmodel→⊕i=1hRdk\mathbb{R}^{d_{\text{model}}} \to \oplus_{i=1}^h \mathbb{R}^{d_k}

  3. 残差连接：多头输出通过残差连接传递到前馈网络，保留原始信息：

    LayerOutput=LayerNorm(X+MultiHead(X))\text{LayerOutput} = \text{LayerNorm}(X + \text{MultiHead}(X))

多头注意力通过这种"分治策略"，显著提升了模型捕捉复杂依赖关系的能力，成为 Transformer 架构的核心创新。� Transformer 架构的核心创新。