2、核心公式拆解：Softmax(QKT/√d)V 的数学推导与几何意义

好，咱们直接进入正题。Self-Attention 的核心，说白了就是下面这个公式：

Attention(Q, K, V) = Softmax(QKT / √dk) · V

我第一次看到这个公式时，心里想的是：「这玩意儿到底在干嘛？」后来在项目中调了无数次模型，才真正理解它的几何直觉。今天我就带你一步步拆开它。

2.1 从「查字典」说起

想象你在查字典。你要找「苹果」这个词的解释。你会怎么做？

  • Query（查询）：就是你要查的词——「苹果」
  • Key（键）：字典里所有的词条标题
  • Value（值）：每个词条对应的解释内容

Self-Attention 干的也是类似的事。每个 token 都会发出一个「查询」，然后跟所有其他 token 的「键」做匹配，最后根据匹配程度加权汇总「值」。

嗯，这里要注意：在 Self-Attention 里，Q、K、V 都是从同一个输入 X 经过线性变换得到的。也就是说，每个 token 既是查询者，也是被查询的对象。

2.2 QKT：相似度计算

先看 QKT 这部分。Q 的形状是 (seq_len, dk)，K 的形状也是 (seq_len, dk)。

QKT 的结果是一个 (seq_len, seq_len) 的矩阵。第 i 行第 j 列的元素，就是第 i 个 token 的 Query 与第 j 个 token 的 Key 的点积。

点积的几何意义：两个向量的点积越大，说明它们的方向越一致，相似度越高。

说白了，QKT 就是在计算「每个 token 对每个其他 token 的关注程度」。数值越大，说明这两个位置越「相关」。

我在项目中遇到过一个问题：当序列长度很长时，QKT 的计算量会爆炸。因为它是 O(n²) 的复杂度。这也是为什么后来有了各种稀疏注意力、线性注意力的变体。

2.3 为什么要除以 √dk？

这是很多初学者会忽略的细节。为什么非要除以根号 dk？

原因很简单：防止梯度消失。

你想想看，dk 是每个头的维度，通常是 64 或 128。两个 dk 维的向量做点积，结果的方差是 dk。也就是说，dk 越大，点积的数值就越大。

如果点积数值很大，Softmax 之后会怎么样？

  • 大的值会更大，接近 1
  • 小的值会更小，接近 0

这就导致 Softmax 的输出几乎成了 one-hot 向量。梯度会变得非常小，模型几乎学不动。

我的习惯：除以 √dk 之后，点积的方差被拉回到 1 左右。这样 Softmax 的输入分布更均匀，梯度更健康。

我曾经在训练一个 Transformer 模型时，忘了加这个缩放因子。结果 loss 死活降不下去。排查了半天才发现是这个问题。嗯，从那以后我再也不敢漏掉 √dk 了。

2.4 Softmax：归一化注意力权重

Softmax 的作用，就是把 QKT/√dk 的结果变成概率分布。每一行的和都是 1。

公式很简单：

Softmax(x_i) = exp(x_i) / Σ exp(x_j)

但这里有个坑：数值稳定性。

如果 x_i 很大，exp(x_i) 会溢出。所以实际实现时，通常会先减去最大值：

def softmax(x):
    # 减去最大值，防止 exp 溢出
    x_max = np.max(x, axis=-1, keepdims=True)
    x_safe = x - x_max
    exp_x = np.exp(x_safe)
    return exp_x / np.sum(exp_x, axis=-1, keepdims=True)

避坑指南：我曾经在 CUDA 实现中直接调用了 exp(x)，结果遇到大数值时直接 NaN。后来才意识到要做 max-shift。这个细节在框架层面已经被封装好了，但如果你自己写底层实现，一定要小心。

2.5 乘以 V：加权求和

最后一步，把 Softmax 得到的权重矩阵乘以 V。

V 的形状是 (seq_len, dv)。Softmax 的结果是 (seq_len, seq_len)。

相乘之后，得到 (seq_len, dv) 的输出。

几何意义是什么？

对于第 i 个 token，它的输出是：

  • 所有 token 的 Value 的加权平均
  • 权重就是第 i 个 token 对其他 token 的注意力分数

说白了，每个 token 的输出，是「看遍了所有其他 token 之后，把相关信息聚合到自己身上」的结果。

这就是 Self-Attention 的「全局感受野」。跟 CNN 那种局部卷积完全不同。

2.6 完整代码实现

下面是一个极简的 PyTorch 实现。去掉了一切花哨的东西，只保留核心逻辑：

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

class ScaledDotProductAttention(nn.Module):
    def __init__(self, d_k):
        super().__init__()
        self.d_k = d_k
        
    def forward(self, Q, K, V, mask=None):
        """
        Q: (batch, seq_len, d_k)
        K: (batch, seq_len, d_k)
        V: (batch, seq_len, d_v)
        mask: (batch, seq_len, seq_len) 可选
        """
        # 1. 计算相似度
        scores = torch.matmul(Q, K.transpose(-2, -1))  # (batch, seq_len, seq_len)
        
        # 2. 缩放
        scores = scores / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))
        
        # 3. 可选：mask（用于 decoder 或 padding）
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        # 4. Softmax
        attn_weights = F.softmax(scores, dim=-1)
        
        # 5. 加权求和
        output = torch.matmul(attn_weights, V)  # (batch, seq_len, d_v)
        
        return output, attn_weights

个人建议：写代码时，一定要把张量的形状注释清楚。我见过太多人因为维度搞混而 debug 半天。尤其是 transpose 和 matmul 的组合，很容易出错。

2.7 一张表总结

步骤 数学操作 几何意义 注意事项
1 QKT 计算所有 token 对之间的相似度 复杂度 O(n²)
2 ÷ √dk 缩放点积，稳定梯度 dk 越大，缩放越重要
3 Softmax 归一化为概率分布 注意数值稳定性
4 × V 加权聚合全局信息 输出维度 = dv

2.8 为什么这个设计这么优雅？

你想想看，整个公式只有三个矩阵乘法加一个 Softmax。但它能做到：

  • 动态权重：权重不是固定的，而是根据输入动态计算
  • 全局交互：每个 token 都能看到所有其他 token
  • 并行计算：所有 token 可以同时计算，不像 RNN 那样串行

这就是为什么 Transformer 能统治 NLP 领域，现在又开始渗透到 CV、多模态等领域。

我记得有一次跟同事讨论，他说 Self-Attention 本质上是在学习一个「动态的图结构」。每个 token 是节点，注意力权重是边的强度。这个视角很有意思，也解释了为什么 Transformer 在关系推理任务上表现那么好。

好，这一章就到这里。下一章我们会深入 Multi-Head Attention，看看为什么多个头比单个头效果更好。

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