共计 2411 个字符，预计需要花费 7 分钟才能阅读完成。
背景：RNN 的局限性
传统 RNN 在处理长序列时存在两个主要缺陷：
- 梯度消失 / 爆炸问题 ：随着序列长度增加，反向传播时梯度会指数级衰减或增长，导致难以训练
- 顺序计算瓶颈 ：必须按时间步顺序计算，无法充分利用 GPU 的并行计算能力
而 self-attention 机制通过全局建模任意位置的关系，完美解决了这些问题：
- 任意两个 token 的直接交互（O(1) 路径长度）
- 可并行计算的矩阵运算
核心技术解析
Q/K/ V 矩阵的数学本质
定义输入序列 $X \in \mathbb{R}^{n\times d_{model}}$，通过三个可学习矩阵变换得到：
$$
\begin{aligned}
Q &= XW^Q, \quad W^Q \in \mathbb{R}^{d_{model}\times d_k} \
K &= XW^K, \quad W^K \in \mathbb{R}^{d_{model}\times d_k} \
V &= XW^V, \quad W^V \in \mathbb{R}^{d_{model}\times d_v}
\end{aligned}
$$
- $Q$ (Query): 表示当前需要关注的内容
- $K$ (Key): 表示可供关注的上下文信息
- $V$ (Value): 实际要聚合的信息
Scaled Dot-Product Attention
完整计算流程：
-
计算 query 和 key 的相似度：
$$\text{Attention}(Q, K, V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V$$ -
缩放因子 $\sqrt{d_k}$ 防止点积结果过大导致 softmax 饱和
- 注意力权重矩阵 $A=\text{softmax}(\frac{QK^T}{\sqrt{d_k}})$ 反映 token 间相关性
多头注意力机制
将 Q /K/ V 拆分为 $h$ 个头（常用 8 -16 个）：
$$
\text{MultiHead}(Q, K, V) = \text{Concat}(head_1,…,head_h)W^O
$$
其中每个头的计算：
$$
head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)
$$
优势：
- 并行计算：各头可独立计算
- 多子空间表示：捕获不同方面的依赖关系
PyTorch 实现细节
基础实现
import torch
import torch.nn as nn
from einops import rearrange
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
self.n_heads = n_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, mask=None):
# x: [batch, seq_len, d_model]
batch_size = x.size(0)
# 1. 线性投影
Q = self.Wq(x) # [batch, seq_len, d_model]
K = self.Wk(x)
V = self.Wv(x)
# 2. 分割多头
Q = rearrange(Q, 'b s (h d) -> b h s d', h=self.n_heads)
K = rearrange(K, 'b s (h d) -> b h s d', h=self.n_heads)
V = rearrange(V, 'b s (h d) -> b h s d', h=self.n_heads)
# 3. 计算注意力
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
attn = torch.softmax(scores, dim=-1)
# 4. 聚合 value
out = torch.matmul(attn, V) # [batch, h, seq_len, d_k]
out = rearrange(out, 'b h s d -> b s (h d)')
return self.Wo(out)关键优化技术
-
梯度检查点 ：
from torch.utils.checkpoint import checkpoint def custom_forward(*inputs): Q, K, V = inputs # 计算注意力... return output out = checkpoint(custom_forward, Q, K, V) -
混合精度训练 ：
with torch.cuda.amp.autocast(): attn_output = mha(x)
性能实测数据
| 头数 | 推理时延 (ms) | GPU 显存 (GB) |
|---|---|---|
| 4 | 12.3 | 2.1 |
| 8 | 15.7 | 2.8 |
| 16 | 22.4 | 4.2 |
工程避坑指南
- LayerNorm 放置位置 ：
- Pre-LN：更稳定但收敛慢
-
Post-LN：需要精细调参
-
数值溢出防护 ：
# 在 softmax 前 clamp 极值 scores = scores.clamp(-50, 50) -
维度对齐检查 ：
assert Q.shape == K.shape == V.shape
延伸思考方向
- 稀疏注意力 ：
- Local window attention
-
Block-sparse patterns
-
混合架构设计 ：
- CNN 提取局部特征 + Attention 建模全局关系
- 计算效率与建模能力的 trade-off
总结
自注意力机制通过灵活的全局交互能力，已成为现代 NLP 架构的核心组件。理解其底层实现细节对模型调优和定制开发至关重要。建议读者在实际项目中多尝试不同头数配置和优化策略，观察对任务性能的具体影响。
