HIGH PRECISION SCORE-BASED DIFFUSION MODELS
Year: 2021

Abstract:
Recent advances in diffusion models bring the state-of-the art performance on image generation tasks. However, image generation is still an arduous task in high resolution, both theoretically and practically. From the theory side, the difficulty arises in estimating the high precision diffusion because the data score goes to∞ as t → 0 of the diffusion time. This paper resolves this difficulty by improving the previous diffusion models from three aspects. First, we propose an alternative parameterization for such an unbounded data score, which theoretically enables the unbounded score estimation. Second, we provide a practical Soft Truncation trick (ST-trick) to handle the extreme variation of the score scales. Third, we design a Reciprocal Variance Exploding Stochastic Differential Equation (RVESDE) to enable the sampling at the high precision of t. These three improvements are applicable to the variations of both NCSN and DDPM, and our improved versions are named as HNCSN and HDDPM, respectively. The experiments show that the improvements result in the state-of-the-art performances in the high resolution image generation, i.e. CelebA-HQ. Also, our ablation study empirically illustrates that all of 1) alternative parameterization, 2) ST-trick, and 3) RVESDE contributes to the performance enhancement.

heading: 1 INTRODUCTION
text: Recent advances in the generative models enable creating of highly realistic images. One direction of such modeling is likelihood-free models (Karras et al., 2019) based on the minimax training. The other direction is likelihood-based models, including VAE Vahdat and Kautz (2020), autoregressive models (Parmar et al., 2018), and flow models (Grcić et al., 2021). Diffusion models (Ho et al., 2020) are one of the most successful likelihood-based models where the generative process is modeled by the reverse diffusion process. The success of diffusion models achieves the state-of-the-art performance in image generation (Dhariwal and Nichol, 2021; Song et al., 2020).
Despite of such success, the score-based diffusion models still struggle in generating realistic images with high resolution (Jolicoeur-Martineau et al., 2021). Rather than improving the network architecture, this paper focuses mostly on the theoretic and the optimization aspects of the diffusion model to improve the sample quality. We first observe that the diffusion time acts in distinctive ways on the sample generation: the diffusion process with the front diffusion time (t ≈ 0) influences the local sample fidelity, and the end diffusion time (t ≈ T ) shapes the sample global structure. However, the unbounded data score at the front diffusion time results in the training instability, and this instability yields the diffusion model with poor estimation limiting the local sample fidelity, which prevents the advancement in the high resolution because of its locality. Hence, this paper constructs a diffusion model capable of training both front time and end time, effectively.
We observe that the current score parametrization is incapable of estimating the data score that blows up as t→ 0, which requires a new theory-driven parametrization for the unbounded score network to enable the successful score estimation near t ≈ 0. This parametrization is universally applicable for any diffusion models including NCSN (Song and Ermon, 2019) and DDPM (Ho et al., 2020), and we call this parametrization as High precision NCSN (HNCSN) and High precision DDPM (HDDPM), respectively. Second, we propose a practical optimization method (ST-trick) to improve the score estimation at the end of the diffusion time (t ≈ T ). This ST-trick provides a distinctive way to optimize the diffusion model parameters, and we argue that this trick successfully trains the score network at the end diffusion time by its nature, which is likely to be ignored due to the minor contribution to the diffusion loss near t ≈ T in the previous studies. Like the new parametrization,
ST-trick fits to any diffusion model, including NCSN/HNCSN and DDPM/HDDPM. Finally, we suggest a new SDE (RVESDE) that 1) mitigates the theoretic dilemma of VESDE that arises as t→ 0 and 2) gives additional performance boost by reducing the Monte-Carlo variance for the diffusion loss. We perform the ablation study on our three improvements, and we find that using 1) the new parametrization, 2) ST-trick, and 3) RVESDE all together improves both density estimation and sample fidelity on benchmark datasets. Although we proceed our arguments based on the continuous diffusion models, we emphasize that both new parametrization and ST-trick are applicable for discrete diffusion models, such as ADM (Dhariwal and Nichol, 2021) and CDM (Ho et al., 2021), as well.

heading: 2 PRELIMINARY
text: The diffusion model trains the network by minimizing the Negative Evidence Lower BOund (NELBO)
(NELBO) Epr(x0) [ − log pθ(x0) ] ≤ L(θ;λ = g2) = 1
2
∫ T
0 g2(t)Lt(θ) dt, (1) where Lt(θ) = Ext [ ‖sθ(xt, t)−∇xt log pt(xt)‖22 ] = Ex0,xt [ ‖sθ(xt, t)−∇xt log p0t(xt|x0)‖22 ] , up to a constant. Here, λ is the weighting function for the diffusion loss ofLt(θ) at each diffusion time, pt(xt) is the marginal distribution of the diffusion process of {xt}Tt=0, and p0t(xt|x0) is the transition probability defined byN (xt;µ(t)x0, σ2(t)I), where µ(t) and σ2(t) are determined by the governing diffusion process. Previous works suggest a couple of diffusion processes: 1) VESDE dxt = g(t) dωt with g(t) = σmin(σmaxσmin ) t √ 2 log (σmaxσmin ), and 2) VPSDE dxt = − 1 2β(t)xt dt + √ β(t) dωt with β(t) = βmin + (βmax − βmin)t. With these diffusion schemes, VESDE has µ(t) ≡ 1 with σ2(t) = σ2min [ (σmaxσmin ) 2t − 1 ] , and VPSDE has µ(t) = e− 1 2 ∫ t 0 β(s) ds with σ2(t) = 1− e− ∫ t 0 β(s) ds. We observe that the transformation of yt = e 1 2 ∫ t 0 β(s) dsxt reduces VPSDE to VESDE of dyt = e 1 2 ∫ t 0 β(s) ds √ β(t) dωt by the Ito’s lemma (Oksendal, 2013), so VPSDE and VESDE are indeed equivalent representations of the linear diffusion. Given their equivalence under the suggested transformation, we build our model based on VESDE as a default diffusion otherwise stated.
2.1 EFFECT OF DIFFUSION TIME ON DIFFUSION MODEL
Density Estimation In practice, previous diffusion models truncate the diffusion time from [0, T ] to [ , T ] to stabilize the training. Figure 1 emphasizes the loss contribution to the NELBO at t→ 0, which signifies lowering in the practice. Moreover, Figure 2 shows the cumulative loss of Figure 1 integrated over [σ, σmax] by changing the integrating variable from t to σ in Eq. 2, and the high precision (i.e., σ ≈ 0 or t ≈ 0) has significant portion on the NELBO.
L(θ; g2) = 1 2 ∫ σmax 0 vLσ−1(v)(θ) dv. (2)
Therefore, lowering σmin = σ( ) guarantees the tighter NELBO given its large portion at t = . Empirically, NCSN++ (Song et al., 2020) performs 3.45 for NELBO trained with σmin = 10−2 on CIFAR-10, but this NELBO is tightened up to 3.04 by lowering σmin to 10−3 in our experiments.
Sample Generation While we discussed the value of the loss at t→ 0, both fron and end diffusion time contribute to the sample quality in two distinctive ways.
End Time (t → T ) : The diffusion model creates a fake sample by reverting the diffusion process with the score estimation. Figure 3 shows two contrastive image generations: the realistic image at the first row, and the unrealistic image at the second row with curly hair on its forehead. This unwanted hair is synthesized at the early steps of the
reverse diffusion, which correspond to the end time of the forward diffusion. Therefore, Figure 3 implies that the overall sample quality is determined on the range of the end diffusion time.
Front Time (t → 0) : The sample generation in Figure 4 shows that the local fidelity heavily depends upon the denoising process at the front of the diffusion time. This intuitive result is grounded by theory: observe that the random variable ‖xt − x0‖22/σ2(t) follows the χ2-distribution of degree d, where d is the data dimension. If the score network exactly esti-
mates the data score, the generated image at time t would follow xt, and the average distance between the real sample (x0) with the fake sample (x ) becomes Epr(x0)Ep0 (x |x0)[‖x − x0‖22] = σ2mind. Therefore, even if the score network perfectly estimates the data score, this leads that minimizing σmin acts in a crucial way for the pixel-wise precision.
From these observations, we identified that the front diffusion time is critical in generating the high resolution image for the pixel-wise details, and we also acknowledge that the end diffusion time contributes to the overall image structure. Now, our question becomes whether the existing models can learn the data score at both front and end diffusion time.

heading: 2.2 PROBLEMS OF TRAINING THE DIFFUSION MODEL
text: WITH HIGH PRECISION
Unbounded Score Figure 5 presents that the data score is unbounded as the diffusion level, σ2(t), diminishes. This requires the score parametrization to handle this unboundedness. For instance, the parametrization of sθ(xt, σ(t)) yields the failure of the score estimation at t = 10−4 on the toy example in Figure 6.
Theoretically, Lemma 1 states that the score estimation with sθ(xt, t) (DDPM) or sθ(xt, σ(t)) (NCSN) fails if the data score is unbounded. See Appendix A for the proof.
Lemma 1. Let H[0,T ] = {s : Rd × [0, T ] → Rd, s is locally Lipschitz}. Suppose a continuous vector field v defined on a subset U of a compact manifold M (i.e., v : U ⊂M → Rd) is unbounded, then there exists no s ∈ H[0,T ] such that limt→0 s(x, t) = v(x) a.e. on U .
The score network, sθ, is an element in H[0,T ] for any neural architecture as long as we put t in the network, so Lemma 1 indicates that the vanilla score network cannot estimate an unbounded data score.
Training Instability Training the diffusion time on the front range is advantageous on the creation of high-quality
samples. However, training the diffusion model with lowered σmin induces the numerical instability, which eventually harms the overall sample quality.
To analyze the logic of the training instability, Figure 7 illustrates the Monte-Carlo samples of the diffusion loss, which blows up as t → 0 (or σ → 0). This leads that the diffusion loss is dominated by the diffusion time with the front range in practice, and the loss at the end diffusion time is barely counted in the gradient signal. This imbalanced loss brings the immatured score estimation at the end diffusion time, which ruins the overall sample shape.
Limitation of Diffusion Process Continuous VESDE has its diffusion variance to be σ2(t) = σ2min [ (σmaxσmin ) 2t − 1 ] . As introduced in Song and Ermon (2020), VESDE has been suggested in its original discrete version based on the geometric property for its smooth transition of the distributional shift. Mathematically, the variance is geometric if ddt log σ
2(t) is a constant, but the continuous VESDE breaks this geometric progression as in Figure 8.
To make it geometric, Song et al. (2020) approximates the VESDE diffusion process by the transition probability of p0t(xt|x0) =
N ( xt;x0, σ 2 min( σmax σmin )2tI ) , so this approximate diffusion process satisfies the geometric progression. However, this approximation leads that xt is not converging to x0 in distribution because σ2min( σmax σmin
)2t → σ2min > 0 as t→ 0. Proposition 1. There is no SDE that has the stochastic process {xt}t∈[0,T ], defined by a transition probability p0t(xt|x0) = N (xt;x0, σ2min(σmaxσmin ) 2tI), as the solution.
Proposition 1 indicates that if we approximate the diffusion process by p0t(xt|x0) = N ( xt;x0, σ 2 min( σmax σmin )2tI ) , then the reverse diffusion cannot be modeled by a reverse SDE. Therefore, VESDE loses its theoretic ground on the continuation of the discrete Markov chain.

heading: 3 HIGH PRECISION SCORE-BASED DIFFUSION MODELS
text: Given the three problems of unbounded score, training instability, and limited diffusion process, this section improves the diffusion model in three corresponding aspects. Section 3.1 introduces the theory-driven new parametrization, called HNCSN and HDDPM, that enables the successful score estimation for the diffusion model with small σmin. Section 3.2 proposes the practical technique, ST-trick, which could estimate the score network at the end of the diffusion time without being dominated by the front diffusion. Section 3.3 suggests the new SDE, RVESDE, which develops the VESDE suitable for small σmin.

heading: 3.1 PARAMETRIZATION MODEL FOR UNBOUNDED SCORE
text: As presented in Lemma 1, the score estimation of baseline models (NCSN/DDPM) fails because the locally Lipschitz with respect to the time argument does not hold anymore at t = 0. From this, we make the time argument extendable to the infinity at t = 0. To proceed, let us define η : t 7→ R as a function that satisfies η(t)→∞ or −∞ as t→ 0. Then, the parametrization of sθ(xt, η(t)) would successfully estimate the data score by Proposition 2. Proposition 2. LetH[1,∞) = {s : Rd× [1,∞)→ Rd, s is locally Lipschitz}. Suppose a continuous vector field v defined on a d-dimensional open subset U of a compact manifold M is unbounded, and the projection of v on each axis is locally integrable. Then, there exists s ∈ H[1,∞) such that limη→∞ s(x, η) = v(x) a.e. on U .
Proposition 2 becomes the foundation of the unbounded parametrization, and the introduction of η is the key to enable the estimation of the unbounded data score. Therefore, we introduce the new parametrization for the high precision score network as sθ(xt, η(t)), rather than sθ(xt, t) (DDPM) or sθ(xt, σ(t)) (NCSN).
Table 1 presents example parametrizations following Proposition 2. We suggest HNCSN with η(t) as the mixture of log σ(t) at the end diffusion time and 1σ(t) at the front diffusion time, in which the inverse function is beneficial on estimating the rapidly blowing data score by orders of 1σ . Note that Song and Ermon (2020) proposed the score parametrization of sθ(xt)σ(t) , and Jolicoeur-Martineau et al. (2020) showed that the optimal score function becomes s∗(xt, t) = Ex0|xt [x0]−xt
σ2(t) ≈ zσ(t) , where z ∼ N (0, I). Therefore, we suggest HNCSN to put 1σ(t) in place of log σ(t) at the front time.
Before we choose optimal η for DDPM, observe that DDPM with the importance sampling has g 2(t) σ2(t) as the importance weight, and this weight forces the sampled diffusion time highly concentrated on t ≈ . Since a small difference between two diffusion times near t1, t2 ≈ results in a large deviation of two data scores in their scales, DDPM cannot approximate this extreme scale deviation with the vanilla parametrization of sθ(xt, t) by Proposition 2. We propose HDDPM with η(t) = ∫ g2(t) σ2(t) dt as the antiderivative of the importance weight. For instance, HDDPM with VPSDE has its antiderivative as ∫ g2(t) σ2(t) dt = log (1− e− ∫ t 0 β(s) ds) +
∫ t 0 β(s) ds. This antiderivative diverges to −∞ as t → 0 for both VPSDE and subVPSDE, so it satisfies the condition of Proposition 2.
The choice of such η(t) for HDDPM helps the enhanced score estimation near t ≈ because the input of η(t) is distributed uniformly when we draw samples from the importance weight. Concretely, when the sampling distribution on the diffusion time is given by p(t) ∝ g
2(t) σ2(t) , the η-distribution from the importance sampling be-
comes p(η) ∝ 1, which is depicted in Figure 9. Therefore, this uniform distribution on η(t) in HDDPM guarantees the high precision to the score estimation.

heading: 3.2 SOFT TRUNCATION TRICK FOR PRACTICAL SCORE ESTIMATION
text: While the unbounded parametrization suggested in Section 3.1 provides the theoretic guarantee of the score estimation, this section proposes a practical optimization method for the high precision score estimation. In order to successfully train the score network without sacrificing the global fidelity, we detour the training instability by introducing the Soft Truncation (ST) trick that is taking a distinctive diffusion truncation for every mini-batch. In formula,
LST (θ;λ,P) = EP(τ)[L(θ;λ, τ)], where P(τ) is the prior distribution for the diffusion truncation. In every mini-batch update, we optimize L(θ;λ, τ) with the sampled τ ∼ P(τ) from the prior. Figure 7 illustrates four cases: blue dots with τ = , green dots with τ = σ−1(0.01), etc. With ST-trick, a mini-batch with high truncation of τ = σ−1(1) (purple) focuses on the score estimation at large t, and mini-batches with medium τ = σ−1(0.1) (green) and = σ−1(0.01) (red) focus on the middle region of [ , T ]. In consequence, ST-trick trains the diffusion model by partitioning its diffusion time for each mini-batch update.
The ST-trick with P(τ) = δ (τ) unconditionally selects for every batch update, so it is equivalent with the original optimization of L(θ;λ, ). Although it is free to choose P, we find that the reciprocal functions work best in practice. Observing that 1 τk 1[ ,T ](τ)/Z → δ (τ) as k → ∞ where Z is the normalizing constant, the recipro-
cal functions are connected to the original optimization. We enumerate few examples of k = 2 and k = 1 in Table 2 for RVESDE and VPSDE, respectively. Note that ST-trick is independent to the choice on the weighting function λ for the batch-wise optimization loss, L(θ;λ, τ). Theoretic Analysis of ST-trick The ST-loss reduces the original diffusion loss to LST (θ;λ,P) = L(θ;λP, ) with λP(t) := ( ∫ t 0 P(τ) dτ ) λ(t) in expectation by the Fubini theorem (Stein and Shakarchi, 2009). While the ST-loss in average is identical to the original diffusion loss, the major difference lies on that 1) L(θ;λP, ) optimizes the diffusion loss in its averaged form and 2) LST (θ;λ,P) with ST-trick optimizes the diffusion loss with the Monte-Carlo sampled truncation, τ . To inspect how soft truncation affects to the optimization, let us assume λ(t) = g2(t). Then, we observe that the theoretic analysis previously suggested in Song et al. (2021) is directly extendable to any τ ∈ (0, T ] by simply switching the positive truncation time, τ , in place of the zero time:
Epτ (xτ ) [ − log pθ,τ (xτ ) ] ≤ L(θ; g2, τ),
where pθ,τ is the marginal probability distribution at time τ of the reverse (generative) process that starts from the prior distribution and flows backwards in time, and pτ is the marginal probability density at time τ of the forward diffusion process. Then, the above inequality is equivalent to
DKL(pτ‖pθ,τ ) ≤ L(θ; g2, τ), after omitting constants independent of θ for the above two inequalities. This KL bound means that our ST-trick simply ignores the time range beneath τ , and optimizes the proxy of the KL divergence between the forward and the reverse diffusion at time τ . On the other hand, no such interpretation can be applied to the diffusion loss without ST-trick, L(θ;λP, ).

heading: 3.3 RECIPROCAL VARIANCE EXPLODING SDE FOR SMALL TRUNCATION
text: This section introduces RVESDE whose variance, σ2(t), is geometric by its nature. RVESDE is defined as the form of dxt = g(t) dwt with
g(t) :=
{ σmax ( σmin σmax ) t
√ 2 log (σmaxσmin )
t if t > 0, 0 if t = 0.1
The transition probability of RVESDE becomes p0t(xt|x0) = N ( xt;x0, σ 2 max( σmin σmax ) 2 t I ) .
Recall that VESDE assumes g(t) = σmin(σmaxσmin ) t √
2 log (σmaxσmin ) and p0t(xt|x0) = N ( xt;x0, σ 2 min [ (σmaxσmin ) 2t− 1 ] I ) . As presented in Figure 8, RVESDE attains the geometric property at the expense of having reciprocated time, 1/t. We could successfully model the forward diffusion process with RVESDE without any approximation, and this exact modeling of the forward process enables applying the stochastic calculus to compute the log-likelihood or to generate new samples.
It is common to apply the importance sampling (Song et al., 2021) in order to minimize the MonteCarlo variance on the diffusion loss as illustrated in Figure 13 (c). RVESDE is advantageous on VE/VPSDEs that the importance sampling is easy to implement: the diffusion loss with RVESDE becomes
L(θ; g2, ) = 1 2
∫ T
g2(t)Lt(θ) dt = log (σmax σmin )∫ 1 1 T σ2(s−1)Ls−1(θ) ds,
by substituting the integrating variable into the reciprocated time. Therefore, the importance sampling for RVESDE is equivalent to drawing the diffusion time uniformly on a reciprocated interval, [ 1T , 1 ].

heading: 4 EXPERIMENTS
text: This section empirically studies our suggestions on benchmark datasets, including CIFAR-10 (Krizhevsky et al., 2009), STL-10 (Coates et al., 2011)2 CelebA (Liu et al., 2015) 64×64, CelebA-HQ (Karras et al., 2017) 256 × 256, LSUN (Yu et al., 2015) Bedroom/Church 256 × 256, and FFHQ (Karras et al., 2019). We compute bits-per-dim, proportional to NLL, by applying the Instantaneous Change of Variable Chen et al. (2018); Song et al. (2020) to the probability flow (Song et al., 2020) with the uniform dequantization (Theis et al., 2015; Hoogeboom et al., 2020), rather than the variational dequantization (Ho et al., 2019) in order to sidestep the need of training an auxiliary flow network. Also, we use the clean-FID introduced in Parmar et al. (2021) for computing the FID (Heusel et al., 2017) score with the predictor-corrector sampler (Song et al., 2020). Training and evaluation details are available in Appendix B.

heading: 4.1 QUANTITATIVE RESULTS
text: Table 3 compares HNCSN++ (RVE) with ST-trick against the current best generative models, and Table 3 shows that our model establishes the state-of-the-art performances. On CIFAR-10, we observe that our model improves NCSN++ (VE) in NLL and IS at the expense of sacrificing FID slightly. In particular, our model surpasses the previous best IS performance (StyleGAN2-ADA+Tuning (Karras et al., 2020)) by performing 10.11 on CIFAR-10. HNCSN++ trained on CelebA is the best performer in terms of the FID performance, but CR-NVAE (Sinha and Dieng, 2021) outperforms our model in NLL. However, we emphasize that HDDPM++ in Table 6 with ST-trick (1.52) outperforms CRNVAE (1.86) in NLL by far. In particular, HNCSN++ with high-dimensional CelebA-HQ 256× 256 performs the state-of-the-art in FID out of the baselines in the diffusion models (NCSN++) and the GAN models (PGGAN (Karras et al., 2017)). Finally, HNCSN++ on STL-10 dataset largely outperforms the baselines by improving the state-of-the-art FID from 15.17 (Park and Kim, 2021) to 7.71, and IS from 11.01 (Park and Kim, 2021) to 13.43.
In Table 3, when computing the Inception Score (IS) (Salimans et al., 2016), there is a minor discrepancy between likelihood-based models and likelihood-free models. IS-50k in diffusion models
1The existence and uniqueness of solution for RVESDE is guaranteed by Theorem 5.2.1 in Oksendal (2013). 2We downsize the dataset from 96× 96 to 48× 48 following Jiang et al. (2021); Park and Kim (2021).
(Song et al., 2020; Ho et al., 2020) computes the score with 50k generated images once, while IS-5k in GAN models (Karras et al., 2020; Park and Kim, 2021) computes the score 10 trials with independently generated 5k samples and report the performance as the average of scores. We report IS-50k following Song et al. (2020), but we note that IS-5k performs almost identical to IS-50k by performing 10.07 on CIFAR-10 and 13.34 on STL-10 in terms of IS-5k.
One particular observation in Table 3 is the IS performance (13.34) of STL-10 that exceeds the number of classes (10) of the dataset. We follow Park and Kim (2021), the current state-of-the-art model on STL-10, to train STL-10 with 105k images by aggregating the labeled (5k) and the unlabeled (100k) images (Jiang et al., 2021). Although the labeled images have 10 classes, the unlabeled images are sampled from a broader distribution of images, so the unlabeled dataset contains the other types of animals such as bears or rabbits, whose classes are different from the labeled classes. Therefore, a well-trained model with 105k images would perform IS that exceeds the number of labeled classes.

heading: 4.2 ABLATION STUDY
text: Throughout the ablation study, we train the model on CelebA as default otherwise stated.
High Precision Parametrization Table 4 presents the model performances after 0.5M of training iterations. First, NCSN++ with lowered σmin effectively reduces the performances on both NLL and FID, as anticipated in Section 2. Second, the new parametrization of HNCSN gives a slight performance gain compared to the original parametrization of NCSN. This performance gain comes from the better estimation for the data score, evidenced by Figure 13 (a). Third, the new parametrization of HDDPM largely improves the sample quality at the expense of slight NLL degradation.
Soft Truncation Figure 10 compares the diffusion model with/without ST-trick in terms of the sample fidelity by training iteration. To compare in a fair setting, we train the model in a pair of comparative settings: first, we train the model with vanilla diffusion loss of L(θ;λP, ) with λP(t) = (1− t )g2(t); and second, we train according to the same loss with ST-trick of λ(t) = g2(t) and P(τ) ∝ 1τ2 . At the initial stage of training, the sample fidelity with ST-trick is worse than that without ST-trick because the model with ST-trick learns the data score of small diffusion slower. However, ST-trick effectively updates the score network as training proceeds,
and Figure 10 shows that the model with ST-trick surpasses the model without ST-trick after the training.
Table 5 presents the performances after 0.5M of training steps. It compares the diffusion models with an identical weighting function, (1 − t )g2, with/without ST-trick, and it shows that ST-trick improves NLL from 1.98 to 1.93 and FID from 3.18 to 1.92, where 1.92 beats the current state-of-the-art FID (2.92) by far.
Table 6: Ablation study for ST-trick.
Model λ ST CelebANLL FID
NCSN++ (VE) σ2 7 2.39 3.95 3 2.41 2.68
HNCSN++ (RVE) g2 7 1.97 3.36 3 1.93 1.92
DDPM++ (VP) σ2 7 1.79 3.03 3 1.75 2.88
HDDPM++ (VP) σ2 7 1.78 3.23 3 1.73 2.22
DDPM++ (VP, IS) g2 7 1.53 5.31 3 1.52 4.50
HDDPM++ (VP, IS) g2 7 1.64 4.65 3 1.52 4.45
Table 6 compares our ST-trick with the vanilla optimization of the hard truncation on weighting functions of λ(t) = g2(t) and λ(t) = σ2(t). Table 6 shows that ST-trick is effective on both sample generation and density estimation. In particular, for any model settings, we conclude that STtrick is a highly effective method that significantly enhances the sample quality without any of model modification.
RVESDE Table 4 shows the efficacy of RVESDE on HNCSN++. Applying RVESDE improves the FID score while preserving the NLL performance, compared to VESDE on HNCSN++. To analyze the performance, recall that Proposition 2 implies that the inequality of NELBO in 1 does not hold anymore, so the NLL performance of VESDE loses its theoretic ground on bounding the log-likelihood. However, our RVESDE allows the inequality of NELBO, so the NLL performance in Table 4 is theoretically grounded. In addition, the experiments in Table 6 shows that RVESDE harmonizes with ST-trick better than VESDE.
Table 7: Ablation study for σmin trained on HNCSN++ (RVE).
ST σmin CIFAR-10
NLL FID-10k
7 10−3 2.96 7.04 10−4 2.97 8.17 10−5 2.97 8.29
3 10−3 2.99 5.09 10−4 2.97 5.54 10−5 2.96 5.98
Optimal Precision Table 7 finds that lowering σmin too small might harm the sample quality due to the training instability. We find that σmin = 10
−3 is a sweet spot for the optimal sample quality. This tendency maintains even when we apply the training stabilization method, such as ST-trick.
The reasoning of such opposite trend could be partially explained by how the data is saved. When we save the data, it is 8-bit quantized in order to downsize the file size, so it is enough to downscale σmin less than 1/256 to fit the model distribution to the data distribution, which is approximately 0.004. Therefore, lowering σmin too small incurs severe numerical instability without gaining additional performance enhancement.
Empirical result in Figure 11 indicates that σmin of 10−3 is indeed enough to create high precision images on high-dimensional FFHQ 256× 256 dataset. Figure 11 illustrates the histogram of the image-byimage quality metric, evaluated by Natural Image Quality Evaluator (NIQE) (Mittal et al., 2012), with 5k samples. Figure 11 implies that the denoised images up to the diffusion levels of 0.1 or 0.01 are not satisfactory in terms of this image-by-image metric. In fact, this discrepancy of the image-to-image quality metric explains the gain in FID when we lower σmin in Table 4, where FID is a distributional metric for sample quality.
High Dimensions To clarify the modeling effect on high dimensional datasets, we train HNCSN++ (RVE) with ST-trick for σmin = 10−3 on datasets of 256 × 256 resolutions. After the training of 3.3M iterations, we draw 5k samples by stopping the reverse process at σ = 10−1/10−2/10−3. When we denoise up to σ = 10−2, Table 8 shows that our model largely improves the sample quality: HNCSN++ outperforms by reporting 13.52 and 14.20 in FID-5k compared to the NCSN++ of 16.42 and 29.14 on FFHQ 256× 256 and LSUN Church 256×256, respectively. Such a huge discrepancy on high-dimensional
datasets is attributed to the poorly synthesized images from NCSN++, which is evidenced by the right
heavy tail in Figure 12. Also, as we denoise further from 10−2 to 10−3, the sample quality improves from 13.52 to 10.35 on FFHQ 256× 256 and 14.20 to 13.75 on LSUN Church 256× 256.
4.3 QUALITATIVE RESULTS
Table 8: Ablation study for denoising effect on highdimensional datasets.
Model σ FFHQ ChurchFID-5k FID-5k
HNCSN++ (RVE) + ST 10−1 60.56 48.06 10−2 13.52 14.20 10−3 10.35 13.75
NCSN++ 10−2 16.42 29.14 This section presents the qualitative experimental results for the proposed methods. Figure 13 (a) compares the score estimation for HNCSN++ and NCSN++ on CelebA. HNCSN++ succeeds in estimating the unbounded score, whereas NCSN++ shows poor estimation at the front diffusion time (or small diffusion level). Figure 13 (b) compares the diffusion model with/without ST-trick in terms of the Monte-Carlo loss density. As expected, the density is bimodal for the model without ST-trick, but it turns out that this heavy tail is largely improved by ST-trick. Figure 13 (c) illustrates that the Monte-Carlo loss variance is reduced with RVESDE, compared to VESDE. We attribute this variance reduction to the exact modeling of the diffusion process.
Figure 14 shows the generated images from various benchmark datasets. We find that our model succeeds generating realistic images on high-dimensional datasets with resolution greater or equal to 256× 256. See Appendix C for the full image generation results.

heading: 5 CONCLUSION
text: This paper tackles the problem of exploding data score with three points: 1) we provide a theorydriven parametrization for the high precision estimation; 2) we introduce a practical optimization method of the diffusion model parameters; and 3) we suggest a new type of SDE that enables to apply the stochastic calculus on small diffusion level. Applying the three improvements to NCSN/DDPM, we presented the high precision version of NCSN and DDPM, which excels in generating the high resolution images.

heading: 6 ETHICS STATEMENT
text: Potential risk from this work is the negative usage of the deep generative models, such as deepfake images. Our contribution enables the creation of the high resolution virtual images, and we acknowledge that there is any chance of misuse for malicious purposes.

heading: 7 REPRODUCIBILITY STATEMENT
text: Complete proofs of the theoretic analyses in the main paper are provided in Appendix A. The code for the experiments would be released to the reviewers at the discussion phase through an anonymous repository for the reproducibility. Most of our code is built based on the released code of the previous work (Song et al., 2020) including the network structure and hyperparameters, as explained in Appendix B.

heading: A PROOFS OF LEMMA 1, PROPOSITION 1, AND PROPOSITION 2
text: Lemma 1. Let H[0,T ] = {s : Rd × [0, T ] → Rd, s is locally Lipschitz}. Suppose a continuous vector field v defined on a subset U of a compact manifold M (i.e., v : U ⊂M → Rd) is unbounded, then there exists no s ∈ H[0,T ] such that limt→0 s(x, t) = v(x) a.e. on U .
Proof. Since U is an open subset of a compact manifold M , ‖x1 − x2‖ ≤ diam(M) for all x1,x2 ∈ U . Also, if t1, t2 ∈ [0, T ], |t1 − t2| is bounded. Hence, the local Lipschitzness of s implies that there exists a positive K > 0 such that ‖s(x1, t1) − s(x2, t2)‖ ≤ K(‖x1 − x2‖ + |t1 − t2|) for any x1,x2 ∈ U and t1, t2 ∈ [0, T ]. Therefore, for any s ∈ H[0,T ], there exists C > 0 such that ‖s(x, t)‖ < C for all x ∈ U and t ∈ [0, T ], which leads no s that satisfies s(x, t)→ v(x) a.e. on U as t→ 0. Proposition 2. LetH[1,∞) = {s : Rd× [1,∞)→ Rd, s is locally Lipschitz}. Suppose a continuous vector field v defined on a d-dimensional open subset U of a compact manifold M is unbounded, and the projection of v on each axis is locally integrable. Then, there exists s ∈ H[1,∞) such that limη→∞ s(x, η) = v(x) a.e. on U .
Proof. Let h be a standard mollifier function. If ht(x) = t−nh(x/t), then vt := ht ∗ v converges to v a.e. on U as t → 0 (Theorem 7-(ii) of Appendix C in Evans (1998)). Therefore, if we define s(x, η) := v1/η(x) on the domain of v1/η(x) and s(x, η) := 0 elsewhere, then s(x, η) = v1/η(x)→ v(x) a.e. on U as η →∞. Now, to show that s(x, η) is locally Lipschitz, let M̃×[η, η] be a compact subset of Rn×[1,∞). From ‖s(x1, η1) − s(x2, η2)‖ = ‖v1/η1(x1) − v1/η2(x2)‖ ≤ ‖v1/η1(x1) − v1/η1(x2)‖ + ‖v1/η1(x2) − v1/η2(x2)‖, if there exists K1,K2 > 0 such that ‖v1/η1(x1) − v1/η1(x2)‖ ≤ K1‖x1 − x2‖ and ‖v1/η1(x1) − v1/η2(x1)‖ ≤ K2|η1 − η2| for all x1,x2 ∈ M̃ and η1, η2 ∈ [η, η], then s(x, η) = v1/η(x) is Lipschitz on M̃ × [η, η]. First, since v1/η is infinitely differentiable on its domain (Theorem 7-(i) of Appendix C in Evans (1998)) and η ∈ [η, η], there exists K1 > 0 such that ‖v1/η(x1) − v1/η(x2)‖ ≤ K1‖x1 − x2‖. Second, the mollifier satisfies the uniform convergence on any compact subset of U (Theorem 7- (iii) of Appendix C in Evans (1998)), which leads that ‖v1/η1(x) − v1/η2(x)‖ ≤ K2| 1η1 − 1 η2 | = K2 |η1−η2| η1η2 ≤ K3|η1 − η2| for some K2,K3 > 0. Therefore, s becomes an element ofH[1,∞).
B IMPLEMENTATION DETAILS
Since the major contribution of this paper is developing the diffusion model in perspectives of theory and optimization method, we use the same model architecture introduced in Song et al. (2020).
B.1 EXPERIMENTAL DETAILS
Training Throughout the experiments, we train our model with a learning rate of 0.0002, warmup of 5000 iterations, gradient clipping by 1. On the experiments for HNCSN++ with σmin = 10−3, we train the diffusion model by truncating σmin to 10−2 at the initial stage of training, and after the saturation in terms of the FID score, we begin to train the score network with the whole range of [σmin, σmax]. We find this way of training works best in practice for the training stability. We follow the convention of Song et al. (2020) to choose σmax. Also, for DDPM-style model training, we use βmin = 0.1 and βmax = 20. For all experiments, we set = 10−5 and T = 1.
We stabilize the training by the Exponential Moving Average (EMA) (Song et al., 2020) with the rate of 0.999 on NCSN++/HNCSN++ and 0.9999 on DDPM++/HDDPM++. In addition, we stabilize the training by applying the dropout (Srivastava et al., 2014) with 0.1 rate. For the optimizer, we use the Adam optimizer (Kingma and Ba, 2014) with (β1, β2) = (0.9, 0.999).
On datasets of resolution 32× 32, we use the batch size of 128, which consumes about 48Gb GPU memory. On STL-10 with resolution 48 × 48, we use the batch size of 192, and on datasets of resolution 64× 64, we experiment with 128 batch size. The batch size for the datasets of resolution
256×256 is 40, which takes nearly 120Gb of GPU memory. On the dataset of 1024×1024 resolution, we use the batch size of 16, which takes around 120Gb of GPU memory. We use five NVIDIA RTX3090 GPU machines to train the model exceeding 48Gb, and we use a pair of NVIDIA RTX-3090 GPU machines to train the model that consumes less than 48Gb.
Evaluation We apply the EMA with rate of 0.999 on NCSN++/HNCSN++ and 0.9999 on DDPM++/HDDPM++. For the density estimation, we obtain the NLL performance by the Instantaneous Change of Variable (Song et al., 2020; Chen et al., 2018). For the sampling, we apply the Predictor-Corrector (PC) algorithm introduced in Song et al. (2020). We set the signal-to-noise ratio as 0.16 on 32 × 32 datasets, 0.17 on 48 × 48 and 64 × 64 datasets, 0.075 on 256×256 sized datasets, and 0.15 on 1024× 1024. On datasets less than 256× 256 resolution, we iterate 1,000 steps for the PC sampler, while we apply 2,000 steps on the other high-dimensional datasets. Throughout the experiments, we use the reverse diffusion (Song et al., 2020) for the predictor algorithm and the annealed Langevin dynamics (Welling and Teh, 2011) for the corrector algorithm.
We compute the FID score (Song et al., 2020) based on the modified Inception V1 network3 using the tensorflow-gan package for CIFAR-10 dataset, and we use the clean-FID (Parmar et al., 2021) based on the Inception V3 network (Szegedy et al., 2016) for the remaining datasets. We note that FID computed by Parmar et al. (2021) reports a higher FID score compared to the original FID calculation4.
B.2 CHOICE OF η(t)
HNCSN As described in Section 3.1, HNCSN assumes the score parametrization of sθ(xt, η(t)) with η(t) being the mixture of log σ(t) and 1σ(t) . Despite of the clear advantage of log σ(t) that the distribution of p(σ) is uniform in VE/RVESDEs, the drastic scale variation at t ≈ 0 prevents the enhanced score estimation. We empirically find that the reciprocal function of 1σ(t) works best on the estimation at the front diffusion time, and we combine these two functions to create a continuously differentiable η(t) by
η(t) := { log σ(t) if σ(t) ≥ σ0 − c1σ(t) + c2 if σ(t) < σ0,
where c1, c2 and σ0 are the hyperparameters. By acknowledging the parametrization of log σ(t) introduced in Song et al. (2020), we choose σ0 as 0.01. Also, to satisfy the continuously differentiability of η(t), the choice of two hyperparameters c1 and c2 satisfies a system of equations with degree 2, so c1 and c2 are fully determined with this system of equations. HDDPM Looking at the antiderivative ∫ g2(t) σ2(t) dt = log (1− e− ∫ t 0 β(s) ds) + ∫ t 0 β(s) ds = log σ2(t)+ ∫ t 0 β(s) ds, the unbounded function reduces to η(t) = 2 log σ(t)+ ∫ t 0 β(s) ds. Therefore, HDDPM resembles the temporal parametrization of log σ(t) (Song et al., 2020) at t ≈ 0 because the second term, ∫ t 0 β(s) ds, is nearly zero at t ≈ 0.

heading: C ADDITIONAL EXPERIMENTAL RESULTS
text: Figure 15 shows how images are created from the trained model on 256-dimensional datasets, and Figures from 16 to 21 present randomly generated samples of the trained model. Also, Tables 9 and 10 present the additional experimental results in terms of FID on LSUN Bedroom and FFHQ 256× 256.
3https://tfhub.dev/tensorflow/tfgan/eval/inception/1 4See https://github.com/GaParmar/clean-fid for the detailed experimental results.
Table 9: LSUN Bedroom
Model Bedroom 256× 256 FID UDM (RVE) + ST 4.57 Likelihood-based Models ADM Dhariwal and Nichol (2021) 1.90 DDPM Ho et al. (2020) 4.90 Likelihood-free Models StyleGAN Karras et al. (2019) 2.65 INR-GAN-bil Skorokhodov et al. (2020) 4.95 COCO-GAN Lin et al. (2019) 6.95
Table 10: FFHQ
Model FFHQ
256× 256 FID
UDM (RVE) + ST 5.54 Likelihood-free Models InsGen Yang et al. (2021) 3.31 Anycost GAN Lin et al. (2021) 3.35 StyleGAN2 ADA+bCR Karras et al. (2020) 3.62 CIPS Anokhin et al. (2020) 4.38 U-Net GAN Schonfeld et al. (2020) 7.48 BigGAN Schonfeld et al. (2020) 11.48

Review summaries:
This paper studies the behaviour of training SDE-based score-matching generative models for time t close to 0. The paper identifies two issues with existing forward SDEs for training such models: 1) the score of the marginal distribution could be large, and 2) the expected squared loss could be large, draining the loss contribution at other times. The paper introduces several remedies to fix this issue.    

Diffusion models have achieved great performance in many applications. However, learning a good diffusion model can be challenging since the data score can go to infinity when $t$ goes to zero. To address this issue, this paper proposes an alternative parameterization for unbounded data scores and provides a practical trick (ST-trick) to handle the variation of the scales of score values. It proposes Reciprocal Variance Exploding Stochastic Differential Equation to sample from the model. The proposed method can be applied to existing NCSN and DDPM models. Empirically, it achieves strong performance on high-resolution image generation.

This paper considers two problems associated with optimizing diffusion models, namely, training instability due to unbounded denoising score functions when the level of noise is very small (at the later stages of the reverse diffusion process), and problems at the earlier stages of the reverse diffusion process where coarse-scale artifacts are sometimes introduced as the corresponding terms of the objective have relatively little influence on the training loss. These problems are addressed through a new parameterization that is better suited to extremely small noise levels, and a soft truncation (ST) trick that better enables the optimization with sufficient influence given to the entire range of noise levels.  The paper has strong technical contributions and very good experimental work, showing improved quantitative measures of performance, namely NLL, FID, and IS.

The paper outlines four problems and proposes a solution for each problem:  (i) Generally, score-based generative models (for example [1]) do not compute the loss on the entire time interval $[0,1]$ but rather only on $[\varepsilon, 1]$, for some $\varepsilon > 0$. The paper claims that choosing $\varepsilon$ too large results in poor details for image generation (see Figure 4), for example, the details (pores) of a human nose.  (ii) The paper shows that the the squared Euclidean norm of the data score (for an unspecified data distribution) empirically diverges as $t \to 0$.  The paper then presents a lemma that states that no neural network can learn an unbounded function as long as it is conditioned on time $t$ or variance $\sigma$. The paper then proposes to condition on $\eta(t)$ where $\eta$ is constructed such that either $\eta(t) \to \infty$ or $\eta(t) \to -\infty$ as $t \to 0$.  (iii) The paper shows that Monte Carlo samples of the likelihood loss have a large variety of scales: values are very large for small times $t$ (or equivalently small noise levels) and very small for large times $t$ (or equivalently large noise levels). Due to the large variety of scales, the authors claim that "the diffusion loss is dominated by the end diffusion time with the front range in practice, and the loss at the end diffusion time is barely counted in the gradient signal. This imbalanced loss brings the immatured score estimation at the end diffusion time, which ruins the overall sample shape." (I am not entirely clear what this quote means.) The paper then proposes a so-called "soft-truncation" (ST) trick. The high-level intuition of the ST trick is described as follows: instead of sampling time points $t$ uniformly over the time range $[\varepsilon, 1]$ in each batch; in each batch it samples a $\tau \sim \mathbb{P}(\tau)$ and then samples $t$ uniformly over the time range $[\tau, 1]$. The authors claim that this has the effect that some batches focus on large time steps (providing a good gradient signal for large noise levels) and some batches focus on small time steps (providing a good gradient signal for small noise levels).  The original setup with fixed $\varepsilon$ can be recovered by setting $\mathbb{P}(\tau)= \delta(\tau = \varepsilon)$. The authors then propose a particular one-parameter ($k$) distribution $\mathbb{P}$ and show different instantiations of $k$. The family is designed in such a way that $\mathbb{P}(\tau) \to \delta(\tau = \varepsilon)$ as $k \to \infty$.  (iv)  The authors show that the continuous VESDE [1],  proposed  as a continuous extension of  [2],  is not actually “geometric” as apparently desired by [2]. The authors then propose the "Reciprocal Variance Exploding" SDE (RVESDE) which can be shown to be "geometric". The authors then claim that importance sampling is simple for the RVESDE.  The paper tests the proposed solutions on CIFAR10 (32x32), CelebA (64x64), CelebA-HQ (256-256), and STL-10 (48x48).  References:  [1] Song et al. Denoising Diffusion Implicit Models. ICLR 2021.  [2] Song & Ermon. Improved Techniques for Training Score-Based Generative Models. NeurIPS 2020.