Robust Diffusion Models via Divergence-Induced Weighted Denoising

Abstract: 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.

• The paper introduces a framework that recasts diffusion model training as divergence minimization, leading to a simple robustness mechanism that significantly improves model performance under contaminated data.
• The proposed method uses divergence-induced weighted denoising, which reduces per-step likelihood ratios to a single scalar parameter called the per-step mismatch $m_t$, leading to more accurate and robust denoising
• On the CIFAR-10 dataset, the method shows a significant reduction in Frechet Inception Distance (FID) when compared to standard losses, particularly under heavy contamination.

Overview

This paper develops a statistical framework that recasts diffusion model training as divergence minimization, and uses this framework to derive a simple robustness mechanism. The central observation is that under the Gaussian reverse-kernel structure of DDPM, the per-step likelihood ratio between the true variational posterior and the model kernel is lognormal with a single scalar parameter—the per-step mismatch $m_t$. Consequently, any f-divergence evaluated at each step reduces to a one-dimensional function $h_G(m_t)$ of the denoising error, and the resulting training objective differs from standard DDPM training only by replacing the weighted MSE loss $m_t$ with $h_G(m_t)$—equivalently, by reweighting each sample's gradient by an influence weight $h_G'(m_t)$. Bounded-influence divergences such as Hellinger (HD) and negative exponential divergence (NED) suppress large-error samples exponentially, connecting diffusion training to classical robust M-estimation. Empirically, on CIFAR-10 with 30% contamination, NED reduces FID from 93.0 (KL) to 77.5 in the multi-seed protocol, and to 55.1 versus 94.0 for Huber loss in a single-seed baseline comparison.

The lognormal likelihood ratio and local divergence construction

The structural foundation is Lemma 2.1: for matched-covariance Gaussian kernels $q_t = \mathcal N(\tilde\mu_t, \Sigma_t)$ and $$p_{\theta,t} = \mathcal N(\mu_{\theta,t}, \Sigma_t)$$, conditional on $(x_0, x_t)$ the likelihood ratio satisfies $\log r_t \sim N(-m_t, 2m_t)$, where $m_t$ is half the Mahalanobis distance between the means. The authors note that while the lognormal form of a Gaussian likelihood ratio is classical, its explicit statement in the diffusion literature appears to be new; existing work computes per-step KL divergences without isolating the distributional form of $h_G(m_t)$0. This lemma is what makes each per-step divergence a scalar function of $h_G(m_t)$1, since $h_G(m_t)$2 with $h_G(m_t)$3 lognormal.

To connect the marginal divergence $h_G(m_t)$4—the quantity one would ideally minimize—to these local terms, Theorem 2.2 shows via the data-processing inequality that the lifted path-space divergence $h_G(m_t)$5 is a tangent upper bound on the marginal divergence, exact when the model reproduces the full path law. For general f-divergences, however, the path likelihood ratio contains an endpoint factor that does not decompose into per-step terms under nonlinear $h_G(m_t)$6. Rather than bounding this factor, the authors use the path-space view as motivation and directly construct the training surrogate

$h_G(m_t)$7

Asymptotic justification comes from a product-to-sum reduction (Theorem 2.3 and Corollary 2.4): under Lindeberg-type moment conditions and pairwise uncorrelatedness of the increments $h_G(m_t)$8, the expectation of $h_G(m_t)$9 applied to the product likelihood ratio equals the sum of per-step divergences up to $m_t$0 error, where $m_t$1. Importantly, the authors are candid that condition (D2)—uncorrelatedness—is an approximation whose validity depends on architecture and training stage, and that even when it fails, the local objective remains well-defined and usable; (D2) is needed only for the asymptotic justification.

A consequence worth emphasizing is that all smooth divergences satisfy $m_t$2 near zero (Theorem 2.5), so at small mismatch (e.g., $m_t$3) the three divergences agree to four digits. The choice of divergence therefore matters primarily in the large-mismatch regime—which is precisely where contamination robustness is needed.

Influence weights and connection to robust M-estimation

Differentiating the objective shows the effective per-sample weight is $m_t$4, where $m_t$5 depends only on the noise schedule. For KL, $m_t$6: every sample contributes equally regardless of residual magnitude, which is exactly why standard DDPM training is non-robust. For HD, $m_t$7, giving closed-form exponential downweighting. For NED, no closed form exists, but Proposition 2.7 establishes the asymptotic rate $m_t$8—the same exponential rate as HD with an additional polynomial correction making it slightly more aggressive. In practice, $m_t$9 is approximated by Gauss–Legendre quadrature over $h_G(m_t)$0; 16 nodes suffice per the ablation, though all experiments conservatively use 64.

The paper makes the robust-statistics connection precise rather than analogical: the residual-space score is $h_G(m_t)$1. For KL this score is unbounded in $h_G(m_t)$2; for HD and NED it is bounded and redescending, paralleling redescending $h_G(m_t)$3-functions in robust regression. Global monotonicity of $h_G(m_t)$4 (Theorem 2.6) is established via a martingale argument: writing $h_G(m_t)$5 as an exponential martingale $h_G(m_t)$6 and applying conditional Jensen's inequality shows $h_G(m_t)$7 is a submartingale, hence $h_G(m_t)$8 is nonnegative and nondecreasing.

Experimental results

All experiments use a deliberately small 35M-parameter DDPM U-Net trained for 100K steps on CIFAR-10/100, with contamination simulated by replacing a fraction $h_G(m_t)$9 of images with a mixture of four corruption types (patch noise, blur, salt-and-pepper, contrast shift). The absolute FID values reflect this controlled budget; the relative ordering across divergences is the quantity of interest.

On clean data ($h_G'(m_t)$0), all three divergences perform comparably, consistent with the first-order equivalence theory. Notably, HD and NED exhibit substantially lower cross-seed variance than KL (standard deviations of 1.7 and 2.4 versus 12.3), suggesting bounded-influence objectives stabilize training—an effect visible even without contamination. Under contamination, the advantage grows with $h_G'(m_t)$1: at $h_G'(m_t)$2, NED achieves FID $h_G'(m_t)$3 versus $h_G'(m_t)$4 for KL over three seeds. On CIFAR-100 (single seed), NED reduces FID from 98.0 to 75.3 at $h_G'(m_t)$5.

Against standard robust losses, the results are stronger but carry a caveat. At $h_G'(m_t)$6 (single seed), HD and NED achieve FID 56.2 and 55.1 versus 94.0 (Huber, $h_G'(m_t)$7) and 90.5 (clipped MSE, $h_G'(m_t)$8)—a 35–40% improvement. However, Huber and clipped MSE actually beat the divergence-induced objectives on clean data (27.8/27.7 versus ~32), and the authors concede that the baseline thresholds were representative defaults rather than grid-searched optima, so systematic tuning might narrow the gap. Computational overhead is negligible: below 1% wall-clock overhead for NED's quadrature, with total training time indistinguishable across divergences (~3 hours per run).

Limitations and open questions

The paper is explicit about several gaps between the framework and its guarantees. First, the local objective is a training surrogate motivated by the path-space bound but is not itself a proven upper bound on the marginal divergence for non-KL f-divergences; quantifying this gap remains open. Second, the product-to-sum reduction is first-order in the local mismatch, so finite-mismatch corrections are unresolved. Third, despite the M-estimation framing, no formal contamination-risk bounds or breakdown-point analysis is provided—a full estimator-level influence-function study is left open. Fourth, the empirical evidence rests on small-scale datasets with substantial cross-seed variance at high contamination, and CIFAR-100 results are single-seed. Finally, the framework assumes Gaussian reverse kernels with matched covariance; extension to learnable covariances, and empirical validation of the proposed flow-matching analogue $h_G'(m_t)$9, remain future work. The authors also note that aggressive downweighting may discard useful signal under mild contamination or clean data, a trade-off governed entirely by the shape of $q_t = \mathcal N(\tilde\mu_t, \Sigma_t)$0.

Conclusion

The paper contributes a clean reduction—lognormal per-step likelihood ratios imply per-step divergences are scalar functions of the denoising mismatch—that unifies diffusion training objectives as divergence-induced weighted denoising and yields a practical, one-line robustness modification with negligible overhead. The theoretical claims are appropriately scoped as asymptotic justifications under stated regularity conditions rather than guarantees, and the strongest empirical claim—that bounded-influence divergences substantially outperform both KL and untuned robust losses under heavy contamination—is supported by multi-seed experiments at small scale. The main open questions are whether the surrogate-to-marginal gap can be bounded, whether formal robustness guarantees can accompany the empirical gains, and whether the framework extends beyond Gaussian reverse kernels and small image benchmarks.