Gaussian Mixture Solvers for Diffusion Models
Authors: Hanzhong Allan Guo
Cheng Lu
Fan Bao
Tianyu Pang
Shuicheng YAN
Chao Du
Chongxuan Li
Venue: NeurIPS 2023 poster (2023)
Keywords: Diffusion models
SDE-based solver
Gaussian mixture
Stroke-based synthesis

Abstract: Recently, diffusion models have achieved great success in generative tasks. Sampling from diffusion models is equivalent to solving the reverse diffusion stochastic differential equations (SDEs) or the corresponding probability flow ordinary differential equations (ODEs). In comparison, SDE-based solvers can generate samples of higher quality and are suited for image translation tasks like stroke-based synthesis. During inference, however, existing SDE-based solvers are severely constrained by the efficiency-effectiveness dilemma. Our investigation suggests that this is because the Gaussian assumption in the reverse transition kernel is frequently violated (even in the case of simple mixture data) given a limited number of discretization steps. To overcome this limitation, we introduce a novel class of SDE-based solvers called \emph{Gaussian Mixture Solvers (GMS)} for diffusion models. Our solver estimates the first three-order moments and optimizes the parameters of a Gaussian mixture transition kernel using generalized methods of moments in each step during sampling. Empirically, our solver outperforms numerous SDE-based solvers in terms of sample quality in image generation and stroke-based synthesis in various diffusion models, which validates the motivation and effectiveness of GMS. Our code is available at https://github.com/Guohanzhong/GMS.

Gaussian Mixture Solvers for Diffusion Models
Hanzhong Guo Cheng Lu", Fan Bao", Tianyu Shuichenn Yan²,
Chao Du Chongxuan Liti,3
GGaoling School of Artificial Intelligence, Renmin University of China
"Sea AI Lab, Singapore
Beijing Key Laboratory of Big Data Management and Analysis Methods, Beijing, China
*Tsinghua University
(allanguo, tianyupang, yansc, duchao Gsea. com; lucheng. 1e15@gmail.......mmm..mmm.aal..m com; bf 19@mailss tsinghua. edu. en; chongxuanl i @rruc. edu. cn
Abstract
Recently, diffusion models have achieved great success in generative tasks. Sam- pling from diffusion models is equivalent to solving the reverse diffusion stochastic differential equations (SDEs) or the corresponding probability flow ordinary differ- ential equations (ODEs). In comparison, SDE-based solvers can generate samples of higher quality and are suited for image translation tasks like stroke-based synthe- sis. During inference, however, existing ADEEbased solvers are severely constrained dilemma. Our investigation suggests that this is be-
by the cause the Gaussian assumption in the reverse transition kernel is frequently violated (even in the case of simple mixture data) given limited number of discretization steps. To overcome this limitation, we introduce a novel class of SDE-based solvers called Gaussian Mixture Solvers (GMS) for diffusion models. Our solver estimates the first three-order moments and optimizes the parameters of a Gaussian mix- ture transition kernel using generalized methods of moments in each step during sampling. Empirically, our solver outperforms numerous SDE-based solvers in terms of sample quality in image generation and stroke-bar based synthesis in various diffusion models, which validates the motivation and effectiveness of GMS. Our code is available https://github.. httgggghww..ooonnouoaaoonnn/////
Introduction
In recent years, deep generative models and especially (score-based) diffusion models [35, 17, 18]
have made remarkable progress in various domains, including image generation [15, 7], audio
generation [22, 32], video generation [16], 3D object generation [31], multi-modal generation [42,
12, 43], and several downstream tasks such as image translation [28, 45] and image restoration [19].
Sampling from diffusion models can be interpreted as solving the reverse-time diffusion stochastic differential equations (SDEs) or their corresponding probability flow ordinary differential equations (ODEs) [37]. SDE-based and ODE-based solvers of diffusion models have very different properties and application scenarios In particular, based solvers usually perform better when given a sufficient number of discretization steps [18]. Indeed, recent empirical study [24] suggests that SDE-baaed solvers can potentially generate high-fidelity samples with realistic details and intricate semantic coherence from pre-trained large-scale text-to-image diffusion models [34]. Besides, SDE- based solvers are preferable in many downstream tasks such as stroke-based synthesis [27], image translation [45], and image manipulation [20].
"Work done during an internship at Sea Al Lab. CCoreespondence to Chao Du and Chongxuan Li.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Ground Truth
Ground Truth
Ground Truth
3.0
SN-DDPM
3.0
SN-DDPM
3.0
DDPM
GMS
GMS
GMS
2.0
2.0
2,0
1.3
1.0
1.0
1.0
0.5
0.0
0.0
0.0
-1.0
-0.5
0.0
0.5
1.0
-1.0
-0.5
0.0
0.5
1.0
(a) # #Step=5
(b) ( #tep==10
(c) # Step=220
mixture of Gaussian. GMS (ours) and SN-DDPM excel in fitting the true distribution when the transition kernel is Gaussian, with sufficient number of sampling steps (c). However, GMS outperforms SN-DDPM [2] when sampling steps are limited and the reverse transition kernel deviates from Gaussian (a-b).
Figure 1: Sampling
Despite their wide applications, SDE-based solvers face significant trade-off between efficiency and effectiveness during sampling, since an insufficient number of steps lead to larger discretization errors. To this end, Bao et al. [3] estimate the optimal variance of the Gaussian transition kernel the reverse process instead of using handcraft variance to reduce the discretization errors. Additionally, Bao et al. [2] explore the optimal diagonal variance when dealing with an imperfect noise network and demonstrate state-of-the-art (SOTA) performance with few steps among other SDE-based solvers. Notably, these solvers all assume that the transition kernel in the reverse process Gaussian.
In this paper, we systematically examine the assumption of the Gaussian transition kernel and reveal that it can be easily violated under limited number of discretization steps even the case of simple mixture data. To this end, we propose a new type of SDE- solver called Gaussian Mixture Solver (GMS), which employs more expressive Gaussian mixture transition kernel in the reverse process for better approximation given limited number of steps (see visualization results on toy data Fig. 1). In particular, we first learn noise prediction network with multiple heads that estimate the high-order moments of the true reverse transition kernel respectively. For sampling, we fit Gaussian mixture transition kernel each step via the generalized methods of the moments [11] using the predicted high-order moments of the true reverse process.
To evaluate GMS, we compare it against a variety of baselines [14, 17, 36] in terms of several widely used metrics (particularly sample quality measured by FID) in the tasks of generation and stroke-based synthesis
multiple datasets. Our results show that GMS outperforms state-of-the-art SDE- solvers [14, 2, 17] in terms of sample quality with the limited number of discretization steps (e.g.. < 100). For instance, GMS improves the FID by 4. 4.44 over the SOTA SDE- solver [2] given 10 steps on CIFAR1O. Furthermore, We evaluate GMS on a stroke-based synthesis task. The findings consistently reveal that GMS achieves higher levels of realism than all aforementioned SDE-based solvers as well as the widely adopted ODE-based solver DDIM [36] while maintaining comparable computation budgets and faithfulness scores (measured by T-2 distance). Such empirical findings validate the motivation and effectiveness of GMS.
2
Background
In this section, we provide brief overview of the (score-based) diffusion models, representative SDE-based solvers for diffusion models, and applications of such solvers.
2.1 Diffusion models
Diffusion models gradually perturb data with forward diffusion process and then lear to reverse such process to recover the data distribution. Formally, let Øy E R" be a random variable with unknown data distribution (zr). Diffusion models define the forward process {2x {zthejo.u) indexed by time i, which perturbs the data by adding Gaussian noise 2o with g(x) ( (t)
(I)

---

In general, the function a(t) and (t) are selected so that the logarithmic signal-tonooie ratio log decreases monotonically with time i, causing the data to diffuse towards random Gaussian noise [21]
20
Furthermore, it has been demonstrated by Kingma et [21] that the following SDE shares an identical transition distribution (00()e)200) with Eq. (1):
dæy = [())) + g(t)du,, zo ~ g((o);
where w E IR" is a standard Wiener process and
f(t) dlog a(t) (²(() (t) 20 d log dllog(t)) dt
dt
dt
Let g(2+)) be the marginal distribution of the above SDE at time t. Its reversal process can be described by a corresponding continuous SDE which recovers the data distribution [37]:
dx = f())=; - o V log ((2+)] dt + g(t)da 21 ~ q(21).
(4)
where @ E Rn is a reverse-time standard Wiener process. The only unknown term in Eq. (4) is the score function v.. log q(zz). To estimate it, existing works [14, 37, 18] train noise network c t) to obtain a scaled score function o Vt) log q(ra) using denoising score matching [38]:
(())= = fo w(t) (t)) (ote) t) <lgidr,
(5)
where w(() is a weighting function, q({) standard Gaussian distribution and s4 ~ q(2, (mr) follows Eq. (1). The optimal solution of the optimization objective Eq. (5) is co () -o(() log g(zz).
Hence, samples can be obtained by initiating the process with a standard Gaussian variable x1, then substituting Vz, log q(zz) with salmedd and discretizing reverse SDE Eq. (4) from & to / = to generate .ro
(t)
2.2 SDE-b solvers for diffusion models
The primary objective of SDE-based solvers lies in decreasing discretization error and therefore minimizing function evaluations required for convergence during the process of discretizing Eq. (4). Discretizing the reverse SDE in Eq. (4) is equivalent to sample from a Markov chain p( = p(z1) I.......ss, p( )t,)) with its trajectory S; = [0, f1, f2...., ti,.., 1]. Song et al. [37] proves that the conventional ancestral sampling technique used in the DPMs [14] that models p( (mx;)
as a Gaussian distribution, can be perceived as a first-order solver for the reverse SDE in Eq. (4). Bao al. [3] finds that the optimal variance of p( (2x,,)
1 11. (zt,)) )) is
02(4)) (1 (1- (eg()
1 d (lo 4:)) 1[])
11 114, (1te.) = X2. + 7t. act;)
(6)
X is the variance of g( [ 20). Analyt- icDPM [3] offers a significant reduction in discretization error during sampling and achieves faster convergence with fewer steps. Moreover, SN-DDPM [2] employs Gaussian transition kernel with an optimal diagonal covariance instead of an isotropic covariance. This approach yields improved sample quality and likelihood compared to other SDE-based solvers within a limited number of steps.
where, Tt. Va((iii))
(1 -ool---))oo((((-)))))) X
(at)))
(2,)
2.3 Applications of SDE-based solvers for stroke-based synthesis
The stroke-based image synthesis is a representative downstream task suitable for SDE-based solvers. involves the user providing a full-resolution image r(()) through the manipulation of RGB pixels, referred as the guided image. The guided image 2(0) possibly contains three levels of guidance: a high-level guide consisting of coarse colored strokes, a mid-level guide comprising colored strokes on a real image, and a low-level guide containing image patches from target image.
SDEdit [27] solves the task by first starting from the guided image x(9)) and adding Gaussian noise to disturb the guided images to 2, and ((x(9) (t))) ~ N (att (oan)),,(())))) (t)(T) same with Eq. (1). Subsequently, it solves the corresponding reverse stochastic differential equation (SDE) up to / = to generate the synthesized image z0) discretizing Eq. (4). Apart from the discretization steps taken by the SDE-based solver, the key hyper-parameter for SDEdit is l. the time step from which we begin the image synthesis procedure in the reverse SDE.
3

---

10-2
Log L2 Norm
Log L2 Norm
10
10-5
10-5
10-9
#Step:10
10-12
#Step:10
#Step:20
#Step:20
11
#Step:100
10-13
#Step:100
Step:1000
#Step:1000
200
800
1000
Time Step
Time Step
(a) logarithm mean of Lg-norm
(b) logarithm median of L2-norm
Figure 2: Empirical evidence of suboptimality of Gaussian kernel on CIFAR1O. (a) and (b) plot the logarithm of the image-wise mean and median L2 Le-norm between caussian-assumed third-order moments and estimated third-order moments. Clearly, as the number of sampling steps decreases, the disparity between the following two third-order moments increases, denoting that the true transition kernel diverges further from the Gaussian distribution. See Appendix D.2 for more details.
3 Gaussian mixture solvers for diffusion models
In this section, we first show through both theoretical and empirical evidence, that the true reverse transition kernel can significantly diverge from Gaussian distributions as assumed in previous SOTA SDE-based solvers [3, 2], indicating that the reverse transition can be improved further by employing more flexible distributions (see Sec. 3.1). This motivates us to propose a novel class of SDE- based solvers, dubbed Gaussian Mixture Solvers, which determines Gaussian mixture distribution via the generalized methods of the moments [1 [1]] using higher-order moments information for the true reverse transition (see Sec. 3.2). The higher-order moments estimated by a noise prediction network with multiple heads on training data, as detailed Sec. 3.3.
3.1 Suboptimality of Gaussian distributions for reverse transition kernels
As described in Sec. existing sfate-of-the-ar SDE-based solvers for DPMs [3, 2] approximate the reverse transition q (I,,) using Gaussian distributions. Such approximations work well when the number of discretization steps these solvers large (e.g.. 1000 steps). However, for smaller discretization steps (such as when faster sampling is required), the validity of this Gaussian assumption will be largely broken. We demonstrate this theoretically and empirically below.
|xa )dzn, which is non-Gaussian for general data distribution ((xa). For instance, for mixture of Gaussian mixture of Dirac g(#o), we can prove that the conditional distributions q( ( (x,)) in the reverse process are non-G Gan-Gaasiin, as characterized by Proposition 3.1, proven in Appendix A.I.
First, observe that ( can be expressed as g( = [f((/// vo (o)g(<o
Proposition 3.1 (Mixture Data Have Gaussann Reverse Kernel). Assume q( is mixture of Dirac or a mixture of Gaussian distribution and the forward process is defined in Eq. (1). The reverse transition kernel q (x);;
: 8 < is Gaussian mixture instead of a Gaussian.
Empirically, we do not know the distributions of the real data (e.g., high-dimensional images) and cannot obtain an explicit form of q( (.te). However, even such cases, it is easy to validate that (x;)) are non-Gaussian In particular, note that the third-order moment of one Gaussian distribution (M) can be represented by first-order moment (M1) and its second-order moment (M2) 2, which motivates relationship induced by the Gaussian assumption. We perform an experiment on CIFAR- 10 and estimate the first three orders of moments M,. M2. M3 for g(a, networks (see details Sec. D.2). As shown Fig. 2, we plot the mean and median of 1g-norm
check whether the first three order moments of ((x))x;) satisfy the
(x;) from data by high-order noise
'To enhance the clarity of exposition, we introduce the notation = 4-11 and / = L, denote two adjacent time steps in the trajectory. Consequently, refer a(2, (x.) the reverse transition kernel in this context. Thh scalar case is M M 3M1(MM Mf). See Appendix D.2 for the vector case.

---

between the estimated third-order moment M3 and third-order moment M calculated under the Gaussian assumption given the different number of steps at different time steps t. In particular, when the number of steps is large (i.e.. #Step=1000), the difference between the two calculation methods is small. As the number of steps decreases, the 12-norm increases, indicating that the true reverse distribution g( (aze) is Gaussinsn With the time step closer t = 0, the 12-norm increases too.
Both the theoretical and empirical results motivate us to weaken the Gaussian assumption in SDE- based solvers for better performance, especially when the step size is large.
3.2 Sampling with Gaussian mixture transition kernel
There are extensive choices of p(z, (az) such that it is more powerful and potentially fits q(z, (az)) better than a Gaussian. In this paper, we choose a simple mixture of the Gaussian model as follows:
M
M
p(
Σ wi ((N(
w =1
where w; is a scalar and (1:(2r) and (xx) are vectors. The reasons for choosing a Gaussian mixture model are fhree-fold. First, Gaussian mixture model is multi-modal (e.g., see Proposition 3.1), potentially leading to a better performance than Gaussian with few steps. Second, when the number of steps is large and the Gaussian is nearly optimal [35, 3], a designed Gaussian mixture model such as our proposed kernel in Eq. (9) can degenerate to Gaussian when the mean of two components are same, making the performance unchanged. Third, a Gaussian mixture model is relatively easy to sample. For the sake of completeness, we also discuss other distribution families in Appendix C.
Traditionally, the EM algorithm is employed for estimating the Gaussian mixture [26]. However, it is nontrivial to apply EM here because we need to learn the reverse transition kernel in Eq. (7) for all time step pairs (s, () by individual EM processes where we need to sample multiple x.. given a Ty to estimate the parameters in the Gaussian mixture indexed by (s, l). This is time-consuming. especially in a nigh-dimensional space (e.g., natural images) and we present an efficient approach.
For improved flexibility in sampling and training. diffusion models introduce the parameterization of the noise network e(ze, t) [14] the data prediction network oo(2,, f) [18]. With such network, the moments under the ( (x)) measure can be decomposed into moments under the g (2)) measure that sampling any Iy I, requires only network whose inputs are Ia and f, such as the decomposition shown in Eq. (10). In previous studies, Gaussian transition kernel was utilized, allowing for the distribution be directly determined after obtaining the estimated first-order moment and handcrafted second-order moment [14] or estimated first-order and second-order moment [3]. In contrast, such feature is not available for the Gaussian mixtures transition kernel in our paper.
Here we present the method to determine the Gaussian mixture given a set of moments and we will discuss how to obtain the moments by the parameterization of the noise network in Sec. 3.3. Assume the length of the estimated moments set is and the number of parameters in the Gaussian mixture is d. We adopt popular and theoretically sound method called the generalized method of moments (GMM) [11] to learn the parameters by:
N Σ g(., 0) W
N Σ g( 0)],
min (1, MN) = min 9
(8)
0
where 0 Rdx1 includes all parameters (e.g., g., the mean of each component) in the Gaussian mixture defined in Eq. (7). For instance, d = 2.M * Damaa +M in Eq. (7), where D represents the number of dimensions of the data and 9 contains M mean vectors with Daata dimensions, M variance vectors of with Daat dimensions (considering only the diagonal elements), and M 1 weight coefficients which are scalar. The component of g( 0) E Nx1 is defined by gn (z,, 0) ( MGM (GM)( where M,, (zx) is the n-th order empirical moments stated in Sec. 3.3 and M1, (G (0) is the n-th order moments of Gaussian mixture under 0, and W is weighted matrix, N₂ is the number of samples.
Theoretically, the parameter ŪaMm obtained by GMM in Eq. (8) consistently converges to the potential optimal parameters 0" for Gaussian mixture models given the moments' condition because
V. 0-) 4 (O, V(DGMM))) as stated in Theorem 3.1 in Hansen [1]].
NNoe that serves as the "training sample" in the corresponding EM process.
5

---

Hence, we can employ GMM to determine a Gaussian mixture transition kernel after estimating the moments of the transition kernel. To strike a balance between computational tractability and expressive power, we specifically focus on the first three-order moments in this work and define
Gaussian mixture transition kernel with two components shown in Eq. (9) whose vectors Pi Mi and of are parameters to be optimized. This degree of simplification is acceptable in terms of its impact. Intuitively, such a selection has the potential to encompass exponential modes throughout the entire trajectory. Empirically, we consistently observe that utilizing a bimodal Gaussian mixture yields favorable outcomes across all experimental configurations.
(1) (2)
( of) p(a,|==) 3 (ul
+ 2 ( of). 3
(9)
meanwhile, under the simplification in our paper, the number of parameters d d= 3 Daata in our Gaussian transition kernel is equal to the number of moments condition N = 3 * D data- According to proposition 3.2, under the selection of arbitrary weighted weights, the asymptotic mean (asymptoti- cally consistent) and asymptotic variance of the estimator remain consistent, proof in Appendix A.2. Hence, any choice of weighted weights is optimal. To further streamline optimization, we set W I.
Proposition 3.2 (Any weighted matrix is optimal for = N). Assume the number of parameters d equals the number of moments condition N, and C3 (z+,t) 2n ( @"-1 c)] (where @" denotes n-fold outers product) which denotes n-th order noise network converging in probability. The asymptotic variance V ((GaM)) and the convergence speed of GMM remain the same malter which weighted matrix is adopted in Eq. (8), Namely, any weighted matrix is optimal.
What's more, we provide a detailed discussion on the selection of parameters such as the choice of different parameters to optimize and the different choices of weight w, in the Gaussian mixture in Appendix E. Combining with Sec. 4.1, our empirical findings illustrate the efficacy of the GMS across various benchmarks via using the Gaussian transition kernel in Eq. (9) fitted by the objective function shown in Eq. (31) in Appendix B.1 via the ADAN [41] as optimization method. Details regarding the parameters for this optimizer are provided Appendix E.5.
3.3 Estimating high-order moments for non-Gaussian reverse process
In Sec. 3.2, we have explored the methodology for determining a Gaussian mixture transition kernel given set of moments Mj...., Mn. In the subsequent section, we will present an approach for estimating these moments utilizing noise networks and elucidate the process of network learning.
Given the forward process described by Eq. (1), it can be inferred that both g( (ar.) and ( mn) follow Gaussian distributions. Specifically, (xx//1=)) NN(neeayps o/ T) [21], where Orls =
a(t) and og = ot)) - o2ioo(()). Consequently, we can deduce that E 1x O" x..] = a(s)
Eg(rallr))(( (o) [x, 8" xz]. where 8 denotes n-fold outer product. Thus, we can first utilize the Gaussian property of g( (rr. 20) and employ an explicit formula to calculate the moments under the measure of q(z,, (xt, What's more, we only consider the diagonal elements of higher-order moments for computational efficiency, similar to Bao et al. [2] since estimating full higher-order moments results escalated output dimensions (e.g., quadratic growth for covariance and cubic for the third-order moments) and thus requires substantial computational demands. The expression of diagonal elements of the third-order moment of the reverse transition kernel can be derived as:
M, E Bg(xs/ze) (zx) [diag(E, (I, & x. & xz)]
Ex( Bg( (xe (diag(a (x, X x. C xx)]
(10)
arjs diag(a, @ xx 8 xx) + 3X3
7 x1
Constant term
3a
a.. 2 82 f|s ( @ xx)) +
a 30
n ④ Ret [co]
of
of
Linear term
2 G010 3 ( |diag( @ .g @ 2o)], ax Aui )2x 3 ( diag(ro @ xo)] +
+ 3
of
2 of
of
Quadratic tenn in
Cubic tern 20g

---

Algorithm Learning of the high order noise network
1: Input: The n-th order noise network 60 (Ie, () and its tunable parameter on.
2: repeat
3: 20 ((<o)
4: f ~ Uniform([.......
5: 6, 1)
6: 2t = a(t)xo + at)e,
7: Take gradient step on 0% 4- V on Ee E0 (Tee e0/
8: until converged
Algorithm Sampling via GMS with the first three order moments
1: Input: The assembled noise network (x t) in Eq. (13)
2: 2T ~ No0, T)
3: for f = -T to do
4: 2, ~ ()
5: M2, M3 = h(ffi) (2,, (21ff (2. (x t), f(a (22, t)) in Appendix
(x,,
6:
Th. o min N2.00 Q(T M4: Tse My, M2, M3) in Sec. 3.2
7: i~ Th. D1-1 M +0.2
8: end for
9:
return 20
where & is the outer product @ is the Hadamard product. Additionally, X corresponds to the variance of (1,, zn), and further elaboration on this matter can be found in Appendix B.
In order compute the n-th order moment Mn, as exemplified by the third-order moment in Eq. (10), it is necessary to evaluate the expectations 1 (x,) [ra].-- (zole,) [diag (zo @=-1 z0)].. Further- more, by expressing Zo as 20 Vo(o (x, (1)e), we can decompose M, into a combination of terms By(=o) ( ==)(E]........ gi (m,) ((e @n-1 e)], which are learned by neural networks. The decom- position of third-order moments M3. as outlined in Eq. (29). is provided to illustrate this concept. Therefore in training, we learn several neural networks {E/ J= by training on the following objective functions:
J
min (eBg(ro)o()) ||e" (z,, 1)112 {.r )
(11)
N where ~ NOO, = a(t) ((t)) and en denotes con-le c. After training, we can use {<] f==1 to replace [Bg(zolz,) 6n1 O e] and estimate the moments J. of reverse transition kernel.
However, in the present scenario, it is necessary infer the network minimum of n times order to make single step of GMS. To mitigate the high cost of the aforementioned overhead in sampling, we adopt the two-stage learning approach proposed by Bao al. [2]. Specifically, in the first stage, we optimize the noise network 6ø (Zz, t) by minimizing the expression min EgEg(ra))te ) ||e cœ(x,, t) or by utilizing pre-trained noise network proposed by [14, 37]. In the second stage, utilize the optimized network as the backbone and keep its parameters fixed, while adding additional heads to generate the n-th order noise network CD.o., (xe, t).
Ce,,. (Ie, NN (co (#),
(12)
where NN is the extra head, which is small network such convolution layers small attention block, parameterized by on. details in Appendix E.3. We present the second stage learning procedure in Algorithm 1. Upon training all the heads,
readily concatenate the outputs of different heads, as the backbone of the higher-order noise network is shared. By doing so, we obtain the assembled noise network y (rr, When estimating the k-th order moments, it suffices extract only the first k components the assembled noise network.
fN (Iz, t) = concat( (x t), E (Ie, t) ... eli (#e,, ..., N. (æ1, t))),
(13)
Required for estimating Mr
7

---

Table 1: Comparison with competitive SDE-based solvers w.r.t FID score + CIFAR10 and ImageNet 64x 64. our GMS outperforms existing SDE-based solvers in most cases. SN-DDPM denotes Extended AnalyticDPM from Bao et al. [2].
CIFARIO (LS)
# TIMESTEPS K
10
20
25
40
50
100
200
1000
DDPM, Be
43.14
25.28
21.63
15.24
15.21
10.94
8.23
5.11
DDPM, Be
233.41
168.22
125.05
82.31
66.28
31.36
12.96
3.04
SN-DDPM GMS (OURS)
21.87
8.32
6.91
4.99
4.58
3.74
3.34
3.71
17.43
7.18
5.96
4.52
4.16
3.26
3.01
2.76
CIFARIO (CS)
IMAGENET x
# TIMESTEPS K
10
25
50
100
200
1000
25
50
100
200
400
4000
DDPM, B
34.76
16.18
11.11
8.38
6.66
4.92
29.21
21.71
19.12
17.81
17.48
16.55
DDPM. BE
205.31
84.71
37.35
14.81
5.74
3.34
170.28
83.86
45.04
28.39
21.38
16.38
SN-DDPM
16.33
6.05
4.19
3.83
3.72
4.08
27.58
20.74
18.04
16.72
16.37
16.22
GMS (OURS)
13.80
5.48
4.00
3.46
3.34
4.23
26.50
20.13
17.29
16.60
15.98
15.79
To this end, outline the GMS sampling process in Algorithm 2, where f/2 (2e, t) represents concat([co ([o t), c..0. (z,, t))). comparison to existing methods with the same network structure, we report the additional memory cost of the assembled noise network in Appendix E.6.
4 Experiment
In this section, first illustrate that GMS exhibits superior sample quality compared to existing SDE-based solvers when using both linear and cosine noise schedules [14, 30]. Additionally, we evaluate various solvers in stroke-bite based image synthesis (i.e., SDEdit) and demonstrate that GMS surpasses other SDE-based solvers, as well as the widely used ODE-based solver DDIM [36].
4.1 Sample quality image data
In this section, we conduct quantitative comparison of sample quality using the widely adopted FID score [13]. Specifically, we evaluate multiple SDE-based solvers, including comparison with DDPM [14] and Extended AnalyticDPM [2] (referred to as SN-DDPM) using the even trajectory.
As shown in Tab. 1, GMS demonstrates superior performance compared to DDPM and SN-DDPM under the same number of steps in CIFARIO and ImageNet 64 64. Specifically, GMS achieves remarkable 4.44 improvement in FID given 10 steps CIFAR 10. Appendix E.7 illustrates in more de- tail the improvement of GMS when the number of sampling steps is limited. Meanwhile, we conduct GMS in ImageNet 256 256 via adopting the U-ViT-Huge [4] backbone the noise network in Ap- pendix E.8. Furthermore, taking into account the additional time required by GMS, our method still ex- hibits improved performance, detailed in Appendix E.9. For integrity, we provide a comparison with other SDE-based solvers based on continuous time diffusion such as Gotta Go Fast [17], EDM [18] and SEED [10] in Appendix 10 and shows that GMS largely outperforms other SDE-based solvers when the number of steps is less than 100. In Appendix G, we provide generated samples from GMS.
4.2 Stroke-based image synthesis based on SDEdit [27]
Evaluation metrics. We evaluate the editing results based on realism and faithfulness similar with Meng et al. [27]. To quantify the realism of sample images, we use FID between the generated images and the target realistic image dataset. To quantify faithfulness, report the Lg distance summed over all pixels between the stroke images and the edited output images.
SDEdit [27] applies noise to the stroke image 2: at time step 4n using N(2 Jo(())>, (²) and discretize the reverse SDE Eq. (4) for sampling. Fig. demonstrates the significant impact of to on the realism of sampled images. As increases, the similarity to real images decreases. choose

---

Realism (FID)
30
GMS (Ours)
DDIM
38
22
21
DUMM SIi DODM DMS
DCW
DOM SN-20NN
54400M9
20
32
10
28
17
12:
0.2
Falthfuiness (Ly-norm
0.3
0.4
0.5
0.6
25
Number stops
50
100
25
Number of steps
50
100
25
Number steps
50
100
(a) Trade-off between realism and faithfulness
(h) to 0.37
(c) 0.47
(d) fo 0.57
Figure 3: Result among different solvers in SDEdit. to denotes the time step of the start of reverse. (a): The points on cach line represent the same to and the number of sampling steps. We select to [300, 400, 500], number of steps [50, 100]. (b): When to 300, the effect of DDIM diminishes more prominently with the increase in the number of steps.
the range from to 0.37 to = 0.57 (T' 1000 in our experiments) for our further experiments since sampled images closely resemble the real images in this range.
Fig. 3(a) illustrates that when using the same to and the same number of steps, edited output images from GMS have lower faithfulness but with higher realism. This phenomenon likely attributed the Gaussian noise introduced by the SDE-based solver during each sampling step. This noise causes the sampling to deviate from the original image (resulting in low faithfulness) but enables the solver transition from the stroke domain to the real image domain. Fig. 3(b) to Fig. 3(d) further demonstrates this phenomenon to certain extent. The realism of the sampled images generated by the SDE-based solver escalates with an increase in the number of sampling steps. Conversely, the realism of the sampled images produced by the of noise, which prevents the ODE-based solver from transitioning from the stroke domain to the real image domain. Additionally, in the SDEdit task, GMS exhibits superior performance compared to SN-DDPM [2] in terms of sample computation cost. Fig. shows the samples using DDIM and GMS when to 400 and the number of steps is 40.
solver diminishes due to the absence
Related work
Faster solvers. In addition to SDE-based solvers, there are works dedicated to improving the efficiency of ODE-based solvers [25, 23, 8]. Some approaches use explicit reverse transition kernels, such those based on generative adversarial networks proposed by Xiao et al. [40] and Wang et al. [39]. Gao et al. [9] employ energy function to model the reverse transition kernel. Zhang and Chen [44] use flow model for the transition kernel.
Non-Gaussian diffusion. Apart from diffusion, some literature suggests using non-Gaussian forward processes, which consequently involve non-Gaussian reverse processes. Bansal al. [1] introduce generalized noise operator that incorporates noise. Nachmani et al. [29] incorporate Gaussian mixture or Gamma noise into the forward process. While these works replace both the forward and reverse processes with non- Gaussian distributions, our approach aims to identify a suitable combination of non-Gaussian distributions model the reverse process.
6 Conclusion
This paper presents a novel Gaussian mixture solver (GMS) for diffusion models. GMS relaxes the Gaussian reverse kernel assumption reduce discretization errors and improves the sample quality under the same sampling steps. Experimental results show that GMS outperforms existing SDE-based solvers, achieving remarkable 4.44 improvement in FID compared to the state-of-the-art SDE-based solver proposed by Bao et al. [2] given 10 steps. Furthermore, due to the presence of noise, SDE-based solvers prove more suitable for stroke-based synthesis tasks and GMS still outperforms satte-of-thh-aae SDE-based solvers.

---

(a) Images
(b) Stroke
(c) GMS
(d) DDIM
Figure 4: SDEdit samples of GMS and DDIM. Due to the presence of noise, SDE-based solvers, such as GMS OURS, generate images with more details.
Limitations and broader impacts. While GMS enhances sample quality and potentially accelerates inference speed compared to existing SDE-based solvers, employing GMS still fall short of real-time applicability. Like other generative models, diffusion models can generate problematic fake content, and the use of GMS may amplify these undesirable effects.
Acknowledgement
This work was supported by NSF of China (Nos. 62076145); Beijing Outstanding Young Scientist Program (No. HWWWZYYJH0120111000000000
Major Innovation Planning Interdisciplinary Platform for the "Double-First Class" Initiative, Renmin University of China; the Fundamental Research Funds for the Central Universities, and the Research Funds of Renmin University of China (No. 22XNKJ13). C. Li was also sponsored by Beijing Nova Program (No. 20220484044).
References
[1] Arpit Bansal, Eitan Borgnia, Hong-Min Chu, Jie S Li, Hamid Kazemi, Furong Huang, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Cold diffusion: Inverting arbitrary image transforms without noise. arXiv preprint arXiv:2208.098: 09392,9992.... 2022.
[2] Fan Bao, Chongxuan Li, Jiacheng Sun, Jun Zhu, and Bo Zhang. Estimating the optimal covari- ance with imperfect mean in diffusion probabilistic models. arXiv preprint arXiv:2206.07309. 2022.
[3] Fan Bao, Chongxuan Li, Jun Zhu, and Bo Zhang. Analytic-dpm: analytic estimate of the optimal reverse variance diffusion probabilistic models. arXiv preprint arXiv:2201.06550, 2022.
10

---

[4] Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pages 22669-22677 2023.
[5] Fan Bao, Shen Nie, Kaiwen Xue, Chongxuan Li, Shi Pu, Yaole Wang, Gang Yue, Yue Cao, Hang Su, and Jun Zhu. One transformer fits all distributions in multi-modal diffusion at scale. arXiv preprint arXh::2333.066 30550.55...5..... 2023.
[6] Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning, arXiv preprint arXiv:2003.04297 2020.
[7] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. Advances in Neural Information Processing Systems, 34:8780-8794, 2021.
[8] Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Genie: Higher-order denoising diffusion solvers. arXiv preprint (rXiv:22/0000545. 2022.
[9] Ruiqi Gao, Yang Song, Ben Poole, Ying Nian Wu, and Diederik Kingma. Learning energy- based models by diffusion recovery likelihood. arXiv preprint arXiv:20/2.088/25 2020.
[10] Martin Gonzalez, Nelson Fernandez, Thuy Tran, Elies Gherbi, Hatem Hajri, and Nader Mas- moudi. Seeds: Exponential sde solvers for fast high-quality sampling from diffusion models. arXiv preprint ar0.v: 14267, 2023.
[11] Lars Peter Hansen. Large sample properties of generalized method of moments estimators. Econometrica: Journal of the econometric society, pages 1029-1054, 1982.
[12] Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfir Aberman, Yael Pritch, and Daniel Cohen-Or. Prompt-to-promp image editing with cross attention control. arXiv preprint urXiv:2208.01626, 2022.
[13] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[14] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6855 2020.
[15] Jonathan Ho. Chitwan Saharia, William Chan, David Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. Mach. Learn. Res., 23(47):1-33, 2022.
[16] Jonathan Ho, Tim Salimans. Alexey Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. arXiv preprint arXiv:2204.0344 2022.
[17] Alexia Lolicoeur-Martineaa, Ke Li, Rémi Piché-Tailleffer Tal Kachman, and Ioannis Mitliagkas. Gotta go fast when generating data with score-based models. arXiv preprint arXiv:2 14080, 2021.
[18] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. arXiv preprint arXiv:2206.00366 2022.
[19] Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models, arXiv preprint arXiv:220..11999; 2022.
[20] Bahjat Kawar, Shiran Zada, Oran Lang, Omer Tov, Huiwen Chang, Tali Dekel, Inbar Mosseri, and Michal Irani. Imagic: Text-based real image editing with diffusion models. arXiv preprint arXiv:2210.09276. 2022.
[21] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances neural information processing systems, 34:21696-21 2021.
[22] Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: versatile diffusion model for audio synthesis. arXiv preprint 207.0060000009611, 2020.
11

---

[23] Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. arXiv preprint arX1v:2202.09778. 2022.
[24] Cheng Lu, 2023.
[25] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. arXiv preprint arXiv:2206.00927. 2022.
[26] Geoffrey J McLachlan and Thriyambakam Krishnan. The EM algorithm and extensions. John Wiley & Sons, 2007.
[27] Chenlin Meng. Yang Song, Jiaming Song, Jiajun Wu, Jun- Yan Zhu, and Stefano Ermon. Sdedit: Image synthesis and editing with stochastic differential equations. arXiv preprint arXiv:2008.01073, 2021.
[28] Eliya Nachmani and Shaked Dovrat. Zero-shot translation using diffusion models. arXiv preprint arXiv:2111.01471, 2021.
[29] Eliya Nachmani, Robin San Roman, and Lior Wolf. Non gaussian denoising diffusion models. arXiv preprint (xXiv:2006.07582.. 2021.
[30] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pages 8162-8177 1171. PMLR, 2021.
[31] Ben Poole, Ajay Jain, Jonathan T Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using
2d diffusion. arXiv preprint arXiv:2209. 14988, 2022.
[32] Vadim Popov, Ivan Vovk. Vladimir Gogoryan, Tasnima Sadekova, and Mikhail Kudinov. Grad- uts: diffusson probabilistic model for text-to-speech. International Conference on Machine Learning, pages 8599-8608. PMLR, 2021.
[33] Prajit Ramachandran, Barret Zoph. and Quoe V Le. Searching for activation functions. arXiv preprint arXiv: 2017.
[34] Alex Shonenkov, Misha Konstantinov, Daria Bakshandaeva, Christoph Schuhmann, Ksenia Ivanova, and Nadiaa Klokova, 2023. attps://githuu com//eep-ffoyd/Ir
[35] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256-2265. PMLR, 2015.
[36] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.005002, 2020.
[37] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011..3355. 2020.
[38] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1661-1674 2011.
[39] Zhendong Wang, Huangjie Zheng, Pengcheng He, Weizhu Chen, and Mingyuan Zhou. Diffusion-gan: Training gans with diffusion. arXiv preprint arXiv:2006.02262, 2022.
[40] Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. arXiv preprint arXiv:2112.07804, 2021.
[41] Xingyu Xie, Pan Zhou, Huan Li, Zhouchen Lin, and Shuicheng Yan. Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models. arXiv preprint arXiv:2208.06677 2022.
[42] Xingqian Xu, Zhangyang Wang, Eric Zhang, Kai Wang, and Humphrey Shi. Versatile diffusion: Text, images and variations all in one diffusion model. arXiv preprint arX1v:2211.08332, 2022.
12

---

[43] Lvmin Zhang and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models. arXiv preprint arXiv:2302.05543, 2023.
[44] Qinsheng Zhang and Yongxin Chen. Diffusion normalizing flow. Advances in Neural Informa- lion Processing Systems, 34: 1220111111,,, 2021.
[45] Min Zhao, Fan Bao, Chongxuan Li, and Jun Zhu. Egsde: Unpaired image-to-image translation via energy-guided stochastic differential equations. arXiv preprint arXiv:2207.06635, 2022.
13

---

A Proof
A.1 Proof of Proposition 3.1
When ((2)) is a mixture of Dirac distribution which means that g Emn w, M
da), Di=1 W; 1, which has total M components, and when the forward process ( is Gaussian distribution as Eq. (1), the reverse process ((x))z)) would be:
(x1a) = ( f0) (o) = w00p://e..0)))22))aa
M
=1/g(at) /
= 1/q(x)) (
According the definition of forward process, the distribution (1 Maral (1 Tz) Due to the Markov property of forward process, when > 8, we have q(C,, Te (mn g(( ( (oo)g((z s.s). The term g( (o) would be viewed Bayesian posterior resulting from prior q (20)) updated with a likelihood term (x)))x,). And therefore
Orl of
a
L
(x))) = (1 20), Mg (21, mo) N2x,
+
of
20.
of
1
It is easy to prove that, the distribution ((x,(11)) (zz) is mixture of Gaussian distribution:
M
x1)
uig(vslzr, (x)
M
wiN((r Ta (1 oz)1)= (T), (2 mo))
When is large, is small, ofe would large, meaning that the influence of zij would be large. Secondly, when (xxO) is mixture Gaussian distribution which means that q( EM w/N(slles E,), E Wi = 1. For simplicity of analysis, we assume that this distribu- tion a one-dimensional distribution or that the covariance matrix is gighdimensional Gaussian distribution with diagonal matrix Σ, = diag, Similar to the situation above, for each dimension in the reverse process:
q( (x,) 1/9(zz) / g(x, (1e (0)g((0))(2))))))))))))))))))))) ( 1o)
ws/a(xxe) / (l 2o) (o) Eg (3r dro
Σ ws/gt (05/g(z)) /
M
(4) (2)
64g =2)
1
dxo
V2TTg
VITO
2T /1 02
i=1
Σ ws/ ((xx)) /
M
mx)
((x=p(z)))= o(2e)
dop
V2TT
And (x.), could be Gaussian mixture which has M component.
A.2 Proof of Proposition 3.2
Recall that our objective function is to find optimal parameters:
i argmin Q QN_
argmin
gg. (0) (N.
(14)
1x1
IXN NXN 1XN
14

---

where (f) is the moment conditions talked about in Sec. 3.2, W is the weighted Matrix, Ne is the sample size. When solving such problems using optimizers, which are equivalent to the one we used when selecting aGMas such that @Qr(0gmar) O, and its first derivative:
09
0QN,(0) (0) 30,
0QN, (0) 00
JQ D0.
No Og(ai, 0)
M Σ 1
DQN. (02
WN. N
g( 0)]
(15)
30g
N 2=1
d0m
8QN (02
dx1
30g
NxN
1xN
xx1
its second derivative (the Hessian matrix):
8² B*QN, (8)
8 PQx,100
8 (QN.110)
28,00
00,00
00,04
@Q.. (0)
8 @'QN, (0)
88201
(16)
802
8 QNQ (0)
dxd
30,0d
82 @2QN. (0)
No
Ög(r,, 0)
No
0g(# 0)
2
00,0,
M
00;
No
00,
1=1
O g 0
No
2
WNN.
No
g 0)}.
00.0;
1=1
By Taylor's expansion of the gradient around optimal parameters 0o have:
aQ N. (0GMM)
dQN.
8QQ.. (0)) (#GMM - 0)
(17)
00
00
0000T
8² QQN. (0o) )-1 (Q).(0o)
(0GMM 0) RJ - 0000T
00
Consider one element of the gradient vector
aQN, (8o)
86.
dQ..
No
Og(I,, 0o)
N
2
r WN.
Σ g(,,, Bu)
(18)
00
M 1
N
ZW
Dalin 20l =Ioom 60m
LElg(E,
Consider one element of the Hessian matrix
08,303
TQN. (0o)
1
No
Ög(zz, 00)
N
0g (2,,,,00 Bo)
= 2
(19)
00:005 Ne
00,
No
00,
i=1
i=1
S
Pg(zz 0o)
N
2
(T WN..
Σ
g( 00)1
F, 21.,. WTo
00,0,
No
1=1
Therefore, is easy to prove that AGMM
Ou P 0. and uses law of large numbers we could obtain,
VT (Qr(6b) 2
1
ag(z; 6o)
T
WM
g( 0o)
1 2 b (. @o).
00mm
M
a0m. m
EW
2,
3o(X.
Tom
(20)
therefore, we have
(²Qr(00)
QQr
VT(BGMM- - -0)) = - T000T
(21)
00
4 V
15

---

When the number of parameters d equals the number of moment conditions N, To becomes (nonsingular) square matrix, and therefore,
VT(GGMM - -60 N
(22)
which means that the foregoing observation suggests that the selection of WN. has no bearing on the asymptotic variance of the GMM estimator. Consequently, it implies that regardless of the specific method employed to determine WN., provided the moment estimates are asymptotically consistent, WN, serves as the optimal weight matrix, even when dealing with small samples.
The moments utilized for fitting via the generalized method of moments in each step are computed based on the noise network's 's first n-th order. To ensure adherence to the aforementioned proposition, necessary to assume that the n-th order of the noise network converges with probability to Bygaulaa ix.) [diag(e @x-11 +)), details in Appendix B. Consequently, the n-th order moments derived from the noise networks converge with probability to the true moments. Therefore, any choice of weight matrix is optimal.
A.3 Non-Gaussian distribution of transition kernel within large discretization steps
we can apply Bayes' rule to the posterior distribution q( (e++ follows:
((11)1+411) = ((1440//1))9(z)) = g(I: (e exp((o(((()() - log(g(((+<()))
(23)
(((++1))
|1t+Ar - T4 - Ael + log p(a)) log(E(+
α exp
20
where At is the step size, ((3+) is the marginal distribution of de- enough, using Taylor expansion for log p()++ we could obtain:
When F++ and It close
0
log logp(xe+sc) Ro logp(x.) + (I++ xr) V p(gp(())) Al logg(xz),
(24)
01
- 2, 24 -
(f( V log (og,((z))) + 0010)
((zg(e(+Ae) x exp
(25)
2g At
By ignoring the higher order terms, the reverse transition kernel will be Gaussian distribution. However, as At increases, the higher-order terms in the Taylor expansion cannot be disregarded, which causes the reverse transition kernel to deviate from a Gaussian distribution.
Empirically, from Fig. 2 in our paper, we observe that as T decreases, the reverse transition kernel increasingly deviates from a Gaussian distribution. For more details, please refer Appendix D.2.
B Calculation of the first order moment and higher order moments
Suppose the forward process is Gaussian distribution same with the Eq. (1) as g(x)(1))) [[=))(())=41
And let 2 1>>>5 8 0 always satisfy, (<<.(=.) = N Na (ags Se, where azes = az/a, and Sqj. of o2 o. V1 a(s).a, -1- att. It's easy prove that the distribution q(xe 2o). q(Iz/<1 |2,, =O) are also a Gaussian distribution [21]. Therefore, the mean x, under the measure gxx,)x)) would be
Egfs, ln.) [Ira] =
(26)
(xt Bele
EE atls
a)l tc
1
Bele E (e
= (xt
arls
16

---

And for the second order central moment CoVgla, my) [r,), we use the total variance theorem, refer to [2]. and similar with [2], we only consider the diagonal covariance.
(xs) Ex(zalzz LoV (o. [a.] +
(27)
X I + (o Bo
as a.
1 +
Cov CoV (((ool)
o
as L of Cov Covg(nolm))led
A +
of
adjo
X T + (Byleola. (e @ ce] Bygcol ( [e)] @ Bo(oole))))))),
0z
since the higher-order moments are diagonal matrix, we use (Ma to represent the diagonal elements that have the same dimensions the first-order moments, such diag(r, x.)
(z) and diag(r, @ 2, xx) M2 have the same dimensions as 2'
Moreover, for the diagonal elements of the third-order moments, we have Byte, (w.) [diag( (I's & 6 x.2] = E [a) Bo [z, @ O p. we could use the fact that Egt,, m.. (xr, ( ro))
(Ia ((xe, vo)) @ (T. - ((x,, mo )) = 0, therefore,
Ma gls (mx) (diag(x.. u. x. @ x.)) Egtrole) (o. (diag(u, (dagg Eg Beces [3diag(æx 13diagg) @ pl
0x.)]
(28)
3diag(r, p(xx, vu) (a zu)) diag( (agg((((a zo) @ u()s,
0 diag( o. ax) 3\ [[
a md
of
Constant term
3 o. a 22
t/e (diag(æ; @ x4)) + 2atosels
+ E (o.) [mo]]
of
of
Linear term in zto
+ A2la03 of
AgloSals
of j2x, )
2 Eg( diag(rn 8 xo)) + GojoBelo )
(xoooo))) (iaag(r 2g @ xo)),
of
Quadratic term in atp
Cubic term in 2gg
where the third equation is derived from the decomposition of higher-order moments of Gaussian distribution, the fourth equation is obtained by splitting It is noted to highlight that in our study, we only consider the diagonal higher-order moments in our method for computational efficiency since estimating full higher-order moments results in escalated output dimensions (e.g.. quadratic growth for covariance and eubic for the third-order moments) and thus requires substantial computational demand and therefore all outer products can be transformed into corresponding element multiplications and we have:
1
(eg(eal (diag(Io @02 20)]
B [diag((r -o()) 8 (x+ a()e))
(29)
a (t)
B /diag(a, 8² - ===30(()(e @ 22)
(t)
30 (1)(1))) (())) x)))
0² [1
E (B
3 (t) (z, @ Eg(
a²(()
+ 30 (t),, (co o o² d].
Therefore, when we need to calculate the third-order moment, we only need to obtain (u(cal=e)(e) Ee ( @ ei] and Eg @ @ eej. Similarly, when we need to calculate
17

---

the n-order moment, we will use Eetzalzy ing a sharing network and using the MSE loss to estimate the network to obtain the above information about different orders of noise.
(nfralay)(ei)) [e)..... )
",,,,,,tr in, Ie -0---___ ee). Bao et al. [2] put forward us-
The function h ii)= p (8) in Algo. 2 is defined h(f (22, 1). (22, 1), Pa (Jzz, ()) M1 (Ie; 1), M2 (Te, 1), (a(S (1), where M1 (Te, ()) = By (F) in Eq. (26), Covg M2(J (Iz, 1)) = Cov q( (m.) [Px] in Eq. (27) and [3] (Tt, ()) = M3 = (e diag((,, @ x. xx)] in Eq. (28)
B.1 Objective function for Gaussian mixture transition kernel with two components
Recall that the general objective function fit Gaussian mixture transition kernel in each sampling step via GMM is shown in Eq. (30).
No
N Σ gxi, 0)) N
min Q00, M1, Mn) = min 1 N I=1
Σ
g(x,, W
(30)
where AT denotes the transpose matrix of matrix A. In our paper, we propose to use the first three
moments to fit a Gaussian mixture transition kernel p(r, (T) = ( (1) oo)) + NN(,
(2) ,o?) in
each sampling step from F2 to Zg. Therefore, the final objective function follow:
Mi - (5 (1)
(1) + 314 (2)
min m ml
M2 - ( + + =l(m, ())2 of)))
PTT
(31)
P of
MM - (5 + of + [14 (2))3 + 3 (2)
0ff)
M1 (Hpp + (2))
M2 - - [1)} 2 + + o
M3 (51(44 11133 3/4 [
3pd (2) of)),
where to simplify the analysis, we use the scalar form of parameters pi Mi of as representation.
(1) (2)
C Modeling reverse transition kernel via exponential family
Analysis in Sec. 3.1 figures out that modeling reverse transition kernel via Gaussian distribution is no longer sufficient in fast sampling scenarios. In addition to directly proposing the use of Gaussian Mixture for modeling, we also analyze in principle whether there are potentially more suitable distributions i.e., the feasibility of using them for modeling.
We would turn back to analyzing the original objective function of DPMs to find a suitable distribution. The forward process (x.)) = N aT Bris I), consistent with the definition in Appendix B. DPMs' goal is to optimize the modeled reverse process parameters to maximize the variational bound L in Ho et al. [14]. And the ELBO in Ho et al. [14] can be re-written to the following formula:
= (KL((((r))ll(pp)))) + F.g
+ H(ro),
(32)
t21
where a: (2xx) is the true distribution and p. p(xz)) the modeled distribution, and the minimum problem could be transformed into sub-problem, proved in Bao et al. [3]:
min L min DKL (g(x, IPpe,, (x (x:)). (0)
(33)
{o.1)--..
We have no additional information besides when the reverse transition kernel is not Gaussian. But Lemma. C.3 proves that when the reverse transition kernel po,,, (I,, (2:)) is exponential family
Pa, (xa/æx) = p(x,, (.))) = A(xx) exp (z)) - a<(0)))). solving the sub-problem Eq. (33) equals solve the following equations, which is match moments between the modeled distribution and true distribution:
E
= Ep [((xx)))
(34)
18

---

When t(x) = (z),.,"")² solving (34) equals to match the moments of true distribution and modeled distribution.
Meanwhile, Gaussian distribution belongs to the exponential family with f(x) = (x, 2² and 0, of 2n? 2z )" details in Lemma. C.2. Therefore, when modeling the reverse transition kernel as Gaussian distribution, the optimal parameters are that make its first two moments equal to the true first two moments of the real reverse transition kernel (2x) which is consistent with the results in Bao et al. [3] and Bao et al. [2].
The aforementioned discussion serves motivation to acquire higher-order moments and identify corresponding exponential family, which surpasses the Gaussian distribution terms of complexity. However, proposition C.I shows that finding such exponential family distribution with higher-order moments is impossible.
Proposition C.1 (Infeasibility of exponential family with higher-ordet moments.). Given the first n-th order moments. It's non-trividl to find an exponential family distribution for min DKL (gllp) when n. is odd. And it's hard to solve min DRL (gllp) when is even.
C.1 Proof of Proposition C.1
Lemma C.2. (Gaussian Distribution belongs to Exponential Family). Gaussian distribution p(z)
exp
(x-n)² 2
is exponential family with t(x) = (2, and O = ( 2 )7
V2Ke
20
Proof. For simplicity, we only prove one-dimensional Gaussian distribution. We could obtain:
1
(x 15-10 mm)2
p(x)
exp
(35)
V2TO
20
1
(x² -2pc + 122)
exp
V
20
122
exp
(log (2TO²)
11/2)
exp
20 (x.2 -2
σ²)
12
exp
(ggg((2000)) (2TT2) -1/2)
20 (-24 1) 2²)T
exp
exp exp " -
112
20 2x
22/²"
+ 20
log logg (2TT²))
where 0 = (64 2 T and t(r) = (2,
Lemma C.3. (The Solution for Exponential Family in Minimizing the KL Divergence). Suppose that p(x) belongs exponential family p(x, 8) h(z) exp(6 t(x))t((=))))))))())) a(0)), and the solution for minimizing the Eg blg is Ex([x)) = Ep(z.0) [xx))
Proof. An exponential family p(=, .7) h(x)) exp(17 (t(x) a(n) x f(xx 1) h(x))()))))))( exp exp)(xx((x)) (t))
with log-partition (11). And we could obtain its first order condition F.yllog r) as:
Vn log S (2, 11) (log h(x)) 17 ((x)) = ((x)
(36)
Vna(T) V, log (f [(a.))dz)
J (n ()r
(37)
J f(x n)dx
=e- / (xx)
/ (x))(a m)dx Epiz,)) [xx)
19

---

In order to minimize the (ql()) glog(g/p) = -Byllogph, we have:
/ dg + / /d(1n" (q(())((x)=))))))))))))))){)))) a())
Elllogpl pl -
0 on
B Bog /
dyl ②
On (nj² -a(n)a]]]
/ dg(z Ep(a,n) [x)))) = E,,((x)]] Ep(arm) [())] =
[,[[x]]] Fp(x.))((())
For the second-order condition, we have the following:
0²
aa)) Om / dp(x,n)t(x)
0
(38)
0np
/ On
h(x)) - - ((n))
/ Mx))(())) 3 exp(1" a( On
/ (x)) - a(n)) F (x))))
(lz,n) t(z)
pp (x)) Ep [x))] = Coupte. [(x)) 2 0
Therefore, the second-order condition for the cross entropy would be:
0²
B [log pl
(E, [(x]] Ep [t())) ((x)))
(39)
0np²
On
/
On p(x, n)((x)dz
8² 3, a()) -Co [xx) <
When we assume that the reverse process Gaussian, the solution Eq. (33) equals match the moment of true distribution and modeled distribution = [z]. E = Cou, [x]]
Lemma C.4. (Infeasibility of the exponential family with higher-onde order moments). Suppose given the first N-th order moments M, I., and modeled solve the minimum problem Eg [log pl when is odd and it's difficult to solve when N even.
exponential family. It is nontrivial to
Proof. While given the mean, covariance, and skewness of the data distribution, assume that could find an exponential family that minimizes the KL divergence, so that the distribution would satisfy the following form:
I(p, 2) log p(x) h(x))
I(p, 2) = DKL. (glp)) - -x1( pt-m) (m)m= = Op
+ X =0 (40)
= p(x) = 1((x)) exp 2 -1
where, ((x) 22 22, when À not zero, f p oo and density can't 't be normalized. The situation would be the same given an odd-order moment.
h((x) exp (No + Aι + Aga2 + Aaæx)) and f dpx Ma. However,
Similarly, given an equation J dpx* M4 and = h(x) exp do + >1 + >22x + >4 Consider such function:
fourth-order moment, could derive that ) above, and we should solve
z(A) = /
0
dx exp(-x2 - Ax A>> 0
(41)
20

---

When À - O, we could obtain lim Z()) V For other cases, the lambda can be expanded and then integrated term by term, which gives Z(A) En - 1'n + 1/2), but this function However, the radius of convergence of this level is O, so when the À takes other values, we need to propose a reasonable expression for the expansion after the analytic extension. Therefore, for solving the equation f dpzz = M4. there is no analytical solution first, and the numerical solution also brings a large computational effort.
D More information about Fig. and Fig. 2
D.I Experiment in Toy-data
To illustrate the effectiveness of our method, we first compare the results of different solvers on one-dimensional data.
The distribution of our toy-data is g( 0.4.( our solvers in each step as p(z,|zz) = (M
0..3V (..00 and we define (1) of) with vectors p Pi (2) and of. (1) o + (M)
which can not overfit the ground truth.
We then train second and third-order noise networks on the one-dimensional Gaussian mixture whose density is multi- modai. We use a simple MLP neural network with Swish activation [33].
Moreover, we experiment with solvers in 8-Gaussian. The result is shown in Tab. 2. GMS outperforms Extended AnalyticDPM (SN-DDPM) [2] as presented in Tab. 2, with bandwidth of I. where is the standard deviation of data and L is the number of samples.
Table 2: Comparison with SN-DDPM w.t.t Likelihood E,,logg Po t on 8-Gaussian. GMS outperforms SN-DDPM.
8-GAUSSIAN
K
5
10
20
40
SN-DDPM
-0.7885
0.0661
0.0258
0.1083
GMS
-0.6304
0.0035
0.0624
0.1127
D.2 Experiment in Fig. 2
In this section, will provide comprehensive explanation of the procedures involved in computing the discrepancy between two third-order moment calculation methods, as depicted in Fig.
The essence of the calculation lies in the assumption that the reverse transition kernel follows Gaussian distribution. By employing the following equations (considering only the diagonal elements of higher-order moments), we can compute the third-order moment using the first two-order moments:
E [z,,) [ O Ft.-, @ It,-ilo = Mc O μ O μ 3p E,
(42)
where is the first-order moment and Σ is the diagonal elements of second order moment, which can be calculated by the Eq. (26) and Eq. (27). Meanwhile, we can calculate the estimated third-order moment M3 by Eq. (28).
We use the pre-trained noise network from Ho et al. [14] and the second-order noise network form Bao al. [2] and train the third-order noise network in CIFAR10 with the linear noise schedule.
Given that all oigher-order moments possess the same dimension as the first-order moment μ., we can directly compare the disparity between different third-order moment calculation methods using the Mean Squared Error (MSE).
Thus, to quantify the divergence between the reverse transition kernel q(x,(z;) and the Gaussian distribution, we can utilize the following equation:
Daje = log (Ey(z,) (2r 3 x. Ka,G M3
(43)
where M3 is obtained via Fq. (28), and we can start different time step and choose a corresponding to calculate the Daje and draw different time step and step size t and we can derive Fig. 2.
21

---

E Experimental details
E.1 More discussion on weight of Gaussian mixture
From Proposition 3.2, we know that when the number of parameters in the Gaussian mixture equals the number of moment conditions, any choice of weight matrix is optimal. Therefore, we will discuss the choice of parameters to optimize in this section. As we have opted for a Gaussian mixture with two components (z)) wiN(ust (1) 20)) + (1
(2) = as foundational solvers, there exist five parameters (considered scalar, with the vector being analogous) available for optimization.
Our primary focus is optimizing the mean and variance of the two components, as optimizing the weight term would require solving the equation multiple times. Additionally, we have a specific requirement that our Gaussian mixture optimization, particularly when the ground truth corresponds to a Gaussian distribution. In Tab. 3, we show the result of different choices of parameters in the Gaussian mixture.
converge a Gaussian distribution at the conclusion of
Table 3: Results among different parameters in CIFARIO (LS), the number of steps is 50. The weight of Gaussian mixture is W1 = 1/ and W2 = 2/
(1)
M
Enjt
(1)) (
(1) (2(2)
CIFARIO (LS)
4.17
10.12
4.22
When parameter is accompanied by superscript, it implies that both components share the same value for that parameter. On the other hand, if parameter is associated with superscript, and only one moment contains that superscript, it signifies that the other moment directly adopts the true value for that parameter.
It is evident that the optimization of the mean value holds greater significance. Therefore, (2)
subsequent choices for optimization are primarily based on the first of parameters (!!)) + u -.je- Another crucial parameter to consider is the selection of weights W;. In Tab. 4, we show the result while changing the weight of the Gaussian mixture and the set of weight w1 3, 42 Y/2 performs best among different weight.
Table 4: Results among different weight choices CIFARIO (LS), the number of steps 50.
99 100
100
5
3
2.
CIFARIO (LS)
4.63
4.20
4.17
4.26
What's more, we observe that such choice of weights (1. 1) consistently yielded superior performance among different datasets. Identifying an optimal weight value remains a promising direction for further enhancing the GMS.
E.2 Details of pre-trained noise networks
In Tab. 5, we list details of pre-trained noise prediction networks used in experiments.
22

---

Table Details of noise prediction networks used our experiments. LS means the linear schedule of (t) [14] in the forward process of discrete time step (see Eq. (1)). CS means the cosine schedule of a (1) [30] in the forward process discrete timesteps (see Eq. (1)).
TIMESTEPS N
NOISE SCHEDULE OPTIMIZER FOR GMM
CIFARIO (LS) CIFARIO (CS)
1000
LS
ADAN
1000
CS
ADAN
IMAGENET 64X64
4000
CS
ADAN
E.3 Details of the structure of the extra head
In Tab. 6, we list structure details of NN1, NN₂ and NN3 of prediction networks used our experiments.
Table 6: NN1 represents noise prediction networks and NN2, NN₃ represent networks for estimating the second- and the third-order of noise, which used in our experiments. Conv denotes the convolution layer. Res denotes the residual block. None denotes using the original network without additional parameters
NN1
NN₂ (NOISE)
NN; (NoISE)
CIFARIO (LS) CIFARIO (CS) IMAGENET 64x64
NONE
CONV CONV RES+CONV
RES+CONV
NONE
RES+CONV RES+CONV
NONE
E.4 Training Details
We use a similar training setting to the noise prediction network in [30] and [2]. On all datasets, we use the ADAN optimizer [41] with learning rate 10 4; we train 2M iterations total for a higher order of noise network; we use an exponential moving average (EMA) with a rate of 0.9999. We use batch size of 64 on ImageNet 64X64 and 128 on CIFAR1O. We save checkpoint every 50K iterations and select the models with the best FID on 50k generated samples. Training one noise network on CIFARIU 10 takes about 100 hours on one A100. Training on ImageNet 64x64 takes about 150 hours on one A100.
E.5 Details of Parameters of Optimizer in Sampling
In Tab. 7, we list details of the learning rate, learning rate schedule, and warm-up steps for different experiments.
Table 7: Details of Parameters of Optimizer used in our experiments. Ir Schedule means the learning rate schedule. min Ir means the minimum learning rate while using the learning rate schedule, 4 is function with the second order growth function of sampling steps t.
LEARNING RATE
LR SCHEDULE
MIN LR
WARM-UP STEPS
CIFAR10
MAX(0. 00.1000.1)
COS COS
0.1 0.04
18
IMAGENET 64x64
MAX(0. 1.,00.1-,.011.00..1 10.100006
18
23

---

where COS represents the cosine learning rate schedule [6]. We find that the cosine learning rate schedule works best. The cos learning rate could be formulated as follows:
Cti
if i Iw
04+1
(44)
m0x((0.5 cos
I- 7-1
+ 1)0r, Obmin)
else
where, a, is the learning rate after steps, I is the warm-up steps, Omin is the minimum learning rate, I is the total steps.
E.6 Details of memory and time cost
In Tab. 8, we list the memory of models (with the corresponding methods) used in our experiments. The extra memory cost higher-order noise prediction network negligible.
Table 8: Model size (MB) for different models. The model of SNDDPM denotes the model that would predict noise and the square of noise; The model of GMDDPM denotes the model that would predict noise, the square of noise, and the third power of noise.
NOISE PREDICTION NETWORK (ALL BASELINES)
NOISE & SN PREDICTION NETWORKS SNDDPM
NOISE & SN PREDICTION NETWORKS (GMDDPM)
CIFARIO (LS) CIFARIO (CS)
50. I MB 50.11 MB 115.46
50.11 MB
50.52 MB (+0.8%)
50.11 MB
50.52 MB (+0.8%) 116.28 (+0.7%)
IMAGENET 64x64
115.87 MB
In Fig. 5, we present a comprehensive breakdown of the time of optimization process within GMS each sampling step. Subfigure (a) in Fig. 5 illustrates that when the batch size is set 100, the time of optimizing approximately 320 steps to fit a Gaussian mixture transition kernel at each step equivalent to the time needed for one network inference, which means that the additional time of GMS for each sampling steps is about 10% when the number of optimizing steps is 30.
Meanwhile, Subfigure (a) in Fig. 5 elucidates the relation between sampling time and the number of sampling steps for both GMS and SN-DDPM. It is noteworthy that the optimization steps employed in GMS remain fixed at 25 per sampling step, consistent with our setting in experiments. Evidently, as the number of sampling steps escalates, GMS demonstrates a proportional increase in computational overhead, consistently maintaining this overhead within 10% margin of the original computational cost.
Opt imize Steps es Time
Sumpling Steps vs Time
2813 1000
2.5
Time (t wcondd)
800
3.2
305
252
1200
Optimizer Stapps
Samping Steps
(a) Optimizing time and steps
(b) Sampling time and steps
Figure 5: Sampling/optimiznng steps and time (in seconds). (a). The correlation between the number of optimizing steps and the corresponding optimizing time for GMS. Notably, the optimizing. (b). The correlation between the number of sampling steps and the corresponding sampling time for a single batch is observed. Notably, the sampling time demonstrates nearly linear growth with the increase in the number of sampling steps for both solvers.
24

---

Since many parts in the GMS introduce additional computational effort, Fig. 6 provides detailed information on the additional computational time the GMS and SN-DDPM relative to the DDPM, assuming that the inference time of the noise network is unit one.
Emphasizing that the extra time is primarily for reference, most pixels can work well with Gaussian distribution. By applying a threshold and optimizing only when the disparity between the reverse transition kernel and Gaussian exceeds it, we can conserve about 4% of computational resources without compromising result quality.
-0.4%
20. -0.8%%
29%
Inference of noise network (DDPM. DDIM, ADPMs)
Inference of 2-nd noise network Inference of network Solve the parameters for GM (SN-DDPM, NPR-DDPM) (GMS)
using different order momeuts
Figure 6: Time cost distribution for GMS.
E.7 The reduction in the FID on CIFARIO for GMS compared to SN-DDPM
To provide a more intuitive representation of the improvement of GMS compared to SN-DDPM at different sampling steps, have constructed Fig. below. As observed from Fig. 7, GMS exhibits a more pronounced improvement when the number of sampling steps is limited. However, the number of sampling steps increases, the improvement of GMS diminishes, aligning with the hypothesis of our non-Gaussian reverse transition kernel as stated in the main text. However, we respectively clarify that due to the nonlinear nature of the FID metric, the relative/absolute FID improvements are not directly comparable across different numbers of steps
E.8
Additional results with latent diffusion
we conduct an experiment on ImageNet 256 x 256 with the backbone noise network as U-ViT- Huge [4]. We train extra the second-order and the third-order noise prediction heads with two transformer blocks with the frozen backbone. The training iterations for each head is 150K and other training parameters the same the training of the backbone of U-ViT-Huge (details in [4]). We will evaluate the sampling efficiency of GMS under conditional sampling (classifier-fre guidance scale is 0.4) and unconditional sampling in Tab.
GMS vs. SN-DDPM Reduction
Reduction
Relative FID Reduction
Sampling Steps
Figure 7: The reduction in the FID on CIFAR10 for GMS compared to SN-DDPM. With a restricted number of sampling steps, the improvement exhibited by GMS surpasses that achieved with an ample number of sampling steps.
25

---

Table Comparison with DDPM and SN-DDPM w.r.t. FID score on ImageNet 256 256 by using latent diffusion. Uncond denotes the unconditional sampling, and Cond denotes sampling combining 90% conditional sampling with classifier- free guidance equals 0.. and 10% unconditional sampling, consistent with Bao et al. [4].
COND, CFG=0.4
UNCOND
# TIMESTEPS K
15
20
25
40
50
100
200
25
40
50
100
DDPM. 3,
6.48
5.30
4.86
4.42
4.27
3.93
3.32
6.47
5.97
5.04
SN-DDPM
4.40
3.36
3.10
2.99
2.97
2.93
2.82
8.19
5.73
5.32
4.60
GMS (OURS)
4.01
3.07
2.89
2.88
2.85
2.81
2.74
7.78
5.42
5.03
4.45
E.9 Additional results with same calculation cost
Since GMS will cost more computation in the process of fitting the Gaussian mixture, we use the maximum amount of computation required (i.e., additional 10% of computation is needed) for comparison, and for a fair comparison, we let the other solvers take 10% more sampling steps. Our GMS still outperforms existing SDE-based solvers with the same (maximum) computation cost in Tab. 10.
Table 10: Fair comparison with competitive SDE-based solvers w.r.t. FID score J on CIFAR10. SN-DDPM denotes Extended AnalyticDPM from [2]. The number of sampling steps for the GMS indicated within parentheses, while for other solvers, it is represented outside of parentheses.
CIFAR10 10 (LS)
# TIMESTEPS K 11(10)
22(20)
28(25)
44(40)
55(50)
110(100)
220(200) 1000(1000)
SN-DDPM'
17.56
7.74
6.76
4.81
4.23
3.60
3.20
3.65
GMS (OURS)
17.43
7.18
5.96
4.52
4.16
3.26
3.01
2.76
In order to provide more intuitive representation for comparing different methods under the same sampling time, we have generated Fig. 8. Subfig. (a) Fig. illustrates that at each step, there an additional computational time cost incurred for GMS to fit the Gaussian mixture transition kernel. This computational overhead exhibits a nearly linear growth pattern with the increase in the number of sampling steps.
Subfig. (b) Fig. 8 offers a valuable perspective on the connection between approximate sampling time and sampling quality for two GMS and SN-DDPM. It becomes apparent that GMS consistently exhibits superior performance when compared to SN-DDPM with identical computational resources. Furthermore, the data reveals that the magnitude of improvement introduced by GMS is more substantial when the number of sampling steps is limited, as opposed to scenarios with higher number of sampling steps.
26

---

Sampling Steps Time
26-2P4
DP0
1224
Time (in Secon00)
1204
F1D
1P Sampling Steps
(a) Sampling steps and sampling time
(b) sampling time and quality
Figure 8: (a). The correlation between the number of sampling steps and the corresponding sampling
time for single batch is observed. Notably, the sampling time demonstrates nearly linear growth with the increase in the number of sampling steps for both solvers.(b) The relation between the sample quality (in FID) and the sampling time (in seconds) of GMS and SN-DDPM on CIFAR10.)
E.10 Compare with continuous time SDE-based solvers
For completeness, we compare the sampling speed of GMS and non-improved reverse transition kernel in Tab. 11, and it can be seen that within 100 steps, our method outperforms other SDE-based solvers. It is worth noting that the EDM-SDE [18] is based on the prediction network, while SEED [10] and GMS are based Ho al. [14]'s pre-trained model which discrete-time noise network.
Table 11: Comparison with SDE-based solvers [18, 10] w.r.t. FID score 4. on CIFAR10.
CIFAR10
# TIMESTEPS K
10
20
40
50
100
SEED-2 [10]
481.09
481
305.88
51.41
11.10
3.19
SEED-3 3 [10]
483.04
462.61
247.44
62.62
3.53
EDM-SDE [18]
35.07
14.04
9.56
5.12
2.99
GMS (OURS)
17.43
7.18
4.52
4.16
3.26
E.11 Codes and License
In Tab. 12, we list the code we used and the license.
Table 12: codes and license.
URL
CITATION
LICENSE
Ho AL. [14] HTTTS//TTTTTEA.COOOOEEEEIITTTTTT NICHOL AND DHARIWAL [30]
WTFPL MIT
SDEdit
Fig. illustrates one of the comprehensive procedures of SDEdit. Given guided image, SDEdit initially introduces noise at the level of at to the original image za. where denotes the noise schedules of forward process diffusion models. Subsequently. using this noisy image It, and then discretizes it via the reverse SDE to generate the final image. Fig. shows that the choice of to will
27

---

can greatly will greatly affect the realism of sample images. With the increase of to, the similarity between sample images and the real image is decreasing. Hence, apart from conducting quantitative evaluations to assess the fidelity of the generated images, it is also crucial to undertake qualitative evaluations examine the outcomes associated with different levels of fidelity. Taking all factors into comprehensive consideration, we have selected the range of to from 0.3T to 0.5T in our experiments.
Real Image
Stroke
co = 100
to =30
=500
= 700
to =900
Figure 9: to denotes the timestep to noise the stroke
Besides experiments on LSUN x 256, we also carried out an experiment SDEdit on ImageNet 64x64. In Tab. 13, show the FID score for different methods in different to and different sample steps. Our method outperforms other SDE-based solvers as well.
Table 13: Comparison with competitive methods in SDEdit w.r.t. FID score on ImageNet 64x64. ODE-based solver is worse than all SDE-based solvers. With nearly the same computation cost, our GMS outperforms existing methods in most cases.
IMAGENET 64X64,
1200
# K
26(28)
51(55)
101(111)
201(221)
DDPM,
21.37
19.15
18.85
18.15
DDIM
21.87
21.81
21.95
21.90
SN-DDPM
20.76
18.67
17.50
16.88
GMS
20.50
18.37
17.18
16.83
G Samples
From Fig. 10 to Fig. 12, we show generated samples of GMS under different number of steps in CIFAR10 and ImageNet 64x64. Here use K to denote the number of steps for sampling.
28

---

(a) GMS (K = 10)
(b) GMS (K 20)
(c) GMS (K = 25)
(d) GMS (K = 10)
(c) GMS (K 50)
(f) GMS (K- 100)
(g) GMS (K 200)
(h) GMS (K 1000)
Figure 10: Generated samples on CIFARIO (LS)
(a) GMS (K 25)
(b) GMS (K 50)
(c) GMS (K 100)
Figure 11: Generated samples ImageNet 64x64.
(a) GMS 200)
GMS (K 400)
(c) GMS (K 4000)
Figure 12: Generated samples on ImageNet 64x
29


Reviews:
[{"review_id": "PSkHPLde2B", "reviewer": "Reviewer_naDW", "summary": "The paper proposes to weaken the Gaussian assumption of the transition probability in the reverse SDE used in deep diffusion models. They first illustrates how and when the Gaussian assumption is wrong. Then they suggest to approximate the non-Gaussian transition probability by a Gaussian Mixture which is adjusted with the method of moments. Finally, they illustrates the superiority of their methods on different standard datasets (CIFAR, ImageNet64) using standard metric (FID). ", "questions": "I did not fully understand why the method of moments is preferred. To me, the issue of estimating a mixture model that depends on both s and t also apply to the estimate of moment. What is the advantage of using the method of moments ? \nEM algorithm is easy to implement and I feel like you might spare computational resources compared to the calculation of higher-order moment. Have you try EM ?\n\nFig 1 : You should plot  the log density to make more visible the differences between the two curves.  ", "limitations": "ok", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The manuscript is well-written and the problem is clearly stated and illustrated.\n- The proposed method allows to reduce the number of time step \n- The results are superior in terms of sample quality (measured by FID).\n\nSpecifically appreciated :\n- The authors account for the computational cost of their method in the supp material.\n- The authors are transparent about real-time applicability of their methods. ", "weaknesses": "- It's unclear what is the interest of the proposed method beyond mathematical and empirical curiosity.\n- The proposed method weakens an assumption at an extra-computational cost. When accounting for this extra-cost improvements are lowered.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your supportive review and valuable suggestions.\n\n\n***Weakness 1: The interest of the proposed method***\n\nWe understand that your concern may relate to the theoretical contributions and practical benefits of our method. (We acknowledge the possibility of potential misunderstanding and welcome any further questions.)\n\nWe find that the Gaussian assumption in the reverse transition kernel can be violated, as shown by our theory and experiments (Fig. 1). Building on this, we introduce the GMS. This new approach improves the efficiency of SDE-based solvers.\n\n\nSDE-based solvers have notable advantages over ODE-based solvers in various downstream applications. Noteworthy examples include stroke-based synthesis, expounded upon in Section 4.2, and image translation, as detailed in [a]. These instances underscore the superior efficacy of SDE-based solvers. Concurrently, when an ample quantity of sampling steps is employed, the performance of SDE-based solvers exhibits superior outcomes in both unconditional and conditional sampling scenarios (as demonstrated in [b]).\n\nOur experiments in Fig. A of the rebuttal PDF show that **GMS achieves superior performacne than competitive SDE-solvers under the same computational budget**. Therefore, we believe that GMS is promising in applications investigated in [a,b].\n\n\n***Weakness 2: Extra-cost during inference***\n\nOur GMS achieves superior performance under the same computational budget. To provide an intuitive comparison, we present the CIFAR10 FID curve plotted against sampling time (in seconds) in Fig. A of the rebuttal PDF. Notably, a consistent improvement margin over SN-DDPM emerges within the same computational budget (sampling time in seconds). Specifically, we observe an FID improvement of \\~0.9 (\\~0.6, \\~0.3, resp.) at sampling time of 40s (60s, 200s, resp.). This observation corroborates our claim that the Gaussian reverse transition kernel deviates when employing fewer discretization steps. We will clarify this in the revision.\n\n\n***Question 1: Advantages of the method of moments compared to the EM algorithm***\n\nThe EM algorithm is not well suited for our method due to the following reasons:\n1. Nontrivial training and loss modifications: Learning the reverse transition kernel $p(x\\_s|x\\_t)=\\\\sum\\_{i=1}\\^{M}w\\_i \\\\mathcal{N}(x\\_s|{\\\\mu\\_i}(x\\_t),{\\\\Sigma\\_i}(x\\_t))$ ($\\\\sum\\_iw\\_i=1$) with EM algorithm requires alternatively performing expectation and maximization steps, where the latter are not differentiable. Moreover, it requires sampling time step pairs $(s, t)$ and many-to-one pairs of $x\\_s$ and $x\\_t$, incurring computational costs. In contrast, our approach seamlessly expands with the noise prediction loss within the diffusion framework and maintains (training) efficiency.\n2. Substantial model architecture changes: EM necessitates nontrivial architectural changes for handling double time inputs $(s, t)$ in modeling $p(x\\_s|x\\_t)$. This poses higher model capacity demands and complex design challenges.\n\nWe will include these in the revision.\n\n\n***Question 2: Re-plot Figure 1***\n\nThank you for your valuable suggestion. We will update it in the final revision.\n\n\n***References***\n\n[a] Zhao et al. Egsde: Unpaired image-to-image translation via energy-guided stochastic differential equations. 2022\\\n[b] Karras et al. Elucidating the Design Space of Diffusion-based Generative Models. 2022"}, {"review_id": "TSIQL76XKA", "reviewer": "Reviewer_iBkM", "summary": "Sampling from diffusion models is equivalent to solving the reverse diffusion SDEs or the corresponding probability flow ODEs. In comparison, SDE-based solvers can generate samples of higher quality and are suited for image translation tasks. However, during inference, existing SDE-based solvers are severely constrained by the efficiency-effectiveness dilemma. To overcome this limitation, this paper introduces a novel class of SDE-based solvers called Gaussian Mixture Solvers (GMS) for diffusion models. Experimental results validate the motivation and effectiveness of GMS solvers.\n\n\n---Post-rebuttal:\nThe authors address some of my concerns, However, I strongly enough the authors to conduct more experiments. I would like to keep my score.", "questions": "Please refer to the weakness", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "a.\tThis paper systematically examines the assumption of the Gaussian transition kernel and reveal that it can be easily violated under a limited number of discretization steps even in the case of simple mixture data. To this end, the authors propose a new type of SDE-based solver called Gaussian Mixture Solvers.\nb.\tThis paper presents an approach for estimating the high-order moments utilizing noise networks.\n", "weaknesses": "a. In Sec. 3.2, these authors claim that a designed Gaussian mixture model can degenerate to a Gaussian, however, it is not clear how to design such a Gaussian mixture model.\nb. As the weight in Eq. 7 is manually setting, how to select an optimal set of weights for a specified testing benchmark?\nc. In Table 1, except for the FID score, extra computational cost and running time should also be compared to verify the effectiveness of the proposed GMS solvers. \nd. From Table 1, we can observe that the improvement of GMS is not so obvious when compared with SN-DDPM. Especially from Table 9 in the supplemental material, we see that with the same computation cost, the gap between GMS and SN-DDPM is small.\ne. In the experiments, the resolution of the dataset is small, how about the improvement of GMS solvers on data with larger scale?\nf. The experimental section is not enough.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your supportive review and valuable comments.\n\n\n***Weakness (a): Design of the Gaussian mixture model & How can it degrade to a Gaussian***\n\nOur design choice of the Gaussian mixture model for the reverse transition kernel is $p(x\\_s|x\\_t)=\\\\frac{1}{3} \\\\mathcal{N}(\\\\mu\\^{(1)}\\_t(x\\_t),\\\\sigma\\^2\\_t(x\\_t))+\\\\frac{2}{3} \\\\mathcal{N}(\\\\mu^{(2)}\\_t(x\\_t),\\\\sigma\\^2\\_t(x\\_t))$ (as described in L182 of Sec. 3.2), where $\\\\mu\\^{(1)}\\_t(x\\_t)$, $\\\\mu\\^{(2)}\\_t(x\\_t)$ and $\\\\sigma\\^2\\_t(x\\_t)$ are the three parameters to be optimized given the first three-order moments. It can degenerate to a Gaussian when $\\\\mu\\^{(1)}\\_t(x\\_t)=\\\\mu\\^{(2)}\\_t(x\\_t)$, in which case we have $p(x\\_s|x\\_t)=\\\\mathcal{N}(\\\\mu\\^{(1)}\\_t(x\\_t),\\\\sigma\\^2\\_t(x\\_t))$. We will clarify in the revision.\n\n\n***Weakness (b): How to select an optimal set of weights?***\n\nFor a Gaussian mixture with two components (i.e., our design choice), determining optimal component weights and parameters from only the first three-order moments is underdetermined. To tackle this, we pragmatically set the weights to ($\\\\frac{1}{3}$, $\\\\frac{2}{3}$) and observe that such choice consistently yielded superior performance among different datasets. Identifying an optimal weight value remains a promising direction for further enhancing the GMS. We will clarify in the revision.\n\n\n***Weakness (c & d): Computational cost and running time should also be compared & The improvement when considering the extra computation time***\n\nThank you for your suggestions. As our GMS enhances sample quality at the expense of additional computation, we provide further insights by presenting the CIFAR10 FID curve plotted against sampling time (in seconds) in Fig. A of the rebuttal PDF, which should be more intuitive than Table 9. Notably, a consistent and demonstrable improvement over SN-DDPM emerges within the same computational budget (sampling time in seconds), showing the superiority of our method. Specifically, we observe an FID improvement of \\~0.9 (\\~0.6, \\~0.3, resp.) at sampling time of 40s (60s, 200s, resp.). We will incorporate this figure adjacent to Table 1 in the revision.\n\nOther analysis and results pertaining to computational cost and running time can be found in the Appendix. E.6. Specifically, Fig. 5 shows that GMS incurs approximately 10% higher computational time per step compared to SN-DDPM, while Fig. 6 offers a breakdown of time allocation across various components.\n\n\n***Weakness (e): Experiments on large-scale data***\n\nGiven the limited time available for rebuttal, we were unable to finish experiments on larger images. Nonetheless, we are actively engaged in experiments on ImageNet 256*256, and we are committed to delivering our best efforts in this regard.\n\n\n***Weakness (f): The experimental section is not enough.***\n\nWe have supplemented the following experiments and results to better demonstrate the effectiveness of GMS:\n- we present the CIFAR10 FID curve plotted against sampling time (in seconds) for SN-DDPM and GMS in Fig. A of the rebuttal PDF.\n- we present the sampling time (in seconds) curve plotted against the NFE (Number of Function Evaluation) for SN-DDPM and GMS in Fig. C of the rebuttal PDF.\n- we compare GMS with more SDE-based solvers such as EDM [a] and SEED [b]. Results show that GMS consistently outperforms within 50 NFEs. Please see details in response to ***Main Weakness 2b*** of ***Reviewer iSEj***.\n\n\nWe are open to conducting additional experiments if they would contribute to a more comprehensive evaluation. Thank you for your feedback.\n\n\n***References***\n\n[a] Karras et al. Elucidating the Design Space of Diffusion-based Generative Models. 2022\\\n[b] Gonzalez et al. SEEDS: Exponential SDE Solvers for Fast High-Quality Sampling from Diffusion Models. 2023"}, {"review_id": "aJgevvC4fP", "reviewer": "Reviewer_iDhp", "summary": "The authors address the efficiency-effectiveness dilemma faced by existing SDE-based solvers in diffusion models during inference. They observe that the Gaussian assumption in the reverse transition kernel is frequently violated, even with a limited number of discretization steps. To overcome this limitation, the authors propose a new class of SDE-based solvers called Gaussian Mixture Solvers (GMS). In this approach, they estimate the first three-order moments and optimize the parameters of a Gaussian mixture transition kernel using generalized methods of moments. They present empirical results that demonstrate that GMS outperforms other SDE-based solvers in terms of sample quality for image generation and stroke-based synthesis in various diffusion models. ", "questions": "1. Can you provide further clarification regarding the claim made in the abstract and various sections of the text that a limited number of discretization steps amplifies the violation of assumptions for the Gaussian transition kernel used in SN-DDPM? How does this relate to the performance of GMS compared to SN-DDPM when considering a small number (e.g., 100) versus a larger number (e.g., 1000) of discretization steps?\n\n2. In Table 1 and Figure 3, it is apparent that the relative (or absolute) improvement of GMS over SN-DDPM is not significantly different when comparing a small number of discretization steps to a larger number. Can you address this discrepancy and explain why the expected noticeable improvement in GMS performance for a smaller number of discretization steps is not observed?\n\n3. The authors mention that GMS is superior to SN-DDPM when considering computational cost (line 273). Could you provide more clarity on this claim and its implications? If GMS does indeed offer computational advantages, please elaborate on this aspect in the text and present corresponding results. Conversely, if this claim is not supported by evidence, please clearly address this discrepancy.\n\n4. Considering the theoretical value of this work, a crucial question arises: Why would one choose to utilize GMS if comparable results can be obtained within the same computational budget using a simpler method like SN-DDPM? Please provide a compelling rationale for using GMS, taking into account its potential advantages and drawbacks compared to alternative approaches.", "limitations": "The authors emphasize that a limited number of discretization steps exacerbates the violation of assumptions for the Gaussian transition kernel in SN-DDPM, as stated in the abstract and various sections of the text (e.g., line 39, 50, 117, and 290). However, it is surprising to note that the expected significant improvement of GMS over SN-DDPM for a smaller number (e.g., 100) versus a larger number (e.g., 1000) of discretization steps is not observed, as evident from table 1 and figure 3. This discrepancy requires clarification.\n\nAdditionally, the authors vaguely assert that GMS outperforms SN-DDPM when considering computational cost (line 273). If this claim holds true, it is crucial to provide clear explanations and present corresponding evidence in the text and results. Conversely, if this claim is not substantiated, it should be explicitly addressed. Although this work holds theoretical value, the fundamental question remains: What is the justification for using GMS if comparable results can be achieved within the same computational budget using a simpler method like SN-DDPM?", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The authors provide a solid theoretical foundation for their approach, highlighting the discrepancy between the empirical data and the assumptions made for the Gaussian transition kernel commonly used in models like SN-DDPM. They tackle this challenge by proposing a novel and effective solution tailored to address this specific issue. Their innovative approach not only takes into account theoretical benefits but also considers computational constraints.", "weaknesses": "Despite introducing GMS as a potential resolution to the efficiency-effectiveness dilemma, the authors fail to provide compelling evidence in their presentation to substantiate their assertion.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments and questions.\n\n\n***Question 1: Clarification regarding the claim that a limited number of discretization steps amplifies the violation of assumptions for the Gaussian transition kernel used in SN-DDPM***\n\nThanks for the insightful comment. We clarify this issue and confirm our claim from both theoretical and empirical standpoints.\n\nTheoretically, we apply Bayes' rule to the posterior distribution $q(x_t|x_{t\\~+\\~\\Delta~ t})$ as follows:\n$q(x_t|x_{t\\~+\\~\\Delta~ t}) = \\dfrac{q(x_{t\\~+\\~\\Delta~ t}|x_{t})q(x_t)}{q(x_{t\\~+\\~\\Delta~ t})} =q(x_{t\\~+\\~\\Delta~ t}|x_{t})\\exp(\\log (q(x_t))-\\log(q(x_{t\\~+\\~\\Delta~ t})))$\n$\\propto \\exp(-\\dfrac{\\left \\\\| x_{t\\~+\\~\\Delta~ t}-x_t-f_t(x_t)\\Delta~ t \\right \\\\|^2 }{2g_t^2\\Delta~ t} +\\log p(x_t)-\\log(x_{t\\~+\\~\\Delta~ t})),$\nwhere $\\Delta~ t$ is the step size, $q(x_t)$ is the marginal distribution of $x_t$. When $x_{t\\~+\\~\\Delta~ t}$ and $x_{t}$ are close enough, using Taylor expansion for $\\log p(x_{t\\~+\\~\\Delta~ t})$, we could obtain:\n$\\log p(x_{t\\~+\\~\\Delta~ t}) \\approx \\log p(x_t)+(x_{t\\~+\\~\\Delta~ t}-x_t)\\nabla_{x_t} \\log p(x_t)+\\Delta~ t \\dfrac{\\partial}{\\partial t} \\log p(x_t)$, \n\n$q(x_t|x_{t\\~+\\~\\Delta~ t})\\propto \\exp(-\\dfrac{\\left \\\\| x_{t\\~+\\~\\Delta~ t}-x_t-[f_t(x_t)-g_t^2\\nabla_{x_t} \\log p(x_t)]\\Delta~ t \\right \\\\|^2 }{2g_t^2\\Delta~ t} +O(\\Delta~ t))$. \nBy ignoring the higher order terms, the reverse transition kernel will be Gaussian distribution. However, as $\\Delta~ t$ increases, the higher-order terms in the Taylor expansion cannot be disregarded, which causes the reverse transition kernel to deviate from a Gaussian distribution. We will include these clarification in the revision.\n\nEmpirically, from Fig. 2 in our paper, we observe that as the number of sampling steps decreases, the backward transition kernel increasingly deviates from a Gaussian distribution. For more details, please refer to Appendix D.2 in our paper.\n\nPlease see our response to ***Question 2*** for analysis of the emprical performance in terms of FID.\n\n\n***Question 2: About the improvement in a limited number of steps***\n\nAs shown in Tab. 1 of the submission (we also plot a curve in Fig. B of the rebuttal PDF for a clearer illustration), GMS achieves an FID improvement of >4.0 with a small number of steps (e.g. 10) and <0.5 with a large number of steps (e.g. 200). \n\nHowever, we respectively clarify that due to the nonlinear nature of the FID metric, the relative/absolute FID improvements are not directly comparable across different number of steps. We will make it clearer in the final version.\n\n\n***Question 3: About the claims that GMS is superior to SN-DDPM when considering computational cost***\n\nThe intended statement here is: \"(In the SDEdit task,) GMS exhibits superior performance compared to SN-DDPM **under the same** computational cost.\" We indeed observed a performance-time curve similar to that of the unconditional generation experiments (Fig. A in the rebuttal PDF), providing evidence for our intended claim.\n\nThank you for pointing this out, and we will correct the original inaccurate claim.\n\n\n***Question 4: Why choosing GMS if comparable results can be obtained within the same computational budget using a simpler method like SN-DDPM***\n\nWe would like to clarify that GMS can achieve superior performance under the same computational budget. To provide an intuitive comparison, we present the CIFAR-10 FID curve against the computation time in Fig. A of the rebuttal PDF, where GMS consisently improves SN-DDPM when varying the computational budgets.\n\nBesides, we provide more results compared to other advanced SDE-based sovlers including EDM [b] and SEED [c], GMS still outperforms within limited number of discretization steps. See details in response to ***Main Weakness 2b*** of Reviewer iSEj.\n\n\n***References***\n\n[a] Song et al. Score-based generative modeling through stochastic differential equations. 2020\\\n[b] Karras et al. Elucidating the Design Space of Diffusion-based Generative Models. 2022\\\n[c] Gonzalez et al. SEEDS: Exponential SDE Solvers for Fast High-Quality Sampling from Diffusion Models. 2023"}, {"review_id": "0ZIWvuFlCW", "reviewer": "Reviewer_iSEj", "summary": "The authors point out that $q(x_s|x_t)$ is not necessarily Gaussian when $t$ is significantly bigger than $s$, and propose to use a mixture of Gaussians in order to model better the reverse process, when the number of integration steps is not large. Such a selection guarantees that with an increasing number of steps, the mixture of Gaussians will subsequently adopt the form of the Gaussian. In order to estimate the parameters of the Gaussians constituting the mixture, the authors estimate the diagonals of higher order moments of the mixture through higher order denoising and integrate them in the generalized method of moments in order to estimate the desired parameters. ", "questions": "What are the advantages of GMS in the light of recent methods such as consistency models [2]?\n\nHave the authors performed experiments to measure how the sampling time per step increases when the number of Gaussians in the mixture (number of moments estimated) increases?\n\n[2] Song et al, Consistency Models, ICML 2023\n", "limitations": "While limitations are presented, the authors could elaborate more on this topic. ", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "The paper is well motivated, as $q(x_s|x_t)$ is not necessarily Gaussian when $t$ is significantly bigger than $s$. \n\nThe experimental results show some improvements when the same number of steps are used. ", "weaknesses": "**Main Weaknesses:**\n\n1) The paper's presentation is not apt for publication.\n\nThe manuscript contains several sentences that are notably difficult to interpret. For instance, the authors' intent remains unclear in the sentence beginning with 'These modeling...' in line 158 and concluding on 160. This aspect only gains clarity in Equation 9, which is located in a different section.\n\nA multitude of vital elements are left undefined. For instance, the parameters $\\theta$ and $\\theta^*$ are not properly concretized in the Gaussian Mixture context, and furthermore, the function $h$ in algorithm 2 lacks a definition in the main manuscript and Appendix B.\n\nThe paper's structure is counter-intuitive. In Algorithm 2, the authors initially compute the moments of the Gaussian mixture (a), and subsequently utilize them to calculate the mixture's defining parameters (b). Contrarily, in Section 3, the authors commence with (b) and proceed to (a). Since this order is not natural it makes the paper difficult to read, and the situation is exacerbated since at this point in a) all the definitions of parameters and moments are general and do not relate to the problem at hand. The connection between these general definitions and the specific problem tackled in the paper is never explicitly stated. It is important to show the connection between the parameters $\\theta$, moments $M_n^{(GM)}(\\theta)$, $M_n(x_i)$ on one hand, and $\\pi_k, \\mu_k, \\sigma_k, \\hat{M}_n$ on the other. This would enable writing the loss in Equation 8 explicitly, which is not done. \n\n2) Sampling times (in seconds) are not provided.\n\n Considering that the premise of paper is enabling generation in fewer steps (that is, speeding up the sampling process), the sampling time curves of the sampling process for different number of steps should be provided in the main paper (Close to Table 1), for both the proposed method as well as the benchmarks. That is, a figure such as Figure 5 in Appendix E6 should be extended for numbers of steps above 50, and placed in the main paper. The values of the curves presented should be in seconds, instead of presenting the ratio between methods. Furthermore, it would be interesting to compare against the stochastic sampler in [1].  \n\n**Minor Weaknesses:**\n\n1) Equation 9 is not derived in Appendix B.\n\n2) Using a mixture of only two Gaussians is a bit underwhelming considering all the material.\n\n3) The authors state that the \"higher-order correlation between two pixels is assumed to be zero\". While it is true that this premise aligns with the work of Bao et al., the authors should justify this in the paper. \n\n4) The Figure in the introduction is misplaced. It does not elucidate the method. It presents results of the performance of the method, hence it should either go to the Experimental section or to the Appendix.\n\nI am willing to increase my score if the authors extend Figure 5 (e.g. 100 200 1000 as in Table 1), and are willing to follow the presentation advices provided in Main Concern 1. Also there are grammar/spelling mistakes in the paper and the authors should ensure they are corrected.\n\n[1] Karraset al, Elucidating the Design Space of Diffusion-Based Generative Models, NeurIPS 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your supportive review and suggestions.\n\n***Main Weakness 1: Presentation***\n\nWe appreciate your careful reading of our paper and your helpful suggestions. We will address these issues in the final version, including but not limited to:\n\n1. We will revise the sentence in line 158 into \"With such a network, the moments under the $q(x\\_s|x\\_t)$ measure can be decomposed into moments under the $q(x\\_0|x\\_t)$ measure, so that sampling any $x\\_t$ to $x\\_s$ requires only a network whose inputs are $x\\_t$ and $t$.\" to make it clearer.\n2. The variable $\\\\theta\\^*$ means the potential optimal parameters for Gaussian mixture models given the first-to-third order of moments and the $\\theta$ represents the variables required to optimize, which are denoted as $\\mu^{(1)}_t$, $\\mu^{(2)}_t$ and $\\sigma^2_t$ in backward transition kernel $p(x_s|x_t) = \\frac{1}{3}\\mathcal{N}({\\mu^{(1)}_t},{\\sigma^2_t})+ \\frac{2}{3}\\mathcal{N}({\\mu^{(2)}_t},{\\sigma^2_t})$.\n3. The function $h$ is defined as $h(f\\^3\\_{[1]}(x\\_t,t),f\\^3\\_{[2]}(x\\_t,t),f\\^3\\_{[3]}(x\\_t,t))=M\\_1(f\\^3\\_{[1]}(x\\_t,t)),M\\_2(f\\^3\\_{[2]}(x\\_t,t)),M\\_3(f\\^3\\_{[3]}(x\\_t,t))$, where $M\\_1(f)=E[x]$ in Eq.(22), $M\\_2(f)=Cov(x)$ in Eq.(23) and $M\\_3(f)=Ske(x)$ in Eq.(24). \n4. We will re-organize the paper, explicitly state the connection between these general definitions and the specific problem tackled and the loss in Equation 8.\n\nWe will also thoroughly correct the grammar/spelling mistakes in our paper.\n\n***Main Weakness 2a: Sampling times (in seconds) are not provided***\n\nWe have included the sampling time curves of the sampling process for different numbers of steps in Fig. C in the rebuttal PDF. Under the same sampling steps, our GMS incurs approximately 10% additional time compared to SN-DDPM due to extra optimization for Gaussian mixture solving. This optimization, however, contributes to improved quality, as shown in Table 1. Besides, we show the optimizing time curve which is required in GMS by changing the batch size and optimizing steps in Fig. C.\n\nFor a more intuitive comparison, we also plot the FID curve w.r.t. the sampling time (in seconds) in Fig. A of the rebuttal PDF. Across the same computational budget (sampling time in seconds), we consistently observe improved performance over SN-DDPM.\n\nWe will include these figures and results in the revision.\n\n\n***Main Weakness 2b: Comparison with more baselines***\n\nFollowing your suggestions, we compare GMS with two additional SDE-based solvers: EDM [a] and SEED [b], with the latter being a higher-order SDE solver. We note that EDM in [a] use $x\\_0$ prediction networks, which differ from ours and are thus not directly comparable. The symbol $\\^*$ indicates network distinctions from the noise network in discrete time.\n\n| Method/NFE | 10 | 20 | 40 | 50 | 100 |\n| :---------------: | :---: |  :----:  | :----: |  :----: |  :----: | \n| GMS | 17.43 | 7.18 | 4.52 | 4.16  |  3.26 |  3.01 | 2.76 |\n| EDM[a]$\\^*$ | 35.07 | 14.04 | 9.56 | 5.12 | 2.99  |   | \n| SEED-2[b] | 481.09 | 305.88 | 51.41 | 11.10 |  3.19 |\n| SEED-3[b] | 483.04 | 462.61 | 247.44| 62.62 | 3.53 | 3.08 |\n\nWe observe that our GMS consistently outperforms within 50 NFEs. We will add these result in the revision.\n\n***Minor Weakness 1: Derivation about Eq. (9)***\n\nThanks for the comments. We will derive Eq. (9) step by step in Appendix B in the final version.\n\n\n***Minor Weakness 2: Reasons for using a mixture of only two Gaussians***\n\nIt's worth noting that in GMS, we use a mixture of two Gaussians (i.e., a two-mode distribution) to effectively capture the reverse transition kernels per sampling step. Intuitively, this choice can potentially encompass exponential modes across the entire trajectory. Empirically, we consistently observe that employing a mixture of two Gaussians yields favorable results across all settings in our experiments.\n\n\n***Minor Weakness 3: Assumption of zero higher-order correlation***\n\nEstimating full higher-order moments results escalated output dimensions (e.g., quadratic growth for covariance and cubic for the third-order moments) and thus requires substantial computational demands. We therefore consider the diagonal higher-order moments in our method for computational efficiency, similar to Bao et al. [c]. We will clarify this in the revision.\n\n\n\n***Minor Weakness 4: The placement and the function of Figure 1***\n\nThe primary purpose of including Fig. 1 in the introduction was to provide readers with a visual insight into the motivation behind our method. The figure showcases the diminishing effectiveness of Gaussian reverse transition kernels (SN-DDPM) when using fewer discretization steps, which motivates the choice of non-Gaussian transition kernel in this paper. We are also open to further discussion on this issue and willing to re-organize it for better presentation.\n\n\n***Question 1: Consistency models***\n\nConsistency models [d] focus on distilled ODE-based solvers, while our work is an SDE-based solver. Notably, SDE-based solvers offer distinct advantages in various downstream tasks, including stroke-based synthesis, image translation, and image manipulation, as outlined in Sec. 1. We will discuss the relation with [d] in the final version.\n\n\n***Question 2: How does the sampling time per step increase when the number of Gaussians in the mixture increases?***\n\nIn our experiments, we did not try increasing the number of Gaussians in the mixture (as well as the number of moments estimated) due to the reason mentioned in our response to ***Minor Weakness 2***.\n\n***References***\n\n[a] Karras et al. Elucidating the Design Space of Diffusion-based Generative Models. 2022\\\n[b] Gonzalez et al. SEEDS: Exponential SDE Solvers for Fast High-Quality Sampling from Diffusion Models. 2023\\\n[c] Bao et al. Estimating the Optimal Covariance with Imperfect Mean in Diffusion Probabilistic Models. 2022\\\n[d] Song et al. Consistency Models. 2023"}]

Author rebuttal:
We thank all reviewers for their valuable and constructive feedback, and we have responded to each reviewer individually. We have also uploaded a rebuttal PDF that includes:
- **Fig. A**: The relation between the sample quality (in FID) and the sampling time (in seconds) of GMS and SN-DDPM on CIFAR10. 
- **Fig. B**: The reduction in the FID on CIFAR10 for GMS compared to SN-DDPM when sampling with different number of steps.
- **Fig. C**: The relation between the sampling/optimizing times (in seconds) and the sampling/optimizing steps of GMS on CIFAR10.

In addition, we have extended our experimental comparisons to include more baselines, such as EDM from [a] and SEED from [b].

***References***

[a] Karras et al. Elucidating the Design Space of Diffusion-based Generative Models. 2022\
[b] Gonzalez et al. SEEDS: Exponential SDE Solvers for Fast High-Quality Sampling from Diffusion Models. 2023

Meta-review: The reviewers appreciated the methodological contribution addressing the issue with Gaussian transition kernel in diffusion models, while emphasising the clarity of the presentation and the thorough discussion of computational contraints. Questions addressed several aspects of this work, among which the need for clarifications on methodological and computational aspects, while requiring additional experimental evidence.

The rebuttal and subsequent discussion were successful in addressing these questions and, as result, reviewers either confirmed or increased their score towards acceptance. Overall, the paper presents a relevant contribution to the conference.
Decision: Accept (poster)