Robust Diffusion Models via Divergence-Induced Weighted Denoising

Standard diffusion models are trained using a weighted mean squared error (MSE) loss, which treats every training sample equally. This approach is sensitive to corrupted or atypical data, which can degrade the quality of generated samples. The authors investigate whether replacing the standard MSE loss with a more robust, divergence-induced objective can improve performance when training data is contaminated.

The authors propose a framework where the standard denoising loss is replaced by a nonlinear transformation $h_G(m_t)$, where $m_t$ is the per-step denoising mismatch and $G$ is an f-divergence. By leveraging the Gaussian structure of the diffusion reverse process, they show that the per-step likelihood ratio follows a lognormal distribution. This allows them to express the training objective as a sum of local conditional divergences. The derivative of this induced divergence, $h'_G(m_t)$, acts as an influence weight that automatically suppresses the contribution of samples with large residuals, effectively connecting diffusion training to robust M-estimation.

The study demonstrates that using bounded-influence divergences, such as Hellinger (HD) or negative exponential (NED) divergence, provides significant robustness against data contamination. In experiments on CIFAR-10 with 30% data contamination, the NED objective reduced the Frechet Inception Distance (FID) from 93.0 (standard KL/MSE) to 77.5. The authors also show that these robust objectives outperform standard alternatives like Huber or clipped MSE losses, while requiring only minimal changes to the training code and adding negligible computational overhead.

This work provides a principled, statistically grounded method for improving the robustness of generative models. By framing diffusion training as divergence-induced weighted denoising, the authors offer a flexible mechanism to handle noisy datasets without needing to manually clean the data or tune complex clipping parameters, making it easier to train high-quality models on real-world, imperfect data.