Covariance-aware sampling for Diffusion Models

2026-05-13

AI Summary

This work addresses the degradation in generation quality of diffusion models under few-step sampling, which stems from modeling only the mean of the reverse process. To overcome this limitation, we introduce a covariance-aware mechanism that explicitly estimates the reverse-process covariance via Tweedie’s formula and efficiently implements it through a structured decomposition in Fourier space. Our approach requires only a single additional Jacobian-vector product computation yet yields substantial improvements in pixel-level generation fidelity. Under identical function evaluation budgets, our method outperforms state-of-the-art second-order samplers—including Heun, DPM-Solver++, and aDDIM—establishing a new, efficient paradigm for few-step diffusion sampling.

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

Problem

• Diffusion Models
• few-step sampling
• covariance
• sampling quality
• reverse process

Innovation

• covariance-aware sampling
• diffusion models
• few-step sampling
• Tweedie's formula
• Fourier-space decomposition