Diffusion Normalizing Flow

1 Introduction . Generative model is a class of machine learning models used to estimate data distributions and sometimes generate new samples from the distributions [ 8 , 35 , 16 , 37 , 7 ] . Many generative models learn the data distributions by transforming a latent variable z with a tractable prior distribution to the data space [ 8 , 35 , 32 ] . To generate new samples , one can sample from the latent space and then follow the transformation to the data space . There exist a large class of generative models where the latent space and the data space are of the same dimension . The latent variable and the data are coupled through trajectories in the same space . These trajectories serve two purposes : in the forward direction x → z , the trajectories infer the posterior distribution in the latent space associated with a given data sample x , and in the backward direction z → x , it generates new samples by simulating the trajectories starting from the latent space . This type of generative model can be roughly divided into two categories , depending on whether these trajectories connecting the latent space and the data space are deterministic or stochastic . When deterministic trajectories are used , these generative models are known as flow-based models . The latent space and the data space are connected through an invertible map , which could either be realized by the composition of multiple invertible maps [ 35 , 8 , 20 ] or a differential equation [ 4 , 14 ] . In these models , the probability density at each data point can be evaluated explicitly using the change of variable theorem , and thus the training can be carried out by minimizing the negative log-likelihood ( NLL ) directly . One limitation of the flow-based model is that the invertible map parameterized by neural networks used in it imposes topological constraints on the transformation from z to x . Such limitation affects the performance significantly when the prior distribution on z is a simple unimodal distribution such as Gaussian while the target data distribution is a well-separated multi-modal distribution , i.e. , its support has multiple isolated components . In [ 6 ] , it is shown that there are some fundamental issues of using well-conditioned invertible functions to approximate such complicated multi-modal data distributions . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . When stochastic trajectories are used , the generative models are often known as the diffusion model [ 38 ] . In a diffusion model , a prespecified stochastic forward process gradually adds noise into the data to transform the data samples into simple random variables . A separate backward process is trained to revert this process to gradually remove the noise from the data to recover the original data distributions . When the forward process is modeled by a stochastic differential equation ( SDE ) , the optimal backward SDE [ 1 ] can be retrieved by learning the score function [ 39 , 40 , 17 , 2 ] . When the noise is added to the data sufficiently slow in the forward process , the backward diffusion can often revert the forward one reasonably well and is able to generate high fidelity samples . However , this also means that the trajectories have to be sufficiently long with a large number of time-discretization steps , which leads to slow training and sampling . In addition , since the forward process is fixed , the way noise is added is independent of the data distribution . As a consequence , the learned model may miss some complex but important details in the data distribution , as we will explain later . In this work , we present a new generative modeling algorithm that resembles both the flow-based models and the diffusion models . It extends the normalizing flow method by gradually adding noise to the sampling trajectories to make them stochastic . It extends the diffusion model by making the forward process from x to z trainable . Our algorithm is thus termed Diffusion Normalizing Flow ( DiffFlow ) . The comparisons and relations among DiffFlow , normalizing flow , and diffusion models are shown in Figure 1 . When the noise in DiffFlow shrinks to zero , DiffFlow reduces to a standard normalizing flow . When the forward process is fixed to some specific type of diffusion , DiffFlow reduces to a diffusion model . In DiffFlow , the forward and backward diffusion processes are trained simultaneously by minimizing the distance between the forward and the backward process in terms of the Kullback-Leibler ( KL ) divergence of the induced probability measures [ 42 ] . This cost turns out to be equivalent to ( see Section 3 for a derivation ) the ( amortized ) negative evidence lower bound ( ELBO ) widely used in variational inference [ 21 ] . One advantage to use the KL divergence directly is that it can be estimated with no bias using sampled trajectories of the diffusion processes . The KL divergence in the trajectory space also bounds the KL divergence of the marginals , providing an alternative method to bound the likelihood ( see Section 3 for details ) . To summarize , we have made the following contributions . 1 . We propose a novel density estimation model termed diffusion normalizing flow ( DiffFlow ) that extends both the flow-based models and the diffusion models . The added stochasticity in DiffFlow boosts the expressive power of the normalizing flow and results in better performance in terms of sampling quality and likelihood . Compared with diffusion models , DiffFlow is able to learn a forward diffusion process to add noise to the data adaptively and more efficiently . This avoids adding noise to regions where noise is not so desirable . The learnable forward process also shortens the trajectory length , making the sampling much faster than standard diffusion models ( We observe a 20 times speedup over diffusion models without decreasing sampling quality much ) . 2 . We develop a stochastic adjoint algorithm to train the DiffFlow model . This algorithm evaluates the objective function and its gradient sequentially along the trajectory . It avoids storing all the intermediate values in the computational graph , making it possible to train DiffFlow for highdimensional problems . 3 . We apply the DiffFlow model to several generative modeling tasks with both synthetic and real datasets , and verify the performance of DiffFlow and its advantages over other methods . 2 Background . Below we provide a brief introduction to normalizing flows and diffusion models . In both of these models , we use τ = { x ( t ) , 0 ≤ t ≤ T } to denote trajectories from the data space to the latent space in the continuous-time setting , and τ = { x0 , x1 , · · · , xN } in the discrete-time setting . Normalizing Flows : The trajectory in normalizing flows is modeled by a differential equation ẋ = f ( x , t , θ ) , ( 1 ) parameterized by θ . This differential equation starts from random x ( 0 ) = x and ends at x ( T ) = z. Denote by p ( x ( t ) ) the probability distribution of x ( t ) , then under mild assumptions , it evolves following [ 4 ] ∂ log p ( x ( t ) ) ∂t = −tr ( ∂f ∂x ) . ( 2 ) Using this relation ( 1 ) ( 2 ) one can compute the likelihood of the model at any data point x . In the discrete-time setting , the map from x to z in normalizing flows is a composition of a collection of bijective functions as F = FN ◦FN−1 · · ·F2 ◦F1 . The trajectory τ = { x0 , x1 , · · · , xN } satisfies xi = Fi ( xi−1 , θ ) , xi−1 = F −1 i ( xi , θ ) ( 3 ) for all i . Similar to Equation ( 2 ) , based on the rule for change of variable , the log-likelihood of any data samples x0 = x can be evaluated as log p ( x0 ) = log p ( xN ) − N� i=1 log |det ( ∂F −1 i ( xi ) ∂xi ) | . ( 4 ) Since the exact likelihood is accessible in normalizing flows , these models can be trained by minimizing the negative log-likelihood directly . Diffusion Models : The trajectories in diffusion models are modeled by stochastic differential equations . More explicitly , the forward process is of the form dx = f ( x , t ) dt+ g ( t ) dw , ( 5 ) where the drift term f : Rd × R → Rd is a vector-valued function , and the diffusion coefficient g : R → R is a scalar function ( in fact , g is often chosen to be a constant ) . Here w denotes the standard Brownian motion . The forward process is normally a simple linear diffusion process [ 38 , 16 ] . The forward trajectory τ can be sampled using ( 5 ) initialized with the data distribution . Denote by pF the resulting probability distribution over the trajectories . With a slight abuse of notation , we also use pF to denote the marginal distribution of the forward process . The backward diffusion from z = x ( T ) to x = x ( 0 ) is of the form dx = [ f ( x , t ) − g2 ( t ) s ( x , t , θ ) ] dt+ g ( t ) dw . ( 6 ) It is well-known that when s coincides with the score function ∇ log pF , and x ( T ) in the forward and backward processes share the same distribution , the distribution pB induced by the backward process ( 6 ) is equal to pF [ 1 ] , [ 29 , Chapter 13 ] . To train the score network s ( x , t , θ ) , one can use the KL divergence between pF and pB as an objective function to reduce the difference between pF and pB . When the difference is sufficiently small , pF and pB should have similar distribution over x ( 0 ) , and one can then use the backward diffusion ( 6 ) to sample from the data distribution . In the discrete setting , the trajectory distributions can be more explicitly written as pF ( τ ) = pF ( x0 ) N� i=1 pF ( xi|xi−1 ) , pB ( τ ) = pB ( xT ) N� i=1 pB ( xi−1|xi ) . ( 7 ) The KL divergence between pF and pB can be decomposed according to this expression ( 7 ) . Most diffusion models use this decomposition , and meanwhile take advantage of the simple structure of the forward process ( 5 ) , to evaluate the objective function in training [ 39 , 40 , 41 ] . 3 Diffusion normalizing flow . We next present our diffusion normalizing flow models . Similar to diffusion models , the DiffFlow models also has a forward process dx = f ( x , t , θ ) dt+ g ( t ) dw , ( 8 ) and a backward process dx = [ f ( x , t , θ ) − g2 ( t ) s ( x , t , θ ) ] dt+ g ( t ) dw . ( 9 ) The major difference is that , instead of being a fixed linear function as in most diffusion models , the drift term f is also learnable in DiffFlow . The forward process is initialized with the data samples at t = 0 and the backward process is initialized with a given noise distribution at t = T . Our goal is to ensure the distribution of the backward process at time t = 0 is close to the real data distribution . That is , we would like the difference between pB ( x ( 0 ) ) and pF ( x ( 0 ) ) to be small . To this end , we use the KL divergence between pB ( τ ) and pF ( τ ) over the trajectory space as the training objective function . Since KL ( pF ( x ( t ) ) |pB ( x ( t ) ) ) ≤ KL ( pF ( τ ) |pB ( τ ) ) ( 10 ) for any 0 ≤ t ≤ T by data processing inequality , small difference between pB ( τ ) and pF ( τ ) implies small difference between pB ( x ( 0 ) ) and pF ( x ( 0 ) ) in terms of KL divergence ( more details are included in Appendix B ) .

Summary: The paper introduces a non-linear drift component to the SDE formulation of diffusion-based generative modelling. Earlier work (DDPM) used a fixed linear drift component. This paper uses a trainable neural network as the drift component, and another one to estimate the score. Both these neural networks are trained by minimizing the KL divergence between the forward and backward trajectories, while accounting for their SDE formulations. It is found that this trained network performs reasonably well compared to prior works, and achieves better sample quality and shorter trajectories compared to DDPM as well as Continuous Normalizing Flows (FFJORD).