Analyzing and Improving the Training Dynamics of Diffusion Models

By Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, Samuli Laine. Computer Vision and Pattern Recognition, 2023

442 citations recorded by the source corpus.

Diffusion models currently dominate the field of data- driven image synthesis with their unparalleled scaling to large datasets. In this paper, we identify and rectify several causes for uneven and ineffective training in the popular ADM diffusion model architecture, without altering its high- level structure. Observing uncontrolled magnitude changes and imbalances in both the network activations and weights over the course of training, we redesign the network layers to preserve activation, weight, and update magnitudes on ex- pectation. We find that systematic application of this philoso- phy eliminates the observed drifts and imbalances, resulting in considerably better networks at equal computational com- plexity. Our modifications improve the previous record FID of 2.41 in ImageNet-512 synthesis to 1.81, achieved using fast deterministic sampling. As an independent contribution, we present a method for setting the exponential moving average (EMA) parameters post-hoc, i.e., after completing the training run. This allows precise tuning of EMA length without the cost of performing several training runs, and reveals its surprising interactions with network architecture, training time, and guidance.

Subfields: Computer Vision and Pattern Recognition, Statistical and Nonlinear Physics

View Paper: https://www.semanticscholar.org/paper/bb857c72e7c75fc3d4b44b1dcaa66c62ea10a2e1

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computer Vision and Pattern Recognition

In Statistical and Nonlinear Physics

Citations in the corpus

Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think (2024, unknown): “Because of this, we do not apply any data augmentation, but we find this does not lead to a big difference, as similarly observed in EDM2 (Karras et al., 2024).”
From Slow Bidirectional to Fast Autoregressive Video Diffusion Models (2024, unknown): “Diffusion models can be trained on either raw data [23, 27, 30] or on a lower-dimensional latent space obtained by a variational autoencoder (VAE) [31, 60, 66, 96, 109].”
Autoregressive Image Generation without Vector Quantization (2024, methodology): “Following [37, 25], we maintain the exponential moving average (EMA) of the model parameters with a momentum of 0.9999.”
Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (2024, background): “Due to their iterative nature and the associated computational costs, as well as the long sampling times during inference, research on formulations for more efficient training and/or faster sampling of these models has increased (Karras et al., 2023; Liu et al., 2022).”
F5-TTS: A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching (2024, methodology): “We use the Exponential Moving Averaged (EMA) [62] weights for inference, and the Euler ODE solver for F5-TTS (midpoint for E2 TTS as described in [36]).”