每日基础课

1. 生活类比 想象你在看一段多人圆桌会议的视频： - 如果只有单头注意力，就像你只带了一只眼睛，你盯住了“谁在说话”，就很难同时分心去观察“旁边人的微表情”和“背景里ppt翻页的动作”。 - 多头注意力就像你请了 8 个专业侦探： - 侦探 1 专看声纹与音调起伏（谁在说话）； - 侦探 2 专看嘴唇动作与音频是否同步（唇音对齐）； - 侦探 3 专看前后几秒的语义上下文（说了什么）。 最后把 8 个侦探的报告汇总拼在一起，得出对当前时刻的完整判断。

2. 核心公式与符号拆解 单个头的计算是缩放点积注意力： $$ \text{head}_i = \text{Attention}(Q W_i^Q, K W_i^K, V W_i^V) = \text{softmax}\left(\frac{Q_i K_i^T}{\sqrt{d_k}}\right) V_i $$ 多头汇总： $$ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W^O $$

符号逐个看： - $Q, K, V$：输入序列的查询、键、值矩阵。如果是自注意力，三者都来自同一个输入特征。 - $W_i^Q, W_i^K, W_i^V$：第 $i$ 个头的线性投影矩阵，负责把原始特征映射到第 $i$ 个专属子空间。 - $d_k$：每个头的特征维度（通常是 总维度 $d_{model} / 头数 h$）。除以 $\sqrt{d_k}$ 是防止点积过大导致 softmax 梯度消失。 - $\text{Concat}$：把所有头的输出按特征维度拼在一起。 - $W^O$：最后的输出线性层，把拼起来的多头信息融合回原始维度。

3. 它解决什么问题？没有它会怎样？ - 解决的问题：单头注意力会对所有关联做加权平均，这会把不同维度的信息“混在一起糊掉”。多头让模型能够同时捕捉多种不同尺度、不同类型的依赖关系。 - 没有它会怎样：在说话人日志（Diarization）中，模型如果只用单头，在遇到两人抢话（重叠语音）时，注意力权重会被迫在两个人之间取折中，导致声纹表征混乱，无法清晰区分“说话人 A”和“说话人 B”。

4. 音视频与说话人日志中的真实位置 - 说话人日志：在 Conformer 或 Transformer 编码器中，多头自注意力层位于局部卷积层之后，用来捕捉说话人跨越数秒甚至数分钟的全局声纹一致性。 - 音视频多模态融合：视频帧序列作 $Q$，音频频谱特征作 $K$ 和 $V$（交叉注意力），不同头分别对齐音视频节拍、嘴型与发音、画面主体与声音来源。

想要建立整体架构直觉，建议去看经典图解资料《The Illustrated Transformer》；想看官方代码如何把张量切块与重组，建议精读《The Annotated Transformer》。

实现思路： 1. 用一个线性层生成全部 $Q, K, V$（比写三个小线性层更高效）。 2. 用 `view` 和 `transpose` 将维度从 `(B, T, D)` 变为 `(B, num_heads, T, head_dim)`。 3. 执行批矩阵乘法计算缩放点积注意力。 4. 将多头结果转置并拼回 `(B, T, D)`。 5. 经过最终输出线性层 $W^O$ 输出。

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim: int, num_heads: int):
        super().__init__()
        assert embed_dim % num_heads == 0, "embed_dim 必须能被 num_heads 整除"
        
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        
        # 一次性投影出 Q, K, V
        self.qkv_proj = nn.Linear(embed_dim, embed_dim * 3)
        # 最后的输出线性融合层
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x: torch.Tensor):
        # x shape: (batch_size, seq_len, embed_dim)
        B, T, D = x.shape
        
        # 1. 线性投影并拆分 Q, K, V
        qkv = self.qkv_proj(x) # (B, T, 3 * D)
        q, k, v = torch.chunk(qkv, 3, dim=-1) # 每个是 (B, T, D)
        
        # 2. 变换维度以支持多头并行计算: (B, T, D) -> (B, num_heads, T, head_dim)
        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        
        # 3. 缩放点积注意力: (B, h, T, head_dim) @ (B, h, head_dim, T) -> (B, h, T, T)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attn_weights = torch.softmax(scores, dim=-1)
        
        # 4. 加权求和: (B, h, T, T) @ (B, h, T, head_dim) -> (B, h, T, head_dim)
        out = torch.matmul(attn_weights, v)
        
        # 5. 拼回多头: (B, num_heads, T, head_dim) -> (B, T, D)
        out = out.transpose(1, 2).contiguous().view(B, T, D)
        
        # 6. 输出线性映射
        output = self.out_proj(out)
        return output, attn_weights

# 测试代码
if __name__ == "__main__":
    torch.manual_seed(42)
    # 模拟 2 段音频，每段 100 帧（如 1 秒钟的语音特征），每帧特征维度 64
    batch_size, seq_len, embed_dim = 2, 100, 64
    num_heads = 4
    
    dummy_audio_feat = torch.randn(batch_size, seq_len, embed_dim)
    
    mha = MultiHeadAttention(embed_dim=embed_dim, num_heads=num_heads)
    out, attn = mha(dummy_audio_feat)
    
    print("输入特征形状:", dummy_audio_feat.shape)
    print("输出特征形状:", out.shape)
    print("注意力权重形状:", attn.shape)
    
    assert out.shape == dummy_audio_feat.shape, "输出与输入形状不一致！"
    assert attn.shape == (batch_size, num_heads, seq_len, seq_len), "注意力矩阵形状不符！"
    print(" 验收通过！多头注意力模块运行正常。")