1. 先从一个大问题说起为什么需要Transformer早年做序列建模大家绕不开的就是RNN、LSTM这套循环神经网络。它们的核心思路是“按时间步一步步处理”当前时刻的隐状态必须等前一个时刻算完才能继续。这种串行特性带来两个很现实的问题一是训练慢长序列等于几百上千步的串联很难并行化二是长期依赖难捕捉即使加了LSTM的门控机制序列一长早期的信息还是会慢慢衰减说白了就是记不住太久以前的事。做机器翻译、文本生成时这种“记性差”特别明显句子一长译文质量就肉眼可见地往下掉。2017年Google那篇《Attention Is All You Need》直接掀了桌子不搞循环不要卷积全靠注意力机制并行度拉满长距离依赖也能一步到位。这就是Transformer的起点。当初我刚看到这个思路的时候第一反应是“这也太激进了”完全抛弃序列顺序的结构光靠注意力匹配全局信息居然能比RNN效果好这么多。但实测下来Transformer在翻译任务上不仅精度更高训练速度还快了一个数量级从那以后基本就成了NLP领域的事实标准。后来BERT、GPT系列一路把它推成AI底座再到ViT把这套结构搬到图像领域整个深度学习的方向都被它改写了一遍。这篇内容不打算只念论文我会把Transformer的原理、架构和注意力机制拆开揉碎讲清楚同时给出一份可直接运行的PyTorch实现从数据构造、模型搭建到训练推理完整走一遍。不管是刚开始接触注意力机制的新手还是想手动复现Transformer做二次开发的工程师都能直接从里面拿到可用的东西。2. 注意力机制Transformer的心脏2.1 自注意力到底在算什么注意力机制的概念其实很好理解一句话让每个位置学会关注序列中其他位置的信息并按重要程度加权聚合。放到翻译场景里就是翻译当前词时模型会自动把注意力分配给它认为相关的源语言词上盯住关键的忽略无关的。那它具体是怎么算的核心是Query、Key、Value这三个向量平时教程里常说的Q、K、V就是它们。可以用一个生活化的场景来类比你脑子里想着一件想买的东西比如“白色T恤”这就是Query超市货架上的每件商品都贴着自己的标签比如“白色T恤”、“蓝色牛仔裤”这些标签就是Key而每件商品本身就是Value。你会先把“白色T恤”这个查询和货架上所有商品的Key做匹配找到标签最匹配的商品然后重点取走那件商品其他商品略过不看。Transformer里的自注意力也是同一套逻辑只不过“脑海里的需求”和“货架上的商品标签”都是从输入序列本身生成的。具体就是输入的每个token都会映射成三个向量然后做三步运算用每个token的Query去和序列里所有token的Key做点积得到一个相似度分数这个分数表示“当前token应该给另一个token多少注意力”把这些分数除以一个缩放因子再做Softmax归一化变成一组加起来等于1的权重拿这组权重去对所有Value向量做加权求和得到当前token的输出向量。写成公式是这样Attention(Q, K, V) softmax(QK^T / √d_k) V其中d_k是Key向量的维度除以√d_k是为了防止点积结果太大把Softmax推入饱和区导致梯度消失。这个细节看起来小实际训练时影响很大不缩放的话模型很容易不稳定。2.2 从“每个词都看所有词”到多头注意力单靠一组Q、K、V其实也能工作但表达能力有限因为所有位置共用一种“注意力模式”。现实里的关系是复杂的有的词要靠得很近才能理解有的词跨了大半个句子才有语义关联有的注意力应该指向语法搭配有的则要指向指代关系。一组注意力头根本顾不过来这么多需求。多头注意力机制的解法是把Q、K、V分别投影到h个不同的子空间里每个头独立做一次注意力计算再把所有头的结果拼起来做一次线性变换。这样每个头就能学到不同的注意力模式有的偏向局部邻近词有的偏向长距离依赖。Google论文里也验证过不同头关注到的语义关系确实差异很大有的头稳定关注相邻词有的头用来处理指代消解。多头注意力流程分四步将输入X分别线性投影成多组Q、K、V假设有h个头每个头有独立的权重矩阵每个头单独计算缩放点积注意力将h个头输出的向量拼接起来经过一个输出投影矩阵得到最终结果。头数h是超参数经典论文里用的是8或者16。我自己的经验是头数太多小模型反而容易“各想各的”信息交融不起来头数太少注意力模式又不够丰富。在中小规模模型上8个头一般是个比较稳的选择。另外提一下自注意力这个名字的意思是Q、K、V都来自同一个输入序列。如果来自不同序列比如Decoder去关注Encoder的输出那就叫交叉注意力。这也是Transformer架构里非常重要的一环后面讲架构时会具体展开。3. Transformer的整体架构拆解3.1 Encoder与Decoder分工明确的编解码结构Transformer原始模型采用的是Encoder-Decoder结构整体分两大块。Encoder负责把输入序列编码成语义丰富的中间表示Decoder负责根据这个表示和已经生成的内容逐步生成目标序列。拿机器翻译举例Encoder读入英文句子输出一组上下文向量Decoder参考这组向量和之前已经翻译出的词逐个预测下一个目标语言的词。Encoder和Decoder各自都由多层堆叠而成每层内部的结构略有不同Encoder层包含两个子层多头自注意力层和逐位置前馈网络层Position-wise Feed-Forward Network每个子层外面都用残差连接和层归一化包一圈Decoder层包含三个子层掩码多头自注意力层、交叉注意力层、前馈网络层同样每层都有残差和归一化。其中Decoder的自注意力层加了掩码Mask这是它和Encoder自注意力最本质的区别。训练时Decoder接收目标序列作为输入如果在某个位置能看到未来位置的词那模型就是在“抄答案”测试时会崩得没法看。掩码的做法很粗暴把未来位置的注意力分数直接设成负无穷这样经过Softmax后这些位置的权重变成0当前token就只看得到自己之前的内容。这就是常说的因果掩码。交叉注意力层则是Decoder与Encoder之间的桥梁。这里的Query来自Decoder前一层的输出Key和Value都来自Encoder的输出。Decoder每生成一个词都要去Encoder编码出的整个源句子里寻找相关信息这个过程就是“对齐”也是Transformer做翻译和生成时语义准确的关键。3.2 位置编码没有顺序信息的序列就是一堆词袋Transformer结构里有个明显的漏洞自注意力计算对位置完全不敏感。如果你把句子里的词序打乱只要每个词的向量不变所有注意力分数和输出结果就完全一样。这显然不行因为“张三打了李四”和“李四打了张三”完全是两码事。最初的论文用的是正弦位置编码公式是PE(pos, 2i) sin(pos / 10000^(2i/d_model))PE(pos, 2i1) cos(pos / 10000^(2i/d_model))这里pos是词在序列中的位置i是位置向量里的维度下标d_model是模型隐藏层的维度。这种编码的特点是不同维度拥有不同的频率模型可以从位置向量中轻易地通过线性变换推导出相对位置信息。另一个好处是它不需要训练任何长度的序列都能算出来没有长度上限的问题。后来的实践中可学习位置编码让模型自己去学一组位置向量也很常见效果和正弦编码差距不大尤其BERT用的是可学习方式。再后来像RoPE这类带旋转的位置编码在LLM里成为主流因为它把位置信息直接融进注意力计算里外推能力更好这也是为什么很多大模型能在训练长度之外继续生成较长内容的秘密之一。新手入门阶段先用正弦位置编码或者可学习位置编码就好不必一上来就折腾RoPE。3.3 残差连接、层归一化与前馈网络的作用Transformer每层结构都少不了三件套残差连接、层归一化、前馈网络。它们各有分工缺一不可。残差连接解决的核心问题是深度网络的退化问题。一般来说网络越深表达能力越强但实际训练时发现层数堆到几十层以后梯度很容易消失模型不升反降。残差连接让每层的输出变成F(x) x即使这一层学的变换没有价值信息也能通过直连通道无损地传到后面梯度也能顺畅地回流。Transformer原始结构里每个子层输出都要加上输入再进下一层就是这个原因。层归一化则是对每个样本的所有特征维度做归一化把数据拉回均值为0、方差为1的分布。它和BatchNorm不一样的是不依赖batch内其他样本的统计量在序列长度变化比较灵活的任务里更稳定。之前做训练的时候我对比过用LayerNorm收敛明显比不用稳学习率可以开得更大也不容易崩。前馈网络在Transformer里被设计成一个两层的全连接结构先升维再降维。原论文里FFN的中间维度设为2048是d_model512的4倍。这个设计思路是让每个位置的向量先映射到更高维空间做非线性变换再压回原来的维度相当于给每个token的表示做了一次特征交叉和抽象。位置上说是“逐位置”意思是同一个FFN对序列里的每个token独立使用权重共享。3.4 Transformer的宏观数据流把以上所有模块拼在一起数据在Transformer内部的完整流程是这样的输入序列先做tokenization每个token通过词嵌入矩阵映射成d_model维的向量加上位置编码模型第一次获得“每个词在哪里”的信息向量进入Encoder经过N层堆叠经典配置是6层。每层内部自注意力聚合全局信息残差和归一化稳定分布FFN逐位置做非线性变换Encoder最终输出一组上下文向量代表整个输入序列的语义编码Decoder接收目标序列的前面部分先经过掩码自注意力限制自己只能看到已生成内容再通过交叉注意力去关注Encoder输出的信息最后经过FFNDecoder每一步输出一个预测概率分布对应词表上每个词的概率取概率最高的词作为当前步的生成结果送入下一步解码。训练阶段可以一次性把整个目标序列喂进Decoder靠掩码保证因果性并行计算效率高。推理阶段只能自回归地一个一个词生成每步把新生成词拼到已有序列里重新输入。4. 从零手写一个Transformer核心组件4.1 用PyTorch实现缩放点积注意力与多头注意力理论容易说真写代码时才会碰到各种细节问题。下面给出一套可以直接运行的PyTorch实现目标是构造一个能完成copy任务的Mini Transformer。copy任务是让模型学会把输入序列原样输出虽然简单但能完整检验模型前向计算、训练和推理整个链路是否正常。先看缩放点积注意力的实现import torch import torch.nn as nn import torch.nn.functional as F import math class ScaledDotProductAttention(nn.Module): def __init__(self, d_k, dropout0.1): super().__init__() self.d_k d_k self.dropout nn.Dropout(dropout) def forward(self, q, k, v, maskNone): # q: [batch, heads, seq_len, d_k] # k: [batch, heads, seq_len, d_k] # v: [batch, heads, seq_len, 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, float(-inf)) attn F.softmax(scores, dim-1) attn self.dropout(attn) output torch.matmul(attn, v) return output, attn这段代码有几个关键细节值得注意scores除以√d_k是论文规定的缩放操作防止点积值太大导致Softmax落到梯度平缓区mask处理用的是masked_fill把需要屏蔽的位置填成负无穷。为什么不是填0因为Softmax操作是非线性的填0不会让那个位置完全失效而负无穷经过Softmax后权重会变成0dropout放在softmax之后加在注意力权重上而不是加在scores上。这个细节是很多实现容易搞错的。作用是在训练时随机衰减一部分注意力权重起到正则化效果防止模型过度依赖某几个位置。多头注意力是在单头基础上做了投影切分和拼接恢复。这里我把所有头的投影用一个矩阵完成再通过view和transpose拆成多头实现上更简洁效率也更高class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout0.1): super().__init__() assert d_model % n_heads 0 self.n_heads n_heads self.d_k d_model // n_heads self.d_model d_model 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) self.attention ScaledDotProductAttention(self.d_k, dropout) def split_heads(self, x): batch_size, seq_len, _ x.size() x x.view(batch_size, seq_len, self.n_heads, self.d_k) return x.transpose(1, 2) # [batch, heads, seq_len, d_k] def forward(self, q, k, v, maskNone): q self.split_heads(self.w_q(q)) k self.split_heads(self.w_k(k)) v self.split_heads(self.w_v(v)) output, attn self.attention(q, k, v, mask) output output.transpose(1, 2).contiguous() output output.view(output.size(0), -1, self.d_model) return self.w_o(output), attn关于维度的设计这里沿用原始Transformer的设定d_model为512时8个head每个head维度64。d_model必须能被n_heads整除不然reshape时会报错。这里的d_k d_model // n_heads也就是每个头的维度。这个设计的妙处在于总参数量没有变所有头的输出拼起来之后恰好还原成d_model维不影响后续网络结构。4.2 前馈网络、位置编码、层归一化与完整Encoder层逐位置前馈网络就是两个线性变换加一个ReLU激活中间维度放大4倍再加回去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) def forward(self, x): return self.linear2(self.dropout(F.relu(self.linear1(x))))d_ff取的是2048。这个放大倍数不是随便定的太小时FFN特征抽象能力有限太大时中间层的参数量会占据整个模型很大比例训练成本明显上升实际效果提升却越来越有限。做中小模型时可以先按4倍设置再根据训练效果微调。位置编码用正弦公式实现。需要注意把位置编码注册成buffer而不是parameter这样它在保存模型时不会算进可训练参数推理时如果遇到更长的序列还能在代码里动态延长class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000): super().__init__() 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) # [1, max_len, d_model] self.register_buffer(pe, pe) def forward(self, x): return x self.pe[:, : x.size(1)]位置编码实现时有一行容易写错div_term用指数方式构造效果等价于对每个偶数维度i计算10000^(-2i/d_model)但数值上更稳定不会因为指数太小产生精度问题。偶数维度用sin奇数维度用cos这样可以保证每个位置得到一个唯一的编码向量。Encoder的每一层可以封装成这样class EncoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, n_heads, dropout) self.ffn 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_out, _ self.self_attn(x, x, x, mask) x self.norm1(x self.dropout1(attn_out)) ffn_out self.ffn(x) x self.norm2(x self.dropout2(ffn_out)) return x这里采用的是Post-LN结构也就是先残差再加LayerNorm这是原始Transformer论文里的方式。后来GPT和很多新模型改用Pre-LN把LayerNorm放在子层之前训练时更稳定但实现上两类都要能写碰到别人代码时才知道为什么结构略有不同。Post-LN对学习率和warmup更敏感需要精心调节Pre-LN的收敛更稳对新手更友好。5. 模型训练与预测的完整闭环5.1 准备数据构建batch并生成掩码构造一个Copy Task需要的训练集非常容易随机生成整数序列输入和目标都是同一个序列。实际代码里还要做padding因为一个batch里的序列长度不一定相同较短的序列要用占位符补齐。Padding的位置在注意力计算时一定要屏蔽掉否则模型会去关注空白位置白白浪费计算还学到噪音。掩码需要分场景设计。Decoder做训练时输入的每个token位置能看到的位置范围受两个条件限制一是不能看未来的token这由因果掩码控制二是不能看padding位置这由padding掩码控制。两个掩码取交集实现上就是把两个布尔矩阵做逻辑与运算。def generate_padding_mask(seq, pad_idx0): return (seq ! pad_idx).unsqueeze(1).unsqueeze(2) # [batch, 1, 1, seq_len] def generate_causal_mask(size): attn_shape (1, size, size) mask torch.triu(torch.ones(attn_shape), diagonal1).bool() # 上三角为True return ~mask # [1, size, size], 下三角含对角线为True因果掩码这里用了triu函数生成一个上三角为True的矩阵然后取反。这样位置i的注意力只能落在位置0到i之间后面的位置全部被masked_fill成负无穷经过Softmax权重为0。写这段代码的时候建议停下来多想想矩阵维度的逻辑后面Debug时很多问题都出在掩码形状对不上。5.2 构建Mini Transformer模型我把Encoder和Decoder整合成一个完整的Transformer类class Transformer(nn.Module): def __init__(self, vocab_size, d_model128, n_heads8, d_ff512, num_layers3, max_len100, dropout0.1, pad_idx0): super().__init__() self.d_model d_model self.pad_idx pad_idx self.embedding nn.Embedding(vocab_size, d_model) self.pos_encoding PositionalEncoding(d_model, max_len) self.encoder_layers nn.ModuleList([ EncoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(num_layers) ]) self.decoder_layers nn.ModuleList([ DecoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(num_layers) ]) self.fc_out nn.Linear(d_model, vocab_size) self.dropout nn.Dropout(dropout) def forward(self, src, tgt): # src: [batch, src_len], tgt: [batch, tgt_len] src_mask generate_padding_mask(src, self.pad_idx) tgt_pad_mask generate_padding_mask(tgt, self.pad_idx) tgt_len tgt.size(1) tgt_causal_mask generate_causal_mask(tgt_len).to(tgt.device) tgt_mask tgt_pad_mask tgt_causal_mask src_emb self.dropout(self.pos_encoding(self.embedding(src) * math.sqrt(self.d_model))) tgt_emb self.dropout(self.pos_encoding(self.embedding(tgt) * math.sqrt(self.d_model))) for layer in self.encoder_layers: src_emb layer(src_emb, src_mask) for layer in self.decoder_layers: tgt_emb layer(tgt_emb, src_emb, src_mask, tgt_mask) logits self.fc_out(tgt_emb) return logitsDecoderLayer的forward签名里有三个输入分别是decoder自身的输入、encoder输出和两个掩码交叉注意力部分的query来自decoderkey和value来自encoderclass DecoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, n_heads, dropout) self.cross_attn MultiHeadAttention(d_model, n_heads, dropout) self.ffn PositionwiseFeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.norm3 nn.LayerNorm(d_model) self.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) self.dropout3 nn.Dropout(dropout) def forward(self, x, enc_output, src_maskNone, tgt_maskNone): attn_out, _ self.self_attn(x, x, x, tgt_mask) x self.norm1(x self.dropout1(attn_out)) cross_out, _ self.cross_attn(x, enc_output, enc_output, src_mask) x self.norm2(x self.dropout2(cross_out)) ffn_out self.ffn(x) x self.norm3(x self.dropout3(ffn_out)) return x代码里有个小细节嵌入层输出乘以√d_model。这是原论文里的做法因为位置编码的数值范围通常在[-1,1]之间如果不放大嵌入值位置信息会在加法中被词嵌入淹没。乘上√d_model之后两者在数值贡献上大致平衡。5.3 训练循环与推理生成训练阶段直接用交叉熵损失padding位置需要计算loss时屏蔽。一种做法是设置ignore_indexCrossEntropyLoss会跳过填充位置的预测def train_model(model, train_loader, epochs20, lr3e-4): criterion nn.CrossEntropyLoss(ignore_index0) # 忽略padding位置 optimizer torch.optim.Adam(model.parameters(), lrlr, betas(0.9, 0.98), eps1e-9) scheduler torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambdalambda step: min((step 1) ** (-0.5), (step 1) * 0.0001 ** (-0.5)) ) model.train() for epoch in range(epochs): total_loss 0 for batch in train_loader: src, tgt_in, tgt_out batch optimizer.zero_grad() logits model(src, tgt_in) # [batch, tgt_len, vocab_size] loss criterion(logits.reshape(-1, logits.size(-1)), tgt_out.reshape(-1)) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() total_loss loss.item() print(fEpoch {epoch1}, Loss: {total_loss / len(train_loader):.4f})代码里已经把tgt_in和tgt_out设计成错位关系训练时Decoder输入目标序列的[0, T-1]部分预测目标是[1, T]部分。这样每个位置的模型都只看到它之前的内容和推理时的行为一致。注意生成tgt_in和tgt_out时要做错位处理如果目标序列是[2, 5, 1, 3]那tgt_in是[2, 5, 1]tgt_out是[5, 1, 3]开头补一个起始符。学习率要加warmup策略这一点非常重要。Transformer对学习率比较敏感直接用固定学习率的Adam会发现训练前期loss很不稳定甚至直接NaN。上面代码里的LambdaLR实现的是论文里的Noam scheduling学习率先线性升到峰值再按步数的平方根倒数衰减。没有这个调度模型很容易在早期就发散。推理阶段的生成用贪心解码就能完成任务演示def greedy_decode(model, src, max_len20, start_token2, end_token3, devicecpu): model.eval() src src.to(device) tgt torch.tensor([[start_token]], devicedevice) with torch.no_grad(): for _ in range(max_len): logits model(src, tgt) next_token logits[:, -1, :].argmax(dim-1).unsqueeze(1) tgt torch.cat([tgt, next_token], dim1) if next_token.item() end_token: break return tgt.squeeze(0).tolist()用贪心解码跑一批测试序列如果模型训练正常几个epoch之后就能观察到模型输出的序列和输入完全一致。这个效果看着简单但能跑通说明整个链路——数据准备、掩码生成、前向计算、损失计算、反向传播、学习率调度——全部没有问题。之后往Translate或者文本生成任务迁移时只需换掉数据集和词表即可。6. 常见问题与排查经验6.1 训练不收敛与Loss为NaN的排查思路训练Transformer时遇到的最典型问题按出现频率排大概是下面这几种问题现象常见原因处理办法Loss不下降学习率太低或太高检查warmup调度调整峰值学习率Loss输出NaN学习率过大、梯度爆炸、attention计算溢出降低学习率加梯度裁剪检查是否除以√d_k训练时好时坏、震荡Adam的epsilon太小或batch size过小设置eps1e-9适当增大batch size模型生成循环重复内容Decoder训练和推理不一致、没有mask检查Decoder的因果掩码是否正确GPU显存不足序列过长或batch太大减少batch size用梯度累积限制max_lenLoss直接输出NaN是我见过最多的情况。排查看三点第一学习率是否过大第二注意力分数是否因为忘了除以√d_k而太大导致Softmax溢出第三Adam的epsilon是否设置得当PyTorch默认eps是1e-8用在Transformer上有时不够稳定论文里用的是1e-9。还有一个容易被忽视的点是梯度裁剪。Transformer深度较大加上自注意力存在长程路径梯度范数很容易超过一个合理的范围。torch.nn.utils.clip_grad_norm_设置max_norm1.0是一个常见且保守的选择基本不会拖慢收敛但能显著降低NaN风险。6.2 推理阶段的坑自回归生成与训练行为的偏差自回归生成时的性能和训练时有明显差异尤其是Decoder部分。训练时我们一次性喂入整个目标序列通过掩码保证因果性所有位置可以并行计算。但推理时只能一个词一个词来每生成一个新词就要把整个序列重新输入模型完整前向一次。这样做有两个问题。一是慢生成100个词就要做100次前向如果模型层数多、序列长推理时延会肉眼可见地涨。优化方案是使用KV Cache把计算过的Key和Value缓存下来避免每步重新计算前面所有token的注意力。原始Transformer代码里没有这个机制现在的GPT系列推理引擎都做了深度优化。二是误差累积。训练时Decoder输入的是真实目标序列但推理时输入的是模型自己生成的前一步输出。一步生成错了后面的所有预测都会受到污染。这就是为什么训练和推理行为要保持一致——如果训练时Decoder能“偷看”到未来位置那模型会形成依赖一旦推理阶段无法偷看就会全面崩溃。所以那个因果掩码绝对不能写错我见过不少复现代码出问题最后定位都是掩码形状或者布尔取反搞反了。6.3 模型参数量与内存的粗略估算写Transformer代码时估算一下模型大小很实用。Transformer参数量主要由四部分构成词嵌入矩阵、注意力层的QKV和输出投影、前馈网络的两个线性层、LayerNorm和位置编码位置编码通常是buffer不占可训练参数。以d_model128、词表大小50、8头、6层EncoderDecoder的小配置估算词嵌入128 * 50 6400每层多头注意力QKV三个矩阵各128128输出投影128128总共65536每层FFN128512 512128 131072每层总参数量约1966086层双层约2.36M再加上其他杂项整体不到3M参数这个规模用CPU跑Copy Task只需要几分钟。如果换成d_model512、词表3万的经典配置词嵌入一项就有1536万参数整体接近6500万。这就是为什么词表大小对模型体积影响巨大——很多模型压缩工作从共享嵌入矩阵入手就是这个原因。7. 几个值得关注的Transformer变体7.1 视觉方向ViT与Swin TransformerViTVision Transformer把Transformer搬进图像领域的方式很直接把图片切成一堆16x16的小patch每个patch拉平后当作一个“词”加上位置编码后送进标准Transformer。这个思路在数据量足够大时效果非常惊艳ImageNet分类上超过了同量级的CNN。但它有个缺点全局注意力在高分辨率图像上计算开销太大patch数量一多内存就撑不住。Swin Transformer的改进思路是引入层次化和窗口化。它把注意力限制在局部窗口内计算窗口之间的信息通过“移位窗口”机制来交换。这样做的好处是计算复杂度从O(N^2)降到O(N)而且能像CNN一样构建多尺度特征金字塔方便做检测和分割这类密集预测任务。7.2 轻量注意力模块SE、CBAM与CA这几个名字频繁出现在图像分类比赛和移动端模型优化里严格来说不是Transformer本体而是在CNN基础上加进去的注意力模块。SE模块Squeeze-and-Excitation做的是通道注意力把每个通道的二维特征图压缩成一个全局描述符再通过两个全连接层学习每个通道的权重最后把权重乘回原始特征图。它关注的是“哪些通道重要”。CBAM在SE的基础上加了空间注意力分支。通道注意力告诉模型“看什么”空间注意力告诉模型“看哪里”两条分支串行组合效果进一步提升参数量增加却很少。CA模块Coordinate Attention的思路更细它把空间注意力分解成水平和垂直两个方向分别编码长程空间依赖对目标定位和细粒度识别更友好。这些模块的共性是轻量、即插即用做图像模型时可以在残差块里直接插入。但要注意的是它们的原理和Transformer里的自注意力有本质区别SE、CBAM、CA更接近“对已有特征做通道或空间上的重标定”而Transformer自注意力做的是“所有位置的信息互相交换”两者解决的问题不在一个层面。7.3 MoE架构与多模态方向MoEMixture of Experts的思路是让模型在推理时只激活一小部分参数而不是每一层都动用全部前馈网络。具体做法是设置多个FFN专家由一个门控网络为每个token选出最合适的几个专家。GPT-4、DeepSeek等大模型走的就是这个路线在总参数不变的前提下大幅降低推理成本让模型可以在更大规模上扩展。多模态Transformer则是把文本、图像、音频等不同模态的数据统一成token序列之后一起送入模型让注意力机制在不同模态之间建立关联。近年来这类模型的进展非常快从图像理解到文生图、视频生成底层几乎都是Transformer架构或其变体。这些方向其实都说明一件事Transformer这套“让所有元素之间互相联系”的框架本身已经超越了具体任务的限制成为一种通用信息处理范式。深度理解基础架构后往上做任何方向的扩展都会轻松很多。8. 写在最后的实操心得手写过一遍Transformer之后最大的体会是看论文和真正实现完全是两码事。论文里的公式只有三行但落在代码上光是一个掩码就要处理好几种形状和叠加逻辑一个维度对不上反向传播阶段就会炸出莫名其妙的错误。建议所有想深入理解这套架构的人都亲手把这个模型从零到一实现一遍——不需要复现多大规模能跑通Copy Task、翻译一个简单句子就足够建立完整的认知框架。再分享一个影响很大的细节训练策略和模型结构同样重要。同一个Transformer代码用固定学习率训练可能怎么都收敛不了加上warmup和梯度裁剪之后几个epoch就开始稳定下降了。第一次跑通时那种感觉是——原来这个模型对训练技巧的要求这么高也难怪原论文里要花那么大的篇幅写训练细节。新手阶段一定要把这些基础训练策略一起学会它们在未来做更大模型时同样不可或缺。 Transformer原理与PyTorch实现：从注意力机制到代码实战
1. 先从一个大问题说起为什么需要Transformer早年做序列建模大家绕不开的就是RNN、LSTM这套循环神经网络。它们的核心思路是“按时间步一步步处理”当前时刻的隐状态必须等前一个时刻算完才能继续。这种串行特性带来两个很现实的问题一是训练慢长序列等于几百上千步的串联很难并行化二是长期依赖难捕捉即使加了LSTM的门控机制序列一长早期的信息还是会慢慢衰减说白了就是记不住太久以前的事。做机器翻译、文本生成时这种“记性差”特别明显句子一长译文质量就肉眼可见地往下掉。2017年Google那篇《Attention Is All You Need》直接掀了桌子不搞循环不要卷积全靠注意力机制并行度拉满长距离依赖也能一步到位。这就是Transformer的起点。当初我刚看到这个思路的时候第一反应是“这也太激进了”完全抛弃序列顺序的结构光靠注意力匹配全局信息居然能比RNN效果好这么多。但实测下来Transformer在翻译任务上不仅精度更高训练速度还快了一个数量级从那以后基本就成了NLP领域的事实标准。后来BERT、GPT系列一路把它推成AI底座再到ViT把这套结构搬到图像领域整个深度学习的方向都被它改写了一遍。这篇内容不打算只念论文我会把Transformer的原理、架构和注意力机制拆开揉碎讲清楚同时给出一份可直接运行的PyTorch实现从数据构造、模型搭建到训练推理完整走一遍。不管是刚开始接触注意力机制的新手还是想手动复现Transformer做二次开发的工程师都能直接从里面拿到可用的东西。2. 注意力机制Transformer的心脏2.1 自注意力到底在算什么注意力机制的概念其实很好理解一句话让每个位置学会关注序列中其他位置的信息并按重要程度加权聚合。放到翻译场景里就是翻译当前词时模型会自动把注意力分配给它认为相关的源语言词上盯住关键的忽略无关的。那它具体是怎么算的核心是Query、Key、Value这三个向量平时教程里常说的Q、K、V就是它们。可以用一个生活化的场景来类比你脑子里想着一件想买的东西比如“白色T恤”这就是Query超市货架上的每件商品都贴着自己的标签比如“白色T恤”、“蓝色牛仔裤”这些标签就是Key而每件商品本身就是Value。你会先把“白色T恤”这个查询和货架上所有商品的Key做匹配找到标签最匹配的商品然后重点取走那件商品其他商品略过不看。Transformer里的自注意力也是同一套逻辑只不过“脑海里的需求”和“货架上的商品标签”都是从输入序列本身生成的。具体就是输入的每个token都会映射成三个向量然后做三步运算用每个token的Query去和序列里所有token的Key做点积得到一个相似度分数这个分数表示“当前token应该给另一个token多少注意力”把这些分数除以一个缩放因子再做Softmax归一化变成一组加起来等于1的权重拿这组权重去对所有Value向量做加权求和得到当前token的输出向量。写成公式是这样Attention(Q, K, V) softmax(QK^T / √d_k) V其中d_k是Key向量的维度除以√d_k是为了防止点积结果太大把Softmax推入饱和区导致梯度消失。这个细节看起来小实际训练时影响很大不缩放的话模型很容易不稳定。2.2 从“每个词都看所有词”到多头注意力单靠一组Q、K、V其实也能工作但表达能力有限因为所有位置共用一种“注意力模式”。现实里的关系是复杂的有的词要靠得很近才能理解有的词跨了大半个句子才有语义关联有的注意力应该指向语法搭配有的则要指向指代关系。一组注意力头根本顾不过来这么多需求。多头注意力机制的解法是把Q、K、V分别投影到h个不同的子空间里每个头独立做一次注意力计算再把所有头的结果拼起来做一次线性变换。这样每个头就能学到不同的注意力模式有的偏向局部邻近词有的偏向长距离依赖。Google论文里也验证过不同头关注到的语义关系确实差异很大有的头稳定关注相邻词有的头用来处理指代消解。多头注意力流程分四步将输入X分别线性投影成多组Q、K、V假设有h个头每个头有独立的权重矩阵每个头单独计算缩放点积注意力将h个头输出的向量拼接起来经过一个输出投影矩阵得到最终结果。头数h是超参数经典论文里用的是8或者16。我自己的经验是头数太多小模型反而容易“各想各的”信息交融不起来头数太少注意力模式又不够丰富。在中小规模模型上8个头一般是个比较稳的选择。另外提一下自注意力这个名字的意思是Q、K、V都来自同一个输入序列。如果来自不同序列比如Decoder去关注Encoder的输出那就叫交叉注意力。这也是Transformer架构里非常重要的一环后面讲架构时会具体展开。3. Transformer的整体架构拆解3.1 Encoder与Decoder分工明确的编解码结构Transformer原始模型采用的是Encoder-Decoder结构整体分两大块。Encoder负责把输入序列编码成语义丰富的中间表示Decoder负责根据这个表示和已经生成的内容逐步生成目标序列。拿机器翻译举例Encoder读入英文句子输出一组上下文向量Decoder参考这组向量和之前已经翻译出的词逐个预测下一个目标语言的词。Encoder和Decoder各自都由多层堆叠而成每层内部的结构略有不同Encoder层包含两个子层多头自注意力层和逐位置前馈网络层Position-wise Feed-Forward Network每个子层外面都用残差连接和层归一化包一圈Decoder层包含三个子层掩码多头自注意力层、交叉注意力层、前馈网络层同样每层都有残差和归一化。其中Decoder的自注意力层加了掩码Mask这是它和Encoder自注意力最本质的区别。训练时Decoder接收目标序列作为输入如果在某个位置能看到未来位置的词那模型就是在“抄答案”测试时会崩得没法看。掩码的做法很粗暴把未来位置的注意力分数直接设成负无穷这样经过Softmax后这些位置的权重变成0当前token就只看得到自己之前的内容。这就是常说的因果掩码。交叉注意力层则是Decoder与Encoder之间的桥梁。这里的Query来自Decoder前一层的输出Key和Value都来自Encoder的输出。Decoder每生成一个词都要去Encoder编码出的整个源句子里寻找相关信息这个过程就是“对齐”也是Transformer做翻译和生成时语义准确的关键。3.2 位置编码没有顺序信息的序列就是一堆词袋Transformer结构里有个明显的漏洞自注意力计算对位置完全不敏感。如果你把句子里的词序打乱只要每个词的向量不变所有注意力分数和输出结果就完全一样。这显然不行因为“张三打了李四”和“李四打了张三”完全是两码事。最初的论文用的是正弦位置编码公式是PE(pos, 2i) sin(pos / 10000^(2i/d_model))PE(pos, 2i1) cos(pos / 10000^(2i/d_model))这里pos是词在序列中的位置i是位置向量里的维度下标d_model是模型隐藏层的维度。这种编码的特点是不同维度拥有不同的频率模型可以从位置向量中轻易地通过线性变换推导出相对位置信息。另一个好处是它不需要训练任何长度的序列都能算出来没有长度上限的问题。后来的实践中可学习位置编码让模型自己去学一组位置向量也很常见效果和正弦编码差距不大尤其BERT用的是可学习方式。再后来像RoPE这类带旋转的位置编码在LLM里成为主流因为它把位置信息直接融进注意力计算里外推能力更好这也是为什么很多大模型能在训练长度之外继续生成较长内容的秘密之一。新手入门阶段先用正弦位置编码或者可学习位置编码就好不必一上来就折腾RoPE。3.3 残差连接、层归一化与前馈网络的作用Transformer每层结构都少不了三件套残差连接、层归一化、前馈网络。它们各有分工缺一不可。残差连接解决的核心问题是深度网络的退化问题。一般来说网络越深表达能力越强但实际训练时发现层数堆到几十层以后梯度很容易消失模型不升反降。残差连接让每层的输出变成F(x) x即使这一层学的变换没有价值信息也能通过直连通道无损地传到后面梯度也能顺畅地回流。Transformer原始结构里每个子层输出都要加上输入再进下一层就是这个原因。层归一化则是对每个样本的所有特征维度做归一化把数据拉回均值为0、方差为1的分布。它和BatchNorm不一样的是不依赖batch内其他样本的统计量在序列长度变化比较灵活的任务里更稳定。之前做训练的时候我对比过用LayerNorm收敛明显比不用稳学习率可以开得更大也不容易崩。前馈网络在Transformer里被设计成一个两层的全连接结构先升维再降维。原论文里FFN的中间维度设为2048是d_model512的4倍。这个设计思路是让每个位置的向量先映射到更高维空间做非线性变换再压回原来的维度相当于给每个token的表示做了一次特征交叉和抽象。位置上说是“逐位置”意思是同一个FFN对序列里的每个token独立使用权重共享。3.4 Transformer的宏观数据流把以上所有模块拼在一起数据在Transformer内部的完整流程是这样的输入序列先做tokenization每个token通过词嵌入矩阵映射成d_model维的向量加上位置编码模型第一次获得“每个词在哪里”的信息向量进入Encoder经过N层堆叠经典配置是6层。每层内部自注意力聚合全局信息残差和归一化稳定分布FFN逐位置做非线性变换Encoder最终输出一组上下文向量代表整个输入序列的语义编码Decoder接收目标序列的前面部分先经过掩码自注意力限制自己只能看到已生成内容再通过交叉注意力去关注Encoder输出的信息最后经过FFNDecoder每一步输出一个预测概率分布对应词表上每个词的概率取概率最高的词作为当前步的生成结果送入下一步解码。训练阶段可以一次性把整个目标序列喂进Decoder靠掩码保证因果性并行计算效率高。推理阶段只能自回归地一个一个词生成每步把新生成词拼到已有序列里重新输入。4. 从零手写一个Transformer核心组件4.1 用PyTorch实现缩放点积注意力与多头注意力理论容易说真写代码时才会碰到各种细节问题。下面给出一套可以直接运行的PyTorch实现目标是构造一个能完成copy任务的Mini Transformer。copy任务是让模型学会把输入序列原样输出虽然简单但能完整检验模型前向计算、训练和推理整个链路是否正常。先看缩放点积注意力的实现import torch import torch.nn as nn import torch.nn.functional as F import math class ScaledDotProductAttention(nn.Module): def __init__(self, d_k, dropout0.1): super().__init__() self.d_k d_k self.dropout nn.Dropout(dropout) def forward(self, q, k, v, maskNone): # q: [batch, heads, seq_len, d_k] # k: [batch, heads, seq_len, d_k] # v: [batch, heads, seq_len, 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, float(-inf)) attn F.softmax(scores, dim-1) attn self.dropout(attn) output torch.matmul(attn, v) return output, attn这段代码有几个关键细节值得注意scores除以√d_k是论文规定的缩放操作防止点积值太大导致Softmax落到梯度平缓区mask处理用的是masked_fill把需要屏蔽的位置填成负无穷。为什么不是填0因为Softmax操作是非线性的填0不会让那个位置完全失效而负无穷经过Softmax后权重会变成0dropout放在softmax之后加在注意力权重上而不是加在scores上。这个细节是很多实现容易搞错的。作用是在训练时随机衰减一部分注意力权重起到正则化效果防止模型过度依赖某几个位置。多头注意力是在单头基础上做了投影切分和拼接恢复。这里我把所有头的投影用一个矩阵完成再通过view和transpose拆成多头实现上更简洁效率也更高class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout0.1): super().__init__() assert d_model % n_heads 0 self.n_heads n_heads self.d_k d_model // n_heads self.d_model d_model 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) self.attention ScaledDotProductAttention(self.d_k, dropout) def split_heads(self, x): batch_size, seq_len, _ x.size() x x.view(batch_size, seq_len, self.n_heads, self.d_k) return x.transpose(1, 2) # [batch, heads, seq_len, d_k] def forward(self, q, k, v, maskNone): q self.split_heads(self.w_q(q)) k self.split_heads(self.w_k(k)) v self.split_heads(self.w_v(v)) output, attn self.attention(q, k, v, mask) output output.transpose(1, 2).contiguous() output output.view(output.size(0), -1, self.d_model) return self.w_o(output), attn关于维度的设计这里沿用原始Transformer的设定d_model为512时8个head每个head维度64。d_model必须能被n_heads整除不然reshape时会报错。这里的d_k d_model // n_heads也就是每个头的维度。这个设计的妙处在于总参数量没有变所有头的输出拼起来之后恰好还原成d_model维不影响后续网络结构。4.2 前馈网络、位置编码、层归一化与完整Encoder层逐位置前馈网络就是两个线性变换加一个ReLU激活中间维度放大4倍再加回去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) def forward(self, x): return self.linear2(self.dropout(F.relu(self.linear1(x))))d_ff取的是2048。这个放大倍数不是随便定的太小时FFN特征抽象能力有限太大时中间层的参数量会占据整个模型很大比例训练成本明显上升实际效果提升却越来越有限。做中小模型时可以先按4倍设置再根据训练效果微调。位置编码用正弦公式实现。需要注意把位置编码注册成buffer而不是parameter这样它在保存模型时不会算进可训练参数推理时如果遇到更长的序列还能在代码里动态延长class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000): super().__init__() 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) # [1, max_len, d_model] self.register_buffer(pe, pe) def forward(self, x): return x self.pe[:, : x.size(1)]位置编码实现时有一行容易写错div_term用指数方式构造效果等价于对每个偶数维度i计算10000^(-2i/d_model)但数值上更稳定不会因为指数太小产生精度问题。偶数维度用sin奇数维度用cos这样可以保证每个位置得到一个唯一的编码向量。Encoder的每一层可以封装成这样class EncoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, n_heads, dropout) self.ffn 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_out, _ self.self_attn(x, x, x, mask) x self.norm1(x self.dropout1(attn_out)) ffn_out self.ffn(x) x self.norm2(x self.dropout2(ffn_out)) return x这里采用的是Post-LN结构也就是先残差再加LayerNorm这是原始Transformer论文里的方式。后来GPT和很多新模型改用Pre-LN把LayerNorm放在子层之前训练时更稳定但实现上两类都要能写碰到别人代码时才知道为什么结构略有不同。Post-LN对学习率和warmup更敏感需要精心调节Pre-LN的收敛更稳对新手更友好。5. 模型训练与预测的完整闭环5.1 准备数据构建batch并生成掩码构造一个Copy Task需要的训练集非常容易随机生成整数序列输入和目标都是同一个序列。实际代码里还要做padding因为一个batch里的序列长度不一定相同较短的序列要用占位符补齐。Padding的位置在注意力计算时一定要屏蔽掉否则模型会去关注空白位置白白浪费计算还学到噪音。掩码需要分场景设计。Decoder做训练时输入的每个token位置能看到的位置范围受两个条件限制一是不能看未来的token这由因果掩码控制二是不能看padding位置这由padding掩码控制。两个掩码取交集实现上就是把两个布尔矩阵做逻辑与运算。def generate_padding_mask(seq, pad_idx0): return (seq ! pad_idx).unsqueeze(1).unsqueeze(2) # [batch, 1, 1, seq_len] def generate_causal_mask(size): attn_shape (1, size, size) mask torch.triu(torch.ones(attn_shape), diagonal1).bool() # 上三角为True return ~mask # [1, size, size], 下三角含对角线为True因果掩码这里用了triu函数生成一个上三角为True的矩阵然后取反。这样位置i的注意力只能落在位置0到i之间后面的位置全部被masked_fill成负无穷经过Softmax权重为0。写这段代码的时候建议停下来多想想矩阵维度的逻辑后面Debug时很多问题都出在掩码形状对不上。5.2 构建Mini Transformer模型我把Encoder和Decoder整合成一个完整的Transformer类class Transformer(nn.Module): def __init__(self, vocab_size, d_model128, n_heads8, d_ff512, num_layers3, max_len100, dropout0.1, pad_idx0): super().__init__() self.d_model d_model self.pad_idx pad_idx self.embedding nn.Embedding(vocab_size, d_model) self.pos_encoding PositionalEncoding(d_model, max_len) self.encoder_layers nn.ModuleList([ EncoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(num_layers) ]) self.decoder_layers nn.ModuleList([ DecoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(num_layers) ]) self.fc_out nn.Linear(d_model, vocab_size) self.dropout nn.Dropout(dropout) def forward(self, src, tgt): # src: [batch, src_len], tgt: [batch, tgt_len] src_mask generate_padding_mask(src, self.pad_idx) tgt_pad_mask generate_padding_mask(tgt, self.pad_idx) tgt_len tgt.size(1) tgt_causal_mask generate_causal_mask(tgt_len).to(tgt.device) tgt_mask tgt_pad_mask tgt_causal_mask src_emb self.dropout(self.pos_encoding(self.embedding(src) * math.sqrt(self.d_model))) tgt_emb self.dropout(self.pos_encoding(self.embedding(tgt) * math.sqrt(self.d_model))) for layer in self.encoder_layers: src_emb layer(src_emb, src_mask) for layer in self.decoder_layers: tgt_emb layer(tgt_emb, src_emb, src_mask, tgt_mask) logits self.fc_out(tgt_emb) return logitsDecoderLayer的forward签名里有三个输入分别是decoder自身的输入、encoder输出和两个掩码交叉注意力部分的query来自decoderkey和value来自encoderclass DecoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, n_heads, dropout) self.cross_attn MultiHeadAttention(d_model, n_heads, dropout) self.ffn PositionwiseFeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.norm3 nn.LayerNorm(d_model) self.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) self.dropout3 nn.Dropout(dropout) def forward(self, x, enc_output, src_maskNone, tgt_maskNone): attn_out, _ self.self_attn(x, x, x, tgt_mask) x self.norm1(x self.dropout1(attn_out)) cross_out, _ self.cross_attn(x, enc_output, enc_output, src_mask) x self.norm2(x self.dropout2(cross_out)) ffn_out self.ffn(x) x self.norm3(x self.dropout3(ffn_out)) return x代码里有个小细节嵌入层输出乘以√d_model。这是原论文里的做法因为位置编码的数值范围通常在[-1,1]之间如果不放大嵌入值位置信息会在加法中被词嵌入淹没。乘上√d_model之后两者在数值贡献上大致平衡。5.3 训练循环与推理生成训练阶段直接用交叉熵损失padding位置需要计算loss时屏蔽。一种做法是设置ignore_indexCrossEntropyLoss会跳过填充位置的预测def train_model(model, train_loader, epochs20, lr3e-4): criterion nn.CrossEntropyLoss(ignore_index0) # 忽略padding位置 optimizer torch.optim.Adam(model.parameters(), lrlr, betas(0.9, 0.98), eps1e-9) scheduler torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambdalambda step: min((step 1) ** (-0.5), (step 1) * 0.0001 ** (-0.5)) ) model.train() for epoch in range(epochs): total_loss 0 for batch in train_loader: src, tgt_in, tgt_out batch optimizer.zero_grad() logits model(src, tgt_in) # [batch, tgt_len, vocab_size] loss criterion(logits.reshape(-1, logits.size(-1)), tgt_out.reshape(-1)) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() total_loss loss.item() print(fEpoch {epoch1}, Loss: {total_loss / len(train_loader):.4f})代码里已经把tgt_in和tgt_out设计成错位关系训练时Decoder输入目标序列的[0, T-1]部分预测目标是[1, T]部分。这样每个位置的模型都只看到它之前的内容和推理时的行为一致。注意生成tgt_in和tgt_out时要做错位处理如果目标序列是[2, 5, 1, 3]那tgt_in是[2, 5, 1]tgt_out是[5, 1, 3]开头补一个起始符。学习率要加warmup策略这一点非常重要。Transformer对学习率比较敏感直接用固定学习率的Adam会发现训练前期loss很不稳定甚至直接NaN。上面代码里的LambdaLR实现的是论文里的Noam scheduling学习率先线性升到峰值再按步数的平方根倒数衰减。没有这个调度模型很容易在早期就发散。推理阶段的生成用贪心解码就能完成任务演示def greedy_decode(model, src, max_len20, start_token2, end_token3, devicecpu): model.eval() src src.to(device) tgt torch.tensor([[start_token]], devicedevice) with torch.no_grad(): for _ in range(max_len): logits model(src, tgt) next_token logits[:, -1, :].argmax(dim-1).unsqueeze(1) tgt torch.cat([tgt, next_token], dim1) if next_token.item() end_token: break return tgt.squeeze(0).tolist()用贪心解码跑一批测试序列如果模型训练正常几个epoch之后就能观察到模型输出的序列和输入完全一致。这个效果看着简单但能跑通说明整个链路——数据准备、掩码生成、前向计算、损失计算、反向传播、学习率调度——全部没有问题。之后往Translate或者文本生成任务迁移时只需换掉数据集和词表即可。6. 常见问题与排查经验6.1 训练不收敛与Loss为NaN的排查思路训练Transformer时遇到的最典型问题按出现频率排大概是下面这几种问题现象常见原因处理办法Loss不下降学习率太低或太高检查warmup调度调整峰值学习率Loss输出NaN学习率过大、梯度爆炸、attention计算溢出降低学习率加梯度裁剪检查是否除以√d_k训练时好时坏、震荡Adam的epsilon太小或batch size过小设置eps1e-9适当增大batch size模型生成循环重复内容Decoder训练和推理不一致、没有mask检查Decoder的因果掩码是否正确GPU显存不足序列过长或batch太大减少batch size用梯度累积限制max_lenLoss直接输出NaN是我见过最多的情况。排查看三点第一学习率是否过大第二注意力分数是否因为忘了除以√d_k而太大导致Softmax溢出第三Adam的epsilon是否设置得当PyTorch默认eps是1e-8用在Transformer上有时不够稳定论文里用的是1e-9。还有一个容易被忽视的点是梯度裁剪。Transformer深度较大加上自注意力存在长程路径梯度范数很容易超过一个合理的范围。torch.nn.utils.clip_grad_norm_设置max_norm1.0是一个常见且保守的选择基本不会拖慢收敛但能显著降低NaN风险。6.2 推理阶段的坑自回归生成与训练行为的偏差自回归生成时的性能和训练时有明显差异尤其是Decoder部分。训练时我们一次性喂入整个目标序列通过掩码保证因果性所有位置可以并行计算。但推理时只能一个词一个词来每生成一个新词就要把整个序列重新输入模型完整前向一次。这样做有两个问题。一是慢生成100个词就要做100次前向如果模型层数多、序列长推理时延会肉眼可见地涨。优化方案是使用KV Cache把计算过的Key和Value缓存下来避免每步重新计算前面所有token的注意力。原始Transformer代码里没有这个机制现在的GPT系列推理引擎都做了深度优化。二是误差累积。训练时Decoder输入的是真实目标序列但推理时输入的是模型自己生成的前一步输出。一步生成错了后面的所有预测都会受到污染。这就是为什么训练和推理行为要保持一致——如果训练时Decoder能“偷看”到未来位置那模型会形成依赖一旦推理阶段无法偷看就会全面崩溃。所以那个因果掩码绝对不能写错我见过不少复现代码出问题最后定位都是掩码形状或者布尔取反搞反了。6.3 模型参数量与内存的粗略估算写Transformer代码时估算一下模型大小很实用。Transformer参数量主要由四部分构成词嵌入矩阵、注意力层的QKV和输出投影、前馈网络的两个线性层、LayerNorm和位置编码位置编码通常是buffer不占可训练参数。以d_model128、词表大小50、8头、6层EncoderDecoder的小配置估算词嵌入128 * 50 6400每层多头注意力QKV三个矩阵各128128输出投影128128总共65536每层FFN128512 512128 131072每层总参数量约1966086层双层约2.36M再加上其他杂项整体不到3M参数这个规模用CPU跑Copy Task只需要几分钟。如果换成d_model512、词表3万的经典配置词嵌入一项就有1536万参数整体接近6500万。这就是为什么词表大小对模型体积影响巨大——很多模型压缩工作从共享嵌入矩阵入手就是这个原因。7. 几个值得关注的Transformer变体7.1 视觉方向ViT与Swin TransformerViTVision Transformer把Transformer搬进图像领域的方式很直接把图片切成一堆16x16的小patch每个patch拉平后当作一个“词”加上位置编码后送进标准Transformer。这个思路在数据量足够大时效果非常惊艳ImageNet分类上超过了同量级的CNN。但它有个缺点全局注意力在高分辨率图像上计算开销太大patch数量一多内存就撑不住。Swin Transformer的改进思路是引入层次化和窗口化。它把注意力限制在局部窗口内计算窗口之间的信息通过“移位窗口”机制来交换。这样做的好处是计算复杂度从O(N^2)降到O(N)而且能像CNN一样构建多尺度特征金字塔方便做检测和分割这类密集预测任务。7.2 轻量注意力模块SE、CBAM与CA这几个名字频繁出现在图像分类比赛和移动端模型优化里严格来说不是Transformer本体而是在CNN基础上加进去的注意力模块。SE模块Squeeze-and-Excitation做的是通道注意力把每个通道的二维特征图压缩成一个全局描述符再通过两个全连接层学习每个通道的权重最后把权重乘回原始特征图。它关注的是“哪些通道重要”。CBAM在SE的基础上加了空间注意力分支。通道注意力告诉模型“看什么”空间注意力告诉模型“看哪里”两条分支串行组合效果进一步提升参数量增加却很少。CA模块Coordinate Attention的思路更细它把空间注意力分解成水平和垂直两个方向分别编码长程空间依赖对目标定位和细粒度识别更友好。这些模块的共性是轻量、即插即用做图像模型时可以在残差块里直接插入。但要注意的是它们的原理和Transformer里的自注意力有本质区别SE、CBAM、CA更接近“对已有特征做通道或空间上的重标定”而Transformer自注意力做的是“所有位置的信息互相交换”两者解决的问题不在一个层面。7.3 MoE架构与多模态方向MoEMixture of Experts的思路是让模型在推理时只激活一小部分参数而不是每一层都动用全部前馈网络。具体做法是设置多个FFN专家由一个门控网络为每个token选出最合适的几个专家。GPT-4、DeepSeek等大模型走的就是这个路线在总参数不变的前提下大幅降低推理成本让模型可以在更大规模上扩展。多模态Transformer则是把文本、图像、音频等不同模态的数据统一成token序列之后一起送入模型让注意力机制在不同模态之间建立关联。近年来这类模型的进展非常快从图像理解到文生图、视频生成底层几乎都是Transformer架构或其变体。这些方向其实都说明一件事Transformer这套“让所有元素之间互相联系”的框架本身已经超越了具体任务的限制成为一种通用信息处理范式。深度理解基础架构后往上做任何方向的扩展都会轻松很多。8. 写在最后的实操心得手写过一遍Transformer之后最大的体会是看论文和真正实现完全是两码事。论文里的公式只有三行但落在代码上光是一个掩码就要处理好几种形状和叠加逻辑一个维度对不上反向传播阶段就会炸出莫名其妙的错误。建议所有想深入理解这套架构的人都亲手把这个模型从零到一实现一遍——不需要复现多大规模能跑通Copy Task、翻译一个简单句子就足够建立完整的认知框架。再分享一个影响很大的细节训练策略和模型结构同样重要。同一个Transformer代码用固定学习率训练可能怎么都收敛不了加上warmup和梯度裁剪之后几个epoch就开始稳定下降了。第一次跑通时那种感觉是——原来这个模型对训练技巧的要求这么高也难怪原论文里要花那么大的篇幅写训练细节。新手阶段一定要把这些基础训练策略一起学会它们在未来做更大模型时同样不可或缺。