High Precision Score-based Diffusion Models

1 INTRODUCTION . 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 . 2 PRELIMINARY . 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 . 2.2 PROBLEMS OF TRAINING THE DIFFUSION MODEL . 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 can not 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 can not be modeled by a reverse SDE . Therefore , VESDE loses its theoretic ground on the continuation of the discrete Markov chain .

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