Robust Diffusion Models via Divergence-Induced Weighted Denoising
This work provides a principled robust training framework for diffusion models, addressing the problem of data contamination for practitioners training generative models.
The paper introduces a robust training method for diffusion models by replacing the standard MSE denoising loss with an f-divergence-induced weighted loss, which improves performance under data contamination. On CIFAR-10 with 30% contamination, the NED variant reduces FID from 93.0 (KL) to 77.5.
We show that replacing the standard MSE denoising loss in diffusion models with a nonlinear transformation induced by an f-divergence yields a simple robust training surrogate that empirically improves performance under data contamination, with small additional computational overhead. The theoretical foundation rests on a local divergence construction: under the Gaussian reverse-kernel structure of DDPM, each per-step likelihood ratio follows a lognormal distribution parameterized by a scalar mismatch, so the conditional f-divergence at each step reduces to a one-dimensional function of the denoising error. Summing these local divergences yields a training objective that unifies diffusion training as divergence induced weighted denoising, where the derivative of the induced divergence acts as a residual-space influence weight that controls the contribution of each sample. Bounded-influence divergences (Hellinger, negative exponential) suppress large error samples, with Hellinger yielding an explicit exponential weight, connecting the framework to robust M-estimation. Empirically, on CIFAR-10 under 30% contamination, NED reduces FID from 93.0 (KL) to 77.5, while also outperforming standard robust losses such as Huber and clipped MSE.