共计 2219 个字符,预计需要花费 6 分钟才能阅读完成。
序列建模的架构演进
传统 RNN 面临梯度消失和顺序计算的固有缺陷。对于长度为 $n$ 的序列,RNN 的时间复杂度为 $O(n)$ 但存在梯度连乘问题:
$$ \frac{\partial L}{\partial h_t} = \prod_{k=t}^{T} \frac{\partial h_{k+1}}{\partial h_k} \cdot \frac{\partial L}{\partial h_T} $$

CNN 通过卷积核局部感知的特性,虽然实现了并行计算(时间复杂度 $O(log_k n)$),但长距离依赖需要堆叠多层。当序列长度达到 2048 时,典型 CNN 需要至少 11 层才能建立全局连接。
Transformer 核心机制
自注意力矩阵运算
给定输入矩阵 $X \in \mathbb{R}^{n \times d_{model}}$,通过线性变换得到 Q /K/V:
$$ Q = XW_Q, \quad K = XW_K, \quad V = XW_V $$
注意力权重通过缩放点积计算:
$$ Attention(Q,K,V) = softmax(\frac{QK^T}{\sqrt{d_k}})V $$
可视化过程:
1. $QK^T$ 矩阵反映 token 间相关性
2. $\sqrt{d_k}$ 缩放避免梯度饱和
3. softmax 归一化得到概率分布
位置编码设计
使用三角函数编码绝对位置:
$$ PE_{(pos,2i)} = sin(pos/10000^{2i/d_{model}}) $$
$$ PE_{(pos,2i+1)} = cos(pos/10000^{2i/d_{model}}) $$
该设计使得模型可以学习到:
– 相对位置关系:存在线性变换 $PE_{pos+k}$ 可表示为 $PE_{pos}$ 的线性函数
– 距离感知:波长形成几何级数,覆盖不同距离尺度
多头注意力并行化
将 $d_{model}$ 维度分割为 $h$ 个头:
$$ head_i = Attention(XW_Q^i, XW_K^i, XW_V^i) $$
$$ MultiHead(X) = Concat(head_1,…,head_h)W_O $$
并行优势体现在:
1. 计算分块后矩阵乘法可完全并行
2. 不同头学习多样化注意力模式
3. 计算复杂度保持在 $O(n^2 \cdot d_{model})$
PyTorch 实现解析
import torch
import torch.nn as nn
import einops
from torch.nn.functional import scaled_dot_product_attention
class MultiHeadAttention(nn.Module):
def __init__(self, d_model=512, h=8):
super().__init__()
self.d_k = d_model // h
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):
q = einops.rearrange(self.W_q(x), 'b n (h d) -> b h n d', h=h)
k = einops.rearrange(self.W_k(x), 'b n (h d) -> b h n d', h=h)
v = einops.rearrange(self.W_v(x), 'b n (h d) -> b h n d', h=h)
# 工业级 mask 处理
if mask is not None:
mask = mask.unsqueeze(1) # 广播到所有头
attn = scaled_dot_product_attention(q, k, v, attn_mask=mask)
else:
attn = scaled_dot_product_attention(q, k, v)
return self.W_o(einops.rearrange(attn, 'b h n d -> b n (h d)'))
GPU 优化技巧:
– 使用 torch.backends.cuda.sdp_kernel() 启用 Flash Attention
– 梯度检查点时禁用 VJP 计算
性能实测数据
| 层数 | 头数 | FLOPs (T) | 显存(GB) |
|---|---|---|---|
| 12 | 8 | 3.2 | 6.4 |
| 24 | 16 | 12.8 | 18.7 |
| 36 | 32 | 28.3 | 34.2 |
KV Cache 优化方案:
1. 推理时缓存 $K_t, V_t \in \mathbb{R}^{n \times d_k}$
2. 计算复杂度从 $O(n^2)$ 降为 $O(n)$
3. 内存增长为 $O(n)$ 而非 $O(n^2)$
工程实践要点
LayerNorm 放置策略
- Pre-LN 结构更稳定:
# 残差连接前标准化 x = x + Dropout(Attention(LayerNorm(x))) - Post-LN 需要精细调参但理论容量更大
混合精度训练
- 使用
torch.cuda.amp自动管理 - 关键位置保持 FP32:
- LayerNorm 输出
- Softmax 输入
- 损失函数计算
开放性问题
- 线性注意力变体(如 Performer)在长序列任务中能否保持性能?
- 理论复杂度 $O(n)$ vs $O(n^2)$
-
实际任务中的近似误差
-
深度与延迟的帕累托前沿:
- 模型压缩技术(蒸馏 / 量化)的收益曲线
- 硬件特性对最优架构的影响
Transformer 架构通过数学上的优雅设计,在工程实现中展现出惊人的扩展性。随着硬件与算法的协同进化,其潜力边界仍在不断拓展。
