3. RNN的梯度消失与梯度爆炸
RNN 训练不起来，八成是梯度出问题了。
我在做分子生成时，第一个模型跑了一周，loss 纹丝不动。检查半天，梯度已经小到可以忽略不计。这就是经典的梯度消失。
这一章，咱们把梯度消失和梯度爆炸的根因、解决手段，一次性讲透。
3.1 为什么RNN会梯度消失或爆炸？
原因藏在时间轴里。
简单说，RNN 在时间步 t 的隐状态 h_t，是由 h_{t-1} 和 x_t 算出来的。反向传播时，梯度要沿着时间步一层一层传回去。传到第 1 步，梯度可能会变得非常小（消失）或非常大（爆炸）。
从数学角度看，这是因为梯度中包含连乘项。你想想看，假设每个时间步的 Jacobian 矩阵的范数都小于 1，连乘几十步，梯度就接近 0 了。反之，如果大于 1，梯度会指数级增长。
- 深度时间展开导致的连乘效应
- 激活函数选择不当（如 sigmoid/tanh）会加剧梯度缩小
- 权重矩阵初始化不当导致爆炸或消失
我曾经在一个 50 步的序列长度上试过 sigmoid + 随机初始化。结果训练 100 轮，loss 几乎无变化。这就是活生生的梯度消失案例。
这里给出一个简单的 RNN 正向传播和梯度更新的示意公式：
h_t = tanh(W_hh * h_{t-1} + W_xh * x_t + b_h)
o_t = W_hy * h_t + b_o
损失 L = sum over t of (o_t - y_t)^2
反向梯度：
dL/dW_hh = sum over t of (dL/dh_t * dh_t/dW_hh)
其中 dh_t/dW_hh 包含对 h_{t-1} 的连乘这就是问题所在。
下面这张图能帮你快速理解梯度传播的路径和瓶颈：
说白了，RNN 的梯度问题本质上是深度神经网络的通病，只是 RNN 的“深度”藏在时间维度里，容易被忽视。
3.2 梯度截断：最简单的急救方案
梯度爆炸相对好处理。最简单的办法就是：梯度截断（Gradient Clipping）。
我个人的习惯是，发现 loss 突然跳到 NaN，第一反应就是梯度爆炸了。这时候，打开梯度裁剪，多半能救回来。
梯度裁剪有两种常见方式：
| 方式 | 公式 | 特点 |
|---|---|---|
| 按值裁剪 | clip(g, -c, c) | 简单粗暴，适合已知阈值 |
| 按范数裁剪 | g = g * c / ||g|| | 保留方向，适合自适应场景 |
我建议在 RNN 分子生成任务中，优先使用按范数裁剪。原因很简单：分子序列的梯度范数变化较大，固定阈值容易误伤。
PyTorch 中的实现也很直接：
import torch.nn.utils as utils
# 假设 model 是你的 RNN，loss.backward() 之后
utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
optimizer.step()合适的裁剪阈值通常从 1.0 到 10.0 之间。你可以从 5.0 开始，观察梯度范数的分布再做调整。
3.3 权重初始化：从源头预防
梯度消失不能只靠裁剪，得从初始化入手。
你想想看，如果权重矩阵一开始就很小，连乘后梯度自然消失。如果太大，又会爆炸。合适的初始化能让梯度的方差在传播过程中保持稳定。
我个人最常用的两种初始化是 Xavier（Glorot）和 He 初始化。
| 初始化方法 | 适用范围 | 激活函数 | 计算公式 |
|---|---|---|---|
| Xavier (Glorot) | 对称激活函数 | tanh, sigmoid | W ~ Uniform(-sqrt(6/(n_in+n_out)), sqrt(6/(n_in+n_out))) |
| He (Kaiming) | 非对称激活函数 | ReLU, Leaky ReLU | W ~ Normal(0, sqrt(2/n_in)) |
在 RNN 中，我建议针对时间步的隐层权重使用 He 初始化，因为分子生成任务常用 ReLU 或其变体。但注意，输出层的权重还是要配合具体损失函数选择。
一个常见的例子：
import torch.nn as nn
def init_weights(m):
if isinstance(m, nn.Linear):
nn.init.kaiming_uniform_(m.weight, mode='fan_in', nonlinearity='relu')
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.RNN):
for name, param in m.named_parameters():
if 'weight_ih' in name:
nn.init.xavier_uniform_(param)
elif 'weight_hh' in name:
nn.init.orthogonal_(param)
model.apply(init_weights)如果你遇到梯度消失，先检查初始化。如果遇到爆炸，先用梯度裁剪。两个配合使用，基本能解决 RNN 在分子生成任务中 90% 的梯度问题。