Improving Probabilistic Diffusion Models With Optimal Diagonal Covariance Matching

2024-06-16

AI Summary

Diffusion models suffer from large approximation errors in diagonal covariance prediction, which hinder both sampling efficiency and generation quality. To address this, we propose Optimal Covariance Matching (OCM): instead of data-driven estimation, OCM directly regresses the analytically optimal diagonal covariance and constructs an unbiased gradient objective. By enforcing moment matching of the covariance and incorporating a Gaussian prior, OCM improves covariance accuracy at the parametrization level, substantially reducing Gaussian approximation bias. Experiments across mainstream diffusion models demonstrate consistent improvements: log-likelihood gains of +0.12–0.28, FID reductions of −2.3–4.1, and 1.3–1.8× faster sampling. Crucially, OCM provides the first empirical validation of a strong positive correlation between covariance modeling fidelity and generative performance. This work establishes a new paradigm for efficient, high-fidelity sampling in diffusion models.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Economics, Online Markets and Human Computation: Data quality aspects of human-annotated datasets

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

Problem

• Enhance diffusion model sampling efficiency.
• Reduce covariance prediction approximation error.
• Improve recall rate and likelihood in models.

Innovation

• Optimal Covariance Matching introduced
• Diagonal covariance regression applied
• Sampling efficiency significantly enhanced