Year: 2021

Abstract:
Whereas diverse variations of diffusion models exist in image synthesis, the previous variations have not innovated on a topic of their diffusing mechanism that was always kept linear. Meanwhile, it is intuitive that there would be more promising diffusion patterns to result in a distribution closer to the approximating data distribution. This paper introduces such adaptive and nonlinear diffusion method for score-based diffusion models. Unlike static and linear Variance Exploring (VE) or Variance Preserving (VP) Stochastic Differential Equations (SDE) of the previous diffusion models, our Parametrized Diffusion Model (PDM) learns the optimal diffusion process by combining a normalizing flow and a diffusion process. Specifically, PDM utilizes the flow to nonlinearly transform a data variable into a latent variable, and PDM applies the diffusion process to the transformed latent variable with the linear diffusing mechanism. Subsequently, PDM enjoys a nonlinear and learned diffusion from the perspective of the data variable. We train PDM with a variational proxy of the log-likelihood, and we prove that variational gap between the variational bound and the log-likelihood becomes tight when the normalizing flow is optimal.

heading: 1 INTRODUCTION
text: 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.

heading: 2 PRELIMINARY
text: 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.

heading: 3 MOTIVATION OF NONLINEAR DIFFUSING MECHANISM
text: 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.

heading: 3.1 VARIATIONAL GAP OF VAE
text: 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.

heading: 3.2 VARIATIONAL GAP OF DIFFUSION MODELS
text: This section constructs an analogy of a diffusion model to VAE. To begin with, we observe that the classical data processing inequality in VAE holds on a stochastic process with bivariate random variables on a finite-dimensional Euclidean space. On the other hand, the time horizon of a diffusion
model is continuous, so it contains an infinite number of random variables in its stochastic process as {xt}Tt=0. Fortunately, the finite-dimensional classical data processing inequality on the bivariate random variables is directly extendable to a stochastic process with a continuous index. Formally, an infinite-variate joint distribution is called a path measure (Oksendal, 2013), which is defined as a probability law on a space of continuous functions induced by {xt}Tt=0. Therefore, the path measure is an alternative concept of a joint distribution for a stochastic process with a continuous index. The path measure for the forward diffusion in a diffusion model corresponds to a joint distribution for the encoder part in VAE, and similarly, the path measure for a generative diffusion corresponds to the decoder part in VAE. If µ is the forward (inference) path measure, and if νθ is the generative path measure; then the data processing inequality becomes
DKL(pr‖pθ) ≤ DKL(µ‖νθ). The real distribution of pr is a marginal distribution of µ at t = 0, and the generative distribution of pθ becomes a marginal distribution of νθ at t = 0. Under this concept, the data processing inequality means that a distributional divergence of DKL(µ‖νθ) between two path measures lower bounded by a divergence of DKL(pr‖pθ) between two sliced path measures at t = 0. Previous works (Song et al., 2021; Huang et al., 2021) regard this variational bound being equivalent to a diffusion loss,
DKL(µ‖νθ) = L ( {xt}Tt=0, g2;θ ) +DKL(pT ‖π),
where G(xt, t) = g(t) is a scaler function on t, and λ(t) is a square of the diffusion term g2(t). Since the previous diffusion models optimize the diffusion loss, the KL divergence of DKL(µ‖νθ) is an equivalent target of optimization on the diffusion loss. Subsequently, we can credit variational gap as a training performance to tighten NELBO to NLL by
Gap(µ,νθ) = DKL(µ‖νθ)−DKL(pr‖pθ). This variational gap satisfies (see Appendix B.2)
Gap(µ,νθ) = Epr(x0) [ DKL ( µ({xt}|x0)‖νθ({xt}|x0) )] ,
so variational gap is minimized if an inference posterior, µ({xt}|x0), estimates the intractable generative posterior, νθ({xt}|x0). However, unlike VAE, the inference posterior is not trainable, which restricts variational gap to be strictly positive in previous diffusion models. We note that a diffusion model provides tractable NLL using the Instantaneous Change of Variable (Song et al., 2020), and this tractable variational gap becomes useful in analyzing a diffusion model quality.

heading: 4 PARAMETRIZED DIFFUSION MODEL

heading: 4.1 NO CLOSED-FORM TRANSITION PROBABILITY ON NONLINEAR DIFFUSION
text: As far as we know, all previously proposed SDEs, including VESDE and VPSDE, are categorized as linear SDEs because they have linear drift term of f(xt, t). This drift term determines the dynamics of the diffusion, so a data-adaptive diffusion could be attained simply by parametrizing drift and diffusion terms. However, a transition probability with nonlinear coefficients for a SDE, p0t(xt|x0), is intractable in general, and a diffusion loss becomes intractable to optimize. In fact, it is worth noting that a closed-form of the transition probability exists only on a certain type of SDEs, like VESDE and VPSDE.

heading: Invertible Path Invertible Path

heading: 4.2 PARAMETRIZED FORWARD DIFFUSION WITH NORMALIZING FLOW
text: To construct a tractable transition probability while keeping nonlinearity, we observe that a nonlinear diffusion can be decomposed into a linear diffusion and a nonlinear invertible transformation, where the transformation provides the flexible nonlinearity on a diffusion process. Figure 1 and 2 provide a schematic view of our model. Normalizing flow transforms a data input to a latent variable, and a diffusion model linearly permeates the latent variable to a noise distribution.
Concretely, let us define z0 to be a transformed latent variable, z0 = hφ(x0), where φ is flow model parameters. Then, a canonical diffusion with linear dynamics of
dzt = − 1
2 β(t)zt dt+ g(t) dωt, z0 = hφ(x0) with x0 ∼ pr,
describes a diffusing mechanism of a latent variable. Since the transformation is invertible, if we define an inverse random variable by xφt = h −1 φ (zt), the dynamics of a stochastic process of {x φ t }Tt=0 is governed by a SDE of
dxφt = fφ(x φ t , t) dt+ Gφ(x φ t , t) dωt, x φ 0 ∼ pr,
where drift and diffusion terms are derived by{ fφ(xt, t) = − 12β [ ∇hφ ]−1 hφ + 1 2g 2tr ( ∇2h−1φ ) Gφ(xt, t) = g [ ∇hφ ]−1 ,
from the Ito’s lemma (Oksendal, 2013). See full derivation and explanation in Appendix B.1. This diffusion process is induced diffusion on the data space that is defined indirectly by transforming back from the latent space to the data space. The nonlinearity of this induced diffusion fully depends on the invertible transformation as illustrated in Figure 2.
It is not straightforward to directly construct a generative reverse process on the data space, xφ,θt . Figures 1 and 2 build the generative process by red arrows using the transformation. Starting from a prior distribution of π, we denoise a random variable from the prior to zθ0 according to the generative process of the linear diffusion on the latent space, and we transform the fully denoised latent of zθ0 back to the data space xφ,θ0 = h −1 φ (z θ 0).

heading: 4.3 TRACTABLE VARIATIONAL BOUND OF PARAMETRIZED DIFFUSION MODEL
text: The joint modeling of normalizing flow and diffusion enables a tractable training of nonlinear dynamics. The log-likelihood on a data space reduces to the log-likelihood on a latent space by adding the log determinant of Jacobian to the latent log-likelihood. This calculation of log-likelihoods provides Theorem 1 that bounds the log-likelihood with a tractable ELBO.
Theorem 1. Suppose that pφ,θ(x0) is the log-likelihood of a generative random variable xφ,θ0 . Then, the log-likelihood is bounded by
Epr(x0) [ − log pφ,θ(x0) ] ≤ L ( {xt}Tt=0, g2; {φ,θ} ) ,
where L ( {xt}Tt=0, g2; {φ,θ} ) =− Epr(x0) [ log ∣∣∣ det(∂hφ
∂x0 )∣∣∣]+ L({zt}Tt=0, g2;θ)− EzT [ log π(zT )], (7)
with L ( {zt}Tt=0, g2;θ ) = ∫ T 0 g2(t)Ez0,zt [ ‖sθ(zt, t) − ∇zt log p0t(zt|z0)‖22 ] dt, up to a constant. Here, p0t(zt|z0) is a transition probability of the linear diffusion process in the latent space.
PDM ELBO in Eq. 7 is different from the vanilla diffusion ELBO in Eq. 3 from three perspectives. First, the normalizing flow loss, i.e., the log determinant of Jacobian, is added to the loss of PDM ELBO. This normalizing flow loss represents a transformation of a data space into a latent space. Second, the remaining terms, related to the diffusion model, is a linear diffusion on the latent space with linear dynamics, which provides tractable optimization, rather than an intractable data diffusion of nonlinear dynamics. Third, the optimization loss in Eq. 7 estimates the score on the latent space by a score network of sθ(zt, t), while the diffusion loss in Eq. 3 estimates the score on the data space by a score network of sθ(xt, t).

heading: 4.4 OPTIMIZATION LOSSES
text: Previous works (Ho et al., 2020; Song et al., 2021; Vahdat et al., 2021) have empirically demonstrated that training a diffusion model with a likelihood weighting, λ(t) = g2(t), could bring worse performances than applying other weighting functions. For instance, Ho et al. (2020); Song and Ermon (2020) introduce a variance weighting function, λ(t) = σ2(t), and Ho et al. (2020) show that this weighting outperforms a likelihood weighting in terms of the sample generation.
Within our framework, we also observe that a variance weighting function is advantageous over a likelihood weighting. Henceforth, we optimize a flow model and a diffusion model with different losses as follows:
Flow: min φ L({xt}Tt=0, g2; {φ,θ})
Diffusion: min θ L({zt}Tt=0, {g2 or σ2};θ),
Note that the first term of the log Jacobian and the third term of the cross entropy in Eq. 7 are independent to score parameters, θ. Therefore, diffusion parameters are not influenced by these terms, in which only the second term remains to update the diffusion parameters.

heading: 4.5 TIGHT VARIATIONAL GAP OF PARAMETRIZED FORWARD DIFFUSION
text: This section theoretically analyzes variational gap of PDM. To construct variational gap of PDM, we enumerate few nomenclatures for path measures (see Table 6). On a data space, µdφ is a path measure for an inference process; and νdφ,θ is a path measure for a generative process. On a latent space, µlφ is a path measure for an inference process; and ν l θ is a path measure for a generative process. Here, νlθ does not depend on φ because the generative diffusion on the latent space starts from the φ-independent noise distribution (π); and because the diffusion proceeds a denoising process by following a score estimation of sθ(zt, t) that depends only on θ.
Variational gap on the data space becomes
Gapd(µ d φ,ν d φ,θ) =NELBO− NLL (8)
=DKL(µ d φ‖νdφ,θ)−DKL(pr‖pφ,θ), (9)
where the subscript in Gapd represents the data space. Recall that training the encoder in VAE guarantees tighter variational gap because the encoder distribution of qφ(x, z) approaches toward the decoder distribution of pθ(x, z) by minimizing NELBO of DKL ( qφ(x, z)‖pθ(x, z) ) . However, this analogy is not valid in variational gap in Eq. 8 because both inference and generative path measures depend on φ. This means that the generative path measure could arbitrarily deviate from the inference path measure when we optimize φ.
Using the change of variable, however, variational gap is equivalent to
Gapd(µ d φ,ν d φ,θ) =DKL(µ d φ‖νdφ,θ)−DKL(pr‖pφ,θ)
=DKL(µ l φ‖νlθ)−DKL(pφ‖pθ)
=Gapl(µ l φ,ν l θ).
With this derivation, optimizing an invertible transformation guarantees tighter variational gap because the generative measure of νlθ is now fixed on φ. Also, we note that this variational gap reduces to a KL divergence between the inference posterior and the generative posterior by
Gapd(µ d φ,ν d φ,θ) = Gapl(µ l φ,ν l θ) = Ez0 [ DKL ( µφ({zt}|z0)‖νθ({zt}|z0) )] , (10)
which provides another analogy to VAE. We enumerate variational gaps of the explained models in Table 5. In fact, we show in Theorem 2 that the optimum of PDM guarantees that the generative distribution pφ∗,θ equals to the data distribution pr, and variational gap reduces to zero in the optimum. See Appendix A for proofs.
Theorem 2. Suppose a flow network is flexible enough to transform pr to arbitrary continuous distribution. For any θ, if∇·sθ(·, t) is symmetric for any t ∈ [0, T ] and ∇· log π(·) ≡ sθ(·, T ), then there exists φ∗ such that Gapd(µ d φ∗ ,ν d φ∗,θ) = 0. Furthermore, we have
φ∗ = argmin φ Gapd(µ d φ,ν d φ,θ) = argmin φ DKL(µ d φ‖νdφ,θ) = argmin φ DKL(pr‖pφ,θ).
To the best of our knowledge, Theorem 2 is the first analysis in the community of diffusion models that the theoretic optimum of a model completely recovers the data distribution. To proceed this argument, let us consider a linear diffusion model in a data space with no flow model. Then, a marginal distribution of the forward process is pT at t = T , which differs from a prior distribution of π. In addition, suppose that the diffusion model reaches to the optimum point, i.e., its score network estimates the exact data score by sθ∗(xt, t) = ∇xt log pt(xt). Then, a generative transition probability equals to the groundtruth (reverse) transition probability of pT0(x0|xT ). Therefore, the generative distribution, pθ∗(x0) = ∫ π(xT )pT0(x0|xT ) dxT , naturally diverges from the real
distribution, pr(x0) = ∫ pT (xT )pT0(x0|xT ) dxT , because of the mismatch of pT and π. This implies that the original linear diffusion model cannot estimate the real distribution, pr(x0). In contrast to the original diffusion model, applying normalizing flow to a diffusion model enables PDM to recover the real distribution by Theorem 2. See Figures 8 and 9 for further illustration.
5 RELATED WORK 𝑥𝑥0
Recently, a diffusion model has been jointly modeled in conjunction with other deep generative models. There are two major previous works on this line of research. First, Vahdat et al. (2021) (LSGM) in
Figure 3 applies a diffusion model to the latent space of VAE. Since VAE does not attain an invertible structure, LSGM is not suggesting the expansion of diffusing mechanisms. In other words, LSGM cannot provide diffused data learning at the middle of the process because LSGM does not induce the diffusion process on the data space due to the lack of VAE’s invertibility. Consequently, the analyzing tool for LSGM should be variational gap of VAE, rather than that of the diffusion model.
Second, Song et al. (2021) (ScoreFlow) in Figure 4 simply concatenate normalizing flow on top of a diffusion model to provide a dequantized data distribution only at the initial diffusion time. Despite of the structural similarity of ScoreFlow and PDM, a major difference of ScoreFlow and PDM lies on their density model constructions: pθ in Scoreflow and pφ,θ in PDM. ScoreFlow constructs
pθ only with a diffusion model, so the flow and the diffusion networks are not jointly combined in the loss function of ∫ pr(x) ∫ [0,1)d
qφ(u|x) log pθ(x+u)qφ(u|x) dudx, where u is a dequantization variable that lies in a unit cube, [0, 1)d. In other words, ScoreFlow adopts the flow outside of a model density, so the flow is independent to the model density. In fact, ScoreFlow can be best understood as a dequantization technique (Ho et al., 2019), whose goal is estimating the optimal dequantizing region
Table 1: Performance comparison on CIFAR-10.
SDE Name Model NLL NELBO Gap FID
VE NCSN++ Diff Baseline 3.45 4.43 0.98 2.20PDM (FID) Diff & Flow 3.04 3.09 0.05 3.36
VP
DDPM++ (FID) Diff Baseline 3.13 3.29 0.16 3.16 PDM (FID) Diff & Flow 3.04 3.09 0.05 6.79 PDM (FID, dec) Diff & Flow-v2 3.11 3.17 0.06 4.43
DDPM++ (NLL) Diff Baseline 3.05 3.22 0.17 9.68 PDM (NLL) Diff & Flow 2.98 2.99 0.01 9.29 PDM (NLL, dec) Diff & Flow-v2 2.96 2.99 0.03 7.23
DDPM++ (NLL, IS) Diff Baseline 2.91 3.10 0.19 5.70 PDM (NLL, IS) Diff & Flow 2.99 3.01 0.02 8.94 PDM (NLL, dec, IS) Diff & Flow-v2 2.94 2.95 0.01 6.84
LSGM-109M (balanced) Diff & VAE - 2.96 - 4.60 LSGM-480M (balanced) Diff & VAE - 2.95 - 2.17 LSGM-480M (NLL) Diff & VAE - 2.87 - 6.89 LSGM-480M (FID) Diff & VAE - 3.43 - 2.10
Table 2: Performance comparison on CelebA.
Model NLL NELBO Gap FID
NCSN++ (VE) 2.39 3.96 1.57 3.95 PDM (VE, FID) 2.00 2.09 0.09 2.50
DDPM++ (VP) 1.79 2.19 0.40 3.03 PDM (VP, FID) 2.04 2.10 0.06 2.23
UDM 1.93 - - 2.78 DDGM - - - 2.92 CR-NVAE - 1.86 - - NCP-VAE - - - 5.25 DenseFlow-74-10 1.99 - - - Styleformer - - - 3.66
on [0, 1)d to minimize dequantization gap (Hoogeboom et al., 2020). Therefore, variational gap of ScoreFlow remains to be strictly positive because training the flow does not optimize the forward diffusing mechanism. On the other hand, PDM constructs flow-embedded model density of pφ,θ , and PDM optimizes the loss of ∫ pr(x) log ∫ [0,1)d
pφ,θ(x + u) dudx that puts normalizing flow inside of the model density. Although we train PDM with a uniform dequantization (Theis et al., 2015) by default, we could apply the variational dequantization to PDM as well.

heading: 6 EXPERIMENTS
text: This section quantitatively and qualitatively analyzes suggested PDM on CIFAR-10 (Krizhevsky et al., 2009) and CelebA (Liu et al., 2015) 64 × 64. We compute NLL for density estimations in bits per dimension scale by solving probability flow (Song et al., 2020) backwards in time, and we compare baseline performances of Song et al. (2020; 2021) computed on a uniform dequantization. For sample generations, we use the Predictor-Corrector algorithm (Song et al., 2020) for PDM on VESDE, and we use the numerical ODE sampler of the probability flow (Song et al., 2020) with the RK45 method on VPSDE. Throughout experimental results, we use NCSN++ (VE) and DDPM++ (VP) (Song et al., 2020) as backbones of diffusion models, and a pair of ResNet-based flow models (Chen et al., 2019; Ma et al., 2020) as backbones of flow models. We give a full description on models for our experiments in Appendix C. We name PDM (FID) and PDM (NLL) for experiments that use weighting functions as σ2 and g2, respectively.

heading: 6.1 QUANTITATIVE RESULTS
text: We compare PDM to baselines in Table 1 on VESDE and VPSDE as linear diffusing mechanisms. Although minor performance degradation on the sample performance, Table 1 empirically demonstrates reduced variational gap. In all experimental settings, variational gaps significantly drop less than 0.1. On the other hand, variational gaps of the baseline diffusions on VPSDE for the FID model and the NLL model are 0.16 and 0.19, respectively. We conclude that PDM effectively decreases variational gap by training a parametrized inference path measure.
Table 2 compares experiments on CelebA to the current state-of-the-art baseline models in various disciplines: the diffusion model of UDM (Kim et al., 2021) and DDGM (Nachmani et al., 2021), the VAE model of CR-NVAE (Sinha and Dieng, 2021), the flow model of DenseFlow (Grcić et al., 2021), and the GAN model of Styleformer (Park and Kim, 2021). We observe a pair of implications. First, training a diffusing mechanism with PDM largely reduces variational gap. In particular, PDM reduces this inference gap from 1.57 to 0.09 on the NCSN++ backbone with VESDE and from 0.40 to 0.05 on the DDPM++ backbone with VPSDE. Also, we emphasize that NELBO of PDM largely outperforms NELBO of the baseline models. Second, our PDM reports the state-of-the-art performance with 2.23 FID score in the CelebA sample generation. This opposite trend on FID (compared to the results in CIFAR-10) indicates that a nonlinear diffusion is more effective on datasets of higher dimensions. In fact, this is a quite intuitive result because as the data dimension increases, the family of feasible diffusion expands exponentially, and the data-optimal diffusing mechanism is more likely to deviate from linear diffusions.

heading: 6.2 QUALITATIVE RESULTS
text: Test NELBO To figure out an effect of a nonlinear diffusing mechanism on the training procedure, Figure 5 (a) shows test NELBO by training iterations on a diffusion model with 35M parameters. Without any modification on hyperparameters, we train PDM with a shallow ResFlow (Chen et al., 2019) of 3M parameters, and Figure 5 (a) shows that the training on PDM surprisingly stabilizes test NELBO. Figure 5 (b) illustrates that such training stability can be attributed to nearly identical score scales. As described in Figure 5 (b), the score scale on a data space is unbounded as the diffusion time goes to zero, and this extreme scale variation eventually leads to training instability on the baseline diffusion models. On the other hand, adopting flow transformation on a data space mitigates such extreme scale variations that destabilize training.
Figure 6 shows test NELBO on a diffusion model with 107.6M number of parameters. To stabilize the diffusion model, we apply the AdamW optimizer (Loshchilov and Hutter, 2017) with 0.01 weight decay, and Exponential Moving Average (Song et al., 2020) of 0.9999. Using the ResFlow model with 10.6M parameters, Figure 6 indicates that PDM consistently outperforms the baseline diffusion in terms of NELBO with smaller variations.
Flow Depth The flow model architecture is a building block that fundamentally determines the family of inference path measures. Table 4 empirically demonstrates that an inflexible variational family results in large variational gap on CIFAR-10. Additionally, the absolute magnitude of NLL performance keeps decreasing as the number of layers (L) increases.
Sample Generation Figure 7 presents samples from (a) a latent space and (b) a data space. Since a sample generation on the latent space is unnormalized, we normalize the sampled latent variables to be bounded on the range of [0, 255], so we can visualize the latent information in image. Figure 7 indicates that a diffusion model and a flow model have distinctive and complementary roles in sample generation: 1) the diffusion model constructs the global context and 2) the flow model enriches this global context by colorizing it into a realistic image.

heading: 7 CONCLUSION
text: This paper expands the linear diffusing mechanism to be nonlinear by combining an invertible transformation and a diffusion model. This nonlinear diffusing mechanism learns the forward diffusion process out of variational family of inference path measures, and such optimization provides tighter variational gap compared to a baseline model with linear diffusion.

heading: 8 ETHICS STATEMENT
text: As our work improves sample quality, our methodology could be a source of fake datasets, i.e. fake images. This is a common threat from the line of deep generative modeling.

heading: 9 REPRODUCIBILITY STATEMENT
text: Every proofs and derivations are explained in details on Appendix A and B. Also, for reproducibility, we would release our code as soon as the discussion period opens through an anonymous repository to reviewers. We note that most of our code is built based on the released code of Song et al. (2020) in order to compare with baselines in a fair setting introduced in Song et al. (2020) and Song et al. (2021). Additionally, we provide the details on training/evaluation/neural architecture in Appendix C.

heading: A PROOFS OF THEOREM 1 AND THEOREM 2
text: Theorem 1. Suppose that pφ,θ(x0) is the log-likelihood of the generative random variable xφ,θ0 . Then, the log-likelihood is bounded by
Epr(x0) [ − log pφ,θ(x0) ] ≤ Lφ,θ ( {xt}Tt=0; g2 ) ,
where
Lφ,θ ( {xt}Tt=0; g2 ) =− Epr(x0) [ log ∣∣∣det(∂hφ
∂x0 )∣∣∣]+ Lθ({zt}Tt=0; g2)− EzT [ log π(zT )] + d
2 ∫ T 0 β(t)− g 2(t) σ2(t) dt,
with σ2(t) = 1− e− ∫ t 0 β(s) ds + ∫ t 0 [g2(s)− β(s)] ds.
Proof of Theorem 1. From the change of variable, the transformation of z0 = hφ(x0) induces
pr(x0) = p0(z0)∣∣∣det(∂hφ∂x0 )∣∣∣−1 ,
and thus the entropy of the data distribution becomes H(pr) = − ∫ pr(x0) log pr(x0) dx0
= − ∫ p0(z0) log
p0(z0)∣∣∣det(∂hφ∂x0 )∣∣∣−1 dz0 = − ∫ pr(x0) log ∣∣∣ det(∂hφ ∂x0
)∣∣∣dx0 − ∫ p0(z0) log p0(z0) dz0 = −Epr(x0) [ log ∣∣∣det(∂hφ
∂x0 )∣∣∣]− ∫ p0(z0) log p0(z0) dz0 = −Epr(x0) [ log ∣∣∣det(∂hφ
∂x0 )∣∣∣]+H(p0). From Theorem 4 of Song et al. (2021), the entropy at t = 0 equals to
H(p0) = H(pT )− 1
2 ∫ T 0 Ept(zt) [ 2∇zt · f(zt, t) + g2(t)‖∇zt log pt(zt)‖22 ] dt,
where f(zt, t) is a drift term of the diffusion for zt and pt is the probability distribution of zt. Therefore, the negative log-likelihood becomes
−Epr(x0) [ log pφ,θ(x0) ] = DKL(pr‖pφ,θ) +H(pr)
≤ DKL(µdφ‖νdφ,θ) +H(pr)
= DKL(µ d φ‖νdφ,θ)− Epr(x0)
[ log ∣∣∣ det(∂hφ
∂x0 )∣∣∣]+H(p0) = DKL(µ l φ‖νlθ)− Epr(x0) [ log ∣∣∣det(∂hφ
∂x0 )∣∣∣]+H(pT ) −1 2 ∫ T 0 Ezφt [ − dβ(t) + g2(t)‖∇zt log pt(zt)‖22 ] dt.
Now, from Theorem 1 of Song et al. (2021), the KL-divergence between the path measures becomes
DKL(µ l φ‖νlθ) = DKL(pT ‖π) +
1
2 ∫ T 0 g2(t)Ept(zt) [ ‖sθ(zt, t)−∇zt log pt(zt)‖22 ] dt,
so if we plug in this into the negative log-likelihood, we yield the following: −Epr(x0) [ log pφ,θ(x0) ] = −Epr(x0) [ log ∣∣∣ det(∂hφ
∂x0
)∣∣∣]+ 1 2 ∫ T 0 g2(t)Ezt [ ‖sθ(zt, t)−∇zt log pt(zt)‖22 ] dt
+DKL(pT ‖π) +H(pT )− 1
2 ∫ T 0 Ezt [ − dβ(t) + g2(t)‖∇zt log pt(zt)‖22 ] dt
= −Epr(x0) [ log ∣∣∣ det(∂hφ
∂x0 )∣∣∣]− EzT [ log π(zT )]+ d2 ∫ T 0 β(t) dt
+ 1
2 ∫ T 0 g2(t)Ezt [ ‖sθ(zt, t)−∇zt log pt(zt)‖22 − ‖∇zt log pt(zt)‖22 ] dt
Also, we have
Ezt [ sθ(zt, t) · ∇zt log pt(zt) ] = ∫ pt(zt)sθ(zt, t) · ∇zt log pt(zt) dzt
= ∫ sθ(zt, t) · ∇ztpt(zt) dzt
= ∫ sθ(zt, t) · ∫ p0(z0)∇ztp0t(zt|z0) dz0 dzt
= Ez0Ezt|z0 [ sθ(zt, t) · ∇zt log p0t(zt|z0) ] Therefore,
1
2 ∫ T 0 g2(t)Ezt [ ‖sθ(zt, t)−∇zt log pt(zt)‖22 − ‖∇zt log pt(zt)‖22 ] =
∫ T 0 g2(t)Ezt [1 2 ‖sθ(zt, t)‖22 − sθ(zt, t) · ∇zt log pt(zt) ] =
∫ T 0 g2(t)Ez0Ezt|z0 [1 2 ‖sθ(zt, t)‖22 − sθ(zt, t) · ∇zt log p0t(zt|z0) ] = 1
2 ∫ T 0 g2(t)Ez0Ezt|z0 [ ‖sθ(zt, t)−∇zt log p0t(zt|z0)‖22 − ‖∇zt log p0t(zt|z0)‖22 ] .
Now, since p0t(zt|z0) = N (zt;µ(t)zt, σ2(t)I) for µ(t) and σ2(t) determined by β(t) and g(t), we have
Ezt|z0 [ ‖∇zt log p0t(zt|z0) ] = Ezt|z0 [∥∥∥zt − µ(t)z0 σ2(t) ∥∥∥2 2 ] = EN (z;0,I) [ ‖z‖22 σ2(t) ] = d σ2(t) ,
and we have the desired result.
Theorem 2. Suppose the flow network is flexible enough to transform pr to arbitrary continuous distribution. For any θ, if∇·sθ(·, t) is symmetric for any t ∈ [0, T ] and ∇· log π(·) ≡ sθ(·, T ), then there exists φ∗ such that Gapd(µ d φ∗ ,ν d φ∗,θ) = 0. Furthermore, we have
φ∗ = argmin φ Gapd(µ d φ,ν d φ,θ) = argmin φ DKL(µ d φ‖νdφ,θ) = argmin φ DKL(pr‖pφ,θ).
Proof of Theorem 2. The symmetry of∇·sθ(·, t) indicates that sθ(·, t) is a closed 1-form. Therefore, from the Poincare’s lemma (Do Carmo, 2012), sθ(·, t) is exact, and there exists a scalar field fθt that satisfies sθ(·, t) = ∇·fθt (·) for any t. If we define by qθt (zt) = exp (fθt (zt))∫ exp (fθt (z)) dz , then we have sθ(zt, t) = ∇zt log qθt (zt, t) for any t. With this construction, the reverse SDE on the transformed latent space becomes
dzθt = [ − 1
2 β(t)zθt − g2(t)∇zθt log q θ t (z
θ t ) ] dt+ g(t) dω̄t, z θ T ∼ π = qθT .
If ρθ is the path measure of the solution of the above SDE, then the marginal distribution of ρθ at t would be qθt . Therefore, if h −1 φ∗ maps pr to q θ 0 , then the forward process on the latent space {zφ ∗
t }Tt=0 becomes the generative process on the latent space {zθt }Tt=0. Thus, we have µlφ∗ = νlθ in distribution, which leads DKL(µlφ∗‖νlθ) = 0. We get the desired result by the change of variable formula: DKL(µdφ∗‖νdφ∗,θ) = DKL(µlφ∗‖νlθ) = 0.

heading: B DERIVATIONS

heading: B.1 DERIVATION OF NONLINEAR DIFFUSION TERMS
text: Throughout this section, we omit φ for the notational simplicity. Beginning with the linear diffusion on the latent space given by
dzt = − 1
2 β(t)zt dt+ g(t) dωt, z0 = h(x0) with x0 ∼ pr, (11)
we have defined xt = h−1(zt) as the induced random variables. Then by the multivariate Ito’s Lemma (Oksendal, 2013), the k-th component of the induced variable satisfies
dxt,k = ∂h−1k ∂t dt+ [ ∇zth−1k (zt) ]T dzt + 1 2 tr ( ∇2zth −1 k (zt) dzt dz T t ) (12)
Plugging the linear diffusion SDE 11 in the SDE of Eq. 12, Eq. 12 is derived by
dxt,k = [ ∇zth−1k (zt) ]T {−1 2 β(t)zt dt+ g(t) dωt } + 1 2 g2(t)tr ( ∇2zth −1 k (zt) ) dt
= { −1 2 β(t) [ ∇zth−1k (zt) ]T zt + 1 2 g2(t)tr ( ∇2zth −1 k (zt) )} dt+ g(t) [ ∇zth−1k (zt) ]T dωt,
(13)
because ∂h −1 k
∂t = 0. Then, Eq. 13 in vector form becomes
dxt = f(xt, t) dt+ G(xt, t) dωt,
where vector-valued drift and diffusion terms are given by{ f(xt, t) = − 12β(t)∇zth −1(zt)zt + 12g 2(t)tr ( ∇2zth −1(zt) )
G(xt, t) = g(t)∇zth−1(zt), (14)
where∇2zth −1(zt) is a 3-dimensional tensor with (i, j, k)-th element to be ( ∇2zth −1 k (zt) ) i,j , and the
trace operator applied on this tensor results in the vector of [ tr ( ∇2zth −1 1 (zt) ) , ..., tr ( ∇2zth −1 d (zt) )]T . From the inverse function theorem (Rudin et al., 1964), the Jacobian of the inverse function ∇zth−1(zt) equals to the inverse Jacobian [ ∇xth(xt)
]−1 . Therefore, Eq. 14 is transformed to{
f(xt, t) = − 12β(t) [ ∇xth(xt) ]−1 h(xt) + 1 2g 2(t)tr ( ∇2zth −1(zt) )
G(xt, t) = g(t) [ ∇xth(xt) ]−1 .
(15)
Now, we derive the second term of f in terms of xt as follows: observe that ∑ k ∂h−1i ∂zt,k ∂hk ∂xt,j
= δi,j , where δi,j = 1 if i = j and 0 otherwise. Differentiating both sides with respect to xt,l, we have∑
k
{ ∂
∂xt,l ( ∂h−1i ∂zt,k )} ∂hk ∂xt,j + ∂h−1i ∂zt,k { ∂ ∂xt,l ( ∂hk ∂xt,j )} = 0,
where the first term is∑ k,m ∂hm ∂xt,l { ∂ ∂zt,m ( ∂h−1i ∂zt,k )} ∂hk ∂xt,j = ∑ k,m ( ∇xth(xt) )T l,m ( ∇2zth −1 i (zt) ) m,k ( ∇xth(xt) ) k,j ,
using the chain rule, and the second term becomes∑ k ∂h−1i ∂zt,k { ∂ ∂xt,l ( ∂hk ∂xt,j )} = ∑ k ( ∇zth−1(zt) ) i,k ( ∇2xthk(xt) ) l,j .
From above, we derive the trace term of f in Eq. 15 as tr ( ∇2zth −1(zt) ) = −tr ([ ∇xth(xt) ]−T ([∇xth(xt)]−1 ∗ ∇2xth(xt)) [∇xth(xt)]−1) ,
Table 5: Comparison of variational gap between the VAE, the diffusion model, and PDM.
Models Diffusion Stochastic Process Variational Gap (=NLL−NELBO) VAE No {x, z} Ex [ DKL(qφ(z|x)‖pθ(z|x)) ] Diffusion Models Data space {xt}Tt=0 Ex0 [ DKL(µ({xt}|x0)‖νθ({xt}|x0))
] PDM Latent space {xt}Tt=0/{zt}Tt=0 Ez0 [ DKL ( µlφ({zt}|z0)‖νlθ({zt}|z0)
)] where ∇2xth(xt) is a 3-dimensional tensor with (i, j, k)-th element to be ( ∇2xthk(xt) ) i,j
. Also, we define ∗ operation as the element-wise matrix multiplication given by([
∇xth(xt) ]−1 ∗ ∇2xth(xt))
i,j := ∇xt
[ h(xt) ]−1(∇2xth(xt)) i,j .
Combining all together, thus, we derive nonlinear drift term in Eq. 15 as a function of xt given by
f(xt, t) =− 1
2 β(t)
[ ∇xth(xt) ]−1 h(xt)
− 1 2 g2(t)tr
([ ∇xth(xt) ]−T ([∇xth(xt)]−1 ∗ ∇2xth(xt)) [∇xth(xt)]−1) .

heading: B.2 DERIVATION OF VARIATIONAL GAP
text: First, the KL divergence between two path measures satisfies DKL(µ‖νθ) = ∫ µ({xt}) log
µ({xt}) νθ({xt}) dx0:T
= ∫ pr(x0)µ({xt}|x0) log
pr(x0)µ({xt}|x0) pθ(x0)νθ({xt}|x0) dx0:T
= ∫ pr(x0)µ({xt}|x0) [ log pr(x0)
pθ(x0) + log µ({xt}|x0) νθ({xt}|x0)
] dx0:T
= ∫ pr(x0) log pr(x0)
pθ(x0) dx0 +
∫ pr(x0)µ({xt}|x0) log
µ({xt}|x0) νθ({xt}|x0) dx0:T
=DKL(pr‖pθ) + Epr(x0) [ DKL ( µ({xt}|x0)‖νθ({xt}|x0) )] ,
where x0:T = {xt}Tt=0. Therefore, variational gap is derived by
Gap(µ,νθ) =DKL(µ‖νθ)−DKL(pr‖pθ) =Epr(x0) [ DKL ( µ({xt}|x0)‖νθ({xt}|x0) )] .
The gap Gapd(µ d φ,ν d φ,θ) in Eq. 10 is derived to be Ez0 [ DKL ( µφ({zt}|z0)‖νθ({zt}|z0) )] , analogously.
C IMPLEMENTATION DETAILS

heading: C.1 MODEL ARCHITECTURE
text: Diffusion Model We implement two diffusion models as backbone: NCSN++ (VE) (Song et al., 2020) and DDPM++ (VP) (Song et al., 2020), where two backbones are one of the best performers in CIFAR-10 dataset. In our setting, NCSN++ assumes the score network with parametrization of sθ(zt, log σ2(t)), where σ2(t) = σ2min( σmax σmin
)2t is the variance of the transition probability p0t(zt|z0) with VESDE. As introduced in Song et al. (2020), we use the Gaussian Fourier embeddings (Tancik et al., 2020) to model the high frequency details across the temporal embedding. DDPM++ models the score network with parametrization of θ(zt, t), which
targets to estimate −σ(t)∇zt log pt(zt). We use the Transformer sinusoidal temporal embedding (Vaswani et al., 2017).
We use the U-Net (Ronneberger et al., 2015) architecture for the score networks on both NCSN++ and DDPM++ based on (Ho et al., 2020). We stack U-Net resblocks of up-and-down convolutions with skip connections that connect the identical-dimensional layers. Also, we follow Ho et al. (2020) by applying the global attention at the resolution of 16× 16. We use four U-Net resblocks with four feature map resolutions (32× 32 to 4× 4). There are eight and four resblocks at each resolution on CIFAR-10 and CelebA, respectively.
Flow Model We implement two normalizing flow models as backbone: the ResFlow (Chen et al., 2019) and the decoupled ResFlow (Ma et al., 2020). Although Ma et al. (2020) makes use of the Glow (Kingma and Dhariwal, 2018) for their flow model, we replace this glow architecture with the ResFlow, mainly because the ResFlow transforms the variables from the identity function. We name the experiments on this decoupled ResFlow as PDM (FID, dec) and PDM (NLL, dec). Otherwise, we train the model with ResFlow as default. For the ResFlow, we drop three components from the original paper: 1) the activation normalization Kingma and Dhariwal (2018), 2) the batch normalization Ioffe and Szegedy (2015), and 3) the fully connected layers.
On the deep ResFlow, we use resolutions of 32× 32 and 16× 16 on both CIFAR-10 and CelebA. Otherwise, we use a single resolution of 32× 32 on other experiments whose models have networks with layers less than 32. We apply the LipSwish activation (Chen et al., 2019) that provides the Lipschitz coefficients less than a unity. For the multi-GPU training, we use the Neumann logdeterminant gradient estimator, instead of the memory efficient estimator (Chen et al., 2019).

heading: C.2 EXPERIMENTAL DETAILS
text: Training All the models are trained with batch size of 128 and EMA rate of 0.999. We train each model for around 10 days on four V-100 GPUs with 128Gb GPU memory for all experiments. At the initial training phase, we train the diffusion model parameters with fixed flow parameters initialized by the identity until the sample generation performance saturates. After this pretraining phase, we start updating both diffusion and flow models. In this learning phase, we apply the learning rate scheduling after the sample performance saturation. For the diffusion model, we drop the learning rate from 2 × 10−4 to 10−5. For the flow model, we drop the learning rate from 10−3 to 10−5 for VPSDE and 5 × 10−5 to 10−5 for VESDE. We emphasize this learning rate scheduling takes additional performance boost.
VESDE and VPSDE have different training details. VESDE assumes σmin = 10−2 on all experiments. VESDE has σmax = 50 on CIFAR-10 and σmax = 90 on CelebA. On the other hand, VPSDE assumes β(t) = βmin + (βmax − βmin)t with βmin = 0.1 and βmax = 20. Both VESDE and VPSDE truncate the diffusion time on [ , T ] in order to stabilize the diffusion model, where = 10−5 and T = 1. On sample generation, however, DDPM++/PDM with VPSDE draws samples by proceeding the generative diffusion process up to = 10−3, which corresponds to the diffusion variance to be σmin = 10−2. On the sample generation, we apply the Predictor-Corrector (PC) sampler Song et al. (2020) for NCSN++/PDM with VESDE and the Probability Flow (PF) ODE sampler Song et al. (2020) for DDPM++/PDM. Throughout the experiments, we provide the identical diffusion model structures to compare the baseline model and the PDM model in a fair setting.
Evaluation While we use the EMA rates of 0.999 for both diffusion and flow models for training, we use this EMA only for the diffusion model on the evaluation mode because the flow model without the EMA on the evaluation gives better samples. One peculiarity occurs in the experiment on the decoupled ResFlow. We empirically find that turning off the EMA for both diffusion model and flow model on the evaluation mode provides better performances of 3.11 in test NLL and 3.17 in test NELBO as reported in Table 1 on PDM (FID, dec). With the EMA applied to the diffusion model, we have 3.15 for test NLL and 3.22 for test NELBO, which shows slight degradation on the density estimation performance.
We compute the FID score based on Song et al. (2020) for CIFAR-10 with the modified Inception V1 network1 in order to compare PDM to the baselines (Song et al., 2020; 2021) in a fair setting. On the
1https://tfhub.dev/tensorflow/tfgan/eval/inception/1
other hand, for the CelebA dataset, we compute the clean-FID (Parmar et al., 2021) that provides consistently antialiased results.
C.3 VARIANCE REDUCTION Flow Training When we train the flow network with Lφ,θ ( {xt}Tt=0; g2), this NELBO contains
the integration of Lθ ( {zt}Tt=0; g2 ) . However, previous works on the diffusion models (Nichol and Dhariwal, 2021) show that the estimation variance is largely reduced with the importance sampling, which leads the performance variations (Song et al., 2021). Concretely, the importance sampling chooses an importance weight that is proportional to g
2(t) σ2(t) , and estimates the integration by Lθ ( {zt}Tt=0; g2 ) = ∫ T 0 g2(t)Lt ( {zt}Tt=0;θ ) dt ≈ ∑N n=1 σ 2(tn)Ltn ( {zt}Tt=0;θ ) , where tn is sampled from the importance distribution.
For VESDE, it satisfies β(t) = 0 and g(t) = σmin(σmaxσmin ) t √
2 log (σmaxσmin ). Throughout the experiments, we select σmin = 0.01, and, σmax = 50, 90 for CIFAR-10 and CelebA, respectively. The transition probability becomes p0t(zt|z0) = N (zt; z0, σ2(t)), where σ2(t) = ∫ t 0 g2(s) ds =
σ2min [ (σmaxσmin ) 2t − 1 ] . Song et al. (2020) approximates this variance to be σ2app(t) = σ 2 min( σmax σmin )2t so that the variance is proportional to g2(t). Therefore, the importance weight follows the uniform distribution, and the importance sampling is equivalent with choosing the uniform t.
On the other hand, VPSDE satisfies β(t) = βmin + (βmax − βmin)t with g(t) = √ β(t). Then, the transition probability becomes p0t(zt|z0) = N (zt;µ(t)zt, σ2(t)I), where µ(t) = e− 1 2 ∫ t 0 β(s) ds and
σ2(t) = 1− e− ∫ t 0 β(s) ds. Thus, VPSDE has the importance weight of g
2(t) σ2(t) =
β(t) 1−e− ∫ t 0 β(s) ds .
The Monte-Carlo sample from this importance weight is the solution of the inverse Cumulative Distribution Function (CDF) of the importance distribution as
t = F−1(u) (16)
⇐⇒ u = F (t) = 1 Z
∫ t
g2(s) σ2(s) ds = 1 Z
( F(t)−F( ) ) , (17)
where u is a uniform sample from [0, 1], F(t) is the antiderivative of the importance weight given by F(t) = log (1− e−0.5t2(βmax−βmin)−tβmin) + 0.5t2(βmax − βmin) + tβmin, and Z is the normalizing constant given by
Z =
∫ T
g2(t) σ2(t) dt
= [ log (1− e−0.5t 2(βmax−βmin)−tβmin) + 0.5t2(βmax − βmin) + tβmin ]T
= log (1− e−0.5T 2(βmax−βmin)−Tβmin)− log (1− e−0.5 2(βmax−βmin)− βmin)
+ 0.5(T 2 − 2)(βmax − βmin) + (T − )βmin =23.86
for T = 1 and = 10−5. The solution for the inverse CDF in Eq. 16 becomes
e ∫ t 0 β(s) ds = 1 + exp (Zu+ F( ))
⇐⇒ ∫ t 0 β(s) ds = 1 2 (βmax − βmin)t2 + βmint = log (1 + exp (Zu+ F( )))
⇐⇒ t = −βmin +
√ β2min + 2(βmax − βmin) log ( 1 + exp (Zu+ F( )) ) βmax − βmin .
The variation of the Monte-Carlo diffusion time depends on the uniform sample of u.
Diffusion Training Now, on the side of training the diffusion network, as described in the main paper, selecting the weighting function heavily influences the model performances. For instance,
Song et al. (2021) compare the weighting functions of g2 and σ2, and Song et al. (2021) conclude that the weighting function of λ(t) = σ2(t) works the best for the sample generation, and λ(t) = g2(t) yields the best NLL performance. We train the diffusion model with σ2(t) for PDM (VE/VP, FID), and with g2(t) for PDM (VP, NLL) models.

heading: D ADDITIONAL EXPERIMENTAL RESULTS

heading: D.1 DETAILED COMPARISON AND ABLATION ON TOLERANCE
text: Table 7 has two implications. First, we find that the reported NLL and NELBO performances are slightly better than what we actually obtain in our pretraining step, denoted by DDPM++ (NLL, IS, ours). Comparing with our results, optimizing PDM is indeed beneficial on density estimation, and variational gap is significantly reduced from 0.17 to 0.01. Second, we find the tolerance of 10−1 for the numerical ODE solver on computing NLL is the sweet spot on obtaining best NLL of 2.86, which beats LSGM-480M (NLL) on density estimation. This is noteworthy because our model uses 120M number of parameters, a quarter of parameters in LSGM. On the other hand, NELBO is independent of the tolerance because it does not use the
numerical ODE solver, and Table 7 shows that variational gap is minimized when the tolerance is small enough.
Figure 10 illustrates the ablation study on the tolerance for the ODE solver. Surprisingly, PDM is robust on the tolerance level up to 10−1, in contrast to LSGM with poor sample quality at 10−1 (Vahdat et al., 2021). From this robustness on the tolerance, we could reduce the sampling wall clock time from 360s (tolerance 10−5) to 72s (tolerance 10−1) for 1,024 samples, which is 5x faster, with a minor sacrifice in FID score from 6.79 to 8.56 in PDM (FID) with VPSDE.

heading: D.2 PRECISION AND RECALL
text: While FID measures the overall sample quality that counts both fidelity and diversity at once, we provide precision and recall (Kynkäänniemi et al., 2019) based on the modified Inception V1 network
Table 10: FID-1k of PDM with VPSDE on CIFAR10.
Temperature FID-1k
0.8 15.60 0.9 15.27 1.0 15.04 1.05 14.95 1.1 14.90 1.15 14.94 1.2 15.05 1.3 15.44
2 with 50k samples in order to separate the sample fidelity and diversity. Precision is a metric that measures the sample fidelity, and the recall is for the sample diversity. Table 8 indicates that both precision and recall remains at the identical scale, except for precision at PDM (FID, dec).

heading: D.3 EFFECT OF SIGNAL-TO-NOISE AND TEMPERATURE
text: As the flow colorizes the latent sample, we experiment the effect of temperature (Kingma and Dhariwal, 2018) on this colorization in Figure 11. With temperature τ , the normalizing flow puts its latent input scaled by τ to the flow network. In our PDM, the image color with a higher temperature tends to be brighter, and we find that the optimal temperature depends on the experimental settings.
On sample generation, we draw latent samples from either the PC sampling or the PF ODE sampling. PC sampling assumes the Signalto-Noise Ratio (SNR) that modifies the Corrector algorithm (Song et al., 2020) of the annealed Langevin dynamics (Song and Ermon,
2019). From the sampled latent variable, we scale by multiplying the temperature τ (Kingma and Dhariwal, 2018) to the latent variable. Since the diffusion model learns the global image structure, the normalizing flow only paints the color of the image, and the temperature controls the brightness of the generated image. This is a contrastive result to the previous flow models (Kingma and Dhariwal, 2018; Chen et al., 2019) that the temperature in the previous models controls both global context (such as the smoothness) and color. Therefore, we can create images with different colors by manipulating this temperature.
2Since we already used this inception network in calculating the FID score, we compute precision and recall based on this inception network. Note that we used this version of inception network to comply the baseline (Song et al., 2020).
CIFAR-10 With these hyperparameters of the SNR and the temperatures, we find optimal values in the below Tables. Table 9 computes the FID-5k score of all combinations for (SNR, τ ) in PDM (VE, FID) on CIFAR-10. Unlike the SNR of 0.16 that is previously suggested in (Song et al., 2020) for the vanilla diffusion model, we find that SNR of 0.145 consistently outperforming other values for SNR in terms of FID. Similarly, Table 9 presents that the optimal temperature is 1.1.
For VPSDE, Table 10 shows that the temperature of 1.3 performs the best for sample generation in PDM (VP, FID). For PDM (VP, FID, dec), we find that the temperature of 1.1 performs the best, and for PDM (VP, NLL, dec), the vanilla temperature of 1.0 generates the best samples.
CelebA Table 11 shows the FID-1k performance of PDM (VE, FID) on CelebA. It shows that the SNR of 0.17 proposed in (Song et al., 2020) performs the best also in our PDM. In addition, it shows that the optimal temperature is either 1.0 or 1.05. Table 12 shows the FID-1k performance of PDM (VP) on CelebA. It is clear in this table that the sample quality is the best on the temperature range of 1.05 to 1.15.

heading: D.4 RANDOM SAMPLES
text: Figure 12 and 13 show the non cherry-picked random samples from PDM (VE, FID) on CIFAR-10 and PDM (VP, FID) on CelebA, respectively.

Review summaries:
The author proposed a framework that combines normalizing flow with diffusion models, allowing non-linear diffusions to be used for inference and generation. Specifically, the author first discussed the similarity between VAE and diffusion models to motivate the importance of using non-linear diffusion and the reducing variational gap. Then, the author proposed to decouple the non-linear diffusion into an invertible transformation with the normalizing flow and typical linear diffusions on the transformed variables. A training objective based on variational lower bound was also proposed. Further, the author also showed that the variational gap between the lower bound and log-likelihood can be reduced with the proposed non-linear diffusions. 

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.

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.