查看全课大纲（第 16 / 40 节）

LSTM/GRU：门控如何缓解梯度问题

LSTM/GRU：门控如何缓解梯度问题？

小象实战讲义 · 深度学习基础：从神经网络到 Transformer

上一节我们学习了循环神经网络（RNN）的基本结构，它通过引入循环连接，让网络具备了处理序列数据的能力。然而，标准的RNN在训练长序列时，会遭遇严重的梯度消失问题，导致其难以学习长距离的依赖关系。本节我们将深入探讨两种革命性的门控循环单元——LSTM和GRU，揭示它们如何通过精巧的“门”机制，让梯度在时间轴上“流动”得更远，从而成为处理长序列任务的基石。学完本节，你将能回答：为什么LSTM能记住更久远的信息？遗忘门、输入门、输出门各自在做什么？GRU相比LSTM做了哪些简化？

💡 核心导读

LSTM的核心是“细胞状态”：一条贯穿时间的“高速公路”，信息可以几乎无衰减地流过，由三个门（遗忘、输入、输出）控制信息的读写。
门控机制是梯度“守护者”：细胞状态采用加法式更新（ $C_t = f_t \odot C_{t-1} + \dots$ ），反向传播时梯度沿这条路径只需乘以遗忘门 $f_t$ ； $f_t$ 越接近 1，梯度衰减越慢——不像标准 RNN 那样每步都要乘 $\tanh'$ （<1）和 $W$ 。
GRU是LSTM的简化版：它将LSTM的三个门合并为两个（重置门和更新门），并将细胞状态与隐状态合并，在多数任务上表现相当且计算更高效。
代码实现揭示形状流动：我们将从零实现LSTM单元的前向传播，并与PyTorch官方实现对比，理解(batch, seq, hidden)形状的张量如何流经门控。

🧭 前置知识

RNN的基本结构与按时间展开（见上一节）。
矩阵乘法形状规则：(batch, dim1) @ (dim1, dim2) -> (batch, dim2)。
Sigmoid和Tanh激活函数的特性（输出范围、导数）。
链式法则与梯度传播的基本概念。

LSTM：长短期记忆网络

标准RNN的隐状态更新公式为 $h_t = \tanh(W h_{t-1} + U x_t + b)$ 。在反向传播时，梯度需要连续乘以 $\frac{\partial h_t}{\partial h_{t-1}} = \text{diag}(\tanh'(a_t)) W$ 。由于 $\tanh'(a_t)$ 的绝对值通常小于1，且 $W$ 的谱范数也可能带来缩放，梯度在沿时间反向传播时会指数级衰减或爆炸，这就是梯度消失/爆炸问题。

LSTM（Long Short-Term Memory）由 Sepp Hochreiter 和 Jürgen Schmidhuber 于1997年提出，其核心思想是引入一个细胞状态（cell state） $C_t$ 作为贯穿时间的“信息高速公路”，并设计三个“门”来精细调控信息的流动。

LSTM 单元的结构

一个LSTM单元在每个时间步 $t$ 接收以下输入：

当前输入 $x_t$ （形状 (batch, input_size)）
上一时刻的隐状态 $h_{t-1}$ （形状 (batch, hidden_size)）
上一时刻的细胞状态 $C_{t-1}$ （形状 (batch, hidden_size)）

它输出：

当前隐状态 $h_t$ （形状 (batch, hidden_size)）
当前细胞状态 $C_t$ （形状 (batch, hidden_size)）

其内部计算涉及四个并行的全连接层（共享输入），产生四个向量：

遗忘门（Forget Gate） $f_t = \sigma(W_f \cdot [h_{t-1}, x_t] + b_f)$
- 决定从旧细胞状态 $C_{t-1}$ 中丢弃多少信息。 $\sigma$ 是sigmoid函数，输出在[0,1]之间。
输入门（Input Gate） $i_t = \sigma(W_i \cdot [h_{t-1}, x_t] + b_i)$
- 决定有多少新信息将被存入细胞状态。
候选细胞状态（Candidate Cell State） $\tilde{C}_t = \tanh(W_C \cdot [h_{t-1}, x_t] + b_C)$
- 根据当前输入和上一隐状态生成的新信息。
输出门（Output Gate） $o_t = \sigma(W_o \cdot [h_{t-1}, x_t] + b_o)$
- 决定基于当前细胞状态，输出多少信息到隐状态 $h_t$ 。

最终，细胞状态和隐状态的更新公式为： $C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t$ $h_t = o_t \odot \tanh(C_t)$

其中 $\odot$ 表示逐元素相乘（Hadamard积）。

手动实现与形状验证

让我们用代码实现一个时间步的LSTM计算，并验证其形状变化。

import torch
import torch.nn as nn
torch.manual_seed(0) # 固定随机种子，确保结果可复现

# 定义超参数
batch_size = 2
input_size = 3
hidden_size = 5

# 初始化参数（模拟一个LSTM单元）
def init_lstm_params(input_size, hidden_size):
    # 将四个门的权重矩阵拼接在一起计算，这是常见的实现方式。
    # 权重形状: (4 * hidden_size, input_size + hidden_size)
    # 偏置形状: (4 * hidden_size,)
    weight = torch.randn(4 * hidden_size, input_size + hidden_size) * 0.01
    bias = torch.zeros(4 * hidden_size)
    return weight, bias

weight, bias = init_lstm_params(input_size, hidden_size)
print(f"权重矩阵 W 形状: {tuple(weight.shape)}") # 应输出 (20, 8)
print(f"偏置向量 b 形状: {tuple(bias.shape)}")   # 应输出 (20,)

# 构造输入和前一时刻状态
x_t = torch.randn(batch_size, input_size)          # 当前输入
h_prev = torch.zeros(batch_size, hidden_size)      # 上一时刻隐状态
C_prev = torch.zeros(batch_size, hidden_size)      # 上一时刻细胞状态

# 手动前向传播一个LSTM时间步
def lstm_cell_forward(x_t, h_prev, C_prev, weight, bias):
    # 1. 拼接输入和上一隐状态
    combined = torch.cat((h_prev, x_t), dim=1)  # 形状: (batch, hidden_size + input_size)
    # 2. 计算所有门的激活值（一次矩阵乘法）
    gates = combined @ weight.T + bias          # 形状: (batch, 4 * hidden_size)
    # 3. 分割为四个部分
    split_size = gates.shape[1] // 4
    i_t, f_t, g_t, o_t = torch.split(gates, split_size, dim=1)
    # 4. 应用激活函数
    i_t = torch.sigmoid(i_t)  # 输入门
    f_t = torch.sigmoid(f_t)  # 遗忘门
    g_t = torch.tanh(g_t)     # 候选值
    o_t = torch.sigmoid(o_t)  # 输出门
    # 5. 更新细胞状态和隐状态
    C_t = f_t * C_prev + i_t * g_t
    h_t = o_t * torch.tanh(C_t)
    return h_t, C_t, (i_t, f_t, g_t, o_t) # 返回状态和门值（用于观察）

h_t, C_t, gates = lstm_cell_forward(x_t, h_prev, C_prev, weight, bias)
print(f"输入 x_t 形状: {tuple(x_t.shape)}")
print(f"输出隐状态 h_t 形状: {tuple(h_t.shape)}")
print(f"输出细胞状态 C_t 形状: {tuple(C_t.shape)}")
print(f"遗忘门 f_t 的均值（期望接近0.5）: {gates[1].mean().item():.4f}")

权重矩阵 W 形状: (20, 8)
偏置向量 b 形状: (20,)
输入 x_t 形状: (2, 3)
输出隐状态 h_t 形状: (2, 5)
输出细胞状态 C_t 形状: (2, 5)
遗忘门 f_t 的均值（期望接近0.5）: 0.5005

门控如何缓解梯度消失？

LSTM缓解梯度消失的关键在于细胞状态 $C_t$ 的更新路径。观察更新公式： $C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t$

在反向传播时，损失函数 $L$ 对 $C_{t-1}$ 的梯度有两条路径：

通过 $C_t$ 直接依赖于 $C_{t-1}$ （第一项）。
通过 $h_t$ 依赖于 $C_t$ ，再间接依赖于 $C_{t-1}$ 。

我们重点关注第一条直接路径的梯度。根据链式法则： $\frac{\partial C_t}{\partial C_{t-1}} = \text{diag}(f_t) + \text{（一些其他项）}$

这里 $\text{diag}(f_t)$ 是一个以遗忘门向量 $f_t$ 为对角元素的对角矩阵。关键在于 $f_t$ 越接近 1，这条直接路径上的梯度衰减越慢。早期 TensorFlow 的实现曾把遗忘门偏置默认初始化为 1.0（forget_bias=1.0），让训练初期的 $f_t \approx 0.73$ 、更倾向保留旧信息；PyTorch 没有这个特殊默认（偏置是小随机数，初始 $f_t \approx 0.5$ ），但在训练过程中学到的 $f_t$ 常常会靠近 1。

在沿时间反向传播梯度时，梯度 $\frac{\partial L}{\partial C_{t-k}}$ 会连续乘以这些接近 1 的因子，从而极大地减缓了梯度的指数衰减。细胞状态 $C$ 充当了一条梯度可以近乎无损流动的“高速公路”。

相比之下，标准RNN中 $\frac{\partial h_t}{\partial h_{t-1}}$ 包含 $\tanh'$ （绝对值小于1）和权重矩阵 $W$ ，梯度衰减要快得多。

梯度流动模拟

让我们模拟一个简化场景，观察不同遗忘门偏置设置下，梯度沿时间步的传播情况。

import torch
torch.manual_seed(0)

def simulate_gradient_flow(seq_len, forget_bias):
    """
    模拟梯度沿时间步的传播。
    forget_bias: 遗忘门偏置，影响遗忘门的值。
    """
    # 假设每个时间步的遗忘门值固定为 sigmoid(forget_bias)
    f = torch.sigmoid(torch.tensor(forget_bias)).item()
    print(f"遗忘门固定值 f ≈ {f:.4f}")

    # 初始梯度（假设损失对最终细胞状态的导数为1）
    grad = 1.0
    grads_over_time = [grad]
    for step in range(seq_len):
        grad = grad * f  # 梯度乘以遗忘门系数
        grads_over_time.append(grad)
    return grads_over_time

# 模拟序列长度
seq_len = 10
print(f"序列长度: {seq_len}")
print("-" * 30)

# 情况1：遗忘门偏置为0 (f ≈ 0.5)
print("情况1: forget_bias = 0.0")
grads1 = simulate_gradient_flow(seq_len, 0.0)
print(f"  初始梯度到第{seq_len}步的梯度: {grads1[0]:.4f} -> {grads1[-1]:.4f}")
print(f"  衰减系数: {grads1[-1]/grads1[0]:.6f}")
print()

# 情况2：遗忘门偏置为1 (f ≈ 0.73)
print("情况2: forget_bias = 1.0")
grads2 = simulate_gradient_flow(seq_len, 1.0)
print(f"  初始梯度到第{seq_len}步的梯度: {grads2[0]:.4f} -> {grads2[-1]:.4f}")
print(f"  衰减系数: {grads2[-1]/grads2[0]:.6f}")
print()

# 情况3：遗忘门偏置为2 (f ≈ 0.88)
print("情况3: forget_bias = 2.0")
grads3 = simulate_gradient_flow(seq_len, 2.0)
print(f"  初始梯度到第{seq_len}步的梯度: {grads3[0]:.4f} -> {grads3[-1]:.4f}")
print(f"  衰减系数: {grads3[-1]/grads3[0]:.6f}")

序列长度: 10
------------------------------
情况1: forget_bias = 0.0
遗忘门固定值 f ≈ 0.5000
  初始梯度到第10步的梯度: 1.0000 -> 0.0010
  衰减系数: 0.000977

情况2: forget_bias = 1.0
遗忘门固定值 f ≈ 0.7311
  初始梯度到第10步的梯度: 1.0000 -> 0.0436
  衰减系数: 0.043604

情况3: forget_bias = 2.0
遗忘门固定值 f ≈ 0.8808
  初始梯度到第10步的梯度: 1.0000 -> 0.2810
  衰减系数: 0.281034

模拟清晰地显示，把遗忘门偏置设得更大（如 1 或 2）、让 $f$ 更接近 1，能显著减缓梯度随时间的衰减：经过 10 个时间步， $f \approx 0.88$ 时梯度仍保留约 28%，而 $f=0.5$ 时只剩不到 0.1%。这正是 LSTM 能学习长距离依赖的数学本质。

要强调的是，这是为了看清机制而做的受控实验（人为固定 $f$ ），不是 LSTM 的默认配置：真实模型的门值由数据学出来，能学出「该记多久就记多久」的自适应遗忘。

GRU：门控循环单元

GRU（Gated Recurrent Unit）由 Cho 等人在2014年提出，是LSTM的一种简化变体。它将LSTM的三个门减少为两个，并合并了细胞状态和隐状态。

GRU 单元的结构

GRU单元在每个时间步 $t$ 的计算如下：

更新门（Update Gate）： $z_t = \sigma(W_z \cdot [h_{t-1}, x_t] + b_z)$
- 控制有多少旧信息被保留，多少新信息被加入。类似于LSTM中遗忘门和输入门的结合。
重置门（Reset Gate）： $r_t = \sigma(W_r \cdot [h_{t-1}, x_t] + b_r)$
- 控制上一时刻隐状态 $h_{t-1}$ 有多少信息被用于计算候选状态。
候选隐状态（Candidate Hidden State）： $\tilde{h}_t = \tanh(W \cdot [r_t \odot h_{t-1}, x_t] + b)$
- 基于重置门过滤后的上一状态和当前输入计算的新状态。
最终隐状态（Final Hidden State）： $h_t = (1 - z_t) \odot h_{t-1} + z_t \odot \tilde{h}_t$
- 这是GRU的核心更新公式，是旧状态和新候选状态的加权和。

GRU 的梯度流分析

观察GRU的更新公式 $h_t = (1 - z_t) \odot h_{t-1} + z_t \odot \tilde{h}_t$ 。在反向传播时， $h_t$ 对 $h_{t-1}$ 的梯度也包含一条直接路径： $\frac{\partial h_t}{\partial h_{t-1}}$ 包含 $(1 - z_t)$ 项。与 LSTM 的 $f_t$ 对称， $(1 - z_t)$ 越接近 1，这条路径上的梯度衰减越慢。如果希望训练初期多保留旧状态，可以把更新门偏置 $b_z$ 设成负值（例如 -1，使 $z_t \approx 0.27$ 、 $1-z_t \approx 0.73$ ）；但 PyTorch 默认不做这种偏置设置，实际效果主要靠训练中学到的门值。

import torch
import torch.nn as nn
torch.manual_seed(0)

# 定义超参数
batch_size = 2
input_size = 3
hidden_size = 5

# 使用 PyTorch 的 GRUCell
gru_cell_pt = nn.GRUCell(input_size, hidden_size)

# 构造输入和初始状态
x_t = torch.randn(batch_size, input_size)
h_prev_gru = torch.zeros(batch_size, hidden_size)

# GRUCell 前向传播
h_gru = gru_cell_pt(x_t, h_prev_gru)

print(f"GRU 输入 x_t 形状: {tuple(x_t.shape)}")
print(f"GRU 输入隐状态 h_prev 形状: {tuple(h_prev_gru.shape)}")
print(f"GRU 输出隐状态 h_gru 形状: {tuple(h_gru.shape)}")
print(f"GRU 参数总数: {sum(p.numel() for p in gru_cell_pt.parameters())}")

GRU 输入 x_t 形状: (2, 3)
GRU 输入隐状态 h_prev 形状: (2, 5)
GRU 输出隐状态 h_gru 形状: (2, 5)
GRU 参数总数: 150

LSTM vs. GRU 对比

特性	LSTM	GRU
门数量	3个 (遗忘门 $f$ , 输入门 $i$ , 输出门 $o$ )	2个 (更新门 $z$ , 重置门 $r$ )
状态数量	2个 (隐状态 $h$ , 细胞状态 $C$ )	1个 (隐状态 $h$ )
核心更新公式	$C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t$ $h_t = o_t \odot \tanh(C_t)$	$h_t = (1 - z_t) \odot h_{t-1} + z_t \odot \tilde{h}_t$
参数数量	较多 ( $4 \times (input+hidden) \times hidden$ )	较少 ( $3 \times (input+hidden) \times hidden$ )
计算效率	稍慢	稍快
常见表现	在需要精细控制记忆的任务上可能略有优势	在许多任务上与LSTM表现相当，是常用的默认选择

使用 PyTorch 的 LSTM 模块

在实际项目中，我们总是使用高度优化过的框架内置模块。PyTorch提供了 nn.LSTM，它封装了多层、双向等复杂功能。

import torch
import torch.nn as nn
torch.manual_seed(0)

# 定义网络参数
batch_size = 2
seq_len = 4
input_size = 3
hidden_size1 = 5
hidden_size2 = 3
num_layers = 2

# 构造一个2层LSTM
lstm = nn.LSTM(input_size=input_size,
               hidden_size=hidden_size1,
               num_layers=num_layers,
               batch_first=True, # 输入形状为 (batch, seq, feature)
               dropout=0.0) # 训练时可启用dropout防止过拟合

# 构造输入序列 (batch, seq_len, input_size)
input_seq = torch.randn(batch_size, seq_len, input_size)

# 初始隐状态和细胞状态 (num_layers, batch, hidden_size)
# 如果不提供，默认为零
h0 = torch.zeros(num_layers, batch_size, hidden_size1)
c0 = torch.zeros(num_layers, batch_size, hidden_size1)

# 前向传播
output, (h_n, c_n) = lstm(input_seq, (h0, c0))

print(f"输入序列形状: {tuple(input_seq.shape)}")
print(f"LSTM输出序列形状 output: {tuple(output.shape)}")
print(f"最后时刻所有层的隐状态 h_n 形状: {tuple(h_n.shape)}")
print(f"最后时刻所有层的细胞状态 c_n 形状: {tuple(c_n.shape)}")
print(f"输出序列 output 的最后一个时间步 == h_n 的最后一层: {torch.allclose(output[:, -1, :], h_n[-1, :, :])}")

输入序列形状: (2, 4, 3)
LSTM输出序列形状 output: (2, 4, 5)
最后时刻所有层的隐状态 h_n 形状: (2, 2, 5)
最后时刻所有层的细胞状态 c_n 形状: (2, 2, 5)
输出序列 output 的最后一个时间步 == h_n 的最后一层: True

形状解读：

output: 形状为 (batch, seq_len, hidden_size)，包含了最后一层LSTM每个时间步的隐状态 $h_t$ 。这对于需要每个时间步输出的任务（如序列标注）非常有用。
h_n: 形状为 (num_layers, batch, hidden_size)，是所有层在最后一个时间步的隐状态。
c_n: 形状为 (num_layers, batch, hidden_size)，是所有层在最后一个时间步的细胞状态。

📝 动手练一练

题目：请修改上面的代码，构造一个双向双层LSTM（bidirectional=True），并打印其输出和最终状态的形状。观察形状与单向LSTM有何不同，并思考双向LSTM的输出维度为什么是 hidden_size * 2。

👉 点击查看参考答案

import torch
import torch.nn as nn
torch.manual_seed(0)

batch_size = 2
seq_len = 4
input_size = 3
hidden_size = 5
num_layers = 2

# 构造双向双层LSTM
bilstm = nn.LSTM(input_size=input_size,
                 hidden_size=hidden_size,
                 num_layers=num_layers,
                 batch_first=True,
                 bidirectional=True)

input_seq = torch.randn(batch_size, seq_len, input_size)
output_bi, (h_n_bi, c_n_bi) = bilstm(input_seq)

print(f"双向LSTM输出形状 output_bi: {tuple(output_bi.shape)}")
print(f"双向LSTM最后时刻隐状态 h_n_bi 形状: {tuple(h_n_bi.shape)}")
print(f"双向LSTM最后时刻细胞状态 c_n_bi 形状: {tuple(c_n_bi.shape)}")

双向LSTM输出形状 output_bi: (2, 4, 10)
双向LSTM最后时刻隐状态 h_n_bi 形状: (4, 2, 5)
双向LSTM最后时刻细胞状态 c_n_bi 形状: (4, 2, 5)

思路与解释：

输出形状：(2, 4, 10)。因为 bidirectional=True，每个时间步的隐状态是前向和后向LSTM隐状态的拼接，所以特征维度是 hidden_size * 2 = 10。
最终状态形状：(4, 2, 5)。第一维 4 对应 num_layers * num_directions = 2 * 2 = 4，排列顺序是「第 1 层前向、第 1 层后向、第 2 层前向、第 2 层后向」（同一层的前向/后向相邻）。注意：后向 LSTM 在 h_n 里的「最终」隐状态对应原序列的开头（它从末尾反向处理到第一个位置），所以是 output[:, 0, hidden_size:] 与 h_n 的后向部分对齐；而 output 最后一个时间步里的后向分量只包含「从末尾那一步出发」的短程信息。

常见错误

错误理解或写法	现象或报错	原因	正确做法
错误理解LSTM输出`output`和`h_n`的关系	误以为`output`包含所有层所有时间步的隐状态，或取`h_n`做序列任务。	`output`仅包含最后一层所有时间步的 $h_t$ ；`h_n`包含所有层最后一个时间步的 $h_t$ 。	需要每个时间步的预测时，使用`output`；只需要序列的汇总信息时，可使用`h_n[-1]`（最后一层最后时刻）。
手动实现LSTM时忘记拼接 $h_{t-1}$ 和 $x_t$	计算出的门控值形状错误，或无法进行矩阵乘法。	四个门的计算需要同时看到上一隐状态和当前输入。	使用`torch.cat((h_prev, x_t), dim=1)`得到组合输入。
在`nn.LSTM`中混淆`batch_first`参数	输入张量形状与模型期望不匹配，导致运行时错误。	PyTorch默认LSTM输入形状为`(seq_len, batch, input_size)`。	若数据形状为`(batch, seq_len, features)`，设置`batch_first=True`。
未正确初始化LSTM/GRU的隐状态，或跨 batch 复用隐状态	训练结果不稳定或难以收敛。	隐状态里装着「上一个序列」的上下文；跨 batch 复用等于把两个无关样本的记忆混在一起。	每个 batch（或每个序列）开始时把 $h_0, C_0$ 重新置零（`torch.zeros(num_layers, batch, hidden_size)`）；零初始化就是标准做法，不需要额外随机化。
将梯度消失问题完全归咎于激活函数tanh	认为LSTM解决了梯度消失只是因为换了激活函数。	标准RNN梯度消失的主因是 $\tanh'$ 与 $W$ 的逐时间步连乘。LSTM 的关键是加法式的细胞状态更新给出了一条只乘 $f_t$ 的直接梯度路径，而不是「换了激活函数」。	理解门控机制对梯度流的保护作用，而不是把 LSTM 当成「tanh 换成 sigmoid」的改进版。

本章小结

LSTM通过引入细胞状态 $C_t$ 和三个门（遗忘、输入、输出），创造了一条梯度可以长期流动的路径，有效缓解了RNN的梯度消失问题。
遗忘门是长距离记忆的关键。通过将其偏置初始化为正数（如1），使门值接近1，保证了梯度在沿时间反向传播时衰减缓慢。
GRU是LSTM的简化高效变体，它将细胞状态与隐状态合并，使用更新门和重置门，参数更少，在许多任务上表现与LSTM相当。
PyTorch的nn.LSTM模块封装了完整功能（多层、双向），输出包含序列输出output和最终状态(h_n, c_n)，需根据任务需求正确选取。
双向RNN通过结合前向和后向信息，能更好地捕捉序列的上下文，输出维度是单向的两倍。

📋 行动清单

运行本节第一个代码块，尝试修改forget_bias的初始化值（如改为-1），观察遗忘门均值的变化，并思考这对记忆能力的影响。
将“动手练一练”中的双向LSTM改为单向（bidirectional=False），再次运行，对比输出形状的差异，加深对双向机制的理解。
（选做）查阅PyTorch文档中nn.LSTMCell和nn.GRUCell的用法，尝试用Cell版本手动循环实现一个序列的前向传播。