很多人背注意力机制的公式背得滚瓜烂熟,什么 Q、K、V,什么 softmax,什么 scale,但真让他自己动手写一个带注意力的模型,立马卡壳。原因很简单,注意力机制不是一个孤立的概念,它必须挂在一个具体的网络结构上才有意义,而多数教程恰恰把“公式推导”和“工程实现”拆成了两件事。这次我就用 Python + PyTorch 从零实现一个 HAN(Hierarchical Attention Network,层级注意力网络),带你从代码层面把“注意力机制”彻底看穿。HAN 特别适合拿来练手,因为它在词和句子两个层级分别做了一次完整的注意力计算,结构清晰、效果直观,跑通一次之后,你再看 Transformer 里的自注意力,会觉得顺眼很多。
这篇博文面向的读者是有一定 Python 基础、想深入理解注意力机制的初学者,也可能是被各种注意力变体绕晕、想找回底盘感的人。我会先讲清楚模型的整体设计,再拆解注意力机制的核心原理,然后给出完整可运行的 PyTorch 代码,最后分享我在实际训练中踩过的坑和调优经验。全程以“能跑起来”为第一目标,代码不长,但每个关键动作我都会解释为什么这么做。
1. 项目概述与设计思路
1.1 HAN 到底在解决什么问题
先聊一下为什么会有 HAN 这种结构。在做文本分类的时候,比如情感分析、新闻分类、垃圾邮件识别,我们面对的不是一句短文本,而是一篇由多个句子组成的文档。很多经典做法是:把整篇文档的词都嵌入到一个向量里,然后对词向量做平均池化或者最大池化,得到整个文档的向量表示,最后再送到分类器里。
这种做法的问题在于,一篇文档里并不是每个词、每个句子都同等重要。比如一篇产品评论里,“电池续航太差”这句话里的“差”显然比“我昨天买了一个”里的“我”更能决定情感极性。如果用平均池化,所有词的贡献被一视同仁地压进同一个向量,关键信息就被稀释掉了。最大池化稍好一点,但它只能抓住单一最强的信号,没办法兼顾多个关键点。
HAN 的思路很直接:我们人在读文章的时候,也并不是每个词都看得一样重,而是先快速浏览,把重要词圈出来,再看看这些重要词集中在哪些句子里,最后综合判断整篇文章在讲什么。HAN 就干了同样的事情:先计算每个词对句子情感的“贡献权重”,加权求和得到句子向量;再计算每个句子对整篇文章情感的“贡献权重”,加权求和得到文档向量。这个“先词后句、层层筛选”的过程,就是“层级注意力”的核心。
1.2 为什么“手写实现”比“调库”更重要
现在 PyTorch 生态里已经有很多封装好的注意力实现,甚至有nn.MultiheadAttention这种一行就能调用的接口。但我强烈建议你至少手动实现一遍 HAN 这种结构再谈“我理解了注意力机制”。原因是:注意力机制的本质极其简单,就是“打分、归一化、加权求和”三步,但一旦被封装成高级 API,你反而看不到内部的数据流动方向了。而我个人也觉得,在nn.MultiheadAttention里塞进去的维度变换,对初学者来说是很大的干扰项。你自己实现的时候,每写一行代码都得想清楚这一步输入是什么形状、输出是什么形状、mask 在哪里起作用,这个思考过程比任何公式推导都有价值。
另外,HAN 还有一个特别适合教学的特点:它把“注意力”这种看似抽象的概念变成了可视化的权重大小。训练结束之后,你能够把每个词的 attention 权重取出来,用热力图直接看到模型把注意力放在了哪一个词上,这种反馈会极大增强你对模型的理解。
选 PyTorch 而不是 TensorFlow 的原因也很简单:动态图的调试体验好,print(shape)到处打就行,代码风格也更接近普通 Python,对新手更友好。何况现在 PyTorch 在学术界和工业界的生态已经非常成熟,学它不亏。
2. 核心机制原理解读
2.1 注意力机制的本质:打分、归一化、加权求和
很多人把注意力机制想复杂了。我换个方式来解释:假如你在读书,书包里只有一张空白的摘抄纸,现在要写一篇三百字的书评,你不可能把全书都抄下来,你会一边读一边给每个句子“打分”——这句重要加两分,那句不重要跳过——最后把分数最高的句子摘抄下来,拼成你想要的资料。这个“打分、挑重点、组合信息”的过程,就是注意力机制在做的事。
在神经网络里,我们面对的是一组向量,比如一个句子里每个词的隐藏状态:(h_1, h_2, ..., h_T)。注意力机制要做三件事:
第一步,打分。给每个向量算一个分数 (e_i),这个分数代表了它和当前任务的相关程度。打分函数可以很简单,比如直接接一个线性层再套 tanh 激活,再映射成标量;也可以复杂一点,比如用点积、加性注意力、多头注意力里的缩放点积等。HAN 用的就是比较经典的多层感知机打分方式。
第二步,归一化。把打分得到的分数 (e_i) 转成一组和为 1 的权重 (\alpha_i)。最常见的就是 softmax。为什么用 softmax?因为它能把分数变成相对大小明显、梯度性质良好的概率分布。你很难手动指定“第 3 个词权重 0.5,第 5 个词权重 0.3”,但 softmax 能根据原始分数自己学出这样的分配。
第三步,加权求和。用归一化后的权重对原始的隐藏状态做加权平均,得到一个新的向量 (c = \sum_i \alpha_i h_i)。这个向量就是模型“带着关注点”去看完整段输入之后,提取出的关键信息。
这三步全程可微,意味着模型可以通过反向传播自动学会什么样的输入更重要。所谓“注意力”,本质上就是一个由上下文动态决定的数据聚合函数。后面你看到的各种注意力变体,大多数只是在“打分”这一步做了文章:CBAM 加了通道维度的注意力,自注意力用了 QKV 点积,多头注意力是把多个注意力头并行算了一遍再拼起来。底层逻辑没变。
2.2 词级注意力:让模型学会“读词”
HAN 的第一层注意力叫作 word-level attention,作用对象是“一个句子内的所有词向量”。结构大概是这样的:先把一个句子的所有词都过一遍词嵌入层,得到一个词向量序列,再放进一个双向 GRU 里,得到每个词对应的隐藏状态 (h_{it}),下标 i 表示第 i 个句子,t 表示第 t 个词。
然后,用一个共享的小型多层感知机对每个词的隐藏状态打分,得到 (u_{it} = \tanh(W_w h_{it} + b_w)),再用一个全局上下文向量 (u_w) 和它做内积,得到最终分 (e_{it} = u_{it}^T u_w)。经过 softmax 得到每个词的注意力权重 (\alpha_{it}),最后加权求和所有隐藏状态,得到这个句子的向量表示 (s_i = \sum_t \alpha_{it} h_{it})。
这个“全局上下文向量” (u_w) 可以理解为一个可学习的“问题”:模型在阅读一个句子时,会一直在心里默念“当前任务关心什么样的词”,然后根据每个词和这个问题的匹配程度来分配注意力。训练的过程中,这个上下文向量会被调整成和数据集最相关的一个语义方向。
2.3 句级注意力:把句子当“词”再来一遍
拿到每个句子的向量 (s_i) 之后,整篇文档就变成了一个“句子向量序列”。接下来,HAN 再做一次几乎完全一样的操作:把句子向量序列送进另一个双向 GRU,让模型在看完整篇文章的上下文后,为每个句子生成隐藏状态 (H_i)。再用另一套打分网络计算每个句子的权重,加权求和得到文档向量 (v)。最后把 (v) 接一个全连接层和 softmax,就可以做分类了。
这就是 HAN 的“层级”所在:最底层是词,中间层是句子,最顶层是整篇文档。每一个层级都独立完成了一次“打分、归一化、加权求和”的注意力计算,但它们的输入完全不同。这个设计的巧妙之处在于,它把“从大量信息中挑出关键信息”这个目标拆解成了两步,让模型在词的层面只关注“哪些词重要”,在句子的层面只关注“哪些句子重要”,而不是让一个注意力层一口气处理所有词的拼接结果。后者不仅计算量大,而且很难学到这种中间层级的抽象语义。
3. 数据准备与 PyTorch 实现
3.1 数据格式与预处理
开始写模型之前,先解决数据的问题。HAN 输入的一个样本是一篇文档,但一篇文档要同时拆成“句子”和“词”两个维度,这一点和普通文本分类很不一样。我用一个简单的产品评论情感分类为例,每一条评论只有两三句话,但足够说明数据格式。
原始数据大概是这样的:
"手机很好用,但电池续航不行。"预处理要做三件事:按句子切分,得到["手机很好用,但电池续航不行。"];按词切分(这里用简单的空格分词,中文场景可以用 jieba),得到["手机", "很好用", "但", "电池", "续航", "不行"];再把词映射成词表中的索引 id。
因为 batch 里每篇文档的句子数和每句的词数都不同,我们必须做 padding。我规定一个固定的max_sentences和max_seq_len,超过的截断,不足的补 0。这里有一个很重要的细节:词的 padding 索引设为 0,并且让 PyTorch 的 Embedding 层忽略它。同时,我们还需要维护两个 mask 张量:一个是句子维度的 masksent_mask,标记哪些句子是真实的、哪些句子是 padding 出来的;另一个是词维度的 maskword_mask,标记每个句子内部的哪些位置是有效词、哪些是 padding。这两个 mask 在后面计算注意力时会直接用到。
3.2 词级编码与注意力模块代码
注意力模块是整个 HAN 的核心,我建议单独封装成一个Attention类。它的输入是inputs和mask,inputs的形状是(batch_size, seq_len, hidden_size),mask的形状是(batch_size, seq_len),1 表示有效位置,0 表示 padding 位置。
打分部分我用一个线性层把hidden_size映射到一个标量,前面套一个 tanh 激活。这一步对标论文里的u = tanh(W h + b)。代码可以精简为:
# attention.py import torch import torch.nn as nn import torch.nn.functional as F class Attention(nn.Module): def __init__(self, hidden_size): super(Attention, self).__init__() self.score_layer = nn.Linear(hidden_size, hidden_size, bias=False) self.context_layer = nn.Linear(hidden_size, 1, bias=False) def forward(self, inputs, mask): # inputs: (batch, seq_len, hidden_size) # mask: (batch, seq_len) 1为有效,0为padding u = torch.tanh(self.score_layer(inputs)) scores = self.context_layer(u).squeeze(-1) # (batch, seq_len) scores = scores.masked_fill(mask == 0, -1e9) weights = F.softmax(scores, dim=1) # (batch, seq_len) weighted = torch.bmm(weights.unsqueeze(1), inputs).squeeze(1) return weighted, weights这里要注意的细节有两个。第一,masked_fill要在 softmax 之前做,把 padding 位置的分数设成一个极小的负数,这样 softmax 之后这些位置的概率会趋近于 0。如果放到 softmax 之后再做,那会破坏概率分布的“和为 1”特性,加权求和就会出问题。第二,torch.bmm这一步做了批量矩阵乘法,把每个时间步的隐藏状态和对应权重乘起来再加总,相当于把“打分、归一化、加权求和”里的最后一步一次性做完。
3.3 句级注意力与 HAN 主类代码
有了Attention这个通用模块,HAN 主类就清晰多了。我用双向 GRU 作为词级和句级的编码器。词级 GRU 的输入维度是embedding_dim,输出维度是hidden_size * 2,因为双向会把前后两个方向的隐藏状态拼接起来。之后词级注意力把每个句子压缩成一个句向量,句向量再进入句级 GRU,句子级注意力再把整个文档压缩成一个文档向量。
代码如下:
# han.py import torch import torch.nn as nn class HAN(nn.Module): def __init__(self, vocab_size, embedding_dim=100, hidden_size=64, num_classes=2, n_layers=1): super(HAN, self).__init__() self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0) self.word_gru = nn.GRU(embedding_dim, hidden_size, num_layers=n_layers, batch_first=True, bidirectional=True) self.word_attn = Attention(hidden_size * 2) self.sent_gru = nn.GRU(hidden_size * 2, hidden_size, num_layers=n_layers, batch_first=True, bidirectional=True) self.sent_attn = Attention(hidden_size * 2) self.fc = nn.Linear(hidden_size * 2, num_classes) def forward(self, x, word_mask, sent_mask): # x: (batch, num_sentences, seq_len) # word_mask: (batch, num_sentences, seq_len) # sent_mask: (batch, num_sentences) batch_size, num_sentences, seq_len = x.size() # 词级处理:把(batch, num_sentences, seq_len)合并成(batch*num_sentences, seq_len) word_input = x.view(-1, seq_len) word_mask_flat = word_mask.view(-1, seq_len) embedded = self.embedding(word_input) # (batch*sent, seq_len, emb_dim) word_out, _ = self.word_gru(embedded) # (batch*sent, seq_len, 2*hidden) sent_vecs, word_weights = self.word_attn(word_out, word_mask_flat) sent_vecs = sent_vecs.view(batch_size, num_sentences, -1) # (batch, num_sentences, 2*hidden) # 句级处理 sent_out, _ = self.sent_gru(sent_vecs) # (batch, num_sentences, 2*hidden) doc_vec, sent_weights = self.sent_attn(sent_out, sent_mask) # (batch, 2*hidden) logits = self.fc(doc_vec) return logits, word_weights, sent_weights这段代码里最关键也最容易懵的地方在word_input的 reshape。我一开始写的时候也经常在这里绕晕:原本输入是“文档 -> 句子 -> 词”的三维结构,但我们想让同一个 GRU 去处理所有句子里的所有词,所以把前两个维度合并成一个大 batch,等注意力取到句向量之后再重新 reshape 回原来的文档维度。只要掌握了这个“展平再还原”的思路,HAN 的实现就没什么神秘的了。
3.4 训练流程与完整训练代码
训练流程和普通文本分类模型差别不大。我用一个最简单的情感分类数据集,只求先把整个流程跑通。核心训练循环如下:
# train.py import torch import torch.nn as nn from torch.utils.data import DataLoader, Dataset class TextDataset(Dataset): def __init__(self, data, max_sentences=8, max_seq_len=20): self.data = data self.max_sentences = max_sentences self.max_seq_len = max_seq_len def __len__(self): return len(self.data) def __getitem__(self, idx): labels, sents = self.data[idx] # 这里 data 中的 sents 已经是索引化后的二维列表 # 截断或补齐句子数 if len(sents) > self.max_sentences: sents = sents[:self.max_sentences] else: sents += [[0] * self.max_seq_len] * (self.max_sentences - len(sents)) sent_mask = [1 if any(s) else 0 for s in sents] # 句子是否有真实词 new_sents = [] word_mask = [] for sent in sents: if len(sent) > self.max_seq_len: sent = sent[:self.max_seq_len] else: sent = sent + [0] * (self.max_seq_len - len(sent)) new_sents.append(sent) word_mask.append([1 if w != 0 else 0 for w in sent]) return torch.tensor(new_sents), torch.tensor(word_mask), torch.tensor(sent_mask), torch.tensor(labels) def train_one_epoch(model, dataloader, optimizer, criterion, device): model.train() total_loss = 0.0 total_correct = 0 total_num = 0 for x, word_mask, sent_mask, y in dataloader: x = x.to(device) word_mask = word_mask.to(device) sent_mask = sent_mask.to(device) y = y.to(device) optimizer.zero_grad() logits, _, _ = model(x, word_mask, sent_mask) loss = criterion(logits, y) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() total_loss += loss.item() * y.size(0) pred = logits.argmax(dim=1) total_correct += (pred == y).sum().item() total_num += y.size(0) return total_loss / total_num, total_correct / total_num model = HAN(vocab_size=vocab_size, embedding_dim=100, hidden_size=64, num_classes=2) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) criterion = nn.CrossEntropyLoss() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) for epoch in range(20): train_loss, train_acc = train_one_epoch(model, train_loader, optimizer, criterion, device) print(f"epoch {epoch + 1}: loss={train_loss:.4f}, acc={train_acc:.4f}")这里我特意加了梯度裁剪clip_grad_norm_,因为 GRU 在反向传播时很容易出现梯度爆炸,尤其在小数据集上训练的时候。裁剪到一个较小的范数(比如 1.0)并不影响模型收敛,但能避免 loss 突然变成 NaN 的尴尬。如果你在自己的数据上跑,发现 loss 震荡严重,可以试试把学习率从1e-3降到5e-4,通常会有明显改善。
4. 实操过程与调优实录
4.1 维度匹配:最容易翻车的地方
很久以前我第一次写 HAN,最崩溃的不是理解注意力,而是torch.Size对不上。这里我以batch_size=16, num_sentences=8, seq_len=20, embedding_dim=100, hidden_size=64为例,完整捋一遍各个阶段的形状变化:
- 输入
x:(16, 8, 20) - 经过 embedding:
(16, 8, 20, 100) view(-1, 20)展平前两维:(128, 20, 100)- 词级双向 GRU 输出:
(128, 20, 128) - 词级注意力输出句向量:
(128, 128) view(16, 8, -1)还原文档结构:(16, 8, 128)- 句级双向 GRU 输出:
(16, 8, 128) - 句级注意力输出文档向量:
(16, 128) - 全连接输出 logits:
(16, 2)
如果你在调试的时候发现word_input = x.view(-1, seq_len)之后维度对不上,大概率是 batch 里的每个样本num_sentences不一致。我建议在 DataLoader 的collate_fn里统一做 padding,而不是在Dataset.__getitem__里只对单个样本处理。最开始我是后者,结果 batch 里句子长度参差不齐,一合并就报错。后来我干脆把所有样本在Dataset内部就截断/补齐成同样的max_sentences和max_seq_len,并且在collate_fn里什么都不做,直接堆叠张量,省去了大量麻烦。
4.2 掩码处理与 padding 的坑
mask 是 HAN 里最容易出 bug 的地方,一旦写错,模型可能还能训练,但效果会变得很差,而且很难定位问题。常见的坑有三个:
第一个坑是 padding 位置的分数没有屏蔽。如果忘了masked_fill,padding 出来的 0 向量也会参与注意力权重的计算,因为它们与打分网络计算出来的分数可能并不小。结果就是模型把大量注意力浪费在根本没有意义的占位符上,文档向量被严重稀释。我管这个问题叫“模型在盯着空气看”。
第二个坑是词级 mask 和句级 mask 搞混。在把(batch, num_sentences, seq_len)展平之后,词级 mask 也要跟着展平到(batch * num_sentences, seq_len),否则维度对不上,或者更隐蔽地,mask 和输入错位了。你在调试的时候可以随机选一个样本,手工把 attention 权重和原始文本对照一下,看看模型是不是真的把高权重给到了“很棒”“差”这类情感词,而不是给了“的”“了”这种虚词。
第三个坑是在评估时取 attention 权重也要小心。因为我用了 padding,所以权重里的 padding 位置虽然是 0,但它们仍然占据概率分布的一部分,真实词上的权重会被压缩。好在它们本来就是从 softmax 里算出来的合法分布,不会影响分类结果。但如果你想拿权重做可视化,最好先把 padding 位置去掉再归一化一次,否则画出来的热力图上会有一大段零值尾巴,观感很怪。
4.3 训练稳定性与效果提升
HAN 在小数据集上其实很容易过拟合。我现在一般会做几件事情来提升稳定性和最终效果。
第一,固定随机种子。PyTorch 里设置torch.manual_seed(42)加上random.seed(42),如果用了 numpy 还要设置np.random.seed(42)。不设种子的话,你复现别人的实验时,可能每次跑出来的模型结果差一两个百分点,你根本无法判断是调参有效还是运气好。
第二,给 embedding 层加 dropout。HAN 的结构不复杂,参数量主要来自词向量和 GRU。在 EMBEDDING 层后接一个nn.Dropout(0.5),能让模型在数据量有限的情况下不那么容易记住训练集里的噪音信号。如果你觉得效果不明显,还可以在 HAN 主类里给两个 GRU 也加上 0.2 到 0.3 的 dropout。
第三,验证集上的早停。训练到一定 epoch 后,验证集 loss 会先降后升,这是典型的过拟合信号。我在脚本里记录了验证集准确率最高的模型参数,训练结束后再load_state_dict回来做测试。不要只在训练集上盯指标,否则你会发现训练准确率接近 100%,测试集却惨不忍睹。
第四,学习率不要贪大。我刚跑这个模型的时候用lr=1e-2,结果 loss 在几个 epoch 内直接飞掉了,变成了 NaN。后来换成1e-3,情况才稳下来。如果数据量很小,甚至可以考虑5e-4。注意 Adam 本身对学习率比较鲁棒,但 GRU 的循环结构仍然对过大的梯度非常敏感,这就是为什么要做梯度裁剪。
5. 常见问题与排查技巧实录
5.1 常见问题速查表
这里整理了一些我在实际操作中遇到的典型问题,以及对应的排查思路。建议收藏,遇到问题先对照一下。
| 症状 | 可能原因 | 排查方法与解决方案 |
|---|---|---|
| loss 变成 NaN | 学习率过高、梯度爆炸 | 调低学习率,加梯度裁剪,检查 embedding 是否有异常值 |
| 训练准确率高但测试准确率低 | 过拟合 | 加 dropout,做早停,增加训练数据或数据增强 |
| 注意力权重几乎均匀分布 | 模型没有学到有效的语义特征 | 换用预训练词向量,增大 embedding_dim,检查是否忘了 mask |
| 词级和句级 GRU 的输出维度对不上 | reshape 或层参数设置错误 | 打印每步 tensor shape,重点检查 view 前后的维度 |
| 训练很慢 | batch_size 太大或序列过长 | 调小 batch_size 和 max_seq_len,先用小数据跑通流程 |
| 结果随机性大 | 没有固定随机种子 | 在train.py开头固定torch.manual_seed(42) |
5.2 HAN 这个框架还能怎么改
既然你已经把 HAN 从零写出来了,接下来完全可以在这个骨架上做各种扩展实验,这也是我推荐动手实现模型的原因:改起来非常方便。
第一个方向是把中文分词换成更好的工具。我这里为了演示用了简单的空格分词,但在实际中文数据集上,你至少应该用 jieba 或者其他工业级分词工具。分词质量直接影响模型的表达能力,词表里的 OOV(out-of-vocabulary)词太多,再好的注意力机制也难有作为。
第二个方向是换编码器。双向 GRU 是论文里的原版配置,但你可以很容易地把nn.GRU换成nn.LSTM,对比一下两者的效果和训练速度。LSTM 在长文本上通常表现更好,代价是参数量和计算量更大。如果你想让实验更有趣,还可以试试把词级编码器换成一个小型 Transformer Encoder,这样就能顺带练习一下自注意力的实现。
第三个方向是做 attention 权重的可视化。如果你训练的是二分类情感模型,可以写一个小函数,输入一句话,输出每个词和每个句子的权重,再用 matplotlib 画成热力图。这一步虽然不影响模型效果,但对理解模型行为太有帮助了,也是你以后给朋友解释“注意力到底是什么”的最好素材。
第四个方向是接入预训练词向量。在nn.Embedding.from_pretrained里加载 GloVe、word2vec 或者中文的腾讯词向量,然后设置freeze=False微调,通常会比随机初始化词向量收敛得更快、效果更好。这也是一种很廉价的“迁移学习”实验。
5.3 我在实际调试中的一点体会
最后聊点实际操作层面的感受。刚开始调 HAN 那会儿,我总喜欢一次性把代码写得非常完备,结果一跑起来报错,来回找 bug 找了一个小时。后来学乖了,先写一个极简的版本,只跑一个 batch,确认前向传播的每个张量形状都对,再写训练循环。现在我的习惯是:任何自定义模型,都会在forward函数的每个关键节点后面加一行print(x.shape),跑通之后再把打印去掉。这个习惯看着笨,但真的能省下大量排查时间。
另外,如果你对 PyTorch 的自动求导机制还不太熟,建议在训练之前手动算一下交叉熵损失对 logits 的梯度,再用loss.backward()之后的model.fc.weight.grad对比验证一下。这个练习做完之后,你会发现模型训练里的“反向传播”不再是一个黑盒子,HAN 里那些 GRU、注意力层对你来说也都只是“可以求导的一串数学运算”而已。