Attention 03：单头注意力的完整计算流程

返回专题总目录 · 上一讲：Q / K / V 直觉

01 > 02 > [03] > 04 > 05 > 06 | 07 > 08 > 09 > 10 > 11 > 12

“如果你不能用一个 3 词的句子手推一遍 Attention，那你还没有真正理解它。“

这一讲要回答什么#

上一讲我们用”查资料”的类比理解了 Q / K / V 的角色。但类比终究是类比——模型真正跑的是矩阵乘法和 softmax。

这一讲只回答一个问题：给定一个短句，Attention 从输入到输出，数学上每一步到底是怎么算的？

目标是把下面这条公式链一次讲完整：

我们不会只给公式——我们会用一个 3 词的短句，手推每一步的维度和数值。

先把维度搞清楚#

这是所有初学者最容易卡住的地方，所以我们先把维度全部写明白。

假设输入序列有 $n = 3$ 个词，每个词的嵌入维度 $d_{\text{model}} = 4$，注意力头的维度 $d_k = d_v = 3$。

用一张表总结所有维度：

维度的关键规律：不管序列多长，$QK^\top$ 永远是 $(n, n)$——这就是 Attention 时间复杂度 $O(n^2)$ 的来源。

第一步：输入矩阵 $X$#

假设我们的句子是 [猫, 坐, 地毯]，每个词已经通过 Embedding 层变成了 4 维向量：

• 第 1 行：$\mathbf{x}_1$ =“猫”的嵌入
• 第 2 行：$\mathbf{x}_2$ =“坐”的嵌入
• 第 3 行：$\mathbf{x}_3$ =“地毯”的嵌入

第二步：线性投影得到 Q、K、V#

接下来，$X$ 分别乘以三个权重矩阵：

投影之后的结果（使用演示中的示例权重）：

对应的 PyTorch 代码：

import torch
import torch.nn as nn

# 参数
n = 3          # 序列长度
d_model = 4    # 嵌入维度
d_k = 3        # 注意力头维度

# 输入：3 个词，每个 4 维
X = torch.tensor([
    [1.0, 0.5, -0.2, 0.8],   # 猫
    [0.3, -0.6, 1.2, 0.1],   # 坐
    [0.7, 0.9, 0.4, -0.5],   # 地毯
])  # shape: (3, 4)

# 三个投影矩阵（实际训练中是可学习参数）
W_Q = nn.Linear(d_model, d_k, bias=False)
W_K = nn.Linear(d_model, d_k, bias=False)
W_V = nn.Linear(d_model, d_k, bias=False)

# 投影
Q = W_Q(X)    # (3, 4) @ (4, 3) → (3, 3)
K = W_K(X)    # (3, 4) @ (4, 3) → (3, 3)
V = W_V(X)    # (3, 4) @ (4, 3) → (3, 3)

print(f"Q shape: {Q.shape}")  # torch.Size([3, 3])
print(f"K shape: {K.shape}")  # torch.Size([3, 3])
print(f"V shape: {V.shape}")  # torch.Size([3, 3])

第三步：计算注意力分数 $QK^\top$#

这是 Attention 最核心的一步——让每个 Query 和所有 Key 做点积，得到一个”分数矩阵”。

矩阵中第 $i$ 行第 $j$ 列的值 $S_{ij} = \mathbf{q}_i \cdot \mathbf{k}_j$，表示位置 $i$ 对位置 $j$ 的原始相关性分数。

为什么是 $QK^\top$ 而不是 $QK$？

$Q$ 的形状是 $(3, 3)$，$K$ 的形状也是 $(3, 3)$。如果直接做 $QK$，那是 $(3, 3) \times (3, 3)$——虽然维度上可行，但含义不对。

我们需要的是让 $Q$ 的每一行（一个 Query）和 $K$ 的每一行（一个 Key）做点积。转置 $K$ 之后：

结果中第 $i$ 行第 $j$ 列 = $\mathbf{q}_i$ 和 $\mathbf{k}_j$ 的点积。这正是我们想要的。

对应的 PyTorch 代码：

# 计算原始注意力分数
scores = Q @ K.transpose(-2, -1)  # (3, 3) @ (3, 3)^T → (3, 3)
print(scores)
# 每个 scores[i][j] = Q 第 i 行 · K 第 j 行 = q_i · k_j

第四步：除以 $\sqrt{d_k}$ ——缩放#

为什么要除以 $\sqrt{d_k}$？这不是可选的 trick，而是必要的数值稳定措施。

直觉解释：当 $d_k$ 很大时，$\mathbf{q} \cdot \mathbf{k} = \sum_{i=1}^{d_k} q_i k_i$ 的结果也会很大（因为更多项求和）。大的分数值传入 softmax 后会导致梯度消失——输出接近 one-hot，模型很难学到”分散注意力”的模式。

数学上：假设 $q_i, k_i$ 独立、均值为 0、方差为 1，那么：

除以 $\sqrt{d_k}$ 后方差变为 1，分数不会因为维度增大而膨胀。

对应的 PyTorch 代码：

import math

d_k = Q.size(-1)  # = 3
scores_scaled = scores / math.sqrt(d_k)
# 等价于：scores / (3 ** 0.5) ≈ scores / 1.732

第五步：Softmax → 注意力权重#

Softmax 对分数矩阵的每一行独立操作，将分数转化为概率分布：

每一行加起来等于 1——这就是”注意力权重”的含义。

对应的 PyTorch 代码：

import torch.nn.functional as F

# softmax 在最后一个维度上（即每一行独立做 softmax）
attn_weights = F.softmax(scores_scaled, dim=-1)  # (3, 3)

print(attn_weights)
# 每行之和 == 1
print(attn_weights.sum(dim=-1))  # tensor([1., 1., 1.])

第六步：加权求和 → 输出#

用注意力权重矩阵 $\alpha$ 左乘 Value 矩阵 $V$，得到最终输出：

输出的第 $i$ 行是所有 Value 向量按第 $i$ 行权重的加权平均：

对应的 PyTorch 代码：

# 加权求和
output = attn_weights @ V  # (3, 3) @ (3, 3) → (3, 3)
print(f"Output shape: {output.shape}")  # torch.Size([3, 3])

# output[0] = α[0,0]*V[0] + α[0,1]*V[1] + α[0,2]*V[2]
# "猫"位置的输出 = 融合了"猫""坐""地毯"信息的新表示

把所有步骤串起来：完整的 Attention 公式#

完整的 PyTorch 实现#

把上面所有步骤合在一起，一个完整的单头注意力就是这么简短：

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

class SingleHeadAttention(nn.Module):
    """单头 Scaled Dot-Product Attention"""

    def __init__(self, d_model: int, d_k: int):
        super().__init__()
        self.d_k = d_k
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_k, bias=False)

    def forward(self, X: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            X: (n, d_model) — 输入序列
        Returns:
            output: (n, d_k)  — 融合上下文后的表示
            attn:   (n, n)    — 注意力权重矩阵
        """
        Q = self.W_Q(X)                              # (n, d_k)
        K = self.W_K(X)                              # (n, d_k)
        V = self.W_V(X)                              # (n, d_k)

        scores = Q @ K.transpose(-2, -1)             # (n, n)
        scores_scaled = scores / math.sqrt(self.d_k) # (n, n)
        attn = F.softmax(scores_scaled, dim=-1)      # (n, n)
        output = attn @ V                            # (n, d_k)

        return output, attn

# ---- 使用 ----
attn_layer = SingleHeadAttention(d_model=4, d_k=3)

X = torch.tensor([
    [1.0, 0.5, -0.2, 0.8],
    [0.3, -0.6, 1.2, 0.1],
    [0.7, 0.9, 0.4, -0.5],
])

output, weights = attn_layer(X)
print(f"Output:\n{output}")
print(f"\nAttention weights:\n{weights}")
print(f"Rows sum to 1: {weights.sum(dim=-1)}")

运行这段代码，你会得到：

• output: $(3, 3)$ 矩阵——每个词融合上下文后的新表示
• weights: $(3, 3)$ 矩阵——每个词对其他词的注意力分配

逐行对照：公式 ↔ 代码#

四个容易搞混的问题#

1. 为什么 $QK^\top$ 是”词与词的关系矩阵”？#

$S_{ij} = \mathbf{q}_i \cdot \mathbf{k}_j$ 是 Query $i$ 和 Key $j$ 的点积。点积衡量两个向量的方向相似性——方向越接近，值越大。

所以 $S_{ij}$ 大 ⟹ 位置 $i$ 认为位置 $j$ 和自己”相关”——注意这里的”相关”是模型通过训练学到的，不是固定的语言学标签。

2. 为什么 softmax 要对每一行做？#

因为我们在为每个 Query 位置生成一个完整的注意力分布。第 $i$ 行表示”位置 $i$ 如何把注意力分配给所有位置”。

如果对列做 softmax，含义就变了——变成”位置 $j$ 被各个 Query 竞争的程度”，这不是我们想要的。

3. 为什么最后输出还是和输入一样的”一个词一个向量”？#

因为 $\alpha V$ 的结果是 $(n, d_v)$——仍然是 $n$ 个向量。但每个向量的含义变了：

• 输入 $\mathbf{x}_i$：词 $i$ 的静态嵌入，不知道上下文
• 输出 $\mathbf{o}_i$：词 $i$ 融合了整个序列上下文后的动态表示

形式相同，但信息量完全不同。

4. $d_k$ 和 $d_{\text{model}}$ 什么关系？#

在多头注意力中，通常 $d_k = d_{\text{model}} / h$，其中 $h$ 是头数。比如 $d_{\text{model}} = 512, h = 8$ 时，$d_k = 64$。

但在单头情况下，$d_k$ 可以自由设定。关键约束只有一个：$d_k$ 越大，$QK^\top$ 的方差越大，所以必须除以 $\sqrt{d_k}$。

一张图总结完整流程#

每一步的形状变化：

试着想一想#

1. 如果把 $\sqrt{d_k}$ 换成 $d_k$（不取根号），softmax 的输出会有什么变化？
2. 在 $QK^\top$ 的结果中，对角线 $S_{ii}$ 表示什么？它一定最大吗？
3. 如果 $n = 1000$，$QK^\top$ 有多少个元素？这就是为什么长序列计算量大的原因。
4. 修改上面的 SingleHeadAttention 类，给它加一个 causal mask 参数，让位置 $i$ 只能看到 $j \leq i$ 的位置。（提示：在 softmax 之前把未来位置的分数设为 $-\infty$。）
5. 运行代码，比较 output[0] 和 V[0]——输出和”猫”自己的 Value 差多少？这个差距就是上下文信息的融合效果。

下一讲#

下一讲我们深入 $QK^\top$，回答一个更本质的问题：

为什么偏偏是点积来衡量词与词的关系？为什么不是别的形式？