Research map / Image generation / Fast diffusion sampling
Improving Probabilistic Diffusion Models With Optimal Diagonal Covariance Matching
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Optimal Covariance Matching (OCM), an unbiased objective that trains a small network on top of a frozen pretrained score model to predict the diagonal of the optimal denoising covariance, improving sampling efficiency, FID and likelihood of DDPM/DDIM models.
Problem
Diffusion models with fixed or heuristic denoising covariances need many sampling steps for good quality and likelihood; exact optimal diagonal covariance requires expensive Hessian diagonal estimation at sampling time.
Approach
Use a generalized analytical covariance identity linking the posterior covariance to the score Hessian, then amortize the Hessian diagonal with a small network trained by an unbiased OCM objective (one Rademacher sample per step), added in parallel to the frozen score network at negligible cost.
Key results
- On CIFAR10 (CS) with 10 steps, OCM-DDPM reaches FID 14.32 versus 34.76 for DDPM with beta-tilde and 16.33 for SN-DDPM.
- On ImageNet 256x256 (DiT), OCM-DPM needs only 10 steps to reach FID around 6 (5.33 reported), the fewest among compared methods, at 1.007x time per step.
- Likelihood: OCM-DDPM gives best or second-best NLL in most settings, e.g. CIFAR10 (LS) 5.32 bits/dim at 10 steps vs 74.95 for DDPM beta-tilde.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace SN-DDPM covariance estimation with OCM covariance estimation | DDPM · CIFAR10 (CS), 10 timesteps | FID: 16.33 -> 14.32 | Helped |
| Replace SN-DDPM covariance estimation with OCM covariance estimation | DDPM · CelebA 64x64, 10 timesteps | NLL (bits/dim): 4.46 -> 4.69 | Hurt |
| Replace SN-DDPM covariance estimation with OCM covariance estimation | DDPM · ImageNet 64x64, 25 timesteps | NLL (bits/dim): 4.56 -> 4.45 | Helped |
| Use OCM covariance estimation instead of SN-DPM | Diffusion sampling · CIFAR10, FID around 6 | least number of timesteps: 1.005×{17} (5.81) -> 1.003×\textbf{16} (5.83) | Helped |
| Use OCM covariance estimation instead of SN-DPM | Diffusion sampling · CelebA 64x64, FID around 6 | least number of timesteps: 1.019×{22} (5.96) -> 1.015×\textbf{21} (5.94) | Helped |
| Use OCM covariance estimation instead of the heuristic $\tilde{\beta}$ covariance | DDPM · CIFAR10 (CS), 5 timesteps | FiD: 34.76 -> 38.88 | Helped |
Surprising findings
- SN-DDPM achieves very poor NLL at small step counts on CIFAR10 (e.g. 30.79 at 10 steps) despite competitive FID, attributed to error amplification in its quadratic term.
- DDPM with beta (fixed isotropic) obtains the best FID at 1000 steps on CIFAR10 but is far worse at small step counts.
Limitations
- Diagonal covariance only; full covariance would be better but is impractical at high dimension, as the paper notes.
- Gains on some settings are small or second-best; NPR-DDPM is better on certain NLL cases and SN-DDPM on some FID cases.
- Experiments rely on pretrained score networks and fixed sampling recipes (covariance clipping, even trajectory), and no code release is stated.
Experimental setup
Datasets: CIFAR10 CelebA 64x64 ImageNet 64x64 LSUN Bedroom ImageNet 256x256 2D mixture of Gaussians toy data
Benchmarks: CIFAR10 CelebA 64x64 ImageNet 64x64 LSUN Bedroom ImageNet 256x256
Metrics: FID NLL (bits/dim) MMD Recall wall-clock time per NFE
Baselines: DDPM (beta) DDPM (beta-tilde) A-DDPM I-DDPM NPR-DDPM SN-DDPM DDIM DiT
Builds on: DDPM DDIM Analytic-DPM Improved DDPM SN-DDPM
Closest papers by meaning
- Free Hunch: Denoiser Covariance Estimation for Diffusion Models Without Extra Costs (ICLR 2025 · similarity 0.92)
- Diffusion Models With Learned Adaptive Noise (NeurIPS 2024 · similarity 0.89)
- O(d/T) Convergence Theory for Diffusion Probabilistic Models under Minimal Assumptions (ICLR 2025 · similarity 0.89)
- Improving Diffusion Models for Inverse Problems Using Optimal Posterior Covariance (ICML 2024 · similarity 0.89)
- Generalization through variance: how noise shapes inductive biases in diffusion models (ICLR 2025 · similarity 0.88)
- A Mixture-Based Framework for Guiding Diffusion Models (ICML 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.