1. 从序列建模的困境说起做NLP或者时序预测的朋友大概率都经历过这样一个阶段一开始用RNN做序列到序列的任务效果还行但一旦序列变长模型就开始“失忆”——前面输入的内容到后面就忘得差不多了。更让人头疼的是RNN的训练没法并行一个时间步算完才能算下一个GPU利用率低得让人想砸键盘。后来LSTM和GRU出来了靠门控机制缓解了长距离依赖的问题但本质上还是串行计算治标不治本。Attention机制的出现在当时看来只是一个“小补丁”。2014年前后Bahdanau等人在机器翻译任务里给Encoder-Decoder结构加了一个注意力层让解码器在生成每个词的时候能够动态地“回看”编码器的所有隐藏状态而不是只依赖最后一个固定长度的上下文向量。这个改动带来的效果提升非常明显翻译质量上了一个台阶。但即便如此RNN仍然是整个架构的骨架Attention只是挂在上面的一個附属模块。真正的转折发生在2017年。那篇著名的论文提出了Transformer架构彻底抛弃了循环结构完全用Attention来建模序列中的依赖关系。这一下子把并行计算的问题解决了训练速度大幅提升而且效果还更好。从这之后BERT、GPT系列、T5等一系列模型全部建立在Transformer之上Attention从“辅助工具”变成了“核心引擎”。这篇文章我会从RNN的局限性讲起一步步拆解Attention的核心原理、Transformer的架构设计、多头注意力的计算细节以及在实际项目中落地时需要注意的那些坑。不管你是刚接触深度学习的新手还是已经用过Transformer但对其内部机制一知半解的开发者应该都能从中找到对自己有用的内容。2. RNN的局限与Attention的动机2.1 循环神经网络到底卡在哪里RNN的核心思想很直观维护一个隐藏状态每读入一个输入就更新一次这样理论上可以捕捉任意长度的历史信息。但实际用起来问题很快就暴露了。最典型的是梯度消失和梯度爆炸。反向传播的时候梯度要沿着时间步一路传回去每经过一个时间步就乘一次权重矩阵。如果权重矩阵的特征值小于1梯度会指数级衰减前面的信息就“传不回来”了如果大于1又会指数级爆炸。LSTM通过门控机制和细胞状态在一定程度上缓解了这个问题但并没有从根本上解决。另一个容易被忽视的问题是固定长度的上下文向量。在标准的Encoder-Decoder架构里编码器把整个输入序列压缩成一个固定维度的向量然后解码器从这个向量出发生成整个输出序列。你可以想象一下把一篇500字的文章压缩成一个256维的向量信息损失有多大。序列越长这个瓶颈就越严重。还有一个工程上的痛点无法并行。RNN的计算图是沿着时间步展开的第t步的计算必须等第t-1步完成。这意味着在GPU上训练的时候大量的计算单元处于空闲状态。对于动辄上千长度的序列训练时间会变得非常夸张。2.2 Attention的核心直觉Attention的思路其实很朴素既然把所有信息压缩到一个向量里会丢东西那不如让解码器在每一步都能直接看到编码器的所有输出然后根据当前的需要动态地决定“看哪里”。打个比方你在读一篇英文文章然后翻译成中文。传统的Encoder-Decoder就像是你先把整篇文章背下来然后凭记忆翻译。而Attention就像是你把原文放在旁边每翻译一个词就扫一眼原文看看当前这个词应该对应原文的哪个部分。显然后者的翻译质量会更高因为你不需要把所有细节都记在脑子里。具体到计算上Attention做的事情就是给定一个查询向量Query计算它和一组键向量Key的相似度然后用这些相似度作为权重对一组值向量Value进行加权求和。这个过程可以用一个公式概括Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V这个公式看起来简单但里面有几个关键的设计决策。除以sqrt(d_k)是为了防止点积结果过大导致softmax梯度消失。softmax的作用是把相似度分数转化成概率分布使得权重之和为1。最终得到的输出是Value的加权组合权重反映了Query和各个Key之间的相关性。2.3 从软注意力到硬注意力在实际应用中Attention有几种不同的变体。最常见的是软注意力Soft Attention也就是上面公式描述的那种对所有位置都分配一个权重然后加权求和。这种方式的优点是可微可以直接用反向传播训练。另一种是硬注意力Hard Attention它只选择一个位置而不是对所有位置加权。这种方式在某些任务上效果更好比如图像描述生成因为它更符合人类“聚焦”的认知方式。但硬注意力不可微需要用强化学习的方法来训练实现起来复杂得多实际项目中用得比较少。还有一种叫自注意力Self-Attention也就是Query、Key、Value都来自同一个序列。这种方式可以让序列中的每个位置都和其他所有位置进行交互捕捉全局依赖关系。Transformer的核心就是自注意力。3. Transformer架构的完整拆解3.1 整体结构一览Transformer的架构可以分为编码器和解码器两大部分。编码器由N个相同的层堆叠而成每层包含两个子层多头自注意力层和前馈神经网络层。解码器也是N层堆叠但每层包含三个子层掩码多头自注意力层、交叉注意力层和前馈神经网络层。每个子层都用了残差连接和层归一化具体来说是输出 LayerNorm(x Sublayer(x))这种设计让梯度能够更顺畅地传播使得训练深层网络成为可能。残差连接解决了深层网络的退化问题层归一化则稳定了训练过程。位置编码是另一个关键设计。因为Transformer完全抛弃了循环结构模型本身不知道序列中元素的顺序。所以需要额外注入位置信息。原始论文用的是正弦和余弦函数PE(pos, 2i) sin(pos / 10000^(2i/d_model)) PE(pos, 2i1) cos(pos / 10000^(2i/d_model))这种编码方式的优点是它可以外推到比训练时更长的序列而且不同位置之间的相对关系可以通过线性变换得到。3.2 多头注意力的计算细节多头注意力是Transformer的核心创新之一。它的思路是与其用一个注意力头不如用多个头每个头学习不同的关注模式。比如在翻译任务中一个头可能关注语法结构另一个头可能关注语义相关性。具体计算过程是这样的首先把输入通过线性变换映射到多个子空间每个子空间对应一个头。然后每个头独立计算注意力最后把所有头的输出拼接起来再过一次线性变换。假设模型维度是512用了8个头那么每个头的维度就是64。计算过程如下import torch import torch.nn as nn import math class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() self.d_model d_model self.num_heads num_heads self.d_k d_model // num_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, query, key, value, maskNone): batch_size query.size(0) Q self.W_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) K self.W_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) V self.W_v(value).view(batch_size, -1, self.num_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_weights torch.softmax(scores, dim-1) context torch.matmul(attn_weights, V) context context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) output self.W_o(context) return output, attn_weights这段代码有几个值得注意的地方。除以sqrt(d_k)是必须的否则当d_k较大时点积结果会很大softmax的梯度会变得极小。mask的处理也很关键在解码器的自注意力层里需要屏蔽掉未来位置的信息防止模型“偷看”答案。3.3 前馈网络与残差连接前馈网络看起来很简单就是两个线性变换中间加一个ReLU激活class PositionwiseFeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout0.1): super().__init__() self.linear1 nn.Linear(d_model, d_ff) self.linear2 nn.Linear(d_ff, d_model) self.dropout nn.Dropout(dropout) self.relu nn.ReLU() def forward(self, x): return self.linear2(self.dropout(self.relu(self.linear1(x))))但这里有一个经验性的设计中间层的维度d_ff通常是d_model的4倍。比如d_model512时d_ff2048。这个比例不是随便定的而是经过大量实验验证的。更大的中间层可以增加模型的表达能力但也会增加参数量和计算量。残差连接和层归一化的位置也有讲究。原始论文用的是Post-LN也就是先做子层计算再加残差最后做层归一化。但后来的研究发现Pre-LN先做层归一化再做子层计算最后加残差训练更稳定不需要学习率预热就能收敛。现在大部分实现都倾向于Pre-LN。4. 从零实现一个简化版Transformer4.1 位置编码的实现位置编码虽然公式看起来简单但实现的时候有一些细节需要注意。首先是编码的维度要和模型维度一致其次是不同位置的编码要能区分开。class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000, dropout0.1): super().__init__() self.dropout nn.Dropout(dropout) pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) 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) self.register_buffer(pe, pe) def forward(self, x): x x self.pe[:, :x.size(1), :] return self.dropout(x)这里用register_buffer把位置编码注册为缓冲区这样它会被保存到模型的状态字典里但不会被当作可训练参数。另外位置编码是在输入嵌入之后加的而不是拼接这样可以保持维度一致。4.2 编码器层的堆叠编码器层就是把多头注意力和前馈网络串起来中间加上残差连接和层归一化class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, num_heads) self.feed_forward PositionwiseFeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) def forward(self, x, maskNone): attn_output, _ self.self_attn(x, x, x, mask) x self.norm1(x self.dropout1(attn_output)) ff_output self.feed_forward(x) x self.norm2(x self.dropout2(ff_output)) return x这里用的是Pre-LN的变体先做注意力计算加残差再做层归一化。实际项目中我一般会先用Post-LN复现论文结果然后再尝试Pre-LN看能不能提升训练稳定性。4.3 解码器的掩码机制解码器的自注意力层需要加掩码防止当前位置看到未来的信息。这个掩码是一个上三角矩阵对角线以上的位置都是0def create_causal_mask(seq_len): mask torch.triu(torch.ones(seq_len, seq_len), diagonal1).bool() return ~mask在实际使用中还需要处理padding的掩码。padding的位置不应该参与注意力计算否则会引入噪声。通常的做法是把padding位置的注意力分数设为一个很大的负数这样softmax之后权重就接近0了。交叉注意力层的Query来自解码器Key和Value来自编码器的输出。这一层不需要因果掩码因为解码器可以看到完整的编码器输出。5. 训练技巧与常见问题排查5.1 学习率预热为什么重要Transformer训练的时候学习率预热几乎是一个必选项。原始论文用的是先线性增加再按步数的平方根倒数衰减的策略def get_lr(step, d_model, warmup_steps4000): return d_model ** (-0.5) * min(step ** (-0.5), step * warmup_steps ** (-1.5))为什么需要预热因为Transformer的参数量大初始化的时候梯度方向可能不太稳定。如果一开始就用大学习率模型很容易发散。预热让模型先在小学习率下“热身”等梯度方向稳定了再逐步增大学习率。我试过跳过预热直接训练结果loss在前几百步就炸了调小学习率虽然能训起来但收敛速度慢了很多。所以这个步骤最好不要省。5.2 梯度裁剪与混合精度训练梯度裁剪是另一个常用的技巧。Transformer的梯度有时候会突然变得很大裁剪可以防止参数更新过猛torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)混合精度训练可以显著减少显存占用和加速计算。PyTorch的amp模块用起来很方便from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): output model(input) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()但混合精度训练有时候会遇到数值不稳定的问题特别是softmax和层归一化这些操作。如果发现loss出现NaN可以先关掉混合精度排查一下。5.3 常见问题速查表问题现象可能原因排查方向Loss变成NaN学习率过大、梯度爆炸检查学习率、加梯度裁剪训练loss下降但验证loss上升过拟合增加dropout、数据增强、早停注意力权重全部均匀初始化不当、温度过高检查初始化、调整softmax温度训练速度慢序列过长、batch过大检查显存占用、优化数据加载模型不收敛学习率预热不足、层归一化位置不当调整预热步数、尝试Pre-LN这个表格是我在实际项目中踩坑之后总结的基本上覆盖了80%的常见问题。遇到问题的时候可以先对照排查能省不少时间。6. 从Transformer到现代大模型6.1 预训练与微调范式Transformer真正爆发是在预训练范式成熟之后。BERT用掩码语言建模做预训练然后在下游任务上微调刷新了多项NLP任务的记录。GPT系列则用自回归语言建模走的是生成路线。这两种范式的核心区别在于注意力掩码的方式。BERT用的是双向注意力每个位置可以看到左右两边GPT用的是因果注意力每个位置只能看到左边。双向注意力适合理解任务因果注意力适合生成任务。预训练的好处是模型可以从大量无标注文本中学习通用的语言表示然后只需要少量标注数据就能适配到具体任务。这大大降低了NLP应用的门槛。6.2 效率优化从稀疏注意力到线性注意力标准自注意力的计算复杂度是O(n^2)n是序列长度。当序列长度达到几千甚至几万的时候计算量和显存占用都会变得不可接受。所以后来出现了很多效率优化的变体。稀疏注意力只计算部分位置的注意力比如只关注局部窗口或者固定的稀疏模式。Longformer和BigBird是这方面的代表。线性注意力则通过核函数近似把复杂度降到O(n)。Performer用的是随机特征映射Linformer则对Key和Value做低秩近似。这些方法各有取舍。稀疏注意力实现简单但可能丢失长距离依赖线性注意力理论上有优势但实际效果有时候不如标准注意力。选择哪种方案要看具体任务对精度和效率的要求。6.3 实际项目中的选型建议如果你现在要做一个NLP项目我的建议是优先用预训练模型微调不要从零训练Transformer。HuggingFace的transformers库提供了大量预训练模型加载和微调都很方便。选模型的时候要考虑任务类型、数据量、计算资源。文本分类任务用BERT-base通常就够了生成任务可以用GPT-2或者T5。如果数据量很大可以考虑更大的模型如果计算资源有限用蒸馏后的小模型或者量化后的模型。序列长度也是一个关键因素。如果序列长度在512以内标准Transformer完全够用如果超过1024就要考虑效率优化方案了。我一般会先用标准注意力跑一个baseline然后再尝试优化方案看能带来多少提升。7. 我踩过的那些坑第一个坑是位置编码的实现。我一开始直接把位置编码加到了词嵌入上但忘了做dropout结果训练的时候过拟合很严重。后来加了dropout才正常。另外位置编码的max_len要设得比实际序列长度大一些否则长序列会报错。第二个坑是掩码的处理。解码器的因果掩码和padding掩码要分开处理我一开始混在一起导致模型在训练的时候能看到padding位置的信息验证集效果很差。正确的做法是先创建因果掩码再把padding位置屏蔽掉。第三个坑是学习率预热。我一开始觉得预热没必要直接用了固定学习率结果loss在前几百步就发散了。后来老老实实加了预热训练就稳定了。预热的步数一般设成总训练步数的10%左右比较合适。第四个坑是混合精度训练。开了amp之后训练速度确实快了不少但有时候会遇到梯度下溢的问题。后来加了GradScaler才解决。如果遇到NaN可以先关掉amp排查确认没问题再开。这些坑看起来都是小问题但每一个都让我多花了好几个小时甚至一两天的时间。希望看到这篇文章的朋友能少走一些弯路。Attention机制和Transformer架构本身并不复杂真正花时间的是这些工程细节。把细节处理好了模型的效果自然就上来了。