SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \hat {Y} _ {t _ {k + 1}} = \underbrace {e ^ {\frac {1}{2} \int_ {t _ {k}} ^ {t _ {k + 1}} g _ {s} ^ {\leftarrow 2} \mathrm {d} s} \hat {Y} _ {t _ {k}} + \left[ \int_ {t _ {k}} ^ {t _ {k + 1}} e ^ {\frac {1}{2} \int_ {t} ^ {t _ {k + 1}} g _ {s} ^ {\leftarrow 2} \mathrm {d} s} \frac {g _ {t} ^ {\leftarrow 2} + h _ {t} ^ {\leftarrow 2}}{2} \mathrm {d} t \right] S _ {k}} _ {=: \mathcal {T} _ {1}} \\ + \int_ {t _ {k}} ^ {t _ {k + 1}} \underbrace {e ^ {\frac {1}{2} \int_ {t} ^ {t _ {k + 1}} g _ {s} ^ {\leftarrow 2} \mathrm {d} s} h _ {t} ^ {\leftarrow}} _ {:= \mathcal {T} _ {2}} \mathrm {d} W _ {t}, \\ \end{array}
$$
$$
\boldsymbol {S} _ {k} := \nabla \log p _ {t _ {k}} ^ {\leftarrow} (\hat {Y} _ {t _ {k}}).
$$
$$
\hat {Y} _ {t _ {k + 1}} = \mathcal {T} _ {1} + \sqrt {\int_ {t _ {k}} ^ {t _ {k + 1}} \left(\mathcal {T} _ {2}\right) ^ {2} \mathrm {d} t} Z _ {k}, \quad Z _ {k} \sim \mathcal {N} (0, I _ {d}).
$$
If we pick
$$
\left\{ \begin{array}{l} h _ {t} ^ {\leftarrow} = \alpha g _ {t} ^ {\leftarrow}, \qquad \alpha \in \mathbb {R} ^ {+}, \\ g _ {t} = \sqrt {\beta_ {0} + (\beta_ {1} - \beta_ {0}) t} \qquad \text {w h i c h i m p l i e s t h a t} \qquad g _ {t} ^ {\leftarrow} = \sqrt {\beta_ {0} + (\beta_ {1} - \beta_ {0}) (T - t)}, \end{array} \right.
$$
then after straightforward calculations,
$$
\hat {Y} _ {t _ {k + 1}} = \gamma_ {k} \hat {Y} _ {t _ {k}} + (1 + \alpha^ {2}) (\gamma_ {k} - 1) \boldsymbol {S} _ {k} + \sqrt {\alpha^ {2} (\gamma_ {k} ^ {2} - 1)} Z _ {k},
$$
where  $\delta_k \coloneqq t_{k+1} - t_k$  and  $\gamma_k \coloneqq \exp\left(\frac{\delta_k(2\beta_0 + (2t_k - 2T + \delta_k)(\beta_0 - \beta_1))}{4}\right)$ .
H.5 Default training weight
Given the fixed  $X_0$ , we can explicitly solve (1) (with the choice  $f_t(x) = -\frac{g_t^2}{2} x$ ):
$$
X _ {t} = \left(G _ {t}\right) ^ {- 1} X _ {0} + \int_ {0} ^ {t} G _ {t} ^ {- 1} G _ {s} g _ {s} \mathrm {d} W _ {s},
$$
where  $G_{t} = \exp \left(\int_{0}^{t}\frac{g_{s}^{2}}{2}\mathrm{d}s\right)$ . The standard deviation of the Brownian motion term above is
$$
\varpi_ {t} := \sqrt {\int_ {0} ^ {t} G _ {t} ^ {- 2} G _ {s} ^ {2} g _ {s} ^ {2} \mathrm {d} s}.
$$
When  $g_{t} = \sqrt{\beta_{0} + (\beta_{1} - \beta_{0})t}$  as used in many literature [30, 17], we can explicitly solve for  $\varpi_t$ :
$$
\varpi_ {t} = \sqrt {1 - \exp \left(- \frac {1}{2} t ^ {2} \left(\beta_ {1} - \beta_ {0}\right) - t \beta_ {0}\right)}.
$$
In literatures, to balance the noise over time in training the score function, a common practice is to use the default weight  $\omega_{t} = \varpi_{t}^{2}$  (5) in the (denoising) score-matching loss function (4).
H.6 Additional numerical results for experiments in  $\S 4$
We present more numerical results on 1D Gaussian mixture, Swiss roll, and MNIST to further verify theoretical results.
1D Gaussian mixture
We evaluate the performances of generative models under different values of  $h$  for 1D Gaussian mixture, and present the visualization and numerical results in Fig. 8 and Fig. 9, respectively. In Fig. 8, a clear trend shows that with increasing  $h$ , the empirical density of generated samples better matches the true density function. This trend is more quantitatively captured in Fig. 9, from which we clearly observe that the distance between the empirical and the true density function decreases to the numerical threshold exponentially fast. Numerical threshold means the error of various distances when  $\epsilon = 0$ ; due to the space discretization when computing various distances, the numerical values of various distances are not exactly zero even when we use the exact score function. However, increasing  $h$  can help us to almost reach this limit and this phenomenon is theoretically described in Prop. 3.4.
Figure 8: Visualization of 1D GMM, where  $\mathcal{E}_t^{\leftarrow}(x) = \nabla \log p_t^{\leftarrow}(x)$ . We can observe that increasing h can help to generate samples with a distribution closer to the true  $p_0$ .
(e)  $\epsilon = 0.2$
(c)  $\mathcal{E}_t^{\leftarrow} = \mathbb{I}_{t < 0.95T}\nabla \log p_t^{\leftarrow}$
Figure 9: Numerical results of 1D 2-mode Gaussian mixture. We can observe that the distance between the empirical density and the true density function decreases as  $h$  increases, when  $\mathcal{E}_T^{\leftarrow}$  is not extremely large.
Swiss roll
In Fig. 10, we provide additional figures to discuss the effect of time-discretization. When the numerical error is negligible, we can observe that the generative process with a larger h can provide a clearer picture of Swiss roll, as shown in Fig. 10c. However, when the discretization error cannot be ignored, the conclusion may be reversed. Therefore, it is necessary to design and employ more accurate numerical methods for models with large h in order to fully benefit from diffusion models with a large diffusion coefficient.
In Fig. 11, we display JS and KL divergences between the true density  $p_0$  and the generated samples. Since the data distribution of Swiss roll is highly localized (data is concentrated on a curve embed in 2D), accurately computing the KL divergence poses a significant numerical challenge. That is why we use Wasserstein distance instead in Fig. 5a. Based on more robust symmetric metrics (i.e., the JS and Wasserstein distance herein), we can observe that a larger h can indeed diminish the error in sample generation.
$h^2 = 0.0g^2$
$h^2 = 1.0g^2$
$h^2 = 2.0g^2$
$h^2 = 3.0g^2$
(a) discretization steps  $= 100$
$h^2 = 4.0g^2$
(b) discretization steps  $= 1,000$
$h^2 = 3.0g^2$
$\mathrm{h}^2 = 4.0\mathrm{g}^2$
$h^2 = 0.0g^2$
Figure 10: Visualization results of Swiss roll with different number of time steps.
$h^2 = 1.0g^2$
(c) discretization steps  $= 10,000$
Figure 11: Numerical results of Jensen-Shannon divergence and Kullback-Leibler divergence of Swiss roll.
We conduct further experiments to explore the effect of  $h^{\leftarrow}$  and the time discretization steps. It is important to note that if one has an extremely well trained score function, then the effect of  $h^{\leftarrow}$  is indeed negligible, as shown in (3). To somewhat magnify the score training error for MNIST (but with a reasonable score function), we increase the time  $T$  and use a smaller architecture with fewer parameters; the reason of the occasional failure to generate clear MNIST images in later figures comes from this deliberate experimental design.
For an existing pre-trained score function with non-negligible error, we notice that increasing  $h^{\leftarrow}$  can almost ubiquitously improve the quality of sample generation; see Fig. 12 (as well as Fig. 15 and Fig. 16 under different training setup). This improvement is supported by our theoretical results, particularly Prop. 3.4. Notably, even when choosing  $h^{\leftarrow} = 0$  (ODE) and  $h^{\leftarrow} = g^{\leftarrow}$  (the default diffusion model in many studies) occasionally fail to generate an image, simply by increasing the magnitude of  $h^{\leftarrow}$  (possibly at the cost of more computational resources), we have a larger chance of generating an image with reasonable quality; for instance, see the last row in Fig. 12 (particularly see the last row in Fig. 15). The ODE-based model sometimes fails to generate hand-writing digits even when the step number is  $10^{3}$ , whereas the diffusion model (with large  $h^{\leftarrow}$ ) does not encounter this issue. This conclusion might be reversed when the time step size is large, which is similarly observed in the Swiss roll.
The computational cost of SDE-based models with large diffusion  $(h^{\leftarrow} > g^{\leftarrow})$  is relatively high due to the necessity of a larger number of time steps: a larger  $h^{\leftarrow}$  has the similar effect as running Langevin for a larger time horizon as discussed in § 3.3 and a longer time simulation is expected to be more expensive and also its accuracy largely relies on a well-chosen time step. However, this can be offset by the ability to use a lightweight architecture that possibly speeds up the generative process. A detailed comparison of various diffusion-based models with the same computational budget constraint is challenging and is slightly beyond the scope of this work, and we will leave this task to future work.
weight type=default, trial=1
Figure 12: Visualization of sample generation in MNIST where we used default weight  $\omega$  (5). We use different random seed to ensure robustness and this is the result for trial 1. Results for trial 2 is not included in this work as the overall tendency remains the same.
I Numerical experiments for adopting different weight in training
We have two reasons to explore the effects of various weight functions in training.
The first reason is that theoretically, any positive scalar-valued functions  $\omega_{t}$  on  $(0,T)$  is a valid candidate. This prompts the question of whether such a default weight function (5) is optimal in designing the loss function. We fully acknowledge that the default choice adopted in most literatures is a very effective one. However, the mathematical reason behind it is still not satisfactory in our opinion. This motivates us to ask whether the default choice is really the optimal one, at least in certain circumstances.
The second (and actually the primary reason) comes from our theoretical predictions discussed in § 3.7. As there is a tight connection between the optimal reverse-time generative process and the time-distribution of the score error (see Prop. 3.4, Prop. 3.5 and Prop. 3.6), if we are willing to train or re-train the score function and are interested in using the ODE-based model (for fast sample generation), it appears that we should focus more on the noise's end comparatively. To achieve this goal, namely, to control the distribution of score error, we adopt different weight schemes in the loss function only for training: default weight in (5), a data-driven case (more weight in the data side) and a noise-driven case (more weight in the noise side):
$$
\omega_ {t} = \left\{ \begin{array}{l l} \varpi_ {t} ^ {3} & \text {n o i s e - d r i v e n , o r s i m p l y r e f e r r e d a s}" \mathbf {n o i s e}"; \\ \frac {\varpi_ {t} ^ {2}}{0 . 2 5 + \varpi_ {t}} & \text {d a t a - d r i v e n , o r s i m p l y r e f e r r e d a s " d a t a "}. \end{array} \right. \tag {45}
$$
There is no theoretical reasons behind the noise-driven and data-driven choices in the last equation. We merely experiment with two reasonable choices and at the same time they are expected to help us control the time distribution of the score error.
After the initial camera-ready version of this work, we notice that such a weight design (in particular the noise-driven weight) has also been studied in [12]. Their experiments seem to provide further evidence about the validity of our arguments. However, we would like to emphasize that the approaches and perspectives towards proposing such a weight design are different: in [12], their conjecture is based on an insightful perception argument, whereas we take a mathematical approach and come up with the design based on the above error accumulation analysis.
I.1 Our guess
Based on our theoretical results, we guess that if a different weight function can really achieve our expected goal (namely, control the time distribution of the score function), the noise-driven case should be more suitable for ODE models, and the data-driven weight should give us the worst performance for ODE models.
We remark that this conjecture is surely not universal, and its validity remains to be fully validated by more benchmark experiments. Nevertheless, our numerical experiments below suggest its potential usefulness and it prompts an interesting question to explore and design the optimal score-matching loss function, which is rarely studied in literature. In what follows, we report numerical experiments for MNIST and CIFAR-10.
I.2 MNIST
We used different weights to train the score function and then visualize their generated samples in Fig. 13. We can clearly observe that the score function trained by the noise-driven weight produces comparatively better samples in ODE models. We plot the denoising score-matching loss for score functions obtained from training with various weights in Fig. 14:
- to make the comparison more straightforward, we visualize the time-distribution of relative score-matching loss rather than the absolute value, namely, we demonstrate:
$$
t \mapsto \frac {\mathbb {E} _ {X _ {0} \sim p _ {0}} \mathbb {E} _ {X _ {t} \sim p _ {t | 0} (\cdot | X _ {0})} \left[ \left\| \left(\mathfrak {S} ^ {(i)}\right) _ {t} (X _ {t}) - \nabla \log p _ {t | 0} (X _ {t} | X _ {0}) \right\| ^ {2} \right]}{\mathbb {E} _ {X _ {0} \sim p _ {0}} \mathbb {E} _ {X _ {t} \sim p _ {t | 0} (\cdot | X _ {0})} \left[ \left\| \left(\mathfrak {S} ^ {(\text {d e f a u l t})}\right) _ {t} (X _ {t}) - \nabla \log p _ {t | 0} (X _ {t} | X _ {0}) \right\| ^ {2} \right]}, i \in \{n o i s e, d a t a \}, \tag {46}
$$
and  $\mathfrak{S}^{(i)}$  refers to the score function obtained from training using weight  $i$  and we use test datasets to approximately represent  $p_0$ ;
- to ensure robustness, we independently conduct two trials with different neural network initializations: For each trial, the initial neural network is the same and the only difference in training is to adopt different weights in the loss function.
In Fig. 14, we can clearly observe that adopting different weights indeed help us to control how the score error is distributed over time as we expect (e.g., noise-driven weight helps us reduce the error near  $t \approx T$  and has an opposite effect near  $t \approx 0$ ); as shown in Fig. 13, their numerical performances in terms of sample generation also match our guess above.
We further visualize how the diffusion coefficient  $h$  and time step size affect the sample generation quality for score functions obtained by data-driven weight (see Fig. 15) and noise-driven weight (see Fig. 16). The conclusion is the same as in the default weight case. This further validates our theoretical results in § 3.
(a) Trial number 1
Figure 13: We show generated samples trained for MNIST using three different weight functions in the loss function: we use  $h^2 / g^2 = 0$  (ODE), time step is 1000, the trial number is the index for independent random initialization.
(b) Trial number 2
Relative score-matching loss
Figure 14: Score-matching loss for data-driven and noise-driven case on MNIST, compared with the default case
weight type  $=$  data, trial  $= 1$
Figure 15: Visualization of sample generation in MNIST where we used data-driven weight  $\omega$  (45). We use different random seed to ensure robustness and this is the result for trial 1. Results for trial 2 is not included in this work as the overall tendency remains the same.
weight type=noise, trial=1
Figure 16: Visualization of sample generation in MNIST where we used noise-driven weight  $\omega$  (45). We use different random seed to ensure robustness and this is the result for trial 1. Results for trial 2 is not included in this work as the overall tendency remains the same.
I.3 CIFAR-10
We carry out a similar experiment for CIFAR-10 to test the effect of training weights  $\omega_{t}$ : we used the same initialization (referred to as the trial number below) and all other hyper-parameters, except that we employ different weights in score-matching loss (4). The same architecture is used as in [17] for CIFAR-10 and the detailed hyper-parameters can be found in source codes.
In Fig. 17, we observe that overall over the whole training period, the noise-driven weight leads into a score function estimate no worsen than that by the default weight: due to stochastic fluctuations and other uncertainties (in particular if we adopt the mixed precision training), there is no guarantee that the noise-driven one is always better, but the overall tendency is still observable and clear. In Fig. 18, the noise-driven one actually has a slightly larger score-matching loss (SML) than the default one (which could be possibly explained by how we measure the SML in Fig. 18). What is interesting is that score functions trained by the noise-driven weight and the data-driven weight have a similar SML, which both decay at the similar pace; however, the FID values for the score function estimate by noise-driven weight are much smaller than that by the data-driven weight. This apparent gap clearly explains that apart from the total score-matching loss (which does matter), the time distribution of the score error plays an important role in determining the final sampling error, echoing our theoretical results in § 3.
For instance, in Fig. 17, if we consider the first experiment (i.e., trial=0) with float16 mixed-precision training (i.e., mixed=True), we notice that relatively near 80k and 120k training iterations, the performance of noise-driven one is much better than the default one, which is consistent with Fig. 19 that the relative loss near  $t \approx T$  is more minimised for iterations 80k and 120k, compared with other iteration stages. Moreover, for the same experiment in Fig. 17, the data-driven one has a much worse FID value at iteration 200k, which is compatible with the increasing relative error near the noise's end (i.e.,  $t \approx T$ ) in the last row of Fig. 19. For the remaining three experimental setup (either different initialization or training precision), we notice a similar consistency between how the time-distribution of the SML behaves and how FID values change. This relation, so far, still cannot be used as a rigorous quantitative indicator to predict one based on the other quantity, but qualitatively, the above explained relationship does appear to be numerically valid and theoretically sound.
In summary, if we have two score functions  $\mathfrak{S}_1$  and  $\mathfrak{S}_2$  from training:
- If the SML for  $\mathfrak{S}_1$  is much larger than the SML for  $\mathfrak{S}_2$ , then we can probably confidently expect that  $\mathfrak{S}_2$  is a more accurate estimate.
- However, when the total SML (4) for both are close, then the time-distribution of the score-matching loss together with which generative dynamics is chosen will play a significant role in determining the final sample generation quality, which is probably largely overlooked in current literature as far as we know. A full investigation and in particular whether it is possible to adapt this observation to achieve the state-of-art models will be left to the next stage of research.
Figure 17: We visualize FIDs of score-training estimates by various weights with respect to the number of training iteration. We consider ODE models (i.e.,  $h = 0$ ) when computing FIDs. Trial = 0, 1 refers to different neural network initialization; mixed=True means we use float16 mixed-precision for training; mixed=False means we use 32-bit precision.
Figure 18: We visualize score-matching loss (4) for score function trained using various weights with respect to the number of training iteration's. Trial  $= 0$ , 1 refers to different training initialization; mixed=True means we use float16 mixed-precision for training; mixed=False means we use 32-bit precision. When we compare SML for different score functions above, we consistently use the "default" weight for fairer comparison, even though some score functions are trained using different weight functions; we also use the test dataset to approximately represent  $p_0$ .
trial=0, mixed=True
Figure 19: The first independent experiment with the float16 mixed precision training: We visualize the relative time distribution of the SML (46) (on the left) and the time distribution of SML (on the right) for various training weights and for various training stages (each row).
trial  $= 0$  ,mixed  $\equiv$  False
Figure 20: The first independent experiment with the full 32-bit precision training: We visualize the relative time distribution of the SML (46) (on the left) and the time distribution of SML (on the right) for various training weights and for various training stages (each row).
trial=1, mixed=True
Figure 21: The second independent experiment with the float16 mixed precision training: We visualize the relative time distribution of the SML (46) (on the left) and the time distribution of SML (on the right) for various training weights and for various training stages (each row).
trial=1, mixed=False
Figure 22: The second independent experiment with the full 32-bit precision training: We visualize the relative time distribution of the SML (46) (on the left) and the time distribution of SML (on the right) for various training weights and for various training stages (each row).