资讯动态

自注意力机制与Transformer文本分类实战:从原理到PyTorch实现

发布时间:2026/10/14 10:41:04 来源:尧图企业网站定制
做深度学习教程做到第六篇很多人其实已经能跑通模型也能调一些参数但一旦要自己设计一个稍微复杂一点的神经网络还是会卡在“结构对不对”和“为什么这样设计”这两个核心问题上。这篇我来把这两个问题一起解决。结合系列标题里的关键词这第六篇主要做两件事一是把深度神经网络中最容易被当成黑盒的“自注意力机制”从底层数学和代码两个层面拆开二是基于 Python 和 torch 生态把原理落成一套能真正在实际任务里跑出效果的小型文本分类系统让你既能看懂它为什么工作也能直接照着改。这一篇承接前面的完整课程链路。前面我们已经解决了数据加载、反向传播、梯度下降、CNN、RNN 这些基础所以这里默认读者已经会搭一个简单 MLP也见过最基本的 Dataset 和训练循环逻辑。如果你的基础还没到这个程度建议先回看系列前几篇否则直接看注意力机制的数学会比较痛苦。如果你属于那种已经会调库、但没深究过每行代码含义的读者这篇的定位是帮你在思维上“解锁”从“我知道这里有个注意力层”变成“我能手写这个注意力层”。我给这篇文章设置的路线是先讲设计思路再逐层拆原理最后落到代码实操和排查技巧。所有代码都是可运行的最小实现不建议直接上大规模数据集因为这篇的重点不是刷分而是把复杂机制吃透。1. 内容整体设计与思路拆解1.1 为什么第六篇要选“自注意力”这个切口如果你回顾一下前面几篇会发现卷积网络的核心假设是局部性循环网络的核心假设是顺序性它们的结构里都带了一个很强的先验。自注意力机制则反过来它一上来就不对位置做任何假设只通过“两两之间算相关性”的方式来重构信息。这个转折彻底改变了我们设计深度网络的方式也是后来很多大模型的基本组件。把它作为“理解深度神经网络工作原理”的最终收口我认为比再讲一个具体网络更划算你弄懂了自注意力就差不多弄懂了今天主流的 Transformer 架构你弄懂了 Transformer再看各种带注意力变体的网络都会轻松很多。在第六篇选这个主题还有一个很实际的原因它帮你把前面的知识串成一条线。你需要用到矩阵乘法需要理解为什么 softmax 里的数值会饱和需要做残差连接和 LayerNorm甚至要做 Mask 来屏蔽无效位置——这些恰好都是系列前几篇反复强调过的细节。所以这是一篇“集大成”的文章适合你在掌握了基础训练技巧之后来读。1.2 用一个小型分类任务串起原理到落地我特意没有选图像任务或大规模语言模型任务来做例子因为小数据集才能把每一层的变化看明白。这篇我会用一个中文短文本二分类任务贯穿全文给出一段产品评价判断它倾向好评还是差评。数据量很小几百条就够演示训练目标也不追求刷分而是让你能清楚看到输入经过自注意力之后哪些位置被赋予了高权重最终的分类结果是如何被这些权重影响的。这个任务虽然小但它具备实际任务的所有关键环节。你需要做文本清洗、构建词表、把句子填充到统一长度、生成 padding mask、构造 batch、定义模型、训练、评估。整个流程和做大规模 NLP 项目是一样的只是网络规模小到可以在 CPU 上几分钟跑完。这样你在复现的时候不会因为资源不够而卡住也能在调试过程中直接打印每一个中间张量的形状。1.3 技术栈与准备工作环境上不需要太高配置。Python 3.9 以上装上 torch 生态的基础库再加一个分词工具就够。如果你的 Python 环境还没有这些依赖可以直接用如下命令补一下pip install python-dotenv tqdm torch tokenizers我这里之所以强调“torch 生态”而不是某一个具体版本是因为后面代码中用到的主要是动态图、自动求导、张量运算这些通用能力版本差异对示例的影响非常小。当然不同版本的接口细节会有差异比如多头注意力模块在不同版本里对batch_first参数的处理略有不同真遇到报错就按报错提示改一下。准备工作里还有一件容易被忽略的事准备好一个可复现的随机种子。深度学习实验如果不固定随机种子同一份代码两次训练的结果可能差不少这会影响我们对“模型有没有真的学对”的判断。我通常会在脚本开头固定随机种子这样调试时比较安心。2. 吃透自注意力的四个底层细节2.1 Q、K、V 各自承担什么信息角色自注意力里最迷的三个字母就是 Q、K、V。我第一次学的时候也一直在想注意力为什么非要分成三个角色后来我用一个生活化的类比才真正理顺你可以把一段文本里的每个 token 想象成一个员工。Q 是这位员工向别人提出的需求标签“我在找什么”K 是这位员工对外展示的标签“我提供什么”V 则是这位员工真正贡献出来的信息内容。注意力计算的过程就是先让所有人的 Q 去匹配所有人的 K匹配度高的组合就被分配更高的权重最后按照这个权重把 V 汇总起来。用这样的方式理解自注意力的本质就是让每个位置在全局范围内重新收集信息。比如在“这家店的牛肉面很好吃”这句话里“牛肉面”这个 token 的 Q 会强烈匹配到“好吃”的 K于是“好吃”对应的 V 就被加到“牛肉面”这个位置的输出表示中形成一个更丰富的语义表示。这个过程是动态的同一个词在不同句子里学到的东西不一样因为上下文不一样。在代码里这一层通常被实现成两个线性变换先让输入经过三个线性层得到 Q、K、V然后把它们投入缩放点积注意力中。实现的时候要注意这三个线性层初始时如果没有特殊处理计算结果不会立刻呈现我们刚才说的语义匹配需要经过训练慢慢学到这个匹配关系。2.2 缩放点积背后的数学直觉自注意力的分数计算公式几乎每个教程都会写Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V很多人背下这个公式却没搞懂为什么要除以根号下的d_k。这里面的数学直觉其实很直接如果让 Q 和 K 的每个维度都服从均值 0、方差 1 的分布那么两个向量的点积的期望是 0但方差会随着维度d_k的增大而增大等于d_k。也就是说维度越高点积的绝对值可能越大进入 softmax 之后概率分布会越来越尖锐——大的分数被放大小的分数被压到接近 0对应的梯度就会很快消失。除以d_k的平方根本质上是把点积的方差重新拉回 1 左右让 softmax 的输入落在梯度比较明显的区间。这个操作不引入额外学习参数计算成本几乎为零所以它的收益是纯免费的数值稳定性更好训练更容易收敛。你再看到“为什么要缩放”这个问题时就可以直接理解为“给 QK^T 的结果做一次尺度归一”和 BatchNorm 里把输入拉到合理尺度是同一个思路。还有一个工程细节计算 QK^T 的时候Q 和 K 都是带了 batch 维和 seq 维的三维张量所以要用矩阵乘法而不要做逐样本循环。矩阵乘法的并行化程度高得多在实际任务里跑起来会有明显差距。2.3 多头拆分与维度压缩多头注意力是我觉得自注意力中最实用也最容易被忽视的设计。单个注意力头只能捕获一种相关性模式比如只关注位置相邻的词或者只关注主语和谓语的关系。多头的思路是并行做多组注意力计算每组都在一个低维子空间里做最后再把结果合并这样同样的计算开销下模型能同时捕获多种类型的关系。实现多头的时候常见的做法不是真的为每个头建立独立的 Q、K、V 层而是先用一个大的线性层把d_model投影成num_heads * d_k再通过 reshape 把它拆成多个头。比如d_model128num_heads4那么每个头的维度就是32。这种做法的好处是用一套参数完成整体投影训练时的矩阵运算高效而且头之间可以共享底层特征。有一个特别容易错的地方多头输出合并之后顺序不能搞错。我的经验是把“头”这一维放到 batch 维之后、seq 维之前变成(batch, num_heads, seq_len, d_k)做完 attention 之后再 reshape 回(batch, seq_len, num_heads * d_k)。如果你把头维放在最后虽然某些操作也能运行但形状不直观后续接 LayerNorm 和全连接层时容易埋下隐患。2.4 位置编码把顺序信息“焊死”进向量自注意力对位置完全不敏感同一个序列正着排和倒着排在注意力机制眼里没有任何区别。所以在 Transformer 架构里位置信息是以“编码”的方式额外加到输入向量上的。原始论文用正弦余弦函数生成位置编码它的好处是可以外推更长的序列但在实际小任务里我更推荐直接使用可学习的位置参数。可学习位置编码的生成逻辑非常简单初始化一个nn.Embedding(max_len, d_model)每次前向计算时根据当前句子的实际长度取前max_len个位置向量加到输入向量上。这个 Embedding 会随着训练一起更新所以模型能自己调整不同位置的重要程度。需要提醒的是位置编码加进去之后输入向量的数值分布会发生变化。如果你的网络后面马上接一个对输入范围很敏感的模块建议在加位置编码之前先做一次 LayerNorm或者对输入向量做适度缩放否则训练初期容易震荡。这不是什么大问题但排查起来很隐蔽。3. 搭一个能跑通的小型 Transformer 并把它用在实际任务上3.1 从原始文本到批量张量实际任务的第一步是怎么把文本变成模型能吃的张量。我通常的做法是先给原始文本做分词建立词表然后给每个 token 分配一个 ID接着把同一批句子填充到一样长。填充方式我用的是右侧补 0也就是在句子末尾补 padding token。这里有一个实操要点padding 之后必须记录哪些位置是真实 token哪些位置是 padding。这个记录就是后面要用的 attention mask。注意 mask 的逻辑在不同框架中可能是反的。以我常用的代码习惯为例有些模块的key_padding_mask里True表示“这个位置是 padding不要参与注意力”所以很多人习惯写成mask(input_ids 0)来生成这个 bool 张量这是对的。我建议写一个足够通用的 datasetimport torch from torch.utils.data import Dataset class TextClassificationDataset(Dataset): def __init__(self, texts, labels, token_to_id, max_len): self.texts texts self.labels labels self.token_to_id token_to_id self.max_len max_len def __len__(self): return len(self.texts) def __getitem__(self, idx): tokens list(self.texts[idx])[: self.max_len - 2] ids [self.token_to_id.get([CLS], 0)] [self.token_to_id.get(t, 1) for t in tokens] [self.token_to_id.get([SEP], 0)] ids ids[: self.max_len] # 右侧补 padding token假设 id0 是 padding mask [1] * len(ids) [0] * (self.max_len - len(ids)) ids ids [0] * (self.max_len - len(ids)) return torch.tensor(ids), torch.tensor(mask), torch.tensor(self.labels[idx])提示上面这段代码只是展示结构如果你直接跑记得把 mask 的逻辑按你选定的框架习惯统一好。这里我生成的 mask 是 1 表示真实 token0 表示 padding后续传给模型时可能需要转成模块要求的格式。很多新手在这个环节会直接把整段文本过一个分词器就算了忽略了长度统计。对于真实任务建议统计一下训练集文本长度的分位数把max_len设在 90% 分位附近。太小会截断有效信息太大则会大幅度拉长训练时间而且 padding 部分也会额外增加注意力计算量。3.2 完整编码器 Block 的实现一个标准的 Transformer 编码器 Block 由两个子层组成第一个是多头自注意力第二个是前馈网络每个子层外面都套着残差连接和 LayerNorm。先给一个直接可用的实现import torch import torch.nn as nn class TransformerBlock(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout0.1): super().__init__() self.attn nn.MultiheadAttention(d_model, num_heads, dropoutdropout, batch_firstTrue) self.ffn nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model), nn.Dropout(dropout), ) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) def forward(self, x, key_padding_maskNone): # pre-norm先归一化再进入子层训练更稳定 x x self.attn(self.norm1(x), self.norm1(x), self.norm1(x), key_padding_maskkey_padding_mask)[0] x x self.ffn(self.norm2(x)) return x这段代码里有几个值得注意的选择。第一我用了 pre-norm 结构也就是把 LayerNorm 放在注意力子层之前。相比 post-norm它在深网络中更多被使用因为梯度传导更平稳训练初期不容易崩。第二注意力子层里 Q、K、V 的输入我都传了同一个x这是自注意力的标准用法。你当然可以先分别做三个线性变换再传进去但nn.MultiheadAttention内部已经包含这些变换没必要重复做。第三batch_firstTrue可以让输入输出形状更符合直觉我们传入的x是(batch, seq_len, d_model)。前馈网络我这里用了两层线性加一个 GELU 激活。GELU 在 Transformer 里效果通常不差但如果你希望代码更简单用 ReLU 也没有问题。d_ff我习惯设成4 * d_model这是最常见的比例既保证了网络容量又不至于让参数量失控。模型最后需要接一个分类头。通常做法是取序列第一个位置的输出也就是[CLS]对应位置的向量过一层nn.Linear(d_model, num_classes)。这里有一个细节对于分类任务取第一个位置不一定总是最优有的任务上取均值池化比取[CLS]更稳。我的建议是都试一下不过本文以取[CLS]为例实现起来更直观。3.3 小规模文本分类任务训练流程搭建完模型接下来就是训练。这里我给你一个直接可以跑的流程重点是看清每个张量的流动。def train_epoch(model, loader, criterion, optimizer, device): model.train() total_loss, correct, total 0.0, 0, 0 for input_ids, padding_mask, labels in loader: input_ids input_ids.to(device) padding_mask padding_mask.to(device) labels labels.to(device) logits model(input_ids, key_padding_maskpadding_mask) loss criterion(logits, labels) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() * input_ids.size(0) correct (logits.argmax(dim-1) labels).sum().item() total labels.size(0) return total_loss / total, correct / total训练之前损失函数的选择要跟任务匹配二分类任务用CrossEntropyLoss它内部会对 logits 做 softmax所以模型最后一层不要手动加 softmax如果你想看概率要先从 logits 里取softmax(logits, dim-1)但argmax不需要 softmax直接比较 logits 的最大值下标即可。另外一点容易被忽视padding_mask不要光依赖模型的 attention 层自动处理。如果你用了自己写的手工注意力一定要确保 padding 位置的注意力分数被置为负无穷否则那些位置会被分配非 0 权重导致模型从 padding token 里学到无意义信息。即便你用的是封装好的多头注意力模块也建议显式传入key_padding_mask。通常训练几十个 epoch 之后准确率就能达到一个比较稳定的水平。因为这是一个偏演示性质的模型我不会把结果吹得太高重点是你打开 TensorBoard 或者简单打印 loss 曲线时能看到它在逐渐下降并且测试集上的准确率和直觉判断一致。3.4 关键训练参数与选择逻辑接着说说我常用的训练参数以及为什么这么选。模型维度d_model我一般从 128 起序列比较短、任务比较简单时可以降到 64num_heads通常取 4过少表达能力受限过多在小型任务上反而容易过拟合d_ff保持 4 倍关系即可dropout 我设 0.1数据集越小的场景dropout 可以适当提高到 0.2 甚至 0.3用来抗过拟合。优化器我默认用 AdamW因为它在带动量的同时对权重衰减做了修正比传统 Adam 更稳。学习率我惯用 2e-4 到 5e-4 之间。注意力模型对学习率比较敏感学习率太大容易出现 loss 突然变成 NaN太小则收敛很慢。如果你只有一个 CPU 机器batch size 建议设 16太大容易爆内存太小则梯度噪声明显。如果显存允许优先把 batch size 翻倍并等比例调低学习率。这里还有一个值得重点说的点注意力矩阵的计算复杂度是序列长度的平方。也就是说序列长度从 100 涨到 200注意力部分的计算量不是翻倍而是涨 4 倍。在实际任务里给模型选一个合理的max_len比选一个更大的d_model更影响训练速度。很多优化方案的本质也就是在想办法跳过那些不重要的位置对降低这个平方复杂度的影响。4. 训练中我反复踩过的深度模型问题4.1 Loss 不降、起不来或爆掉这类问题在深度模型里非常常见我想单独拿出来仔细聊聊。loss 不下降先排查最容易犯的错学习率太小或者数据预处理环节有问题。我曾经调试一个很小的情感分类任务模型结构完全没问题但 loss 就是稳定在一个固定值附近不降。后来发现我的标签在构建 dataset 时全部被写成了同一个值模型在学一个恒定的输出loss 当然不会因为训练步数增多而变化。在开始调结构之前先打印几个 batch 出来看看标签分布这是最便宜也最有效的排查手段。loss 起不来也就是每一步 loss 都在上下震荡常见原因是模型没有收敛或者优化器参数设置不对。我的经验是先固定随机种子把学习率调小两个量级再看 loss 有没有稳定下降。如果调小学习率之后 loss 很稳定说明模型本身没问题只是学习率不在最优区间。loss 变 NaN 的原因更直接通常是数值溢出。注意力分数的绝对值因为某种原因变得特别大softmax 之后出现极端的概率分布梯度也跟着爆炸。出现这种情况的时候先检查输入有没有包含异常大的值然后检查是不是漏了缩放因子也就是前面反复说的sqrt(d_k)。此外AdamW 的默认 epsilon 比较小也可以适当调到 1e-8 以上。4.2 长度 padding 带来的隐性偏差padding 的坑实在太多了我再单独展开一次。很多模型的训练 loss 正常推理结果也正常但一眼看过去却总觉得模型对短句子格外偏好。这个现象背后往往是 mask 用错了。如果你没有把 padding 位置的注意力权重清 0绝大多数句子在训练时都把注意力的一部分倾斜到 padding token 上。padding token 是同一个 id对所有句子贡献相同的噪声模型如果偷懒就很容易从这些噪声里找到一种“最保险”的输出方式导致它的分类依据变得很奇怪。排查 mask 问题有一个很直观的方法训练后随便取几个句子打印出最后一层注意力权重看 padding 位置有没有接近 0 的权重。如果 padding 位置权重明显大于 0说明 mask 没有生效。另一种问题更隐蔽你只 mask 了 padding 部分却没有 mask 掉未来信息这在做自回归任务时会导致信息泄露。虽然本文主要讲分类但如果你后面转做生成模型一定要记得因果 mask 和 padding mask 是两回事。4.3 注意力矩阵数值溢出注意力矩阵的数值问题分几种。最常见的是 softmax 的输入太大导致指数函数溢出。即使没有溢出如果输入分布非常不均匀也会让梯度变小。解决办法除了缩放点积还有在计算 softmax 时使用 max 来稳定数值很多框架的 softmax 内部已经做了这一步但你手工实现时就要留意。另一种溢出发生在多头合并时。如果把多个头的注意力输出直接拼接然后进入下一个线性层线性层的权重初始化范围如果不匹配结果也可能过大。我的做法是在每个子层后面都接一个 dropout并且把前馈网络的第二层初始化设置得稍微小一点这样即使中间结果偏大经过残差连接后也不会直接传导到下一层。最后我想提醒的是不要只盯着训练 loss 看。训练 loss 降到很低不代表模型真的理解了任务在小数据集上模型完全可能直接背题把整个训练集记下来然后测试集上一塌糊涂。所以评估环节一定要留独立的验证集而且验证集要和训练集分开处理不能有文本重叠。4.4 常见训练事故速查我整理了一份平时调试时经常对照的速查表按症状分列如下症状常见原因排查方向loss 恒定在某个固定值附近标签全是一个类别 / 输出被池化成了常数检查数据集里的标签分布打印 logits 统计loss 出现 NaN学习率过大 / 数值溢出 / 输入有 NaN降低 lr检查数据清洗检查注意力缩放训练 loss 下降但验证 loss 上升过拟合 / 数据泄漏增加 dropout检查验证集文本是否和训练集重复短句子预测偏向某一类padding mask 未生效打印注意力权重检查 padding 位置权重推理比训练慢很多没有设 eval 模式 / 动态图缓存写model.eval()推理时用torch.no_grad()或torch.inference_mode()多卡训练结果和单卡不一致随机种子没固定 / 数据采样顺序不一致固定种子统一数据采样逻辑提示表格里这些问题的排查顺序基本固定先查数据再查 mask最后才轮到调结构。不要一上来就改网络那样往往会把问题越改越复杂。5. 我的几点实战心得我自己在最初调这些深度模型的时候吃过最多亏的地方不是写不出结构而是“结构能跑但不理解它为什么这样跑”。后来我养成一个习惯每写一个网络组件都会强制自己用一个小到能打印张量的例子跑一遍然后把中间结果拿过来看。比如自注意力我就用一句话长度的输入把每个 head 的注意力矩阵可视化出来看看观察模型到底在关注哪些词。这一步对我理解 Q、K、V 的帮助比读十遍论文都大。另一个心得是关于“泛化能力”这件事。很多人在小任务上把准确率调得很高之后换一个数据分布就崩原因往往不是模型不够深而是数据预处理、mask、标签分布这些细节不干净。深度网络本身不过就是一个大规模可微函数你喂进去什么规律它就会拟合什么规律。所以我会在日常工作中尽量把数据管线的正确性放在模型结构之前先确保喂进去的是我想要的规律。最后这一篇的代码没有覆盖高效推理、量化、分布式训练这些工程方向但对你理解核心原理来说已经够用。如果你照着前面的步骤把一个小型 Transformer 跑通并且能清楚解释每一步形状变化和梯度流向那你已经超过了绝大多数只会调库的学习者。后续照着这个思路扩展到真正的生产环境比如接上更完整的数据管线、使用更长的序列、加入更强的训练策略都会顺很多。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑