Reproduce “Variance-Aware Adaptive Weighting for Diffusion Model Training” on OpenResearch

About this paper

Nanlong Sun, Lei Shi

Kennesaw State University researchers introduce a variance-aware adaptive weighting strategy for diffusion model training, addressing imbalanced training dynamics across noise levels. The method achieves a lower FID score of 13.58 on CIFAR-10 compared to the baseline's 14.21, while also reducing performance variance and maintaining negligible computational overhead.

Key insights

• Per-sample training loss variance is highly heterogeneous across log-SNR levels, concentrated in mid-to-high SNR regions and spanning nearly two orders of magnitude, indicating imbalanced training signals.
• The problem of imbalanced training in diffusion models can be framed in terms of stochastic gradient variance, where the optimal sampling density for minimizing variance should be proportional to the conditional standard deviation of the gradients.
• A simple, adaptive reweighting of the training loss can effectively mitigate the observed variance imbalance, leading to more efficient and stable diffusion model training.

Summary by alphaXiv.