AI注意力即一切实践笔记

Transformer Attention机制 深度学习
作者: 醉月思 发布日期: 2024-08-17 分类: 学习笔记

当时我在做一个机器翻译的项目，用的还是基于 LSTM 的 seq2seq 模型，那个下午训练曲线又震荡得厉害，我盯着屏幕看了二十分钟，想的是：这玩意儿到底什么时候能收敛。

“Attention is all you need.” — Vaswani et al., 2017

这句话现在已经被说到滥大街了，但说实话，第一次读这篇论文时，我以为只是又一篇"换个架构刷个排行榜"的典型学术作。

从 RNN 的注意力说起

2017 年之前，NLP 基本是 RNN 的天下。LSTM/GRU 解决了梯度消失的问题，但序列依赖还在。句子越长，前面信息丢得越严重。那时候做机器翻译，我们用带注意力机制的 encoder-decoder：

# 经典的 Bahdanau attention（2014）
class BahdanauAttention(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.Wa = nn.Linear(hidden_size, hidden_size)
        self.Ua = nn.Linear(hidden_size, hidden_size)
        self.va = nn.Linear(hidden_size, 1)

    def forward(self, query, keys):
        # query: decoder hidden state [batch, hidden]
        # keys: encoder outputs [batch, seq_len, hidden]
        scores = self.va(torch.tanh(
            self.Wa(query).unsqueeze(1) + self.Ua(keys)
        ))  # [batch, seq_len, 1]
        scores = scores.squeeze(2)  # [batch, seq_len]
        weights = F.softmax(scores, dim=1)  # [batch, seq_len]
        context = torch.bmm(weights.unsqueeze(1), keys)  # [batch, 1, hidden]
        return context.squeeze(1), weights

这个公式现在看起来很朴素，但在当时解决了实际问题：decoder 不再只能看 encoder 最后一个状态，而是可以"回看"整个输入序列。权重可视化出来，确实能看到模型在翻译每个词时关注原文的哪些位置。

但这套架构有个硬伤：依然要串行计算。第一个 token 算完才能算第二个，第 N 个要等前 N-1 个都算完。训练时还好，推理时这就是个显眼的瓶颈。

2017 年那篇论文到底改变了什么

《Attention Is All You Need》这篇论文最刺眼的不是注意力公式，而是它直接把 RNN 扔了。

传统的 encoder-decoder 是这样：

graph LR A[Input] --> B[Encoder RNN] B --> C[Hidden State] C --> D[Decoder RNN] D --> E[Output]

Transformer 去掉了 RNN，全部用注意力：

graph LR A[Input] --> B[Self-Attention] B --> C[Feed Forward] C --> D[Cross Attention] D --> E[Feed Forward] E --> F[Output]

这看起来只是模块替换，但本质变化是：Self-Attention 是并行的。你不需要等前一个时间步算完，整个序列可以一次性扔进去。

但论文有个地方写得不够清楚：Multi-Head Attention 的具体实现。我当时照着公式写了一版，跑出来效果很差，排查了三天才发现是维度对齐搞错了。

踩坑实录：多头注意力的实现陷阱

论文里说 Multi-Head Attention 是把输入分别投影到不同的子空间：

$$ \text{MultiHead}(Q, K, V) = \text{Concat}(head_1, …, head_h)W^O $$ $$ \text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) $$

听着简单，但我第一次写出来这样：

# 错误版本
class MultiHeadAttentionWrong(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.num_heads = num_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)

    def forward(self, query, key, value):
        B, T, C = query.shape
        # 这里的拆分有问题
        q = self.Wq(query).view(B, T, self.num_heads, -1)
        k = self.Wk(key).view(B, T, self.num_heads, -1)
        v = self.Wv(value).view(B, T, self.num_heads, -1)

        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(k.size(-1))
        # ... 后续代码

问题在哪？维度没对上。view 操作默认不会 reorder 维度，导致注意力矩阵计算时 num_heads 维度混进了序列长度维度，结果算出来的注意力完全错误。

正确的做法是先 transpose：

# 正确版本
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.num_heads = num_heads
        self.d_k = d_model // num_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, query, key, value, mask=None):
        B, T, C = query.shape  # batch, seq_len, d_model

        # 投影并分割头
        q = self.Wq(query).view(B, T, self.num_heads, self.d_k).transpose(1, 2)
        k = self.Wk(key).view(B, T, self.num_heads, self.d_k).transpose(1, 2)
        v = self.Wv(value).view(B, T, self.num_heads, self.d_k).transpose(1, 2)
        # 现在是 [B, num_heads, T, d_k]

        # 计算注意力分数
        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_weights = F.softmax(scores, dim=-1)
        context = torch.matmul(attn_weights, v)  # [B, num_heads, T, d_k]

        # 合并头
        context = context.transpose(1, 2).contiguous().view(B, T, C)
        return self.Wo(context)

这个 transpose 在论文里只提了一行，但实际实现中漏掉就会产生一个很难调试的 bug：模型训练也能收敛，但效果就是差。我当时是用随机初始化的权重做单元测试才发现的。

位置编码的另一个坑

RNN 的序列信息是隐含在时间步里的，Transformer 用位置编码显式注入位置信息。论文给了两种方案：sinusoidal 和可学习的。

我一开始想当然地用了可学习的，因为那样实现简单：

# 可学习的位置编码
class PositionalEncodingLearned(nn.Module):
    def __init__(self, d_model, max_len=512):
        super().__init__()
        self.pos_embedding = nn.Embedding(max_len, d_model)

    def forward(self, x):
        B, T, C = x.shape
        positions = torch.arange(T, device=x.device).unsqueeze(0)
        return x + self.pos_embedding(positions)

这样训练是能收敛，但后来在做长文本外推时发现：超过训练序列长度后，性能急剧下降。因为可学习的位置编码本质上只是在训练区间内插值，对未见过的位置没有任何"先验"。

换回 sinusoidal 后问题缓解了不少：

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # [1, max_len, d_model]
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

但这也不是万能药。sinusoidal 在长距离上的衰减太快，我在做长文档理解时，超过 1024 长度后，注意力权重就开始"飘"了。后来改用了 RoPE（Rotary Positional Embedding），这又是另一个故事。

完整的 Encoder Block 实现

把上面的组件拼起来，一个完整的 Transformer Encoder Block 大概是这样：

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model)
        )

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Self-attention + residual + norm
        attn_out = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_out))

        # Feed-forward + residual + norm
        ffn_out = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_out))

        return x

这里有个细节：LayerNorm 放在 residual 之前还是之后？原始论文是放在之后（Pre-LN），但后来很多实现改成了之前（Post-LN）。我试过两种方式，Pre-LN 在训练稳定性上确实更好，尤其是深度网络。

实际使用中的几个发现

我在几个项目里用 Transformer 代替 LSTM，有些观察不一定有理论支持，但确实有用：

1. 批次大小比序列长度更敏感

RNN 时代，我们习惯把 batch_size 设小一点，序列长度长一点。但 Transformer 吃 GPU 带宽很厉害，batch_size 太小的话，GPU 利用率只有 30% 左右。我后来把 batch_size 从 32 改到 128，训练时间几乎没增加，但收敛速度明显提升。

2. 梯度裁剪依然必要

理论上 Attention 的梯度不会像 RNN 那样爆炸，但我在训练一个深度 Transformer 时，还是遇到了梯度爆炸。加了梯度裁剪后稳定了不少：

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

3. 学习率预热几乎是必须的

原论文用了一个学习率预热策略，一开始没在意，直接用固定的 1e-4，结果训练直接发散。后来加了预热：

def get_lr_schedule(optimizer, warmup_steps, d_model):
    def lr_lambda(step):
        if step < warmup_steps:
            return step / warmup_steps * (d_model ** -0.5)
        return (d_model ** -0.5) * (warmup_steps ** 0.5) / (step ** 0.5)
    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

这个公式看着玄乎，但确实管用。

从论文到代码的几个距离

《Attention Is All You Need》这篇论文写得不错，但从读到跑通，中间还是有些细节要补：

  • Multi-Head Attention 的维度操作：论文一行带过，但实现时 transpose/view 的顺序错了就会挂
  • Mask 的处理：padding mask 和 look-ahead mask 的合并方式，论文没细说，但两者混在一起会影响训练效果
  • 初始化策略：Xavier 还是 He？论文没提，但用了不合适的初始化会导致训练前几百步损失不降
  • 残差连接的缩放：深层 Transformer 有时会对残差输出做缩放（比如乘 0.5），这在原论文里没写

这些细节你可以在后来的实现（比如 fairseq、HuggingFace）里看到，但如果你从零开始复现，大概会踩一遍。

今天的注意力机制已经不只是 “Attention Is All You Need” 了

2017 年的注意力机制现在看来只是起点。后来出现了各种变体：

  • Linear Attention：把注意力复杂度从 $O(n^2)$ 降到 $O(n)$，适合长序列
  • Flash Attention：通过内存访问优化加速计算，实际应用很关键
  • Sparse Attention：只计算部分位置的注意力，比如 Longformer、BigBird
  • RoPE：用旋转矩阵做位置编码，现在基本成了标准

这些变种在论文里看着都挺漂亮，但真正用起来时，环境、硬件、数据分布都会影响效果。我做长文本分类时试了 Linear Attention，理论上更快，但在我的数据集上效果反而不如原始的；换成 Flash Attention 后，训练速度提升了 30%，而且效果没有损失。

这一路折腾下来的一些判断

注意力机制确实改变了 NLP 的格局，但它不是银弹。有些场景下，RNN + Attention 的组合反而更合适——比如小数据集、计算资源受限、或者需要强序列依赖的任务。

Transformer 的优势在于可扩展性和泛化能力，但这优势建立在足够的算力和数据基础上。如果你只有几百条训练数据和一个家用 GPU，未必能比调得好的 LSTM 好多少。

另一个感受是：读论文和写代码是两码事。论文里的公式写得再漂亮，实现时也要关心维度对齐、内存访问、梯度稳定这些"脏活"。真正把一个模型从论文搬到生产环境，中间要填的坑远比公式多。

现在回头看 2017 年那篇论文，它之所以能产生影响，不是因为写出了一个完美的模型，而是因为它指出了一个方向：注意力机制本身可能比序列架构更本质。后来的发展也确实证明了这一点——但那是另一篇文章的事了。

版权声明: 本文首发于 指尖魔法屋-AI注意力即一切实践笔记（https://blog.thinkmoon.cn/post/999-attention-deep-dive-from-paper-to-practice/） 转载或引用必须申明原指尖魔法屋来源及源地址！

相关阅读

根据标签、分类和主题为你挑选的相近文章