""" ============================================================================= 第 6 讲 · 代码 5：注意力机制 + 束搜索 ============================================================================= 对应讲义：lectures/lecture06_lstm_seq2seq_attention.md 第 6、8、9 节 对应作业：CS224N Winter 2025 及更早的作业 3（带乘性注意力的神经机器翻译） ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 代码 04 的 seq2seq 把整个源句子压缩成一个向量，句子越长翻译越差（信息瓶颈）。 注意力（Bahdanau 等 2015；Luong 等 2015）的想法：**解码的每一步，都直接看一眼编码器的所有隐藏状态， 挑出和当前要生成的词最相关的那几个，加权平均后一起用来预测。** 解码第 t 步（解码器隐藏状态 h_t，编码器隐藏状态 h^enc_1 ... h^enc_S）： 1. 打分：e_{t,i} = h_tᵀ W_att h^enc_i （乘性注意力，和作业 3 一样） 2. 归一化：α_t = softmax(e_t) （注意力分布：在源句子各个位置上的权重） 3. 加权求和：a_t = Σ_i α_{t,i} h^enc_i （注意力输出 / 上下文向量） 4. 组合：o_t = tanh(W_u [a_t; h_t])，P(y_t) = softmax(W_vocab o_t) 本文件的模型和代码 04 **完全相同**（同样的双向 LSTM 编码器、同样的解码器），唯一的区别就是加了注意力。 演示 1：训练带注意力的模型 演示 2：按源句子长度比较 seq2seq 和注意力模型 —— 瓶颈消失了吗？ 演示 3：注意力热图 —— 模型学会了"对齐"（解码到句末的动词时，注意力跳回源句子的动词） 演示 4：束搜索 vs 贪心解码 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 05_attention.py # 完整运行（会加载或训练 checkpoints/seq2seq.pt 和 attention.pt） python 05_attention.py --quick # 冒烟测试 需要 PyTorch。完整运行约 5~10 分钟（CPU；如果 04 号文件已经保存了 seq2seq.pt，会快一些）。 """ import os # 必须在 import torch 之前设置：限制底层数学库（OpenMP）的线程数，原因见下面 set_num_threads 处的说明 os.environ.setdefault("OMP_NUM_THREADS", "1") import importlib import sys import time import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from bleu import corpus_bleu from toy_translation import Vocab, make_dataset s2s = importlib.import_module("04_seq2seq") torch.set_num_threads(1) # 小模型、逐步计算：单线程最快（见 04 号文件的说明） def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 第 1 部分：带注意力的 seq2seq # ============================================================================= class AttentionSeq2Seq(s2s.Seq2Seq): """ 带乘性注意力的 seq2seq（继承代码 04 的 Seq2Seq，编码器完全相同）。 【新增的参数】 W_att : (2H → 2H) 线性层，用于计算 e_{t,i} = h_tᵀ W_att h^enc_i W_u : (4H → 2H) 线性层，把 [a_t; h_t] 组合成 o_t 解码器的输入改为 [上一个词的嵌入; 上一步的 o_{t-1}]（"输入喂入"，input feeding，作业 3 也这样做）， 让模型在决定下一步看哪里时，知道上一步看了哪里。 【Python 小知识】 class AttentionSeq2Seq(s2s.Seq2Seq)：继承。子类自动拥有父类的所有方法（encode、forward）， 只需要重写不同的部分（这里是 __init__ 和 decode_step）。super().__init__(...) 调用父类的构造函数。 """ def __init__(self, src_V, tgt_V, emb=64, hidden=128, dropout=0.2): super().__init__(src_V, tgt_V, emb, hidden) H2 = 2 * hidden self.decoder = nn.LSTMCell(emb + H2, H2) # 输入 = 词嵌入 + 上一步的 o self.W_att = nn.Linear(H2, H2, bias=False) self.W_u = nn.Linear(2 * H2, H2, bias=False) self.dropout = nn.Dropout(dropout) self.out = nn.Linear(H2, tgt_V) def decode_step(self, y_prev, state, enc_out, mask, extra=None): """ 解码一步。 【参数】 y_prev : (B,) 上一个目标词 state : (h, c) 解码器状态 enc_out : (B, S, 2H) 编码器每个位置的隐藏状态 mask : (B, S) True = 真实的词 extra : 上一步的 o_{t-1}，(B, 2H)；第一步为 None（用 0） 【返回】(logits (B, V), 新状态, o_t, 注意力分布 α_t (B, S)) 【边界】 - 填充位置的打分设为 -inf，softmax 后权重为 0 —— 注意力不会落在 上。 【Python 小知识】 torch.bmm：批量矩阵乘法。enc_out (B, S, 2H) 与 query (B, 2H, 1) 相乘得 (B, S, 1)， 即对每个句子、每个源位置算一个点积。 """ if extra is None: extra = torch.zeros(y_prev.shape[0], enc_out.shape[2]) x = torch.cat([self.tgt_emb(y_prev), extra], dim=1) h, c = self.decoder(x, state) scores = torch.bmm(self.W_att(enc_out), h.unsqueeze(2)).squeeze(2) # e_{t,i}，(B, S) scores = scores.masked_fill(~mask, float("-inf")) alpha = torch.softmax(scores, dim=1) # α_t a = torch.bmm(alpha.unsqueeze(1), enc_out).squeeze(1) # a_t，(B, 2H) o = self.dropout(torch.tanh(self.W_u(torch.cat([a, h], dim=1)))) # o_t return self.out(o), (h, c), o, alpha # ============================================================================= # 第 2 部分：束搜索 # ============================================================================= @torch.no_grad() def beam_search(model, src, src_len, tgt_vocab, beam_size=5, max_len=40, length_norm=True): """ 对**一个**源句子做束搜索，返回得分最高的译文（词列表）。 【思路】（讲义第 6.3 节） 贪心解码每一步只保留最好的 1 个前缀，一旦选错无法回头。 束搜索每一步保留得分最高的 k 个前缀（"束"），得分 = 对数概率之和： 1. 对束里的每个前缀，算出下一个词的分布，扩展出 k × V 个候选； 2. 从所有候选中选出得分最高的 k 个，作为新的束； 3. 生成了 的前缀被移到"已完成"列表，不再扩展。 结束后，从已完成的译文中选最好的。因为每多一个词就多乘一个小于 1 的概率， 长译文的总得分天然更低，所以通常用**长度归一化**：得分 / 长度。 【参数】 model : 带 encode / decode_step 接口的模型 src, src_len : 一个句子（批大小 1） tgt_vocab : 目标词表 beam_size : 束宽 k（k = 1 就是贪心） max_len : 最大长度 length_norm : 是否用长度归一化 【返回】list[str] 【边界】 - 束搜索不保证找到全局最优的译文（那需要枚举 V^T 种可能），只是比贪心更不容易"走错路"。 - 达到最大长度仍未完成的前缀，也会作为候选参与最后的比较。 """ model.eval() enc_out, mask, state = model.encode(src, src_len) beams = [([1], 0.0, state, None)] # (编号序列, 对数概率之和, 解码器状态, extra) finished = [] for _ in range(max_len): candidates = [] for seq, score, st, extra in beams: logits, new_st, new_extra, _ = model.decode_step(torch.tensor([seq[-1]]), st, enc_out, mask, extra) logp = F.log_softmax(logits[0], dim=-1) top_lp, top_id = logp.topk(beam_size) for lp, wid in zip(top_lp.tolist(), top_id.tolist()): candidates.append((seq + [wid], score + lp, new_st, new_extra)) candidates.sort(key=lambda x: -x[1]) beams = [] for cand in candidates: if cand[0][-1] == 2: # 生成了 finished.append(cand) else: beams.append(cand) if len(beams) == beam_size: break if len(finished) >= beam_size or not beams: break pool = finished or beams best = max(pool, key=lambda x: x[1] / (len(x[0]) - 1) if length_norm else x[1]) return tgt_vocab.decode(best[0][1:]) # ============================================================================= # 第 3 部分：画注意力热图 # ============================================================================= def ascii_attention(alpha, src_words, tgt_words): """ 在终端里画注意力热图：行是生成的目标词，列是源句子的词，字符越"重"权重越大。 【参数】 alpha : np.ndarray (T, S)，每一行是一个解码步的注意力分布 src_words : 源句子的词（含 ） tgt_words : 生成的目标词（含 ） 【返回】str """ shades = " .:-=+*#%@" w = max(len(x) for x in src_words) + 1 lines = [" " * 10 + "".join(f"{s[:w - 1]:>{w}s}" for s in src_words)] for t, tw in enumerate(tgt_words): row = "".join(f"{shades[min(int(a * 10), 9)] * 2:>{w}s}" for a in alpha[t, :len(src_words)]) lines.append(f"{tw:>9s} {row}") return "\n".join(" " + l for l in lines) # ============================================================================= # 演示 # ============================================================================= def main(quick=False): train, test = make_dataset() src_vocab, tgt_vocab = Vocab([s for s, _ in train]), Vocab([t for _, t in train]) epochs = 1 if quick else 10 if quick: train, test = train[:2000], test[:100] # ---------------- 演示 1：训练 ---------------- print_header(f"演示 1：训练带注意力的 seq2seq（{len(train)} 对句子，{epochs} 轮）") torch.manual_seed(0) if quick: att = AttentionSeq2Seq(len(src_vocab), len(tgt_vocab)) s2s.train_nmt(att, train, src_vocab, tgt_vocab, epochs=epochs) base = s2s.Seq2Seq(len(src_vocab), len(tgt_vocab)) s2s.train_nmt(base, train, src_vocab, tgt_vocab, epochs=epochs, verbose=False) else: t0 = time.perf_counter() att = s2s.load_or_train(AttentionSeq2Seq, train, src_vocab, tgt_vocab, "attention", epochs=epochs) print(f" 用时 {time.perf_counter() - t0:.1f} 秒（取决于你的电脑；如果是加载的检查点则很快）") print(" 加载或训练不带注意力的 seq2seq（代码 04）作为对照：") base = s2s.load_or_train(s2s.Seq2Seq, train, src_vocab, tgt_vocab, "seq2seq", epochs=epochs, verbose=False) n_att = sum(p.numel() for p in att.parameters()) n_base = sum(p.numel() for p in base.parameters()) print(f"\n 参数量：seq2seq {n_base:,}，注意力模型 {n_att:,}") # ---------------- 演示 2：按长度比较 ---------------- print_header("演示 2：seq2seq vs 注意力，按源句子长度比较") for name, m in [("seq2seq", base), ("注意力", att)]: bleu, exact = s2s.evaluate(m, test, src_vocab, tgt_vocab) print(f"\n {name}：整个测试集 BLEU {bleu:.3f}，完全正确 {exact:.1%}") print() buckets = s2s.length_buckets(n=100 if quick else 300) s2s.report_by_length(base, src_vocab, tgt_vocab, buckets, label="seq2seq") print() s2s.report_by_length(att, src_vocab, tgt_vocab, buckets, label="注意力") print("\n 比训练时更长的句子（4 个介词短语），两个模型各错在哪里？") for name, m in [("seq2seq", base), ("注意力", att)]: hyps, refs = s2s.translate(m, buckets[4], src_vocab, tgt_vocab) diff = np.mean([len(h) - len(r) for h, r in zip(hyps, refs)]) print(f" {name}：译文平均比参考短 {-diff:.1f} 个词。例子：") print(f" 参考：{' '.join(refs[0])}\n 模型：{' '.join(hyps[0])}") print(""" 观察： - 在训练见过的长度上（0~3 个介词短语），注意力模型几乎全对，而 seq2seq 随长度明显下降 —— 解码器不再需要把整个源句子"背"在一个向量里，每一步都能直接取用源句子中需要的那部分信息。 - 在比训练时更长的句子上，注意力模型的 BLEU 更高，"完全正确"却是 0：它的译文非常流畅， 但总是**漏掉一个介词短语**（训练时最多只见过 3 个，模型学到了"3 个之后就该输出动词了"）。 seq2seq 的译文长度大致对，但内容和顺序乱了。 - 结论：注意力解决了信息瓶颈，但**不会自动推广到训练分布之外的长度**。长度泛化至今仍是 序列模型（包括 Transformer）的一个难题。（具体数字以你运行的结果为准。）""") # ---------------- 演示 3：注意力热图 ---------------- print_header("演示 3：注意力热图（行：生成的 Zhoric 词；列：英语源句子；符号越重权重越大）") examples = [s for s in test if len(s[0]) >= 10][:2] for src_words, ref in examples: b = s2s.make_batches([(src_words, ref)], src_vocab, tgt_vocab, 1, shuffle=False)[0] ids, attns = s2s.greedy_decode(att, b["src"], b["src_len"]) hyp = tgt_vocab.decode(ids[0]) T = len(hyp) + 1 alpha = torch.stack(attns[:T], dim=0)[:, 0].numpy() print(f"\n 英语：{' '.join(src_words)}\n 参考：{' '.join(ref)}\n 模型：{' '.join(hyp)}\n") print(ascii_attention(alpha, src_words + [""], hyp + [""])) srcw = src_words + [""] pairs = [f"{w}→{srcw[int(alpha[i].argmax())]}" for i, w in enumerate(hyp + [""])] print("\n 每个生成词注意力最大的源词：" + "，".join(pairs)) print(""" 观察： - 实词基本对齐到对应的英语词：rega → kings、miro → cats、gron → old、goru → farmer…… 没有任何人告诉模型哪个词对应哪个词，它在学习翻译的过程中自己学会了这种"对齐"。 - 但有些"对齐"和人的直觉不同：生成地点名词和后置词时，注意力常常落在前面的介词（in、near、behind）上； 生成句末的动词时，注意力没有跳回英语的动词，而是落在源句子末尾的 上。 原因是注意力看的是**编码器的隐藏状态**，而双向 LSTM 的每个状态都已经混入了上下文： 前向 LSTM 在 处已经读完了整个句子（包括动词），介词位置的后向状态里也有后面那个名词的信息。 模型只要能从某个位置取到需要的信息就行，不一定非要看"对应"的那个词。 - 这提醒我们：注意力权重可以给我们一些直观的线索，但**不一定**是模型决策的忠实解释（第 17 讲）。""") # ---------------- 演示 4：束搜索 ---------------- n_beam = 50 if quick else 300 print_header(f"演示 4：束搜索 vs 贪心（注意力模型，测试集前 {n_beam} 句）") sub = test[:n_beam] for k in [1, 3, 5]: t0 = time.perf_counter() hyps, refs = s2s.translate(att, sub, src_vocab, tgt_vocab, decode_fn=lambda m, s, l, k=k: beam_search(m, s, l, tgt_vocab, beam_size=k)) bleu, _ = corpus_bleu(hyps, [[r] for r in refs]) exact = np.mean([h == r for h, r in zip(hyps, refs)]) print(f" 束宽 k = {k}：BLEU {bleu:.3f}，完全正确 {exact:.1%}，用时 {time.perf_counter() - t0:.1f} 秒" + ("（k = 1 等价于贪心）" if k == 1 else "")) print(""" 观察：束搜索通常比贪心略好或持平（在这个简单的任务上，模型已经很确定，贪心本身就很少走错）； 代价是解码时间大约增加到 k 倍。在真实的机器翻译中，束宽通常取 4~10。 （具体数字以你运行的结果为准。）""") if __name__ == "__main__": main(quick="--quick" in sys.argv)