Covariance-aware sampling for Diffusion Models
For practitioners using pixel-space diffusion models, this method offers a practical improvement in few-step sampling quality with minimal computational overhead.
The authors propose a covariance-aware sampler for pixel-space Diffusion Models that improves sample quality in the few-step regime by explicitly modeling the reverse-process covariance. Their method outperforms state-of-the-art second-order samplers (Heun, DPM-Solver++) and aDDIM at the same number of function evaluations.
We present a covariance-aware sampler that improves the quality of pixel-space Diffusion Model (DM) sampling in the few-step regime. We hypothesize that in the few-step regime samplers fail because they rely solely on the predicted mean of the reverse distribution, while our solution explicitly models the reverse-process covariance. Our method combines Tweedie's formula to estimate the covariance with an efficient, structured Fourier-space decomposition of the covariance matrix. Implemented as an extension of DDIM, our method requires only a minimal overhead: one extra Jacobian-Vector Product (JVP) per step. We demonstrate that for pixel-based DMs, our method consistently produces superior samples compared to state-of-the-art second order samplers (Heun, DPM-Solver++) and the recent aDDIM sampler, at an identical number of function evaluations (NFE).