Title: Gaussian Mixture Solvers for Diffusion Models
Venue: NeurIPS 2023
Accepted: yes
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

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

1 Introduction

In recent years, deep generative models and especially (score-based) diffusion models [35, 14, 37, 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, 5, 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, SDE-based solvers usually perform better when given a sufficient number of discretization steps [18]. Indeed, a recent empirical study [24] suggests that SDE-based 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, SDEbased solvers are preferable in many downstream tasks such as stroke-based synthesis [27], image translation [45], and image manipulation [20].

(a)

#

Step=5

(b)

#

Step=10

(c)

#

Step=20

3.5

3.5

3.5

Ground Truth

Ground Truth

Ground Truth

3.0

SN-DDPM

3.0

SN-DDPM

3.0

SN-DDPM

GMS

GMS

GMS

2.5

2.5

2.5

2.0

2.0

2.0

1.5

1.5

1.5

1.0

1.0

1.0

0.5

0.5

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

1.0

-0.5

0.0

0.5

1.0

Despite their wide applications, SDE-based solvers face a 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 in the reverse process instead of using a 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 a few steps among other SDE-based solvers. Notably, these solvers all assume that the transition kernel in the reverse process is Gaussian.

In this paper, we systematically examine 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, we propose a new type of SDE-based solver called Gaussian Mixture Solver (GMS) , which employs a more expressive Gaussian mixture transition kernel in the reverse process for better approximation given a limited number of steps (see visualization results on toy data in Fig. 1). In particular, we first learn a noise prediction network with multiple heads that estimate the high-order moments of the true reverse transition kernel respectively. For sampling, we fit a Gaussian mixture transition kernel in 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, 2, 17, 36] in terms of several widely used metrics (particularly sample quality measured by FID) in the tasks of generation and stroke-based synthesis on multiple datasets. Our results show that GMS outperforms state-of-the-art SDE-based 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.44 over the SOTA SDE-based solver [2] given 10 steps on CIFAR10. 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 L 2 distance). Such empirical findings validate the motivation and effectiveness of GMS.

2 Background

In this section, we provide a 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 a forward diffusion process and then learn to reverse such process to recover the data distribution. Formally, let x 0 ∈ R n be a random variable with unknown data distribution q ( x 0 ) . Diffusion models define the forward process { x t } t ∈ [0 , 1] indexed by time t , which perturbs the data by adding Gaussian noise to x 0 with

In general, the function a ( t ) and σ ( t ) are selected so that the logarithmic signal-to-noise ratio log a 2 ( t ) σ 2 ( t ) decreases monotonically with time t , causing the data to diffuse towards random Gaussian noise [21]. Furthermore, it has been demonstrated by Kingma et al. [21] that the following SDE shares an identical transition distribution q t | 0 ( x t | x 0 ) with Eq. (1):

where ω ∈ R n is a standard Wiener process and

Let q ( x t ) 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]:

where ¯ ω ∈ R n is a reverse-time standard Wiener process. The only unknown term in Eq. (4) is the score function ∇ x t log q ( x t ) . To estimate it, existing works [14, 37, 18] train a noise network ϵ θ ( x t , t ) to obtain a scaled score function σ ( t ) ∇ x t log q ( x t ) using denoising score matching [38]:

where w ( t ) is a weighting function, q ( ϵ ) is standard Gaussian distribution and x t ∼ q ( x t | x 0 ) follows Eq. (1). The optimal solution of the optimization objective Eq. (5) is ϵ θ ( x t , t ) = -σ ( t ) ∇ x t log q ( x t ) .

Hence, samples can be obtained by initiating the process with a standard Gaussian variable x 1 , then substituting ∇ x t log q ( x t ) with -ϵ θ ( x t ,t ) σ ( t ) and discretizing reverse SDE Eq. (4) from t = 1 to t = 0 to generate x 0 .

2.2 SDE-based 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 ( x 0:1 ) = p ( x 1 ) ∏ t i -1 ,t i ∈ S t p ( x t i -1 | x t i ) with its trajectory S t = [0 , t 1 , t 2 , ..., t i , .., 1] . Song et al. [37] proves that the conventional ancestral sampling technique used in the DPMs [14] that models p ( x t i -1 | x t i ) as a Gaussian distribution, can be perceived as a first-order solver for the reverse SDE in Eq. (4). Bao et al. [3] finds that the optimal variance of p ( x t i -1 | x t i ) ∼ N ( x t i -1 | µ t i -1 | t i , Σ t i -1 | t i ( x t i )) is

where, γ t i = √ α ( t i -1 ) -√ σ 2 ( t i -1 ) -λ 2 t i √ α ( t i ) σ 2 ( t i ) , λ 2 t i is the variance of q ( x t i -1 | x t i , x 0 ) . AnalyticDPM [3] offers a significant reduction in discretization error during sampling and achieves faster convergence with fewer steps. Moreover, SN-DDPM [2] employs a 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.

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. It involves the user providing a full-resolution image x ( g ) through the manipulation of RGB pixels, referred to as the guided image. The guided image x ( g ) 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 a target image.

SDEdit [27] solves the task by first starting from the guided image x ( g ) and adding Gaussian noise to disturb the guided images to x t and q ( x ( g ) ( t 0 ) | x ( g ) ) ∼ N ( x t 0 | a ( t 0 ) x ( g ) , σ 2 ( t ) I ) same with Eq. (1). Subsequently, it solves the corresponding reverse stochastic differential equation (SDE) up to t = 0 to generate the synthesized image x (0) discretizing Eq. (4). Apart from the discretization steps taken by the SDE-based solver, the key hyper-parameter for SDEdit is t 0 , the time step from which we begin the image synthesis procedure in the reverse SDE.

(a) logarithm mean of

L

2

-norm

(b) logarithm median of

L

2

-norm

10-3

10-2

Norm

10-6

10-5

ION

2

10-9

10-8

Log

#Step:10

g

10-12

#Step:10

#Step:20

5

#Step:20

10-11

#Step:100

10-15

#Step:100

#Step:1000

#Step:1000

0

200

400

600

800

1000

0

200

400

600

800

1000

TimeStept

Time Step t

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 a Gaussian mixture distribution via the generalized methods of the moments [11] using higher-order moments information for the true reverse transition (see Sec. 3.2). The higher-order moments are estimated by a noise prediction network with multiple heads on training data, as detailed in Sec. 3.3.

3.1 Suboptimality of Gaussian distributions for reverse transition kernels

As described in Sec. 2, existing state-of-the-art SDE-based solvers for DPMs [3, 2] approximate the reverse transition q ( x t i -1 | x t i ) using Gaussian distributions. Such approximations work well when the number of discretization steps in these solvers is 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.

First, observe that q ( x s | x t ) 1 can be expressed as q ( x s | x t ) = ∫ q ( x s | x t , x 0 ) q ( x 0 | x t )d x 0 , which is non-Gaussian for a general data distribution q ( x 0 ) . For instance, for a mixture of Gaussian or a mixture of Dirac q ( x 0 ) , we can prove that the conditional distributions q ( x s | x t ) in the reverse process are non-Gaussian, as characterized by Proposition 3.1, proven in Appendix A.1.

Proposition 3.1 (Mixture Data Have non-Gaussian Reverse Kernel) . Assume q ( x 0 ) is a mixture of Dirac or a mixture of Gaussian distribution and the forward process is defined in Eq. (1) . The reverse transition kernel q ( x s | x t ) , s < t is a 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 ( x s | x t ) . However, even in such cases, it is easy to validate that q ( x s | x t ) are non-Gaussian. In particular, note that the third-order moment of one Gaussian distribution ( M ( G ) 3 ) can be represented by its first-order moment ( M 1 ) and its second-order moment ( M 2 ) 2 , which motivates us to check whether the first three order moments of q ( x s | x t ) satisfy the relationship induced by the Gaussian assumption. We perform an experiment on CIFAR-10 and estimate the first three orders of moments ˆ M 1 , ˆ M 2 , ˆ M 3 for q ( x s | x t ) from data by high-order noise networks (see details in Sec. D.2). As shown in Fig. 2, we plot the mean and median of l 2 -norm

between the estimated third-order moment ˆ M 3 and third-order moment M ( G ) 3 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 l 2 -norm increases, indicating that the true reverse distribution q ( x s | x t ) is non-Gaussian. With the time step closer to t = 0 , the l 2 -norm increases too.

Both the theoretical and empirical results motivate us to weaken the Gaussian assumption in SDEbased 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 ( x s | x t ) such that it is more powerful and potentially fits q ( x s | x t ) better than a Gaussian. In this paper, we choose a simple mixture of the Gaussian model as follows:

where w i is a scalar and µ i ( x t ) and Σ i ( x t ) are vectors. The reasons for choosing a Gaussian mixture model are three-fold. First, a 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 a 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, t ) by individual EM processes where we need to sample multiple x s 3 given a x t to estimate the parameters in the Gaussian mixture indexed by ( s, t ) . This is time-consuming, especially in a high-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 ϵ ( x t , t ) [14] or the data prediction network x 0 ( x t , t ) [18]. 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 , such as the decomposition shown in Eq. (10). In previous studies, Gaussian transition kernel was utilized, allowing for the distribution to 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 a 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 N and the number of parameters in the Gaussian mixture is d . We adopt a popular and theoretically sound method called the generalized method of moments (GMM) [11] to learn the parameters by:

where θ ∈ R d × 1 includes all parameters (e.g., the mean of each component) in the Gaussian mixture defined in Eq. (7). For instance, d = 2 M ∗ D data + M -1 in Eq. (7), where D data represents the number of dimensions of the data and θ contains M mean vectors with D data dimensions, M variance vectors of with D data dimensions (considering only the diagonal elements), and M -1 weight coefficients which are scalar. The component of g ( x i , θ ) ∈ R N × 1 is defined by g n ( x i , θ ) = M n ( x i ) -M ( GM ) n ( θ ) , where M n ( x i ) is the n -th order empirical moments stated in Sec. 3.3 and M ( GM ) n ( θ ) is the n -th order moments of Gaussian mixture under θ , and W is a weighted matrix, N c is the number of samples.

Theoretically, the parameter ˆ θ GMM obtained by GMM in Eq. (8) consistently converges to the potential optimal parameters θ ∗ for Gaussian mixture models given the moments' condition because √ N c ( ˆ θ GMM -θ ∗ ) d →N (0 , V ( θ GMM )) , as stated in Theorem 3.1 in Hansen [11].

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 a Gaussian mixture transition kernel with two components shown in Eq. (9) whose vectors µ (1) t , µ (2) t , and σ 2 t 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.

meanwhile, under the simplification in our paper, the number of parameters d = 3 ∗ D data 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 (asymptotically 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 d = N ) . Assume the number of parameters d equals the number of moments condition N , and ϵ n θ ( x t , t ) p → E q ( x 0 | x t ) [ diag ( ϵ ⊗ n -1 ϵ )] (where ⊗ n denotes n-fold outers product) which denotes n -th order noise network converging in probability. The asymptotic variance V ( θ GMM ) and the convergence speed of GMM remain the same no matter 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 i in the Gaussian mixture in Appendix E.1. 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 in 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 a set of moments M 1 , ..., M n . 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 q ( x t | x s ) and q ( x s | x t , x 0 ) follow Gaussian distributions. Specifically, q ( x t | x s ) ∼ N ( x t | a t | s x s , σ 2 t | s I ) [21], where a t | s =

α ( t ) α ( s ) and σ 2 t | s = σ 2 ( t ) -a 2 t | s σ 2 ( s ) . Consequently, we can deduce that E q ( x s | x t ) [ x s ⊗ n x s ] = E q ( x 0 | x t ) q ( x s | x t ,x 0 ) [ x s ⊗ n x s ] , where ⊗ n denotes n -fold outer product. Thus, we can first utilize the Gaussian property of q ( x s | x t , x 0 ) and employ an explicit formula to calculate the moments under the measure of q ( x s | x t , x 0 ) . 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 in 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:

Algorithm 1 Learning of the high order noise network

1: Input: The n -th order noise network ϵ n θ,ϕ n ( x t , t ) and its tunable parameter ϕ n

2: repeat

3: x 0 ∼ q ( x 0 )

4: t ∼ Uniform ([1 , ..., T ])

5: ϵ t ∼ N (0 , I )

6: x t = α ( t ) x 0 + σ ( t ) ϵ t

7: Take a gradient step on ϕ ( t ) n ←∇ ϕ n ∥ ∥ ∥ ϵ n t -ϵ n θ,ϕ n ( x t , t ) ∥ ∥ ∥ 2 2

8: until converged

Algorithm 2 Sampling via GMS with the first three order moments

1: Input: The assembled noise network f 3 ( x t , t ) in Eq. (13)

2: x T ∼ N (0 , I )

3: for t = T to 1 do

4: z t ∼ N (0 , I )

5: ˆ M 1 , ˆ M 2 , ˆ M 3 = h ( f 3 [1] ( x t , t ) , f 3 [2] ( x t , t ) , f 3 [3] ( x t , t )) in Appendix B

6: π ∗ k , µ ∗ k , σ ∗ k = min π k ,µ k ,σ k Q ( π k , µ k , σ k , ˆ M 1 , ˆ M 2 , ˆ M 3 ) in Sec. 3.2 7: i ∼ π ∗ k , x t -1 = µ ∗ i + σ ∗ i z t

8: end for

9: return x 0

where ⊗ is the outer product ⊙ is the Hadamard product. Additionally, λ 2 t corresponds to the variance of q ( x s | x t , x 0 ) , and further elaboration on this matter can be found in Appendix B.

In order to compute the n -th order moment ˆ M n , as exemplified by the third-order moment in Eq. (10), it is necessary to evaluate the expectations E q ( x 0 | x t ) [ x 0 ] , . . . , E q ( x 0 | x t ) [ diag ( x 0 ⊗ n -1 x 0 )] . Furthermore, by expressing x 0 as x 0 = 1 √ α ( t ) ( x t -σ ( t ) ϵ ) , we can decompose ˆ M n into a combination of terms E q ( x 0 | x t ) [ ϵ ] , . . . , E q ( x 0 | x t ) [( ϵ ⊙ n -1 ϵ )] , which are learned by neural networks. The decomposition of third-order moments ˆ M 3 , as outlined in Eq. (29), is provided to illustrate this concept. Therefore in training, we learn several neural networks { ϵ n θ } N n =1 by training on the following objective functions:

where ϵ ∼ N (0 , I ) , x t = α ( t ) x 0 + σ ( t ) ϵ , and ϵ n denotes ϵ ⊙ n -1 ϵ . After training, we can use { ϵ n θ } N n =1 to replace { E q ( x 0 | x t ) [ ϵ ⊙ n -1 ϵ ] } N n =1 and estimate the moments { ˆ M n } N n =1 of reverse transition kernel.

However, in the present scenario, it is necessary to infer the network a minimum of n times in order to make a 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 et al. [2]. Specifically, in the first stage, we optimize the noise network ϵ θ ( x t , t ) by minimizing the expression min E t E q ( x 0 ) q ( ϵ ) ∥ ϵ -ϵ θ ( x t , t ) ∥ 2 2 or by utilizing a pre-trained noise network as proposed by [14, 37]. In the second stage, we utilize the optimized network as the backbone and keep its parameters θ fixed, while adding additional heads to generate the n -th order noise network ϵ n θ,ϕ n ( x t , t ) .

where NN is the extra head, which is a small network such as convolution layers or small attention block, parameterized by ϕ n , details in Appendix E.3. We present the second stage learning procedure in Algorithm 1. Upon training all the heads, we can 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 f N ( x t , t ) . When estimating the k -th order moments, it suffices to extract only the first k components of the assembled noise network.

To this end, we outline the GMS sampling process in Algorithm 2, where f 3 [2] ( x t , t ) represents concat ([ ϵ θ ( x t , t ) , ϵ 2 θ,ϕ 2 ( x t , t )]) . In 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, we 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-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 on image data

In this section, we conduct a quantitative comparison of sample quality using the widely adopted FID score [13]. Specifically, we evaluate multiple SDE-based solvers, including a 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 CIFAR10 and ImageNet 64 × 64. Specifically, GMS achieves a remarkable 4.44 improvement in FID given 10 steps on CIFAR10. Appendix E.7 illustrates in more detail 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 as the noise network in Appendix E.8. Furthermore, taking into account the additional time required by GMS, our method still exhibits improved performance, as 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 E.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, we report the L 2 distance summed over all pixels between the stroke images and the edited output images.

SDEdit [27] applies noise to the stroke image x g at time step t 0 using N ( x g t 0 | α ( t ) x g , σ 2 ( t ) I ) and discretize the reverse SDE in Eq. (4) for sampling. Fig. 9 demonstrates the significant impact of t 0 on the realism of sampled images. As t 0 increases, the similarity to real images decreases. we choose

(a) Trade-off between realism and faithfulness

(b)

t

0

= 0

.

3

T

(c)

t

0

= 0

.

4

T

(d)

t

0

= 0

.

5

T

DDIM

GMS (Ours)

38

22

16

(FID)

30

36

21

15

DDIM

SN-DDPM

GMS

Realism

25

FID

34

SN-DDPM

DDIM

GMS

20

DDIM

GMS

SN-DDPM

14

20

32

19

13

30

18

15

28

17

12

0.2

0.3

0.4

0.5

0.6

25

50

100

25

50

100

25

50

100

FaithfulnessScore(L2-norm↓)

Number of steps

Numberofsteps

Numberofsteps

the range from t 0 = 0 . 3 T to t 0 = 0 . 5 T ( 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 t 0 and the same number of steps, edited output images from GMS have lower faithfulness but with higher realism. This phenomenon is likely attributed to 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 to transition from the stroke domain to the real image domain. Fig. 3(b) to Fig. 3(d) further demonstrates this phenomenon to a 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 ODE-based solver diminishes due to the absence 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. 4 shows the samples using DDIM and GMS when t 0 = 400 and the number of steps is 40 .

5 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 as those based on generative adversarial networks proposed by Xiao et al. [40] and Wang et al. [39]. Gao et al. [9] employ an energy function to model the reverse transition kernel. Zhang and Chen [44] use a 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 et al. [1] introduce a 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 to 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 to reduce discretization errors and improves the sample quality under the same sampling steps. Experimental results show that GMS outperforms existing SDE-based solvers, achieving a 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 state-of-the-art SDE-based solvers.

(a) Real Images

(b) Stroke

(c) GMS

(d) DDIM

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.

References

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.09392 , 2022.

Fan Bao, Chongxuan Li, Jiacheng Sun, Jun Zhu, and Bo Zhang. Estimating the optimal covariance with imperfect mean in diffusion probabilistic models. arXiv preprint arXiv:2206.07309 , 2022.

Fan Bao, Chongxuan Li, Jun Zhu, and Bo Zhang. Analytic-dpm: an analytic estimate of the optimal reverse variance in diffusion probabilistic models. arXiv preprint arXiv:2201.06503 , 2022.

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/CVF Conference on Computer Vision and Pattern Recognition , pages 22669-22679, 2023.

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 arXiv:2303.06555 , 2023.

Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. arXiv preprint arXiv:2003.04297 , 2020.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems , 34:8780-8794, 2021.

Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Genie: Higher-order denoising diffusion solvers. arXiv preprint arXiv:2210.05475 , 2022.

Ruiqi Gao, Yang Song, Ben Poole, Ying Nian Wu, and Diederik P Kingma. Learning energybased models by diffusion recovery likelihood. arXiv preprint arXiv:2012.08125 , 2020.

Martin Gonzalez, Nelson Fernandez, Thuy Tran, Elies Gherbi, Hatem Hajri, and Nader Masmoudi. Seeds: Exponential sde solvers for fast high-quality sampling from diffusion models. arXiv preprint arXiv:2305.14267 , 2023.

Lars Peter Hansen. Large sample properties of generalized method of moments estimators. Econometrica: Journal of the econometric society , pages 1029-1054, 1982.

Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfir Aberman, Yael Pritch, and Daniel Cohen-Or. Prompt-to-prompt image editing with cross attention control. arXiv preprint arXiv:2208.01626 , 2022.

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.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. J. Mach. Learn. Res. , 23(47):1-33, 2022.

Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. arXiv preprint arXiv:2204.03458 , 2022.

Alexia Jolicoeur-Martineau, Ke Li, Rémi Piché-Taillefer, Tal Kachman, and Ioannis Mitliagkas. Gotta go fast when generating data with score-based models. arXiv preprint arXiv:2105.14080 , 2021.

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. arXiv preprint arXiv:2206.00364 , 2022.

Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. arXiv preprint arXiv:2201.11793 , 2022.

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.

Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems , 34:21696-21707, 2021.

Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. arXiv preprint arXiv:2009.09761 , 2020.

Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. arXiv preprint arXiv:2202.09778 , 2022.

Cheng Lu, 2023. https://github.com/huggingface/diffusers/pull/3344 .

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.

Geoffrey J McLachlan and Thriyambakam Krishnan. The EM algorithm and extensions . John Wiley & Sons, 2007.

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:2108.01073 , 2021.

Eliya Nachmani and Shaked Dovrat. Zero-shot translation using diffusion models. arXiv preprint arXiv:2111.01471 , 2021.

Eliya Nachmani, Robin San Roman, and Lior Wolf. Non gaussian denoising diffusion models. arXiv preprint arXiv:2106.07582 , 2021.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning , pages 8162-8171. PMLR, 2021.

Ben Poole, Ajay Jain, Jonathan T Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. arXiv preprint arXiv:2209.14988 , 2022.

Vadim Popov, Ivan Vovk, Vladimir Gogoryan, Tasnima Sadekova, and Mikhail Kudinov. Gradtts: A diffusion probabilistic model for text-to-speech. In International Conference on Machine Learning , pages 8599-8608. PMLR, 2021.

Prajit Ramachandran, Barret Zoph, and Quoc V Le. Searching for activation functions. arXiv preprint arXiv:1710.05941 , 2017.

Alex Shonenkov, Misha Konstantinov, Daria Bakshandaeva, Christoph Schuhmann, Ksenia Ivanova, and Nadiia Klokova, 2023. https://github.com/deep-floyd/IF .

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pages 2256-2265. PMLR, 2015.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502 , 2020.

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.13456 , 2020.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation , 23(7):1661-1674, 2011.

Zhendong Wang, Huangjie Zheng, Pengcheng He, Weizhu Chen, and Mingyuan Zhou. Diffusion-gan: Training gans with diffusion. arXiv preprint arXiv:2206.02262 , 2022.

Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. arXiv preprint arXiv:2112.07804 , 2021.

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.

Xingqian Xu, Zhangyang Wang, Eric Zhang, Kai Wang, and Humphrey Shi. Versatile diffusion: Text, images and variations all in one diffusion model. arXiv preprint arXiv:2211.08332 , 2022.

Lvmin Zhang and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models. arXiv preprint arXiv:2302.05543 , 2023.

Qinsheng Zhang and Yongxin Chen. Diffusion normalizing flow. Advances in Neural Information Processing Systems , 34:16280-16291, 2021.

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.

A Proof

A.1 Proof of Proposition 3.1

When q ( x 0 ) is a mixture of Dirac distribution which means that q ( x 0 ) = ∑ M i =1 w i δ ( x -x i ) , ∑ M i =1 w i = 1 , which has total M components, and when the forward process q ( x t | x 0 ) is a Gaussian distribution as Eq. (1), the reverse process q ( x s | x t ) would be:

According to the definition of forward process, the distribution q ( x t | x i 0 ) = N ( x t | √ ¯ a t x i 0 , (1 -¯ a t ) I ) . Due to the Markov property of forward process, when t > s , we have q ( x s , x t | x 0 ) q ( x s | x 0 ) q ( x t | x s ) . The term q ( x s | x t , x 0 ) would be viewed as a Bayesian posterior resulting from a prior q ( x s | x 0 ) , updated with a likelihood term q ( x t | x s ) . And therefore

It is easy to prove that, the distribution q ( x s | x t ) is a mixture of Gaussian distribution:

When t is large, s is small, σ 2 t | s would be large, meaning that the influence of x i 0 would be large.

Secondly, when q ( x 0 ) is a mixture of Gaussian distribution which means that q ( x 0 ) = ∑ M i =1 w i N ( x i 0 | µ i , Σ i ) , ∑ M i =1 w i = 1 . For simplicity of analysis, we assume that this distribution is a one-dimensional distribution or that the covariance matrix is a high-dimensional Gaussian distribution with a diagonal matrix Σ i = diag i ( σ 2 ) . Similar to the situation above, for each dimension in the reverse process:

And q ( x s | x t ) could be a Gaussian mixture which has M component.

A.2 Proof of Proposition 3.2

Recall that our objective function is to find optimal parameters:

where g M ( θ ) is the moment conditions talked about in Sec. 3.2, W is the weighted Matrix, N c is the sample size. When solving such problems using optimizers, which are equivalent to the one we used when selecting θ GMM such that ∂Q T ( θ GMM ) ∂θ = 0 , and its first derivative:

its second derivative (the Hessian matrix):

By Taylor's expansion of the gradient around optimal parameters θ 0 , we have:

Consider one element of the gradient vector ∂Q Nc ( θ 0 ) ∂θ m

Consider one element of the Hessian matrix ∂ 2 Q Nc ( θ 0 ) ∂θ i ∂θ T j

Therefore, it is easy to prove that θ GMM -θ 0 p → 0 , and uses law of large numbers we could obtain,

therefore, we have

When the number of parameters d equals the number of moment conditions N , Γ 0 becomes a (nonsingular) square matrix, and therefore,

which means that the foregoing observation suggests that the selection of W N c has no bearing on the asymptotic variance of the GMM estimator. Consequently, it implies that regardless of the specific method employed to determine W N c , provided the moment estimates are asymptotically consistent, W N c 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 first n -th order. To ensure adherence to the aforementioned proposition, it is necessary to assume that the n -th order of the noise network converges with probability to E q ( x 0 | x t ) [ diag ( ϵ ⊗ n -1 ϵ )] , 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 ( x t | x t +∆ t ) as follows:

where ∆ t is the step size, q ( x t ) is the marginal distribution of x t . When x t +∆ t and x t are close enough, using Taylor expansion for log p ( x t +∆ t ) , we could obtain:

By ignoring the higher order terms, the reverse transition kernel will be Gaussian distribution. However, as ∆ 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.

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 to Appendix D.2.

B Calculation of the first order moment and higher order moments

Suppose the forward process is a Gaussian distribution same with the Eq. (1) as q ( x t | x s ) = N ( x t | a ( t ) x t -1 , σ ( t ) I ) .

And let 1 ≥ t > s ≥ 0 always satisfy, q ( x t | x s ) = N ( x t | a t | s x s , β t | s I ) , where a t | s = a t /a s and β t | s = σ 2 t -a 2 t | s σ 2 s , σ s = √ 1 -a ( s ) , σ t = √ 1 -a ( t ) . It's easy to prove that the distribution q ( x t | x 0 ) , q ( x s | x t , x 0 ) are also a Gaussian distribution [21]. Therefore, the mean of x s under the measure q ( x s | x t ) would be

And for the second order central moment Cov q ( x s | x t ) [ x s ] , we use the total variance theorem, refer to [2], and similar with [2], we only consider the diagonal covariance.

since the higher-order moments are diagonal matrix, we use diag ( M ) to represent the diagonal elements that have the same dimensions as the first-order moments, such as diag ( x s ⊗ x s ) = Cov q ( x s | x t ) [ x s ] and diag ( x s ⊗ x s ⊗ x s ) = ˆ M 3 have the same dimensions as x s

Moreover, for the diagonal elements of the third-order moments, we have E q ( x s | x t ) [ diag ( x s ⊗ x s ⊗ x s )] = E q ( x 0 | x t ) E q ( x s | x t ,x 0 ) [ x s ⊙ x s ⊙ x s ] , we could use the fact that E q ( x s | x t ,x 0 ) [( x s -µ ( x t , x 0 )) ⊙ ( x s -µ ( x t , x 0 )) ⊙ ( x s -µ ( x t , x 0 ))] = 0 , therefore,

0

Linear term in

x

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 cubic 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:

Therefore, when we need to calculate the third-order moment, we only need to obtain E q ( x 0 | x t ) [ ϵ t ] , E q ( x 0 | x t ) [ ϵ t ⊙ ϵ t ] and E q ( x 0 | x t ) [ ϵ t ⊙ ϵ t ⊙ ϵ t ] . Similarly, when we need to calculate

the n -order moment, we will use E q ( x 0 | x t ) [ ϵ t ] , ..., E q ( x 0 | x t ) [ ϵ t ⊙ n -1 ϵ t ] . Bao et al. [2] put forward using a sharing network and using the MSE loss to estimate the network to obtain the above information about different orders of noise.

The function h ( f 3 [1] , f 3 [2] , f 3 [3] ) in Algo. 2 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 3 [1] ( x t , t )) = E q ( x s | x t ) [ x s ] in Eq. (26), M 2 ( f 3 [2] ( x t , t )) = Cov q ( x s | x t ) [ x s ] in Eq. (27) and M 3 ( f 3 [3] ( x t , t )) = ˆ M 3 = E q ( x s | x t ) [ diag ( x s ⊗ x s ⊗ x s )] in Eq. (28)

B.1 Objective function for Gaussian mixture transition kernel with two components

Recall that the general objective function to fit a Gaussian mixture transition kernel in each sampling step via GMM is shown in Eq. (30).

where A T 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 ( x s | x t ) = 1 3 N ( µ (1) t , σ 2 t ) + 2 3 N ( µ (2) t , σ 2 t ) in each sampling step from x t to x s . Therefore, the final objective function as follow:

where to simplify the analysis, we use the scalar form of parameters µ (1) t , µ (2) t , σ 2 t as representation.

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

Wewould turn back to analyzing the original objective function of DPMs to find a suitable distribution. The forward process q ( x t | x s ) = N ( x t | a t | s x s , β t | s 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:

where q t . = q ( x t ) is the true distribution and p t . = p ( x t ) is the modeled distribution, and the minimum problem could be transformed into a sub-problem, proved in Bao et al. [3]:

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 p θ s | t ( x s | x t ) is exponential family p θ t ( x s | x t ) = p ( x t , θ s | t ) = h ( x t ) exp ( θ T s | t t ( x t ) -α ( θ s | t ) ) , solving the sub-problem Eq. (33) equals to solve the following equations, which is to match moments between the modeled distribution and true distribution:

When t ( x ) = ( x, .., x n ) T , solving Eq.(34) equals to match the moments of true distribution and modeled distribution.

Meanwhile, Gaussian distribution belongs to the exponential family with t ( x ) = ( x, x 2 ) T and θ t = ( µ t σ 2 t , -1 2 σ 2 t ) T , 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 q ( x s | x t ) , which is consistent with the results in Bao et al. [3] and Bao et al. [2].

The aforementioned discussion serves as a motivation to acquire higher-order moments and identify a corresponding exponential family, which surpasses the Gaussian distribution in terms of complexity. However, proposition C.1 shows that finding such exponential family distribution with higher-order moments is impossible.

Proposition C.1 (Infeasibility of exponential family with higher-order moments.) . Given the first n -th order moments. It's non-trivial to find an exponential family distribution for min D KL ( q || p ) when n is odd. And it's hard to solve min D KL ( q || p ) when n is even.

C.1 Proof of Proposition C.1

Lemma C.2. (Gaussian Distribution belongs to Exponential Family). Gaussian distribution p ( x ) = 1 √ 2 πσ exp ( -( x -µ ) 2 2 σ 2 ) is exponential family with t ( x ) = ( x, x 2 ) T and θ = ( µ σ 2 , -1 2 σ 2 ) T

Proof. For simplicity, we only prove one-dimensional Gaussian distribution. We could obtain:

Lemma C.3. (The Solution for Exponential Family in Minimizing the KL Divergence). Suppose that p ( x ) belongs to exponential family p ( x, θ ) = h ( x ) exp ( θ T t ( x ) -α ( θ ) ) , and the solution for minimizing the E q [log p ] is E q [ t ( x )] = E p ( x,θ ) [ t ( x )] .

Proof. An exponential family p ( x, η ) = h ( x ) exp ( η T t ( x ) -α ( η ) ) ∝ f ( x, η ) = h ( x ) exp ( η T t ( x ) ) with log-partition α ( η ) . And we could obtain its first order condition on E q [log p ] as:

In order to minimize the D KL ( q || p ) = ∫ q log( q/p ) = -E q [log p ] , we have:

For the second-order condition, we have the following:

Therefore, the second-order condition for the cross entropy would be:

When we assume that the reverse process is Gaussian, the solution to Eq. (33) equals to match the moment of true distribution and modeled distribution µ = E q [ x ] , Σ = Cov q [ x ] .

Lemma C.4. (Infeasibility of the exponential family with higher-order moments). Suppose given the first N -th order moments M i , i = 1 , .., N and modeled p as an exponential family. It is nontrivial to solve the minimum problem E q [log p ] when N is odd and it's difficult to solve when N is even.

Proof. While given the mean, covariance, and skewness of the data distribution, assume that we could find an exponential family that minimizes the KL divergence, so that the distribution would satisfy the following form:

where, t ( x ) = ( x, x 2 , x 3 ) , p = h ( x ) exp ( λ 0 + λ 1 x + λ 2 x 2 + λ 3 x 3 ) and ∫ dpx 3 = M 3 . However, when λ 3 is not zero, ∫ p = ∞ and density can't be normalized. The situation would be the same given an odd-order moment.

Similarly, given a more fourth-order moment, we could derive that λ 3 = 0 above, and we should solve an equation ∫ dpx 4 = M 4 and p = h ( x ) exp ( λ 0 + λ 1 x + λ 2 x 2 + λ 4 x 4 ) . Consider such function:

When λ -→ 0 , we could obtain lim λ → 0 Z ( λ ) = √ π For other cases, the lambda can be expanded and then integrated term by term, which gives Z ( λ ) ∼ ∑ ∞ n =0 ( -λ ) n n ! Γ(2 n +1 / 2) , but this function However, the radius of convergence of this level is 0 , 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 ∫ dpx 4 = M 4 , there is no analytical solution first, and the numerical solution also brings a large computational effort.

D More information about Fig. 1 and Fig. 2

D.1 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 q ( x 0 ) = 0 . 4 N ( -0 . 4 , 0 . 12 2 ) + 0 . 6 N (0 . 3 , 0 . 05 2 ) and we define our solvers in each step as p ( x s | x t ) = 1 3 N ( µ (1) t , σ 2 t ) + 2 3 N ( µ (1) t , σ 2 t ) with vectors µ (1) t , µ (2) t and σ 2 t , 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-modal. We use a simple MLP neural network with Swish activation [33].

Moreover, we experiment with our solvers in 8-Gaussian. The result is shown in Tab. 2. GMS outperforms Extended AnalyticDPM (SN-DDPM) [2] as presented in Tab. 2, with a bandwidth of 1 . 05 σL -0 . 25 , where σ is the standard deviation of data and L is the number of samples.

D.2 Experiment in Fig. 2

In this section, we will provide a comprehensive explanation of the procedures involved in computing the discrepancy between two third-order moment calculation methods, as depicted in Fig. 2.

The essence of the calculation lies in the assumption that the reverse transition kernel follows a 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:

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 ˆ M 3 by Eq. (28).

Weuse the pre-trained noise network from Ho et al. [14] and the second-order noise network form Bao et al. [2] and train the third-order noise network in CIFAR10 with the linear noise schedule.

Given that all higher-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 s | x t ) and the Gaussian distribution, we can utilize the following equation:

where ˆ M 3 is obtained via Eq. (28), and we can start at different time step t and choose a corresponding s to calculate the D s | t and draw different time step and step size t -s and we can derive Fig. 2.

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 q ( x s | x t ) = ω 1 N ( µ (1) s | t , Σ (1) s | t ) + ω 2 N ( µ (2) s | t , Σ (2) s | t ) as our foundational solvers, there exist five parameters (considered scalar, with the vector cases being analogous) available for optimization.

Our primary focus is on 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 can converge to a Gaussian distribution at the conclusion of 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.

When a parameter is not accompanied by a superscript, it implies that both components share the same value for that parameter. On the other hand, if a parameter is associated with a 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, our subsequent choices for optimization are primarily based on the first set of parameters µ (1) s | t , µ (2) s | t , Σ s | t . Another crucial parameter to consider is the selection of weights ω i . In Tab. 4, we show the result while changing the weight of the Gaussian mixture and the set of weight ω 1 = 1 3 , ω 2 = 1 2 performs best among different weight.

What's more, we observe that such choice of weights ( 1 3 , 1 2 ) 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 our experiments.

E.3 Details of the structure of the extra head

In Tab. 6, we list structure details of NN 1 , NN 2 and NN 3 of prediction networks used in our experiments.

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 a learning rate of 10 -4 ; we train 2M iterations in total for a higher order of noise network; we use an exponential moving average (EMA) with a rate of 0 . 9999 . We use a batch size of 64 on ImageNet 64X64 and 128 on CIFAR10. We save a checkpoint every 50K iterations and select the models with the best FID on 50k generated samples. Training one noise network on CIFAR10 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.

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:

where, a t is the learning rate after t steps, I w is the warm-up steps, α min 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 is negligible.

In Fig. 5, we present a comprehensive breakdown of the time of optimization process within GMS at each sampling step. Subfigure (a) in Fig. 5 illustrates that when the batch size is set to 100, the time of optimizing approximately 320 steps to fit a Gaussian mixture transition kernel at each step is 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 a 10 % margin of the original computational cost.

0

200

400

600

800

1000

Optimizer Steps

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

Time (seconds)

Optimizer Steps vs Time

Batch=100

Batch=500

Batch=1000

Nosie network inference time (Batch=100)

(a) Optimizing time and steps

10 1

10 2

10 3

Sampling Steps

0

200

400

600

800

1000

1200

1400

Time (in seconds)

Sampling Steps vs Time

SN-DDPM

GMS

(b) Sampling time and steps

Since many parts in the GMS introduce additional computational effort, Fig. 6 provides detailed information on the additional computational time of 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 a 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%

~0.4%-0.8%

~9%

Inferenceofnoisenetwork

Inferenceof2-ndnoisenetwork

：Inferenceof3-rdnoisenetwork

Solve the parameters for GM

(DDPM,DDIM,ADPMs)

(SN-DDPM,NPR-DDPM)

(GMS)

using different order moments

E.7 The reduction in the FID on CIFAR10 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, we have constructed Fig. 7 below. As observed from Fig. 7, GMS exhibits a more pronounced improvement when the number of sampling steps is limited. However, as 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 × 256 with the backbone noise network as U-ViTHuge [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 are the same as the training of the backbone of U-ViT-Huge (details in [4]). We will evaluate the sampling efficiency of GMS under conditional sampling (classifier-free guidance scale is 0.4) and unconditional sampling in Tab. 9

25

50

75

100

125

150

175

200

Sampling Steps

1

2

3

4

Relative FID Reduction

GMS vs. SN-DDPM FID Reduction

Relative FID Reduction

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

In order to provide a more intuitive representation for comparing different methods under the same sampling time, we have generated Fig. 8. Subfig. (a) of Fig. 8 illustrates that at each step, there is 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) of 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 a higher number of sampling steps.

10 1

10 2

10 3

Sampling Steps

0

200

400

600

800

1000

1200

1400

Time (in seconds)

Sampling Steps vs Time

SN-DDPM

GMS

(a) Sampling steps and sampling time

(b) sampling time and quality

SN-DDPM

8

GMS

~0.9

FID

5

4

~0.3

4 × 101

6×101

102

2 × 102

3 × 102

Time (seconds)

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 x 0 prediction network, while SEED [10] and GMS are based on Ho et al. [14]'s pre-trained model which is a discrete-time noise network.

E.11 Codes and License

In Tab.12, we list the code we used and the license.

F SDEdit

Fig. 9 illustrates one of the comprehensive procedures of SDEdit. Given a guided image, SDEdit initially introduces noise at the level of σ t 0 to the original image x 0 , where σ denotes the noise schedules of forward process in diffusion models. Subsequently, using this noisy image x t 0 and then discretizes it via the reverse SDE to generate the final image. Fig. 9 shows that the choice of t 0 will

can greatly will greatly affect the realism of sample images. With the increase of t 0 , 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 to examine the outcomes associated with different levels of fidelity. Taking all factors into comprehensive consideration, we have selected the range of t 0 from 0.3T to 0.5T in our experiments.

Real Image

Stroke

to =100

to =300

to =500

to =700

to =900

Besides experiments on LSUN 256 × 256, we also carried out an experiment of SDEdit on ImageNet 64 × 64. In Tab. 13, we show the FID score for different methods in different t 0 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 64 × 64. ODE-based solver is worse than all SDE-based solvers. With nearly the same computation cost, our GMS outperforms existing methods in most cases.

G Samples

From Fig. 10 to Fig. 12, we show generated samples of GMS under a different number of steps in CIFAR10 and ImageNet 64 × 64. Here we use K to denote the number of steps for sampling.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n**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.\n\nQUESTIONS:\nWhat 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\nLIMITATIONS:\nWhile limitations are presented, the authors could elaborate more on this topic.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nDespite 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe 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?", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nSampling 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.\n\nSTRENGTHS:\na.\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\nWEAKNESSES:\na. 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\nQUESTIONS:\nPlease refer to the weakness\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: 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).

STRENGTHS:
- The manuscript is well-written and the problem is clearly stated and illustrated.
- The proposed method allows to reduce the number of time step 
- The results are superior in terms of sample quality (measured by FID).

Specifically appreciated :
- The authors account for the computational cost of their method in the supp material.
- 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.
- The proposed method weakens an assumption at an extra-computational cost. When accounting for this extra-cost improvements are lowered.

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 ? 
EM 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 ?

Fig 1 : You should plot  the log density to make more visible the differences between the two curves.

LIMITATIONS:
ok
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06