Diffusion Normalizing Flow

1 Introduction

• Deterministic flow models face topological constraints when mapping a simple Gaussian prior to data with well-separated modes.They use invertible maps and can train by directly minimizing negative log-likelihood.Source: §1 · p. 1
• Diffusion models use a fixed stochastic forward process and a trained backward process, but long trajectories slow training and sampling.Because noise addition is independent of the data distribution, the learned model may miss complex but important data details.Source: §1 · p. 2
• DiffFlow makes both forward and backward processes trainable and stochastic, extending normalizing flows with noise and diffusion models with a learnable forward process.When noise shrinks to zero, DiffFlow reduces to normalizing flow; fixing its forward process to a specific diffusion makes it a diffusion model.Source: §1 · p. 2, Figure 1 · p. 2
• DiffFlow jointly trains its forward and backward diffusion processes by minimizing the KL divergence between their induced trajectory distributions.This objective is equivalent to an amortized negative evidence lower bound and can be estimated without bias from sampled trajectories.Source: §1 · p. 2
• DiffFlow reports a 20 times speedup over diffusion models without much decrease in sampling quality.The introduction attributes the faster sampling to a learnable forward process that can shorten trajectories and add noise more selectively.Source: §1 · p. 2
• DiffFlow, a density-estimation model extending flow- and diffusion-based models, is applied to several generative-modeling tasks on synthetic and real datasets, where its performance and advantages over other methods are verified.Source: §1 · pp. 2–3

2 Background

• A trajectory τ denotes a path from data space to latent space, represented continuously as x(t) or discretely as a sequence of states.The notation applies to both normalizing flows and diffusion models.Source: §2 · p. 3
• Normalizing flows use parameterized differential equations or compositions of bijective functions to map data to latent variables.Under the change-of-variable relation, they can evaluate data likelihood and train by minimizing negative log-likelihood directly.Source: §2 · p. 3
• Diffusion models describe trajectories with stochastic differential equations, using a forward process and a backward diffusion to recover data from noise.The forward SDE has drift f and diffusion coefficient g, with Brownian motion w; the backward process can be trained using a score function.Source: §2 · p. 3
• The KL divergence between forward and backward trajectory distributions can serve as a training objective, and diffusion models commonly use its decomposition.Their simple forward-process structure supports objective evaluation during training.Source: §2 · p. 3

3 Diffusion normalizing flow

• The forward and backward SDEs are trained by minimizing trajectory-space KL divergence, which bounds the KL divergence between their distributions at any time.The forward drift is learnable, and the forward process starts from data while the backward process starts from a noise distribution.Source: §3 · p. 4

3.1 Implementation

• The discretized trajectory-KL loss is optimized with Monte Carlo gradient estimates, and the terminal backward density is typically a standard Gaussian.The terminal term is easy to calculate because the backward process ends in a simple distribution.Source: §3.1 · p. 4

3.2 Stochastic Adjoint method

• Stochastic adjoint training addresses N-step memory growth by caching intermediate states and exactly reproducing trajectories; cached states use about 2% of image-training memory.The approach reproduces the forward noise and drift and score values from cached states.Source: §3.2 · p. 5

3.3 Time discretization and progressive training

• The fixed Lβ schedule provides finer time resolution near x0, while randomized L̂β supports progressive increases in N and speeds training by up to 16 times.The authors report that progressive training yields lower loss and better stability than fixed large N throughout training.Source: §3.3 · pp. 5–6

3.4 Learnable forward process

• As diffusion coefficients g_i approach zero, DiffFlow reduces to normalizing flow and its objective becomes negative log-likelihood; learnable noise also relaxes bijective-flow expressivity limits.The learnable forward process adapts across modes and datasets, retaining topology and sharp details for reverse-process learning.Source: §3.4 · pp. 6–7
• Unlike DDPM’s data-invariant noising, DiffFlow adapts diffusion across modes and datasets, retaining topological information and sharp details for reverse-process learning.DDPM can destroy density details quickly, while DiffFlow’s data-dependent transformation preserves some topology.Source: §3.4 · pp. 6–7

4 Experiments

• The experiments assess likelihood and sample quality, using a marginal-distribution-equivalent ODE for likelihood evaluation and SDE sampling with λ = 1 performing best empirically.The synthetic 2-D comparison uses models with comparable network sizes, around 90k learnable parameters.Source: §4 · p. 7, §4.1 · p. 7

4.1 Synthetic 2D examples

• On challenging 2-D datasets, DiffFlow produces the sharpest samples, while FFJORD leaks samples from modes and DDPM blurs fine density details.All three methods perform reasonably on smooth distributions; DDPM's blurred details include ring intersections, fractal-tree leaves, and Sierpiński Carpet boxes.Source: §4.1 · p. 7, Figure 5 · p. 8

4.2 Density estimation on real data

• DiffFlow outperforms FFJORD on all five tabular datasets except HEPMASS and excels NAF on all but GAS.It performs better on most datasets than most approaches trained by directly minimizing negative log-likelihood, including normalizing flows and autoregressive models.Source: §4.2 · p. 7

4.3 Image generation

• On CIFAR-10, DiffFlow has competitive NLL, but its FID is worse than several DDPM variants and DDPM++.It performs better than normalizing flows and DDPM models on NLL, but is slightly worse than DDPM++(sub, deep, sub-vp) and Improved DDPM; the latter use deeper and wider networks with multiple architectural improvements.Source: §4.3 · p. 8, Table 3 · p. 9
• DiffFlow retains better sampling quality as the number of sampling steps decreases, with the comparison reported at N = 100.The relative FIDs degeneracy ratios are compared at N = 100.Source: §4.3 · p. 9, Table 4 · p. 9

• Flow-based models use invertible maps and tractable density evaluation, but bijectivity limits modeling complex data.Stochastic normalizing flows can improve expressiveness in low dimensions, but their known-energy requirement challenges density learning.Source: §5 · p. 9
• DDPM uses a fixed linear forward diffusion, while score-based methods also frame diffusion and reverse sampling as SDEs.The fixed linear process provides closed-form forward marginals, but diffusion models can require thousands of steps for diffusion and generation.Source: §5 · pp. 9–10
• Existing neural-SDE training approaches face scaling costs: solver backpropagation uses linear memory in steps, while pathwise methods scale poorly in computation.DiffFlow's trajectory-state caching permits coarser discretization for larger-dimensional problems, with negligible additional memory relative to network memory.Source: §5 · p. 10

• DiffFlow loses an analytical form for pF (xt|x0), making training less efficient than score-based loss.Source: §6 · p. 10
• Backpropagation through trajectories makes DiffFlow training significantly slower than diffusion models with affine drift.Without progressive training, it is reported as about 6 times slower than DDPM in 2-D toy examples, 55 times in MNIST, and 160 times in CIFAR10.Source: §6 · p. 10
• Even with stochastic adjoint and progressive training savings, DiffFlow training remains more expensive than DDPM and its variants.Source: §6 · p. 10

• DiffFlow jointly trains forward and backward neural SDEs by minimizing their KL divergence.Source: §7 · p. 10
• Added noise relaxes the bijectivity condition of most normalizing flows, while a learnable forward diffusion adapts the transformation from data to noise.Source: §7 · p. 10
• DiffFlow models complex details missed by representative flow and diffusion models, but has worse CIFAR-10 FID than DDPM.The authors identify alternative architectures, time discretization, and time intervals as possible avenues for improvement.Source: §7 · p. 10
• The algorithm learns high-dimensional data distributions and generates samples, with potential use in producing misleading images or videos.Source: §7 · p. 10

A Proof of Theorem 1: Normalizing Flow as a special case of DiffFlow

• As g(t) approaches zero, DiffFlow's stochastic trajectories become deterministic, and its forward SDE reduces to a normalizing-flow ODE.In the discrete implementation, it reduces to a normalizing flow composed of discrete layers.Source: A Proof of Theorem 1: Normalizing Flow… · p. 13
• For a fixed forward process, KL minimization makes the backward conditional the posterior of the forward transition when diffusion intensity is small.The proof uses the change-of-variable rule for the inverse map to connect this result to the normalizing-flow likelihood objective.Source: A Proof of Theorem 1: Normalizing Flow… · p. 13
• In the zero-noise limit, minimizing the trajectory KL divergence is equivalent to maximizing log probability, as in the normalizing-flow objective.The discrete forward map is F_i(x, θ) = x + f(x, t_i, θ)∆t_i.Source: A Proof of Theorem 1: Normalizing Flow… · p. 13

B.1 Proof of Eq (10)

• Eq. (10) follows from measure disintegration and the non-negativity of KL divergence.The proof begins by applying the disintegration theorem.Source: §B.1 · p. 13

B.2 Proof for Eq (14) and log-likelihood of pB(xi|xi+1)

• The backward conditional density is reformulated using the backward dynamics in Eq. (9), yielding its negative log-likelihood term.The passage identifies the conditional distributions as encoding the dynamics in Eq. (9).Source: §B.2 · p. 14

C Derivation of the loss function

• The KL divergence between the forward- and backward-induced distributions is rewritten as a loss, with model-independent terms removed to simplify the objective.The derivation uses the discrete ratio of forward and reverse conditional processes; the expectation of independent Gaussian noise is constant.Source: C Derivation of the loss function · p. 14

D Implemtation of Stochastic Adjoint in PyTorch [34]

• The appendix provides a PyTorch implementation of the stochastic adjoint method, including forward-state storage and backward gradient computation.The implementation saves trajectory states and computes gradients for the state and model parameters.Source: D Implemtation of Stochastic Adjoint… · pp. 15–16

E.1 Time discretization for DDPM

• DDPM uses a fixed forward diffusion with no learnable parameters, discretized with ∆t_i = β_i.Its drift is f(x, t) = −1/2 x.Source: §E.1 · p. 16

E.2 Trajectories on 2D points

• 30 DiffFlow sampling steps versus 500 for DDPM; DiffFlow's 2D trajectories diffuse efficiently while retaining distribution details.The models use similar architectures, with DDPM widened to make model sizes comparable.Source: §E.2 · p. 16

E.3 NLL comparison on 2D points and Sharp Datasets

• DiffFlow performs much better than the other methods on all reported 2D examples, with a larger advantage on distributions concentrated on a 1D submanifold.Without noise, FFJORD more readily transforms smooth than sharp distributions; Figure 8 reports aligned forward and backward trajectory distributions for DiffFlow, unlike the mismatched FFJORD trajectories and detail loss in DDPM.Source: §E.3 · p. 17, Figure 8 · p. 17

F Density estimation

• The density-estimation setup uses ReZero networks, N = 30 for tabular training, normalized inputs, and a five-epoch plateau stopping rule.The authors report standard deviations over three runs and training times of 14 minutes to 2 hours on one RTX 2080 Ti.Source: F Density estimation · pp. 17–18

G.1 Settings

• Terminal time T is described as crucial to training and sampling quality; the best tested values are T = 0.05 for CIFAR10 and T = 0.1 for CelebA.Larger T injects more noise, which helps diffuse data into Gaussian but requires the backward process to remove more noise.Source: §G.1 · p. 18

G.2 Image Completion

• Image completion completes full images from partial masked images using the whole masked forward trajectory rather than only the masked image.The trajectory signals guide backward completion, which iteratively refines masked regions based on unmasked regions.Source: §G.2 · p. 19

G.3 More samples

• Figures 13–15 provide uncurated generated samples for MNIST, CIFAR10, and CelebA.The appendix also labels Figures 11–12 as CIFAR10 and CelebA image-completion samples.Source: §G.3 · p. 19, Figure 11 · p. 21, Figure 12 · p. 22

H Marginal Equivalent SDEs

• Equations (27) and (29) share the same marginal distributions when their distributions at T are identical.The derivation uses the Fokker–Planck–Kolmogorov equation to establish this equivalence.Source: H Marginal Equivalent SDEs · p. 19
• With small T, reconstruction is easier but diffusion is harder; more noise improves diffusion while challenging reconstruction and sample quality.Figure 10 contrasts xN, its reconstruction, and samples from three models.Source: Figure 10 · p. 20