多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

VLA 系统学习第 15 课:为什么一个 Attention Head 还不够?——Multi-Head Attention 与 Transformer Block

VLA 系统学习第 15 课:为什么一个 Attention Head 还不够?——Multi-Head Attention 与 Transformer Block 第十四课标准答案先把上一课的 12 道题收掉再把一些常见问题凝练成一段给出问题与答案最后进入第 15 课。1. Q、K、V 分别怎么理解最适合当前阶段的理解是\[ Q\text{我想找什么} \]\[ K\text{我拿什么和你进行匹配} \]\[ V\text{如果你关注我我真正提供什么信息} \]所以 Attention 不是直接拿 \(Q\) 和 \(K\) 当最终信息而是先用它们决定谁应该关注谁、关注多少。最终真正被加权汇总的是 \(V\)。2. Self-Attention 中 Q、K、V 都来自 \(X\)为什么还不是同一个 Tensor因为经过了三套不同的可训练参数\[ QXW_Q \]\[ KXW_K \]\[ VXW_V \]虽然原始来源都是\[ X \]但是\[ W_Q\neq W_K\neq W_V \]所以得到的数值不同功能也不同。3. 如果\[ X:[32,10,64] \]并且暂时统一\[ D64 \]那么\[ Q:[32,10,64] \]\[ K:[32,10,64] \]\[ V:[32,10,64] \]4. 为什么\[ Q:[32,10,64] \]乘\[ K^T:[32,64,10] \]得到\[ [32,10,10] \]因为矩阵乘法\[ [A,B]\times[B,C]\rightarrow[A,C] \]所以每个 Batch 内\[ [10,64]\times[64,10] \]得到\[ [10,10] \]加上 Batch\[ \boxed{[32,10,10]} \]5.[32,10,10]中最后两个 10 分别是什么它们都表示 Token 位置但扮演的角色不同。第一个\[ 10 \]表示10 个 Query Token。第二个\[ 10 \]表示每个 Query 都要比较的 10 个 Key Token。所以元素\[ S_{ij} \]表示第 \(i\) 个 Query Token 和第 \(j\) 个 Key Token 的匹配程度。6. 为什么 Softmax 对最后一个维度做因为对于一个固定 Query我们希望它对所有 Key 的注意力权重加起来等于 1。例如\[ [0.1,0.6,0.2,0.1] \]表示当前 Query10% 关注 Token 160% 关注 Token 220% 关注 Token 310% 关注 Token 4。所以weights torch.softmax(scores, dim-1)是在“所有 Key”这一维做 Softmax。7. 为什么除以\[ \sqrt D \]因为\[ Q\cdot K \]需要对 \(D\) 个元素进行乘加。当 \(D\) 很大时点积的数值尺度容易变大。如果很大的分数直接进入 SoftmaxSoftmax 可能变得过于极端不利于稳定训练。所以进行缩放\[ \frac{QK^T}{\sqrt D} \]控制数值尺度。严格写法通常是 \(\sqrt{d_k}\)但在我们当前单头且统一维度的阶段先用 \(D\) 理解即可。8. 为什么 Attention Weight 后还必须乘 \(V\)因为 Attention Weight\[ [B,T,T] \]只告诉我们应该从谁那里拿多少信息。真正的信息内容在\[ V:[B,T,D] \]里面。所以\[ AttentionWeights\times V \]才真正完成按照注意力关系重新汇总 Token 信息。最终\[ [B,T,T]\times[B,T,D] \rightarrow [B,T,D] \]9. Self-Attention 和 Cross-Attention 的核心区别是什么Self-Attention\[ Q,K,V \]来自同一个 Sequence。可以理解成一组 Token 内部相互交流。Cross-Attention\[ Q \]来自一组 Token而\[ K,V \]来自另一组 Token。可以理解成一组 Token 去另一组 Token 中检索信息。后面 ACT Decoder 会再次遇到这个区别。10. Padding Mask 为什么通常在 Softmax 前作用因为我们希望 Padding 的注意力权重最终为\[ 0 \]一种常见做法是在 Softmax 前把对应 Score 设置成非常大的负数例如\[ -\infty \]那么经过 Softmax\[ e^{-\infty}\approx0 \]因此模型不会从 Padding Token 中读取信息。11. Attention 中主要学习哪些 Parameter最核心的是\[ W_Q,\quad W_K,\quad W_V \]以及真实 Multi-Head Attention 中常见的最终输出投影\[ W_O \]这些参数都会通过\[ Loss \rightarrow Backward \rightarrow Gradient \rightarrow Optimizer \]进行学习。12. 完整解释 Attention 公式\[ Attention(Q,K,V) softmax \left( \frac{QK^T}{\sqrt D} \right)V \]可以直接翻译成\[ QK^T \]每个 Query 和所有 Key 做匹配。↓\[ \frac{QK^T}{\sqrt D} \]控制匹配分数的数值尺度。↓\[ softmax \]把匹配分数变成注意力比例。↓\[ \times V \]按这些比例把真正的信息重新汇总。最终\[ \boxed{ 每个Token得到融合其他Token信息后的新表示 } \]常见问题的凝练这段可以直接放在这一课 Attention 公式后面。问题为什么计算 Attention 时要把 \(K\) 转置使 \(QK^T\) 从 \([B,T,D]\) 得到 \([B,T,T]\)既然 \([B,T,T]\) 已经是一张注意力矩阵为什么还要继续与 \(V:[B,T,D]\) 相乘最后重新得到的 \([B,T,D]\) 又有什么作用答案可以浓缩成下面这一条因果链第一步为什么要 \(K^T\)对首先就是矩阵乘法的需要但更深层原因是我们想让“每个 Query 和每个 Key 两两做点积”。因为\[ Q:[T,D] \]每一行是一条 Query\[ K:[T,D] \]每一行是一条 Key。我们想计算\[ q_i\cdot k_j \]对于所有\[ i,j \]所以把\[ K:[T,D] \]转成\[ K^T:[D,T] \]这样\[ [T,D]\times[D,T] [T,T] \]矩阵乘法一次就把所有\[ q_i\cdot k_j \]都算出来了。所以\[ \boxed{ K转置不是单纯为了凑Shape而是为了让每个Q和每个K完成两两点积 } \]第二步为什么[B,T,T]还不是最终结果因为\[ [B,T,T] \]只是一张Token 和 Token 之间的关系表。例如某一行\[ [0.1,0.7,0.2] \]只说明当前 Token 应该分别从三个 Token 中拿 10%、70%、20% 的信息。但是这张表本身没有告诉你那三个 Token 到底有什么信息。真正的信息放在\[ V \]里。因此还需要\[ AttentionWeights\times V \]第三步为什么乘完又回到[B,T,D]因为\[ [B,T,T] \times [B,T,D] \]本质上是在对 \(T\) 个 Value Token 做加权求和。对于第 \(i\) 个 Token\[ o_i \sum_{j1}^{T} a_{ij}v_j \]因此第 \(i\) 个位置仍然得到一个\[ D \]维向量。一共有 \(T\) 个位置所以重新得到\[ [B,T,D] \]但是这个新的\[ [B,T,D] \]和原来的已经不一样。原来的 Token主要表示自己的信息。Attention 后的 Token已经把其他 Token 的相关信息融合进来了。因此整个 Attention 最重要的 Shape 变化可以记成\[ \boxed{ [B,T,D] \overset{QK^T}{\longrightarrow} [B,T,T] \overset{\times V}{\longrightarrow} [B,T,D] } \]含义分别是\[ \boxed{ Token信息 \rightarrow Token关系 \rightarrow 融合关系后的新Token信息 } \]而最后这个新的\[ [B,T,D] \]还能继续进入下一层 Transformer、Feed Forward Network 或最后的任务 Head。这一点就是 Transformer 能不断堆叠的基础。VLA 系统学习第 15 课为什么一个 Attention Head 还不够——Multi-Head Attention 与 Transformer Block我们现在已经懂了单头 Attention\[ X \rightarrow Q,K,V \rightarrow QK^T \rightarrow AttentionWeights \rightarrow V \rightarrow 新的X \]但这里还有一个限制如果所有 Token 之间只有一套 Q/K/V 表示那么模型每次只能在这一套表示空间里判断 Token 关系。真实 Sequence 里的关系可能很多。例如一段机器人历史中一个 Token 和另一个 Token 之间可能同时存在时间上的关系状态变化关系视觉对应关系任务相关关系。我们不应该人工规定Head 1 一定学时间Head 2 一定学空间。真实网络学到什么由训练决定。但是可以让模型并行使用多套不同的 Q/K/V 投影从多个表示子空间建立 Attention。这就是Multi-Head Attention一、先用一个非常具体的 Shape 开始假设\[ X:[B,T,D] \]具体取\[ X:[32,10,64] \]现在我们设置\[ H4 \]也就是 4 个 Attention Head。我们把\[ D64 \]平均拆成\[ 4 \]份\[ d_h\frac{D}{H} \frac{64}{4} 16 \]所以\[ \boxed{ DH\times d_h } \]这里开始重新引入小写\[ d_h \]但这次含义非常明确一个 Head 自己处理的特征维度。所以\[ D64 \]是整个模型的 Embedding Dimension而\[ d_h16 \]是单个 Head 的维度。二、是不是直接把原来的 64 个特征硬切四块从 Shape 上看确实会\[ 64 \rightarrow 4\times16 \]但在此之前一般先经过可训练 Projectionself.q_proj nn.Linear(64, 64)self.k_proj nn.Linear(64, 64)self.v_proj nn.Linear(64, 64)输入\[ X:[32,10,64] \]得到\[ Q,K,V:[32,10,64] \]然后才重新组织 Shape\[ [32,10,64] \rightarrow [32,10,4,16] \]这里32Batch10Token4Head16每个 Head 的维度。代码q self.q_proj(x)q q.reshape(B, T, 4, 16)这一步没有增加或者减少数据数量\[ 644\times16 \]只是重新解释了最后一个维度。三、为什么还要再transpose现在\[ Q:[B,T,H,d_h] \]也就是\[ [32,10,4,16] \]但我们希望每一个 Head 独立计算自己的一套 Attention。所以更方便的排列是\[ [B,H,T,d_h] \]也就是\[ [32,4,10,16] \]代码q q.transpose(1, 2)K、V 同样处理q q.reshape(B, T, H, d_h).transpose(1, 2)k k.reshape(B, T, H, d_h).transpose(1, 2)v v.reshape(B, T, H, d_h).transpose(1, 2)现在每个 Head 都拥有\[ [T,d_h] \]的一套 Q、K、V。四、现在每个 Head 独立做我们上一课完全一样的事情现在\[ Q:[B,H,T,d_h] \]\[ K:[B,H,T,d_h] \]K 交换最后两个维度\[ K^T:[B,H,d_h,T] \]于是\[ QK^T \]得到\[ \boxed{ [B,H,T,T] } \]例如\[ [32,4,10,10] \]这意味着32 个 Sample每个 Sample 有 4 张不同的 Attention Matrix每张都是 10×10。所以单头时\[ [B,T,T] \]多头以后\[ \boxed{ [B,H,T,T] } \]只是多了\[ H \]这个 Head Dimension。五、代码已经可以和数学完全对应scores q k.transpose(-2, -1)Shape\[ [32,4,10,16] \times [32,4,16,10] \]得到\[ [32,4,10,10] \]然后scores scores / math.sqrt(self.head_dim)weights torch.softmax(scores, dim-1)Shape 仍然\[ [32,4,10,10] \]然后out weights v因为\[ [32,4,10,10] \times [32,4,10,16] \]得到\[ \boxed{ [32,4,10,16] } \]也就是每个 Head 都完成了一次独立的信息聚合。六、四个 Head 算完以后怎么办现在\[ out:[B,H,T,d_h] \]即\[ [32,4,10,16] \]Transformer 最终还是希望回到\[ [B,T,D] \]所以先交换回来\[ [B,H,T,d_h] \rightarrow [B,T,H,d_h] \]代码out out.transpose(1, 2)得到\[ [32,10,4,16] \]然后把\[ H\times d_h \]重新合并\[ 4\times1664 \]代码out out.reshape(B, T, D)得到\[ \boxed{ [32,10,64] } \]所以整个 Multi-Head Attention Shape 特别值得记住\[ [B,T,D] \]↓\[ [B,H,T,d_h] \]↓\[ [B,H,T,T] \]↓\[ [B,H,T,d_h] \]↓\[ \boxed{ [B,T,D] } \]七、为什么最后还有一个out_proj标准 Multi-Head Attention 通常还会有self.out_proj nn.Linear(D, D)所以多头拼回来以后out self.out_proj(out)Shape\[ [B,T,D] \rightarrow [B,T,D] \]不变。但是这个 Linear 可以进一步学习多个 Head 拼接回来的信息应该如何重新混合。所以完整思想多个Head分别读取信息 ↓ 把多个Head结果拼回来 ↓ 再经过Output Projection重新融合八、现在写出一个最小可读版本import mathimport torchimport torch.nn as nnclass SimpleMultiHeadAttention(nn.Module): def __init__(self, d_model64, num_heads4): super().__init__() self.d_model d_model self.num_heads num_heads self.head_dim d_model // num_heads self.q_proj nn.Linear(d_model, d_model) self.k_proj nn.Linear(d_model, d_model) self.v_proj nn.Linear(d_model, d_model) self.out_proj nn.Linear(d_model, d_model) def forward(self, x): B, T, D x.shape q self.q_proj(x) k self.k_proj(x) v self.v_proj(x) q q.reshape( B, T, self.num_heads, self.head_dim ).transpose(1, 2) k k.reshape( B, T, self.num_heads, self.head_dim ).transpose(1, 2) v v.reshape( B, T, self.num_heads, self.head_dim ).transpose(1, 2) scores q k.transpose(-2, -1) scores scores / math.sqrt(self.head_dim) weights torch.softmax(scores, dim-1) out weights v out out.transpose(1, 2) out out.reshape(B, T, D) out self.out_proj(out) return out如果输入\[ x:[32,10,64] \]输出仍然\[ \boxed{[32,10,64]} \]这点特别重要。Multi-Head Attention 不是把 Sequence 变没了。而是让原来的 Token Sequence 完成一次多视角的信息交换然后仍然返回同样结构的 Token Sequence。九、Attention 做完以后为什么还不是完整 Transformer现在我们只有\[ X \rightarrow MultiHeadAttention \rightarrow Y \]但是 Transformer Block 还要解决另外一个问题。Attention 擅长的是\[ \boxed{ Token之间交换信息 } \]例如Token 1 从 Token 2、3、4 读取什么。但每个 Token 拿到这些信息以后还需要进一步在自己的 Feature Dimension 内进行非线性加工。这就是Feed Forward Network简称\[ FFN \]十、FFN 和 Attention 的职责非常不同这是 Transformer 特别需要理解的一点。Attention主要进行\[ \boxed{ Token\leftrightarrow Token } \]的信息交互。FFN则对每个 Token分别使用同一个 MLP\[ D \rightarrow D_{ff} \rightarrow D \]例如\[ 64 \rightarrow 256 \rightarrow 64 \]代码ffn nn.Sequential( nn.Linear(64, 256), nn.GELU(), nn.Linear(256, 64))如果输入\[ [32,10,64] \]第一层\[ [32,10,64] \rightarrow [32,10,256] \]GELU\[ [32,10,256] \]第二层\[ [32,10,256] \rightarrow [32,10,64] \]注意\[ T10 \]一直没变。因为 FFN 对每一个 Token 独立进行特征变换。十一、所以 Attention 和 FFN 可以这样区分这是这一课非常重要的一组条件反射\[ \boxed{ Attention 不同Token之间交流 } \]而\[ \boxed{ FFN 每个Token内部继续加工Feature } \]可以粗略画成Token 1 ─┐ Token 2 ─┼→ Attention → Token之间交换信息 Token 3 ─┘ 然后 Token 1 → FFN → 自己内部加工 Token 2 → FFN → 自己内部加工 Token 3 → FFN → 自己内部加工两部分配合起来才形成 Transformer Block 的核心计算。十二、但为什么还有 Residual Connection如果我们直接\[ X \rightarrow Attention \rightarrow Y \]那么新结果完全替换旧 \(X\)。Transformer 通常会做\[ YXAttention(X) \]这叫Residual Connection残差连接。代码x x attention_outputShape\[ [B,T,D] [B,T,D] \]得到\[ [B,T,D] \]十三、Residual 为什么重要可以先抓两个直觉。第一Attention 新学到的是“补充信息”原来的 Token 信息不需要全部扔掉。所以\[ XAttention(X) \]有一种原始信息 新信息的意味。第二Residual 能为深层网络提供更直接的梯度传播路径。所以堆很多 Transformer Block 时更容易优化。以后学 ResNet 时还会系统碰到 Residual当前理解这两个作用就够。十四、LayerNorm 又是什么Transformer 中还会频繁看到nn.LayerNorm(D)它用于对 Feature 做归一化。假设\[ X:[B,T,D] \]例如\[ [32,10,64] \]那么nn.LayerNorm(64)通常是对每一个 Token 自己的 64 个 Feature进行归一化处理。所以 Shape 不变\[ [32,10,64] \rightarrow [32,10,64] \]它主要帮助网络内部数值尺度更加稳定。LayerNorm 的数学细节后面遇到需要时再进一步展开现在先理解它的位置和职责。十五、终于拼成一个完整 Transformer Block一种经典的 Post-Norm 表达可以写成\[ Y LayerNorm( XMHA(X) ) \]然后\[ Z LayerNorm( YFFN(Y) ) \]完整链X ↓ Multi-Head Attention ↓ X ↓ LayerNorm ↓ FFN ↓ 前面的结果 ↓ LayerNorm ↓ OutputShape 从头到尾\[ [B,T,D] \]保持不变。十六、代码里完整长什么样把刚才的 Multi-Head Attention 放进去class SimpleTransformerBlock(nn.Module): def __init__(self, d_model64, num_heads4): super().__init__() self.attention SimpleMultiHeadAttention( d_modeld_model, num_headsnum_heads ) self.norm1 nn.LayerNorm(d_model) self.ffn nn.Sequential( nn.Linear(d_model, d_model * 4), nn.GELU(), nn.Linear(d_model * 4, d_model) ) self.norm2 nn.LayerNorm(d_model) def forward(self, x): attn_out self.attention(x) x x attn_out x self.norm1(x) ffn_out self.ffn(x) x x ffn_out x self.norm2(x) return x输入\[ [32,10,64] \]Attention\[ [32,10,64] \]Residual\[ [32,10,64] \]Norm\[ [32,10,64] \]FFN 中间\[ [32,10,256] \]最后回来\[ [32,10,64] \]最终整个 Block\[ \boxed{ [32,10,64] \rightarrow [32,10,64] } \]但是里面的信息已经进行了Token 间交互 Token 内特征变换。十七、真实 Transformer 的 Norm 顺序可能与你看到的不一样以后你可能看到Attention → Add → Norm也可能看到Norm → Attention → Add前者通常叫Post-Norm后者通常叫Pre-Norm所以不要以后发现代码顺序不一样就认为哪个一定写错了。不同 Transformer 实现可能选择不同结构。尤其我们进入真实 ACT / DETR 风格代码以后会直接检查它到底用的是哪一种。现在重点不是背哪个唯一正确而是知道 Transformer Block 的核心零件仍然是\[ \boxed{ Attention FFN Residual Normalization } \]十八、为什么 Transformer 可以堆很多 Block因为每一个 Block 输入\[ [B,T,D] \]输出仍然\[ [B,T,D] \]所以可以\[ X \rightarrow Block_1 \rightarrow Block_2 \rightarrow Block_3 \rightarrow \cdots \]第一层 Token 完成一次信息交流。第二层在新的 Representation 上继续交流。第三层再继续。于是 Representation 越来越丰富。十九、现在nn.TransformerEncoderLayer已经不应该是纯黑盒了以后看到layer nn.TransformerEncoderLayer( d_model64, nhead4)你应该先想到它内部大致包含Multi-Head Self-Attention Residual LayerNorm Feed Forward Network Residual LayerNorm而d_model64就是\[ D64 \]nhead4就是\[ H4 \]于是每 Head 大致\[ d_h\frac{64}{4}16 \]这时候 API 就已经能和数学对应起来了。二十、我们现在距离 ACT 到底还差什么其实已经非常近了。我们已经掌握\[ [B,T,D] \]Token / EmbeddingPositionQ/K/VSelf-AttentionCross-Attention 的基本思想Multi-Head AttentionTransformer Block。现在还差两个关键问题第一Transformer Encoder 和 Decoder 到底有什么区别第二一个 Action Query 如何通过 Decoder 去读取 Image / Robot State 的信息然后产生 Action Sequence这两个问题恰好就是理解 ACT 架构非常重要的最后一段 Transformer 基础。所以接下来不会再继续向外扩散。第十五课自测如果\[ D128,\quad H8 \]那么单个 Head Dimension\[ d_h \]是多少输入\[ X:[32,20,64] \]设置\[ H4 \]拆 Head 后为什么得到\[ [32,4,20,16] \]Multi-Head Attention 中 Score Matrix 的 Shape 为什么从单头的\[ [B,T,T] \]变成\[ [B,H,T,T] \]如果\[ AttentionWeights:[32,4,20,20] \]\[ V:[32,4,20,16] \]相乘以后 Shape 是什么为什么多个 Head 的结果最后还要重新合并成\[ [B,T,D] \]out_proj的作用是什么Attention 和 FFN 的核心职责分别是什么为什么 FFN 通常不改变 Sequence Length \(T\)Residual Connection\[ XAttention(X) \]为什么要求两者 Shape 相同nn.LayerNorm(64)面对\[ [32,10,64] \]主要归一化哪个维度Transformer Block 为什么能够连续堆很多层为什么不能机械认为“第 1 个 Head 一定负责时间第 2 个 Head 一定负责视觉”
返回列表