3、自注意力机制：Self-Attention原理，QKV矩阵的计算过程，多头注意力的拆分与合并

好，咱们进入正题。自注意力机制，说白了就是让模型学会「自己看自己」。你想想看，在处理一句话时，每个词都需要知道跟其他词的关系有多重要。比如「我喜欢吃苹果」——「苹果」跟「吃」关系近，跟「我」关系就远一点。自注意力就是干这个的。

3.1 从「查字典」理解QKV

我第一次接触QKV时，觉得这名字起得挺玄乎。其实没那么复杂。我习惯用一个比喻：你拿着一个问题（Query），去一堆资料（Key）里找答案，最后拿到结果（Value）。

  • Query（Q）：当前词在问「谁跟我有关？」
  • Key（K）：其他词在回答「我跟你有多相关？」
  • Value（V）：真正要提取的信息内容

嗯，这里要注意：Q、K、V其实都是从同一个输入X通过线性变换得到的。说白了就是三个不同的投影矩阵在干活。

3.2 QKV矩阵的计算过程

假设输入序列长度为n，每个词的维度是d。那么输入X的形状就是(n, d)。

计算过程分三步走：

  1. 生成Q、K、V：分别乘以三个权重矩阵Wq、Wk、Wv，形状都是(d, d_k)。
  2. 计算注意力分数：Q乘以K的转置，再除以√d_k做缩放。
  3. Softmax归一化：得到注意力权重，再乘以V得到输出。

公式长这样：

Attention(Q, K, V) = softmax(Q × K^T / √d_k) × V

为什么要除以√d_k？我在项目中遇到过，如果不做这个缩放，当d_k很大时，点积结果会非常大，softmax会进入梯度饱和区。说白了就是模型学不动了。

核心要点：Q和K的点积衡量的是「相似度」。相似度越高，注意力权重越大。这个权重再乘以V，就是加权后的信息。

3.3 多头注意力：拆分与合并

一个头只能学到一种关系模式。但现实中，词与词之间的关系是多样的。比如「苹果」既跟「吃」有动作关系，又跟「红」有属性关系。一个头搞不定，那就多来几个头。

拆分过程：

  • 假设有h个头，每个头的维度是d_k = d / h
  • 把Q、K、V分别拆成h份，每份形状是(n, d_k)
  • 每个头独立计算自注意力

合并过程：

  • 把h个头的输出拼起来，形状变回(n, d)
  • 再乘以一个输出投影矩阵Wo，得到最终结果

代码实现起来其实很直观：

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        self.Wq = nn.Linear(d_model, d_model)
        self.Wk = nn.Linear(d_model, d_model)
        self.Wv = nn.Linear(d_model, d_model)
        self.Wo = nn.Linear(d_model, d_model)
    
    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        
        # 生成QKV并拆分成多头
        Q = self.Wq(x).view(batch_size, seq_len, self.num_heads, self.d_k)
        K = self.Wk(x).view(batch_size, seq_len, self.num_heads, self.d_k)
        V = self.Wv(x).view(batch_size, seq_len, self.num_heads, self.d_k)
        
        # 转置为 (batch, heads, seq, d_k)
        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)
        
        # 计算注意力
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn, V)
        
        # 合并多头
        output = output.transpose(1, 2).contiguous()
        output = output.view(batch_size, seq_len, d_model)
        
        return self.Wo(output)

个人经验：我建议num_heads一般设为8或16。头数太少，表达能力不够；头数太多，每个头的维度太小，学不到啥东西。d_k至少要有32，这是我踩坑踩出来的经验。

3.4 避坑指南

我曾经在做一个文本分类项目时，把num_heads设成了32，d_model才256。结果每个头的维度只有8，模型怎么训都训不上去。后来改成8个头，效果立竿见影。

还有一点：多头注意力的计算量是O(n²·d)，序列长度n一长，显存就炸了。我建议在实际应用中，如果序列超过512，可以考虑用稀疏注意力或者窗口注意力来优化。

参数 推荐值 说明
num_heads 8 ~ 16 头数适中，平衡表达力和计算量
d_k ≥ 32 每个头的维度不能太小
d_model 256 ~ 1024 根据任务复杂度调整

嗯，总结一下：自注意力就是让每个词去「关注」所有词，多头注意力就是让模型从多个角度去「关注」。拆分是为了并行学习不同关系，合并是为了把学到的信息整合起来。这个设计，说实话，真的很巧妙。

← 上一章 📖 返回目录 下一章 →