每日基础课

一句话定义：自注意力机制是一种让序列中的每个元素，都能根据上下文内容，动态决定给其他所有元素分配多少关注度的计算方式。

记忆口诀： “Q 查 K 算分，Softmax 归一化，乘上 V 汇聚成新自己。”

生活类比：学术研讨会自报家门 想象一个包含 50 人的闭门研讨会，每个人代表音频里的一个时间帧（比如 20ms 的声音）。 1. Query（Q，提问者）：“我正在说这句话，我需要弄清楚当前声音和谁最相关？” 2. Key（K，名片）：在场每个人胸口挂着的名片，上面写着自己的特征（比如“我是说话人 A 的辅音”、“我是背景噪声”）。 3. 点积计算（匹配度）：你拿着你的 Q，去和全场所有人的 K 逐个握手比对。匹配度越高，得分越高。 4. Softmax（归一化）：把全部分数变成百分比，加起来刚好是 100%。 5. Value（V，干货内容）：每个人发言的实际内容。 6. 加权求和：你把每个人的发言内容（V）按照刚才算出的百分比打包带走。最终，你手里的新特征既保留了你自己，又融合了全场对你最重要的上下文。

想用图直观建立 Q/K/V 的空间几何直觉，非常推荐阅读《The Illustrated Transformer》，它能帮你把这个握手过程彻底可视化。

核心数学公式： $$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V$$

符号逐个拆解： - $Q$（Query，查询矩阵）：形状为 `(序列长度 T, 特征维度 d_k)`。代表“主动寻找依赖”的一方。 - $K^T$（Key 的转置）：形状为 `(特征维度 d_k, 序列长度 T)`。转置是为了让 $Q$ 和 $K$ 能做矩阵乘法。 - $Q K^T$：两个矩阵相乘，得到一个 `(T, T)` 的方阵。第 $i$ 行第 $j$ 列的数值，就是第 $i$ 个时间步对第 $j$ 个时间步的原始相关性打分。 - $\sqrt{d_k}$（缩放因子）：$d_k$ 是特征维度的长度。当维度很大时，点积的结果会非常大，导致 Softmax 的梯度几乎为 0（梯度消失）。除以 $\sqrt{d_k}$ 能把数值拉回稳定区间。 - $\text{softmax}(\cdot)$：在最后一维做归一化，把打分变成每一行和为 1 的注意力权重概率分布。 - $V$（Value，内容矩阵）：形状为 `(序列长度 T, 输出维度 d_v)`。用算好的权重对 $V$ 做矩阵乘法，输出更新后的上下文表征。

它解决什么问题？没有它会怎样？ 在自注意力出现之前，序列建模主要靠两招： 1. CNN（卷积）：视野局限。卷积核一次只能看局部几帧，如果要让第 1 帧和第 1000 帧建立联系，必须堆叠极深的卷积层，信息传递容易衰减。 2. RNN/LSTM（循环神经网络）：一步一步按时间往后传。前面的信息传到后面早就稀释或遗忘了；更致命的是它必须串行计算，显卡算力完全跑不满。

没有自注意力：在音视频或长语音中，5 秒前的声纹和现在的声纹如果跨越了数百帧，模型很难直接对齐它们，必须费力地在多层网络中反复传递隐状态。自注意力把两点之间的距离直接缩短为 1 步计算。

实现思路： 1. 准备三个线性变换层（`nn.Linear`），把输入投影为 $Q, K, V$。 2. 矩阵乘法计算相关度：$Q \times K^T$，结果维度为 `(B, T, T)`。 3. 缩放并做 Softmax：除以 $\sqrt{d_k}$ 后在最后一维做 `torch.softmax`。 4. 汇聚内容：用注意力矩阵乘以 $V$，输出更新后的张量。

想对照工业级代码逐行精读每一处细节，可以参考《The Annotated Transformer》，它是代码级理解经典架构的极佳教材。

完整可运行代码：

import math
import torch
import torch.nn as nn


class SingleHeadSelfAttention(nn.Module):

    def __init__(self, embed_dim: int):
        super().__init__()
        self.embed_dim = embed_dim

        # 1. 映射生成 Q, K, V 的线性层
        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False)

    def forward(self, x: torch.Tensor):
        # x 形状: (batch_size, seq_len, embed_dim)
        B, T, D = x.shape

        # 线性投影
        Q = self.q_proj(x)  # (B, T, D)
        K = self.k_proj(x)  # (B, T, D)
        V = self.v_proj(x)  # (B, T, D)

        # 2. 计算点积相似度 scores = Q * K^T
        # K 转置最后两个维度: (B, T, D) -> (B, D, T)
        # scores 形状: (B, T, T)
        scores = torch.matmul(Q, K.transpose(-2, -1))

        # 3. 缩放 (Scale)
        scaled_scores = scores / math.sqrt(D)

        # 4. Softmax 归一化为注意力权重
        # dim=-1 保证每一行的所有列加起来等于 1
        attn_weights = torch.softmax(scaled_scores, dim=-1)  # (B, T, T)

        # 5. 用权重对 V 进行加权求和
        # (B, T, T) @ (B, T, D) -> (B, T, D)
        output = torch.matmul(attn_weights, V)

        return output, attn_weights


# --- 验证测试 ---
if __name__ == "__main__":
    torch.manual_seed(42)

    # 模拟场景：2 条语音切片，每条截取 5 个时间帧，每帧特征维度 64
    batch_size = 2
    seq_len = 5
    embed_dim = 64

    fake_audio_frames = torch.randn(batch_size, seq_len, embed_dim)

    # 实例化并运行
    attn_layer = SingleHeadSelfAttention(embed_dim)
    out, weights = attn_layer(fake_audio_frames)

    print("=== 形状验证 ===")
    print(f"输入形状:   {fake_audio_frames.shape}")  # (2, 5, 64)
    print(f"输出形状:   {out.shape}")  # (2, 5, 64)
    print(f"权重矩阵:   {weights.shape}")  # (2, 5, 5)

    print("\n=== 数值正确性验证 ===")
    # 验证权重矩阵第 1 个样本的每一行之和是否为 1
    row_sums = weights[0].sum(dim=-1)
    print(f"第 0 个样本各行求和值: {row_sums.detach().numpy()}")
    assert torch.allclose(
        row_sums, torch.ones_like(row_sums)
    ), "Softmax 归一化异常！"
    print("测试全部通过！")