自然语言处理 Transformer 系列写到了第五篇。前四篇分别把注意力机制的来龙去脉、Transformer 整体架构、训练细节,以及如何用图解的方式把论文里那些绕口的术语翻译成人话讲了一遍。按原计划这一篇应该开始聊 BERT 和 GPT 这类预训练模型的演进,但我临时改了个主意,先停下来把手写实现这件事做掉。原因很简单——连续几篇都在“看懂”层面打转,再不碰代码,很多读者对 Transformer 的理解就一直浮在表面。
这篇文章的核心任务很明确:从零手写一个最小可用的 Transformer,并用正弦数据预测这个经典场景做验证。听起来好像是从 NLP 跑偏到了时间序列,但实际上这是检验模型结构是否真正学明白的最佳试金石。正弦数据没有词表、没有复杂的预处理,却能同时验证位置编码、自注意力、Encoder-Decoder 交互这些核心组件是否真的在工作。文章适合两类读者:一类是已经看过不少 Transformer 图解和理论文章、但一直没动手写过代码的人;另一类是写过简单分类模型、想理解序列到序列任务到底怎么搭的训练者。读完这篇文章,你会带着一个自己能跑、能改、能继续往 NLP 方向迁移的最小实现离开。
1. 组件拆解与架构逻辑——先搞清楚每个零件为什么存在
手写代码之前,我建议先把整个结构在脑子里重新过一遍。很多人一开始就盯着代码逐行看,结果被多头注意力的 reshape 绕晕了,根本原因是脑子里没有一张结构总图。Transformer 的代码并不难,难的是没有建立“每个模块到底解决什么问题”的认知,一旦理解到位,代码就只是一层窗户纸。
1.1 Encoder 与 Decoder 的分工:两座塔各干各的活
原版 Transformer 是 Encoder-Decoder 结构。Encoder 这一侧负责把完整的输入序列编码成语义表示,它的特点是能看到整个序列的所有位置,这就是所谓的双向上下文。比如在机器翻译里,源语言句子中的所有词都可以互相参考,因为翻译一个词时本来就需要参考整个句子的意思。Decoder 这一侧则完全不同,它负责逐个生成目标序列,每一次只能看到当前位置之前已经生成的内容,不能偷看后面的词,否则就变成了“开卷考试”。
这个分工直接决定了两个塔在代码实现上的差异:Encoder 使用普通自注意力,所有位置互相可见;Decoder 的自注意力必须加一个上三角掩码,把未来位置全部遮住。而两个塔之间还有一层交叉注意力,Decoder 的每一层都会去 Encoder 的输出里“查询”信息——这在代码里对应的是 cross-attention 的 key 和 value 来自 Encoder 输出,query 来自 Decoder 自身隐状态。
我在刚开始写代码时犯过一个认知错误:以为交叉注意力和自注意力在实现上需要两套完全不同的代码。后来才想明白,多头注意力模块本身是通用的,区别只在传入的 query、key、value 来自哪里。自注意力调用时传三次同一个输入,交叉注意力调用时传入的是两个不同来源的张量。这个“同一个模块,不同输入”的思路,是理解 Transformer 代码结构的一把钥匙。
1.2 位置编码:没有这个,序列就失去了顺序
Transformer 和 RNN、LSTM 最本质的区别,就是它没有天然的时序结构。RNN 是一个一个词按顺序处理的,位置信息隐含在循环次数里。Transformer 是并行处理的,所有词一次性进入模型,如果不额外把位置信息告诉模型,那么“我爱你”和“你爱我”在模型看来就是完全相同的输入,只是词向量排列顺序不同但整体集合相同。这显然不行。
所以位置编码的作用,就是给每个位置生成一个独特的、可学习的或固定的向量,加到词嵌入上。原版论文用的是正弦函数式的位置编码,核心思想是用不同频率的正弦和余弦波组合,让每个位置获得唯一的编码模式。频率的选择是有讲究的:低频分量在不同位置之间变化平缓,高频分量变化剧烈,组合起来就能在低维空间里区分任意长度的位置。
有个细节值得多说一句:词嵌入在加位置编码之前,通常会乘以一个缩放系数,比如根号 d_model。这么做的原因是位置编码的数值范围是固定的[-1,1],而词嵌入经过随机初始化后数值通常很小,如果不缩放,位置编码会把词嵌入的信息完全淹没。缩放之后两者在数值量级上更加匹配,模型一开始训练时就能同时看到语义和位置信息,而不是先花大量时间重新平衡两者的比例。
1.3 多头注意力:“多”在哪里,为什么不是单头
把注意力机制想象成你在阅读一道复杂题目时的思考过程。单头注意力意味着你只能从一个角度去分析题目,比如只看关键词的相似度。而多头注意力相当于同时从多个角度审视这个问题——一个头关注语法关系,一个头关注指代关系,一个头关注语义相近的词——最后把多个角度的结论拼接起来,形成更全面的判断。
代码实现上,多头并不是真的调用了多次独立的注意力函数然后拼结果,那样效率太低。实际做法是把 d_model 维度的向量切成 n_heads 份,每一份通过独立的线性变换得到自己的 query、key、value,然后并行计算。严格来说,PyTorch 的张量操作里这一步是通过 view 和 transpose 把维度重新排列实现的,本质上利用了矩阵运算的并行性。
头数的选择是个工程问题。头数太少,每个头要负责的表示空间太大,模型表达能力受限;头数太多,单个头的维度太小,反而学不到有效信息。实际经验是 d_model 和 n_heads 的比例关系比绝对数值更重要,d_k = d_model / n_heads 一般在 16 到 64 之间。如果发现注意力矩阵可视化出来一片均匀,所有位置权重差不多,多半是头数太多或者训练不充分,导致每个头都没学到专精的模式。
2. 手写实现:从零搭一个最小可用的 Transformer
理论捋清楚之后,代码其实只是把上面这些概念翻译成张量操作。我用的工具是 PyTorch,版本 2.x 都行,不需要任何第三方 NLP 库。整个实现大约三百行,包含位置编码、多头注意力、前馈网络、Encoder 层、Decoder 层和最终的组装。下面我会按模块讲,每个模块都附上完整代码和关键解释。
2.1 Embedding 与位置编码的具体实现
首先实现位置编码模块。这里有几个值得注意的工程细节,我直接放到注释和代码后的解释里。
import math import torch import torch.nn as nn class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len=512): super().__init__() 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) self.register_buffer('pe', pe) def forward(self, x): return x + self.pe[:, :x.size(1)]代码里 div_term 计算看起来有点绕,其实是把论文公式里的1 / 10000^(2i/d_model)换成了指数形式exp(-2i * log(10000) / d_model)。这样做的目的是防止当 i 很大时直接算幂次导致数值不稳定或者速度变慢。pe[:, 0::2]表示取所有偶数维列,pe[:, 1::2]取所有奇数维列,分别填入正弦和余弦值。这样每个位置的编码都不是孤立的,相邻位置在部分维度上拥有连续性变化,模型就能感知到相对距离关系。
这里还要注意一点,我用 register_buffer 来存储位置编码矩阵,而不是直接赋给某个属性。这样做的正确性在于,buffer 会随模型一起自动迁移到 GPU、一起保存和加载,但又不会被当作需要梯度的模型参数。位置编码不是训练出来的,不需要更新,所以不能放进 nn.Parameter。
2.2 多头注意力机制手写完整代码
多头注意力是整个模型的绝对核心,我见过很多人在这里卡住。其实只要记住一条主线:线性投影、分头、缩放点积注意力、合并头、输出投影,就永远不会乱。
class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout=0.1): super().__init__() assert d_model % n_heads == 0 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) def forward(self, query, key, value, mask=None): batch_size = query.size(0) # 线性投影 + 分头 Q = self.w_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) K = self.w_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) V = self.w_v(value).view(batch_size, -1, self.n_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 = torch.softmax(scores, dim=-1) attn = self.dropout(attn) output = torch.matmul(attn, V) # 合并所有头 + 输出投影 output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) return self.w_o(output)分头那一步很多人看不懂,我拆开讲。query 经过线性投影后形状是(batch, seq_len, d_model),view 把它变成(batch, seq_len, n_heads, d_k),也就是说把原来的 d_model 维拆成了 n_heads 组,每组 d_k 维。然后 transpose(1, 2) 把 n_heads 提到第二位,得到(batch, n_heads, seq_len, d_k)。这一步做完之后,每个头在自己的 d_k 维空间里做矩阵乘法,互不干扰。
scores 除以根号 d_k 这个缩放操作,很多人只是照抄,不清楚原因。如果不缩放,当 d_k 较大时,点积结果会很大,softmax 会进入饱和区,梯度趋近于零,训练直接卡死。除以根号 d_k 可以理解成把点积结果拉回到一个合理的数值范围,让 softmax 有足够的梯度信号。这是一个写错不会立即报错、但会严重影响训练效果的细节。
mask 的处理方式是关键。NLP 任务里常见两种 mask,一种是 padding mask,用来遮住填充的无效 token;一种是 look-ahead mask,上三角矩阵,让 Decoder 的每个位置只能看到自己前面的位置。代码里masked_fill(mask == 0, -1e9)把被遮住的位置填成一个很大的负数,softmax 之后这些位置的权重就趋近于零,等于不让模型看到这些位置。这里的 -1e9 是经验值,不能填 -inf,因为某些框架里 -inf 经过 softmax 可能出现 NaN。
2.3 前馈网络、残差连接与层归一化
注意力层的输出不会直接传递到下一层,而是先过一个前馈网络、加上残差连接、再做层归一化。前馈网络本身结构很简单:线性层把 d_model 维升到 d_ff 维,经过 ReLU 激活,再降回 d_model 维。这个非线性变换的引入是有道理的,因为注意力本质上是线性加权求和,如果不加前馈网络,多层 Transformer 叠加之后仍然是线性变换,表达能力会大打折扣。
残差连接解决的是深层网络的梯度传播问题。Transformer 动辄六层、八层甚至更深,没有残差连接,梯度在反向传播过程中很容易消失。残差连接让梯度可以跳过子层直接回传,相当于给深度网络加了一条“高速公路”。
层归一化在代码里直接调用 nn.LayerNorm 就行,但有个结构设计值得注意:原版论文用的是 Post-LN 结构,也就是先执行子层计算,残差相加之后再归一化;后来很多实践发现 Pre-LN 结构,也就是先归一化再进子层,在深层模型上训练更稳定。我在手写代码里用的是 Post-LN,因为更贴近原论文,如果你发现训练深层模型时 loss 震荡非常剧烈,可以改成 Pre-LN 试试,很多情况下问题直接解决。
class SublayerConnection(nn.Module): def __init__(self, d_model, dropout=0.1): super().__init__() self.norm = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, sublayer): return x + self.dropout(sublayer(x)) def forward_pre(self, x, sublayer): return x + self.dropout(sublayer(self.norm(x)))2.4 组装 Encoder 和 Decoder
有了以上所有基础模块,就可以像搭积木一样组装出完整的 Transformer 了。每个 Encoder 层包含两个子层:自注意力子层和前馈网络子层。Decoder 层包含三个子层:带掩码的自注意力、交叉注意力、前馈网络。
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.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout) self.sublayer1 = SublayerConnection(d_model, dropout) self.sublayer2 = SublayerConnection(d_model, dropout) def forward(self, x, mask): x = self.sublayer1(x, lambda x: self.self_attn(x, x, x, mask)) x = self.sublayer2(x, self.feed_forward) 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.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout) self.sublayer1 = SublayerConnection(d_model, dropout) self.sublayer2 = SublayerConnection(d_model, dropout) self.sublayer3 = SublayerConnection(d_model, dropout) def forward(self, x, memory, src_mask, tgt_mask): x = self.sublayer1(x, lambda x: self.self_attn(x, x, x, tgt_mask)) x = self.sublayer2(x, lambda x: self.cross_attn(x, memory, memory, src_mask)) x = self.sublayer3(x, self.feed_forward) return x交叉注意力那行代码值得再强调一遍:self.cross_attn(x, memory, memory, src_mask)里,query 来自 Decoder 侧的 x,但 key 和 value 都来自 Encoder 侧的 memory。这就是机器翻译里 Decoder 在生成每个词时“回看”源语言句子含义的实现方式。
Encoder 和 Decoder 本身还需要管理层的堆叠顺序。一个完整的模型包含词嵌入层(或输入投影层)、N 个 EncoderLayer、N 个 DecoderLayer、输出投影层。每一层输出会直接作为下一层输入,层与层之间不需要额外处理,因为每个 EncoderLayer 内部已经完成了残差和归一化。
3. 实战验证:用 Transformer 预测正弦数据
模型写完之后,找什么任务来验证它是否真正实现了设计目标?我选择的是正弦函数预测,这是 Transformer 作者本人在论文中用来验证模型泛化能力的任务之一,也是社区里很多实现选择的标准验证场景。
3.1 为什么拿正弦函数做验证
正弦函数看起来简单,但对序列模型来说并不简单。它要求模型既记住历史片段的具体数值,又从中推断出周期性规律,再准确外推未来走势。这个任务同时考验了位置编码和自注意力:位置编码必须让模型知道当前处于周期内的哪个阶段,注意力必须让模型找到与当前点最相关的历史参考点。
对比一下为什么不用文本任务来验证。文本需要分词、建词表、padding、embedding 等一堆预处理,而且文本数据的质量依赖于数据集的规模和清洗程度。如果模型效果不好,很难判断是 Transformer 本身实现有问题,还是数据处理环节出了问题。正弦数据就没有这些问题——用 numpy 生成一组数值,做简单的滑动窗口切分,就能直接开训。任何异常都能追溯到模型内部,这对验证一个手写实现来说价值极大。
整个验证过程设计为序列到序列任务:输入过去 60 个时间步的数值,预测未来 30 个时间步的数值。Encoder 读入历史序列,Decoder 逐步生成未来序列。这模拟的是机器翻译里“读入源语言句子,输出目标语言句子”的结构,只是把离散 token 换成了连续数值。
3.2 数据构造与训练参数配置
数据构造的代码很短,但细节决定了任务难度。
import numpy as np import torch from torch.utils.data import DataLoader, TensorDataset # 生成正弦序列 t = np.linspace(0, 30, 1801) x = np.sin(t) # 滑动窗口构造样本 src_len, tgt_len = 60, 30 samples_x, samples_y = [], [] for i in range(0, len(x) - src_len - tgt_len, tgt_len): samples_x.append(x[i:i+src_len]) samples_y.append(x[i+src_len:i+src_len+tgt_len]) samples_x = torch.tensor(np.array(samples_x), dtype=torch.float32).unsqueeze(-1) samples_y = torch.tensor(np.array(samples_y), dtype=torch.float32).unsqueeze(-1) dataset = TensorDataset(samples_x, samples_y) dataloader = DataLoader(dataset, batch_size=64, shuffle=True)这里有一个我踩过的坑:滑窗步长不要等于 1,否则相邻样本高度重叠,验证集和训练集信息混在一起,模型在验证集上看起来性能极好,但实际是“背”了训练集的延展,而不是学到了规律。我用的步长是 tgt_len,也就是每个窗口不重叠,保证训练样本之间有足够的独立性。
训练配置上我推荐一组很稳的初始参数:d_model=64,n_heads=4,num_layers=2,d_ff=256,dropout=0.1。这些参数足够小,普通笔记本电脑的 CPU 都能几分钟跑完一个实验。d_model=64 配 n_heads=4,每个头是 16 维,处于合理区间;层数设为 2 而不是原论文的 6,是因为正弦预测任务复杂度有限,层数太多反而容易过拟合,训练时间也成倍增加。
这里替大家解释一下:Transformer 预测正弦数据时,Embedding 层并不像 NLP 那样使用查表式的 nn.Embedding,因为输入是连续数值而不是离散 token id。需要用一个线性层做输入投影,把单维数值映射到 d_model 维;输出端也用一个线性层把 d_model 维映射回单维预测值。这个适配本身不会影响模型的验证效果。
3.3 训练、评估与推理的完整流程
训练循环整体比较简单,但有两个环节需要特别设计:Teacher Forcing 和推理时的自回归。我先把训练循环写出来,再详细解释这两点。
def train_step(model, src, tgt, optimizer, criterion, device): model.train() optimizer.zero_grad() src = src.to(device) # (batch, src_len, 1) tgt = tgt.to(device) # (batch, tgt_len, 1) # 构造 Decoder 输入:从 tgt 中取前 tgt_len-1 个时间步 tgt_input = tgt[:, :-1, :] # (batch, tgt_len-1, 1) tgt_output = tgt[:, 1:, :] # (batch, tgt_len-1, 1) # 构造 look-ahead mask tgt_len_cur = tgt_input.size(1) tgt_mask = torch.triu(torch.ones(tgt_len_cur, tgt_len_cur), diagonal=1).bool().to(device) output = model(src, tgt_input, src_mask=None, tgt_mask=tgt_mask) loss = criterion(output, tgt_output) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() return loss.item()代码里的核心逻辑是交叉熵砍头尾:Decoder 输入取目标序列的前 29 个点,训练目标取目标序列的后 29 个点。这样每个时间步都是“看到前 29 个点,预测第 30 个点”的学习任务,和机器翻译里“输入前一个词,预测下一个词”完全一致。
推理时完全不同。训练时可以喂真实历史值让模型快速收敛,但推理阶段没有真实值可用,只能把模型上一步的输出拼接到输入末尾,再预测下一步。这就是自回归推理。我遇到的常见问题是,训练效果很好,推理时却越预测越偏,原因是训练时模型习惯了真实值的偏移纠正,没有学会在自身错误之上继续预测。解决方法是训练时引入 scheduled sampling 策略——以一定概率用模型自己的输出替代真实输入——或者通过调整目标序列长度和教师强制比例来缓解。
推理代码不复杂,但需要注意每次预测只取最后一个时间步的输出,然后拼接成一个动态增长的序列。
def predict(model, src, max_len=30): model.eval() with torch.no_grad(): memory = model.encode(src) dec_input = src[:, -1:, :] # 用序列最后一个值作为起点 outputs = [] for _ in range(max_len): tgt_mask = torch.triu(torch.ones(dec_input.size(1), dec_input.size(1)), diagonal=1).bool() out = model.decode(dec_input, memory, tgt_mask=tgt_mask) next_val = out[:, -1:, :] outputs.append(next_val) dec_input = torch.cat([dec_input, next_val], dim=1) return torch.cat(outputs, dim=1)我在实测中跑出来的结果是:训练 2000 个 step 左右 loss 降到很低,预测曲线肉眼可见跟上了正弦的走势,相位略有漂移,但整体形状正确。如果你画出来是一条直线或者振幅缩得很小,优先排查位置编码是否生效、学习率是否太高、数据是否被错误归一化。
4. 训练细节与工程避坑
模型实现好之后,真正的考验其实在训练环节。Transformer 这种结构对训练技巧的敏感度相当高,很多参数差一点就导致训练失败。这一节我把自己踩过的坑和排查经验一次性整理出来。
4.1 Teacher Forcing:训练为什么这样喂数据
Teacher Forcing 是序列生成模型训练的标准做法,但很多人第一次接触时不太理解。简单来说,训练时不需要让模型完全靠自己的预测来生成下一个值,而是把真实的目标序列喂给 Decoder 作为输入,让模型在每个位置学习“如果前面这些值是真实值,我该如何预测下一个值”。这样做的好处是收敛快、训练稳定,坏处是可能导致训练和推理之间的分布不一致,也就是前面提到的误差累积。
在正弦预测任务里,Teacher Forcing 的实现就是把目标序列整体前移一位作为 Decoder 输入。这样模型在训练时始终看到的是干净的真实值,而不是自己上一轮的五错误预测。代码里用切片操作就能完成,不需要循环逐时间步处理。PyTorch 的张量操作天然支持这种并行处理方式,这也是 Transformer 比 RNN 训练效率高的重要原因。
4.2 标签平滑、学习率预热与梯度裁剪
这些技巧单独看都很简单,但每个都在实战中救过我的模型。学习率预热是 Transformer 中最重要的训练技巧之一。因为模型刚初始化时,各层的参数都还比较随机,尤其是多头注意力里的线形层,如果在训练一开始就用大的学习率,容易把注意力分数的数值分布推到极端饱和区,模型就很难恢复。预热阶段先用很小学习率让模型在“平坦区域”走稳,再提高到预设的学习率最大值,最后再逐步衰减。
梯度裁剪是我强烈建议加上的操作。一个 step 就加一行代码:torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)。手写实现时很容易忽略,但实践里正是这一步防止了训练中 loss 偶发性的巨大跳变。
标签平滑这个技巧适合分类任务,正弦预测是回归任务用 MSE 损失,不需要标签平滑,但如果后续迁移到文本分类,标签平滑能明显改善模型的泛化能力。标准做法是把 one-hot 向量中的 0 改为一个小量 0.05,再把正确标签位置改为 0.9,防止模型过度自信、输出过于极端的概率分布。
4.3 常见问题速查表
我把实战中遇到的高频问题整理成了一张速查表,按照“现象—原因—解决办法”的顺序排列,方便读者照着排查。
| 现象 | 可能原因 | 解决办法 |
|---|---|---|
| loss 不下降或下降极慢 | 学习率过小 / 位置编码未生效 | 调大学习率到 1e-3 量级;检查位置编码是否加到输入上 |
| loss 很快降到 0,但验证集效果极差 | 数据滑动窗口重叠 / 过拟合 | 增大滑窗步长;降低模型层数或增大 dropout |
| 注意力权重几乎均匀分布 | 训练不充分 / 头数过多 | 增加训练步数;减少 n_heads 到 2 或 3 |
| 预测曲线相位漂移明显 | Teacher Forcing 导致的误差累积 | 推理时用自回归;训练时加入计划采样 |
| loss 出现 NaN | 学习率过高 / 数据未归一化 | 降低学习率;对输入数据做标准化 |
| 多层堆叠训练极不稳定 | Post-LN 结构问题 | 换成 Pre-LN 即先归一化再进子层 |
还有一个容易忽略的细节:训练和推理时的 mask 形状要完全一致。Encoder 的 padding mask 形状是(batch, 1, 1, src_len),Decoder 的 look-ahead mask 形状是(1, 1, tgt_len, tgt_len)。如果手写实现时把两者的维度搞乱了,PyTorch 不会直接报错,但广播机制会导致 mask 作用在错误的维度上,注意力分数被遮错位置,模型输出几乎随机。这种问题极难排查,我的经验是,怀疑 mask 时直接打印 scores 的数值分布,看有没有大量极小的负数值,有的话说明 mask 生效,没有的话就说明 mask 根本没有起到遮掩作用。
5. 扩展方向:从正弦预测到文本分类
验证完正弦预测之后,你已经拥有了一套完整的 Encoder-Decoder 模型。但自然语言处理领域里,更多任务并不需要同时使用 Encoder 和 Decoder,而是只用其中的一部分。
5.1 架构调整:序列分类只需 Encoder 塔
文本分类、情感分析、意图识别这类任务,核心需求是理解整个文本的含义,然后输出一个类别标签。这类任务使用 Encoder 单独工作就足够了。具体做法是:将文本分词、映射为 token id、经过嵌入层和位置编码、通过多层 Encoder 层编码得到每个 token 的隐向量,然后需要把所有 token 的信息聚合为一个向量。
聚合方式有两种主流选择:一种是在序列开头加入一个特殊的 CLS token,训练的时候让这个 token 对应位置的输出学习到整个序列的语义摘要,分类时直接取 CLS 位置的输出接分类头;另一种是全局池化,把整个序列所有位置的隐向量取平均或取最大值,再接分类头。BERT 用的就是前者,很多实际项目为了省事用后者,效果差异往往不大。
架构上只需要给 Encoder 的输出加一个线性分类头就完成了。如果我们继续复用第 2 节实现的代码,会发现整个模型从原来的 Encoder-Decoder 六层结构简化成了一层 Encoder 接一层线性层,代码量反而更少了。
5.2 文本分类的关键变化点
从正弦预测迁移到文本分类,有三个关键变化点值得单独拿出来讲。
第一,输入处理从数值序列变成了离散 token 序列。这意味着需要建立一个词表,把文本拆分并映射成整数 id,模型里要把原来的线性输入投影层改成 nn.Embedding。这个词表的大小需要根据数据集规模来决定,一般英文任务用 30000 到 50000 个词就足够覆盖绝大多数需求,中文则需要更大的词表或者干脆按字建表。
第二,损失函数从均方误差变成了交叉熵。正弦预测是回归任务,输出层是单个神经元;分类任务则要用 nn.Linear(d_model, num_classes) 输出类别的 logits,再配合 CrossEntropyLoss。
第三,训练推理机制完全不同。文本分类不需要自回归生成,不需要 Teacher Forcing,也不需要逐时间步预测。模型一次前向就能得到整体的分类概率分布,训练速度比序列到序列任务快得多。
我在做分类迁移实验时的一个经验是:先用正弦预测的权重初始化新模型的 Encoder 部分,然后再用文本分类数据微调。虽然正弦数据和文本数据分布差异极大,但这种预训练加微调的方式确实能加快收敛速度。原因是正弦预测让模型学会了通用的相对位置感知能力和注意力模式的组合方式,这些能力本身和具体数据类型无关。
5.3 这些方向也值得关注
Transformer 的扩展已经远远超出了自然语言处理的范围,最近几年在遥感影像语义分割、交通流量预测等领域也出现了大量 Transformer 变体。比如有些遥感领域的模型把多模态融合和 Transformer 结合,解决不同类型传感器数据的对齐问题;交通领域的 PDFormer 则关注传播延迟感知的动态长时序建模,本质上是在时间维度上引入注意力机制。这些方向虽然和 NLP 细节不同,但核心的 QKV 注意力框架一脉相承。
对我们学习者来说,这意味着手写一个最小实现的价值被放大了。因为你一旦彻底掌握了这个基础框架,再去看 Swin Transformer、Vision Transformer 这些改进了空间建模方式的变体时,会发现它们的核心依然是 QKV、位置编码、前馈网络这几板斧,只是根据领域特性调整了输入组织方式和位置编码形式。先把这个最小实现写通,后续的学习曲线会平滑非常多。
我个人做了这么多年模型实现和调优,最大的体会是:读十篇论文不如手写一遍代码。很多细节只有在写的时候才会真正暴露出来——mask 的维度问题、位置编码的数值范围问题、Teacher Forcing 的误差累积问题,所有这些在你只是一遍遍看论文图解时是不可能意识到的。这个最小实现就是你深入 Transformer 世界的一根拐杖,建议在此基础上多做实验:改一改 d_model 看效果如何变化、在训练时加入计划采样、把 Post-LN 换成 Pre-LN,每一次改动都会加深你对这个模型的理解。最后再分享一个小技巧:把训练过程中各层的注意力权重保存下来做可视化,你会直观地看到不同头在不同的层里关注模式完全不同,那一刻你会觉得自己确实把这个模型“拿捏”了。