梯度裁剪方法

训练大型神经网络 (neural network)，特别是多层Transformer模型，有时会造成数值不稳定。一个常见的问题是“梯度爆炸”现象，即在反向传播 (backpropagation)时，梯度的数值变得过大。这可能导致模型参数 (parameter)进行巨大的更新，可能引发发散（即损失值飙升至无穷大或NaN——非数字）或阻止收敛的震荡。这种不稳定在多层网络中，或在使用某些激活函数 (activation function)或初始化方案时，可能尤其常见，即便使用了层归一化 (normalization)等方法。由于低精度格式的数值范围有限，在混合精度训练（第20章讨论）中，这种情况也可能加剧。

梯度裁剪是一种直接而有效的方法，它通过限制梯度的大小来缓解这个问题，在优化器使用梯度更新模型权重 (weight)之前进行操作。主要思路不是改变梯度更新的方向，而是当梯度大小超过预设阈值时，限制其大小。

按梯度范数裁剪

对于大型语言模型来说，最常用的方法是按梯度的L2范数（欧几里得范数）进行裁剪。这种做法将所有模型参数 (parameter)的全部梯度（或有时是每个参数组的梯度）视为一个向量 (vector)，计算其L2范数，如果范数超过给定阈值 $c$，则重新缩放该向量。

从数学上讲，令 $\mathbf{g}$ 代表所有梯度串联在一起的向量。L2范数计算如下：

$$\|\mathbf{g}\| = \sqrt{\sum_{i} g_i^2}$$

裁剪操作则应用如下：

$$\mathbf{g} \leftarrow \frac{c}{\|\mathbf{g}\|} \mathbf{g} \quad \text{如果 } \|\mathbf{g}\| > c$$

如果范数 $\|\mathbf{g}\|$ 小于或等于阈值 $c$，梯度保持不变。如果范数大于 $c$，梯度向量 $\mathbf{g}$ 会被按 $c / \|\mathbf{g}\|$ 的比例缩小，从而确保其新范数恰好为 $c$。这在限制梯度大小的同时，保持了梯度更新的方向。

在PyTorch中，这通常使用 torch.nn.utils.clip_grad_norm_ 函数实现。它在反向传播 (backpropagation)（计算梯度）之后、优化器步骤（根据梯度更新权重 (weight)）之前应用。

import torch
from torch.nn.utils import clip_grad_norm_

# 假设模型、损失、优化器已定义
# ... 在训练循环内部 ...
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward() # 计算梯度

# 定义最大梯度范数阈值
max_grad_norm = 1.0

# 裁剪梯度
total_norm = clip_grad_norm_(
    model.parameters(),
    max_norm=max_grad_norm,
    norm_type=2.0
)

# 可选：如果需要，在裁剪前记录梯度范数
# print(f"裁剪前的梯度范数: {total_norm}")

optimizer.step() # 使用（可能已裁剪的）梯度更新权重
# ... 训练循环的其余部分 ...

在此代码段中，clip_grad_norm_ 计算传递给它的参数（model.parameters()）的所有梯度的总L2范数。如果此范数超过 max_norm（我们的阈值 $c$），它会通过重新缩放 原地 修改梯度。该函数返回裁剪前的原始总范数，这有助于观察训练过程。设置 norm_type=2.0 明确指定了L2范数。

按梯度值裁剪

另一种方法是按值裁剪，尽管在训练大型Transformer模型时不如常用。这种方法独立地裁剪每个梯度分量 $g_i$，如果它落在指定范围 $[-c, c]$ 之外。

$$g_i \leftarrow \text{max}(\text{min}(g_i, c), -c)$$

这意味着任何大于 $c$ 的梯度分量都被设为 $c$，任何小于 $-c$ 的分量都被设为 $-c$。与范数裁剪不同，这种方法可以改变总梯度向量 (vector)的 方向，因为不同分量可能被不同程度地裁剪或根本不裁剪。

在PyTorch中，这可以使用 torch.nn.utils.clip_grad_value_ 完成：

import torch
from torch.nn.utils import clip_grad_value_

# ... 在训练循环内部，在 loss.backward() 之后 ...

# 为每个梯度分量定义最大绝对值
clip_value = 0.5

# 按值裁剪梯度
clip_grad_value_(model.parameters(), clip_value=clip_value)

optimizer.step()
# ... 训练循环的其余部分 ...

尽管更简单，但与范数裁剪相比，按值裁剪在深度学习 (deep learning)优化中通常被认为理论依据较少，因为它不将梯度视为一个统一的方向向量。对于大多数大型语言模型训练情况，clip_grad_norm_ 是更受推荐的方法。

梯度裁剪的实际考量

• 选择阈值 $c$： 裁剪阈值（max_norm 或 clip_value）是一个超参数 (parameter) (hyperparameter)，通常需要凭经验调整。在大型语言模型训练中，max_norm 的一个常见起始值是 1.0。
  • 如果设置过高，它可能无法有效防止偶然出现的巨大梯度引起的不稳定。
  • 如果设置过低，它可能通过不必要地缩小有用梯度而阻碍学习，减慢收敛速度。
  • 在训练期间监测梯度范数（clip_grad_norm_ 在裁剪 发生前 返回的值）有助于确定这个选择。如果范数频繁达到阈值，你可能要考虑学习率是否过高，或者阈值是否可以略微提高。如果范数很少接近阈值，它可能没有产生太大影响。
• 与优化器和学习率的相互作用： 梯度裁剪与优化器和学习率调度配合使用。它起到安全机制的作用，但不能取代对合适的学习率和调度器的需求。如果梯度持续过大并被裁剪，这可能表明学习率对于当前的训练阶段来说过高。
• 必要性： 尽管梯度裁剪通常有助于稳定性，尤其是在非常大的模型中或在训练的初始阶段（或微调 (fine-tuning)时），但如果其他因素（适当的初始化、像Pre-LN这样的归一化 (normalization)、调整得当的学习率调度、AdamW优化器）有助于稳定训练，它可能并非总是严格必要。然而，考虑到大型语言模型训练的成本，它经常被用作一种预防措施。

下图说明了按范数裁剪的效果。位于圆圈外部（代表范数阈值 $c$）的梯度向量 (vector)会沿径向按比例缩小，趋向于圆周边界，同时保持其原始方向。

梯度向量 $g_2$ 和 $g_3$ 原始范数大于阈值 $c$。按范数裁剪将它们按比例缩小（蓝色虚线箭头），使其位于由 $c$ 定义的边界上，而 $g_1$ 已经在阈值内，保持不变。

通过防止过大的更新，梯度裁剪显著提升了成功训练大型语言模型所需的稳定性，使AdamW等优化器和精心设计的学习率调度能有效应对复杂的损失曲面。