第 4 章：LSTM 长短期记忆网络 —— 门控机制、细胞状态与遗忘门/输入门/输出门

说到 RNN，你是不是也有这种感觉？
明明刚读完句子前半段，后半段就开始“失忆”了。
梯度消失，说白了就是网络记不住远距离的依赖。

我刚开始碰 LSTM 时，总觉得三个门控机制挺玄乎的。但拆开来看，
其实就是给 RNN 装了一套“智能阀门系统”。

4.1 从 RNN 到 LSTM：为什么需要长短期记忆？

普通 RNN 在处理长序列时，梯度很容易在反向传播中指数级衰减。
你想想看，一个分子 SMILES 序列可能有几十甚至上百个字符，
让网络记住 50 步之前的原子信息，基本不可能。

我在项目中遇到过用普通 RNN 做分子性质预测，结果跑了三天三夜，
损失就是不降。后来发现网络只学到了最近 3-4 个字符的局部模式。
遇到长距离依赖的任务，比如预测分子骨架的末端官能团，它就彻底懵了。

核心痛点：标准 RNN 只能依赖短期上下文，无法处理时间跨度大的依赖关系。 LSTM 通过引入细胞状态和门控机制，把“记忆有效期”从几步延长到上百步。

4.2 细胞状态：信息传递的“高速公路”

细胞状态是 LSTM 的核心。你可以把它想象成一条贯穿整个序列的传送带。
信息沿着这条传送带一路向前传，只有在门控打开的地方才会被修改。

我个人习惯把细胞状态比作河流的主航道。门控就是水闸，
决定让哪些支流汇入主航道，又让哪些水排出去。
主航道的水流一直保持稳定——这就是 LSTM 能避免梯度消失的秘密。

概念类比作用
细胞状态河流主航道长期记忆的载体，信息流基本不变
遗忘门排水闸决定丢弃哪些旧信息
输入门进水闸决定存入哪些新信息
输出门取水闸基于当前状态输出有用信息

4.3 三扇门的精妙协作

4.3.1 遗忘门：决定忘记什么

遗忘门是个 Sigmoid 层。输入是当前的 x_t 和上一时刻的隐藏状态 h_{t-1}。
输出是一个 0 到 1 之间的向量。

f_t = σ(W_f · [h_{t-1}, x_t] + b_f)

输出值接近 1 表示“完全保留”，接近 0 表示“完全丢弃”。
嗯，这里要注意：遗忘门的输入是拼接向量 [h_{t-1}, x_t]，
很多初学者会遗漏这个拼接步骤。

我曾经写过一个 LSTM 做分子生成，忘记拼接隐藏状态和输入，
结果网络直接退化成了普通感知器。排查了一整天，最后发现遗忘了遗忘门的输入。

遗忘门的"遗忘"不是删除整段记忆，而是按元素缩放细胞状态。每个维度独立决定保留多少。

4.3.2 输入门：决定记住什么

输入门由两部分组成：
1) Sigmoid 层：决定更新哪些值。
2) Tanh 层：生成候选细胞状态 Ĉ_t。

i_t = σ(W_i · [h_{t-1}, x_t] + b_i)
Ĉ_t = tanh(W_C · [h_{t-1}, x_t] + b_C)

最终新的细胞状态 C_t 由遗忘门和输入门共同决定：

C_t = f_t ⊙ C_{t-1} + i_t ⊙ Ĉ_t

你看，这里用的是逐元素乘法（⊙）。
旧的记忆被遗忘门过滤，新的候选记忆被输入门筛选。
两者相加，就是更新后的细胞状态了。

4.3.3 输出门：决定输出什么

输出门决定了当前时刻的隐藏状态 h_t。
细胞状态先经过 Tanh 压缩到 [-1, 1] 范围，再与输出门的结果相乘：

o_t = σ(W_o · [h_{t-1}, x_t] + b_o)
h_t = o_t ⊙ tanh(C_t)

输出门可以理解为：根据当前任务的需要，从细胞状态中提取有用信息。
在分子生成任务中，这一步直接决定了下一个原子是什么。

4.4 LSTM 的完整前向传播过程

我个人建议把公式按时间步串联起来理解。给定当前输入 x_t 和上一状态 h_{t-1}, C_{t-1}：

  1. 遗忘门计算：f_t = sigmoid(W_f * [h_{t-1}, x_t] + b_f)
  2. 输入门计算：i_t = sigmoid(W_i * [h_{t-1}, x_t] + b_i)
  3. 候选状态计算：Ĉ_t = tanh(W_C * [h_{t-1}, x_t] + b_C)
  4. 更新细胞状态：C_t = f_t * C_{t-1} + i_t * Ĉ_t
  5. 输出门计算：o_t = sigmoid(W_o * [h_{t-1}, x_t] + b_o)
  6. 生成隐藏状态：h_t = o_t * tanh(C_t)
注意参数的维度一致性！假设 x_t 维度为 D，h_t 维度为 H，那么 [h_{t-1}, x_t] 的维度为 D+H。权重矩阵 W_f 的形状是 [H, D+H]，偏置 b_f 的形状是 [H]。写代码时最容易踩的坑就是矩阵乘法形状不匹配。

4.5 LSTM 的门控结构图

下面这张图，我把它叫做“LSTM 记忆管理三件套”。
箭头表示数据流向，中间的圆圈就是那三个关键的阀门。

LSTM 记忆管理三件套 —— 门控与细胞状态协作 C_{t-1} h_{t-1} 遗忘门 f_t = σ(...) 输入门 i_t = σ(...) 输出门 o_t = σ(...) C_t h_t x_t 数据流向 实线：门控输出 虚线：输入提供给门控 C_t 传承长期记忆

4.6 代码怎么敲？

光看公式不写代码，永远学不会。下面是一个用 PyTorch 手写的 LSTM 单元：

import torch
import torch.nn as nn

class LSTMCell(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.hidden_dim = hidden_dim

        # 把四个门的权重合并在一个矩阵里，并行计算效率更高
        self.W = nn.Linear(input_dim + hidden_dim, 4 * hidden_dim)

    def forward(self, x, h_prev, c_prev):
        # 拼接输入
        combined = torch.cat([h_prev, x], dim=1)

        # 计算四个门
        gates = self.W(combined)
        i, f, o, g = gates.chunk(4, dim=1)

        # 激活
        i = torch.sigmoid(i)   # 输入门
        f = torch.sigmoid(f)   # 遗忘门
        o = torch.sigmoid(o)   # 输出门
        g = torch.tanh(g)      # 候选细胞状态

        # 更新细胞状态和隐藏状态
        c_new = f * c_prev + i * g
        h_new = o * torch.tanh(c_new)

        return h_new, c_new
参数合并在一个线性层里，同时计算四组变换，可以减少矩阵乘法的次数。这是工业界最常用的优化手段。

4.7 在分子生成中，LSTM 具体怎么用？

分子生成本质上是一个字符级别的时间序列生成任务。
分子式或 SMILES 序列作为输入，LSTM 每一步输出一个字符。

具体流程：

  • 输入层：将分子字符（如 "C", "N", "O"）转换为独热编码或词嵌入。
  • LSTM 层：负责捕捉序列中的长程依赖。比如苯环的结构模式。
  • 输出层：用 Softmax 预测下一个字符的概率分布。

我记得有一次优化分子生成模型，把 LSTM 的隐藏维度从 128 提到 256，
生成的分子有效比例直接从 62% 涨到 84%。当然，训练时间也翻了一倍。
你想想看，隐藏维度翻倍，权重矩阵的参数量翻了 4 倍。

LSTM 参数量 = 4 × (hidden_dim × (input_dim + hidden_dim) + hidden_dim) 公式里那个 4 就是来自三个门加一个候选状态，总共四组权重。

4.8 避坑指南

我踩过的坑，希望你别踩：

  • 梯度爆炸：门控可以缓解梯度消失，但无法防止梯度爆炸。一定要加梯度裁剪。
  • 初始状态：不要初始化隐藏状态为全零。用小的均匀分布初始化效果更好。
  • 序列长度：在分子生成中，设置最大序列长度。防止模型一直生成无效分子。
  • Dropout：别在循环连接上加 Dropout，PyTorch 的 LSTM 默认就不允许。

我曾经在一个项目里忘了做梯度裁剪，训练三个小时后面损失突然变成 NaN。
排查半天发现是梯度爆炸导致权重溢出。从那以后，我每次写 RNN/LSTM 代码，
第一件事就是把梯度裁剪写进去。

# 正确的梯度裁剪写法 —— 每个批次反向传播后调用
nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

好了，关于 LSTM 的核心概念和实战注意点，我就讲到这里。
你如果有机会上手跑一跑这个代码，一定会有更深的体会。
算法看起来复杂，但写出来其实就是那几行代码的配合。

← 上一章 📖 返回目录 下一章 →