PaperlayerView exact revision on arXiv

Source-linked AI summary

Likelihood Training of Schrödinger Bridge using Forward-Backward SDEs Theory

Tianrong Chen, Guan-Horng Liu, Evangelos A. Theodorou

arXiv:2110.11291v5stat.MLcs.LGmath.APmath.OC

TL;DR

SB offers flexible generative modeling, but its relationship to modern likelihood-based training remains unclear, while existing approaches face difficult PDEs and restrictive procedures. The paper uses Forward-Backward SDEs to construct SB likelihood objectives, yielding SB-FBSDE, which achieves comparable image-generation results and enables likelihood computation for SB.

  • Problem

    The relationship between SB’s optimization principle and modern likelihood-based generative training remains relatively unexplored, while SGM and prior SB approaches have restrictive or costly requirements.

  • Method

    The paper develops a Forward-Backward SDE framework that represents SB’s nonlinear PDE solutions and constructs log-likelihood objectives for SB.

  • Results

    SB-FBSDE achieves 2.98 bits/dim and a 3.18 FID score on CIFAR10, comparable to leading methods and outperforming prior optimal-transport methods.

  • Takeaways & Limitations

    The framework retains SB’s mathematical flexibility while supporting modern generative-training techniques and likelihood computations that prior optimal-transport methods could not feasibly provide.

  • Takeaways & Limitations

    SB-FBSDE typically requires 2–2.5 times more memory than SGM because it computes divergences and maintains two distinct networks.

Abstract

from arXiv · show

Schrödinger Bridge (SB) is an entropy-regularized optimal transport problem that has received increasing attention in deep generative modeling for its mathematical flexibility compared to the Scored-based Generative Model (SGM). However, it remains unclear whether the optimization principle of SB relates to the modern training of deep generative models, which often rely on constructing log-likelihood objectives.This raises questions on the suitability of SB models as a principled alternative for generative applications. In this work, we present a novel computational framework for likelihood training of SB models grounded on Forward-Backward Stochastic Differential Equations Theory - a mathematical methodology appeared in stochastic optimal control that transforms the optimality condition of SB into a set of SDEs. Crucially, these SDEs can be used to construct the likelihood objectives for SB that, surprisingly, generalizes the ones for SGM as special cases. This leads to a new optimization principle that inherits the same SB optimality yet without losing applications of modern generative training techniques, and we show that the resulting training algorithm achieves comparable results on generating realistic images on MNIST, CelebA, and CIFAR10. Our code is available at https://github.com/ghliu/SB-FBSDE.

1 INTRODUCTION

The paper examines whether Schrödinger Bridge optimization can support modern likelihood-based generative training while retaining SB's flexibility. It introduces a Forward-Backward SDE framework that connects SB optimality to SGM training and reports strong image-generation results.

  • Motivation: SGMs diffuse data to noise and learn the reverse process through score-function regression.The score is the gradient of the log probability density at each time step.
  • Motivation: SGMs face restrictions from simple diffusion forms and long diffusion horizons, which can make generation notoriously slow.The terminal distribution must become approximately Gaussian before generation can proceed effectively.
  • Motivation: Schrödinger Bridge formulates an entropy-regularized optimal transport problem that learns two policies transforming between arbitrary distributions over a finite horizon.This flexibility motivates SB-inspired alternatives to fixed SGM diffusion processes.
  • Research gap: The relationship between SB optimization and modern likelihood training remains underexplored, while prior SB approaches use multi-stage optimization or traditional IPF algorithms.The paper identifies this gap despite the related distribution-transforming formulations of SB and SGM.
  • Contribution: Forward-Backward SDEs connect SB's PDE optimality conditions to an exact log-likelihood expression that generalizes SGM to fully nonlinear diffusion.The framework also yields a training principle retaining SB flexibility while using modern generative-training techniques.
  • Contribution: SB-FBSDE outperforms prior optimal-transport-inspired baselines and is comparable to other existing models on high-fidelity image synthesis.Experiments cover MNIST, CelebA, and CIFAR10.

2 PRELIMINARIES

SGM diffuses data toward a prior and learns a reverse process through score regression, while SB optimizes flexible finite-horizon transformations between arbitrary endpoint distributions. Forward-Backward SDEs connect their stochastic-control formulations and provide a scalable route from SB optimality PDEs to likelihood training.

  • 2.1 SCORE-BASED GENERATIVE MODEL (SGM): SGM progressively diffuses data toward approximately Gaussian noise, then learns the reverse process by regressing a network to the score function.The forward diffusion uses drift f, scalar diffusion g(t), and a Wiener process; the learned score replaces the exact score during generation.
  • 2.1 SCORE-BASED GENERATIVE MODEL (SGM): SGM requires simple diffusion structures for analytic conditional scores and sufficiently long horizons for the terminal distribution to approximate the prior.The diffusion design also affects the training weighting and can lead to unstable training; the prior need not equal the actual terminal marginal.
  • 2.2 SCHRÖDINGER BRIDGE (SB): Schrödinger Bridge seeks optimal policies that transform between arbitrary endpoint distributions over a finite horizon, using coupled PDEs to characterize optimality.Its forward and backward SDE representations use drift corrections based on ∇x log Ψ and ∇x log bΨ, respectively.
  • 2.2 SCHRÖDINGER BRIDGE (SB): Both SGM and SB can be viewed as stochastic-control problems seeking time-varying policies under control-affine SDEs with additive noise.The SGM policy is the conditional score, whereas the SB policy is the backward drift correction; their forward and backward processes have equivalent marginals across time.
  • 2.2 SCHRÖDINGER BRIDGE (SB): SB's generative pipeline is difficult because obtaining its optimal drift requires solving coupled PDEs that are hard even for low-dimensional systems.Prior approaches therefore use multi-stage constructions, half-bridge optimization, or IPF; the framework here targets scalable computation of the optimality conditions.
  • 2.2 SCHRÖDINGER BRIDGE (SB): Forward-Backward SDEs theory links SB optimality to parameterized likelihood objectives, and the SB objective collapses to the SGM objective under a specified choice of auxiliary processes.SB additionally permits nonlinear drifts and can compensate when the designed SGM diffusion does not reach the prior exactly.

3 APPROACH

The approach uses Forward-Backward SDEs to represent Schrödinger Bridge PDE optimality conditions and derive likelihood objectives. This establishes connections to SGM and flow-based training while supporting flexible practical optimization.

  • Unified stochastic-control view: Both SGM and SB can be viewed as control-affine SDEs seeking time-varying control policies under forward and backward stochastic processes.The controls are the score function in SGM and ∇x log bΨ in SB.
  • FBSDE representation: Forward-Backward SDE theory represents nonlinear PDE solutions through coupled forward and backward SDEs, offering a computational alternative for high-dimensional optimal-control problems.The representation uses transformations between PDE solutions and backward-SDE variables, with conditional expectations recovering the PDE quantities.
  • SB likelihood computation: For SB, the backward SDE solutions at t = 0 provide an unbiased estimation of a data point's log-likelihood along paths sampled from the optimal forward SDE.The forward process is the optimal SB SDE after substituting Zt = g∇x log Ψ.
  • SB likelihood computation: Theorem 4 expresses SB log-likelihoods using the FBSDE system, yielding a parameterized lower bound trained with divergence-based objectives for two policy networks.The policies Zt and bZt are approximated by parameterized models such as DNNs.
  • Connection to SGM: SB likelihood training collapses to the SGM objective when the forward control is zero and the backward control becomes the score function.When the terminal SDE marginal equals pprior, no forward control is needed; otherwise SB uses forward and backward policies to steer and compensate the process.
  • Connections and implementation: The same SB likelihood objective can be recovered through flow-based training of an associated ordinary differential equation, while SB-FBSDE supports joint or alternate training procedures.Joint training can converge faster per iteration but has higher memory complexity.

4 EXPERIMENTS

Experiments validate SB-FBSDE on synthetic continuous and discontinuous distributions and on MNIST, resized CelebA, and CIFAR10. The method produces high-fidelity images, improves forward-diffusion alignment with the prior, and benefits from Langevin correction.

  • Toy datasets: SB-FBSDE learns nonlinear, dissimilar forward and backward vector fields that transport samples between prior and data distributions on continuous and discontinuous toy datasets.The forward field drives diffusion toward pprior, while the backward field drives generation toward pdata.
  • Image datasets: SB-FBSDE synthesizes high-fidelity images on MNIST, resized CelebA, and CIFAR10.Uncurated samples are presented for all three image datasets.
  • Image datasets: 2.98 bits/dim and 3.18 FID score on CIFAR10 are comparable to leading generative models and outperform prior optimal-transport methods in sample quality.The NLL is measured on the test set, while FID is measured against the training set.
  • Validity of SB forward diffusion: SB-FBSDE learns terminal distributions closer in KL divergence to pprior than SGM under both degenerate and linear base drifts.Unlike SGM, whose diffusion is fixed, SB-FBSDE progressively shortens the KL gap through its learned forward policy.
  • Effect of Langevin corrector: Adding Langevin correction uniformly improves FID scores on CelebA and CIFAR10 throughout training.The correction helps redistribute samples toward the desired density after discretization errors from Euler–Maruyama propagation.

5 CONCLUSION

The paper presents SB-FBSDE as a Forward-Backward SDE framework for computing Schrödinger Bridge log-likelihoods and validating the approach on image-generation tasks. It reports encouraging high-fidelity samples while retaining a rigorous mathematical framework.

  • SB-FBSDE computes Schrödinger Bridge log-likelihoods using a framework grounded in Forward-Backward SDE theory.
  • The framework generalizes previous theoretical results for Score-based Generative Models and supports modern generative training techniques for SB.
  • Experiments on MNIST, CelebA, and CIFAR10 show encouraging results in synthesizing high-fidelity samples.

REPRODUCIBILITY STATEMENT

The paper provides algorithmic and objective details for reproducing its training procedures, while placing implementation details and proofs in appendices.

  • Algorithms 1–3 specify the training procedures and objectives in the main text.The objectives are given in equations (16), (18), and (19).
  • Data preprocessing and other implementation details are provided in Appendix D.
  • Theoretical proofs are placed in Appendix B because of space constraints.

A INTRODUCTION OF SCHRÖDINGER BRIDGE

The Schrödinger Bridge is formulated as a stochastic optimal-control problem whose optimality conditions can be represented through coupled PDEs and transformed into forward-backward SDEs. This formulation supplies the theoretical bridge to score-based generative modeling.

  • The associated stochastic-control formulation minimizes control energy subject to a terminal boundary condition.
  • SB optimality is characterized by coupled PDEs with boundary conditions linking the two endpoint marginals.The relevant equations are Kolmogorov backward and forward PDEs, also called Hamilton–Jacobi–Bellman and Fokker–Planck equations.
  • A Hopf–Cole transformation converts the stochastic-control value function and optimal marginal density into the SB system.
  • A special SB SDE coincides with the SDE used in score-based generative models, recovering the corresponding result as a special case.
  • Forward-Backward SDE theory connects the SB PDE system to coupled SDEs, yielding optimal forward and backward processes.The backward process follows from reversing the optimal forward SDE and applying the SB factorization principle.

B PROOFS AND REMARKS IN SECTION 3

The section establishes FBSDE representations of SB optimality and derives log-likelihood and probability-flow formulations. These results connect coupled forward-backward SDEs with the SB PDE system and support likelihood-based training.

  • Assumptions: The paper states regularity assumptions ensuring SDE existence and uniqueness, Itô calculus, backward-SDE definitions, and exponential limiting behavior.These include smooth densities, continuous coefficients, Lipschitz and growth conditions, and regularity of the transformed value functions.
  • FBSDE representation: Nonlinear Feynman-Kac relations connect coupled forward-backward SDEs to the nonlinear PDEs underlying SB optimality.The backward variables are time-varying functions of the forward process, with Y_t and Z_t obtained from the value function.
  • Remarks: The classical-solution derivations can be extended to viscosity solutions when classical solutions do not exist.The extension is formulated through uniform limits of classical solutions on compact sets.
  • Log-likelihood: Theorem 4 expresses the SB model’s log-likelihood at a data point as an expectation over the forward SDE trajectory.The objective contains the policy norms, a divergence term, and a cross term involving the forward and backward policies.
  • Probability flow: A probability-flow ODE characterizes the optimal SB processes and recovers the same log-likelihood objective through flow-based change-of-variables calculations.The derivation uses the factorization Z_t + bZ_t = g∇_x log p_SB,t.
  • Symmetric formulation: The likelihood computation extends symmetrically to the terminal-state objective through a time-coordinate transformation that exchanges forward and backward SB policies.Theorem 11 derives the corresponding FBSDEs and log-likelihood expression for L_SB(x_T).

C COMPARISON WITH PRIOR SB WORKS

Compared with prior SB approaches, SB-FBSDE uses divergence-based likelihood objectives and supports direct continuous-time connections to SGM. It also changes corrector sampling and permits joint policy optimization.

  • Training objective: Prior SB methods use mean-matching regression or IPF, whereas SB-FBSDE uses divergence-based objectives derived for likelihood training.The divergence objective avoids directly approximating the generally intractable SB score ∇_x log p_SB,t.
  • SDE model class: SB-FBSDE connects to SGM directly in continuous time because both use the same SDE model class, avoiding dependence on numerical step sizes.The paper contrasts this with classical SB connections that emerge only after discretization.
  • Model parametrization: Unlike Vargas et al.’s non-parametric models, SB-FBSDE uses DNNs to approximate both SB policies, as does De Bortoli et al.Vargas et al.’s Gaussian-process models are therefore not directly comparable.
  • Training algorithm: SB-FBSDE can jointly optimize both policies, while concurrent methods emphasize alternating training through IPF.Joint optimization resembles diffusion flow-based training, whereas alternating optimization remains available when computational budgets are constrained.
  • Corrector sampling: SB-FBSDE’s Langevin corrector uses Z_t + bZ_t to estimate the optimal-density score, unlike prior sampling based on a single policy network.The paper reports major quantitative improvement of up to 4 FID for its corrector, while prior work reported negligible improvement.

D EXPERIMENT DETAILS

The experiments use dataset-specific architectures, preprocessing, training schedules, and Langevin-corrected sampling. The implementation includes both alternating and joint optimization procedures.

  • Data preprocessing: MNIST is padded to 32×32, CelebA is resized to 3×32×32, and CelebA and CIFAR10 use random horizontal flips.These preprocessing choices are intended to prevent degenerate feature maps, accelerate training, and increase diversity.
  • Generative process: Sampling starts from the prior and propagates the backward policy, then applies Langevin correction using the estimated score from Z(t, X_t) + bZ(t, X_t).The procedure uses a predefined signal-to-noise ratio r = 0.05 for all experiments.
  • Network architectures: The model uses fully connected swish blocks for toy data and a reduced U-net for MNIST, with dataset-specific architectures summarized in Table 4.The MNIST U-net uses one residual block, two attention heads, and 32 channels.
  • Evaluation conditions: Table 1 results omit the Langevin corrector because of computational constraints, while other experiments use it because it generally improves performance.The paper notes that CIFAR10 results may therefore improve further with the corrector.
  • Evaluation conditions: FID scores in the corrector experiment are computed with 10k samples, and improvement remains consistent across larger architectures and increased sample counts.The reduced U-net is used for the CelebA acceleration experiment.

E ADDITIONAL EXPERIMENTS

Additional experiments compare SB-FBSDE with prior SB results under matched settings and examine staged training on CIFAR10. The method is also illustrated with generated samples on MNIST, CelebA, and CIFAR10.

  • Comparison with prior work: The comparison with De Bortoli et al. uses the same reduced U-net, preprocessing, N=50 steps, and 0.5-second horizon.This setup is intended to provide a fair comparison of generated images.
  • Training stages: Figure 9 contrasts SGM-regression training, backward-policy refinement, and full SB-FBSDE training at successive stages.The stages correspond to 50k SGM iterations, 5k backward-policy updates, and full forward-backward SB training.
  • Staged training: CIFAR10 training first uses SGM regression for 50k iterations, then applies SB divergence-based training to refine the result.Table 5 labels these stages as “SGM/50k” and “SB/{f,b}/5k” according to the training loss and policy updates.
  • Staged training: 13.35 FID after 5k divergence-training iterations improves from 33.68, and full training lowers FID further to 11.85.The qualitative comparison reports substantially reduced visible flaws and noise after the 5k refinement stage.
  • Generated samples: Figures 10–12 show uncurated SB-FBSDE samples generated on MNIST, resized CelebA, and CIFAR10.These figures provide qualitative sample outputs for the three image datasets.
Loading 2110.11291v5…