Covariance-aware sampling for Diffusion Models

TL;DR

This paper introduces a covariance-aware sampling method for diffusion models that enhances sample quality in few-step regimes by explicitly modeling the reverse process covariance with minimal computational overhead.

Contribution

The authors propose a novel covariance-aware sampler that improves diffusion model sampling quality by explicitly modeling covariance, extending DDIM with efficient Fourier-space decomposition.

Findings

Our method outperforms state-of-the-art second order samplers at the same NFE.

It requires only one extra Jacobian-Vector Product per step.

The approach consistently produces superior samples in pixel-based DMs.

Abstract

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).