引言
2017年，Vaswani等人在《Attention Is All You Need》中提出的Transformer架构彻底改变了自然语言处理领域。它完全摒弃了循环结构，代之以自注意力（Self-Attention）机制，实现了序列中任意位置之间的直接信息交互。本文从数学角度深入拆解Transformer的核心组件——Scaled Dot-Product Attention、Multi-Head Attention、位置编码和前馈网络。
理论基础
自注意力的数学定义
给定输入序列 $X \in \mathbb{R}^{n \times d}$（$n$ 个 token，维度 $d$），通过三个可学习矩阵 $W^Q, W^K, W^V \in \mathbb{R}^{d \times d_k}$ 计算查询（Query）、键（Key）、值（Value）：
$$Q = X W^Q, \quad K = X W^K, \quad V = X W^V$$
注意力权重矩阵为：
$$\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^T}{\sqrt{d_k}} \right) V$$
缩放因子 $\sqrt{d_k}$ 的作用：当 $d_k$ 较大时，$Q K^T$ 的元素量级约为 $\mathcal{O}(\sqrt{d_k})$，经 softmax 后进入梯度饱和区（接近 0 或 1），缩放可保持梯度更新量。
Multi-Head Attention
将 $Q/K/V$ 切分为 $h$ 个头，每个头学习不同的注意力模式：
$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W^O$$
其中 $\text{head}_i = \text{Attention}(Q W_i^Q, K W_i^K, V W_i^V)$。
前馈网络
每个 Transformer 层后接一个 Position-wise Feed-Forward Network：
$$\text{FFN}(x) = \max(0, x W_1 + b_1) W_2 + b_2$$
位置编码
Transformer 本身不具备序列顺序感知。加法位置编码公式：
$$PE_{(pos, 2i)} = \sin(pos / 10000^{2i/d_{\text{model}}})$$
$$PE_{(pos, 2i+1)} = \cos(pos / 10000^{2i/d_{\text{model}}})$$
数值实现
import torch
import torch.nn as nn
import math
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.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
self.W_o = nn.Linear(d_model, d_model)
def forward(self, x, mask=None):
batch, seq_len, d_model = x.shape
Q = self.W_q(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
K = self.W_k(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
V = self.W_v(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
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)
out = torch.matmul(attn, V).transpose(1, 2).contiguous()
return self.W_o(out.view(batch, seq_len, d_model))参数统计与复杂度分析
以 GPT-3（175B 参数）为例：
| 组件 | 参数量 | 复杂度 |
|---|---|---|
| Embedding | $V d_{\text{model}}$ | — |
| 每层 MHA | $4 d_{\text{model}}^2$ | $\mathcal{O}(n^2 d_{\text{model}})$ |
| 每层 FFN | $8 d_{\text{model}}^2$ | $\mathcal{O}(n d_{\text{model}}^2)$ |
| 总层数 | 96 层 | — |
自注意力复杂度 $\mathcal{O}(n^2 d)$ 是长序列处理的瓶颈，FlashAttention（Dao et al., 2022）通过 IO-aware 分层计算将内存需求从 $\mathcal{O}(n^2)$ 降至 $\mathcal{O}(n)$。
结论与展望
Transformer 的核心优势在于并行化训练和全局信息交互。当前前沿包括：FlashAttention 加速、Longformer/BigBird 等稀疏注意力变体、Transformer + RLHF 的对齐训练范式。理解 Self-Attention 的数学本质，是驾驭大模型时代的关键。
发表回复