Improving Probabilistic Diffusion Models With Optimal Diagonal Covariance Matching
Zijing Ou, Mingtian Zhang, Andi Zhang, Tim Z. Xiao, Yingzhen Li, David Barber
Abstract
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 covariances. Unlike traditional data-driven covariance approximation approaches, our method involves directly regressing the optimal 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 both diffusion models and latent diffusion models.
BibTeX
@inproceedings{
ou2025improving,
title={Improving Probabilistic Diffusion Models With Optimal Diagonal Covariance Matching},
author={Zijing Ou and Mingtian Zhang and Andi Zhang and Tim Z. Xiao and Yingzhen Li and David Barber},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025},
url={https://openreview.net/forum?id=fV0t65OBUu}
}