Techniques for Training Stability (Gradient Clipping, EMA)

Training deep generative models, such as diffusion models, frequently presents complex optimization challenges over extended periods. While advanced architectures and loss functions establish the foundation, achieving stable convergence and optimal performance often necessitates supplementary techniques to control the training dynamics. Gradient clipping and Exponential Moving Average (EMA) of model weights are two primary methods used for improving training stability.

Gradient Clipping

Deep networks, including the U-Nets or Transformers used in diffusion models, can sometimes suffer from the exploding gradient problem. During backpropagation, gradients can accumulate and become excessively large, leading to drastic updates in model weights. These large updates can destabilize training, causing the loss to diverge or oscillate wildly, preventing the model from converging to a good solution.

Gradient clipping directly addresses this by limiting the magnitude of the gradients before they are used to update the model parameters. The most common method is clipping by norm, typically the L2 norm (Euclidean norm).

Let $\mathbf{g} = \nabla_{\theta} L$ be the gradient of the loss $L$ with respect to the model parameters $\theta$. We calculate the L2 norm of the entire gradient vector across all parameters: $\|\mathbf{g}\| = \sqrt{\sum_i g_i^2}$. If this norm exceeds a predefined threshold $c$, the gradient vector is rescaled to have a norm equal to $c$.

$$\mathbf{g} \leftarrow \begin{cases} \mathbf{g} & \text{if } \|\mathbf{g}\| \le c \\ \frac{c}{\|\mathbf{g}\|} \mathbf{g} & \text{if } \|\mathbf{g}\| > c \end{cases}$$

This ensures that the magnitude of the weight update step is bounded, preventing extreme changes caused by outlier gradients.

Practical Notes

• Choosing the Threshold $c$: This is a critical hyperparameter. A value that's too high might rarely activate the clipping, failing to prevent instability. A value that's too low might excessively shrink gradients, slowing down convergence or preventing the model from reaching optimal performance. Typical values often range between 1.0 and 10.0, but the best value depends on the model architecture, dataset, batch size, and learning rate. It usually requires some experimentation.
• Implementation: Most deep learning frameworks provide built-in functions for gradient clipping. For instance, in PyTorch, you would typically use torch.nn.utils.clip_grad_norm_ after the loss.backward() call and before the optimizer.step() call.

# Example snippet (PyTorch)
optimizer.zero_grad()
loss = compute_loss(model, batch)
loss.backward()
# Clip gradients before the optimizer step
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # max_norm is 'c'
optimizer.step()

Gradient clipping acts as a safety mechanism, particularly useful during the initial phases of training or when using high learning rates.

Exponential Moving Average (EMA) of Weights

Stochastic Gradient Descent (SGD) and its variants introduce noise into the training process due to the mini-batch nature of updates. While this stochasticity helps escape poor local minima, it also means the model parameters can fluctuate significantly around an optimal point, even late in training. Using the parameters from the very last training step might not yield the best possible model performance, as it could represent a transient, noisy state.

Exponential Moving Average (EMA) provides a way to obtain a more stable and often better-performing set of model weights by averaging the parameters over recent training history. It maintains a "shadow" copy of the model parameters, which are updated slowly based on the current training parameters.

Let $\theta_t$ be the model parameters after the update at step $t$, and let $\theta'_t$ be the corresponding EMA parameters. The EMA parameters are updated using a decay factor $\beta$:

$$\theta'_{t} = \beta \theta'_{t-1} + (1 - \beta) \theta_t$$

Here, $\beta$ is a hyperparameter typically set close to 1 (e.g., 0.99, 0.999, or even 0.9999). A higher $\beta$ means the EMA weights change more slowly and incorporate a longer history of parameters (more smoothing). A lower $\beta$ makes the EMA weights track the current training weights more closely.

Diagram illustrating the relationship between regular training weights ($\theta_t$, blue nodes) updated via optimization steps, and the EMA weights ($\theta'_t$, green nodes) updated as a weighted average of the previous EMA weight and the current training weight. Dashed lines indicate the influence of the current training weight on the next EMA weight.

Practical Notes

• When to Use EMA Weights: The EMA parameters $\theta'$ are typically not used during the training loop itself (i.e., they don't influence the gradient calculation or the optimizer's updates on $\theta$). Instead, the EMA weights are maintained separately and used after training is complete, or periodically during training for evaluation and especially for generating samples from the diffusion model. Using EMA weights often leads to noticeably better sample quality compared to using the final raw weights $\theta_T$.
• Choosing the Decay $\beta$: The optimal decay rate $\beta$ depends on the training duration and learning rate schedule. Higher values like 0.999 or 0.9999 are common for long training runs, providing significant smoothing. For shorter training or rapidly changing learning rates, a slightly smaller $\beta$ might be preferred.
• Implementation: Many training libraries or codebases implement EMA callbacks or utilities. This typically involves creating a copy of the model parameters at the start and updating this copy after each optimizer step according to the EMA formula.

# Snippet (often part of a callback or utility class)
ema_decay = 0.999
ema_model_params = [p.clone().detach() for p in model.parameters()]

# Inside training loop, after 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)

# For inference/sampling, load the ema_model_params into the model
# ... load ema params ...
# model.eval()
# generate_samples(model, ...)

Combined Use

Gradient clipping and EMA are often used together. Clipping prevents catastrophic divergence due to isolated large gradients, while EMA smooths out the general stochasticity of the optimization process, leading to a final model state. By employing these techniques, you can significantly improve the reliability of training advanced diffusion models and often achieve superior generative performance.