每日基础课

一句话定义：自注意力机制，就是让一段序列里的每个位置，去“查看”同一段序列里的所有位置，然后按重要程度加权汇总信息。

记忆口诀： “Q 去找 K，分数当权重，最后加权 V。”

生活类比： 你在开会。 每个人都要写一份会议纪要。 写的时候，他不会只看自己说了什么。 他会看所有人说了什么。 谁的话和自己的任务更相关，就多参考谁。 谁不相关，就少参考谁。

在自注意力里： 每个 token、每一帧音频、每一段视频片段，都像会议中的一个人。 它会问： “我应该重点看谁？” 然后把别人的信息按权重合成自己的新表示。

必要公式：

Self-Attention 的核心公式是：

\[ Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_k}})V \]

逐个解释符号：

1. \(X\) 输入序列。 可以是一句话的 token 向量。 也可以是一段音频的帧特征。 也可以是一段视频的帧/片段特征。 形状常见为：

\[ X \in \mathbb{R}^{T \times d} \]

其中： - \(T\)：序列长度，比如 token 数、音频帧数、视频片段数 - \(d\)：每个位置的特征维度

2. \(Q\)、\(K\)、\(V\) 它们都由输入 \(X\) 线性变换得到：

\[ Q=XW_Q \]

\[ K=XW_K \]

\[ V=XW_V \]

其中： - \(Q\)：Query，查询。表示“我想找什么信息” - \(K\)：Key，键。表示“我有什么特征可被别人匹配” - \(V\)：Value，值。表示“如果别人关注我，我实际贡献什么内容” - \(W_Q, W_K, W_V\)：可学习的参数矩阵

一个直觉： Q 像搜索词。 K 像文档标签。 V 像文档正文。

3. \(QK^T\) 这是在算相似度。 每个位置的 Q，和所有位置的 K 做点积。 点积越大，说明越相关。

如果序列长度是 \(T\)，那么：

\[ QK^T \in \mathbb{R}^{T \times T} \]

这就是注意力分数矩阵。 第 \(i\) 行表示： 第 \(i\) 个位置在看所有位置时，各自有多重要。

4. \(\sqrt{d_k}\) 这是缩放因子。 \(d_k\) 是 Key/Query 的维度。 如果不除以 \(\sqrt{d_k}\)，点积可能太大。 softmax 会变得太尖。 模型训练会不稳定。

5. \(softmax\) 把分数变成权重。 每一行的权重加起来等于 1。 可以理解为： “我把注意力预算分给所有位置。”

6. 最后乘以 \(V\) 得到加权求和后的新表示。 也就是： 每个位置根据自己关注的对象，重新整理了一份信息。

它解决什么问题：

自注意力解决的是： “序列中不同位置之间，如何直接建立关系。”

比如一句话： “他说他今天不能来。” 第二个“他”指谁？ 要看前文。

比如音频： 当前一帧到底是说话声、背景声、音乐，还是重叠说话？ 不能只看当前帧。 要看前后上下文。

比如说话人日志 diarization： 某一段声音属于说话人 A 还是 B。 不能只看这一小段。 还要看更早、更晚的片段。 因为同一个说话人的声音特征会在时间上反复出现。

没有自注意力会怎样：

没有它也能做序列建模。 但会麻烦一些。

以前常用 RNN/LSTM。 它按时间一步一步读。 远距离信息要经过很多步传递。 长序列时容易变弱。 并行也不方便。

也常用 CNN。 它看局部窗口。 想看远处，要堆很多层。 或者加大卷积感受野。 但“谁和谁相关”不是显式算出来的。

自注意力的关键好处是： 任意两个位置可以一步建立联系。 而且可以并行计算。

在音视频理解里，这很重要。 因为视频事件、说话人变化、语音上下文，经常跨越较长时间。

思路：

1. 准备一个小序列。 可以把它想成一小段音频的 5 个时间帧。 每帧有 4 维特征。

• 用三个线性层生成 Q、K、V。

• 计算注意力分数：

\[ scores = \frac{QK^T}{\sqrt{d_k}} \]

4. 对最后一维做 softmax。 得到每个时间步看其他时间步的权重。

5. 用权重乘以 V。 得到新的上下文表示。

完整可运行代码如下：

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class SingleHeadSelfAttention(nn.Module):
    def __init__(self, input_dim, attn_dim):
        super().__init__()
        self.input_dim = input_dim
        self.attn_dim = attn_dim

        # Q, K, V 都从同一个输入 X 变换而来
        self.w_q = nn.Linear(input_dim, attn_dim, bias=False)
        self.w_k = nn.Linear(input_dim, attn_dim, bias=False)
        self.w_v = nn.Linear(input_dim, attn_dim, bias=False)

    def forward(self, x, mask=None):
        """
        x: [batch, time, input_dim]
        mask: 可选，[batch, time]，1 表示有效位置，0 表示 padding
        """

        # 1. 生成 Q, K, V
        # q/k/v: [batch, time, attn_dim]
        q = self.w_q(x)
        k = self.w_k(x)
        v = self.w_v(x)

        # 2. 计算 QK^T
        # k.transpose(-2, -1): [batch, attn_dim, time]
        # scores: [batch, time, time]
        scores = torch.matmul(q, k.transpose(-2, -1))

        # 3. 缩放，避免点积过大
        scores = scores / math.sqrt(self.attn_dim)

        # 4. 如果有 padding mask，就把无效位置设成很小
        # 这样 softmax 后权重接近 0
        if mask is not None:
            # mask: [batch, time]
            # key_mask: [batch, 1, time]
            key_mask = mask.unsqueeze(1)
            scores = scores.masked_fill(key_mask == 0, -1e9)

        # 5. 对最后一维做 softmax
        # attn_weights[b, i, j] 表示：
        # 第 b 个样本中，第 i 个时间步关注第 j 个时间步的权重
        attn_weights = F.softmax(scores, dim=-1)

        # 6. 加权求和 V
        # output: [batch, time, attn_dim]
        output = torch.matmul(attn_weights, v)

        return output, attn_weights


def main():
    torch.manual_seed(42)

    batch_size = 2
    time_steps = 5
    input_dim = 4
    attn_dim = 8

    # 假设这是两段很短的音频特征
    # 每段有 5 个时间帧，每帧 4 维
    x = torch.randn(batch_size, time_steps, input_dim)

    model = SingleHeadSelfAttention(input_dim=input_dim, attn_dim=attn_dim)

    output, attn_weights = model(x)

    print("Input shape:", x.shape)
    print("Output shape:", output.shape)
    print("Attention weights shape:", attn_weights.shape)

    print("\nAttention matrix of sample 0:")
    print(attn_weights[0])

    # 检查每一行注意力权重是否加和为 1
    row_sums = attn_weights.sum(dim=-1)
    print("\nRow sums of attention weights:")
    print(row_sums)

    print("\nAre row sums close to 1?")
    print(torch.allclose(row_sums, torch.ones_like(row_sums), atol=1e-6))

    # 再演示一个带 padding mask 的例子
    # 第 0 个样本 5 帧都有效
    # 第 1 个样本只有前 3 帧有效，后 2 帧是 padding
    mask = torch.tensor([
        [1, 1, 1, 1, 1],
        [1, 1, 1, 0, 0]
    ])

    output_masked, attn_weights_masked = model(x, mask=mask)

    print("\nWith padding mask:")
    print("Masked output shape:", output_masked.shape)
    print("Masked attention matrix of sample 1:")
    print(attn_weights_masked[1])

    print("\nRow sums with mask:")
    print(attn_weights_masked.sum(dim=-1))


if __name__ == "__main__":
    main()