984 字

5 分钟

多头注意力

2026-04-22

Transformer

/

Multi-Head Attention

/

多头注意力

多头注意力 (Multi-Head Attention)#

1. 为什么需要多头？#

类比：分析一篇文章时，一个人可能关注语法结构，另一个关注情感色彩，第三个关注事实关系。多头注意力就是让模型同时从多个角度分析同一句话，每个”头”学习捕捉不同类型的关系。

单头注意力只有一组 $W_Q, W_K, W_V$ ，只能学到一种关注模式。实际上，词与词之间存在多种关系：

关系类型	示例
语法关系	”猫”→“坐在”（主语→谓语）
指代关系	”它”→“猫”（代词→指代对象）
修饰关系	”美丽的”→“花园”（形容词→名词）
语义关系	”国王”→“王后”（同领域词）

一个头很难同时捕捉所有这些关系——多头让每个头专注于一种模式。

2. 计算流程#

公式#

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) \cdot W_O$

其中每个头：

$\text{head}_i = \text{Attention}(X W_{Q_i},\; X W_{K_i},\; X W_{V_i})$

3. 参数量分析#

原论文中 $d_{model} = 512$ ， $h = 8$ 个头：

$d_k = d_v = \frac{d_{model}}{h} = \frac{512}{8} = 64$

参数矩阵	形状	数量	参数量
$W_{Q_i}$	$(512, 64)$	8 个	$8 \times 512 \times 64 = 262,144$
$W_{K_i}$	$(512, 64)$	8 个	$262,144$
$W_{V_i}$	$(512, 64)$	8 个	$262,144$
$W_O$	$(512, 512)$	1 个	$262,144$
合计			1,048,576 ≈ 1M

巧妙的设计
每个头的维度，总参数量 = ，和用单个大头完全一样！多头注意力不增加计算量，但表达能力更强——这是一笔”免费的午餐”。

4. 实际中怎么高效计算？#

实际实现中不会真的创建 $h$ 组独立的投影矩阵，而是用一个大矩阵一次投影，再 reshape 拆分：

1
# 伪代码：高效的多头注意力
2
Q = X @ W_Q  # (T, d_model) → (T, d_model)
3
Q = Q.reshape(T, h, d_k).transpose(0, 1)  # → (h, T, d_k)
4
# K, V 同理
5

6
# 所有头并行计算注意力
7
attn_output = scaled_dot_product_attention(Q, K, V)  # (h, T, d_v)
8

9
# 拼接并投影回 d_model
10
output = attn_output.transpose(0, 1).reshape(T, d_model) @ W_O

5. 代码实现#

1
import subprocess
2
subprocess.check_call(["pip", "install", "numpy"])
3
import numpy as np
4

5
def softmax(x, axis=-1):
6
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
7
    return e_x / np.sum(e_x, axis=axis, keepdims=True)
8

9
def multi_head_attention(X, W_Q, W_K, W_V, W_O, h, mask=None):
10
    """Multi-Head Attention (纯 NumPy 实现)
11

12
    Args:
13
        X: (T, d_model) 输入
14
        W_Q, W_K, W_V: (d_model, d_model) 投影矩阵
15
        W_O: (d_model, d_model) 输出投影
16
        h: 头数
17
        mask: (T, T) 可选掩码
18

19
    Returns:
20
        output: (T, d_model)
21
    """
22
    T, d_model = X.shape
23
    d_k = d_model // h
24

25
    # 投影
26
    Q = X @ W_Q  # (T, d_model)
27
    K = X @ W_K
28
    V = X @ W_V
29

30
    # 拆分成 h 个头: (T, d_model) → (h, T, d_k)
31
    Q = Q.reshape(T, h, d_k).transpose(1, 0, 2)  # (h, T, d_k)
32
    K = K.reshape(T, h, d_k).transpose(1, 0, 2)
33
    V = V.reshape(T, h, d_k).transpose(1, 0, 2)
34

35
    # 每个头独立计算注意力
36
    scores = Q @ K.transpose(0, 2, 1) / np.sqrt(d_k)  # (h, T, T)
37
    if mask is not None:
38
        scores = np.where(mask == 0, -1e9, scores)
39
    weights = softmax(scores, axis=-1)
40
    attn_out = weights @ V  # (h, T, d_k)
41

42
    # 拼接: (h, T, d_k) → (T, d_model)
43
    concat = attn_out.transpose(1, 0, 2).reshape(T, d_model)
44

45
    # 输出投影
46
    output = concat @ W_O  # (T, d_model)
47

48
    return output, weights
49

50
# ========== 测试 ==========
51
np.random.seed(42)
52
T, d_model, h = 4, 8, 2  # 4个词, 8维, 2个头
53

54
X = np.random.randn(T, d_model)
55
W_Q = np.random.randn(d_model, d_model) * 0.1
56
W_K = np.random.randn(d_model, d_model) * 0.1
57
W_V = np.random.randn(d_model, d_model) * 0.1
58
W_O = np.random.randn(d_model, d_model) * 0.1
59

60
output, weights = multi_head_attention(X, W_Q, W_K, W_V, W_O, h)
61

62
print(f"输入形状: {X.shape}")
63
print(f"输出形状: {output.shape}")
64
print(f"\n头 1 的注意力权重:\n{weights[0].round(3)}")
65
print(f"\n头 2 的注意力权重:\n{weights[1].round(3)}")
66
print("\n观察：两个头学到了不同的注意力模式")

6. 不同头学到了什么？#

研究发现，多头注意力中的不同头确实会专注于不同类型的模式：

注意力头的冗余性
并非所有头都同样重要。研究（Michel et al., 2019）发现，在训练好的模型中修剪掉 20-40% 的头，性能下降很小。这说明存在一定冗余，但”冗余”本身也是一种鲁棒性保障。

多头注意力 (Multi-Head Attention)#

1. 为什么需要多头？#

2. 计算流程#

公式#

3. 参数量分析#

4. 实际中怎么高效计算？#

5. 代码实现#

6. 不同头学到了什么？#

相关笔记#