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 paper proposes to combine Normalizing Flows with generative Diffusion Models in such a way that the target data is first non-linearly transformed via the flow and then the distribution over latent embeddings is modeled with the diffusion model. The authors call their model *Parametrized Diffusion Model (PDM)*. Since the flow is invertible and defined by a deterministic function, it is possible to formally relate the combined flow and diffusion process to a non-linear diffusion directly in data space, leveraging Ito's Lemma. The paper validates the PDM on image modeling benchmarks (CIFAR-10 and CelebA 64x64). The experiments show that the method leads to reduced variational gaps between the data likelihood and the evidence lower bound that is used for training the model, compared to several baselines.