共计 2487 个字符，预计需要花费 7 分钟才能阅读完成。
背景痛点
刚接触 Transformer 时，我最头疼的就是各种教程一上来就贴大段代码，却没人说清楚为什么自注意力要计算 $QK^T$，或者位置编码为什么要用三角函数。直到看到 3Blue1Brown 的视频，才发现原来这些操作背后都有直观的几何意义。这篇文章就想用这种可视化思维，带大家从数学原理推导到 PyTorch 实现。
原理可视化
位置编码的波长之谜
传统 RNN 通过时间步记录位置信息，而 Transformer 的位置编码公式看起来有点神秘：
$$
PE_{(pos,2i)} = \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)
$$
其实这个设计非常巧妙：
- 波长衰减 ：$10000^{2i/d_{\text{model}}}$ 使得不同维度 $i$ 的波长呈指数增长，低频（大波长）捕获长距离依赖，高频（小波长）记录局部位置
- 相对位置 ：通过三角函数的和角公式可以证明，$PE_{pos+k}$ 能表示为 $PE_{pos}$ 的线性组合，这正是模型能学到相对位置的关键
自注意力的几何视角
计算 $QK^T$ 时，很多人以为是简单的矩阵乘法。其实从几何看：
- 每个查询向量 $q$ 与所有键向量 $k$ 做点积，本质上是在计算向量夹角余弦值（相似度）
- 缩放因子 $\sqrt{d_k}$ 防止点积结果过大导致 softmax 饱和（后面会详细解释）
- 可视化注意力权重时，你会看到类似「光照强度」的分布，亮的区域就是模型聚焦的位置
PyTorch 实现
多头注意力实现
先看最核心的 MultiHeadAttention 类（已精简非关键代码）：
import torch
import torch.nn as nn
import torch.nn.functional as F
class MultiHeadAttention(nn.Module):
def __init__(self, d_model=512, n_heads=8):
super().__init__()
assert d_model % n_heads == 0
self.d_k = d_model // n_heads
# 用一个大矩阵并行计算所有头的 QKV
self.qkv_proj = nn.Linear(d_model, 3*d_model)
self.out_proj = nn.Linear(d_model, d_model)
def forward(self, x, mask=None):
batch_size = x.size(0)
# 1. 生成 QKV 并分头 [B,L,D] -> [B,L,3*D] -> 3×[B,H,L,D_k]
qkv = self.qkv_proj(x)
q, k, v = torch.chunk(qkv, 3, dim=-1)
q = q.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
# 2. 用 einsum 计算注意力分数 [B,H,L,L]
scores = torch.einsum('bhqd,bhkd->bhqk', q, k) / self.d_k**0.5
# 3. 掩码处理（解码器需要）if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
# 4. softmax 归一化
attn = F.softmax(scores, dim=-1)
# 5. 加权求和 [B,H,L,D_k]
out = torch.einsum('bhql,bhld->bhqd', attn, v)
# 6. 合并所有头
out = out.transpose(1, 2).contiguous()
out = out.view(batch_size, -1, self.d_model)
return self.out_proj(out)关键技巧解析
- einsum 魔法 ：用
torch.einsum('bhqd,bhkd->bhqk', q, k)一句搞定多头矩阵乘法，避免手动 permute 维度 - 合并计算 ：通过
nn.Linear(3*d_model)一次性计算 QKV，比分开三次计算更快 - contiguous()：转置操作后必须调用以保证内存连续性，否则 view 会报错
避坑指南
梯度爆炸问题
当 $d_k$ 较大时，$QK^T$ 的值可能非常大，导致 softmax 进入饱和区。解决方法：
- 一定要用缩放因子 $1/\sqrt{d_k}$
- 初始化时适当调小线性层的权重（如用
nn.init.xavier_normal_(layer.weight, gain=0.02)）
验证注意力权重
调试时可以检查：
- 每行的注意力权重求和是否为 1（
torch.allclose(attn.sum(dim=-1), torch.ones_like(...))） - 掩码位置的值是否真的是负无穷（
attn.masked_fill(mask==0,0).sum() == attn.sum()）
延伸思考
缩放因子的数学原理
为什么是 $\sqrt{d_k}$ 而不是其他值？推导如下：
假设 $q$ 和 $k$ 的分量是独立同分布的随机变量，均值为 0，方差为 1，则点积 $q \cdot k = \sum_{i=1}^{d_k} q_i k_i$ 的方差就是 $d_k$。缩放 $1/\sqrt{d_k}$ 正好使方差回归到 1，保持梯度稳定性。
可视化复现方法
在 Jupyter Notebook 中可视化注意力矩阵：
# 示例：绘制第一个头的注意力热图
import matplotlib.pyplot as plt
plt.matshow(attn[0,0].detach().numpy(), cmap='Reds')
plt.colorbar()
plt.title("Attention Head 1")总结
通过 3Blue1Brown 的几何视角，我们理解了 Transformer 那些「反直觉」设计背后的数学之美。实践中建议：
1. 先用小规模数据调试注意力机制
2. 善用 einsum 处理复杂张量运算
3. 官方 nn.Transformer 已经高度优化，生产环境建议直接使用
下次见到 $QK^T/\sqrt{d_k}$ 时，希望你能想起这篇文章——它不只是公式，更是一幅动态的几何画卷。
