Demystifying Diffusion Objectives: Reweighted Losses are Better Variational Bounds
This work provides a theoretical justification for widely used reweighting schemes in diffusion models, which is important for researchers developing more efficient generative models.
The authors tackled the problem of improving training objectives for diffusion models by deriving a theoretical interpretation of reweighted losses, showing they provide better variational bounds and reduce data-model KL-divergences. They demonstrated significant improvements in pixel-space image modeling with masked diffusion, achieving sample quality comparable to continuous diffusion models.
We derive a new theoretical interpretation of the reweighted losses that are widely used for training diffusion models. Our method is based on constructing a cascade of time-dependent variational lower bounds on the data log-likelihood, that provably improves upon the standard evidence lower bound and results in reduced data-model KL-divergences. Combining such bounds gives rise to reweighted objectives that can be applied to any generative diffusion model including both continuous Gaussian diffusion and masked (discrete) diffusion models. Then, we showcase this framework in masked diffusion and report significant improvements over previous training losses in pixel-space image modeling, approaching sample quality comparable to continuous diffusion models. Our results also provide a theoretical justification for the simple weighting scheme widely used in masked image models.