宽度在头之间分割
单个头看到一种模式
活动 67 介绍了缩放点积注意力:查询向量 Q、键向量 K、值向量 V;计算 Q·Kᵀ/√d_k,掩码,softmax,加权 V。一个头学习一种关系模式:可能是主谓一致,可能是标点配对,也可能没什么用。
多头注意力并行运行相同的操作,多次,在令牌表示的不同切片上。十二个并行头。十二种可能的关系模式。头通过训练压力单独专业化;没有架构师告诉头 4 去关注动词时态。
分割关系
ANDREA-120M 设置 d_model = 768 & n_head = 12。多头注意力将 768 分割成 12 个每个 64 的块:
head_dim = d_model / n_head
64 = 768 / 12
每个头都操作64维向量。分割应用得很干净:d_model 必须能被 n_head 整除,没有余数。违反此规则的配置会在配置验证时失败,而不是运行时。
三个模型,三个分割
| 变体 | d_model | n_head | head_dim |
|---|---|---|---|
| ANDREA-12M | 384 | 12 | 32 |
| ANDREA-120M | 768 | 12 | 64 |
| ANDREA-480M | 1536 | 24 | 64 |
注意:ANDREA-12M 和 ANDREA-120M 保持 n_head=12 恒定;仅 d_model 及因此的 head_dim 缩放。ANDREA-480M 将头数加倍至 24,保持 head_dim=64 与 ANDREA-120M 匹配。
计算 head_dim
每个头三个矩阵,或一个大矩阵
每个头的视图
每个头都需要自己的查询投影、键投影和值投影。对于头 h:
Q_h = X · W_Q^h 其中 W_Q^h 的形状为 [d_model, head_dim]
K_h = X · W_K^h 其中 W_K^h 的形状为 [d_model, head_dim]
V_h = X · W_V^h 其中 W_V^h 的形状为 [d_model, head_dim]
X 携带输入形状 [batch, seq_len, d_model]。投影后,Q_h、K_h、V_h 各自携带形状 [batch, seq_len, head_dim]。
融合视图
每个头的矩阵在内存中并排放置。单个融合矩阵 W_Q 形状为 [d_model, d_model],一次性生成所有头:
Q_fused = X · W_Q # [batch, seq_len, d_model]
Q_per_head = reshape(Q_fused) # [batch, n_head, seq_len, head_dim]
融合的矩阵乘法只需一个 BLAS 调用,而不是 12 个。CUDA 张量核心在这种大小的矩阵乘法上达到峰值吞吐量;每个头的矩阵乘法会低效利用硬件。
参数数量
三个融合矩阵 W_Q、W_K、W_V,每个 d_model × d_model。加上输出投影 W_O,也是 d_model × d_model。对于 ANDREA-120M:
每层注意力参数 = 4 × 768² = 2,359,296 ≈ 2.36M
12 层总参数 = 12 × 2.36M ≈ 28.3M
ANDREA-120M 总参数的大约四分之一位于注意力投影中。其余四分之三位于 MLP 子层和嵌入层中。
命名投影矩阵
十二个向量变为一个
每个头计算后
每个头产生一个形状为 [batch, seq_len, head_dim] 的输出张量。十二个头产生十二个这样的张量。沿特征维度进行拼接将它们重新堆叠在一起:
concat_output = concat(head_1, head_2, ..., head_12)
shape = [batch, seq_len, n_head × head_dim]
= [batch, seq_len, 768] # for ANDREA-120M
Concat 将 split 逆转。总特征维度恢复为 d_model。维度上没有信息丢失;区别在于每个块包含的内容:head 1 的块反映了 head 1 学习到的注意力模式。
输出投影 W_O
仅靠 Concatenation 会让 heads 保持隔离:head 4 的输出紧邻 head 7 的输出,它们彼此不知晓。形状为 [d_model, d_model] 的输出投影 W_O 将它们混合:
attention_output = concat_output · W_O
shape = [batch, seq_len, d_model]
在 W_O 之后,每个输出维度都携带所有十二个头的学习线性组合。信息通过这个单一的矩阵乘法在头之间自由流动。
为什么头会专业化
架构中没有任何东西强制头 4 学习动词时态或头 9 学习配对标点。专业化源于梯度压力:在训练过程中,贡献冗余的头接收到的梯度比独特贡献的头小。在数千步后,每个头都会进入一个最有效地降低总损失的利基。
经验上,训练好的 transformer 显示出处理以下内容的头:位置模式(头查看前一个 token)、句法模式(头查看匹配的闭括号)、语义模式(头查看最近的命名实体)。没有标签训练这种专业化。仅训练信号,通过 W_O 传播,就会对头进行排序。
为什么是十二个头,而不是一个更宽的头
CUDA 如何存储 Heads
一个单一的张量,经过重塑
ANDREA 的训练引擎 microgpt_cuda.cu 不会为 12 个 heads 分配 12 个单独的缓冲区。它分配一个融合张量,并将 head 维度视为步幅模式:
// 在 Q = X · W_Q(一个矩阵乘法,跨 heads 融合)之后
// Q 的形状为 [batch, seq_len, d_model]
// 重塑为 [batch, seq_len, n_head, head_dim]
// 转置为 [batch, n_head, seq_len, head_dim]
// 每个头现在在内两个维度上连续
转置将 n_head 移到 seq_len 前面。为什么?因为下一个操作 (Q_h · K_h^T) 需要每个头的 seq_len × head_dim 切片在内存中连续。CUDA 矩阵乘法在连续张量上运行更快。
一个内核,多头
单个注意力 CUDA 内核并行运行于所有头。各线程块处理一个 (batch, head) 对;块内的线程协作处理 seq_len × head_dim 瓦片。内核从不知道它处理多个头;启动网格处理并行性。
配置反映硬件
ANDREA-120M 选择 n_head=12, head_dim=64 与 RTX 4090 张量核心对齐,后者偏好 matmul 瓦片为 16 的倍数。head_dim=64 = 4 × 16 精确匹配瓦片形状。head_dim=32 (ANDREA-12M) 也匹配但未充分利用瓦片。head_dim=72 不匹配且会强制回退内核。
最终图景
| 步骤 | 操作 | 输出形状 |
|---|---|---|
| 1. 项目 | Q = X · W_Q (同理 K, V) | [batch, seq, d_model] |
| 2. 重塑 & 转置 | 将 d_model 分割 → (n_head, head_dim) | [batch, n_head, seq, head_dim] |
| 3. 每个头的注意力 | 每个头上的缩放点积 | [batch, n_head, seq, head_dim] |
| 4. 转置 & 重塑 | 合并 (n_head, head_dim) → d_model | [batch, seq, d_model] |
| 5. 输出投影 | output = concat · W_O | [batch, seq, d_model] |
五个步骤。三次矩阵乘法涉及输入(Q、K、V 投影)。一次矩阵乘法涉及连接的头(W_O)。一个注意力内核并行处理所有头。ANDREA-120M 每层运行所有五个步骤,十二层深,每前向传播一次。