Transformer原理与PyTorch实现：从自注意力到编码器-解码器架构详解

1. Attention Is All You Need——从“一句话”到一场架构革命

2017年，一篇题为《Attention Is All You Need》的论文被放到arXiv上，彼时自然语言处理领域还在RNN、LSTM、GRU的统治之下。序列建模的常规打法是“一步步走”：当前时刻的输出依赖上一时刻的隐状态，像一条流水线，前面的工序没做完，后面的就没法开工。这种循环结构天生存在两个痛点，一是难以并行，训练速度被死死摁住；二是长距离依赖问题——句子太长时，靠一个隐状态向量把前面所有信息“压缩”进去，信息必然大量丢失，梯度也容易在反向传播中消失或爆炸。

Transformer就是在这种背景下被提出的。它把注意力机制从“辅助组件”升级成了“主架构”，彻底放弃了循环和卷积，只用注意力完成输入输出之间的全局依赖建模。最初它是为机器翻译设计的，但随后的几年里，BERT、GPT、T5、ViT、Swin Transformer等一系列模型彻底改变了大模型和AI的生态格局。

这篇文章，我想从一个“想真正搞懂Transformer内部到底发生了什么”的视角，把它的原理、架构设计和实现细节逐步拆开讲一遍。适合刚接触Transformer、读过论文但没完全吃透、或者想从公式和代码层面加深理解的读者。

2. 整体设计思路与核心思想拆解

2.1 为什么循环结构不是唯一解——Transformer的设计哲学

在Transformer出现之前，处理序列数据的典型工具是RNN及它的变体LSTM、GRU。RNN的核心思路是“状态传递”：一个隐状态向量在时间步之间流转，每读入一个新单词，就把这个单词的信息“融合”进状态里。

这个设计的类比就是一条单行隧道：每辆车（单词）都必须按顺序进入隧道，当前车辆的位置由前面所有车辆决定。隧道里的车越多，后面的车越难看到前面的路况——这就是长距离依赖问题。而且这个隧道内部是逐车放行的，效率极低。

Transformer的哲学完全不同：它把句子里的所有单词同时摆出来，任意两个单词之间直接建立联系，不再依赖“中间人”式的状态传递。你可以把它的自注意力机制理解成一场全员大会：每个参会者（单词）都可以直接和所有其他参会者对话，互相了解对方在聊什么。这种全连接的信息交互方式，既解决了远距离信息衰减的问题，又天然可以并行计算。

这个“全连接替代顺序传递”的核心转变，带来一个连锁结果：模型不再需要按顺序处理输入，所有位置可以同时计算。于是GPU的大规模并行能力终于被完全释放出来，训练效率比RNN高了一个数量级。

2.2 编码器-解码器框架——让Transformer跑起来的整体蓝图

原版Transformer沿用了序列到序列任务中的编码器-解码器结构。编码器负责把源语言句子“理解”成一组中间表示，解码器负责根据这些表示，逐词生成目标语言句子。

编码器由6个相同的层堆叠而成，每层包含两个子层：多头自注意力层和前馈神经网络层。每个子层外面套了一层残差连接，随后做层归一化。

解码器也是6层堆叠，但每层比编码器多一个子层。第一个子层是带掩码的多头自注意力，因为生成第t个词时，只能看到已经生成的词，不能偷看未来的词。第二个子层是编码器-解码器注意力，负责让解码器从编码器的输出中“提取”与当前生成相关的信息。第三个子层依旧是前馈网络。

输入序列先经过词嵌入层映射成向量，再叠加位置编码，携带上顺序信息后进入编码器。解码器的输入同样是目标语言的词嵌入加位置编码。

整体数据流向大概是这样的：

• 源语言句子 -> 嵌入 + 位置编码 -> 6层编码器 -> 编码器输出
• 目标语言句子 -> 嵌入 + 位置编码 -> 带掩码的注意力 -> 编码器-解码器注意力 -> 前馈网络 -> 线性层 -> Softmax -> 输出概率

这套结构后来成为各类Transformer变体的基础母版。

2.3 与传统序列模型的核心差异一览

我整理了一张对比表，方便你直观感受Transformer和RNN/LSTM之间的本质区别：

O(n^2) 复杂度是Transformer最大的软肋，这也是后来各种稀疏注意力、线性注意力研究的出发点。但在大多数序列长度适中的场景里，这个代价是值得的。

3. 核心细节拆解——从嵌入到自注意力的层层解构

3.1 输入表示：词嵌入与单词的“语义坐标”

Transformer无法直接处理离散的单词，需要先把每个词映射成一个稠密向量，这就是词嵌入。

词嵌入层本质上是一个查找表：每个词对应一个固定维度的向量（原论文中维度 d_model = 512）。训练过程中，这个查找表会不断更新，最终让语义相近的词在向量空间中距离更近。比如“猫”和“狗”的向量距离，会明显小于“猫”和“汽车”的距离。

代码实现很直接，在PyTorch里就是一行：

import torch.nn as nn d_model = 512 vocab_size = 32000 embedding = nn.Embedding(vocab_size, d_model)

使用嵌入层得到的向量本身不包含任何位置信息。一个句子把词序完全打乱，得到的嵌入向量集合是一模一样的。这对一句“我打你”和“你打我”意思完全不同的话来说，显然不能接受。所以Transformer引出了位置编码。

3.2 位置编码的玄机——为什么用正弦和余弦函数

原论文采用了一种基于三角函数的固定位置编码方式。对于位置 pos 和维度 i，编码值定义为：

PE(pos, 2i) = sin(pos / 10000^(2i / d_model)) PE(pos, 2i + 1) = cos(pos / 10000^(2i / d_model))

也就是说，偶数维度用正弦函数，奇数维度用余弦函数。不同位置上的编码向量各不相同，词嵌入和位置编码逐元素相加后，每个词就同时携带了语义信息和位置信息。

为什么偏偏选正弦和余弦？这背后有几个很巧妙的原因：

第一，这种编码方式可以用相对位置关系来表达。通过三角函数的和差化积公式，PE(pos+k) 可以表示为 PE(pos) 的线性函数。换句话说，模型能够学到“位置pos的词”和“位置pos+k的词”之间的相对位置关系，这对语言理解非常有价值。

第二，三角函数的取值范围天然固定在[-1, 1]之间，不会出现数值爆炸，和词嵌入相加时也不会破坏原有特征的量级。

第三，它是固定的，不需要训练参数，实现简单且能处理任意长度的序列。

位置编码还可以采用可学习的方式——把位置也当作一个嵌入层，训练时一并更新。BERT用的就是这个方案。两种方式实验效果相差不大，固定三角函数在序列长度泛化上略有优势，可学习方式在短序列任务上有时调参更方便。

3.3 自注意力机制：Q、K、V到底是干嘛的

自注意力机制是Transformer的心脏。为了说明白它，我习惯用一个生活化类比：你在一个满是人的房间里要找合作伙伴，你会做三件事——

• Query（查询/问句）：你心里有一个需求，比如“谁懂Python”。
• Key（键/标签）：每个人的胸前挂着一个牌子，写着各自的技能标签。
• Value（值/内容）：每个人真正的能力细节。

你拿着自己的需求（Q），和房间里所有人的标签（K）逐一比对，计算匹配度，然后根据匹配度加权汇总所有人的能力（V），于是你得到了一个最符合需求的人的综合画像。

注意力机制就是这个过程的数学化。

具体操作上，输入向量 x_i 分别经过三个权重矩阵 W_Q、W_K、W_V 映射，得到 q_i、k_i、v_i。随后进行如下计算：

Attention(Q, K, V) = softmax(Q * K^T / sqrt(d_k)) * V

这里的 d_k 是每个头的维度。除以 sqrt(d_k) 用于缩放，目的是防止Q和K的点积数值过大，把softmax推进饱和区导致梯度消失。

关于为什么除以 sqrt(d_k)，有一个简单的统计学解释：如果 Q 和 K 的每个分量都是均值为0、方差为1的独立随机变量，那么它们的点积的均值是0，但方差等于 d_k。方差越大，点积的绝对值就可能越大，softmax输出的分布就越尖锐——几乎所有的注意力权重都集中到一个位置上，其他位置的梯度会变得很小。除以 sqrt(d_k) 之后，点积的方差就被重新拉回1附近，softmax的输入落在梯度平缓区之外的概率大大降低。

我见过的不少同学在复现时直接忽略了这个缩放因子，结果训练初期损失函数震荡剧烈，加了这一步后立刻稳定下来，可见这个细节不是可选项，是必需品。

3.4 多头机制——模型为什么要“多几双眼睛”

单套注意力只能学习一种“匹配模式”。但一个词在句子中可能同时扮演多种角色，和不同词之间存在不同层面的关系。比如“苹果”这个词，在“苹果发布了新手机”里它和“手机”是品牌-产品关系，在“我吃了一个苹果”里它和“吃”是动作-对象关系。

多头注意力的核心思想，是并行地用多套 Q、K、V 映射矩阵，让每套从不同子空间学习不同的关系模式。原论文使用 8 个头，每个头的维度 d_k = d_model / 8 = 64。每个头独立计算注意力后，输出拼接在一起，再经过一个线性层投影回 d_model 维。

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) * W_O 其中 head_i = Attention(Q * W_Q_i, K * W_K_i, V * W_V_i)

实验经验表明，8 个头在大部分任务上表现良好。头数太少，模型的表达能力受限，难以同时捕捉多种关系；头数太多，每个头的维度被压缩得过小，单个头能承载的信息容量不足，而且训练开销和显存占用成倍增加。

3.5 前馈网络：每个位置的“深度思考”

自注意力子层负责信息交互，算的是“词与词之间的关系”。但模型还需要一个独立的非线性变换层来提取更抽象的特征——这就是位置感知的前馈网络。

它是一个单隐层的全连接网络，公式如下：

FFN(x) = max(0, x * W_1 + b_1) * W_2 + b_2

中间层的维度被扩展到 d_ff = 2048，是 d_model 的四倍。激活函数用的是ReLU，后来一些变体把它换成了GELU或SwiGLU，效果略有提升。这个子层的本质是给模型增加非线性表达能力，把注意力层输出的表示映射到更高维空间做特征提取，再压缩回原维度。

有意思的是，这个前馈网络是逐位置独立应用的——也就是说，同一个FFN作用在每个位置上的向量，但不同位置之间没有信息交换。所以整个Transformer的信息交换只发生在注意力层，前馈层是做逐点的非线性变换。

3.6 残差连接与层归一化——让训练稳定的两块基石

残差连接来自ResNet的思想：把子层的输入和输出直接相加，让梯度可以无损地跨层传播。没有残差连接，深层Transformer几乎无法训练，因为梯度经过多层叠加后会越来越小，或者越来越大。

层归一化的操作对象是每个样本的每个位置上的特征向量，对这个 d_model 维向量计算均值方差，然后做归一化，再引入两个可学习的缩放和平移参数。和BatchNorm的区别在于：LayerNorm对整个特征维度做归一化，和批量大小无关，因此不受batch_size变化的连锁影响，在NLP任务中大放异彩。

原论文的层归一化位置在残差连接之后，也就是“先相加再归一化”。后来GPT等模型改成“先归一化再相加”，理由是这样对训练的稳定性更友好。两种方式都有人用，但Pre-LN在实践中更常见。

3.7 解码器里的掩码机制——为什么不能偷看未来

解码器在自回归生成时需要保证一个规则：预测第t+1个词时，只能看到第1到第t个词，不能看到后面的词。不然就相当于考试时偷看了答案，模型不会真正学会生成。

实现方式很简单：在计算Q和K的点积之后，把矩阵中对应“未来位置”的数值替换成一个极小的负数（比如 -1e9），这样经过softmax后这些位置的权重趋近于零，就相当于被遮住了。

mask: 一个上三角矩阵，右上角全为-1e9，其余位置为0

这里要特别注意：编码器里没有掩码，因为源语言句子是一整句同时输入的，模型要考虑全部词之间的关系。解码器里两种注意力的情况也不同：第一个注意力子层要加掩码，第二个编码器-解码器注意力子层不需要掩码（解码器的Q要访问编码器输出的所有K和V）。

4. 动手实现——从零手写一个简化版Transformer

4.1 环境与准备工作

在动手写代码前，先把环境准备妥当。我用的配置如下：

• Python 3.10+
• PyTorch 2.0+
• GPU显存建议至少8GB（没有GPU也能跑，只是训练会慢很多）

代码组织上，我建议按模块拆分：注意力层单独一个文件，编码器层、解码器层各放一个，方便调试和维护。

4.2 核心模块的PyTorch实现

注意力层是第一个要写的模块。这里我把多头注意力封装成一个子类：

import torch import torch.nn as nn import torch.nn.functional as F import math class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout=0.1): super().__init__() assert d_model % n_heads == 0, "d_model must be divisible by n_heads" self.d_model = d_model self.n_heads = n_heads self.d_k = d_model // 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) self.dropout = nn.Dropout(dropout) self.scale = math.sqrt(self.d_k) def forward(self, q, k, v, mask=None): batch_size = q.size(0) # 1. 线性映射后拆分多头：从(batch, seq_len, d_model) # 变成(batch, n_heads, seq_len, d_k) Q = self.w_q(q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) K = self.w_k(k).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) V = self.w_v(v).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) # 2. 计算注意力分数 scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale # 3. 掩码处理 if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) # 4. softmax归一化 + 加权求和 attn_weights = F.softmax(scores, dim=-1) attn_weights = self.dropout(attn_weights) context = torch.matmul(attn_weights, V) # 5. 合并多头，线性投影 context = context.transpose(1, 2).contiguous().view( batch_size, -1, self.d_model ) output = self.w_o(context) return output

这里有个很关键的操作是view和transpose的顺序。很多人手写时容易出错：在切成多头时，要先把最后一维拆成(n_heads, d_k)，再用transpose调到正确维度，最后要记得调用contiguous()做内存重排，否则后面的view会报错。

接下来实现前馈网络：

class PositionWiseFeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout=0.1): super().__init__() self.fc1 = nn.Linear(d_model, d_ff) self.fc2 = nn.Linear(d_ff, d_model) self.dropout = nn.Dropout(dropout) def forward(self, x): return self.fc2(self.dropout(F.relu(self.fc1(x))))

然后是位置编码：

class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len=5000, dropout=0.1): super().__init__() self.dropout = nn.Dropout(dropout) pe = torch.zeros(max_len, d_model) position = torch.arange(0, max_len, dtype=torch.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): x = x + self.pe[:, : x.size(1), :] return self.dropout(x)

注意register_buffer的使用——它让位置编码矩阵随模型移动设备（GPU/CPU）时自动切换，同时不会被视为需要更新梯度的模型参数，这是一个很省心的细节。

编码器层和解码器层就是把这些模块依次拼起来：

class EncoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout=0.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, mask=None): # 残差连接 + 层归一化 x = self.norm1(x + self.dropout1(self.self_attn(x, x, x, mask))) x = self.norm2(x + self.dropout2(self.ffn(x))) return x class DecoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout=0.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_mask=None, tgt_mask=None): x = self.norm1(x + self.dropout1(self.self_attn(x, x, x, tgt_mask))) x = self.norm2(x + self.dropout2(self.cross_attn(x, enc_output, enc_output, src_mask))) x = self.norm3(x + self.dropout3(self.ffn(x))) return x

注意解码器的cross_attn里，Q来自解码器自身（目标语言），K和V来自编码器输出（源语言）。这是信息在两个语言之间流动的关键通道。

4.3 训练一个微型翻译模型

整体模型搭建好之后，还要处理数据加载、损失计算和训练循环。语言任务里最核心的一个细节就是label smoothing和损失函数。

这里我直接推荐用PyTorch自带的CrossEntropyLoss并配合ignore_index参数：

criterion = nn.CrossEntropyLoss(ignore_index=0, label_smoothing=0.1)

其中0是padding的索引，label_smoothing是Label Smoothing技术，它不把真值标签的one-hot分布当作绝对目标，而是稍微“抹平”一些概率，把一部分概率均匀分给所有错误类别。这个技巧能在训练前期防止模型过度自信，提升泛化能力。

训练批量大小根据显存来定，原文用的是128，我在单卡环境下通常设32或者64。学习率调度的模式用的是论文里的Noam式：先线性增加到峰值，再按平方根的倒数衰减。简单起见，也可以用Warmup + Adam默认学习率1e-3：

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.98), eps=1e-9)

训练循环不多赘述。这里给出两个我在实际调通代码时踩过的坑：

• mask的维度对齐。在多头注意力中，mask先要reshape成(batch, 1, 1, seq_len)，才能正确广播到(batch, n_heads, seq_len, seq_len)，否则要么报维度错，要么该掩的位置没掩上。
• 解码器训练时的输入是右移一位的目标序列。意思是：训练时给解码器输入“起始符+正确答案的前t个词”，让它预测“答案的第t+1个词”。这个shift逻辑不处理好，训练和推理会出现严重的不一致。

4.4 推理阶段的逐词生成

训练结束后，推理是一个自回归过程：初始输入只有起始符，模型输出第一个预测词，把这个词追加到输入里，再喂回解码器，循环直到输出结束符或达到最大长度。

这里有一个常见的优化手段叫缓存KV：在生成第t+1个词时，前t个词对应的K和V其实没有变化，不需要重新计算。把每一层的K、V缓存下来，能节省大量重复计算，显存換速度，是实际部署时必做的优化。

5. 常见问题与排查技巧实录

5.1 模型训练不收敛，loss一直在原地打转

这是我被问得最多的一个问题。造成这种状况的原因通常是以下几个：

• 学习率过大。Transformer对学习率异常敏感，过大时loss直接震荡不降。建议先降到原来的1/10试试，再配合Noam式warmup。
• 缩放因子没加。前面强调过的除以sqrt(d_k)如果漏了，softmax容易饱和，梯度消失，loss半天掉不动。
• 初始化问题。Linear层默认初始化有时候在深层次Transformer里不够稳。可以手动用Xavier初始化替换前馈层和注意力投影层的权重。
• 数据里有大量padding。padding的位置如果不被mask掉，模型会学到一堆噪声。确保padding mask和注意力mask都传进了注意力层。

5.2 显存爆炸，OOM频繁

• 最直接的调整是把batch_size调小。
• 确认用的是自动混合精度训练。PyTorch 2.0里的torch.autocast配GradScaler，用FP16训练，正常能省一半左右的显存。
• 序列长度如果是个位数和上百位的情况混杂，可以用动态批次（按长度分桶），避免整个batch都被最长句子拖累。

5.3 解码器输出出现“重复词”或者“循环词”

推理时生成结果经常出现“我我我我”或者“喜欢喜欢喜欢”这种状况，本质是模型的自回归误差叠加导致的。常见解法有：

• 做beam search（束搜索），不要用贪心解码。
• 对重复n-gram做惩罚（penalty）。
• 使用top-k采样配合temperature调节，在生成任务里效果明显。
• 降低温度参数让分布更集中，减少随机性过大导致的生成漂移。

5.4 位置编码选固定三角函数还是可学习参数

原版的三角函数式位置编码在绝大部分场景下已经够用。可学习位置编码在短文本任务上有时能小幅提升效果，但序列长度一旦超过训练时的max_len，泛化会明显变差。如果任务里存在超长序列场景，还有相对位置编码（如T5的相对位置偏置）和旋转位置编码（RoPE）可选——后者因为对长序列外推更友好，已经被大量主流大模型采用。

我个人的经验是：短任务优先用固定三角函数，省心省力不踩坑；涉及大规模预训练的动态序列场景，直接考虑RoPE或ALiBi这类更新式的方案。

5.5 训练与推理行为不一致的坑

训练时解码器输入是一整段ground truth（教师强制），而推理时输入是模型自己的预测，这种不一致叫exposure bias。缓解手段是计划采样——训练过程中以一定概率把真实词替换成模型预测的词，让模型慢慢适应自己的错误分布。不过在实际落地中，大多数场景下做不做这一步差别不大，优先把模型规模和训练数据搞上去，收益更高。

6. 从原版到百花齐放——Transformer的演进方向与个人实操体会

Transformer从一出生就呈现极强的可塑性。BERT把编码器拿来做双向语言理解，GPT把解码器改造成单向自回归生成器，ViT把图像切块当作句子序列喂进编码器，Swin Transformer引入分层和窗口注意力处理视觉高分辨率问题。短短几年，它从NLP边界扩展到计算机视觉、语音、推荐系统、蛋白质结构预测等各种领域。

这些演变背后的几个主线其实很清楚：

• 一是注意力效率的提升：从标准O(n^2)注意力到稀疏注意力、滑动窗口、线性注意力；
• 二是位置编码方案更迭：固定式、可学习式、相对位置、旋转位置；
• 三是归一化和激活函数的微调：Pre-LN、Post-LN、GELU、SwiGLU；
• 四是训练技巧的成熟：混合精度、梯度累积、模型的规模扩展法则。

我个人在这几年的工程实践中，最大的一个体会是：不要急着堆层数和参数，先把一个小的Transformer在你的数据上跑通、吃透，观察loss曲线和注意力热力图，你对模型内部工作机制的理解和直觉建立会非常快。这个模型虽然看起来复杂，但它的每一个设计决策都有清晰的逻辑支撑——残差解决梯度、缩放解决数值、多头解决关系、掩码解决因果。搞懂了这些，后面的BERT、GPT、ViT对你来说就是换皮和微调的差异问题。

如果这篇文章能帮你在“原理上理清楚、代码里跑明白”这两个维度上都有所收获，我这几天的整理就没白费。现在不妨打开PyTorch，亲手写完那个简化版，然后把注意力权重可视化出来，你会看到完全不一样的世界。