Maximum Likelihood Training of Parametrized Diffusion Model

1 INTRODUCTION . Diffusion models have recently achieved successes on a task of sample generations , and some researches claim the state-of-the-art performance over Generative Adversarial Networks ( GAN ) ( Karras et al. , 2019 ) . This success is highlighted particularly in the community of likelihood-based models , including normalizing flows ( Grcić et al. , 2021 ) , autoregressive models ( Parmar et al. , 2018 ) , and variational autoencoders ( VAE ) ( Vahdat and Kautz , 2020 ) . Moreover , this success is noteworthy because it is achieved merely using linear diffusing mechanisms , such as VESDEs ( Song and Ermon , 2020 ) and VPSDEs ( Ho et al. , 2020 ) . This paper expands these linear diffusing mechanisms of VE/VPSDE to a data-adaptive trainable nonlinear diffusion . To motivate the expansion , though there are structural similarities between diffusion models and VAEs , the forward diffusion in a diffusion model has not been trained in existing literature while its counterpart , which is the encoder of VAE , is trainable . Rather , the current diffusion models assume the linear diffusing mechanism to be fixed throughout training procedure . Because of this static nature of the diffusing mechanism , variational gap ( Cremer et al. , 2018 ) between the log-likelihood and the Evidence Lower BOund ( ELBO ) remains to be strictly positive unless the score perfectly estimates the data score . This variational gap prevents the score training from being Maximum Likelihood Estimation ( MLE ) . This gap is a fundamental motivation to develop a trainable encoder of the diffusion model , so that the gap can be tight . As we tighten variational gap by training a nonlinear forward diffusion , the biggest challenge comes from an intractable optimization loss . The denoising diffusion loss requires a transition probability to be closed-form in order to achieve fast optimization , but a nonlinear diffusing mechanism , in general , has no closed-form perturbation probability . Hence , our innovation becomes designing a diffusion model with a tractable loss while the diffusion is learnable and flexible , which we achieve by merging previous diffusion models and the tractable variable transformation . Our innovation concentrates on the theoretic and the practical aspects of a new diffusion , which is tractable , learnable , and flexible . Theoretically , we observe that a nonlinear diffusing mechanism can be transformed to a linear diffusion under an invertible transformation , and vice versa . In practice , we implement PDM by merging a normalizing flow model and a diffusion model , which constructs a tractable , learnable , and flexible nonlinear diffusing mechanism . From this construction , we prove that variational gap can be further tightened by training the forward diffusion , nonlinearly . Also , we demonstrate the state-of-the-art performance in CelebA on Fréchet Inception Distance ( FID ) ( Heusel et al. , 2017 ) . We summarize our key contributions , as below . • PDM expands the scope of a forward diffusion process from linear and fixed dynamics to nonlinear and trainable dynamics . • PDM minimizes variational gap between the log-likelihood and ELBO , lower than the gap of linear diffusion models because PDM learns a data-adaptive nonlinear diffusion . • PDM explicitly proves in Theorem 2 that the optimal generative distribution equals to the real distribution , which was not the case for a linear diffusion model . 2 PRELIMINARY . A diffusion model is constructed with bidirectional stochastic processes . The forward direction diffuses an input variable to a noise variable , and the reverse direction denoises the random noise to construct a realistic input data instance . The diffusion model learns this reverse direction ( generative process ) by estimating data score . We provide a brief summary of diffusion models , as below . Forward Diffusion At the beginning of model build-up , a diffusion process requires a diffusing mechanism on an input variable . This paper assumes that the diffusion process is governed by a SDE of dxt = f ( xt , t ) dt + G ( xt , t ) dωt with x0 ∼ pr , where pr is the distribution from a real-world dataset ; and xt is the solution of the SDE . Reverse Diffusion The theory of stochastic calculus guarantees that we could create an identical diffusion process , { xt } Tt=0 , by solving an associated reverse SDE backwards in time ( Anderson , 1982 ) . The associated reverse SDE is dxt = [ f ( xt , t ) −G ( xt , t ) GT ( xt , t ) ∇xt log pt ( xt ) ] dt+ G ( xt , t ) dω̄t , xT ∼ pT , ( 1 ) where ωt and ω̄t are standard Wiener processes with time flows forward and backward , respectively ; and pt is a probability law of xt . Generative Diffusion A diffusion model approximates the above SDE 1 to eventually yield an estimation on the data distribution of pr . In SDE 1 , previous literature setup that drift and diffusion terms , f and G , are determined a-priori in the forward diffusion . However , data score , ∇xt log pt ( xt ) , is intractable to compute , so we estimate this data score with a score network of sθ ( xt , t ) in order to mimic the reverse diffusion with our generative process . This score network approximates the reverse diffusion by plugging the estimated score in place of data score with the below generative process : dxθt = [ f ( xθt , t ) −G ( xθt , t ) GT ( xθt , t ) sθ ( xθt , t ) ] dt+ G ( xθt , t ) dω̄t , x θ T ∼ π . ( 2 ) The generative process starts from a prior distribution ( π ) , and it constructs time-continuous random variables xθt by solving a SDE 2 backwards in time . The generated stochastic process is denoted by { xθt } Tt=0 , and we omit θ in the superscript if no confusion arises . With this generative process , we define a generative distribution , xθ0 ∼ pθ , as the probability density of the generated random variable at time t = 0 . Score Estimation We train the score network by a variational bound of the log-likelihood , given by Ex0 [ − log pθ ( x0 ) ] ≤L ( { xt } Tt=0 , λ = g2 ; θ ) − ExT [ log π ( xT ) ] = ∫ T 0 g2 ( t ) Lt ( { xt } Tt=0 ; θ ) dt− ExT [ log π ( xT ) ] , ( 3 ) where Lt ( { xt } Tt=0 ; θ ) = Ex0 , xt [ ‖sθ ( xt , t ) − ∇xt log p0t ( xt|x0 ) ‖22 ] up to a constant , where p0t ( xt|x0 ) is a transition probability from x0 to xt . Here , λ is a weighting function that determines the level of contribution for each diffusion time on the overall diffusion loss , L ( { xt } Tt=0 , λ ; θ ) ( Song et al. , 2020 ) . Variational bound holds when the weighting function is the likelihood weighting ( g ) ( Song et al. , 2021 ) , where G is a scalar-valued g function . 3 MOTIVATION OF NONLINEAR DIFFUSING MECHANISM . Though it has long been theoretically and empirically grounded to train the encoder part in VAE , such solid ground is not accomplished in diffusion models . This section analyzes structural similarities between a VAE model and a diffusion model , which brings the foundation of a data-adaptive nonlinear diffusing mechanism . 3.1 VARIATIONAL GAP OF VAE . Given Negative ELBO ( NELBO ) of Negative Log-Likelihood ( NLL ) in VAE as − log pθ ( x ) ≤Eqφ ( z|x ) [ − log pθ ( x|z ) ] +DKL ( qφ ( z|x ) ‖p ( z ) ) ( 4 ) =− log pθ ( x ) +DKL ( qφ ( z|x ) ‖pθ ( z|x ) ) , ( 5 ) we have a pair of interpretations on this NELBO . First , when we focus on Eq . 4 , NELBO 1 ) aids data reconstruction by optimizing Eqφ ( z|x ) [ − log pθ ( x|z ) ] ; and NELBO 2 ) regularizes the inference distribution of the encoder to a prior distribution by DKL ( qφ ( z|x ) ‖p ( z ) ) . On the other hand , if we concentrate on Eq . 5 , NELBO bounds NLL by approximating an intractable decoder posterior of pθ ( z|x ) with a tractable encoder posterior of qφ ( z|x ) . A vanilla VAE ( Kingma and Welling , 2013 ) assumes this approximate posterior to be a Gaussian distribution with mean and diagonal covariance estimated by amortized inference : qφ ( z|x ) = N ( z ; µφ ( x ) , σ 2 φ ( x ) I ) . By expanding the flexibility of this approximate posterior into a variational family of general distributions , for instance , Rezende and Mohamed ( 2015 ) resulted in tighter NELBO , which leads the optimization of VAE closer to MLE , and their choice of a generalizable model was normalizing flow . To connect this NELBO to a diffusion loss , we restate NELBO in the language of a stochastic process . Having that VAE attains a stochastic process of bivariate random variables , { x , z } , NELBO is reformulated to a KL divergence between two joint distributions modeled in bidirectional ways : DKL ( pr‖pθ ) ≤DKL ( pr‖pθ ) + Epr ( x ) [ DKL ( qφ ( z|x ) ‖pθ ( z|x ) ) ] =Epr ( x ) [ log pr ( x ) pθ ( x ) + Eqφ ( z|x ) [ log qφ ( z|x ) pθ ( z|x ) ] ] =Epr ( x ) qφ ( z|x ) [ log pr ( x ) qφ ( z|x ) pθ ( x ) pθ ( z|x ) ] =DKL ( qφ ( x , z ) ‖pθ ( x , z ) ) , ( 6 ) The inequality 6 , i.e. , DKL ( pr‖pθ ) ≤ DKL ( qφ ( x , z ) ‖pθ ( x , z ) ) , is well-known by itself in the field of information theory ( Duchi , 2016 ) by the name of data processing inequality . This restated bound interprets that VAE is essentially a bimodeling approach of a joint distribution on the bivariate stochastic process , { x , z } . On the forward direction ( x → z ) , a latent variable is conditioned on a data variable , and the joint distribution is modeled by qφ ( x , z ) = pr ( x ) qφ ( z|x ) . On the reverse direction ( z → x ) , a generative data variable is conditioned on the latent variable , and the joint distribution is modeled by pθ ( x , z ) = p ( z ) pθ ( x|z ) . Under this equivalent framework , we present an analytic tool to measure the closeness of NELBO and NLL ( see Table 5 ) by Gap ( qφ , pθ ) =NELBO− NLL =DKL ( qφ ( x , z ) ‖pθ ( x , z ) ) −DKL ( pr‖pθ ) =Epr ( x ) [ DKL ( qφ ( z|x ) ‖pθ ( z|x ) ) ] , and we denote this quantity as variational gap ( Cremer et al. , 2018 ) . Though it is one of central tasks in VAE to minimize this variational gap ( Kingma et al. , 2016 ) , only limited works ( Burda et al. , 2015 ; Neal , 2001 ) have estimated variational gap due to its computational burden . We introduce the tractable computation of variational gap in the next section on diffusion models .

Summary: The core contribution of the paper is to show how a nonlinear diffusion model may be constructed by cascading a normalizing flow with a linear diffusion in the resulting latent domain. The paper comes with a strong theoretical flavour and the proposed approach is well grounded. Applications in classic image generation tasks are provided.