训练深度生成模型,例如扩散模型,经常面临长时间复杂的优化挑战。虽然精巧的网络结构和损失函数是主要组成部分,但要获得稳定的收敛和最佳性能,通常需要额外技术来调节训练动态。梯度裁剪和模型权重的指数移动平均 (EMA) 是提高训练稳定性的两种主要方法。梯度裁剪深度网络,包括扩散模型中使用的U-Net或Transformer,有时会遇到梯度爆炸问题。在反向传播过程中,梯度可能累积并变得非常大,导致模型权重发生剧烈更新。这些大幅更新会使训练不稳定,导致损失发散或剧烈震荡,从而阻止模型收敛到一个好的解。梯度裁剪直接处理此问题,通过限制梯度的幅度,在它们被用来更新模型参数之前。最常用的方法是按范数裁剪,通常是L2范数 (欧几里得范数)。设 $\mathbf{g} = \nabla_{\theta} L$ 是损失 $L$ 关于模型参数 $\theta$ 的梯度。我们计算整个梯度向量(所有参数)的L2范数:$|\mathbf{g}| = \sqrt{\sum_i g_i^2}$。如果此范数超过预设的阈值 $c$,梯度向量将被重新缩放,使其范数等于 $c$。$$ \mathbf{g} \leftarrow \begin{cases} \mathbf{g} & \text{如果 } |\mathbf{g}| \le c \ \frac{c}{|\mathbf{g}|} \mathbf{g} & \text{如果 } |\mathbf{g}| > c \end{cases} $$这确保了权重更新步长的幅度受到限制,避免了由异常梯度引起的极端变化。实际考量阈值 $c$ 的选择: 这是一个重要的超参数。过高的值可能很少触发裁剪,从而无法阻止不稳定。过低的值可能会过度缩小梯度,减缓收敛速度或阻止模型达到最佳性能。典型值通常在1.0到10.0之间,但最佳值取决于模型结构、数据集、批量大小和学习率。它通常需要一些实验。实现: 大多数深度学习框架都提供了梯度裁剪的内置函数。例如,在PyTorch中,通常会在调用 loss.backward() 之后、调用 optimizer.step() 之前使用 torch.nn.utils.clip_grad_norm_。# 示例代码 (PyTorch) optimizer.zero_grad() loss = compute_loss(model, batch) loss.backward() # 在优化器步骤前裁剪梯度 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # max_norm 即是 'c' optimizer.step()梯度裁剪作为一种安全机制,在训练的初始阶段或使用较高学习率时尤其有用。权重的指数移动平均 (EMA)随机梯度下降 (SGD) 及其变体由于小批量更新的特性,会给训练过程引入噪声。尽管这种随机性有助于摆脱不好的局部最小值,但也意味着模型参数即使在训练后期也可能在最优解附近大幅波动。使用最后一步训练的参数可能无法带来最佳模型性能,因为它可能代表一个瞬时、有噪声的状态。指数移动平均 (EMA) 提供了一种方法,通过对近期训练历史中的参数进行平均,来获得更稳定且通常性能更好的模型权重集。它维护模型参数的“影子”副本,这些副本根据当前训练参数缓慢更新。设 $\theta_t$ 是第 $t$ 步更新后的模型参数,设 $\theta'_t$ 是对应的 EMA 参数。EMA 参数使用衰减因子 $\beta$ 进行更新:$$ \theta'{t} = \beta \theta'{t-1} + (1 - \beta) \theta_t $$这里,$\beta$ 是一个超参数,通常设为接近1的值(例如0.99、0.999,甚至0.9999)。较高的 $\beta$ 意味着EMA权重变化更慢,并包含了更长的参数历史(平滑效果更好)。较低的 $\beta$ 使EMA权重更紧密地跟随当前训练权重。digraph G { rankdir=LR; node [shape=circle, style=filled, fillcolor="#a5d8ff", fontname="sans-serif", fontsize=10]; edge [arrowhead=vee, color="#495057"]; subgraph cluster_0 { label = "训练步骤"; bgcolor="#e9ecef"; style=filled; node [shape=circle, style=filled, fillcolor="#a5d8ff"]; t0 -> t1 -> t2 -> t3 -> t4 [label=" SGD 更新", fontsize=8, fontcolor="#495057"]; t0 [label="θ_t-4"]; t1 [label="θ_t-3"]; t2 [label="θ_t-2"]; t3 [label="θ_t-1"]; t4 [label="θ_t"]; } subgraph cluster_1 { label = "EMA 权重"; bgcolor="#e9ecef"; style=filled; node [shape=circle, style=filled, fillcolor="#96f2d7"]; et0 -> et1 -> et2 -> et3 -> et4 [label=" EMA 更新 (β)", fontsize=8, fontcolor="#495057"]; et0 [label="θ'_t-4"]; et1 [label="θ'_t-3"]; et2 [label="θ'_t-2"]; et3 [label="θ'_t-1"]; et4 [label="θ'_t"]; } // 显示 EMA 依赖关系的连接 t0 -> et1 [style=dashed, color="#adb5bd"]; t1 -> et2 [style=dashed, color="#adb5bd"]; t2 -> et3 [style=dashed, color="#adb5bd"]; t3 -> et4 [style=dashed, color="#adb5bd"]; // 如果需要,添加不可见边缘进行对齐 // t0 -> et0 [style=invis]; // t1 -> et1 [style=invis]; // t2 -> et2 [style=invis]; // t3 -> et3 [style=invis]; // t4 -> et4 [style=invis]; }图表说明了通过优化步骤更新的常规训练权重 ($\theta_t$,蓝色节点) 与作为前一个 EMA 权重和当前训练权重的加权平均值更新的 EMA 权重 ($\theta'_t$,绿色节点) 之间的关系。虚线表示当前训练权重对下一个 EMA 权重的影响。实际考量何时使用 EMA 权重: EMA 参数 $\theta'$ 通常 不 用于训练循环本身(即它们不影响梯度计算或优化器对 $\theta$ 的更新)。相反,EMA 权重是单独维护的,并在训练完成后使用,或在训练期间定期用于评估,尤其是用于从扩散模型生成样本。使用 EMA 权重通常会带来明显更好的样本质量,优于使用最终原始权重 $\theta_T$。衰减 $\beta$ 的选择: 最佳衰减率 $\beta$ 取决于训练时长和学习率调度。0.999 或 0.9999 等较高的值常用于长时间训练,提供显著的平滑效果。对于较短的训练或学习率快速变化的情况,可能更偏好稍小的 $\beta$ 值。实现: 许多训练库或代码库都实现了 EMA 回调或工具。这通常包括在开始时创建模型参数的副本,并在每个优化器步骤后根据 EMA 公式更新此副本。# 代码片段(通常是回调或实用程序类的一部分) ema_decay = 0.999 ema_model_params = [p.clone().detach() for p in model.parameters()] # 在训练循环中,optimizer.step() 之后 with torch.no_grad(): for ema_p, current_p in zip(ema_model_params, model.parameters()): ema_p.mul_(ema_decay).add_(current_p, alpha=1 - ema_decay) # 用于推断/采样时,将 ema_model_params 加载到模型中 # ... 加载 EMA 参数 ... # model.eval() # generate_samples(model, ...)组合使用梯度裁剪和 EMA 经常一起使用。裁剪可以防止因孤立的大梯度引起的灾难性发散,而 EMA 则平滑了优化过程中的整体随机性,从而获得一个稳定的最终模型状态。通过应用这些方法,您可以明显提高训练高级扩散模型的可靠性,并常能获得卓越的生成性能。