共计 2915 个字符，预计需要花费 8 分钟才能阅读完成。
RNN/LSTM 的困境
在处理长序列任务（如机器翻译、语音识别）时，传统 RNN/LSTM 架构存在两个致命缺陷：
-
梯度消失问题：随着序列长度增加，反向传播时梯度会指数级衰减，导致模型难以学习远距离依赖关系。实验显示，当序列长度超过 50 步时，LSTM 对早期信息的记忆保留率不足 30%。
-
无法并行计算：RNN 必须按时间步顺序计算，无法利用现代 GPU 的并行计算能力。即便使用 LSTM 的变体，处理 1000 长度的序列仍需要约 3 倍的实时计算时间。
自注意力机制原理
Transformer 提出的自注意力机制通过以下方式突破上述限制：
-
并行计算：所有位置的注意力得分可同时计算，公式简化为：
$$\text{Attention}(Q,K,V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V$$
其中 $Q$（查询）、$K$（键）、$V$（值）均来自同一输入序列的线性变换。 -
长程依赖建模：任意两个位置的距离均为 1 步矩阵运算，彻底解决梯度消失问题。实验表明，在 WMT14 英德翻译任务上，自注意力模型对 50 词以上依赖关系的捕捉准确率比 LSTM 高 42%。
PyTorch 实现
基础 MultiHeadAttention
import torch
import torch.nn as nn
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)
@torch.jit.script
def scaled_dot_product_attention(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
mask: torch.Tensor = None
) -> torch.Tensor:
"""
执行缩放点积注意力计算
参数:
q: [batch, heads, seq_len, d_k]
mask: [batch, 1, 1, seq_len] (可选)
"""
attn_scores = torch.matmul(q, k.transpose(-2, -1)) / (q.size(-1) ** 0.5)
if mask is not None:
attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
attn_weights = torch.softmax(attn_scores, dim=-1)
return torch.matmul(attn_weights, v)
def forward(self, x, mask=None, kv_cache=None):
batch_size = x.size(0)
# 线性投影 + 分头
q = self.W_q(x).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
k = self.W_k(x).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
v = self.W_v(x).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
# 使用 KV 缓存（推理优化）if kv_cache is not None:
k = torch.cat([kv_cache['k'], k], dim=2)
v = torch.cat([kv_cache['v'], v], dim=2)
# 计算注意力
attn_output = self.scaled_dot_product_attention(q, k, v, mask)
# 合并多头输出
output = attn_output.transpose(1, 2).contiguous() \
.view(batch_size, -1, self.n_heads * self.d_k)
return self.W_o(output), {'k': k, 'v': v}关键实现细节
- KV 缓存：在自回归生成时缓存历史 K /V，避免重复计算
- Mask 处理：
- 填充 mask（pad_mask）：忽略无效位置
- 因果 mask（causal_mask）：防止未来信息泄露
- 数值稳定性：缩放因子 $\sqrt{d_k}$ 防止点积结果过大导致 softmax 饱和
性能优化
Flash Attention
通过分块计算和算子融合，将 HBM 访问量从 $O(N^2)$ 降至 $O(N)$。PyTorch 2.0+ 原生支持：
with torch.backends.cuda.sdp_kernel(enable_flash=True):
attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)显存占用实验
| 序列长度 | 原始 Attention(MB) | Flash Attention(MB) |
|---|---|---|
| 512 | 1203 | 687 |
| 1024 | 4812 | 1354 |
| 2048 | 19248 | 2531 |
使用 torch.profiler 定位瓶颈
with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]
) as prof:
output = model(inputs)
print(prof.key_averages().table(sort_by="cuda_time_total"))典型输出会显示 matmul 和 softmax 操作的耗时占比。
避坑指南
- 数值稳定性：
- 必须使用缩放因子 $1/\sqrt{d_k}$
-
混合精度训练时建议使用
torch.nn.functional.scaled_dot_product_attention -
分布式训练：
- 使用
nn.parallel.DistributedDataParallel时，注意不同 GPU 间 Attention mask 的同步 -
推荐在
forward()开始时调用broadcast(mask, src=0) -
超长序列处理：
- 当序列长度 >1 万时，考虑使用内存高效的 Attention 变体
- 示例配置：
attn_impl = "flash" if seq_len < 8192 else "memory_efficient"
开放性问题
当序列长度突破 10 万量级时，稀疏注意力成为必选项。如何选择模式？
- 固定模式(Fixed)：适合有规律间隔的任务（如 DNA 序列）
- 跨步模式(Strided)：平衡局部和全局注意力（推荐默认配置）
- 随机模式(Random)：适合无预设结构的数据（如点云）
实际应用中，可先用小规模数据测试不同模式的效果，再结合 torch.profiler 分析计算开销。
