""" ============================================================================= 第 6 讲 · 代码 2：LSTM 与 GRU —— 从零实现、梯度流动、长距离记忆任务、门在做什么 ============================================================================= 对应讲义：lectures/lecture06_lstm_seq2seq_attention.md 第 3 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- LSTM（Hochreiter & Schmidhuber 1997）给 RNN 加了一个"细胞状态" c_t 和三个"门"： f_t = σ(W_f [h_{t-1}; x_t] + b_f) 遗忘门：旧的记忆保留多少 i_t = σ(W_i [h_{t-1}; x_t] + b_i) 输入门：新内容写入多少 o_t = σ(W_o [h_{t-1}; x_t] + b_o) 输出门：记忆读出多少 c̃_t = tanh(W_c [h_{t-1}; x_t] + b_c) 候选的新内容 c_t = f_t ∘ c_{t-1} + i_t ∘ c̃_t 细胞状态：**加法**更新 h_t = o_t ∘ tanh(c_t) 关键在 c_t 的更新是"加法"：∂c_t/∂c_{t-1} = diag(f_t)，没有反复乘同一个矩阵 W， 只要遗忘门接近 1，信息和梯度就能沿着细胞状态几乎无损地流过很多步。 演示 1：从零实现 LSTM 单元和 GRU 单元，和 PyTorch 的 nn.LSTMCell / nn.GRUCell 对比 演示 2：初始化时的梯度流动 —— ‖∂h_T/∂x_1‖ 随序列长度 T 的变化（RNN vs LSTM vs GRU） 演示 3：长距离记忆任务 —— "记住第一个词，在 T 步之后说出来"（RNN vs LSTM vs GRU，以及门偏置的影响） 演示 4：信号在状态里能保留多久？—— 用探针比较 LSTM 的细胞状态和 RNN 的隐藏状态 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 02_lstm_gru.py # 完整运行 python 02_lstm_gru.py --quick # 冒烟测试（训练步数很少） 需要 PyTorch。完整运行约 8~12 分钟（CPU，演示 3 要训练 12 个模型）。 """ import os # 必须在 import torch 之前设置：限制底层数学库（OpenMP）的线程数，原因见下面 set_num_threads 处的说明 os.environ.setdefault("OMP_NUM_THREADS", "1") import sys import time import numpy as np import torch import torch.nn as nn import torch.nn.functional as F # 【Python 小知识】这些模型很小、而且是一步一步算的，多线程反而更慢（线程之间的协调开销大于收益， # 实测单线程比 4 线程快约 2 倍），还会和同时运行的其他程序争抢 CPU。所以只用 1 个线程。 torch.set_num_threads(1) def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 第 1 部分：从零实现 LSTM 和 GRU 单元 # ============================================================================= class MyLSTMCell(nn.Module): """ 手写的 LSTM 单元（一个时间步）。 【思路】 把 4 个门的权重叠在一起：一次矩阵乘法算出 4 × H 个数，再切成 i、f、g（候选内容）、o 四段。 （顺序 i, f, g, o 和 PyTorch 的 nn.LSTMCell 一致，方便对比。） 【参数】（构造函数）input_size : 输入维度；hidden : 隐藏状态维度 【forward 的参数】 x : (B, input_size) state : (h, c)，各为 (B, hidden) 【返回】 (h_new, c_new, gates)：gates 是一个字典，包含这一步的 i、f、g、o（用于演示 4 观察门的取值） """ def __init__(self, input_size, hidden): super().__init__() self.hidden = hidden self.W_x = nn.Linear(input_size, 4 * hidden) # 作用在输入 x 上的权重（含偏置） self.W_h = nn.Linear(hidden, 4 * hidden) # 作用在 h_{t-1} 上的权重（含偏置） def forward(self, x, state): h, c = state z = self.W_x(x) + self.W_h(h) i, f, g, o = z.chunk(4, dim=1) # 【Python 小知识】chunk 把张量平均切成几块 i, f, o = torch.sigmoid(i), torch.sigmoid(f), torch.sigmoid(o) g = torch.tanh(g) c_new = f * c + i * g # 加法更新！ h_new = o * torch.tanh(c_new) return h_new, c_new, {"i": i, "f": f, "g": g, "o": o} class MyGRUCell(nn.Module): """ 手写的 GRU 单元（Cho et al. 2014），按 PyTorch nn.GRUCell 的公式： r = σ(W_ir x + b_ir + W_hr h + b_hr) 重置门 z = σ(W_iz x + b_iz + W_hz h + b_hz) 更新门 n = tanh(W_in x + b_in + r ∘ (W_hn h + b_hn)) 候选的新状态 h' = (1 - z) ∘ n + z ∘ h 新旧状态的"插值" 【思路】 GRU 没有单独的细胞状态，用更新门 z 在"保留旧状态"和"换成新内容"之间插值。 z 接近 1 时，h' ≈ h：信息原样保留，梯度也原样传回 —— 和 LSTM 的遗忘门作用相同。 参数比 LSTM 少（3 组而不是 4 组），速度更快。 """ def __init__(self, input_size, hidden): super().__init__() self.W_x = nn.Linear(input_size, 3 * hidden) self.W_h = nn.Linear(hidden, 3 * hidden) def forward(self, x, h): xr, xz, xn = self.W_x(x).chunk(3, dim=1) hr, hz, hn = self.W_h(h).chunk(3, dim=1) r = torch.sigmoid(xr + hr) z = torch.sigmoid(xz + hz) n = torch.tanh(xn + r * hn) return (1 - z) * n + z * h def copy_from_torch(my_cell, torch_cell): """把 PyTorch 单元的权重复制到手写的单元里（两者的门的排列顺序相同）。""" with torch.no_grad(): my_cell.W_x.weight.copy_(torch_cell.weight_ih) my_cell.W_x.bias.copy_(torch_cell.bias_ih) my_cell.W_h.weight.copy_(torch_cell.weight_hh) my_cell.W_h.bias.copy_(torch_cell.bias_hh) def demo_cells(): """演示 1：手写单元与 PyTorch 单元的输出完全一致。""" print_header("演示 1：手写 LSTM / GRU 单元 vs PyTorch") torch.manual_seed(0) B, D, H = 3, 5, 4 x, h, c = torch.randn(B, D), torch.randn(B, H), torch.randn(B, H) ref_lstm, my_lstm = nn.LSTMCell(D, H), MyLSTMCell(D, H) copy_from_torch(my_lstm, ref_lstm) h1, c1 = ref_lstm(x, (h, c)) h2, c2, gates = my_lstm(x, (h, c)) print(f"\n LSTM：h 最大误差 {torch.max(torch.abs(h1 - h2)).item():.1e}，c 最大误差 {torch.max(torch.abs(c1 - c2)).item():.1e}") ref_gru, my_gru = nn.GRUCell(D, H), MyGRUCell(D, H) copy_from_torch(my_gru, ref_gru) print(f" GRU ：h 最大误差 {torch.max(torch.abs(ref_gru(x, h) - my_gru(x, h))).item():.1e}") n_rnn = sum(p.numel() for p in nn.RNNCell(D, H).parameters()) n_gru = sum(p.numel() for p in ref_gru.parameters()) n_lstm = sum(p.numel() for p in ref_lstm.parameters()) print(f" 参数量（输入 {D} 维、隐藏 {H} 维）：RNN {n_rnn}，GRU {n_gru}（3 倍），LSTM {n_lstm}（4 倍）") print(f" 这一步的门的平均取值：输入门 {gates['i'].mean():.2f}，遗忘门 {gates['f'].mean():.2f}，输出门 {gates['o'].mean():.2f}" "（随机初始化时都在 0.5 左右）") # ============================================================================= # 演示 2：初始化时的梯度流动 # ============================================================================= def make_rnn(cell, input_size, hidden, gate_bias=1.0): """ 创建一个单层循环网络，并设置"记忆门"的初始偏置。 【参数】 cell : "rnn"、"lstm" 或 "gru" input_size : 输入维度；hidden : 隐藏维度 gate_bias : LSTM 遗忘门 / GRU 更新门的初始偏置（PyTorch 有两个偏置，各设一半，加起来等于 gate_bias） 【思路】（讲义第 3.3 节） 门的输出是 σ(... + 偏置)。随机初始化时"..."部分接近 0，所以门的初始值约为 σ(偏置)： 偏置 0 → 0.50；偏置 1 → 0.73；偏置 3 → 0.95。 对 LSTM 的遗忘门，0.73 意味着每一步只保留 73% 的记忆，50 步后只剩 0.73⁵⁰ ≈ 10⁻⁷； 0.95 则剩 0.95⁵⁰ ≈ 0.08 —— 信息和梯度都能传过来，模型才有机会学会利用它。 PyTorch 的 GRU 公式是 h' = (1 - z) n + z h，z 是更新门，z 接近 1 表示保留旧状态， 所以对 GRU 也是把 z 的偏置调大。普通 RNN 没有门，这个参数对它无效。 【返回】nn.RNN / nn.LSTM / nn.GRU """ net = {"rnn": nn.RNN, "lstm": nn.LSTM, "gru": nn.GRU}[cell](input_size, hidden, batch_first=True) if cell in ("lstm", "gru"): with torch.no_grad(): for name, p in net.named_parameters(): if "bias" in name: # LSTM 门的顺序是 i, f, g, o；GRU 是 r, z, n。两者的"记忆门"都在第二段 [H, 2H) p[hidden:2 * hidden].fill_(gate_bias / 2) return net def demo_gradient_flow(): """演示 2：‖∂h_T/∂x_1‖ 随 T 的变化（未训练的网络）。""" print_header("演示 2：初始化时，最后的隐藏状态对第一个输入的梯度 ‖∂h_T/∂x_1‖") D, H = 16, 64 Ts = [5, 10, 25, 50, 100] print(f"\n {'模型':<18s}" + "".join(f"{'T=' + str(T):>11s}" for T in Ts)) for cell, gb, label in [("rnn", 0, "RNN"), ("gru", 0, "GRU（偏置 0）"), ("gru", 3, "GRU（偏置 3）"), ("lstm", 1, "LSTM（遗忘偏置 1）"), ("lstm", 3, "LSTM（遗忘偏置 3）")]: row = [] for T in Ts: vals = [] for seed in range(5): torch.manual_seed(seed) net = make_rnn(cell, D, H, gate_bias=gb) x = torch.randn(1, T, D, requires_grad=True) out, _ = net(x) out[0, -1].sum().backward() vals.append(x.grad[0, 0].norm().item()) row.append(np.mean(vals)) print(f" {label:<18s}" + "".join(f"{v:>11.2e}" for v in row)) print(""" 观察（5 个随机初始化取平均）： - RNN 的梯度随 T 指数衰减，100 步后变成 0。 - 门控模型的梯度衰减得慢，而且**取决于记忆门的初始偏置**：偏置越大，门越接近 1，梯度传得越远。 原因是 LSTM 的细胞状态 c_t = f_t ∘ c_{t-1} + ...，反向传播时 ∂c_t/∂c_{t-1} = diag(f_t)， 是逐元素乘以遗忘门，而不是反复乘同一个矩阵 W_h 再乘 tanh 的导数。 （这是在**初始化时**的比较；训练过程中门的取值会变化，见演示 4。）""") # ============================================================================= # 演示 3：长距离记忆任务 # ============================================================================= N_SIGNAL, N_NOISE = 8, 8 def memory_batch(B, T, gen): """ 生成一批"记忆任务"的数据：第一个词是信号（8 种之一），后面跟 T 个随机的噪声词，最后要求说出信号是什么。 【参数】B : 批大小；T : 噪声词的个数（= 信号与"提问"之间的距离）；gen : torch.Generator 【返回】(X, y)：X 形状 (B, T+1)，词编号在 0..15（0~7 是信号，8~15 是噪声）；y 形状 (B,)，取值 0..7 """ y = torch.randint(0, N_SIGNAL, (B,), generator=gen) noise = torch.randint(N_SIGNAL, N_SIGNAL + N_NOISE, (B, T), generator=gen) return torch.cat([y[:, None], noise], dim=1), y class MemoryModel(nn.Module): """嵌入 → 循环层 → 用最后一个隐藏状态分类。""" def __init__(self, cell, hidden=64, gate_bias=1.0): super().__init__() self.emb = nn.Embedding(N_SIGNAL + N_NOISE, 16) self.rnn = make_rnn(cell, 16, hidden, gate_bias) self.out = nn.Linear(hidden, N_SIGNAL) def forward(self, X): h, _ = self.rnn(self.emb(X)) return self.out(h[:, -1]) def train_memory(cell, T, steps, seed=0, gate_bias=1.0): """ 在记忆任务上训练一个模型，返回 (测试准确率, 训练用时)。 【边界】 - 随机猜的准确率是 1/8 = 12.5%。 - 梯度消失时，第一个词几乎收不到梯度，模型学不会，准确率停在 12.5% 左右。 """ torch.manual_seed(seed) gen = torch.Generator().manual_seed(seed) model = MemoryModel(cell, gate_bias=gate_bias) opt = torch.optim.Adam(model.parameters(), lr=3e-3) t0 = time.perf_counter() for _ in range(steps): X, y = memory_batch(128, T, gen) opt.zero_grad() F.cross_entropy(model(X), y).backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() X, y = memory_batch(1000, T, torch.Generator().manual_seed(999)) with torch.no_grad(): acc = (model(X).argmax(1) == y).float().mean().item() return acc, time.perf_counter() - t0, model def demo_memory(steps): """演示 3：记忆任务。""" print_header(f"演示 3：长距离记忆任务（每个模型训练 {steps} 步，批大小 128）") print(""" 输入：[信号, 噪声, 噪声, ..., 噪声]（T 个噪声词），要求在最后一步说出信号是 8 种中的哪一种。 这是最纯粹的"长距离依赖"：要输出的信息只在第一个时间步出现过。随机猜的准确率是 12.5%。 """) Ts = [20, 50, 100] configs = [("rnn", 0, "RNN"), ("lstm", 1, "LSTM（遗忘偏置 1）"), ("lstm", 3, "LSTM（遗忘偏置 3）"), ("gru", 3, "GRU（更新偏置 3）")] print(f" {'模型':<18s}" + "".join(f"{'T=' + str(T):>14s}" for T in Ts)) models = {} for cell, gb, label in configs: row = [] for T in Ts: acc, secs, model = train_memory(cell, T, steps, gate_bias=gb) models[(cell, gb, T)] = model row.append(f"{acc:.1%}({secs:.0f}s)") print(f" {label:<18s}" + "".join(f"{r:>14s}" for r in row)) print(""" 说明：括号里是训练用时（取决于你的电脑）。每个设置只训练了 1 次（随机种子 0）。 观察： - 距离较短（T = 20）时，普通 RNN 也学得会；距离一长（T = 50、100），它的准确率停在随机水平 —— 第一个词带来的梯度要穿过几十次矩阵连乘，早就消失了。 - 门控本身并不自动解决问题：遗忘门偏置为 1 的 LSTM 在 T = 50 上同样学不会。 它的遗忘门初始约为 σ(1) ≈ 0.73，50 步之后信息只剩 0.73⁵⁰ ≈ 10⁻⁷。 - 把记忆门的偏置调到 3（初始约 0.95），LSTM 和 GRU 都能很快学会 T = 50；GRU 甚至能学会 T = 100。 门控结构让"几乎原样保留信息"变得**可能**；合适的初始化让模型一开始就处在这种状态， 于是梯度能传回去，模型才学得会。（后来有工作根据任务需要的记忆长度来设置遗忘门偏置， 叫"chrono 初始化"。具体数字以你运行的结果为准。）""") return models # ============================================================================= # 演示 4：信号能在状态里保留多久？ # ============================================================================= def probe_accuracy(states, y, n_train=300): """ 探针：训练一个线性分类器（softmax 回归），从状态向量中读出信号是 8 种中的哪一种。 【参数】 states : (N, H) 张量，每个序列在某个时间步的状态 y : (N,) 信号类别 n_train: 前多少个样本用于训练，其余用于测试 【返回】float，测试准确率（随机猜是 12.5%） """ torch.manual_seed(0) X = (states - states[:n_train].mean(0)) / (states[:n_train].std(0) + 1e-6) # 标准化 probe = nn.Linear(X.shape[1], N_SIGNAL) opt = torch.optim.Adam(probe.parameters(), lr=0.05) for _ in range(300): opt.zero_grad() F.cross_entropy(probe(X[:n_train]), y[:n_train]).backward() opt.step() with torch.no_grad(): return (probe(X[n_train:]).argmax(1) == y[n_train:]).float().mean().item() def demo_gates(lstm_model, rnn_model, T): """演示 4：信号信息在状态中保留了多久？LSTM 的门取什么值？""" print_header(f"演示 4：信号在状态里能保留多久？（记忆任务 T = {T}，训练好的模型）") lstm = lstm_model.rnn cell = MyLSTMCell(16, 64) with torch.no_grad(): cell.W_x.weight.copy_(lstm.weight_ih_l0) cell.W_x.bias.copy_(lstm.bias_ih_l0) cell.W_h.weight.copy_(lstm.weight_hh_l0) cell.W_h.bias.copy_(lstm.bias_hh_l0) X, y = memory_batch(600, T, torch.Generator().manual_seed(5)) with torch.no_grad(): e = lstm_model.emb(X) h, c = torch.zeros(600, 64), torch.zeros(600, 64) cs, fs, iss = [], [], [] for t in range(T + 1): h, c, g = cell(e[:, t], (h, c)) cs.append(c.clone()) fs.append(g["f"].mean().item()) iss.append(g["i"].mean().item()) acc = (lstm_model.out(h).argmax(1) == y).float().mean().item() rnn_h, _ = rnn_model.rnn(rnn_model.emb(X)) # (600, T+1, 64) print(f"\n 用手写单元重放 LSTM 的准确率：{acc:.1%}（应该和 PyTorch 模型的结果一致）") print(" 在每个时间步训练一个线性探针，从状态向量中读出信号（300 个序列训练、300 个测试；随机猜 12.5%）：\n") print(f" {'时间步':<8s}{'LSTM 细胞状态 c_t':>18s}{'RNN 隐藏状态 h_t':>18s}{'LSTM 平均遗忘门':>16s}{'LSTM 平均输入门':>16s}") for t in [0, 1, 2, 5, 10, 20, T // 2 if T // 2 not in (10, 20) else T - 10, T]: a_lstm = probe_accuracy(cs[t], y) a_rnn = probe_accuracy(rnn_h[:, t], y) tag = " ← 读到信号" if t == 0 else "" print(f" t = {t:<4d}{a_lstm:>18.1%}{a_rnn:>18.1%}{fs[t]:>16.3f}{iss[t]:>16.3f}{tag}") print(""" 观察： - LSTM 的细胞状态在每一步都能几乎完美地读出信号 —— 信息从 t = 0 一直保留到最后。 它的遗忘门平均值在 0.95 以上，每一步只"忘掉"很少一点。 - 没学会这个任务的 RNN：刚读到信号时（t = 0）当然能读出来，但几步之后探针准确率就迅速下降， 很快接近随机猜 —— 信号被后面的噪声词一步步"冲掉"了。RNN 在前向传播中保不住信息， 在反向传播中梯度也传不回去，两者是同一件事的两面。 （具体数值以你运行的结果为准。）""") def main(quick=False): demo_cells() demo_gradient_flow() models = demo_memory(steps=150 if quick else 2000) demo_gates(models[("lstm", 3, 50)], models[("rnn", 0, 50)], 50) if __name__ == "__main__": main(quick="--quick" in sys.argv)