SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} L (p, \hat {\lambda}) = D _ {\mathrm {K L}} (q | | p) - \hat {\lambda} ^ {T} (\int p t - m) \Rightarrow \frac {\partial}{\partial p} L (p, \hat {\lambda}) = \log \frac {p (x)}{h (x)} + 1 - \hat {\lambda} ^ {T} t = 0 \tag {40} \\ \Rightarrow p (x) = h (x) \exp \left(\hat {\lambda} ^ {T} t - 1\right) \\ \end{array}
$$
where,  $t(x) = (x,x^2,x^3), p = h(x)\exp \bigl (\lambda_0 + \lambda_1x + \lambda_2x^2 +\lambda_3x^3\bigr)$  and  $\int dp x^{3} = M_{3}$ . However, when  $\lambda_{3}$  is not zero,  $\int p = \infty$  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  $\lambda_3 = 0$  above, and we should solve an equation  $\int dp x^4 = M_4$  and  $p = h(x)\exp \bigl (\lambda_0 + \lambda_1x + \lambda_2x^2 +\lambda_4x^4\bigr)$ . Consider such function:
$$
Z (\lambda) = \int_ {- \infty} ^ {\infty} d x \exp \left(- x ^ {2} - \lambda x ^ {4}\right), \lambda > 0 \tag {41}
$$
When  $\lambda \longrightarrow 0$ , we could obtain  $\lim_{\lambda \to 0} Z(\lambda) = \sqrt{\pi}$ . For other cases, the lambda can be expanded and then integrated term by term, which gives  $Z(\lambda) \sim \sum_{n=0}^{\infty} \frac{(-\lambda)^n}{n!} \Gamma(2n + 1/2)$ , but this function However, the radius of convergence of this level is 0, so when the  $\lambda$  takes other values, we need to propose a reasonable expression for the expansion after the analytic extension. Therefore, for solving the equation  $\int dp x^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\mathcal{N}(-0.4,0.12^2) + 0.6\mathcal{N}(0.3,0.05^2)$  and we define our solvers in each step as  $p(x_{s}|x_{t}) = \frac{1}{3}\mathcal{N}(\mu_{t}^{(1)},\sigma_{t}^{2}) + \frac{2}{3}\mathcal{N}(\mu_{t}^{(1)},\sigma_{t}^{2})$  with vectors  $\mu_t^{(1)},\mu_t^{(2)}$  and  $\sigma_t^2$ , 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-DDDPM) [2] as presented in Tab. 2, with a bandwidth of  $1.05\sigma L^{-0.25}$ , where  $\sigma$  is the standard deviation of data and  $L$  is the number of samples.
Table 2: Comparison with SN-DDPM w.r.t. Likelihood  $\mathbb{E}_q[\log p_\theta (x)]$  ↑ on 8-Gaussian. GMS outperforms SN-DDPM.
#K	8-GAUSSIAN
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, 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:
$$
\mathbb {E} _ {q \left(x _ {t _ {i - 1}} \mid x _ {t _ {i}}\right)} \left[ x _ {t _ {i - 1}} \odot x _ {t _ {i - 1}} \odot x _ {t _ {i - 1}} \right] _ {\mathrm {G}} \doteq M _ {G} = \mu \odot \mu \odot \mu + 3 \mu \odot \Sigma , \tag {42}
$$
where  $\mu$  is the first-order moment and  $\Sigma$  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  $\hat{M}_3$  by Eq. (28).
We use 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  $\mu$ , 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:
$$
\mathrm {D} _ {s \mid t} = \log \left(\mathbb {E} _ {q \left(x _ {s} \mid x _ {t}\right)} \left[ x _ {s} \odot x _ {s} \odot x _ {s} \right] _ {\mathrm {G}} - \hat {M} _ {3}\right) ^ {2}, \tag {43}
$$
where  $\hat{M}_3$  is obtained via Eq. (28), and we can start at different time step  $t$  and choose a corresponding  $s$  to calculate the  $\mathrm{D}_{s|t}$  and draw different time step and step size  $t - s$  and we can derive Fig. 2.
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}) = \omega_{1}\mathcal{N}(\mu_{s|t}^{(1)},\Sigma_{s|t}^{(1)}) + \omega_{2}\mathcal{N}(\mu_{s|t}^{(2)},\Sigma_{s|t}^{(2)})$  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.
Table 3: Results among different parameters in CIFAR10 (LS), the number of steps is 50. The weight of Gaussian mixture is  $\omega_{1} = \frac{1}{3}$  and  $\omega_{2} = \frac{2}{3}$
μs|t(1), μs|t(2), Σs|t	μs|t(1), Σs|t, Σs|t	μs|t, Σs|t, Σs|t
CIFAR10 (LS)	4.17	10.12	4.22
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  $\mu_{s|t}^{(1)},\mu_{s|t}^{(2)},\Sigma_{s|t}$ . Another crucial parameter to consider is the selection of weights  $\omega_{i}$ . In Tab. 4, we show the result while changing the weight of the Gaussian mixture and the set of weight  $\omega_{1} = \frac{1}{3},\omega_{2} = \frac{1}{2}$  performs best among different weight.
Table 4: Results among different weight choices in CIFAR10 (LS), the number of steps is 50.
ω1=1/100, ω2=99/100	ω1=1/5, ω2=4/5	ω1=1/3, ω2=2/3	ω1=1/2, ω2=1/2
CIFAR10 (LS)	4.63	4.20	4.17	4.26
What's more, we observe that such choice of weights  $\left(\frac{1}{3},\frac{1}{2}\right)$  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.
Table 5: Details of noise prediction networks used in our experiments. LS means the linear schedule of  $\sigma (t)$  [14] in the forward process of discrete time step (see Eq. (1)). CS means the cosine schedule of  $\sigma (t)$  [30] in the forward process of discrete timesteps (see Eq. (1)).
# TIMESteps N	NOISE SCHEDULE	OPTIMIZER FOR GMM
CIFAR10 (LS)	1000	LS	ADAN
CIFAR10 (CS)	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  $\mathrm{NN}_1$ ,  $\mathrm{NN}_2$  and  $\mathrm{NN}_3$  of prediction networks used in our experiments.
Table 6:  $\mathrm{NN}_1$  represents noise prediction networks and  $\mathrm{NN}_2$ ,  $\mathrm{NN}_3$  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	NN2 (NOISE)	NN3 (NOISE)
CIFAR10 (LS)	NONE	CONV	RES+CONV
CIFAR10 (CS)	NONE	CONV	RES+CONV
IMAGENET 64x64	NONE	RES+CONV	RES+CONV
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.
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,  $\iota_t$  is a function with the second order growth function of sampling steps  $t$ .
LEARNING RATE	LR SCHEDULE	MIN LR	WARM-UP STEPS
CIFAR10	MAX(0.16-μt*0.16,0.12)	COS	0.1	18
IMAGENET 64×64	MAX(0.1-μt*0.1,0.06)	COS	0.04	18
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:
$$
\alpha_ {i + 1} = \left\{ \begin{array}{c c} \frac {i}{I _ {w}} \alpha_ {i} & \text {i f} \quad i \leq I _ {w} \\ \max  \left(\left(0. 5 \cos \left(\frac {i - I _ {w}}{I - I _ {w}} \pi\right) + 1\right) \alpha_ {t}, \alpha_ {\min }\right) & \text {e l s e} \end{array} \right. \tag {44}
$$
where,  $a_{t}$  is the learning rate after  $t$  steps,  $I_{w}$  is the warm-up steps,  $\alpha_{\mathrm{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.
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 BASELINE)	NOISE & SN PREDICTION NETWORKS SNDDPM	NOISE & SN PREDICTION NETWORKS (GMDDPM)
CIFAR10 (LS)	50.11 MB	50.11 MB	50.52 MB (+0.8%)
CIFAR10 (CS)	50.11 MB	50.11 MB	50.52 MB (+0.8%)
IMAGENET 64×64	115.46	115.87 MB	116.28 (+0.7%)
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.
(a) Optimizing time and steps
(b) Sampling time and steps
Figure 5: Sampling/optimizing 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.
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.
Figure 6: Time cost distribution for GMS.
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 \times 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 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
Figure 7: The reduction in the FID on CIFAR10 for GMS compared to SN-DDDPM. With a restricted number of sampling steps, the improvement exhibited by GMS surpasses that achieved with an ample number of sampling steps.
Table 9: 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.4 and 10% unconditional sampling, consistent with Bao et al. [4].
# TIMESTEPS K	COND, CFG=0.4	UNCOND
15	20	25	40	50	100	200	25	40	50	100
DDPM, βt	6.48	5.30	4.86	4.42	4.27	3.93	3.32	8.62	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., 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.
Table 10: Fair comparison with competitive SDE-based solvers w.r.t. FID score  $\downarrow$  on CIFAR10. SN-DDPM denotes Extended AnalyticDPM from [2]. The number of sampling steps for the GMS is indicated within parentheses, while for other solvers, it is represented outside of parentheses.
# TIMESTEPS K	CIFAR10 (LS)
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 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.
(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 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.(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  $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.
Table 11: Comparison with SDE-based solvers [18, 10] w.r.t. FID score  $\downarrow$  on CIFAR10.
# TIMESTEPS K	CIFAR10
10	20	40	50	100
SEED-2 [10]	481.09	305.88	51.41	11.10	3.19
SEED-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
HTTPS://GIRHUB.COM/W86763777/PYTORCH-DDPM	HO ET AL. [14]	WTFPL
HTTPS://GIRHUB.COM/OPENAI/IMPROVED-DIFFUSION	NICHOL AND DHARIWAL [30]	MIT
F SDEdit
Fig. 9 illustrates one of the comprehensive procedures of SDEdit. Given a guided image, SDEdit initially introduces noise at the level of  $\sigma_{t_0}$  to the original image  $x_0$ , where  $\sigma$  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.
Figure 9:  $t_0$  denotes the timestep to noise the stroke
Besides experiments on LSUN  $256 \times 256$ , we also carried out an experiment of SDEdit on ImageNet  $64 \times 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  $\downarrow$  on ImageNet  ${64} \times  {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.
#K	IMAGENET 64X64, t0 = 1200
26(28)	51(55)	101(111)	201(221)
DDPM,βn	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 a different number of steps in CIFAR10 and ImageNet  $64 \times 64$ . Here we use  $K$  to denote the number of steps for sampling.
(a) GMS  $(K = 10)$
(b) GMS  $(K = 20)$
(c) GMS  $(K = 25)$
(d) GMS  $(K = 40)$
(e) GMS  $(K = 50)$
(f) GMS  $(K = 100)$
(g) GMS  $(K = 200)$
(h) GMS  $(K = 1000)$
Figure 10: Generated samples on CIFAR10 (LS)
(a) GMS  $(K = 25)$
(b) GMS  $(K = 50)$
(c) GMS  $(K = 100)$
Figure 11: Generated samples on ImageNet  $64 \times 64$ .
(a) GMS  $(K = 200)$
Figure 12: Generated samples on ImageNet  $64\times 64$
(b) GMS  $(K = 400)$
(c) GMS  $(K = 4000)$