深度学习｜表示学习｜多头注意力在计算时常见的张量维度变换总结｜28

如是我闻： 以下是多头注意力（Multi-Headed Attention）在计算时常见的张量维度变换总结，帮助理解从输入到输出是如何一步步处理的。为了方便，令：

• $B$ 表示 batch size（批量大小）
• $S$ 表示 sequence length（序列长度）
• $m$ 表示 num_heads（注意力头数）
• $h$ 表示 head_size（每个头的维度）
• $d_{\mathrm{model}}=m\times h$ 表示模型隐层维度

1. 输入（queries、keys、values）
   $$\text{形状}=(B,S,d_{\mathrm{model}}).$$
   在「自注意力」（self-attention）场景下，三者通常是同一个张量；在「交叉注意力」（cross-attention）场景下，$\text{queries}$ 和 $\text{keys, values}$ 可能来自不同子网络。

2. 线性映射（$W_Q,W_K,W_V$）

   • 对 $\text{queries}$ 做线性变换得到 (Q)：形状仍为 $(B,S,d_{\mathrm{model}})$
   • 对 $\text{keys}$ 做线性变换得到 $K$：形状同上
   • 对 $\text{values}$ 做线性变换得到 $V$：形状同上

3. 拆分 heads（split heads）

   • 将 $(B,S,d_{\mathrm{model}})$ reshape + transpose 成 $(B,m,S,h)$。
   • 这样每个 batch、每个序列位置上就可以拆出 $m$ 个“头”，每个头维度为 $h$。
   • 拆分后：
     $$Q,K,V\stackrel{\text{split}}{\to }(B,m,S,h).$$

4. 计算注意力分数（scores）

   • 使用 scaled dot-product：
     $$\text{scores}=\frac{Q\times K^T}{\sqrt{h}}\text{形状}=(B,m,S,S).$$
   • 此时会应用「下三角 mask」（causal mask）以保证自回归：只关注「过去和当前」位置，屏蔽「未来」位置。
   • 对 $\text{scores}$ 做 $\mathrm{softmax}$ 得到注意力权重 $attn\_weights$。

5. 加权求和（attended values）
   $$\text{attended\_values}=attn\_weights\times V,\text{形状}=(B,m,S,h).$$
   这样就得到每个 head 对原值向量的加权结果。

6. 合并 heads（merge heads）

   • 将 $(B,m,S,h)$ 还原到 $(B,S,m\times h)$，即 $(B,S,d_{\mathrm{model}})$。
   • 合并之后，相当于将所有 head 的信息拼接到最后一个维度上。

7. 可选的最终线性映射 ${\mathbf{W}}_O$

   • 多数实现会继续用一个线性层 ${\mathbf{W}}_O$（同样是 $(d_{\mathrm{model}},d_{\mathrm{model}})$）把拼接后的多头输出再次投影，形状保持 $(B,S,d_{\mathrm{model}})$。

通过以上步骤，多头注意力便可将序列的上下文信息捕获到不同的 head（不同的子空间），再合并形成新的隐层表示。

cao!