Improving Probabilistic Diffusion Models With Optimal Diagonal Covariance Matching
Authors: Zijing Ou
Mingtian Zhang
Andi Zhang
Tim Z. Xiao
Yingzhen Li
David Barber
Venue: ICLR 2025 Oral (2025)
Keywords: Diffusion Model
Generative Model
Probalistic Modelling
Primary area: generative models
TL;DR: We introduce Optimal Covariance Matching (OCM), a novel method that improves sampling efficiency and accuracy in diffusion models by directly regressing optimal analytic covariances.

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.

Published as a conference paper at ICLR 2025
IMPROVING PROBABILISTIC DIFFUSION MODELS WITH OPTIMAL DIAGONAL COVARIANCE MATCHING
Zijing Ou'; Mingtian Zhang?: Andi Zhang Tim Z. Xiao Yingzhen Li', David Barber?
1 Imperial College London, University College London, University of Manchester,
Max Planck Institute for Intelligent Systems, Tubingen UUniversity of Tubingen, SIMPRSIS. @ iee@@mmprrrr imperial. wwiiaaaa...... ac. uk
2. ouo2u
m. zhangêcs ucl.ac.uk
ABSTRACT
The probabilistic diffusion model has become highly effective across various do- mains. Typically, sampling from diffusion model involves using a denoising distribution characterized by a Gaussian with a learned mean and either fixed learned covariances. In this paper, we leverage the recently proposed covariance moment matching technique and introduce a novel method for learning the diago- nal covariance. Unlike traditional data-driven diagonal covariance approximation approaches, our method involves directly regressing the optimal diagonal analytic covariance using new, unbiased objective named Optimal Covariance Matching (OCM). This approach can significantly reduce the approximation error in covari- ance prediction. We demonstrate how method can substantially enhance the sampling efficiency, recall rate and likelihood of commonly used diffusion models.
INTRODUCTION
Diffusion models (Sohl-Dicksted et al., 2015; Ho et al., 2020; Song & Ermon, 2019) have achieved remarkable success in modeling complex data across various domains (Rombach et al., 2022; Li et al., 2022: Poole et al., 2022; Ho et al., 2022; Hoogeboom et al., 2022; Liu et al., 2023). A conventional diffusion model operates in two stages: a forward noising process, indexed by E [1, which progressively corrupts the data distribution into Gaussian distribution via Gaussian convolution, and a reverse denoising process, which generates images by gradually transforming Gaussian noise back into coherent data samples. In traditional diffusion models, the generation process typically predicts only the mean of the denoising distribution while using a fixed, pre-defined variance (Ho al., 2020). This approach often requires very large number of steps, T, to produce high-quality, diverse samples or to achieve reasonable model likelihoods, leading to inefficiencies during inference.
To address this inefficiency, several works have proposed estimating the diagonal covariance of denoising distributions rather than relying on pre-defined variance values. For instance, Bao et al. (2022b) introduces an analytical form of isotropic, state state-indepent independent covariance that can be estimated from the data. This analytical solution achieves optimal form under the constraints of isotropy and state independence. A more flexible approach involves learning a state-dependent diagonal covariance. Nichol & Dhariwal (2021) propose learning this form by optimizing the variational lower bound (VLB). Bao et al. (2022a) explore learning a deteeddeendent diagonal covariance directly from data, also examining its analytical form. These methods have demonstrated improved image quality with fewer denoising steps. Learning the covariance through VLB optimization (Nichol & Dhariwal, 2021) has become widely adopted strategy in state-of-the-art image and video generative models.
Building on this line of research, the goal of our paper is to develop an improved denoising covariance strategy that enhances both the generation quality and likelihood evaluation while reducing the number of total time steps. Recently, Zhang et al. (2024b) derived the optimal state-dependent full covariance for the denoising distribution. While this method offers greater flexibility than state- dependent diagonal covariance, it requires O(D²) storage for the Hessian matrix and O network evaluations per denoising step, where D is the data dimension. This makes it impractical for high- dimensional applications, such as image generation. To address this limitation, we propose novel,
"qqual contribution. Code is available at: https: http//ggtttteete thub. com/u- zi 1iliii////0/ ww/CC//MPPPPPP

---

Published as a conference paper at ICLR 2025
unbiased covariance matching objective that enables neural network to match the diagonal of the optimal deleedeennddnt diagonal covariance. Unlike previous methods (Bao et al., 2022a; Nichol & Dhariwal, 2021), which learn the diagonal covariance directly from the data, our approach estimates the diagonal covariance from the learned score function. We show that this method significantly reduces covariance estimation errors compared to existing techniques. Moreover, demonstrate that our approach can be applied to both Markovian (DDPM) and non-Markovian (DDIM) diffusion models, well as latent diffusion models. This results in improvements generation quality, recall, and likelihood evaluation, while also reducing the number of function evaluations (NFEs).
2
BACKGROUND OF PROBABILISTIC DIFFUSION MODELS
We first introduce two classes of diffusion models that will be explored further in our paper.
2.1 MARKOVIAN DIFFUSION MODEL: DDPM
Let ((20) be the true data distribution, denoising diffusion probabilistic models (DDPM) (Ho et al., 2020) constructs the following Markovian process
T
((vor)) = ((xo) ) II (( (=1
where (( = W(V1- - 1,, and ßir is the pre-defined noise schedule. We can also derive the following skip-step noising distribution in closed form
q(=/|<0)) = NN -ā)).),
where a. = [1 B.). When T sufficiently large, we have the marginal distribution q(xr) + NO( I). The generation process utilizes an initial distribution g(zr) and the denoising distribution q(=at-1 (x+)) We assume for large T and variance preversing schedule discussed in Equation 1, we can approximate ((TT) R p(æT) (O, I). The true q( (xx) is intractable and
approximated as variational distribution ( (x)) within a Gaussian family, which defines the reverse denoising process
Po (o(2)-1(z)) = -N 0), [x-r
With Tweedie's Lemma (Efron, 2011; Robbins, 1992), can have the following score representation
( 0) (xt + V. (og))())))/el
Bi,
where the approximated score function Z., log (x:) = Vz, log g(x)) can be learned by the denoising score matching (DSM) (Vincent, 2011: Song & Ermon, 2019).
For the covariance 1 0), two heuristic choices were proposed in the original DDPM paper: 1. Σ, (xr; 0) 3r, which is equal to the variance of g( (rr- and 2. [x=(xx;;)) B, where
3. (1 at-1)/(1 āi) 3; is the variance of q(xx-1 1 z0). Although heuristic, Ho et al. (2020)
show that when large, both options yield similar generation quality.
2.2 NON-MARKOVIAN DIFFUSION MODEL: DDIM
In additional the Markovian diffusion process, denoising diffusion implicit models (DDIM) (Song
et al., 2020) only defines the condition q( (zo) = (V (1 āz) and let
(( R 2opo-o/tro
(5)
where q(2,-1/01, 2o) = NN 0Z with
Mt-1 = Vār-120 + V
(6)
When 0,-1= /(1 ā, a,-1)/(1 33 the diffusion process becomes Markovian and equivalent
to DDPM. Specifically, DDIM chooses - O. which implicitly defines mon-Markovian diffusion
2

---

Published as a conference paper at ICLR 2025
process. In the original paper (Song et al., 2020), the ((z))zz) is heuristically chosen as delta function po (z)) = 6(20 - Ho (zz: 8)) where
[o( 0) = (æt + - logp)))))
In both DDPM and DDIM, the covariance of po(xx-1(zx) or (o are chosen based on heuristics. Nichol & Dhariwal (2021) have shown that the choice of covariance makes a big impact when 7 is small. Therefore, for the purpose of accelerating the diffusion sampling, our paper will focus how to improve the covariance estimation quality in these cases. We will first introduce our method in the next section and then compare it with other methods in Section
DIFFUSION MODELS WITH OPTIMAL COVARIANCE MATCHING
Recently, Zhang et al. (2024b) introduce the optimal covariance form of the denoising distribution (x)]²) x q (z) ) for the Gaussian convolution ((ā(z) = Wxx, o² I)), which can be seen as a high-dimensional extension of the second Tweedie's Lemma (Efron, 2011; Robbins, 1992). We
further extend the formula to scaled Gaussian convolutions in the following theorem.
Theorem 1 (Generalized Analytical Covariance Identity). Given joint distribution q(ā, x)
q((()))) ()) (z) with ((](x) = N(ax, I). then the covariance of the true posterior q(xlī) x
((x) (z z), which is defined as [() B,(=(a) |20] - Ex(=)=)(<]] has a closed form:
Σ = (o'V2 logg((5) +² /a?.
(8)
See Appendix A. for a proof. This covariance form can also be shown as the optimal covariance form under the KL divergences (Bao al., 2022b; Zhang et al., 2024b). We can see that the exact covariance in this case only depends on the score function, which indicates the exact covariance can be derived from the score function. In general, the score function already contains all the information of the denoising distribution gxx| see Zhang et al. (2024b) for further discussion. We here only consider the covariance for simplicity.
In the case of the diffusion model, we
Table 1: Comparison of different covariance choices on CI- FARIO (CS). We use Rademacher estimator (M = 100) for the diagonal covariance approximation in OCM-DDPM.
use the learned score function as plug-in approximation in Equation (8). Although the optimal covariance can be directly calculated from the learned score function, but requires calculat- ing the Hessian matrix, which is the Jacobian of the score function. This requires O(D²) storage and D num- ber of network evaluation (Martens
FiD 4
NLL
timesteps
10
15
20
10
15
20
DDPM,
$8.28
34.76
24.02
19.00
203.29
74.95
44
32.20
DDPM,
254.07
205.31
149.67
109.81
7.33
6.51
6.06
5.77
OCM-DDPM 38.88
21.60
13.35
9.75
6.82
4.98
4.62
4.43
et al., 2012) for each denoising step the time t. Zhang et al. (2024b) propose to use the following consistent diagonal approximation (Bekas et al., 2007) to remove the O(D²) storage requirement
diag(H (H))) x 1/M Σ m=1 Um O (i)
where Hi))= V log ( and Um ~(()) is Rademacher random variable (Hutchinson, 1990) with entries +1 and @ denotes the element- wise product. In Table 1, we compare the generation quality of DDPM using different covariance choices. The results demonstrate that generation quality improves significantly when using the Rademacher estimator to estimate the optimal covariance, compared to the heuristic choices of /3 and B. However, as shown in Figure 5 in the Appendix, achieving a desirable approximation on the CIFARIO dataset necessitates M 2 100 Rademacher samples. Each calculation of um & H (ă) (T)o,, requires a forward pass and a backward propagation, leading to roughly 2M network evaluations in total. This significantly slows down the generation speed, making impractical for diffusion models. Inspired by Nichol & Dhariwal (2021); Bao et al. (2022a) and also the amortization technique used in variational inference (Kingma & Welling, 2013; Dayan et al., 1995), we propose to use network to match the diagonal Hessian, which only requires one network pass to predict the diagonal Hessian in the generation time and can be done in parallel with the score/mean predictions with extra time cost. In the next section, we introduce novel unbiased objective to learn the diagonal Hessian from the learned score, which improves the covariance estimation accuracy and leads to better generation quality and higher likelihood estimations.
3

---

Published as a conference paper at ICLR 2025
Algorithm 1 Sampling procedure from L' in OCM-DDPM
Algorithm Sampling procedure from in OCM-DDIM
1: Compute (1r( with Equation (15);
1: Compute (zrs) with Equation (18);
2: Compute E,, (are) with Equation (16);
2: Compute Eo(xx) with Equation (19);
3: Sample (0, I);
3: Sample 3o (H)(E)))))
4: Le + per (24) + (
4: Let + Vāg vI ae Varao
3.1 UNBIASED OPTIMAL COVARIANCE MATCHING
To train network ha match the Hessian diagonal diag(HI a straightforward solution is to directly regress Equation 9 for all the noisy data
M
min E ((E) I
Σ 2m m=1
(10)
M
where Um ~ p(v). Although this objective is consistent when M + co, it will introduce additional bias when M is small. To avoid the bias, we propose the following unbiased optimal covariance matching (OCM) objective
Truem (6) (z)) (() - O (G)elll,
(11)
which does not include expectation within the non-linear L2 norm. The following theorem shows the validity of the proposed OCM objective.
Theorem 2 (Validity of the OCM objective). The objective in Equation (11) upper bounded the base objective (i.e., Equation (10) with M + oo). Moreover, attains optimal when hz( = diag( H (3))
for all ē g(z).
See the Appendix A.2 for proof. The integration over e in Equation (11) can unbiasedly approximated by the Monte-Carlo integration given M Rademacher samples we found
p(v). In practice, works well (see Table 11 in the appendix for the ablation study on varying M values), which also shows the training efficiency of the proposed OCM objective. The learned hg can form the covariance approximation
[(;; Φ) (o 4 ( + 1/?
(12)
We then discuss how to apply the learned covariance to diffusion models.
3.2 DIFFUSION MODELS APPLICATIONS
Given access a learned score function, Vz, log po (2z)., from any pre-trained diffusion model, we denote the Jacobian of the score as (xz)) Assuming M = 1 in the OCM training objective, the covariance learning objective for diffusion models can be expressed as follows:
min Σ bptr - (.
-
(13)
where v ~ p(u) and ha (2.) is a network that conditioned on the state x; and time After training this objective, the learned (xx) can be used to form the diagonal Hessian approximation hā( R diag(H [ (z)) which further forms our approximation covariance. We then derive its use cases in skip-step DDPM and DDIM for Diffusion acceleration.
Skip-Step DDPM: For the general skip-step DDPM with denoising distribution q( (xx)) with < t. When = - 1, this becomes the classic one-step DDPM. We further denote ōva = II Cg. and thus ag.t = ã.,, ae ar.. We can write the forward process
g( (rr) (a (1 ā..)().
(14)
The corresponding Gaussian denoising distribution pe,o (x.) = N (pe 0), Ee o)) has the following mean and covariance functions:
Pe (x 0) = (It + (1 - āea) Vz, log (x:))) Voru.
(15)
Ey (2re; Φ) ((1 -ār "h (o + -
(16)
4

---

Published as a conference paper at ICLR 2025
Table 2: Overview of different covariance estimation methods. Methods are ranked by increasing modeling capability from top to bottom. We also include the intuition of the methods and how many additional network passes are required for estimating the covariance.
Covariance Type
+assss
Intuition
Modeling Capability
#y-independent Isotropic /3 (Ho et al., 2020)
Cov. of g(
#y-independent Isotropic 3 (Ilo al., 2020)
0
Cov. of q[ vo)
#,-independent Isotropic Estimation (Bao al., 2022b)
Estimate from data
27-dependent Diagonal VLB (Nichol & Dhariwal, 2021)
Learn from data
2/-dependent Diagonal NS (Bao al., 2022a)
1
Learn from data
2y- 2x-dependent Diagonal OCM (Ours)
Learn from score
2: dedeennddnt Diagonal Estimation (Zhang et al., 2024b)
2M
Estimate from score
dependentt Diagonal Analytic (Zhang et al., 2024b)
D
Calculate from score
dedeenddent Full Analytic (Zhang al., 2024b)
Calculate from score
The skip-step denoising sample zy (e. can be obtained by Ze (x): 12/ ())
Skip-Step DDIM: Similarly, we give the stip-step formulation of DDIM, where we set the at-1 O in Equation (6) which is also used in the original paper Song et al. (2020). We can use the approximated covariance of pe, (xooZz) to replace the delta function used in the vanilla DDIM, which gives the skip-steps DDIM sample
Vāeso + vI ae .ervvi-ue 04 (x+ - V Vā,zo),
(17)
where xo (3 NNuol 0); Eo (30 ()) and
00(2;;0) = (xt + (1 a. V., log (oggo((x))) V Eo(or; Φ) = ((1 -z) (xc: t) + (1 āx 1t)
(18)
(19)
The sample zo (1) can be obtained by xu #o 0) 0) e (2.;3 Ø). In the next section, we will compare the proposed method to other covariance estimation methods in practical examples.
3.3 DETAILS OF TRAINING AND INFERENCE
Our model comprises two components: score prediction network 8a and diagonal Hessian prediction network hg. In line with Bao et al. (2022a), we parameterize the score prediction network using a pretrained diffusion model, and the Hessian prediction network is parameterized by sharing parameters as follows:
(xx.) NN, (Ba,sNee(zr,
ha(xx) (x,
(20)
where BaseNet represents the commonly used architecture in diffusion models, such as UNet and DiT (Ho et al., 2020; Peebles & Xie, 2023). This parameterization approach only requires additional small neural network, NN2, resulting in negligible extra computational and memory costs compared to the original diffusion models (see Appendix B.1 for more details). In our experiment, fix the parameter = {01, 02) and train the Hessian prediction network exclusively with the proposed OCM objective. After training, samples can be generated using Algorithms I and 2.
4 RELATED COVARIANCE ESTIMATION METHODS
We then discuss different choices of Eg (2x, used in the diffusion model literature, see also Table for an overview. For brevity, we mainly focus on DDPM here, in which the mean py (2x) can be computed as in Equation (15).
1. zr-inndeeeneentt isotropic covariance: DDPM (Ho et al., 2020). The B-DDPM uses the variance of p(zy(1,) which is Ex((x,) = (1 Ge/0p J,, when t'= 1, we have 1x-1(x,) = 3c- The 3-DDPM uses the covariance of p(ær (2o, 21), which is (1-ā.)) (1 av.)).
(I-ac)
5

---

Published as a conference paper at ICLR 2025
2. zj-independent isotropic covariance: A-DDPM (Bao et al., 2022b). A-DDPM assumes state-independent isotropic covariance Ev((x)) = ofd with the following analytic form of
o = 1-ār' (1 (1-0w'x) By(z))
V log po (z))122
Qea
dâvi
where the integration g(xx)) requires a Monte Carlo estimation before conducting generation. This variance choice is optimal under the KL divergence within constrained isotropic sate-indeepndent posterior family.
xe-dependent diagonal covariance: I-DDPM (Nichol & Dhariwal, 2021). In I-DDPM, the diagonal covariance matrix is modelled as the interpolation between 3r and 3.
Ee (zy; v) = exp(ee (xt) log 3r + 11--40 vo(x)))) 3r),
(21)
where vy is parametrized via neural network. The covariance is learned with the variational lower bound (VLB). In the optimal training, the covariance learned by VLB will recover the true covariance in (8). Notably, Nichol & Dhariwal (2021) heuristically obtain the covariance of skip sampling
Ey (Ze; 1 by rescaling Br and 3, accordingly: 3r + Cyt Br - (ي=يا) (1 .v:r). However, when t' 0, Eo(x;; 16) is ill-defined; thus, iDDPM is inapplicable within the DDIM framework.
4. dependent diagonal covariance: SN-DDPM (Bao et al., 2022a). SN-DDPM learns the covariance by training neural network Io to estimate the second moment of the noise E (x: varzo)/VI- āgi
min gg 11 (xx))1
(22)
&
After training, the covariance
Ep (Ire: eb)
can be estimated via
(1 (1-ar) av) (x,)
82
go (2) Vz, loggo(t:)))
En (xri
3
(23)
Ot 1-
where Be. = 1 ar In optimal, Eg (x2; 2b) in Equation (23) will recover the true covariance in Equation (8). We demonstrate the equivalence between OCM-DDPM and SN-DDPM Ap- pendix A.3. However, due to the appearance of the quadratic term in Equation (23), SN-DDPM tends to amplify the estimation error as - O, leading to suboptimal solutions. To mitigate this issue, Bao et al. (2022a) also propose NPR-DDPM, which models the noise prediction residual instead. We recommend referring to their paper for detailed explanations.
Notably, almost all these methods can be applied within the DDIM framework by setting E' = 0. Specifically, p p (zs) can be sampled using Equation (17) with xg ~ (N)(( Eo(x)), where (2)) is the same in (18), but Eg (zx) differs for various methods as discussed previously.
EXPERIMENTAL RESULTS
To support theoretical discussion, first evaluate the performance of optimal covariance matching by training diffusion probabilistic models on 2D toy examples. We then demonstrate its effectiveness in enhancing image modelling in both pixel and latent spaces, focusing the comparison between optimal covariance matching and other covariance estimation methods, and showing that the proposed approach has the potential to scale large image generation tasks.
5.1 TOY DEMONSTRATION
We first demonstrate the effectiveness of our method by considering the data distribution two- dimensional mixture fonty Gaussians (MoG) with means uniformly distributed over [-40, 40] & [-40, 40) and a standard deviation of = /40, where & denotes Cartesian product (see Figure la for visualization). In this case, both the true score Vz. p(xz) and the optimal covariance (xx) are available, allowing us to compare the covariance estimation error. We then learn the covariance using the true scores by different methods conduct DDPM and DDIM sampling for this MoG

---

Published as a conference paper at ICLR 2025
DDPM-6
DDPM-8
A-DDPM
NPR-DDPM
SN-DDPM
OCM-DDPM (Ours)
K=25
K=50
K= 100
K=200
10
FID
RID
NLL
Figure 2: The results of FID NLL for different methods with varying numbers of sampling steps on CIFARIO (CS). Our method consistently achieves the best trade-off between FID and NLL.
problem. For evaluation, we employ the Maximum Mean Discrepancy (MMD) (Gretton et al., 2012). which utilizes five kernels with bandwidths {2- 2,2-1, 22, 2 2²). In Figures and 4d, we show the MMD comparison among different methods. Specifically, we choose to conduct different diffusion steps with the skip-step scheme as discussed in Section 3.2.
The results, shown in Fig- ures 1b and Ic, demonstrate that the proposed method outperforms other covariance learning approaches. Addi- tionally, we include two meth- ods utilizing the true diago- nal and full covariance, serv- ing benchmarks for the best achievable performance. No- tably, in this case, using the full covariance yields better performance compared to the diagonal covariance. This ob- servation highlights the importance of accurate covariance approximation, suggesting that improved methods for learning the off-diagona, terms could lead to even better results. For further analysis, present an additional toy experiment Appendix C. to demonstrate the efficacy of our approach.
DOPK, GT DNGS
DOML DO
DUPM,
IDDP
ADCPM
GEN-OOI OOM (0um)
(a) Training Data and Density (b) DDPM MMD Sieps
(c) DDIM MMD Steps
Figure I: Comparisons of different covariance estimation methods.
Figure (a) demonstrates the training data and the ground truth density. Figures (b) and (c) present the MMD evaluation against the total sampling steps in the DDPM (b) and DDIM (c) settings.
5.2 IMAGE MODELLING WITH DIFFUSION MODEL.
Following the experimental setting Bao al. (2022a), we then evaluate our method across varying pre-trained score networks provided by previous works (Ho et al., 2020; Nichol & Dhariwal, 2021; Song et al., 2020; Bao al., 2022b; Peebles & Xie, 2023). In this experiment, we mainly focus on four datasets: CIFAR 10 (Krizhevsky et al., 2009) with the linear schedule (LS) of 8 (Ho et al., 2020) and the cosine schedule (CS) of 3e (Nichol & Dhariwal, 2021); CelebA (Liu et al., 2015); LSUN Bedroom (Yu et al., 2015). The details of the experimental setting can be found in Appendix B.
IMPROVING LIKELIHOOD RESULTS
We first evaluate the negative log- likelihood (NLL) of the diffusion models by calculating the negative evidence lower bound (ELBO): log (zx)) < Ey log p(za:r) Leeloo (Io) with the same mean yet varying covariances. As per Bao al. (2022b), we report results only for the DDPM forward process, as L = oo under the DDIM forward process. As shown in Table 3. our OCM-DDPM demonstrates the best or second-best performance most scenarios, highlighting the effectiveness of our covariance estimation approach. We also note that SN-DDPM performs poorly on likelihood results in small-scale datasets like CIFARIO and CELEBA. likely due to the amplified error of the quadratic term in Equation (22). Although NPR-DDPM achieves slightly better performance in certain scenarios, it falls short in terms of sample quality, as measured by FID, shown in Table 5. In Appendix C.2, also compare our method to Improved DDPM (Nichol & Dhariwal, 2021) in terms of likelihood on ImageNet 64x64, The results show that our method achieves better likelihood estimation with small number of timesteps, while achieving comparable results with full timesteps.
q(>1:7/20)
7

---

Published as a conference paper at ICLR 2025
Table 3: The NLL (bits/dim) & across various datasets using different sampling steps.
CTFAR 10 (I.S)
CIFARTO (CS)
# TIMESTEPS K
10
25
50
100
200
1000
10
25
50
100
200
1000
DDPM, 3
74.95
24.98
12.01
7.08
5.03
3.73
75.96
24.94
11.96
7.04
4.95
3.60
DDPM, 8
6.99
6.11
5.44
4.86
4.39
3.75
6.51
5.55
4.92
4.41
4.03
3.54
A-DDPM
5.47
4.79
4.38
4.07
3.84
3.59
5.08
4.45
4.09
3.83
3.64
3.42
NPR-DDPM
5.40
4.64
4.25
3.98
3.79
3.57
5.03
4.33
3.99
3.76
3.59
3.41
SN-DDPM
30.79
11.83
7.13
5.24
4.39
3.74
90.85
19.81
9.72
5.58
4.73
OCM-DDPM
5.32
4.63
4.25
3.97
3.78
3.57
4.99
4.34
3.99
3.76
3.59
3.41
CELEBA 64X64
IMAGENET 64X64
TIMESTEPS K
10
25
50
100
200
1000
2,5
50
100
200
400
4000
DDPM, 3
33.42
13.09
7.14
4.60
3.45
2.71
105.87
46.25
22.02
12.10
7.59
3.89
DDPM, 3
6.67
5.72
4.98
4.31
3.74
2.93
5.81
5.20
4.70
4.31
4.04
3.65
A-DDPM
4.54
3.89
3.48
3.16
2.92
2.66
4.78
4.42
4.15
3.95
3.81
3.61
NPR-DDPM
4.46
3.78
3.40
3.11
2.89
2.65
4.66
4.22
3.96
3.80
3.71
3.60
SN-DDPM
18.09
8.05
5.29
4.05
3.40
2.84
4.56
4.18
3.95
3.80
3.71
3.63
OCM-DDPM
4.69
3.86
3.43
3.13
2.90
2.66
4.45
4.15
3.93
3.79
3.70
3.59
IMPROVING SAMPLE QUALITY
Next, we compare sample quality quantitatively using the FID score, with results reported for both the DDPM and DDIM forward processes. As shown in Table 5, our OCM-DPM consistently achieves either the best second-best results across most settings. The qualitative results of the samples generated by our method can be found in Appendix C.6.
Notably, while SN-DDPM often demonstrates slightly better FID performance in certain cases, struggles with likelihood estimation, as shown in Table 3. Conversely, as discussed in the previous section, PPR-DDPM sometimes achieves superior likelihood results in Table 3, but its image generation quality is significantly worse than OCM-DDPM. Therefore, the proposed OCM-DDPM provides a more balanced model, offering a better trade-off between generation quality and likelihood estimation. To better highlight this benefit, we include visualization in Figure 2.
5.3 IMAGE MODELLING WITH LATENT DIFFUSION MODELS
In this section, we apply our meth- ods to latent diffusion models (Vahdat et al,, 2021; Rombach et al., 2022) and with ImageNet 256 demon- strate the scalability of our method. Specifically, we compare our methods to other approaches within the DiT architecture (Peebles & Xie, 2023), evaluating sample quality using the FID score and diversity using the Re- call score (Sajjadi al., 2018). We focus on conditional generation with classifier-free guidance (CFG), which is consistent Peebles & Xie (2023).
0.5
DDPM, $
100PH
DDIM
GCM- DDPM
OCM-DDIM
6:90
0.2
1.59
6.50
23
CPG
CFG
Figure 3: Results of DiT training on ImageNet 256x256.
We generate samples using 10 timesteps with varying CFG coefficients (see Table 12 for exact numerical values).
is important to note that DIT was initially trained using the Improved DDPM (I-DDPM) algorithm (Nichol Dhariwal, 2021). Therefore, the results I-DDPM reflect the original performance of the pre-trained DiT model' provided by Peebles & Xie (2023). For DDPM and DDIM sampling, we discard the learned covariance DiT and use only the learned noise prediction neural network for the posterior mean. In contrast, our OCM-DPM retains the same mean function while learning covariance using the optimal covariance matching objective. We then generate samples with 10 sampling steps and report FID and Recall scores across different CFG coefficients the ImageNet 256x256 dataset. The results are displayed in Figure 3, with DDPM-3 excluded due to its poor performance. It shows that OCM-DPM method achieves the best FID performance with CFG 2.0. Although DDPM-3 and DDIM perform better in terms of FID at CFG 1.5, their Recall scores
We the DiT XI.2 checkpoint relcased in https::///// //g wtt.:/oco/aaacccaaoae faw.book
8

---

Published as a conference paper at ICLR 2025
Table 4: The least number of timesteps 4 required to achieve an FID around 6 (along with the corresponding FID). To account for the additional time cost incurred by the covariance prediction network, we multplyy the results by the ratio of the time cost per single timestep, reflecting the extra computational overhead (see Appendix B.1 for details).
METHOD
CIFARIO
CELEBA 64x64
LSUN BEDROOM IMAGENET 256X256
DDPM
90 90 (6.12)
> 200
130 (6.06)
21 (5.89)
DDIM
30 (5.85)
> 100
BEST FID > 6
11 (5.58)
IMPROVED DDPM
45 (5.96)
MISSING MODEL.
90 (6.02)
22 (6.08)
ANALYTIC-DPM
25 (5.81)
55 (5.98)
100 (6.05)
MISSING MODEL. MISSING MODEL
NPR-DPM
1.002 (5.76)
1.013x50 (6.04)
1.021 (6.01)
SN-DPM
1.005 x (5.81)
1.019x22 (5.96)
1.114x92 (6.02)
MISSING MODEL 1.007 x (5.33)
OCM-DPM (OURS)
1.003 x (5.83)
1.015x21 (5.94) 1.112x90 (6.04)
are lower than those of the proposed OCM methods. This discrepancy arises partly because DDPM and DDIM model the inverse process with deterministic variance, whereas OCM methods explicitly learn the variance from data, allowing for more accurate density estimation and generating more diverse results. As Peebles & Xie (2023) report that DiT achieves the best FID performance with CFG = 1.5, we further evaluate performance with different sampling steps under CFG = 1.5, as shown in Table 13. The results indicate that our OCM-DPM methods offer a great improvement in sample quality and diversity.
To provide overview of the superiority of our method, compare the minimum sampling step to achieve the FID score (Heusel et al., 2017) close to 6 across different approaches (see Appendix B. for details). As shown in Table 4, method requires the fewest denoising steps in most settings. Moreover, we emphasize that compared to recent advanced methods like SN-DPM and NPR-DPM our approach is the only one that consistently delivers both competitive likelihood and FID.
RELATED WORK AND FUTURE DIRECTIONS
In this paper, we show that improving diagonal covariance estimation can significantly enhance the performance of diffusion models. This raises a natural question: can more flexible covariance structures further improve these models? In Figure 1, we demonstrate that a full covariance structure achieves better generation quality with fewer NFEs in a toy 2D problem. However, for high- dimensional problems, the quadratic growth in parameter scale makes full covariance approach computationally impractical. To address this, low-rank or block -diagonaa covariance approximations offer promising alternatives. Developing effective training objectives for flexible covariance structures remains a compelling direction for future research.
In addition to the covariance estimation methods discussed in Section 4, there are other approaches to accelerate sampling in diffusion models. One approach involves using faster numerical solvers for differential equations with continuous timesteps
Marttinnaa et al., 2021; Liu al., 2022; Lu et al., 2022). Another strategy, inspired by Schrödinger bridge (Wang et al., 2021; De Bortoli al., 2021), is introduce nonlinear, trainable forward diffusion process. Additionally, replacing Gaussian modelling of py (zo ( with more expressive alternatives, such as GANs (Xiao et al., 2021), 1), distributional models (Bortoli al., 2025), latent variable models (Yu et al., 2024), or energy-based models (Xu et al., 2024), can also accelerate the sampling of diffusion models.
Recently, distillation techniques have gained popularity, achieving -tate-of-the-art in one-step genera- tion (Zhou et al., 2024). There are two prominent types of distillation methods. The first is trajectory distillation (Salimans & Ho, 2022: Berthelot et al., 2023; Song et al.. 2023; Heek al., 2024: Kim et al., 2023; Li & He, 2024), which focuses on accelerating the process of solving differential equa- tions. The second type involves distillation techniques that utilize a one-step implicit latent variable model as the student model (Luo et al., 2024; Salimans et al., 2024; Xie et al., 2024; Zhou et al., 2024; Zhang et al., 2025) and distills the diffusion process into the student model by minimizing the spread divergence family (Zhang et al., 2020; 2019) through score estimation (Poole et al., 2022; Wang et al., 2024).
Although these distillation methods typically offer faster generation speeds (fewer NFEs) compared to covariance estimation methods, they often lack tractable density or likelihood estimation. This presents challenge for generative modeling applications where likelihood density estimation is

---

Published as a conference paper at ICLR 2025
Table 5: FID score J across various datasets using different sampling steps.
CIFARIO (LS)
CIFARIO (CS)
# TIMESTEPS K
10
25
50
100
200
1000
10
25
50
100
200
1000
DDPM, B
44.45
21.83
15.21
10.94
8.23
5.11
34.76
16.18
11.11
8.38
6.66
4.92
DDPM.
233.41
125.05
66.28
31.36
12.96
3.04
205.31
84.71
37.35
14.81
5.74
3.34
A-DDPM
34.26
11.60
7.25
5.40
4.01
4.03
22.94
8.50
5.50
4.45
4.04
4.31
NPR-DDPM
32.35
10.55
6.18
4.52
3.57
4.10
19.94
7.99
5.31
4.52
4.10
4.27
SN-DDPM
24.06
6.91
4.63
3.67
3.31
3.65
16.33
6.05
4.17
3.83
3.72
4.07
OCM-DDPM
24.94
9.19
5.95
4.36
3.48
3.98
14.32
5.54
4.10
3.84
3.75
4.18
DDIM
21.31
10.70
7.74
6.08
5.07
4.13
34.34
16.68
10.48
7.94
6.69
4.89
A-DDIM
14.00
5.81
4.04
3.55
3.39
3.74
26.43
9.96
6.02
4.88
4.92
4.66
NPR-DDIN
13.34
5.38
3.95
3.53
3,42
3.72
22.81
9.47
6.04
5.02
5.06
4.62
SN-DDIM
12.19
4.28
3.39
3.23
3.22
3.65
17.90
7.36
5.16
4.63
4.63
4.51
OCM-DDIM
10.66
4.35
3.48
3.27
3.29
3.74
16.70
6.71
4.72
4.30
4.54
4.53
CELEBA 64x64
IMAGENET 64x64
# TIMESTEPS K
10
25
50
100
200
1000
25
50
100
200
400
4000
DDPM, 3
36.69
24.46
18.96
14.31
10.48
5.95
29.21
21.71
19.12
17.81
17.48
16.55
DDPM, 3
294.79
115.69
53.39
25.65
9.72
3.16
170.28
83.86
45.04
28.39
21.38
16.38
A-DDPM
28.99
16.01
11.23
8.08
6.51
5.21
32.56
22.45
18.80
17.16
16.40
16.34
NPR-DDPM
28.37
15.74
10.89
8.23
7.03
5.33
28.27
20.89
18.06
16.96
16.32
16.38
SN-DDPM
20.60
12.00
7.88
5.89
5.02
4.42
27.58
20.74
18.04
16.61
16.37
16.22
OCM-DDPM
21.55
12.71
9.24
6.97
5.92
5.04
28.02
20.81
17.98
16.74
16.32
16.31
DDIM
20.54
13.45
9.33
6.60
4.96
3.40
26.06
20.10
18.09
17.84
17.74
19.00
A-DDIM
15.62
9.22
6.13
4.29
3.46
3.13
25.98
19.23
17.73
17.49
17.44
18.98
NPR-DDIM
14.98
8.93
6.04
4.27
3.59
3.15
28.84
19.62
17.63
17.42
17.30
18.91
SN-DDIM
10.20
5.48
3.83
3.04
2.85
2.90
28.07
19.38
17.53
17.23
17.23
18.89
OCM-DDIM
10.28
5.72
4.42
3.54
3.17
3.03
28.28
19.62
17.71
17.42
17.26
19.02
crucial. For example, in A14Science problems involving energy- based data, accurate model density estimation is essential for enabling importance sampling to correct sample bias, discussed in (Zhang al., 2024a; Vargas et al., 2023; Chen et al., 2024; Akhound- Sadegh et al., 2024). Another example is diffusion model-based data compression, where better likelihood corresponds to better compression rates (Townsend et al., 2019; Zhang et al., 2021; Ho et al., 2020; Kingma et al., 2021). 1). In these cases, the proposed OCM-DDPM be used to provide better likelihood estimates, resulting in improved task performance. Additionally, a recent paper (Salimans et al., 2024) shows how momenttmatchin improves distillation, achieving state-of-the-an diffusion quality with first-order matching. Our method could extend this to higher-order moment matching, potentially accelerating distillation training. We leave
promising direction for future work.
Beyond image modelling, our method can be straightforwardly applied to accelerate large-scale video diffusion models (Blattmann al., 2023; Chen et al., 2023), which are based on latent diffusion models. A recent study (Zhao et al., 2024) shows that covariance estimation is crucial in mitigating the image-leakage issue in image-to- -o-idoo generation problems, presenting another promising application of our method. Moreover, our method can be applied to solve inverse problems, where the optimal covariance can enhance the accuracy of posterior sampling (Chung et al., 2022; Rozet et al., 2024). While this work focuses on the DDPM and DDIM sampler, incorporating the improved covariance into the general stochastic and deterministic differential solver (Song et al., 2021; Karras al., 2022) also represents exciting directions for future research.
CONCLUSION
In this paper, proposed method for learning the diagonal covariance of the denoising distribution, which offers improved accuracy compared to other covariance lcarning approaches. Our results demonstrate that enhanced covariance estimation leads to better generation quality and diversity, all while reducing the number of required generation steps. We validated the scalability of our method across different problem domains and scales and discussed its connection to various acceleration techniques and highlighted several promising application areas for future exploration.
10

---

Published as a conference paper at ICLR 2025
8 ACKNOWLEDGMENTS
ZO is supported by the Lee Family Scholarship. MZ and DB acknowledge funding from AI Hub in Generative Models, under grant EP/Y028805/1 and funding from the Cisco Centre of Excellence. We want to thank Wenlin Chen for the useful discussions.
REFERENCES
Akhound- Akhound-Sadegh, T., Rector-Brooks, J., Bose, A. J., Mittal, S., Lemos, P., Liu, C.-H., Sendera, M.. Ravanbakhsh, S., Gidel, G., Bengio, Y., et al. Iterated denoising energy matching for sampling from boltzmann densities. arXiv preprint arXiv:2402.06. 06 221, 121,, 2024.
Bao, F., Li, C., Sun, J., Zhu, J., and Zhang, B. Estimating the optimal covariance with imperfect mean diffusion prohabilistic models. arXiv preprint arXiv:2206.07755 007020 2022a.
Bao, F., Li, C., Zhu, J., and Zhang, B. Analytic-dpm: analytic estimate of the optimal reverse variance in diffusion probabilistic models. arXiv preprint arXiv:2201.06503. 2022b.
Bekas, C., Kokiopoulou, E., and Saad, Y. An estimator for the diagonal of matrix. Applied numerical mathematics, 57(1 57(11-12)12144-2223 2007.
Berthelot, D., Autef, A., Lin, J., Yap, D. A., Zhai, S., Hu, S., Zheng, D., Talbott, W., and Gu, E. Tract: Denoising diffusion models with transitive closure time-distillation. arXiv preprint arXiv:2303.04244 04248, 2023.
Blattmann, A., Dockhorn, T., Kulal, S., Mendelevitch, D., Kilian, M., Lorenz, D., Levi, Y., English, Z., Voleti, V., Letts, A., et al. Stable video diffusion: Scaling latent video diffusion models to large datasets. arXiv preprint arXiv:2311.15127 2023.
Bortoli, V. D., Galashov, A., Guntupalli, S., Zhou, G., Murphy, K., Gretton, A., and Doucet, A. Distributional diffusion models with scoring rules, 2025. http: org abs/ 2502. 02483.
Chen, H., Xia, M. He, Y., Zhang, Y., Cun, X., Yang, S., Xing, J., Liu, Y., Chen, Q. Wang, X., et al. Videocrafterl: Open diffusion models for high-quality video generation. arXiv preprint arXiv:2310.19511 2023.
Chen. W., Zhang, M., Paige, B., Hermández-Lobato, J. M., and Barber, D. Diffusive gibbs sampling. arXiv preprint ar arXiv: 202.:0240,.03008, 2024.
Chung, H., Kim, J., Mccann, M. T., Klasky, M. L., and Ye, J. C. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209. 14009, 2022.
Dayan, P., Hinton, G. E., Neal, R. M., and Zemel, R. S. The helmholtz machine. Neural computation, 7(5):889-9904 904, 1995.
De Bortoli, V., Thortoon, J., Heng, J., and Doucet, A. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34: 17695-17709, 2021.
Efron, B. Tweedie's formula and selection bias. Journal of the American Statistical Association, 106 1496):1602-111,, 2011.
Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A. kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Heek, J., Hoogeboom, E., and Salimans, T. arXiv:2403.c 06600, 2024.
Multistep consistency models. arXiv preprint
11

---

Published as a conference paper at ICLR 2025
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Ho, J., Jain, A., and Abbeel, Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
Ho, J., Salimans, T., Gritsenko, A., Chan, W., Norouzi, M., and Fleet, D. Video diffusion models. Advances in Neural Information Processing Systems, 35:8633-8646, 2022.
Hoogeboom, E., Satorras, V. G., Vignac, C., and Welling, M. Equivariant diffusion for molecule generation in 3d. In International conference on machine learning, pp. 8867-8887. PMLR, 2022.
Hutchinson, M. F. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications NkatisticsSSmmllation and Computation, 19(2):33-450, 1990.
Jolicoeur-Martinenn A., Li, K., Piché-Taillefer R., Kachman, T., and Mitliagkas, Gotta go fast when generating data with score-based models. arXiv preprint arXiv:2105. 2021.
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion- generative models. Advances in Neural Information Processing Systems, 55:265655-66577 2022.
Kim, D., Lai, C.-H., Liao, W.-H., Murata, N., Takida, Y., Uesaka, T., He, Y., Mitsufuji. Y., and Ermon, S. Consistency trajectory models: Learning probability flow ode trajectory of diffusion. arXiv preprint ar arXiv:2310.022799 2023.
Kingma, D., Salimans, T., Poole, B., and Ho, Variational diffusion models. Advances in neural information processing systems, 34:21696-2119 2021.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. arXiv preprint arXiv: 2013.
Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images. 2009.
Li, I. and He, J. Bidirectional consistency models. arXiv preprint arXiv:2403. /8035, 2024.
Li, X., Thickstun, J., Gulrajani, I., Liang, P. S., and Hashimoto, T. B. Diffusion -m improves controllable text generation. Advances in Neural Information Processing Systems, 35:4328-4343, 2022.
Liu, H., Chen, Z., Yuan, Y., Mei, X., Liu, X., Mandic, D., Wang, W., and Plumbley, M. Audioldm: Text-to-audio generation with latent diffusion models. arXiv preprint (rXiv:2301.12503 2023.
Liu, L., Ren, Y., Lin, Z., and Zhao, Z. Pseudo numerical methods for diffusion models on manifolds. arXiv preprint arXiv:2202.09778, 2022.
Liu, Z., Luo, P., Wang, X., and Tang, X. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference on computer vision, pp. 3730-3738, 2015.
Loshchilov, I. Decoupled weight decay regularization. arXiv preprint arXiv:1711.051011 2017.
Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., and Zhu, Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022.
Luo, W., Hu, T.. Zhang, S.. Sun, J., Li, Z., and Zhang, Diff-instruct: A universal approach for transferring knowledge from pre-trained diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
Martens, J.. Sutskever, I., and Swersky, K. Estimating the hessian by back paek-propgggting curvature. arXiv preprint arXiv: 120. 6664,444, 2012.
12

---

Published as a conference paper at ICLR 2025
Meng, C., Song, Y., Li, W., and Ermon, S. Estimating high order gradients of the data distribution by denoising. Advances in Neural Information Processing Systems, 01125559-25366 2021.
Nichol, A. Q. and Dhariwal, P. Improved denoising diffusion probabilistic models. In International conference on machine learning, pp. 8162-8171. PMLR, 2021.
Peebles, W. and Xie, S. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 11195-2200 2023.
Poole, B., Jain, A., Barron, J. T., and Mildenhall, Dreamfusion: Text-to-3d using 2d diffusion. arXiv preprint arXiv: 102009944980 14988, 2022.
Robbins, H. E. An empirical bayes approach to statistics. In Breakthroughs Statisticx: Foundations and basic theory, pp. 388-394. Springer, 1992.
Rombach, R., Blattmann. A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10684-10695, 2022.
Rozet, F., Andry, G., Lanusse, F., and Louppe, G. Learning diffusion priors from observations by expectation maximization. arXiv preprint arXiv:2405.13712 2024.
Sajjadi, M. S., Bachem, O., Lucic, M., Bousquet, O., and Gelly, S. Assessing generative models via precision and recall. Advances in neural information processing systems, 31, 2018.
Salimans, T. and Ho, J. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512. 2022.
Salimans, T., Mensink, T., Heek, J., and Hoogeboom, E. Multistep distillation of diffusion models via moment matching. arXiv preprint arXiv:2406.04103, 2024.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pp- 2256-2265. PMLR, 2015.
Song, J., Meng, C., and Ermon, Denoising diffusion implicit models. arXiv:2010.02502, 2020.
arXiv preprint
Song. Y. and Ermon, Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based genera- tive modeling through stochastic differential equations. In International Conference on Learning Representations, 2021. URL https: : oppprevieeeeeeeeeeeeeee net/ forum met:/oorumiddpPxTIi.... 11d=xXIGG2RRRSS
Song. Y., Dhariwal, P., Chen, M., and Sutskever, I. arXiv:2303.01466, 2023.
Consistency models. arXiv preprint
Townsend, J., Bird, T., and Barber, D. Practical lossless compression with latent variables using bits back coding. arXiv preprint arXiv:1901 048666, 2019.
Vahdat, A., Kreis, K., and Kautz, Score-based generative modeling in latent space. Advances in neural information processing systems, 34:1 1137- 2021.
Vargas, F. Grathwohl, W., and Doucet, A. Denoising diffusion samplers. arXiv preprint arXiv:2302. 13834, 2023.
Vincent, P. connection between score matching and denoising autoencoders. Neural computation, 23(7):1661-177, 2011. 2011.
Wang, G., Jiao, Y. Xu, Q., Wang, Y., and Yang, C. Deep generative learning via schrödinger bridge. In International conference machine learning, pp. 10794-10804, PMLR, 2021.
13

---

Published as a conference paper at ICLR 2025
Wang, Z., Lu, C., Wang, Y., Bao, F., Li, C., Su, H., and Zhu, J. Prolificdreamer: Highhfidelity and diverse text-to-3d generation with variational score distillation, Advances in Neural Information Processing Systems, 36, 2024.
Xiao, Z., Kreis, K., and Vahdat, A. Tackling the generative learning trilemma with denoising diffusion gans. arXiv preprint arXiv:21 2021.
Xie, S., Xiao, Z., Kingma, D. P., Hou, T., Wu, N., Murphy, K. P., Salimans, T., Poole, B., and Gao,
R. Em distillation for one-step diffusion models. arXiv preprint arXiv:2405. 2024.
Xu, M., Geffner, T., Kreis, K., Nic, W., Xu, Y., Leskovee, J., Ermon, S., and Vahdat, A. Energy-based diffusion language models for text generation. arXiv preprint arXiv:2410.2/357 2024.
Yu, F, Seff, A., Zhang, Y., Song, S., Funkhouser, T., and Xiao, Lsun: Construction of large-scale image dataset using deep learning with humans in the loop. arXiv preprint SKKV::5500003666 2015.
Yu, L., Xie, T., Zhu, Y., Yang. T., Zhang, X., and Zhang, Hierarchical semi-implicit variational inference with application to diffusion model acceleration. Advances in Neural Information Processing Systems, 36, 2024.
Zhang, F., He, J., Midgley, L. I., Antorán, J., and Hemández--obato, J. M. Efficient and unbiased sampling of boltzmann distributions via consistency models. arXiv preprint arXhv:2409.07:32 073.23333. 2024a.
Zhang, M., Bird, T., Habib, R., Xu, T., and Barber, D. Variational f-divergence minimization. arXiv preprint 1907.1907.11999 2019.
Zhang, M., Hayes, P., Bird, T., Habib, R., and Barber, D. Spread divergence. International Conference on Machine Learning, pp. 11106-11116 PMLR, 2020.
Zhang, M., Zhang, A., and McDonagh, S. On the out-of-distribution generalization of probabilistic image modelling. Advances in Neural Information Processing Systems, 34:3811-3822 2021.
Zhang, M., Howkins-Hookee A., Paige, B., and Barber, D. Moment matching denoising gibbs sampling. Advances in Neural Information Processing Systems, 36, 2024b.
Zhang, M., He, J., Chen, W. C., Ou, Z., Hernández-Lobato, M., Schölkopf, B., and Barber, D.
Towards training one-step diffusion models without distillation. arXiv preprint arXiv:2502.08005. 2025.
Zhao, M., Zhu, H., Xiang, C., Zheng, K., Li, C., and Zhu, J. Identifying and solving conditional image leakage in image-to- video diffusion model. arXiv preprint 4rXiv:2406.15335 2024.
Zhou, M., Zheng, H., Wang, Z., Yin, M., and Huang, H. Score identity distillation: Exponentially fast distillation of pretrained diffusion models for one-step generation. In Forty-first International Conference on Machine Learning, 2024.
14

---

Published as a conference paper at ICLR 2025
Appendix for "Improved Diffusion Models with Optimal Covariance Matching"
CONTENTS
A Abstract Proof and Derivations
15
A.1 Derivations of the Analytical Full Covariance Identity
15
A.2 Validity of the OCM Objective
16
A.3 Connection SN-DDPM
16
Details Experiments
18
B.1 Details Model Architoctures
18
B.2 Details of Training, Inference, and Evaluation.
19
Additional Experimental Results
20
C.I Additional Toy Demonstration
20
C.2 Additional Likelihood Comparison
21
C.3 Impact of the Numbers of Rademacher
21 "a..............................
C.4 Mean and Variance of Performance
:r..............................
C.5 More Results on Latent Diffusion Models
22
C.6 Generated Samples
22
A ABSTRACT PROOF AND DERIVATIONS
A.1 DERIVATIONS OF THE ANALYTICAL FULL COVARIANCE IDENTITY
Theorem 1 (Generalized Analytical Covariance Identity). Given a joint distribution q(š, x)
( (((() with (((xx)) ((ax,, o2 [), then the covariance of the true posterior q(π(ā)
((r) (T)) which is defined as { (x( - Byg(aaa)( [x]] has a closed form:
[G) = (o V (()) + o² /a2.
(8)
Proof. This proof generalizes the original analytical covariance identity proof discussed in Zhang et al. (2024b). Using the fact that (zq()(r) = 2 (2 ar ax))g(2 (z/z) and the Tweedie's formula expand the hessian of the log go()):
Va log q())
/ (alE (aLEg(=1a((>] Z).
logg( = σ² /
())))))())) (())
V (8 oz)
dr
1 ())=))()) dr
σ² / (ã ax)
1
(=)())z))(()g(1) V= dæ
/ ( ox)
Va9(2lz)q(1) V log q ())(() dx
22 log q( + 1 =
(())
= =-/ (E axx) (ā (() + _ aB ( (1
dr
(())
- の ( + 0²1 001== / -
ax)² - (ā - az) (2 (B [x])) g(xl)) de
= = - o² E [x] a²x(T)
Therefore, we obtain the analytical full covariance identity: [T) = (o" + 021)/11 /a2.
15

---

Published as a conference paper at ICLR 2025
A.2 VALIDITY OF THE OCM OBJECTIVE
Theorem 2 (Validity of the OCM objective). The objective in Equation (11) upper bounded the base objective (i.e., Equation (10) with M oo). Moreover, it attains optimal when h diag H (T)) for all & ~ g(z).
Proof. Recall that to learn the optimal covariance, we can minimize the following loss function
min (() Ihe (i) Eo @ H
(24)
We call this the grounded objective because remains unbiased and consistent, ie.. ha- (Z)), when the inner expectation is approximated with infinite Monte Carlo samples. Using Jensen's 's inequality, we can show that the OCM objective defined (11) provides an upper bound for the grounded objective
BB( (+) Egg H())u)
By( (hs v H S (g))) (T) o H H
(o)
Thus, minimizing the OCM objective also minimizes the grounded objective, leading to a more accurate approximation of the diagonal Hessian. We then show that these two objectives are equivalent when attaining their optimal. To see this, we can expand the OCM objective
(Ø) Ep(r)a(=) (b) c (
Bx 2F. Ph (v H(S)o)] = (#)) , (np (3) diag(H())) (H)))] +
= (xaa)llhs - diag(H (1) +d.
where e, are constants w.r.t the parameter Φ, and line 2 to line follows from the fact that Ep(u] [v @ Hx)vo] diag(H(x)). Therefore, Laum( attains optimal when hz = diag H (T)) for all ā ~ (())
A.3 CONNECTION TO SN-DDPM
In this section, we showcase the connection between OCM-DDPM and SN-DDPM through the second-order Tweedie's formula. Before delving into it, we present the following lemmas, which are essential for our derivation.
Lemma 1 (First order Tweedie's s formula (Efron, 2011)). mean of the inverse density equals
3I), we have the
E
[x(
(5 + vV log (7)).
(25)
Va
Lemma 2 (Second order Tweedie's formula). Let ((]|z) N( Vaa, BI), we have the second moment of the inverse density equals
1 (5xT + (z) + 35s (8) + { + (3)) + BT). (26)
E [x(=)))<=="]
a
where 8 (T) V log
q()) and (z)
logg(ā).
Proof. The proof follows that in (Meng et al., 2021, Appendix B) with the generalization to the scaled Gaussian convolutions. Specifically, we first reparametrized q() as a exponential distribution
q( in) = " a-v(()) go (ū),
(27)
Va where 3 x. (ā) x e and ((n) denotes the partition function. By applying the Bayes rule (()/z) 9( mn)) we have the corresponding posterior
a(x)
(n] " 8 v(n) -(n) p()),
(28)
16

---

Published as a conference paper at ICLR 2025
where ((Ø) log q()) - log (ū)) Since [(n(])dn 1, by taking the derivative w.r.t. & on both sides, we have
/ (11 V=\(E))" ((1)3)= =0 0,
(29)
which implies that E[n/a] vz)) Taking the derivative w.r.t. 3 on both sides again, we have
/ n (m V>)(E)) ((lā) (2)
(30)
which implies that E/mm² (=) V By substituting "8 Vzi(T)
Jo
81 (z) + T. and V=) -
5 (5) + 11 we get the result as desired.
Lemma 3 (Convert the covariance of (((x)) to the hessian of q())). Let (etg()))))= = / (ax, 3I), we have the covariance of the inverse density equals
Cov = -
BV log (T)).
(31)
Proof. Let (ū) = V log (()) and $2 (t)
= V log ( We have
3²
i Voz
Cov (o
Cov ((alaa)
a
3
2 Vav
i Vax
i- Vax
T Vax
B2(((1))
-Bgcll))
(((((()
a
B
B
/8
1
By(e)z) (5 (Va) (o vax) 81 (1)))110)
a
3a
g2
(RET 2ū + B8n (T)) + aF (( (100 1) -1 (T))
a
Ba
82
(B² (2 + 82 81 (1) (i) + BT) - -1
a
32
3 (T + 5 (3))
1 82
=
a
where line to line 3 follows Lemma and line 4 to line 5 follows Lemma 2.
noteworthy that line 3 in the proof of Lemma 3 also showcases the connection between the covariance of ((((z) and the score of q(²):
3² 1 (
Cov Covg(u|a)
V log log q(ā)²)
(32)
where = (ā - / // /3. Now we can apply Equations (31) and (32) to establish the connection between OCM-DDPM and SN-DDPM, as demonstrated in the following theorem.
Theorem 3 (Connection between OCM-DDPM and SN-DDPM). Suppose (( is defined as Equa- tion (1), and the pre-trained score function is well-learned, i.e., Vz, log pe (a:)) Vz, log q(xe). Vry Let h&s 9y be parameterized neural networks. DCM-DDPM learns h& by minimizing the objective LoCM ( as in Equation (13), and SN-DDPM learns ge by minimizing the objective (sw(w)) as in Equation (22). Then in optimal training with o² argmin Locu(() and th" = argmin (sw(w), we (2.))) and
have the optimal diagonal covariance of OCM-DDPNN [ 6²)
that of SN-DDPM -x-1(16;(*)) 1-81-1 Br
(96- V.. logpe(z;)²) are identical. 82
de Toc
Proof. optimal training, we know that ha- (2x) diag(V; V2 log log)(x;)) and go- (xx)) Egtzo 12) [e]]. Bao et al. (2022b) show that the covariance of the denoising density ((01-1/z)) has closed form (see Lemma 13 in Bao et al. (2022b))
Covg(rr-112))
X2 + oo Cov (( Co [x)0])
(33)
17

---

Published as a conference paper at ICLR 2025
where X3 1-8c 1 Br and Yt = vāi-1 vI- ā X
á
vaa- a.
Since
T-a,
V
1-ū
I-aa
(xre)rr)) = -N(voro (1 - --))). applying Lemma 3 gives
(1 (1-aa) ār)
diag(Cov (( (zul) =
(I + (1
ā
Substituting it into Equation (33), we have
diag( glo) (x-1I) (1 (1- (v. ha- (2i) + (1 (1= ox))
Ex- (10e; x*)
az
as desired. Alternatively, applying Equation (32), we have
diag(Cov, o(o) (mo)) (1 - ā.)²
(2)(
-z, logpo(12)²)
āi
1 - αe
Substituting it into Equation (33) gives
diag(Cov,ga, .lz,) -lx (x)-1 1 αt
04-1 Be
8 9 (xx)
Va, (eppo = -x-1(10;)))
Therefore, diag(Cov,ts [[)]) = Ex=1(Z(; Ix-1(x1: as desired.
Remark. Theorem 3 establishes the connection between OCM-DDPM and SN-DDPM using the second-order Tweedie's 's formula. This connection allows OCM-DDPM to utilize the same covariance clipping trick as in SN-DDPM (see Appendix B.2 for details). Additionally, as highlighted in Bao al. (2022a), SN-DDPM suffers from error amplification in the quadratic term, limitation not present in
As shown in Figure 4 and Table 3, OCM-DDPH demonstrates superior covariance estimation accuracy and likelihood performance compared SN-DDPM, underscoring the advantages of the proposed optimal covariance matching objective.
B DETAILS OF EXPERIMENTS
B.1 DETAILS OF MODEL ARCHITECTURES
In this section, we describe our model architectures in detail. Following the approach in Bao et al. (2022a), our model comprises pretrained score neural network with fixed parameters, along with trainable diagonal Hessian prediction network built on top of it.
Details of Pretrained Score Prediction Networks. Table 6 lists the pretrained neural networks utilized in our experi- ments. These models parameterize the noise prediction (z(z), allowing to de- rive the score prediction as se (xx)
Table 6: Source of pretrained score prediction networks used in our experiments.
PROVIDED BY
CIFAR (L.))
BAO ET AL.. (2022B)
CIFARIO (CS)
BAO ET AL. (2022B)
<a(as) following the lo, = - VI-a."
CELEBA 64X64
SONG ET AL. (2020)
IMAGENET 64X64
NICHOL & DHARIWAL (2021)
forward process defined in Equation (1). It is important to note that the pretrained net-
LSUN BEDROOM
Ho AL.. (2020)
IMAGENET 256x256
PEEBLES & XIE (2023)
works for ImageNet 64x64 and 256x256
include both noise prediction networks and covariance networks. In our model architectures, we
utilize only the noise prediction networks.
Details of Diagonal Hessian Prediction
Table 7: Architecture details of our models.
Networks. For fair comparisons, we fol-
NN1
NN2
low the parameterization as per Bao et al.
(2022a) for all models excluding the one on ImageNet 256x256, which was not ex- plored in their paper. The architecture de- tails of NN, and NN, provided in Ta- ble 7, where Conv denotes the convolu- tional layer, Res denotes the residual block
CIFAR 10 (L.S)
Conv
Conv
CIFAR10 (CS)
Conv
Conv
CELEBA 64X64
Conv
Conv
IMAGENET 64X64
Conv
Res+Conv
LSUN BEDROOM IMAGENET 256x256
Conv
Res+Conv
AdaLN
Adnnar
18

---

Published as a conference paper at ICLR 2025
Table 8: The number of parameters and the averaged time (ms) run model function evaluation. All are evaluated with a batch size of 64 on one A100-80GB GPU.
SCORE PREDICTION NETWORK
SCORE & SN PREDICTION NETWORKS
SCORE & DIAGONAL HESSIAN
PREDICTION NETWORKS
CI (LS)
52.54 M/ 44.37 MS
52.55 M / 44.61 MS (+0.5%)
52.55 M / 44.51 MS (+0.3%)
CIFARIO (CS)
52.54 M 45. 13 MS
52.55 M 45. 24 MS (+0.2%)
52.55 M /45. 23 MS (+0.2%)
CELEBA 64x64
78.70 M 67 88 MS
78.71 M 69. 15 MS (+1.9%)
78.71 M 68, 68.89 MS (+1.5%)
IMAGENET 64x64
121,06 106.58 MS
121.49 M 112.53 MS (+5.6%)
121.49 M 112.23 MS (+5.3%)
LSUN BEDROOM
113.67 M 692.58 vss 114.04 M / 771 1.55 MS (+11.4%)
114.04 / 770.73 MS (+11.2%)
IMAGENET 256x256
675.13
M/ 22.84 MS
MISSING MODEL
677.80 M 23.01 MS (+0.7%)
(He et al., 2016), AdaLN denotes the adaptive layer norm block (Peebles & Xie, 2023), and Linear denotes the linear layer.
Cost of Memory and Inference Time. In Table 8, we present the number of parameters and inference time for various models. It is evident that the additional memory cost of the diagonal Hessian prediction network is minimal compared to the original diffusion models, which only include the score prediction. Regarding the extra inference cost, it is negligible for CIFAR1O 100, CelebA 64x64, and ImageNet 256x256. The additional time is at most 5.3% on ImageNet 64x64, and 11.2% on LSUN Bedroom, but this is offset by the benefit of requiring fewer sampling steps to achieve an FID around 6, as shown in Table 4. Notably, although the model size on ImageNet 256x256 is the largest, has the shortest inference time because DiT learns the diffusion model in the latent space.
Details of Table 4. In Table 4, the FID results of baselines are taken from Bao et al. (2022a). Notably, the time cost for a single neural function evaluation is identical across DDPM, DDIM, and Analytic-DPM, all of which use fixed isotropic covariance in the backward Markov process. The ratio of the time cost to the baselines is based on Table 8. For the FID results of our model, OCM-DPM, report performance using the DDPM forward process for LSUN Bedroom and the DDIM forward process for the other datasets.
B.2 DETAILS OF TRAINING, INFERENCE, AND EVALUATION
Our training and inference recipes largely follow those outlined in Bao al. (2022a); Peebles & Xie (2023). Below, we provide a detailed description.
Training Details. We use the AdamW optimizer (Loshchilov; 2017) with learning rate of 0.0001 and train for 500K iterations across all datasets. The batch sizes are set to 64 for LSUN Bedroom, 128 for CIFARIO. CelebA 64x64, and ImageNet 64x64, and 256 for ImageNet 256x256. During training. checkpoints are saved every 10K iterations, and we select the checkpoint with the best FID on 2048 samples generated with full sampling steps". We train our models using one A A100-80G GPU for CIFARIO, CelebA 64x64, and ImageNet 64x64; four A100-80G GPUs for LSUN Bedroom; and eight A100-80G GPUs for ImageNet 256x256.
Sampling Details. As per Bao et al. (2022b), covariance clipping is crucial to the performance in diffusion models with unfixed variance in the backward Markovian. Leveraging the connection between OCM-DPM and SN-DPM established in Theorem 3, we can apply the same clipping strategies outlined in Bao et al. (2022::b). Specifically, we only display the mean of p(zg (1)) at the last sampling and clip the covariance Σ₁ (I2) of p(æ1 (22) such that ||E, (z2) (lo Elc| 255 3. where Il- denotes the infinity norm and e is the standard Gaussian noise. Following Bao et al. (2022b), we use y = 2 on CIFARIO 10 (LS) and CelebA 64x64 under the DDPM forward process, and use y 1 for other cases. For all skip-step sampling methods, we employ the even trajectory (Nichol & Dhariwal, 2021) for selecting the subset of sampling steps.
2
Evaluation Details. The performance is evaluated on the exponential moving average (EMA) model with a rate of 0.9999. For computing the negative likeliiel,hood, we follow Ho et al. (2020); Bao et al. (2022a) by discretizing the last sampling step p(zo|=1) to obtain the likelihood of discrete image data and report the upper bound of the negative log- likelihood on the entire test set. The FID score is computed on 50K generated samples. Following Nichol & Dhariwal (2021); Bao et al.
We use 2048 samples instead of 1000 to ensure that the covariance of the Inception features has full rank.
19

---

Published as a conference paper at ICLR 2025
Table 9: The corresponding mean and standard deviation of NLL and FID reported in Tables and 5. Note that the standard deviation is reported under the scale of percentage (%)
CIFARIO (LS)
CIFARIO (CS)
# TIMESTEPS K
10
25
50
100
200
1000
10
25
50
100
200
1000
MEAN-DDPM (NLL)
5.33
4.63
4.35
3.97
3.78
3.57
4.99
4.34
3.99
3.76
3.59
3.41
STD-DDPM (NLL) %
0.04
0.03
0.03
0.04
0.04
0.01
0.02
0.01
0.03
0.03
0.02
0.01
MEAN-DDPM (FID)
25.07
9.30
5.90
4.37
3.54
4.00
14.46
5.49
4.09
3.83
3.81
4.22
STD-DDPM (FID)
9.05
7.54
3.48
2.77
4.25
2.30
10.40
7.74
2.66
1.32
5.26
3.08
MEAN-DDIM (FID)
10.61
4.31
3.48
3.26
3.25
3.75
16.85
6.70
4.73
4.33
4.56
4.59
STD-DDIM (FID)
5.17
3.55
3.19
1.87
3.97
0.99
13.22
4.23
2.28
2.94
1.61
5.30
CELEBA 64x64
IMAGENET 64x64
# TIMESTEPS K
10
25
50
100
200
1000
25
50
100
200
400
4000
MEAN-DDPM (NLL)
4.69
3.86
3.43 3.13 2.90
2.66
4.45
4.15
3.93
3.79
3.70
3.59
STD-DDPM (NLL) %
0.01
0.01
0.01
0.00
0.01
0.00
0.01
0.00
0.01
0.00
0.01
0.01
MEAN-DDPM (FID)
21.58
12.65
9.24
6.96
6.00
5.02
28.01
20.85
18.01
16.76
16.35
16.33
STD-DDPM (FID)
2.24
7.65
2.29 2.32 5.50
1.37
0.77
4.66
6.82
3.28
3.04
3.21
MEAN-DDIM (FID)
10.36
5.69
4.37
3.50
3.11
3.03
28.34
19.68
17.84
17.41
17.36
18.91
STD-DDIM (FID) %
7.40
4.04
5.12
4.88
6.46
1.55
6.59
3.77
9.15
4.08
7.88
5.26
DOWN, CON DANG
COPN,
DOPN,
DOPN, CON TULL
DEWK SOV FULL
103PM
DENN.
DCW
18
A-COMI
100PM
A-20TH
OCN-DEIN (Oun))
SN CCPN
OCH-COFN (OUn)
(a) Cov Error with True Score
(b) Cov Frror with L Learned Score
(c) DDPM MMD V.s. Steps
(d) DDIM MMD v.s. Steps
Figure 4: Comparisons of different covariance estimation methods based on estimation error and sample generation quality. Figures (a) and (b) show the mean square error of the estimated diagonal covariance under the assumptions: (a) access to the true score, and (b) learned score of the data distribution at various noise levels. Figures (c) and (d) present the MMD evaluation against the total sampling steps in the DDPM (c) and DDIM (d) settings. We can find the proposed OCM method can achieve the lowest estimation error and consistently outperform other baseline methods when fewer generation steps are applied.
(2022a); Peebles & Xie (2023), the reference distribution statistics for FID are calculated using the full training set for CIFARIO and ImageNet, and 50K training samples for CelebA and LSUN Bedroom. Performance results are reported using the same random seed as in Bao et al. (2022a). Additionally, the variance across different random seeds is provided in Appendix C.4.
C ADDITIONAL EXPERIMENTAL RESULTS
In this section, present additional experimental results for comparison. First, we present an additional toy experiment in Appendix C.1 and likelihood comparison between I-DDPM, NPR- DDPM, and OCM-DDPM in Appendix C.2 and present results using different numbers of Rademacher samples in Appendix C.3. We also analyze the performance variance in Appendix C.4. Additionally, we include further experimental results on latent diffusion models in Appendix C.5 and conduct qualitative studies by showcasing samples generated by our models in Appendix C.6.
C.1 ADDITIONAL TOY DEMONSTRATION
To further verify the effectiveness of our method, we include an additional toy example. In this case, consider another two-dimensional mixture of nine Gaussians (MoG) with means located at
'e use the official code from https://github. com homm/baa==///am /eep=//eeexteededdedoo Ana Tytic-
20

---

Published as a conference paper at ICLR 2025
(-3,0,3) @ (-3,0,3) and a standard deviation of a 0.1. To assess different approaches, first learn the covariance using the true scores. Specifically, we train the covariance networks for these methods over 50,000 iterations with a learning rate of 0.001 using an Adam optimizer (Kingma & Ba, 2014). Figure 4a shows the L2 error of the estimated diagonal covariance [ and our method, OCM-DDPM, consistently achieves the lowest error compared the other methods. In practice, the true score is not accessible. Therefore, we also conduct a comparison using a score function learned with DSM. We then apply different covariance estimation methods under the same settings. In Figure 4b, we plot the same L2 error for the learned score setting and find that our method achieves the lowest error for all values.
We further use the learned score and the covariance estimated by different methods to conduct DDPM and DDIM sampling for this MoG problem. Figures 4c and 4d presents the MMD comparison across various methods. It shows that our method outperforms the baselines when the total time step is small, demonstrating the importance of accurate covariance estimation the context of diffusion acceleration. We also include two methods utilizing the true diagonal and full covariance as benchmarks, representing the best achievable performance. Notably, these two methods exhibit similar performance because the MoG used in this setting has symmetric components, in which the covariance is dominated by the diagonal entries.
C.2 ADDITIONAL LIKELIHOOD COMPARISON
Here, we provide additional likelihood comparisons against I-DDPM (Nichol & Dhariwal, 2021) on the ImageNet 64x64 dataset. As discussed in Section 4, I- DDPM parameterizes the diagonal covari- ance by interpolating between 3 and and learns the covariance by maximizing the variational lower bound. For ease of com- parison, we also include the performance of NPR-DDPM and SN-DDOM, which em- ploy alternative MSE loss as detailed in Equation (22) to learn the covariance. Table 10 presents the results, showing that OCM-DDPM achieves the best likelihood with fewer sampling steps while maintaining performance comparable to I-DDPM when using full sampling steps. This highlights the superiority of the proposed optimal covariance matching objective in learning diffusion model for data density estimation.
Table 10: Upper bound on the negative log-likelihood (bits/dim) on the ImageNet 64x64 dataset.
# TIMESTEPS K
25
50
100
200
400
4000
I-DDPM NPR-DDPM
18.91
8.46
5.27
4.24
3.86
3.57
4.66
4.22
3.96
3.80
3.71
3.60
SN-DDPM OCM-DDPM
4.56
4.18
3.95
3.80
3.71
3.63
4.45
4.15
3.93
3.79
3.70
3.59
C.3 IMPACT OF THE NUMBERS OF RADEMACHER SAMPLES
In this section, include an ablation study examining the effect of varying the number of Rademacher samples M. Specifically, we conduct experiments on CIFAR10 10 (CS) using the proposed OCM-DDPM method and evaluate performance based on FID and NLL. The empirical results, presented in Table 11, indicate that larger M (e.g., M = 3) can give a small improvement in FID. This improvement is likely due to the reduced gradient estimation variance during training with larger M. However, in most cases, different values of M yield consistent performance. This is practically desirable as, in practice, setting M 1 allows for efficient training while maintaining strong performance
Table 11: Results OCM-DDPM on CIFARIO (CS) with varying numbers of Rademacher Samples.
FID 4.
NLL (%) f
K
10
25
50
100
10
25
50
100
M 1
14.32
5.54
4.10
3.84
4.99
4.34
3.99
3.76
14.18
5.51
4.11
3.82
4.99
4.34
3.99
3.76
M = 5 14.17
5,51
4.11
3.82
4.99
4.34
3.99
3.76
M = 10 14.16
5.51
4.11
3.82
4.99
4.34
3.99
3.76
C.4 MEAN AND VARIANCE OF PERFORMANCE
In Tables 3 and 5, we evaluate models using the same random seed as Bao al. (2022a). To minimize the impact of randomness, we report the mean and standard deviation in Table by repeating the evaluation three times with different seeds.
21

---

Published as a conference paper at ICLR 2025
Table 12: Results with varying CFG coefficients using 10 sampling steps on ImageNet 256x256.
FID !
RECALL (%)
CFG
1.5
1.75
2.0
3.0
4.0
1.5
1.75
2.0
3.0
4.0
DDPM, 9
30.41
19.26
13.53
11.52
14.76
39.74
35.29
32.08
24.93
19.91
DDPM, 8
210.28
182.78
158.16
89.86
58.16
16.98
22.28
25.13
23.07
19.48
I-DDPM
44.96
20.71
13.01
9.76
13.75
50.32
43.75
41.08
31.77
25.49
DDM-
30.55
17.57
11.12
9.52
13.74
49.39
45.87
42.33
32.13
24.96
DDIM
9.41
6.54
6.12
10.49
14.18
49.15
45.82
41.15
29.27
23.16
OCM-DDIM
11.30
6.67
5.33
9.20
13.49
54.50
50.32
46.58
34.53
25.67
Table 13: Results with CFG=1.5 across different sampling steps ImageNet 256x256.
FID 1
RECALL (%) f
TIMESTEPS K
50
100
200
25
100
200
250
DDPM,
30.41
4.79
3.55
3.09
3.05
2.97
39.74
49.02
51.79
53.54
54.21
54.34
DDPM,
210.28
42.09
9.43
3.63
2.91
2.75
16.98
46.76
51.67
55.01
55.50
55.21
I-DDPM
44.96
9.01
3.70
2.48
2.25
2.74
50.32
54.78
56.45
57.88
58.76
54.69
OCM-DDPM
30.55
4.96
2.71
2.50
2.75
49.39
53.75
54.68
55.85
54.97
54.75
DDIM
9.41
2.70
2.33
2.25
2.23
2.20
49.15
56.41
57.00
57.83
57.81
$7.99
OCM-DDIM
11.30
3.26
2.56
2.28
2.23
2.18
54.50
58.33
58.09
58.80
58.16
58.56
C.5 MORE RESULTS ON LATENT DIFFUSION MODELS
We report the performance using 10 sampling steps with varying CFG coefficients in Table 12. The results indicate that our methods perform best at CFG = 2.0. While DDPM-3 and DDIM show strong FID scores, their Recall is lower due fixed variance, suggesting less diversity in the generated samples. In Table 13, we further compare the performance across different sampling steps at CFG 1.5. The results again demonstrate that our methods, which estimate the optimal covariance from the data, produce more diverse samples while maintaining comparable image quality.
C.6 GENERATED SAMPLES
In this section, we conduct qualitative studies by showcasing the generated samples from our models using different sampling steps K. The results summarized follows:
In Figure 6, we visualize the generated samples using varying numbers of Monte Carlo samples with the Rademacher estimator (refer to Table 1).
In Figure 7, we visualize the training data and the generated samples using the minimum number of sampling steps required to achieve an FID of approximately 6 (refer to Table 4).
In Figures to 11, we visualize the generated samples of our models using different number of sampling steps on CIFARIO (LS), CIFARIO (CS), CelebA 64x64, and ImageNet 64x64, respectively (refer to Table 5).
In Figure 12, we visualize the generated samples on ImageNet 256x256, using 10 sampling steps with different CFG coefficients. (refer Figure and Table 12).
In Figure 13, we visualize the generated samples on ImageNet 256x256, using varying number of sampling steps with CFG set to 1.5 (refer to Table 13).
22

---

Published as a conference paper at ICLR 2025
Daia
M
100
M 1000
Figure 5: Diagonal covariance estimation visualisation with different Rademacher sample numbers.
(a) OCM-DDPM (M=
(b) OCM. DDPM (M = 10)
(c) OCM DDPM (M 15)
(d) OCM DDPM (M 20)
Figure 6: Generated samples with 10 sampling steps using ifferent Rademacher sample numbers on
CIFARIO (CS) (ref: Table 1).
(a) CIFARIO Data
(b) ColebA 64x64 Data
(c) LSUN Bedroom Data
(d) ImageNet 256x256 Data
(c) CI 16)
(f) CalehA 64x64 (K 17)
(g) L.SUN Bedroom (K = 90) (h) ImageNet 256x256 (K 10)
Figure 7: The training data and generated samples of OCM-DPM with minimum steps to achieve an FID around 6. (ref: Table 4)
23

---

Published as a conference paper at ICLR 2025
(a) OCM-DDPM (K 10)
(b) OCM-DDPM (K 25)
(c) OCM-DDPM (K 50)
(d) OCM-DDPM (K
100)
(c) OCM-DDIM (K= = 10)
(f) OCM-DDIM (K 25)
(g) OCM-DDIM (K = 50)
(h) OCM-DDIM (K = 100)
Figure 8: Generated samples with different sampling steps CIFARIO (LS) (ref: Table 5).
(a) OCM-DDPM (K 10)
(b) OCM- (K 25)
(c) OCM-DDPM (K 50)
(d) OCM-DDPM (K 100)
(e) UCM-DDIM (K 10)
(i) OCM-DDIM (K 25)
(K = 50)
(h) OCM-DDIM (K = 100)
Figure 9: Generated samples with different sampling steps CIFAR10 (CS) (ref Table 5).
24

---

Published as a conference paper at ICLR 2025
(a) OCM-DDPM (K 10)
(b) OCM-DDPM (K 25)
(c) OCM-DDPM (K 50)
(d) OCM-DDPM (K
100)
(c) OCM-DDIM (K= = 10)
(f) OCM-DDIM (K 25)
(g) OCM-DDIM (K = 50) (h) OCM-DDIM (K = 100)
Figure 10: Generated samples with different sampling steps CelebA 64x64 (ref: Table 5).
(a) OCM-DDPM (K 10)
(b) OCM-DDPM (K 25)
(c) OCM-DDPM (K 50)
(d) OCM-DDPM (K 100)
(e) OCM-DDIM (K = 10)
(I) OCM-DDIM (K- 25)
(g) OCM-DDIM (K = 50)
(h) OCM- (K = 100)
Figure 11: Generated samples with different sampling steps ImageNet 64x64 (ref: Table 5).
25

---

Published as a conference paper at ICLR 2025
OCM-DDPM (CFG- 1.5) (b) OCM-DDPM (CFG 1.75)
(c) OCM-DDPM (CFG
(d) OCM-DDPM (CFG 3.0)
(c) OCM-DDIM (c5)=1.5) (1) OCM-DDIM (CFG= 1.75) (g) OCM-DDIM (CFG=
(h) OCM-DDIM (CFG=330)
Figure 12: Generated samples with 10 sampling steps using different CFG coefficients on ImageNet
256x256 (ref: Table 12)
OCM OCM-DDPM (K 25)
(b) OCM DDMM (K 50)
(c) OCM- (K 100)
(d) OCM-DDPM (K 200)
(c) OCM-DDIM (K- 25)
(f) OCM-DDIM (K 50)
(g) OCM-DDIM (K 100) (h) OCM-DDIM (K 200)
Figure 13: Generated samples with CFG=1.5 using different sampling steps on ImageNet 256x256 (ref: Table 13)
26


Reviews:
[{"review_id": "uFezPOM0Um", "reviewer": "Reviewer_mSjs", "summary": "The paper propose to build on existing moment matching diffusion progress and focus on the estimation of the optimal covariance. They propose to improve over previous methods for estimating the diagonal of covariance with an additional network that amortizes the cost of computing explicitly the covariance during sampling. They show that their method is challenging with selected comparisons.", "questions": "See Weaknesses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Thank you for your work!\n\n- The paper demonstrates the application of OCM across different models (classical and latent diffusion), making the approach adaptable for a range of diffusion-based applications.\n\n- The paper target efficient diagonal estimation to reach efficient sampling (leading to reduced number of step and more accurate samples).\n\n- The proposed method is timely as it is based on recent covariance matching methods. It correctly try to improve efficiency over existing estimation techniques by amortizing the covariance estimation with an additional neural network. It still bases its estimation on learned diffusion model, making it comparable to previous ideas.\n\n- It looks like very reasonable overhead for the Hutchinson estimator (M = 1 as noted in line 199) provide good performance which is in favor of the method that seeks efficiency.\n\n- They propose an unbiased objective for diagonal covariance matching.", "weaknesses": "In general, the paper is good. However the results miss visualization (in the core text). A toy example would be welcome for visualization as well and easy results interpretation/comparison. The performances reported do not exhibit high improvement over other methods. Although it shows to be challenging! Maybe other performances should be reported to strengthen the defense of this method over others (the authors often point efficiency as a main goal, but no convincing observations in this favor\n\n**General Comments** \n\n- The title is a bit misleading. You do not make optimal covariance matching, only diagonal matching. I also found the choice of diagonal only matching a little bit skipped. I do not say that it's a bad approximation, but it has its limitations. You should motivate more 'why' diagonal is already a good approximation (more than the complexity motivation).\n\n- Figure 1 is not really relevant as so. Ok the number of Rademacher samples improves the diagonal structure, but we don't know the target. Could you estimate it for comparison? Or doing the same for a toy problem with known covariance? Can you quantify or give an idea of information loss induced by only estimating the diagonal (It depends on the problem of course)? Figure 1 do not provide more information than Appendix C4. Maybe the latter can be replaced by something that focus on your method. The number of Rademacher samples do not seems to be the main point of your method.\n\n- Concerning Rademacher samples, it would be nice to see a metric (like FID and IS) evolving with M (the number of samples) to strengthen your claim line 199.\n\n- Echoing with the previous point, it would be nice to show more extensive comparison with other covariance matching methods. Either that target full covariance or more costly diagonal estimation. Insist more on the advantage of your estimator concerning the estimation speed/accuracy ratio. In the same trend as table 4. Maybe you can have a small toy problem (MoG, ...) for which it is better to interpret results (and visualize them) + having access to GT covariance.\n\n- I'm a little bit worried about the 'diagonal' only estimation. What if you have strong local covariance? Could be good to discuss, give hints about potential failure of your method in such setting. I take the example of [1] where they have to define a new sampling strategy as existing ones do not capture correctly covariance (or assumptions are too strong) leading to poor sample accuracy. This method also provides moment matching sampling, efficient by using vjp in spite of modelling the full covariance explicitly. Could you compare to similar method? For example, how do you compare to low-rank estimation which are also less demanding but can model non-diagonal elements?\n\n- In the same idea as [1], current diffusion model and sampling strategies hardly comes without posterior sampling. It would be nice to study the benefit of your estimation for posterior problem. I think that is needed for a sampling method. \n\n- In equation (10), you still need $H(\\tilde{x})$ which can be large as you mention. Do you model it explicitly for training the covariance network or do you model it implicitly using jvp/vjp benefits? This should be discussed. How does it slow down training compared to classical denoising-score matching or other methods?\n\n- Can the introduction of an additional network (meaning additional approximation) lead to cumulative errors? This is not discussed at all, but the network not being well aligned with the covariance of the trained diffusion will surely lead to poor performances even if the diffusion model was good. \n\n**Spotted typos**\n\n- Figure 1: The $K$ corresponds to $M$ in the main text, no? Could be changed for consistency.\n\n- Line 195 and 812: ~Validity of OCM objetcive~ -> Validity of OCM objective.\n\n- Line 309 and 427 ~COM-DDPM~ and ~OCM-DPM~ -> OCM-DDPM, no? \n\n- Appendix A.2: ~Vadality~ -> Validity.\n\nI will be pleased to increase my score once my questions and concerns are answered, discussed and addressed.\n\n**[1] Learning Diffusion Priors from Observations by Expectation Maximization - 2024**\nI took this one as example, you can recursively check references which are relevant for this paper also.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jO8EJQ9gSe", "reviewer": "Reviewer_tC1T", "summary": "The paper proposes an approach to diffusion modelling in which the covariance matrix of the time-reversal of the \"noising\" process is set to a diagonal approximation of the correct form estimated from the approximation of the \"score\". \n\nIn order to do this a pair of neural networks are trained; one to estimate the score as usual and a second to estimate this diagonal matrix. The majority of parameters are shared between the two networks, allowing this to be accomplished with only a small additional computational cost. The properties of the objective function used to learn this covariance are explored to an extent in theoretical work.\n\nNumerical results are presented for a toy model and for some more substantial image generation scenarios.", "questions": "Line 36: in what sense is this a \"posterior\" covariance? It seems from later in the paper that the covariance of the semigroup of the time reversal of the noising process is what's intended; I don't see that there's anything Bayesian going on.\n\nLine 199: \"In practice, we found M = 1 works well,\" raises some obvious questions:\n1. How did you find this?\n2. What does \"works well\" mean and how does increasing M to larger numbers change behaviour? A single sample provides an unbiased estimate, but its variance will clearly be large; generally one would expect performance to increase with larger M and so the question is how does the performance/cost tradeoff behave and this offhand comment doesn't give the reader much information which is a shame if you've done a proper study of this.\n\nHave you done any comparison with the full covariance estimation approach of Zhang et al. (2024b)? While I appreciate it may sometimes be prohibitively costly it would have been nice to see what price is paid for making the diagonal approximation in at least as toy scenario. It seems as though strong correlation between coordinates in the true distribution might be reflected by strong correlation in the time reversal of the noising process, for example, and in those settings approximating it with a diagonal covariance although quite possibly still better than a naive choice might suffer relative to a more flexible approach whereas in settings in which there is not such strong dependency modelling it in the covariance matrix might not be particularly important.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "To the best of my knowledge this approach to approximating the covariance matrix is novel and seems well motivated.\n\nThe numerical results do appear to show that this approach outperforms a number of alternative strategies and the settings of the image generation setting do seem reasonably challenging.\n\nIt seems that the additional implementation and training requirements of this approach over and above those of a more basic diffusion generative model are modest and so this approach could be incorporated into implementations of such approaches at small cost (both human and computational) which means that even modest improvements are potentially valuable.", "weaknesses": "The paper doesn't seem to propose anything radically new; the introduction of sensible approximations of the covariance of the time reversal has been explored quite widely in recent years and so the contribution of this paper is the particular approach to doing that which, while well motivated and seemingly having good properties, does not seem particularly radical. \n\nThere are some odd choices made presentationally at times which to my mind make the paper a bit less readable and hence a bit less accessible. For example, the description of diffusion processes in Section 2.1 takes as a starting point a discrete time process; (1) is not a diffusion process although it is described as one. It might be arrived at as the time discretisation of a diffusion, which may in this case be exact, but that's not the same thing.\nThese are minor points and haven't substantially influenced my view of the paper, but this seems the best box to put this in. The paper does contain quite a number of typos, strangely phrased sentences and similar; careful proofreading would make it more readable. E.g.\n* Abstract: \"diagonal covariances\" is an odd phrase; covariances are the off-diagonal elements of the matrix. \"diagonal covariance matrix\" or similar might be clearer.\n* Theorem 2: \"objetcive\"\n* Theorem 2: \"The objective in Equation (11) upper bounded the grounded objective.\" Didn't immediately seem to make sense to me; I /thought/ the intention was probably something like \"The objective in Equation (11) upper bounds the base objective.\" or something of that sort, but it's not completely clear what the intended base objective was. It turns out that aside from the odd tense and phrasing, the statement of this theorem is using terms which are defined in the supplementary material!\n\nThe main evidence presented in support of the proposed approach is empirical and this evidence isn't completely overwhelming. E.g. from Table 3 it' seems that this method doesn't actually dominate other approaches in the literature although it often does well and doesn't seem to be outperformed by much when it is beaten. Of course, methods can be interesting without beating the state of the art uniformly, and there is the interesting suggestion that when other methods are outperforming this approach in terms of NLL they do less well in FID terms. That, of course, raises the question \"why is this and what does it actually tell us?\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LGIRJ9gCCT", "reviewer": "Reviewer_hfvy", "summary": "The authors introduce an optimal covariance matching objective based on score functions from pertained diffusion models to learn the covariance of the Gaussian transition kernel. \nThe proposed method is superior at covariance estimation in terms of  less storage consumption and less number of network evaluations. \nEmpirical results indicate that the proposed objective enhances performance in both sampling quality (NLL, FID) and efficiency (number of sample steps).", "questions": "- Is the proposed OCM able to help with other diffusion samplers? e.g. EDM [1] stochastic and deterministic samplers? \n- Any insights or intuitions why comparing to baselines, learning the covariance via OCM objective gives better covariance estimation over all noise levels? \n- Just out of curiosity, in Fig 3 (a) and (b), why the covariance error goes up for all methods when t is roughly in range [0, 400]? \n\n[1] Karras, Tero, et al. \"Elucidating the design space of diffusion-based generative models.\" Advances in neural information processing systems 35 (2022): 26565-26577.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well-organized, results are clearly presented, and the technique derivations are solid. \n- Overall paper organization is very smooth, the methodology part is easy to follow, and experiments are well-designed and coupled with the proposed methodology. e.g. from figure3, table 3 and table 5, I can see that the improved sample quality can be directly attributed from better covariance estimation. \n- Unlike the previous works which are of high requirements on storage and network evaluation in covariance estimate (especially the Hessian term), the proposed OCM barely needs 1 network evaluation for each sampling step, so that the method improves number of sampling steps without sacrificing too much sampling speed on covariance estimate.", "weaknesses": "- The authors show that the proposed OCM improves covariance estimation for DDPM and DDIM, which results in improved sample quality and better efficiency. There are many other stochastic and deterministic samplers with different sampling schedules, does the improved covariance helps with all these samplers?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ihwa1QSjJD", "reviewer": "Reviewer_BQbT", "summary": "The authors propose an $x_t$-dependent diagonal covariance estimation method that directly learns from score function. This leads to improved efficiency and more stable/robust denoising, which in turn improves the image generation quality.", "questions": "I have no more questions. I like the presentation (as well as the soundness of the idea, of course) and think this alone makes the paper okay to accept. I am just simply not impressed by the results.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. Very clear and comprehensive presentation, easy to read and follow.\n2. The proposed method is a general method that applies to various settings.\n3. The description of experimental details is generally comprehensive.", "weaknesses": "1. The idea seems both incremental and way-too simple, which is not a fundamental change in the covariance estimation that would lead to a revolutionized/big change in the performance and/or efficiency. It will have a limited impact on the community.\n\n2. The experimental results seem not extraordinary. It does not fundamentally improve the performance (generation quality), which is understandable since covariance is more about sampling efficiency and sample diversity. However, even in terms of the sampling diversity and efficiency, such as indicated by Table 4, the proposed method does not significantly beat strong baselines. Moreover, there are no error bars or uncertainty in the experimental results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)