Transformer与自注意力机制深度解析：从数学原理到PyTorch实现

1次阅读

没有评论

共计 2721 个字符，预计需要花费 7 分钟才能阅读完成。

在 Transformer 出现之前，RNN（循环神经网络）是处理序列数据的主流架构。但 RNN 存在两个致命缺陷：

顺序计算的瓶颈 ：必须逐个时间步处理序列，无法充分利用 GPU 的并行计算能力。即使使用 LSTM/GRU 缓解梯度消失，长距离依赖问题仍难以解决。
信息传递效率低 ：隐藏状态需要像接力棒一样依次传递，早期时间步的信息在传递过程中不断衰减。

Transformer 通过自注意力机制彻底改变了这一局面，使模型能够直接计算任意两个位置的关系，实现了：

完全并行的序列处理
恒定步数的远程依赖捕获
动态权重分配（不同位置关注不同内容）

自注意力的计算可分为三步：

投影得到 Q /K/V：
```
Q = XW^Q, \quad K = XW^K, \quad V = XW^V
```
其中 $X \in \mathbb{R}^{n \times d_{model}}$ 是输入序列，$W$ 为可学习参数矩阵。
计算注意力分数 ：
```
\text{Attention}(Q,K,V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V
```
除以 $\sqrt{d_k}$ 是为了防止点积结果过大导致 softmax 梯度消失。
多头注意力拼接 ：
```
\text{MultiHead} = \text{Concat}(head_1,...,head_h)W^O
```
其中每个 $head_i$ 独立进行自注意力计算。

Q/K/ V 分解 ：将输入映射到不同子空间，分别承担 ” 查询 ”、” 键 ”、” 值 ” 角色
并行计算 ：每个注意力头可独立计算，GPU 利用率显著提升
动态权重 ：softmax 输出的注意力权重随着输入内容动态变化

import torch
import torch.nn as nn
import math

class SelfAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"

        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 /K/V [batch, seq_len, d_model]
        q = self.wq(x)
        k = self.wk(x)
        v = self.wv(x)

        # 2. 切分为多头 [batch, seq_len, n_heads, d_k]
        q = q.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        k = k.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        v = v.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

        # 3. 计算缩放点积注意力
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)

        # 4. 应用 mask（解码器用）if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # 5. softmax 归一化
        attn = torch.softmax(scores, dim=-1)

        # 6. 加权求和并拼接多头
        out = torch.matmul(attn, v)  # [batch, n_heads, seq_len, d_k]
        out = out.transpose(1, 2).contiguous()  # [batch, seq_len, n_heads, d_k]
        out = out.view(batch_size, -1, self.d_model)  # [batch, seq_len, d_model]

        return self.wo(out)

关键点说明：

维度变换 ：通过 view 和 transpose 操作实现多头切分
mask 处理 ：解码时防止看到未来信息
contiguous()：确保内存连续提升计算效率

自注意力的内存消耗与序列长度 $n$ 呈 $O(n^2)$ 关系，这是因为：

$QK^T$ 矩阵形状为 $[n, n]$
当 $n=2048$ 时，单精度浮点矩阵占用 16GB 内存

分块计算（Memory-efficient Attention）
将大矩阵拆分为小块计算
需要手动实现累积 softmax
稀疏注意力（Sparse Transformer）
只计算特定位置的注意力（如局部窗口 / 固定步长）
牺牲部分精度换取内存
FlashAttention（推荐）
利用 GPU 共享内存减少 HBM 访问
官方实现速度提升 2 - 4 倍

梯度控制
在残差连接后立即使用 LayerNorm
初始阶段可设置梯度裁剪（clip_grad_norm_=1.0）

超参数验证

assert d_model % n_heads == 0, "特征维度必须能被头数整除"
assert head_size >= 32, "过小的头尺寸影响表达能力"

混合精度训练

scaler = torch.cuda.amp.GradScaler()

with torch.autocast(device_type='cuda', dtype=torch.float16):
    outputs = model(inputs)
    loss = criterion(outputs, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

实际部署时发现几个易错点：