English· Español· Deutsch· Nederlands· Français· 日本語· ქართული· 繁體中文· 简体中文· Português· Русский· العربية· हिन्दी· Italiano· 한국어· Polski· Svenska· Türkçe· Українська· Tiếng Việt· Bahasa Indonesia

un

访客
1 / ?
返回课程列表

宽度在头之间分割

单个头看到一种模式

活动 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_modeln_headhead_dim
ANDREA-12M3841232
ANDREA-120M7681264
ANDREA-480M15362464

注意:ANDREA-12M 和 ANDREA-120M 保持 n_head=12 恒定;仅 d_model 及因此的 head_dim 缩放。ANDREA-480M 将头数加倍至 24,保持 head_dim=64 与 ANDREA-120M 匹配。

计算 head_dim

假设一个假设的 ANDREA-240M 变体使用 `d_model=1024` 和 `n_head=16`。计算 `head_dim`。然后用一句话解释如果一位同事提议 `d_model=1024` 和 `n_head=15`,在配置验证时会发生什么。

每个头三个矩阵,或一个大矩阵

每个头的视图

每个头都需要自己的查询投影、键投影和值投影。对于头 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_QW_KW_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 子层和嵌入层中。

命名投影矩阵

ANDREA-120M 具有 `d_model=768` 和 `n_head=12`。请命名一个 Transformer 块注意力子层中的四个投影矩阵,并给出每个矩阵的形状(行 × 列)。然后说明其中哪个使用输入 X,哪个使用串联的头输出。

十二个向量变为一个

Multi-Head Attention Flow


每个头计算后

每个头产生一个形状为 [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 传播,就会对头进行排序。

为什么是十二个头,而不是一个更宽的头

想象一个替代方案:用 1 个 768 维的头替换每个 64 维的 12 个头。投影中的总参数保持不变。给出两个不同的理由,解释为什么 12 头配置在语言建模任务上优于 1 头配置。在你的回答中引用注意力模式或训练动态。

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 每层运行所有五个步骤,十二层深,每前向传播一次。

阅读管道

逐步说明单个标记的表示 X[i](一个 768 维向量)在 ANDREA-120M 的一个多头注意力子层中流动时会发生什么。提及:它经历的投影、如何分割、注意力对它做了什么、以及如何返回到 768 维向量。在每个步骤显示维度。