微调LLM前你需要了解的一些概念2--多头注意力机制
1. 先区分两个概念多头和多层Multi-Head Attention 和多层 Transformer Block 不是一回事。一句话区分Multi-Head Attention同一层里多个 attention head 并行看上下文。 多层 Transformer Block很多层 block 串行堆叠一层处理完再传给下一层。更短地记多头 横向多视角 多层 纵向深加工图示如下Block 1 内部Multi-Head AttentionHead 1合并多个 headHead 2Head 3输入 hidden states X0Transformer Block 1Transformer Block 2Transformer Block 3最终 hidden statesBlock 之间是串行的X1 Block1(X0) X2 Block2(X1) X3 Block3(X2)因为第 2 层必须等第 1 层输出之后才能计算。但在一个 Block 内部多个 head 可以并行计算Head 1 ┐ Head 2 ├─ 同时计算 → 合并 → 输出 Head 3 ┘2. 为什么需要 Multi-Head AttentionSelf-Attention 的核心问题是当前 token 应该关注上下文里的哪些 token如果只有一个 head它只能产生一套注意力分布。例如句子小明 把 苹果 放进 书包 因为 它 太小 了模型处理“它”时需要同时理解很多关系1. “它”指代谁 2. “太小”描述的是哪个物体 3. “苹果”和“书包”之间是什么动作关系 4. 句子的主语是谁 5. 哪些词是局部搭配如果只有一个 head它可能给出一套注意力权重书包: 0.70 苹果: 0.15 小明: 0.05 其他: 0.10这套权重主要解决了“它指代书包”这个问题。但同一层里模型还可能需要同时关注动作关系、语法关系、局部搭配等。一个 head 只有一套 softmax 权重很难同时清晰表达多种关系。Multi-Head Attention 的目的就是让同一个 token 在同一层里从多个子空间、多个角度同时理解上下文。可以把多个 head 想成多个小专家Head 1关注指代关系 Head 2关注动作关系 Head 3关注局部搭配 Head 4关注语法结构 Head 5关注长距离依赖同一个 token “它”不同 head 可能关注不同位置Head 1看指代 书包: 0.80, 苹果: 0.10 Head 2看动作链 放进: 0.50, 苹果: 0.20, 书包: 0.20 Head 3看局部搭配 太小: 0.70, 了: 0.20 Head 4看主语结构 小明: 0.40, 把: 0.30, 放进: 0.20最后这些 head 的结果合并起来“它”的表示就会带着更丰富的信息它 ≈ 指代书包 和“太小”形成属性关系 处在“苹果放进书包”的动作结构中 受前面“小明把...”句法结构影响3. 多层 Transformer Block 不能替代多头吗不能完全替代。它们解决的问题不同。如果只有单头但有很多层模型当然也能逐步加工信息第 1 层用一套注意力看上下文 第 2 层再用一套注意力看上一层结果 第 3 层继续加工这像一个人反复读文章很多遍。Multi-Head Attention 像一组人同时从语法、语义、指代、位置等角度读同一句话。两者互补多头同一层内提供多视角并行理解 多层层与层之间逐步抽象、深化、整合图示同一层内多个 head 并行捕捉多种关系多层之间上一层输出作为下一层输入逐步深化表示横向多视角纵向深加工共同提升模型表达能力所以多头不是单纯为了加速而是为了增强表达能力。它让模型在同一层里拥有多套“关注谁”的方案而不是把所有关系挤进一套注意力权重里。4. 每一层都有相同的多头吗通常标准 LLM 中每一层 Transformer Block 都有 Multi-Head Attention。例如一个模型配置可能是num_layers 32 hidden_size 4096 num_attention_heads 32 head_dim 128这意味着一共有 32 层 Transformer Block 每层都有 32 个 attention heads 每个 head 的维度是 128但是要注意每层结构通常相同但参数不共享。也就是说Block 1: 32 个 heads 使用 Wq1, Wk1, Wv1, Wo1 Block 2: 32 个 heads 使用 Wq2, Wk2, Wv2, Wo2 Block 3: 32 个 heads 使用 Wq3, Wk3, Wv3, Wo3第 1 层的 Head 1 和第 20 层的 Head 1 不是同一个 head。它们只是编号相同参数不同输入也不同。因为不同层处理的表示阶段不同第 1 层输入更接近原始 token embedding 第 10 层输入已经融合了一些上下文关系 第 30 层输入更接近高级语义和预测任务所以不同层需要不同的 Q/K/V 参数来处理不同阶段的信息。可以粗略理解为低层 head更偏局部位置、相邻 token、标点、短语搭配 中层 head更偏语法关系、实体关系、指代关系 高层 head更偏任务意图、长距离依赖、输出预测相关信息这不是人工规定而是训练后自然形成的倾向。5. 一个大矩阵如何一次性算出所有 head 的 Q/K/V这是多头注意力最容易卡住的地方。先记住一句话大矩阵负责一次性算出所有 head 的 Q/K/Vreshape 负责把“混在一起的 head 维度”重新分组让程序知道哪一段属于哪个 head。假设seq_len 3 hidden_size 8 num_heads 2 head_dim 4因为hidden_size num_heads × head_dim 2 × 4 8输入 X 的形状是X shape [seq_len, hidden_size] [3, 8]可以理解为 3 个 token每个 token 是 8 维向量token1: [8 维] token2: [8 维] token3: [8 维]如果不用大矩阵你可以想象成分别计算两个 headQ_head1 X Wq_head1 → [3, 4] Q_head2 X Wq_head2 → [3, 4]但工程实现不会一个 head 一个 head 地算而是把两个 head 的投影矩阵拼成一个大矩阵Wq shape [8, 8]一次性计算Q X Wq Q shape [3, 8]这个[3, 8]的最后一维其实可以理解为前 4 维Head 1 的 Q 后 4 维Head 2 的 Q例如Q token1 [h1_1, h1_2, h1_3, h1_4, h2_1, h2_2, h2_3, h2_4]reshape 就是把它从[3, 8]变成[2, 3, 4]含义变成2 个 head 3 个 token 每个 token 在每个 head 里是 4 维6. reshape 的意义是什么reshape 不改变数值只改变数据的组织方式。reshape 前Q shape [3, 8]可以看成token1: [a1 a2 a3 a4 b1 b2 b3 b4] token2: [a5 a6 a7 a8 b5 b6 b7 b8] token3: [a9 a10 a11 a12 b9 b10 b11 b12]reshape 后Q shape [2, 3, 4]变成Head 1: token1: [a1 a2 a3 a4] token2: [a5 a6 a7 a8] token3: [a9 a10 a11 a12] Head 2: token1: [b1 b2 b3 b4] token2: [b5 b6 b7 b8] token3: [b9 b10 b11 b12]所以 reshape 的意义是把原来连在一起的 hidden_size 维向量拆成num_heads × head_dim让每个 head 拿到属于自己的那一段 Q/K/V。K 和 V 也一样K X Wk → [3, 8] → reshape 为 [2, 3, 4] V X Wv → [3, 8] → reshape 为 [2, 3, 4]于是每个 head 都有自己的Head 1: Q1, K1, V1shape 都是 [3, 4] Head 2: Q2, K2, V2shape 都是 [3, 4]7. reshape 后每个 head 如何计算 attention对 Head 1Q1 shape [3, 4] K1 shape [3, 4] V1 shape [3, 4]先计算注意力分数scores1 Q1 K1^T形状是[3, 4] × [4, 3] [3, 3]这个[3, 3]是一张注意力图每个 query token 对每个 key token 的关注分数。例如key token1 key token2 key token3 query token1 0.8 0.1 0.1 query token2 0.2 0.6 0.2 query token3 0.3 0.2 0.5然后缩放并做 softmaxweights1 softmax(scores1 / sqrt(head_dim))再用权重加权 Vout1 weights1 V1形状是[3, 3] × [3, 4] [3, 4]Head 2 同理out2 shape [3, 4]最后把两个 head 的输出拼回去concat(out1, out2) → [3, 8]再经过输出投影矩阵output concat_heads Wo最终输出形状仍然是[3, 8]这样它就可以继续进入残差连接、MLP、下一层 Transformer Block。完整流程如下X: [seq_len, hidden_size] [3, 8]Q X Wq: [3, 8]K X Wk: [3, 8]V X Wv: [3, 8]reshape Q: [num_heads, seq_len, head_dim] [2, 3, 4]reshape K: [2, 3, 4]reshape V: [2, 3, 4]Head 1 AttentionHead 2 Attentionconcat heads: [3, 8]输出投影 WoAttention 输出: [3, 8]8. 加上 batch 后的真实形状真实模型中通常有 batch。输入形状X shape [batch, seq_len, hidden_size]例如batch 2 seq_len 3 hidden_size 8 num_heads 2 head_dim 4那么X shape [2, 3, 8]一次性计算 QQ X Wq Q shape [2, 3, 8]reshapeQ shape [2, 3, 2, 4]通常还会 transpose 一下把 head 维度提前Q shape [2, 2, 3, 4]含义是[batch, num_heads, seq_len, head_dim]也就是2 个样本 2 个 head 3 个 token 每个 head 4 维K 和 V 同理K shape [2, 2, 3, 4] V shape [2, 2, 3, 4]然后批量计算scores Q K.transpose(-2, -1)形状变化[2, 2, 3, 4] [2, 2, 4, 3] [2, 2, 3, 3]含义是每个 batch、每个 head都有一张 seq_len × seq_len 的注意力图。再做weights softmax(scores / sqrt(head_dim)) out weights V形状[2, 2, 3, 3] [2, 2, 3, 4] [2, 2, 3, 4]最后转置并 reshape 回[2, 2, 3, 4] → [2, 3, 2, 4] → [2, 3, 8]这就回到了[batch, seq_len, hidden_size]9. 为什么这样能并行reshape 后head 维度变成了张量中的一个独立维度[batch, num_heads, seq_len, head_dim]GPU 可以把batch 维度 head 维度 token 维度一起打包成大规模矩阵运算。它不是在 Python 里慢慢循环for head in heads: attention(head)而是一次性执行所有 batch 的所有 head 的 QK^T 所有 batch 的所有 head 的 softmax 所有 batch 的所有 head 的 weights V这就是多头注意力中的“并行”。图示Q/K/V: [batch, num_heads, seq_len, head_dim]批量矩阵乘法 QK^Tscores: [batch, num_heads, seq_len, seq_len]softmaxweights Vhead outputs: [batch, num_heads, seq_len, head_dim]合并 heads: [batch, seq_len, hidden_size]10. 为什么不是直接用一个更大的 head一个常见疑问是如果总维度一样为什么不用一个大 head而要拆成多个小 head例如一个 head维度 12和三个 head每个维度 4总维度看起来一样。关键区别在于一个大 head 通常只有一套注意力分布多个 head 有多套独立的注意力分布。一个 head只有一张 attention map多个 headHead 1 一张 attention map Head 2 一张 attention map Head 3 一张 attention map这意味着多头不是简单拆维度而是让模型同时拥有多套“关注谁”的方案。对同一个 token “它”Head 1关注“书包”解决指代 Head 2关注“太小”解决属性 Head 3关注“放进”解决动作关系 Head 4关注“小明”解决主语结构如果只有一个大 head这些关系要挤在一张注意力图里表达会更受限。11. 多头注意力完整流程总结完整流程可以压缩成1. 输入 X: [batch, seq_len, hidden_size] 2. 用大矩阵一次性算 Q/K/V 3. Q/K/V reshape 成 [batch, num_heads, seq_len, head_dim] 4. 每个 head 独立计算 attention map 5. 每个 head 得到自己的输出 6. 把所有 head 的输出 concat 回 hidden_size 7. 经过输出投影 Wo 8. 得到 Multi-Head Attention 输出图示X: [B, S, H]线性投影得到 Q/K/V: [B, S, H]reshape transpose: [B, heads, S, head_dim]QK^T / sqrt(head_dim): [B, heads, S, S]可选 causal masksoftmax 得到注意力权重weights V: [B, heads, S, head_dim]transpose reshape 合并 heads: [B, S, H]输出投影 WoAttention 输出: [B, S, H]其中B batch S seq_len H hidden_size head_dim H / num_heads12. 面试回答模板Q1Multi-Head Attention 和多层 Transformer Block 是一回事吗不是。Multi-Head Attention 是同一层内部的多个 attention head 并行计算用于从多个子空间捕捉不同关系多层 Transformer Block 是多个 block 串行堆叠后一层依赖前一层输出用于逐层深化表示。Q2为什么需要多头注意力一个 attention head 只有一套注意力分布很难同时表达指代、语法、局部搭配、实体关系等多种关系。多头注意力允许模型在同一层中并行学习多套注意力模式再合并成更丰富的上下文表示。Q3多层 Transformer Block 能替代多头吗不能完全替代。多层提供纵向的逐步抽象多头提供横向的多视角信息交互。单头多层像一个人反复读文章多头像多人同时从不同角度读文章两者是互补关系。Q4每一层都有相同的多头吗多数标准 LLM 中每一层通常有相同数量的 attention heads例如每层 32 个 head。但不同层的 Q/K/V/O 参数不共享因此第 1 层的 Head 1 和第 20 层的 Head 1 不是同一个东西它们会学习不同层次的关系。Q5reshape 的意义是什么大矩阵一次性算出的 Q/K/V 形状通常是[batch, seq_len, hidden_size]。reshape 把最后的 hidden_size 维拆成num_heads × head_dim变成[batch, num_heads, seq_len, head_dim]让每个 head 拿到自己对应的 Q/K/V 子空间并独立计算 attention。Q6为什么这种实现可以并行reshape 后head 维度成为张量的批处理维度之一。GPU 可以同时计算所有 batch、所有 head 的矩阵乘法、softmax 和加权求和而不是一个 head 一个 head 地循环执行。13. 最短记忆版Multi-Head Attention 同一层里多个 head 并行计算多套注意力。 多层 Transformer Block 多个 block 串行堆叠逐层深化表示。 为什么需要多头 一个 head 只有一套注意力分布多头可以同时关注指代、语法、动作、局部搭配等关系。 每层是否相同 通常每层 head 数量相同但参数不共享学到的功能也不同。 大矩阵算 QKV 一次性算出所有 head 的 Q/K/V提高计算效率。 reshape 把 hidden_size 拆成 num_heads × head_dim让每个 head 拿到自己的子空间。 并行 GPU 同时计算所有 batch、所有 head 的 attention。