PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling

Valentin De Bortoli, James Thornton, Jeremy Heng, Arnaud Doucet

arXiv:2106.01357v5stat.MLcs.LGmath.PR

TL;DR

Standard score-based generation requires a long noising process to approach a Gaussian prior. The paper formulates generation as a Schrödinger bridge and proposes DSB, an IPF approximation using score-based diffusions. DSB supports shorter intervals, improves endpoint matching over iterations, and is supported by convergence analysis and experiments.

  • Problem

    Standard score-based generation requires sufficiently long forward noising for its final distribution to approximate the Gaussian prior, making generation computationally expensive.

  • Method

    DSB approximates Iterative Proportional Fitting for the Schrödinger bridge using score-based forward and backward diffusions.

  • Results

    DSB provides quantitative convergence results, improves generative models with multiple iterations, and supports generation using shorter time intervals.

  • Takeaways & Limitations

    DSB extends score-based generative modeling to finite-time Schrödinger bridges and also provides a continuous-state computational optimal transport tool.

  • Takeaways & Limitations

    The theoretical framework assumes a time-homogeneous diffusion with drift f(x) = −αx and relies on conditions on the data distribution and score approximation.

Abstract

from arXiv · show

Progressively applying Gaussian noise transforms complex data distributions to approximately Gaussian. Reversing this dynamic defines a generative model. When the forward noising process is given by a Stochastic Differential Equation (SDE), Song et al. (2021) demonstrate how the time inhomogeneous drift of the associated reverse-time SDE may be estimated using score-matching. A limitation of this approach is that the forward-time SDE must be run for a sufficiently long time for the final distribution to be approximately Gaussian. In contrast, solving the Schrödinger Bridge problem (SB), i.e. an entropy-regularized optimal transport problem on path spaces, yields diffusions which generate samples from the data distribution in finite time. We present Diffusion SB (DSB), an original approximation of the Iterative Proportional Fitting (IPF) procedure to solve the SB problem, and provide theoretical analysis along with generative modeling experiments. The first DSB iteration recovers the methodology proposed by Song et al. (2021), with the flexibility of using shorter time intervals, as subsequent DSB iterations reduce the discrepancy between the final-time marginal of the forward (resp. backward) SDE with respect to the prior (resp. data) distribution. Beyond generative modeling, DSB offers a widely applicable computational optimal transport tool as the continuous state-space analogue of the popular Sinkhorn algorithm (Cuturi, 2013).

1 Introduction

Score-based generative modeling reverses diffusion-based noising, but standard methods require long diffusion horizons for the noisy distribution to approach a Gaussian prior. The paper reformulates generation as a Schrödinger bridge and introduces DSB, whose iterations improve finite-time matching while supporting theory and experiments.

  • Motivation: Score-based generative models add noise toward an easy-to-sample prior, then learn reverse-time denoising dynamics with neural networks and score matching.The reverse-time SDE drift depends on scores, the logarithmic gradients of perturbed data distributions.
  • Motivation: Standard score-based generation requires sufficiently long forward diffusion and small numerical steps to approach the prior and accurately approximate the SDE.This makes sample generation computationally expensive.
  • Schrödinger bridge formulation: The Schrödinger bridge finds the closest reference diffusion whose endpoint marginals are the data distribution at t = 0 and the prior at t = T.It is formulated as an entropy-regularized optimal transport problem on path spaces.
  • Method: Diffusion SB implements Iterative Proportional Fitting with score-based diffusions, without state-space discretization, potential regression, or kernel density estimation.Its first iteration recovers Song et al. (2021), while later iterations reduce endpoint-marginal discrepancies.
  • Theory: The paper provides quantitative convergence results for score-based modeling and continuous-state IPF, including a path-space interpretation of DSB as discretized dynamic IPF.The analysis does not rely on classical compactness assumptions.
  • Experiments: Experiments on MNIST and CelebA show that multiple DSB steps improve the generative model, and DSB can interpolate between data distributions.Figure 1 reports convergence after 5 DSB iterations when reference forward and reverse diffusions fail within T = 0.2 and N = 20 steps.

2 Denoising Diffusion, Score-Matching and Reverse-Time SDEs

Score-based generative modeling estimates scores of progressively noised data to approximate reverse-time diffusion, but accuracy depends on sufficiently small discretization steps and forward evolution toward a prior. The section develops continuous-time interpretations and total-variation convergence bounds for this procedure.

  • Score estimation and reverse-time sampling: The forward process is designed so its final marginal approximates an easy-to-sample prior, enabling reverse ancestral sampling to approximate the data distribution.Sampling starts from X_N drawn from the prior and proceeds through reverse-time transitions; the approximation depends on p_N being close to p_prior.
  • Score estimation and reverse-time sampling: Small step sizes are required because reverse-transition approximations rely on nearby successive states and Taylor expansions of the score.The approximation is justified when the increment between successive states is small, which is ensured by choosing γ_{k+1} sufficiently small.
  • Score estimation and reverse-time sampling: Score-based generative modeling estimates perturbed-data scores with neural networks and uses them to approximate reverse-time transitions for sampling.The reverse-time Markov chain is interpreted as an Euler–Maruyama discretization whose score functions are replaced by learned approximations.
  • Continuous-time formulation: The continuous-time reverse process is obtained by time reversal of an SDE, while the associated reverse Markov chain discretizes it with learned score functions.The framework includes Brownian and Ornstein–Uhlenbeck processes through the drift choice f(x) = −αx.
  • Convergence analysis: Theorem 1 bounds total-variation error through score approximation, discretization, and forward-process mixing terms, with separate forms for α > 0 and α = 0.For α > 0 the bound includes C_α(M + γ̄^1/2) exp[D_αT] + B_α exp[−α^1/2T]; for α = 0 it includes C_0(M + γ̄^1/2) exp[D_0T] + B_0(T^−1 + T^−1/2).
  • Convergence analysis: The convergence analysis identifies a trade-off: stronger drift improves continuous-time mixing but worsens discrete-time approximation unless smaller steps are used.The constants in the bound grow with α, while increasing α also accelerates convergence of the continuous-time process.

3 Diffusion Schrödinger Bridge and Generative Modeling

This section formulates generative modeling as a Schrödinger bridge and develops Diffusion Schrödinger Bridge as a score-based approximation to Iterative Proportional Fitting. It also establishes convergence properties for IPF and describes how DSB produces samples while using fewer diffusion steps than standard score-based modeling.

  • Generative modeling as a Schrödinger bridge: The Schrödinger bridge minimizes path-space KL divergence to a reference diffusion while matching pdata at time 0 and pprior at time T.Its dynamic formulation has a static endpoint analogue and can be interpreted as an entropy-regularized optimal transport problem.
  • Diffusion Schrödinger Bridge: DSB samples by initializing at pprior and applying learned reverse-time transitions, with the resulting terminal sample approximately distributed as pdata.The method learns network parameters through empirical score-related losses and fine-tunes later iterations from earlier ones.
  • Iterative Proportional Fitting: Iterative Proportional Fitting alternately enforces the prior and data endpoint marginals, converging toward the Schrödinger bridge under suitable assumptions.The IPF iterates are well-defined under a finite-KL feasibility condition and admit a representation through alternating forward and reverse processes.
  • Diffusion Schrödinger Bridge: Diffusion SB implements IPF with score-based diffusions, avoiding state-space discretization, potential-function regression, and kernel density estimation.The first iteration recovers the score-based reverse-process methodology, while later iterations refine the forward and backward dynamics.
  • Diffusion Schrödinger Bridge: DSB requires a sequence of learned network parameters but uses fewer diffusion steps than standard score-based generative modeling.The initial networks can be trained similarly to prior score-based methods, and subsequent networks can be refined from previous iterations.
  • Convergence of Iterative Proportional Fitting: Theoretical analysis establishes quantitative IPF convergence in non-compact settings and convergence in total variation of the joint distributions toward an IPF limit.The results include monotonicity properties and, under additional conditions, identify the limiting distribution with a solution of the Schrödinger bridge problem.

4 Experiments

Experiments show that DSB converges toward Schrödinger bridges and improves generative modeling across Gaussian, toy, image, and interpolation settings. The method remains effective with shorter diffusion horizons and fewer steps than initial score-based methods, although performance does not yet match state of the art.

  • Gaussian example: In Gaussian experiments, DSB converges toward the analytically known Schrödinger bridge across successive iterations.Figure 2 tracks estimated mean, variance, and first-component covariance against ground truth.
  • Two dimensional toy experiments: With T = 0.2, the first DSB iteration is unsatisfactory because p_N is not close to p_prior, whereas multiple iterations improve synthesis quality.This setting directly illustrates the advantage of iterating beyond the score-based first step.
  • Generative modeling: The implementation is the first practical high-dimensional approximation of the SB problem, but it does not yet compete with state-of-the-art generative methods.Experiments include MNIST and CelebA, with CelebA having dimension d = 3072.
  • Generative modeling: On CelebA, DSB produces qualitative improvements over eight iterations using as few as N = 12 diffusion steps.Higher N gives better results while remaining below the N = 100 used in the cited original SGM procedures; FID improves with DSB iterations.
  • Dataset interpolation: DSB is demonstrated for interpolation between Swiss-roll and S-curve distributions and between EMNIST and MNIST.The demonstrations use iteration 9 with T = 1, N = 50 for the first pair and iteration 10 with T = 1.5, N = 30 for the second.

5 Discussion

DSB extends score-based generative modeling through Schrödinger bridges, enabling shorter diffusion horizons while retaining a diffusion-based formulation. The discussion also identifies theoretical regularity assumptions and practical trade-offs in time horizon, memory, and neural-network evaluations.

  • Discussion: DSB extends score-based generative modeling by interpreting SGM as the first stage of solving a Schrödinger bridge problem.This interpretation motivates generative modeling with fewer diffusion steps and connects DSB to accelerated SGM methods.
  • Discussion: Choosing the final time T too close to 0 decreases generative-model quality, despite DSB permitting arbitrarily small positive horizons.The discussion attributes this behavior partly to the forward-process endpoint remaining too far from the prior.
  • Discussion: The reverse-time SDE approach estimates drift from score networks and uses numerical SDE integrators for sampling.This provides the score-based diffusion mechanism underlying the first DSB iteration.
  • Discussion: For fixed time horizons, the analysis combines bounds that are useful at small and large times, although individual bounds may diverge at one endpoint.One bound grows as t approaches infinity, while another remains bounded there but diverges as t approaches 0.
  • Discussion: The theoretical estimates require smoothness, boundedness, or growth conditions on the data density and its logarithmic derivatives.The stated assumptions include differentiability and bounds on the score and radial score growth.

C.2 Proof of Theorem 1

The proof establishes regularity and convergence estimates for score-based diffusions by combining measure-theoretic identities, Gaussian-process bounds, and diffusion arguments. It uses curvature and growth assumptions to control logarithmic derivatives and distributional distances over a fixed time horizon.

  • C.2 Proof of Theorem 1: The proof uses Markov-kernel identities and total-variation contraction to relate diffusion marginals across time.These identities support the later comparison between the target and Gaussian reference distributions.
  • C.2 Proof of Theorem 1: Girsanov’s theorem compares diffusion laws under different drifts when the required drift-square integrability conditions hold.The proof applies this comparison to processes driven by path-dependent drifts.
  • C.2 Proof of Theorem 1: Curvature assumptions on the data density yield sub-Gaussian behavior and moment bounds for conditional distributions.The argument introduces drift conditions and Lyapunov functionals to control these moments.
  • C.2 Proof of Theorem 1: The proof transfers Ornstein–Uhlenbeck estimates to the considered diffusion through scaling relations and semigroup arguments.This connects the transformed process to Gaussian equilibrium estimates over the fixed horizon T.
  • C.2 Proof of Theorem 1: Derivative estimates for the score are obtained by differentiating the density evolution equation and controlling spatial and temporal growth.The resulting bounds support regularity and approximation arguments for the learned score.

D.3 Proof of Proposition 2

The proof connects IPF updates to constrained maximum-likelihood and ELBO problems, then develops variational formulas used to approximate successive bridge distributions. It also compares the computational costs of alternative parameterizations.

  • D.3 Proof of Proposition 2: The IPF sequence is well-defined under finite Kullback–Leibler divergence and positivity conditions on the relevant densities.The proof recursively constructs densities for alternating IPF iterates.
  • D.3 Proof of Proposition 2: The first IPF step is equivalent to maximizing an ELBO associated with a constrained maximum-likelihood problem.Subsequent IPF steps admit analogous ELBO formulations with alternating endpoint constraints.
  • D.3 Proof of Proposition 2: Three variational formulations are identified for score, backward-kernel, and drift quantities in successive IPF iterates.These formulations provide alternative routes for constructing the distributions used by the algorithm.
  • D.3 Proof of Proposition 2: One formulation requires O(nN) neural-network passes and growing memory, whereas the other two require O(N) passes and fixed memory for each n.The approach associated with Proposition 3 is reported to perform better in practice.

E.1 Proof of Proposition 3

The proof section develops variational representations linking IPF iterates with score-matching quantities and places these results within the paper’s broader theoretical study of Schrödinger bridges.

  • E.1 Proof of Proposition 3: The proof derives variational formulas for the quantities appearing in successive IPF updates.These formulas connect the IPF procedure with score-based representations used by DSB.
  • E.1 Proof of Proposition 3: The variational representation is obtained through conditional-expectation and differentiation arguments under the stated regularity conditions.The proof uses dominated convergence to justify the relevant operations.
  • E.1 Proof of Proposition 3: The surrounding theory studies monotonicity properties of IPF and establishes quantitative convergence bounds.These results are presented as part of the theoretical analysis of Schrödinger bridges.

F.1.1 Monotonicity results

Under B1, the IPF sequence is well-defined and exhibits monotone decreases in successive divergences and total-variation changes. These properties support convergence in total variation to a limiting probability measure.

  • Well-posedness: B1 ensures that the IPF sequence is well-defined, with positive potential functions at every iteration.The construction proceeds recursively from finite initial KL divergence and equivalent endpoint marginals.
  • Monotonicity: KL(π_n+1|π_n) ≤ KL(π_n−1|π_n) and KL(π_n|π_n+1) ≤ KL(π_n|π_n−1) for n ≥ 1.The same monotonicity holds separately along the even and odd subsequences.
  • Monotonicity: The total-variation increments ∥π_n+1 − π_n∥TV are non-increasing, and the corresponding symmetric KL quantity J(π_n+1, π_n) is also non-increasing.The total-variation result follows from marginal identities and data processing, while J inherits monotonicity from the KL bounds.
  • Convergence: There exists a probability measure π∞ such that lim_n→+∞ ∥π_n − π∞∥TV = 0.The proof establishes summability of successive total-variation increments and then uses completeness of finite signed measures.

F.2 Proof of Proposition 5

The proof identifies the total-variation limit of IPF as the Schrödinger bridge under progressively broader sufficient conditions. It combines convergence, factorization of the limiting density, and uniqueness arguments.

  • Identification of the limit: If h is continuous and positive with suitable product bounds, the limiting density factorizes and π∞ equals the Schrödinger bridge π⋆.The factorization is matched to the Schrödinger system, whose solution is unique under the stated conditions.
  • Factorization: The proof uses a measurable factorization result showing that the limiting density is either zero or a product a(x)b(y) almost everywhere.This product structure is the key bridge between the IPF limit and the Schrödinger bridge characterization.
  • Convergence: The IPF sequence converges in total variation to a probability measure π∞ under B1.The broader convergence theorem establishes existence of π∞ before identifying it with the Schrödinger bridge.
  • Identification of the limit: Alternatively, π∞ = π⋆ follows when KL(π∞|µ) < +∞ and L1(ν0) ⊕ L1(ν1) is closed in L1(π∞).This proposition avoids the stronger continuity and product-bound assumptions used in the preceding identification result.

G.1 Geometric convergence rates

In a Gaussian setting, IPF converges geometrically with a rate determined by problem constants rather than dimension. The rate improves when the reference marginals are weakly coupled or the target concentrates near a point.

  • Geometric rate: KL(π⋆|π_n) ≤ Cκ^2n with κ = ρ/(1 + ρ) < 1 and ρ = 2α/β^2.The bound holds for α ∈ (0, 1) and β > 0, with C ≥ 0 explicit in the proof.
  • Parameter dependence: IPF improves as κ approaches 0, which occurs when α ≈ 0 or β ≈ +∞; in the corresponding independent or Dirac-mass limits, convergence occurs in two iterations.The figure is described as showing the evolution of κ^2 with α and β.
  • Parameter dependence: The geometric convergence rate depends on α and β, not on the dimension.Here α controls dependence between the reference marginals, while β controls the Gaussian target concentration.
  • Proof structure: The Gaussian recursion alternates updates of quadratic coefficients, whose even and odd subsequences converge to the fixed point γ⋆.For β^2 > 1 − α^2, the even subsequence is non-decreasing and the odd subsequence is non-increasing.
  • Special cases: When β^2 = 1 − α^2, the IPF sequence is stationary because the reference measure is invariant under IPF.In this case, the limiting density equals the reference density.

H.1 Proof of Proposition 6

The continuous-time IPF construction alternates KL projections of path measures onto terminal and initial marginal constraints. Under Markov references, these projections remain Markov and correspond to forward or reverse diffusions.

  • Dynamic Schrödinger bridges: The dynamic Schrödinger bridge exists exactly when its static counterpart exists, and it is Markov whenever the reference measure is Markov.This property transfers static IPF well-posedness to the path-space formulation.
  • Diffusion representation: The projected path measures are Markov and satisfy the martingale problems associated with the corresponding diffusion processes.The recursion preserves the diffusion representation across alternating forward and reverse updates.
  • Continuous-time IPF: Each odd IPF step minimizes KL divergence subject to the terminal prior marginal, while each even step minimizes it subject to the initial data marginal.The alternating constraints are imposed on path measures over the full trajectory space.
  • Potential formulation: The time-reversal formulation is connected to the continuous-time potential formulation through an explicit identity between potentials and half-bridge densities.The paper notes that each iteration can also be viewed through a Doob h-transform.
  • Likelihood computation: Schrödinger-bridge likelihoods can be computed with the instantaneous change-of-variable formula and a Skilling–Hutchinson trace estimator.In practice, the associated dynamics are discretized and solved as forward or backward ODEs using the learned networks.

I Training Techniques

DSB training uses simulated or cached diffusion trajectories to optimize score-based losses, with specialized shortcuts and refinements aimed at reducing computation and training time.

  • Trajectory-based training: DSB trains score networks by minimizing empirical losses computed from simulated diffusion trajectories.A single uniformly sampled time-step per sample is used for each gradient step because trajectory points are correlated.
  • Trajectory-based training: Closed-form Gaussian sampling accelerates the first DSB iteration when the drift is linear.For α(x) = −αx, trajectory points can be sampled directly from a Gaussian with appropriate mean and covariance.
  • Trajectory-based training: Cached trajectory training stores and resamples diffusion trajectories after the initial iterations to avoid repeatedly simulating expensive full paths.The cache can be refreshed periodically while reused trajectories provide samples for the loss computation.
  • Performance refinements: The implementation recommends a Gaussian prior variance slightly larger than the target-data variance and matching prior and data means.The recommendation is motivated by the observed effect of prior parameters on IPF convergence.
  • Performance refinements: Network refinement initializes iteration n from iteration n−1, reducing the cost of retraining large networks from scratch.The authors identify fine-tuning previous network parameters as a way to reduce training time.
  • Performance refinements: An exponential moving average of network parameters with rate 0.999 improved performance.This refinement is reported as analogous to a technique used in prior score-based generative modeling work.

J.1 Two-dimensional experiments

Two-dimensional experiments examine DSB architectures, convergence, prior-variance effects, and the trade-off between few Langevin steps and stable iterative refinement.

  • Architecture: The two-dimensional networks use MLP blocks with positional encoding and 26,498 total parameters.The architecture includes blocks with shapes (2, 16, 32), (1, 16, 32), and (64, 128, 128, 2).
  • Experimental setup: Two-dimensional experiments use γk = 10^-2, batch size 512, a cache of 10^4 samples, and 10^4 training iterations per DSB step.The cache is refreshed every 10^3 iterations, and the prior mean and variance are matched to the data distribution.
  • Convergence: A prior variance larger than the target-data variance is recommended because prior variance affects DSB convergence speed.The authors report that this choice follows the practical guidance discussed in Technique 4.
  • Convergence: Using N = 1 is undesirable because the terminal marginal poorly approximates the prior and experiments require substantially larger networks to match N = 10.The reported low-N failure is observed even after many IPF iterations.
  • Sample behavior: DSB produces more diverse samples when the prior variance exceeds the data variance, even after few iterations.This observation links prior-scale selection to sample diversity in the two-dimensional experiments.

J.2 Generative Modeling

High-dimensional experiments apply DSB to MNIST and CelebA, with implementation choices for noise schedules, caching, embeddings, and latent-space exploration.

  • Generative modeling: The noise schedule is symmetric under time reversal, with γ0 = 10^-5 and midpoint value γ̄ = 10^-1.The schedule is defined for even N and mirrored for the second half of the sequence.
  • Generative modeling: DSB is evaluated on MNIST with dimension 784 and CelebA with dimension 3072 using Gaussian priors.The experiments use cached trajectories and Adam optimization, with dataset-specific batch sizes and diffusion-step settings.
  • Limitations: The low-N experiment reports significant deterioration after five iterations when N = 2.Each DSB iteration in this experiment uses 30,000 training steps.
  • Latent-space analysis: DSB supports deterministic embeddings by fixing the Gaussian random variables used during sampling.The paper also notes that a neural ordinary differential equation can provide a deterministic embedding with exact likelihood computation.
  • Latent-space analysis: Latent-space experiments include interpolation, temperature scaling, and exploration using an Ornstein–Uhlenbeck process targeting the prior.The exploration generates samples at several times before applying the DSB generative model.
  • Generative modeling: MNIST experiments obtain high-quality samples using N = 30 time steps and 15 × 10^3 training iterations.The baseline configuration instead uses N = 10 and 5 × 10^3 training iterations.

J.3 Dataset interpolation

Dataset interpolation experiments use DSB with modified diffusion settings and examine transitions between handwritten-image datasets and two-dimensional distributions.

  • Experimental setup: Dataset interpolation increases Langevin steps to 50 for two-dimensional examples and 30 for EMNIST/MNIST interpolation.The reference dynamics are changed to those obtained with DSB using a Gaussian prior to speed training.
  • EMNIST/MNIST: EMNIST is reduced to five letter classes so it matches MNIST’s number of classes for translation experiments.Uppercase and lowercase letters are distinguished in the reduced EMNIST dataset.
  • Two-dimensional interpolation: Figure 17 shows iteration 10 of IPF over T = 1.5 using 30 diffusion steps at times 0, 0.4, 1.25, and 1.5.The panels are ordered from the initial time to the terminal time.
  • Two-dimensional interpolation: The two-dimensional experiments include dataset interpolation across several classical datasets.Additional interpolation results are presented separately from the EMNIST/MNIST translation setting.
  • Two-dimensional interpolation: Figure 18 presents dataset interpolation at t = 0, 0.15, 0.30, and 0.5 using DSB iteration 9.The panels show the interpolation sequence at four specified times.
Loading 2106.01357v5…