# Improving Probabilistic Diffusion Models With Optimal Diagonal Covariance Matching

arXiv: https://arxiv.org/abs/2406.10808 · HTML: https://mlsanity.com/paper/2406.10808/

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 52 of its 72 references that are in the map.
- First posted on arXiv 16 Jun 2024.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 52)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [DreamFusion: Text-to-3D using 2D Diffusion](https://mlsanity.com/paper/2209.14988/) (2022)
- [Stable Video Diffusion: Scaling Latent Video Diffusion Models to Large Datasets](https://mlsanity.com/paper/2311.15127/) (2023)
- [Video Diffusion Models](https://mlsanity.com/paper/2204.03458/) (2022)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [ProlificDreamer: High-Fidelity and Diverse Text-to-3D Generation with Variational Score Distillation](https://mlsanity.com/paper/2305.16213/) (2023)
- [Variational Diffusion Models](https://mlsanity.com/paper/2107.00630/) (2021)
- [Diffusion-LM Improves Controllable Text Generation](https://mlsanity.com/paper/2205.14217/) (2022)
- [Pseudo Numerical Methods for Diffusion Models on Manifolds](https://mlsanity.com/paper/2202.09778/) (2022)
- [Equivariant Diffusion for Molecule Generation in 3D](https://mlsanity.com/paper/2203.17003/) (2022)
- [Score-based Generative Modeling in Latent Space](https://mlsanity.com/paper/2106.05931/) (2021)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [AudioLDM: Text-to-Audio Generation with Latent Diffusion Models](https://mlsanity.com/paper/2301.12503/) (2023)
- [VideoCrafter1: Open Diffusion Models for High-Quality Video Generation](https://mlsanity.com/paper/2310.19512/) (2023)
- [Tackling the Generative Learning Trilemma with Denoising Diffusion GANs](https://mlsanity.com/paper/2112.07804/) (2021)
- [Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion](https://mlsanity.com/paper/2310.02279/) (2023)
- [Assessing Generative Models via Precision and Recall](https://mlsanity.com/paper/1806.00035/) (2018)
- [Analytic-DPM: an Analytic Estimate of the Optimal Reverse Variance in Diffusion Probabilistic Models](https://mlsanity.com/paper/2201.06503/) (2022)

## Built on by (most cited first, 13 of 13)

- [Diffusion and Flow-based Copulas: Forgetting and Remembering Dependencies](https://mlsanity.com/paper/2509.19707/) (2025)
- [Improved probabilistic regression using diffusion models](https://mlsanity.com/paper/2510.04583/) (2025)
- [Few-Step Boltzmann Generators via Scalable Likelihood Flow Maps](https://mlsanity.com/paper/2606.29110/) (2026)
- [Outlier-robust Diffusion Posterior Sampling for Bayesian Inverse Problems](https://mlsanity.com/paper/2602.02045/) (2026)
- [Diffusion Alignment Beyond KL: Variance Minimisation as Effective Policy Optimiser](https://mlsanity.com/paper/2602.12229/) (2026)
- [Covariance-aware sampling for Diffusion Models](https://mlsanity.com/paper/2605.13910/) (2026)
- [Spatially Adaptive Noise Injection](https://mlsanity.com/paper/2609.18466/) (2026)
- [Improved Distributional Diffusion Models](https://mlsanity.com/paper/2609.37147/) (2026)
- [Likelihood Matching for Diffusion Models](https://mlsanity.com/paper/2508.03636/) (2025)
- [Demystifying Transition Matching: When and Why It Can Beat Flow Matching](https://mlsanity.com/paper/2510.17991/) (2025)
- [The Value of Covariance Matching in Gaussian DDPMs and the Lanczos Sampler](https://mlsanity.com/paper/2605.22723/) (2026)
- [Localizing Memorized Regions in Diffusion Models via Coordinate-Wise Curvature Differences](https://mlsanity.com/paper/2605.26756/) (2026)
- [Is your uncertainty map wrong, or is its target? Exact diagnostics for the Tweedie diagonal, and a gradient-free alternative](https://mlsanity.com/paper/2609.33786/) (2026)

## Abstract (arXiv)

The probabilistic diffusion model has become highly effective across various domains. Typically, sampling from a diffusion model involves using a denoising distribution characterized by a Gaussian with a learned mean and either fixed or learned covariances. In this paper, we leverage the recently proposed covariance moment matching technique and introduce a novel method for learning the diagonal covariance. Unlike traditional data-driven diagonal covariance approximation approaches, our method involves directly regressing the optimal diagonal analytic covariance using a new, unbiased objective named Optimal Covariance Matching (OCM). This approach can significantly reduce the approximation error in covariance prediction. We demonstrate how our method can substantially enhance the sampling efficiency, recall rate and likelihood of commonly used diffusion models.