Analyzing and Improving the Training Dynamics of Diffusion Models
Authors: name: Tero Karras; name: Miika Aittala; name: Jaakko Lehtinen; name: Janne Hellsten; name: Timo Aila; name: Samuli Laine
Published: 2023-12-05 11:55:47

Abstract: 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 expectation. We
find that systematic application of this philosophy eliminates the observed
drifts and imbalances, resulting in considerably better networks at equal
computational complexity. 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.