- The paper identifies that uncontrolled magnitude drifts in weights and activations hinder training, and it demonstrates how standardization techniques can resolve these issues.
- The paper introduces modifications, including the removal of group normalization and the adoption of pixel normalization, achieving record FID scores on ImageNet synthesis.
- The paper presents a novel post-hoc EMA method that allows fine-tuning after training, offering deeper insights into optimal model averaging for improved image synthesis.
In the paper titled "Analyzing and Improving the Training Dynamics of Diffusion Models," researchers from NVIDIA propose essential improvements to the training process of denoising diffusion probabilistic models (DDPMs), which are extensively used for data-driven image synthesis. The improvements target key challenges in the training dynamics that previously hindered the models' performance.
The team identified that the magnitudes of weights, activations, and subsequent updates within popular diffusion models experience uncontrolled drifts over the course of training, leading to imbalances that degrade model quality. To tackle this, they introduced a series of architecture and training modifications that systematically preserve the expected magnitudes of these components without altering the model's high-level structure.
One significant alteration involves standardizing the magnitude-preserving design of learned layers. By applying this standardization, they ensure that the adjustments made by the optimizer are uniform across the model, preventing any unchecked growth of individual weight vectors and enabling more predictable training behavior.
Further enhancements include the systematic removal of group normalizations. This modification, along with other tweaks such as introducing pixel normalization in encoder blocks and simplified architecture, bolstered the models' fidelity. After these optimizations, the newly designed networks showed considerable quality improvements, setting new records for the Fréchet Inception Distance (FID) in ImageNet synthesis. Notably, the optimized models yielded these results using fast deterministic sampling rather than the stochastic sampling approach common in prior methods.
An independent contribution from the paper is the introduction of a method for setting the exponential moving average (EMA) post hoc. Typically an indispensable technique in refining image synthesis, EMA parameters can often be cumbersome and cost-ineffective to tune since their effects become apparent only when the training nears completion. With the new post-hoc EMA method, researchers can fine-tune the EMA profile even after training runs are completed. This capability allows for the exploration of EMA's interactions with other aspects of training, offering insights that could inform future improvements to model averaging techniques.
The presented improvements aim to aid researchers and practitioners in producing high-quality synthetic images efficiently and provide tools for better comprehension and control of training dynamics in diffusion models. The paper concludes with the team's intention to make their implementation and pre-trained models publicly accessible, giving others in the field the opportunity to leverage and build upon their work.