PaperlayerView exact revision on arXiv

Source-linked AI summary

Maximum Likelihood Training of Implicit Nonlinear Diffusion Models

Dongjun Kim, Byeonghu Na, Se Jung Kwon, Dongsoo Lee, Wanmo Kang, Il-Chul Moon

arXiv:2205.13699v3cs.LG

TL;DR

Linear diffusion is common, while nonlinear diffusion has been explored by relatively few works. INDM combines a normalizing flow with linear latent diffusion to induce a nonlinear data-space process. It reports a state-of-the-art FID of 1.75 on CelebA and identifies training and evaluation time as a limitation.

  • Problem

    Nonlinear extensions of linear diffusion have been investigated by relatively few works, and their effects are not well understood.

    Source: §1 · p. 1

  • Method

    INDM uses a normalizing flow to transform a linear latent diffusion into a nonlinear diffusion on the data space.

    Source: §1 · p. 1, §4 · p. 3, §4.1 · p. 4

  • Results

    FID of 1.75 on CelebA is reported as state of the art.

    Source: §7.2 · p. 10

  • Takeaways & Limitations

    INDM provides a data-adaptive nonlinear diffusion that learns both drift and volatility and trains with MLE.

    Source: §1 · pp. 1–2

    INDM has a limitation in training and evaluation time.

    Source: §8 · p. 10

Abstract

from arXiv · show

Whereas diverse variations of diffusion models exist, extending the linear diffusion into a nonlinear diffusion process is investigated by very few works. The nonlinearity effect has been hardly understood, but intuitively, there would be promising diffusion patterns to efficiently train the generative distribution towards the data distribution. This paper introduces a data-adaptive nonlinear diffusion process for score-based diffusion models. The proposed Implicit Nonlinear Diffusion Model (INDM) learns by combining a normalizing flow and a diffusion process. Specifically, INDM implicitly constructs a nonlinear diffusion on the \textit{data space} by leveraging a linear diffusion on the \textit{latent space} through a flow network. This flow network is key to forming a nonlinear diffusion, as the nonlinearity depends on the flow network. This flexible nonlinearity improves the learning curve of INDM to nearly Maximum Likelihood Estimation (MLE) against the non-MLE curve of DDPM++, which turns out to be an inflexible version of INDM with the flow fixed as an identity mapping. Also, the discretization of INDM shows the sampling robustness. In experiments, INDM achieves the state-of-the-art FID of 1.75 on CelebA. We release our code at https://github.com/byeonghu-na/INDM.

1 Introduction

INDM extends linear diffusion models with a trainable, data-adaptive nonlinear forward process. It uses a normalizing flow to map linear latent diffusion into the data space and is presented as fast, likelihood-trained, and robust to sampling discretization.

1 Introduction

  • INDM implicitly constructs a trainable nonlinear diffusion on the data space by transforming a linear latent diffusion through a normalizing flow.The flow's invertibility is key to learning the nonlinear inference part.Source: §1 · p. 1
  • INDM achieves fast and tractable optimization with implicit modeling.Source: §1 · p. 2
  • INDM learns both drift and volatility coefficients of the forward SDE.Source: §1 · p. 2
  • INDM trains its network with Maximum Likelihood Estimation (MLE).Source: §1 · p. 2
  • INDM is robust to sampling discretization.Source: §1 · p. 2

2 Preliminary

Diffusion models define forward corruption and reverse generation through stochastic processes, with a score network estimating the data score. The preliminary section frames the paper's extension beyond the usual linear drift and volatility choices.

2 Preliminary

  • A forward SDE diffuses data toward noise, while its corresponding reverse SDE denoises noise to regenerate data.Source: §2 · p. 2
  • A score network estimates the data score, which is used to construct a generative SDE that runs backward from a prior distribution.Source: §2 · p. 2
  • With likelihood weighting, the denoising score loss gives a variational upper bound on negative log-likelihood, up to a constant.Source: §2 · p. 2
  • The standard diffusion setup restricts the drift to a linear function of the state and the volatility to a time-scaled identity matrix.Source: §2 · p. 2

3 Motivation of Nonlinear Diffusion Process

The paper contrasts the structure-preserving deformation of linear VPSDE with the motivation for more flexible nonlinear diffusion paths.

3 Motivation of Nonlinear Diffusion Process

  • The VPSDE diffusion path preserves the initial data manifold's overall structure as it deforms the distribution toward N(0, I).Source: §3 · p. 3
  • The VPSDE drift vector field linearly deforms the data distribution.Source: §3 · p. 3
  • The illustrated spiral toy dataset provides the setting for comparing diffusion processes during deformation toward N(0, I).Source: §3 · p. 3

4 Implicit Nonlinear Diffusion Model

INDM induces a nonlinear data-space diffusion by mapping a linear latent diffusion through a normalizing flow. Its training combines flow and denoising losses, and sampling can use a feed-forward path through the flow.

4 Implicit Nonlinear Diffusion Model

  • INDM implicitly parameterizes nonlinear drift and volatility through a linear latent diffusion connected to data by a normalizing flow, avoiding intractable transition-probability calculations.The implicit construction is intended to make optimization fast and tractable.Source: §4 · p. 3

4.1 Data and Latent Diffusion Processes

  • The invertible flow maps data into latent space, where a linear forward diffusion is defined, then maps the latent diffusion back to induce a data-space process.When the flow is the identity, the induced diffusion reduces to a linear diffusion.Source: §4.1 · p. 4

4.2 Model Training and Sampling

  • INDM trains the flow using a tractable loss equivalent to the intractable path-measure KL, combining flow and denoising losses; its score target is the latent forward score.For sampling, the paper chooses the latent path that feeds the flow forward only once.Source: §4.2 · p. 4

5 Related Work

The section compares INDM with prior nonlinear and latent diffusion approaches, emphasizing how its flow-based construction differs in modeling, training, and sampling. INDM avoids a speed–flexibility trade-off described for DiffFlow and retains a forward data diffusion unavailable to LSGM.

5 Related Work

  • INDM keeps latent and data dimensions equal and has a forward data diffusion, unlike LSGM, whose non-invertible encoder prevents constructing one when dimensions differ.On CIFAR-10, LSGM’s latent dimension is 40,080 versus 3,072 for data; INDM keeps the data dimension.Source: §5 · p. 5
  • LSGM is reported as unstable at its best-FID setting λ = σ^2, whereas INDM is reported stable across training configurations.Source: §5 · p. 5
  • DiffFlow models only the drift nonlinearly, while INDM also nonlinearizes volatility; DiffFlow’s explicit parameterization trades training speed against diffusion flexibility.DiffFlow’s discretized forward sampling requires O(N) flow evaluations per update, whereas INDM’s training is invariant to N and sampling uses one flow evaluation.Source: §5 · p. 5
  • SBP uses a bi-constrained path-measure KL minimization, and its CIFAR-10 norm ratio indicates that the nonlinear effect is nearly ignorable for most diffusion times.The reported ratio approaches zero except near t ≈ 0; the authors therefore characterize the diffusion as nearly linear for most of the diffusion time.Source: §5 · p. 6

6 Discussion

INDM is reported to train faster than DDPM++ and to approach MLE training, while its sampling is more robust to discretization step sizes. The discussion attributes these properties to trainable flow nonlinearity and relatively linear diffusion trajectories.

6.1 Benefit of INDM in Training

  • INDM is faster than DDPM++ and asymptotically MLE-like on CIFAR-10 with VPSDE, while DDPM++ is the identity-flow version of INDM.The authors attribute the moving score target in INDM to bidirectional attraction between score functions during flow learning.Source: §6.1 · p. 6
  • The variational gap is zero if and only if the score network belongs to Ssol, the class of forward scores of linear diffusions with an initial distribution.The authors state that the gap is irrelevant to flow parameters and that MLE training keeps the score network nearby Ssol.Source: §6.1 · p. 6
  • For any score network in Ssol, a sufficiently flexible optimal flow can make latent forward and generative SDEs coincide, yielding infinitely many global optima for INDM versus one for DDPM++.The authors suggest this difference potentially explains INDM’s faster convergence.Source: §6.1 · pp. 6–7

6.2 Benefit of INDM in Sampling

  • INDM is more robust than DDPM++ and NCSN++ to sampling discretization step sizes in FID, and its sampler becomes more robust as a increases at tolerance 10^-2 on CIFAR-10.Source: Figure 5 · p. 7, §6.2 · p. 7
  • The authors attribute INDM’s lower discretization sensitivity to its loss design and relatively linear, efficient diffusion trajectory, unlike DDPM++’s trajectory that reverses near its end.INDM’s E[∥z∥^4]^(1/4) is reported as 1.3x DDPM++’s; Figure 7-(c) measures trajectory cosine similarity to the straight line between initial and final points.Source: §6.2 · p. 8

7 Experiments

Experiments evaluate INDM on CIFAR-10 and CelebA 64×64 using diffusion backbones, flow models, and alternative score-training weightings. They report likelihood corrections, image-generation comparisons, and an unpaired dataset-interpolation application.

7.1 Correction on Likelihood Evaluation

  • For likelihood evaluation, the authors use a corrected NLL/NELBO by default; the correction changes VPSDE NLL by about 0.03–0.06.They fix the time truncation at ϵ = 10^-5 and account for the reconstruction probability and residual term.Source: §7.1 · pp. 8–9, Table 4 · p. 9

7.2 Quantitative Results on Image Generation

  • After 500k pre-training steps, INDM performs better than DDPM++ on both NLL and FID.Pre-training is used to search for data nonlinearity near well-trained linear diffusions.Source: §7.2 · p. 9, Table 3 · p. 9
  • INDM improves NELBO across the tested settings and minimizes the variational gap to zero with likelihood weighting.Table 4 compares flow training across linear diffusions and weighting functions.Source: §7.2 · p. 9, Table 4 · p. 9
  • FID 1.75 on CelebA is the reported state-of-the-art, while INDM-121M (NLL) also outperforms LSGM-269M (NLL) in FID with identical NLLs.INDM-118M (FID) is on par with LSGM-476M (FID) on CIFAR-10; INDM (NLL) + ST outperforms LSGM-109M (balanced) in FID with comparable NLL.Source: §7.2 · pp. 9–10, Table 6 · p. 9

7.3 Application Task: Dataset Interpolation

  • INDM’s flow network constructs an interpolated bridge between EMNIST letters and MNIST digits; separately estimated densities are 1.56 BPD for EMNIST and 1.10 BPD for MNIST, without an additional neural network.Source: §7.3 · p. 10, Figure 9 · p. 10

8 Conclusion

The paper combines an invertible transformation with diffusion to train a nonlinear forward diffusion beyond the variational family of inference measures. It identifies training/evaluation time as a limitation and notes potential misuse of deep generative models.

8 Conclusion

  • INDM combines an invertible transformation and diffusion to learn a nonlinear forward process beyond the variational family of inference measures.The conclusion presents this as an expansion of linear diffusion to trainable nonlinear diffusion.Source: §8 · p. 10
  • Training and evaluation time are identified as a limitation of INDM.Source: §8 · p. 10
  • The paper identifies deepfake images as a potential negative use of deep generative models.Source: §8 · p. 10

A.1 Derivation of Variational Bound for Nonlinear Diffusion

The appendix derives a variational bound for a general diffusion SDE and explains why direct optimization of nonlinear drift and volatility can be intractable. It also derives INDM’s nonlinear drift and volatility from a linear latent-space SDE.

A.1 Derivation of Variational Bound for Nonlinear Diffusion

  • The variational bound is extended to a general SDE with state-dependent drift and volatility, beyond the setting where G(xt, t) = g(t).The derivation uses reverse-SDE and path-measure relationships before deriving a bound on model log-likelihood.Source: §A.1 · pp. 18–19
  • Directly optimizing nonlinear drift and volatility is generally intractable when transition probabilities require Feynman–Kac expectations over sample paths.Source: §A.1 · p. 19
  • Even with tractable transitions, computing the bound may not scale because GGT multiplication costs O(d^2) and divergence computation adds a bottleneck as dimension increases.Source: §A.1 · p. 19

A.2 Derivation of Nonlinear Drift and Volatility Terms for INDM

  • Applying Itô’s lemma to the linear latent-space SDE through the inverse flow yields INDM’s data-space drift and volatility terms.The appendix derives the drift trace term and uses inverse-function Jacobian relationships in the calculation.Source: §A.2 · pp. 19–20

B Details on Section 6.1

The analysis characterizes which score networks attain zero variational gap and contrasts INDM’s flexible optimum set with DDPM++’s single optimum. It also presents a score regularizer and states the assumptions and qualification attached to the robustness theorem.

B Details on Section 6.1

  • Theorem 2 characterizes zero-gap score networks as Ssol, and the variational gap remains strictly positive while the score network has a nonzero rotation part.Theorem 2 states the zero-gap condition as an if-and-only-if characterization.Source: B Details on Section 6.1 · p. 22
  • Theorem 3 makes every score network in Ssol a candidate INDM optimum, whereas DDPM++ with identity flow has a unique optimum.The paper describes INDM as having infinitely many optimalities, in contrast to DDPM++’s single optimal point.Source: B Details on Section 6.1 · p. 23

B.1 Restricting Search Space of sθ into Sdiv

  • The regularizer targets Sdiv by measuring Jacobian asymmetry with a random-vector estimator, reducing the stated computation from O(d) to O(1).The paper identifies membership in Sdiv with a symmetric score-network Jacobian and describes the estimator as minimum variance.Source: §B.1 · pp. 23–24
  • Regularization reduces DDPM++’s variational gap at 4× slower training speed, with only marginal NLL improvement in certain scenarios; INDM’s NLL gain is attributed to φ-training and its expanded solution space.The authors say the regularized loss is unnecessary for INDM, which already achieves nearly MLE training.Source: §B.1 · p. 24

C.1 Full Statement of Theorem 4

  • Theorem 4 assumes bounded score-estimation error and C2 smoothness; its inequality may not be tight, and empirical robustness is connected to initial-distribution smoothness.Source: §C.1 · p. 24

C.2 Geometric Interpretation of Latent Diffusion

The appendix relates diffusion-path geometry to scale and to the learned latent manifold’s position relative to the prior. Its examples show that flow training can produce straighter trajectories associated with robust sampling.

C.2 Geometric Interpretation of Latent Diffusion

  • On two-moons and ring data, trajectories are more straightforward when the initial-data scale is smaller relative to the prior, showing that scale matters even for the same linear SDE.Source: §C.2 · p. 24, Figure 12 · p. 25
  • On CIFAR-10, the data’s average squared radius is 776 versus the prior shell’s 3,072, while latent radii range from 5,385 to 31,399, placing the latent manifold outside the prior.The prior shell’s squared radius is d = 3,072.Source: §C.2 · pp. 25–26
  • When the initial manifold lies inside the prior, drift contraction can impede movement toward the endpoint; in the latent trajectory outside the prior, drift instead helps produce a straighter path.Source: §C.2 · p. 26
  • After flow training inflates the two-moons latent manifold outside the data manifold, the probability-flow ODE becomes nearly linear and supports robust sampling.The paper also notes that Figure 7 supports the linearized-path illustration on CIFAR-10.Source: §C.2 · p. 26, Figure 14 · p. 26
  • The paper describes the VPSDE encoder map as nearly optimal transport under squared Euclidean cost [46], with trained diffusion trajectories becoming more straight-like.Figure 14 visualizes Monge trajectories between the latent initial distribution and the prior.Source: §C.2 · p. 26, Figure 14 · p. 26

D.1 Latent Score-based Generative Model (LSGM)

This section examines INDM in relation to other diffusion models, covering latent-space objectives and dimensionality, discretization and nonlinearity, and an INDM data-interpolation update with a combined loss. It also describes likelihood and sampling corrections, implementation choices, and appendix comparisons, including how data-adaptive sampling and final-step selection affect FID.

D.1 Latent Score-based Generative Model (LSGM)

  • Although LSGM optimizes a latent-space denoising score loss, that loss is not a data-space path-measure KL, and differing latent and data dimensions preclude a data-space forward diffusion.The stated reason is that a continuous encoder with a continuous inverse cannot map between spaces of different dimensions [47].Source: §D.1 · p. 27
  • LSGM training fails with variance weighting, and the paper links its instability to a forward process starting at an untrained reconstruction rather than the data.The reconstruction xrec differs from x0 during training; the authors contrast this instability with INDM’s stability across training configurations and cite Vahdat et al. [9] and Dockhorn et al. [20].Source: Table 8 · p. 27, §D.1 · p. 27
  • On CIFAR-10, LSGM’s latent dimension is 15x the data dimension, whereas INDM preserves the data dimension while remaining invertible.The table compares the latent dimensions of LSGM and INDM.Source: Table 9 · p. 28, §D.1 · p. 28

D.2 Diffusion Normalizing Flow (DiffFlow)

  • INDM uses O(1) flow evaluations per optimization step, while DiffFlow uses O(N) and relies on a finite Euler–Maruyama chain whose approximation may be poor for strongly nonlinear drift.The cited Lipschitz bound can be arbitrarily large, and for fφ(xt, t) = x^2_t no finite Lipschitz constant satisfies the bound.Source: §D.2 · pp. 29–30
  • INDM jointly nonlinearizes drift and volatility, whereas DiffFlow is limited to nonlinear drift because tractable discretized inference requires omitting a non-Gaussian volatility term.The paper says this omission can incur large approximation error when volatility depends nonlinearly on xt, limiting DiffFlow to Gφ(xt, t) = gφ(t).Source: §D.2 · pp. 30–31
  • DiffFlow’s sample-path loss estimator has prohibitively large variance as N increases, unlike INDM’s analytic-transition-style Monte Carlo estimation.The two losses coincide in expectation but are estimated differently; Table 10 compares their estimator variances.Source: §D.2 · p. 31, Table 10 · p. 32

D.3 Schrödinger Bridge Problem (SBP)

  • For linear VE or VP reference diffusions whose endpoint is already close to the prior, SBP’s optimal auxiliary vector field is nearly zero, limiting its nonlinearity.SBP characterizes the solution as the least-energy admissible auxiliary vector field.Source: §D.3 · p. 32
  • IPF solves Schrödinger bridge problems by alternating half-bridge problems, but training is slow because sampling requires numerically solving an SDE.The cited convergence result concerns the alternating path-measure marginals.Source: §D.3 · p. 33

E.1 Equivalent Reverse SDEs / E.2 Log-Likelihood for Diffusion Models with Time Truncation

  • For reverse SDEs parameterized by λ, substituting the score estimate yields generative SDEs with distinct marginal distributions; truncated-time likelihood estimation uses variational inference for the unobserved early-time scores.Diffusion time is truncated to [ϵ, 1] because the score loss is unbounded, while the omitted small-time range contributes substantially to log-likelihood.Source: §E.1 · p. 34, §E.2 · p. 34

E.3 NELBO Correction / E.5 Calculating the Residual Term

  • The corrected NELBO includes a residual term omitted in prior work, and the reported bound uses the generative SDE at λ = 1.The residual term accounts for the truncated interval; its calculation assumes Gaussian transition distributions for the forward and generative conditionals.Source: §E.3 · p. 35, §E.5 · p. 36

E.4 NLL Correction

  • For NLL correction, the λ = 0 generative ODE makes the sample path deterministic, and likelihood integration must start from the perturbed xϵ rather than x0.The paper approximates λ = 1 log-likelihood with λ = 0 at the cost of a slight difference between their model distributions; starting from x0 causes a larger NLL discrepancy.Source: §E.4 · pp. 35–36

F.1 Model Architecture

  • The flow implementation replaces Glow with ResFlow after empirical tests found better performance, and sine activations improved training stability and FID while maintaining NLL.The authors drop activation normalization, batch normalization, and fully connected layers from ResFlow.Source: §F.1 · p. 37

F.2 Experimental details / F.2.1 Variance Reduction / F.2.3 Interpolation Task

  • The experiments use regenerated baseline performance for fair comparison; importance sampling is applied for NLL, and interpolation loss uses 0.2Gb of GPU memory versus 2.5Gb for the INDM loss in EMNIST-MNIST.For VESDE, importance sampling is equivalent to uniform time sampling, so no VE NLL experimental setting is used.Source: §F.2 · p. 38, §F.2.1 · p. 38, §F.2.3 · p. 41

F.2.2 Sampling Tricks to Improve FID

  • A final denoising-step search found a sweet spot at x^-0.5 to x^-0.75 for VESDE FID.The paper reports line-searched FID performance for each VESDE setting.Source: §F.2.2 · p. 40
  • For INDM (VE), the data-adaptive prior reduced FID-5k from 8.14 to 7.52; for INDM (VP), it was not effective, so the vanilla prior was used.Source: §F.2.2 · pp. 40–41

F.2.3 Interpolation Task

  • The interpolation procedure sets Ltot = LINDM + Lint and updates φ using the gradient of Ltot.Source: §F.2.3 · p. 41

F.3 Effect of Pre-training

The appendix examines pre-training, training cost, latent-space visualizations, learned diffusion coefficients, and relative transport energy. It also points to full quantitative comparisons and non-cherry-picked samples.

F.3 Effect of Pre-training

  • Pre-training the DDPM++ score network before 350k further flow-and-score training improves FID, and the authors report an empirical advantage to searching near the linear process.The ablation varies pre-training from 100k to 500k steps; quantitative performances in Section 7 use pre-training.Source: §F.3 · pp. 41–42, Table 17 · p. 41

F.4 Training Time

  • INDM’s training time is invariant to discretization count N, unlike SBP and DiffFlow, whose per-step training costs scale as O(N).Source: §F.4 · p. 42, Table 18 · p. 42
  • Total CIFAR-10 training times for INDM, DiffFlow, and LSGM are reported at a similar scale, while SBP is fastest; comparisons are not fully fair across differing setups.The table reports single-GPU time; the authors estimate some baselines from released configurations and note GPU differences.Source: §F.4 · p. 42, Table 19 · p. 42

F.5.1 Visualization of 2d Latent Manifold

  • Training removes the singularity at the origin in the 2D checkerboard data manifold, which is absent from the resulting latent manifold.Source: §F.5.1 · p. 42

F.5.2 Visualization of High-dimensional Latent Vector on Benchmark Datasets

  • Figure 18 presents benchmark data-space and latent-space samples, with latent vectors normalized to [0, 1]^d for visualization.Source: §F.5.2 · p. 42, Figure 18 · p. 42

F.6 Nonlinear Diffusion Coefficient

  • INDM’s learned covariance eigenvalues disperse over diffusion time and vary across data instances, indicating nonlinear, location-dependent diffusion.Source: §F.6 · p. 43, Figure 19 · p. 43

F.7 Relative Energy

F.8 Full Quantitative Tables

  • The appendix provides full quantitative comparisons to baselines and non-cherry-picked INDM samples for CIFAR-10 (VE) and CelebA (VP).Source: §F.8 · p. 43, §F.9 · p. 43

G Proofs of Theorems and Propositions

The appendix proves results about likelihood bounds, score-error gaps, and conditions for zero variational gap, then establishes propositions and bounds used in the discretization analysis.

Theorem 1

Theorem 2

Theorem 3

Propositions 2–3

Proposition 5

Theorem 4

Loading 2205.13699v3…