用PyTorch从零实现注意力机制与Transformer文本分类

这个系列总算写到第六篇了。前五篇我们把深度神经网络这条路从头走了一遍：感知机怎么变成多层网络，反向传播怎么让损失一路传回去，激活函数和优化器各自解决什么问题，再到图像领域的卷积网络和序列领域的循环网络。按这个进度，这一篇该补上我认为整个深度学习版图里最关键的一块拼图了——注意力机制和Transformer架构。文章还是延续之前的节奏，用Python和PyTorch纯手工把核心模块搭一遍，然后落在一个实际任务上：用Transformer做文本情感分类。

先说明一下，本文默认你已经知道基本的张量操作、梯度下降和PyTorch的数据流，这些在前几篇里都写过。不会的地方往回翻一翻就好。

1. 为什么深度学习绕不开注意力机制

1.1 从人的注意力说起

假设你在看一张全家福，你不会把每个像素都看得一样仔细。人脸、表情、谁站在中间，这些信息会天然占据你大量的注意力。等你要描述这张照片的时候，你说出来的内容也基本来自这些重点区域。

深度学习里的注意力机制，本质就是把这种“选择性聚焦”变成一种可学习的计算。网络不再被迫对输入的所有部分一视同仁，而是根据当前的任务需求，自己决定哪些位置更重要、该给多大的权重。这个思想在Transformer出现之前其实就已经存在了，比如早期的机器翻译模型里就已经有人尝试把注意力当成RNN的辅助增强模块。真正让它改变整个行业格局的，是2017年那篇提出Transformer架构的论文。从那之后，图像、语音、文本、推荐系统，几乎所有领域的SOTA模型都开始往注意力靠拢。

1.2 RNN的致命短板

在Transformer之前，处理序列数据的常规方案是RNN或者它的进化版LSTM、GRU。它们的核心逻辑是沿着时间步逐步读取输入，用一个隐状态向量把“记忆”从头带到尾。

这个思路直觉上是对的，但实际用起来有两个很痛的短板。

第一个是长距离依赖问题。序列一长，信息要从第1个位置传到第50个位置，中间每走一步都要经过一次非线性变换和矩阵运算。信号在逐层传递中很容易衰减，就像一群人站成一排传话，传到队尾往往只剩一点模糊的印象。LSTM虽然通过门控机制缓解了梯度消失，但本质上还是没解决信息衰减的问题。

第二个是串行计算。RNN必须等第t步算完才能算第t+1步，因为隐状态是逐步更新的。这意味着GPU并行计算的优势完全发挥不出来。在数据量和模型规模都爆炸的今天，这个限制非常致命。

我当年用LSTM跑长文本分类任务的时候，训练一个中规模模型基本要按“天”来计时，而且文本稍微长一点，效果就明显往下掉。这些痛点不是靠调参能解决的，必须换一种更激进的处理方式。

1.3 Transformer的核心思路

Transformer做了一个看起来很简单的决定：去掉循环结构，让序列里的每一个位置，直接和所有其他位置建立联系。每个词在编码的时候，都会对整句话里的其他词做一次“相关度打分”，然后把相关度当成权重，把所有词的信息按权重融合在一起。

这就是自注意力机制的基本形态。它有两个天然优势：

• 并行性极强。每个位置的计算互不依赖，可以一次性算完，GPU能完全吃饱。
• 信息通路短。任何两个位置之间只隔了一次计算，不存在逐步传递的衰减问题。长距离依赖被直接拉平了。

代价是计算复杂度变高了。序列长度为n的时候，自注意力需要对n×n的相互关系打分，复杂度是O(n²)。这个问题的优化是后话，对入门来说先掌握基础形态更重要。

1.4 Transformer宏观架构速览

完整的Transformer是一个编码器-解码器结构，最初为机器翻译设计。编码器负责把输入序列转换成连续表示，解码器负责根据这个表示逐步生成输出序列。

不过在实际任务里，绝大多数的需求只需要其中一个部分。比如文本分类、语义匹配、情感分析这种任务，只用编码器就够了；文本生成类任务才需要解码器。本文的场景是情感分类，所以我后面会把重点放在编码器部分，让你逐步搭出一个可用的Encoder。

有一点值得注意：Transformer这个名词最容易被误解的地方在于，它不是一个“模型”，而是一套组件组合的框架。你真正要理解的，是里面那几个核心设计各自解决什么问题。接下来这一章，我们把最底层、也最核心的自注意力机制彻底拆开。

2. 自注意力机制的数学直觉与代码实现

2.1 先搞清楚输入是什么

自注意力机制接收的输入通常是三维张量，形状是[batch_size, seq_len, d_model]。batch_size是当前批次里的样本数，seq_len是序列长度，d_model是这个模型内部的向量维度。

一个句子进来以后，首先要通过词表索引查到一个词对应的嵌入向量，形状为[seq_len, d_model]。这个过程可以理解为给每个词一个“高维语义坐标”。在这个坐标空间里，语义越接近的词，向量的方向往往也越接近。

实际处理时，我们往往还会给序列补上一些特殊符号，比如句子不够长就补到固定长度，这个后面说数据预处理的时候再展开。

2.2 Query、Key、Value：更像一个搜索引擎

自注意力里最劝退新手的，就是 Query、Key、Value 这三个概念。其实没必要想得多玄，它和你平时用搜索引擎的流程一模一样。

你搜“如何选显卡”，搜索引擎会先把你输入的这句话变成一个查询向量（Query），然后用这个查询向量去和数据库里每篇文章的标题、关键词（Key）做相关度匹配。匹配结果就是分数，分数越高说明这些文章（Value）越可能包含你要的东西。最终返回的结果是每篇文章的内容按相关度加权组合后的答案。

在自注意力里面，每个位置都会同时扮演三种角色：

• Query（查询）：它代表“我现在想找什么信息”。
• Key（索引）：它代表“我身上有哪些信息可以被别人找到”。
• Value（内容）：它代表“我实际能提供的信息内容”。

三个角色不是凭空来的，而是对原始输入做三次不同的线性变换得到的：

Q = X W_Q K = X W_K V = X W_V

其中W_Q、W_K、W_V都是可学习的权重矩阵，形状都是[d_model, d_k]。d_k是Query和Key的维度，通常比 d_model 小一些，用来控制计算规模。

值得注意的一个设计细节：如果直接拿原始向量做匹配，所有位置共用一份表示，匹配能力和内容表达能力会互相干扰。拆成三个分支以后，网络可以针对“如何被匹配”和“应该提供什么信息”分别学习各自的参数，表达能力更强。

2.3 注意力分数的计算

自注意力的第一步，是计算任意两个位置之间的相关度。这里用到的公式是缩放点积注意力：

[ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V ]

先把 Q 和 K 的转置做矩阵乘法。这一步的物理含义是：第 i 个位置的 Query 和第 j 个位置的 Key 做点积，点积结果越大，说明这两个位置在语义上越相关。所有位置两两计算一遍，得到一个形状为[seq_len, seq_len]的相似度矩阵，我习惯把它叫注意力矩阵。

举个例子，句子“我喜欢这只可爱的猫”。编码“猫”的时候，注意力矩阵里，“猫”这一行和“喜欢”对应的分数可能比较高，和“可爱”也有一定关联。这说明模型在处理“猫”的时候，会优先参考“喜欢”和“可爱”这两个词的信息。

2.4 为什么要除以根号d_k

这是新手最容易忽略、但很重要的一个细节。当d_k比较大的时候，点积的结果区间也会随之变大。想象一下 d_k 等于128，假设 Q 和 K 里面的元素都是均值为0方差为1的随机变量，那么它们的点积期望是0，但方差是128，标准差超过11。这意味着很多点积分数会落在 -30 到 +30 甚至更大的范围里。

softmax 对输入非常敏感。输入一旦绝对值很大，softmax 算出来的概率会变得极度尖锐——最大的那个值直接趋近于1，其他值趋近于0。这样梯度会变得非常小，训练的时候几乎学不动。

所以除以 (\sqrt{d_k}) 的本质，是做一个分布修正，让点积结果的方差回到1左右，softmax入口的分布不至于过于极端。这也解释了为什么缩放因子不用别的数，而刚好是根号d_k。

2.5 Softmax归一化与加权融合

拿到相似度矩阵之后，对每一行做softmax，让每一行的所有分数加起来等于1。这样，每个位置对整句话其他位置的关注程度就是一组概率分布。

最后一步是把归一化后的权重和 Value 矩阵相乘，做加权求和。每个位置最终得到的新向量，相当于是“我参考了所有其他位置的信息，但参考多少由注意力权重决定”。这个过程完成之后，序列里的每个向量都带着全句的上下文信息。

2.6 单头注意力的PyTorch实现

把上面的逻辑写进PyTorch，核心就几十行：

import torch import torch.nn as nn import torch.nn.functional as F class ScaledDotProductAttention(nn.Module): def __init__(self, d_k): super().__init__() self.d_k = d_k def forward(self, q, k, v, mask=None): # q, k, v: [batch, seq_len, d_k] scores = torch.matmul(q, k.transpose(-2, -1)) scores = scores / (self.d_k ** 0.5) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) weights = F.softmax(scores, dim=-1) out = torch.matmul(weights, v) return out, weights

注意mask参数。在真实任务里，序列往往会被填充到固定长度，填充位是没有实际语义的空白，不能让它们参与注意力计算。做法是把填充位置的注意力分数设成一个极大的负数，这样softmax之后它们的权重会趋于0。这个细节到实际任务的章节还会再遇到。

3. 多头注意力与Transformer关键细节

3.1 多头机制到底在干吗

单头自注意力只能学习一种“相关度”的定义方式。但语言是很复杂的：一个词和前面的词可能存在语法关系、语义关联、指代关系等多种不同维度的联系。单一的相关度评分很难面面俱到。

多头注意力的做法，是不再让所有头共用一组 Q/K/V 变换，而是把 d_model 拆成 n_heads 个子空间，每个头用自己独立的一组线性变换，并行地做自注意力。这样每个头可以学到不同的关注模式。有的头可能更关注语法关系，有的头更关注指代消解，有的头可能关注位置远近。最后把多个头的结果拼回去，再通过一个线性层融合。

一个经常被问的问题：如果有多个头，计算量会不会翻很多倍？答案是不会。因为每个头的维度是 (d_model / n_heads)，所有头的总维度和单头的时候保持一致。比如 d_model 是512，8个头，每个头的维度就是64。虽然注意力计算有n_heads份，但每份矩阵乘法规模更小，总体计算量和单头版基本持平。

多头注意力完整实现如下：

class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads): super().__init__() assert d_model % n_heads == 0, "d_model必须能被n_heads整除" 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.out_proj = nn.Linear(d_model, d_model) def forward(self, x, mask=None): batch, seq_len, d_model = x.shape q = self.w_q(x).view(batch, seq_len, self.n_heads, self.d_k) k = self.w_k(x).view(batch, seq_len, self.n_heads, self.d_k) v = self.w_v(x).view(batch, seq_len, self.n_heads, self.d_k) q = q.transpose(1, 2) k = k.transpose(1, 2) v = v.transpose(1, 2) # 现在 q/k/v 形状为 [batch, n_heads, seq_len, d_k] attn, _ = ScaledDotProductAttention(self.d_k)(q, k, v, mask) attn = attn.transpose(1, 2).contiguous().view(batch, seq_len, d_model) return self.out_proj(attn)

这里的关键操作是view之后再transpose。把最后两维拆成[n_heads, d_k]，然后通过转置把头的维度挪到batch后面，方便后续并行计算。每一步操作前最好都打印一下张量形状，我在调试这种重排代码的时候经常靠这个避免维度错乱。

3.2 位置编码：给序列加上一把坐标尺

Transformer有一个天然缺陷：由于所有位置是并行计算的，模型根本不知道第1个词和第5个词谁在前谁在后。你把“猫追老鼠”和“老鼠追猫”两句话按相同的词顺序打乱后喂进来，模型看到的内容完全一样。没有位置信息，这就麻烦了。

解决方式是给每个位置的向量加上一份“位置标记”。Transformer论文里用的是三角函数位置编码：

import math 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).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)]

你在初学阶段可以把这套三角函数当成一种“精心设计的固定浮标”，不用过度纠结。更直接的替代方案是用nn.Embedding训练一组可学习的位置向量，效果在很多任务上也差不多。使用register_buffer是因为位置编码不参与更新，但需要跟着模型一起搬到GPU上。

实操提示：max_len需要大于训练数据的最大序列长度。如果测试阶段出现了超过max_len的样本，后续位置没有编码，等于直接送进模型一堆没有坐标的向量，效果会断崖式下跌。

3.3 残差连接与层归一化

很多Transformer实现里，注意力子层和前馈网络子层外面都会包一层“虚线连接”。这个“虚线”就是残差连接。原理很简单：让梯度有一条从深层直接传回浅层的通路，避免网络太深之后梯度在逐层传递中消失。

另一个必备组件是层归一化。Transformer里比较惯用的做法是把层归一化放在每个子层最前面，这种结构叫Pre-LN。好处是训练更稳定，对学习率的敏感度低不少。后加一层Dropout防止过拟合。结合起来就是一个完整的Transformer编码器块：

class TransformerEncoderBlock(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout=0.1): super().__init__() self.attn = MultiHeadAttention(d_model, n_heads) self.ffn = nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model), ) 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 = x + self.dropout1(self.attn(self.norm1(x), mask)) x = x + self.dropout2(self.ffn(self.norm2(x))) return x

3.4 前馈网络：给模型提供非线性变换

注意力层做的事情本质上是线性加权和。如果整个模型只有注意力机制，不管叠多少层，组合起来仍然是一堆线性变换的嵌套，表达能力严重受限。Transformer在注意力层后面接了一个两层全连接网络，第一层把维度从 d_model 放大到 d_ff（通常是4倍），经过ReLU激活之后，第二层再投影回 d_model 维度。

可以把它理解成“每个位置独立的一个小大脑”。注意力层负责从全句收集信息，前馈网络负责对这些收集来的信息做一步深加工。两者交替堆叠，构成了Encoder的基本循环：交换信息、局部加工、再交换、再加工。

Encoder块可以叠加N层，每一层的输出又作为下一层的输入。浅层可能学到词的局部组合信息，深层就能接触到越来越全局的语义结构。任务越复杂，需要的层数往往越多。

这个Encoder Block已经是Transformer编码器的完整积木了。接下来我用它搭一个能解决实际问题的分类器。

4. 实际任务：用Transformer做文本情感二分类

4.1 任务设定与数据准备

我选的任务是电影评论情感分类：输入一段短评，输出正面或负面两类结果。这个任务足够简单，模型训练不会太久，同时又能把前面讲的所有组件串起来，非常适合作为Transformer的入门实战。

数据来源不限定，网上有很多开放的电影评论数据集。我这里为了快速演示，准备一份简易的小规模样本进行说明，格式长这样：

这部电影看得我热血沸腾，情节紧凑。 1 剧情拖沓，演员演技一言难尽。 0

后面那个数字就是标签，1是正面，0是负面。实际操作中请使用规模更大的标准数据集来训练。

预处理流程大概是：切词 -> 构建词表 -> 把句子转成ID序列 -> 统一填充到固定长度。

切词这步对于中文文本来说，最简单的方式是按单字切分，因为很多词的边界不好判断。字级别模型在短文本分类里效果也不错。具体操作时我会先用一个简单的分词函数切出所有字，统计频次，保留出现次数大于2的字符构建词表。接着给词表加上几个特殊标记：<pad>用来填充、<unk>用来代表词表外的字。

4.2 把原始文本转成模型输入

模型最后的输入是一个整数矩阵，形状是[batch_size, max_len]。每个整数对应词表里的一个位置。同时为了处理padding，还要构造一个注意力mask，把填充的位置标为0，真实文本位置标为1。

这里有一个新手很容易踩的坑：padding 用的ID不能是0，而必须是词表里第一个特殊符号<pad>对应的ID。PyTorch的nn.Embedding有一个padding_idx参数，指定之后这个位置并不会被更新，方便统一处理。

4.3 模型搭建

用前面实现的组件，我现在把它们拼成一个完整的分类器：

class TextClassifier(nn.Module): def __init__(self, vocab_size, d_model, n_heads, n_layers, d_ff, num_classes, max_len, dropout=0.1): super().__init__() self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=0) self.pos_encoding = PositionalEncoding(d_model, max_len) self.encoder_layers = nn.ModuleList([ TransformerEncoderBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_layers) ]) self.classifier = nn.Linear(d_model, num_classes) def forward(self, input_ids, mask=None): x = self.embedding(input_ids) x = self.pos_encoding(x) for layer in self.encoder_layers: x = layer(x, mask) # 全局平均池化，把 [batch, seq_len, d_model] 压成 [batch, d_model] x = x.mean(dim=1) return self.classifier(x)

为什么用全局平均池化而不是直接用最后一个位置的输出？因为在分类任务中，我们并不关心信息具体出现在哪个位置，而是希望综合全句信息做判断。平均池化把整个序列压缩成一个向量，再用它做分类。

一个迷你Transformer的基本配置大概是：d_model=128，n_heads=8，n_layers=2，d_ff=512。这个规模在小数据集上已经够用。

4.4 训练循环与关键参数

训练循环和普通PyTorch模型没有太大区别，但有几个参数需要特别注意：

optimizer = torch.optim.Adam(model.parameters(), lr=2e-4) criterion = nn.CrossEntropyLoss() for epoch in range(epochs): model.train() for batch in train_loader: input_ids, labels = batch input_ids = input_ids.to(device) labels = labels.to(device) mask = (input_ids != 0).unsqueeze(1).unsqueeze(2) # mask 形状 [batch, 1, 1, seq_len]，广播后匹配 [batch, n_heads, seq_len, seq_len] logits = model(input_ids, mask) loss = criterion(logits, labels) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step()

做两个特别说明。

第一个是学习率。Transformer对学习率非常敏感。我用Adam优化器时，经验值通常落在 2e-4 到 5e-4 之间。太高容易loss直接崩掉，太低又收敛很慢。如果训练初期loss持续震荡不下降，先别急着改结构，把学习率调低一个量级试试。

第二个是梯度裁剪。clip_grad_norm_做的事情是限制梯度的最大范数。Transformer模型在训练初期偶尔会出现梯度暴涨，这个操作能有效防止loss变成NaN。虽然大部分时候用不上，但加上之后训练稳定性明显提升，属于“加了不亏”的防御性操作。

4.5 实验效果观察

在小规模数据集上，这个模型大概训练几轮就能跑到85%以上的准确率。如果你不断把语料增大，效果还会继续往上走。我建议你训练之后做两件事来检验模型是否真正学到了东西：

第一，保留模型每层的注意力权重，挑几条样本可视化注意力矩阵。你会看到某些head确实能捕捉到形容词和中心词之间的对应关系。第二，构造一个稍微复杂的测试集，比如带转折的句子“电影前半段很无聊，但是结局让我感动”，看看模型能不能识别出这种复杂情绪。Transformer通常表现得不错，因为每个词都能直接感知整句话的语境。

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

写代码练习Transformer的时候，有几类坑我基本每次都会遇到。这里整理成速查表，按出现频率排序：

5.1 维度错误

按我的经验，80%的维度错误出在多头注意力的reshape环节。一个比较稳的调试方法是：在forward内部每个矩阵运算前后都加一行print(x.shape)，跑一个batch看一遍，确认张量在每一个环节的形状都符合预期。新手阶段不要嫌麻烦，这么做能节省大量排查时间。

5.2 mask维度对不上

注意力mask的形状经常写错。masked_fill要求mask的形状可以广播到scores的形状[batch, n_heads, seq_len, seq_len]。标准做法是构建形状为[batch, 1, 1, seq_len]的mask，然后在最后一个维度做广播。

另外要注意填充ID必须是整数0对应的位置，如果文本里的字恰好也是0号索引，需要检查一下词表和padding标记的顺序是否搞对。

5.3 训练慢且效果差

很多情况下不是模型写错了，而是预处理出了问题。比如句子没有统一长度的策略，导致batch内部长短差异极大，Transformer的序列长度会被拉得很长，计算量跟着暴涨，但有效信息很少。遇到这种情况，可以先统计一下数据长度分布，选取一个覆盖绝大多数样本的截断长度，把极长样本做截断处理。

5.4 动手扩展建议

如果上面的基础版本跑通了，有几个值得一试的扩展方向：

• 把单层Encoder加深，观察深层模型在小数据集上会不会过拟合。
• 在分类任务里加上不同数量的head，比较收敛速度和最终效果。
• 把mean池化换成max池化或者直接取第一个位置的输出，比较差异。
• 训练时逐步加大batch，配合学习率warmup策略，观察loss曲线变化。

这些改动每个都能帮你加深对Transformer内部机制的理解，比只照着跑一遍代码要有效得多。

我个人在实际操作中的体会是：Transformer看起来组件很多，但核心就是“加权求和”四个字。你真正把自注意力那条数学式子吃透了，后面什么BERT、GPT的变体，本质上都是在回答同一个问题——怎么更好地决定该往哪看。建议你把这个Encoder块多搭几遍，从零写、不看参考代码、跑通一个分类任务，然后去调一调head数、层数、dropout，对比结果。等你能根据任务直觉地判断该往哪个方向调参，这一章的东西才算真正变成你自己的功底。