PaperlayerView exact revision on arXiv

Source-linked AI summary

A Variational Perspective on Diffusion-Based Generative Models and Score Matching

Chin-Wei Huang, Jae Hyun Lim, Aaron Courville

arXiv:2106.02808v2cs.LG

TL;DR

Score matching for plug-in reverse SDEs lacked a theoretical link to likelihood and sampling. The paper derives a variational framework for continuous-time diffusion, showing that a particular score-matching objective maximizes a likelihood lower bound for a family of plug-in reverse SDEs.

  • Problem

    The relationship between minimizing score-matching loss and the statistical behavior of the resulting plug-in reverse SDE was unclear.

  • Method

    The paper combines Feynman-Kac and Girsanov theorems to construct a continuous-time variational likelihood framework extending hierarchical VAEs to infinite depth.

  • Results

    Score matching with a particular mixture weighting is equivalent to maximizing a marginal-likelihood lower bound for a family of plug-in reverse SDEs, including Song et al.'s SDE and an equivalent ODE.

  • Takeaways & Limitations

    The framework provides a likelihood-based theoretical foundation for score matching and permits comparison of continuous-time flows with score-trained plug-in reverse SDEs.

  • Takeaways & Limitations

    The paper cautions that score matching is not generally equivalent to maximum likelihood because the relevant inference density need not match the generative density or learned score.

Abstract

from arXiv · show

Discrete-time diffusion-based generative models and score matching methods have shown promising results in modeling high-dimensional image data. Recently, Song et al. (2021) show that diffusion processes that transform data into noise can be reversed via learning the score function, i.e. the gradient of the log-density of the perturbed data. They propose to plug the learned score function into an inverse formula to define a generative diffusion process. Despite the empirical success, a theoretical underpinning of this procedure is still lacking. In this work, we approach the (continuous-time) generative diffusion directly and derive a variational framework for likelihood estimation, which includes continuous-time normalizing flows as a special case, and can be seen as an infinitely deep variational autoencoder. Under this framework, we show that minimizing the score-matching loss is equivalent to maximizing a lower bound of the likelihood of the plug-in reverse SDE proposed by Song et al. (2021), bridging the theoretical gap.

1 Introduction

Generative modeling can be formulated as inverting an inference process, with diffusion models fixing the inference path and learning a reverse process. This paper develops a variational framework connecting score matching to likelihood estimation for plug-in reverse SDEs.

  • Generative modeling inverts an inference process, ranging from invertible transformations to stochastic variational autoencoders.
  • Diffusion-based models freeze a discrete-time inference chain that transforms data into a tractable prior, then train another chain to reverse it.
  • Song et al. derive a reverse-time SDE using the score of perturbed data and define a generative model by substituting a learned score.
  • The paper combines Feynman-Kac and Girsanov results to derive a continuous-time variational likelihood framework and an infinitely deep ELBO.

2 Background

The background introduces inference SDEs that transform data toward a tractable prior and reverse SDEs that recover the data distribution when the score is known. It then frames the unresolved question as how score-matching optimization affects the plug-in reverse SDE.

  • The inference SDE transforms data toward a tractable prior, while its time-reversed SDE has matching marginals when initialized from that prior.
  • Access to the score ∇log q allows the reverse SDE to generate samples from the transformed data process.
  • Equivalent reverse SDEs form a family indexed by λ, with the displayed reverse SDE corresponding to λ = 0.
  • Score matching offers ESM, ISM, SSM, and DSM losses, which are equivalent up to additive constants; Λ(s) weights the loss.
  • The plug-in reverse SDE replaces the true score with sθ, but its relationship to score-matching optimization and statistical closeness to data was previously unclear.

3 Marginal density and stochastic instantaneous change of variable

The paper formulates generative diffusion densities through the Fokker-Planck equation and estimates them using a Feynman-Kac probabilistic representation. With position-independent diffusion, stochastic dynamics can be analyzed as infinitesimal volume-preserving transformations plus drift effects.

  • The generative SDE starts from a tractable prior and induces densities whose evolution follows the Kolmogorov forward, or Fokker-Planck, equation.
  • Assuming position-independent diffusion simplifies the Fokker-Planck equation to drift, divergence, and Hessian terms.
  • The Feynman-Kac formula provides a probabilistic representation for estimating solutions of the resulting second-order linear PDE.
  • Changing variables from v(x,T−t) converts backward Feynman-Kac evolution into forward SDE evolution and represents the terminal marginal density.
  • For a fixed Brownian path, position-independent diffusion acts as a constant additive, volume-preserving transformation, so density changes are computed through the drift.
  • The framework also covers position-dependent diffusion, but the formulas must include spatial partial-derivative terms.

4 Inferring latent Brownian motion

The paper treats Brownian paths as infinite-dimensional latent variables and uses a change of measure to derive a tractable continuous-time ELBO. The resulting inference SDE has an optimal drift characterized by the generative diffusion's score.

  • Latent Brownian inference: Brownian motion is treated as a latent variable, turning likelihood estimation into inference over sample paths in an infinite-dimensional VAE.The goal is to assign higher probability to paths more likely to generate the observation.
  • Change of measure: Girsanov’s theorem supplies the change of measure needed to replace an intractable path-space likelihood with a finite lower bound and a reparameterizable Brownian process.The framework addresses both numerical computation of the density ratio and gradient estimation under the new law.
  • ELBO computation: The continuous-time ELBO is numerically estimable with black-box SDE solvers by accumulating ||a||2 and ∇· µ, with Hutchinson estimation available for the divergence.The generative and inference model parameters are decoupled from the random variables used for estimation.
  • Inference drift: The inference drift can be amortized to encode datum-specific information, while its ||a||2 regularization keeps the changed measure close to the centered Wiener measure.The framework also permits non-Markovian extensions, although the paper uses a Markovian model for simplicity.
  • Optimal inference SDE: The variational gap vanishes exactly when the inference drift is Markovian in the state and equals σ⊤∇log p(y, T −s) almost everywhere.This characterizes the optimal inference SDE and makes a non-amortized Markovian process sufficiently expressive.

5 Infinitely deep hierarchical VAE

The continuous-time ELBO is the limit of the ELBO for discretized hierarchical VAEs, supporting its interpretation as training an infinitely deep stochastic autoencoder. Figure 2 further illustrates the relationship between score matching and likelihood bounds across plug-in reverse SDEs.

  • 5 Infinitely deep hierarchical VAE: Discretized generative and inference SDEs form a hierarchical VAE with L stochastic layers and a corresponding discrete-time ELBO.The construction uses step size ∆t = T/L.
  • 5 Infinitely deep hierarchical VAE: The discrete-time reconstruction error converges to the continuous-time terms ∇· µ and ||a||2 as ∆t approaches zero.The divergence arises from a finite-difference approximation, while the drift regularization comes from comparing Gaussian transition means.
  • 5 Infinitely deep hierarchical VAE: Maximizing the continuous-time ELBO can therefore be viewed as training an infinitely deep hierarchical VAE.Taylor’s theorem controls the polynomial approximation error, which vanishes as the layer count increases.
  • Figure 2: Figure 2 uses rows for different λ values and columns for iterates, while its right panel shows lower bounds across a continuum of plug-in reverse SDEs.The caption states that the lower bound is optimized when the score-matching loss is minimized.

6 Score-based generative modeling

The paper connects score matching to likelihood maximization for a family of plug-in reverse SDEs, including the original reverse SDE and an equivalent ODE. Experiments examine the resulting computation–estimation trade-offs and bias–variance choices.

  • Score-based generative modeling: Matching the score of q(y, s) is equivalent to maximizing a lower bound on the plug-in reverse SDE’s marginal likelihood.The network sθ approximates the score ∇log q(y, s), and the connection follows by comparing the ELBO integrand with implicit score matching.
  • Score-based generative modeling: The reparameterized generative and inference SDEs define a λ-indexed family whose special cases include Song et al.’s plug-in reverse SDE and an equivalent ODE.The original reverse SDE corresponds to λ = 0, while the ODE is obtained as λ = 1.
  • Score-based generative modeling: Minimizing implicit score matching loss maximizes the CT-ELBO for a continuum of plug-in reverse SDEs, including the ODE as λ approaches 1.The corresponding ELBOs and score-matching objectives agree up to multiplying and additive constants; only λ = 0 has a directly estimable constant in practice.
  • Computational trade-off: Continuous-time flows converge rapidly by parameter updates, but their runtime-normalized convergence is nearly indistinguishable from plug-in reverse SDEs.Flows spend computation on numerical integration, whereas reverse SDE training uses random time steps, enabling more updates but noisier gradients within a fixed runtime.
  • Bias and variance trade-off: The likelihood integral can use Hutchinson trace estimation for divergence, while denoising score matching avoids the high variance of implicit estimation when ∇sθ is large.The denoising estimator uses a reweighted loss; its practical bias–variance behavior depends on the time-sampling distribution.
  • Bias and variance trade-off: Uniform time sampling emphasizes noisier data, whereas debiased sampling favors smaller perturbations and can reduce the associated weighting bias.The experiments compare parameterizing sθ versus a on MNIST and CIFAR10, with uncertainty estimated across three random seeds.

7 Related work

The paper situates its contribution among diffusion models, reverse-time processes, score matching, de Bruijn identities, and neural SDEs. Its distinction is a marginal-density variational analysis connecting implicit score matching to plug-in reverse dynamics.

  • Diffusion-based generative models: The work gives a theoretical foundation for viewing diffusion models and conditional denoising score matching as learning to reverse an inference process.Unlike the discrete-time ELBO-to-DSM connection, it derives implicit score matching from the marginal density’s Fokker–Planck equation using Feynman–Kac and Girsanov arguments.
  • Time-reversal of diffusion processes: Its marginal-equivalent reverse processes differ from classical reverse-time diffusions because equivalence is asserted through marginals when sθ equals the true score.The comparison concerns the distributions over time rather than necessarily identical pathwise processes.
  • De Bruijn’s identity: Prior work represents KL divergence through weighted Fisher divergence, whereas this paper links score matching to a likelihood lower bound for plug-in reverse SDEs.The paper argues that a direct maximum-likelihood equivalence does not follow from the prior KL identity under its generative-process construction.
  • Learning SDEs: Compared with neural-SDE approaches using an emission probability, this framework directly expresses the marginal density as an expectation via Feynman–Kac.This avoids smoothing the density with an emission probability that would be a Dirac point mass in the considered setting.

8 Conclusion and Discussion

The paper develops a continuous-time variational likelihood framework for diffusion models and connects score matching to likelihood lower bounds. It also identifies computational tools and assumptions needed to make the framework tractable.

  • Conclusion and Discussion: The framework estimates marginal likelihoods for continuous-time diffusion models, encompassing continuous-time normalizing flows and score-based generative models.It treats the Brownian motion as an infinite-dimensional latent variable and performs inference over likely sample paths.
  • Conclusion and Discussion: Both generative and inference SDE parameters can be learned with a numerical solver having constant memory cost.The framework can reduce training time when f and g are simple enough that perturbed samples can be obtained without numerical integration.
  • Conclusion and Discussion: The CT-ELBO is numerically estimable with black-box SDE solvers by augmenting the dynamics with norm and divergence accumulators.The divergence can be estimated efficiently with Hutchinson’s trace estimator and reverse-mode automatic differentiation.
  • Conclusion and Discussion: The score-matching identities rely on regularity conditions and a Feynman–Kac assumption on the relevant functions.The proof also uses integration by parts, Stokes’ theorem, and Hutchinson estimation for sliced score matching.
  • Conclusion and Discussion: The continuous-time change-of-variables derivation extends finite-step Jacobian products to an Itô-integral limit under a uniformly Lipschitz drift.This yields the continuous-time flow formula with the drift evaluated along the Itô solution.

F Proofs

The proofs establish the continuous-time ELBO, characterize its variational gap, and prove convergence from the discrete-time ELBO under regularity assumptions.

  • Continuous-time ELBO: The continuous-time ELBO is obtained by rewriting the likelihood lower bound using stochastic-calculus representations.The derivation uses a reparameterized process and the Feynman-Kac, Itô, and Fokker–Planck identities.
  • Variational gap: The variational gap vanishes exactly when the inference drift equals the score-based optimal drift almost everywhere.The condition is a(Y_s(ω),s) = σ^T∇log p(Y_s(ω),T−s).
  • Variational gap: The gap proof subtracts the lower bound from the marginal likelihood and eliminates stochastic terms through expectation, Fokker–Planck expansion, and conditional Fubini.The proof also uses Itô’s formula and the covariance relation dY_s dY_s^T = σσ^Tds.
  • Consistency: Under bounded fourth-order derivatives and non-singular diffusion, the discrete-time ELBO converges to the continuous-time ELBO as L →∞.The proof controls approximation errors with dominated convergence and Euler–Maruyama mean-square convergence.
  • Consistency: The discrete-to-continuous proof expands drift differences with Taylor’s theorem and cancels determinant terms through Hutchinson’s identity, Jacobi’s formula, and telescoping.The finite-layer reconstruction error therefore approaches the continuous-time expression as the step size decreases.

G Equivalent SDEs

The paper constructs a family of equivalent diffusion SDEs sharing the same Fokker–Planck evolution, then derives their equivalent reverse processes parameterized by λ.

  • Definitions: Equivalent SDEs are defined by equality of the distributions induced by their processes at every time.Reverse equivalence compares a process with the time-reversal of another process.
  • Equivalent SDE construction: A family indexed by λ is constructed by rearranging the Fokker–Planck equation while preserving the original diffusion’s evolution.The construction assumes a position-independent diagonal diffusion coefficient for simplicity.
  • Reverse processes: Rearranging the Fokker–Planck equation of each member yields an equivalent reverse SDE with corresponding drift and diffusion coefficients.The reverse process starts from the terminal density q(·,T).
  • Special cases: The plug-in reverse SDE of Song et al. corresponds to λ = 0, while the equivalent plug-in reverse ODE corresponds to λ = 1.The paper refers to Figure 5 for simulations across λ values.

H Score matching and plug-in reverse SDEs

The paper connects score matching to likelihood maximization for a family of plug-in reverse SDEs by matching generative and inference diffusion coefficients.

  • Equivalent family: The same equivalence extends to the family of equivalent plug-in reverse SDEs, provided their generative and inference diffusion coefficients match.Matching diffusion coefficients ensures a finite Radon–Nikodym derivative.
  • ELBO connection: For λ < 1, the continuous-time ELBO of the plug-in reverse SDE is expressed through the score-matching discrepancy.The derivation applies the score-matching loss identity after substituting the generative and inference SDEs.
  • ELBO connection: Averaging the ELBO over the data distribution yields the connection between score matching and maximum likelihood.The result applies the identity relating the score-matching objective to the expected ELBO.
  • λ dependence: The average continuous-time ELBO is maximized at λ = 0, whereas λ = 1 is excluded because the expression divides by zero.The paper notes that λ values close to 1 produce effectively deterministic SDEs and relates this to the low BPD of the plug-in ODE.
  • Experimental setup: The experiments use a variance-preserving SDE and non-uniform sampling to debias the denoising score-matching loss weighted by σ².The variance-preserving schedule uses β(s) = (βmax−βmin)s + βmin.

J Experiments

The experiments use variance-preserving diffusion models on CIFAR10 and MNIST with dataset-specific network configurations and standard optimization procedures.

  • Model configuration: CIFAR10 uses a modified U-Net, while MNIST uses three feature-map resolutions, fewer channels, and no dropout.Both experiments use the variance-preserving SDE with βmin = 0.1, βmax = 20, and T = 1.
  • Optimization and preprocessing: Optimization uses Adam with learning rate 0.0001 and minibatch size 128 for all experiments.The preprocessing applies uniform dequantization and a logit transform with α = 0.05.
  • Optimization and preprocessing: CIFAR10 additionally uses random horizontal flipping for regularization.The implementation details are supplemented by the project repository referenced by the paper.
Loading 2106.02808v2…