PaperlayerView exact revision on arXiv

Source-linked AI summary

Solving Schrödinger Bridges via Maximum Likelihood

Francisco Vargas, Pierre Thodoroff, Neil D. Lawrence, Austen Lamacraft

arXiv:2106.02081v9stat.MLcs.LG

TL;DR

The paper addresses scalable estimation of Schrödinger bridges under general diffusion priors. It reformulates the problem through maximum likelihood and proposes IPML with Gaussian-process drift estimation; the method converges theoretically and is demonstrated on synthetic and biological data.

  • Problem

    Scalable numerical methods for estimating Schrödinger bridges, especially with general diffusion priors and high-dimensional data, remain limited.

  • Method

    IPML recasts dynamic IPFP iterations as regression-based maximum likelihood and solves the resulting drift-estimation problems using Gaussian processes.

  • Results

    IPML converges in probability at each iteration and produces physically interpretable bridge dynamics, with experiments spanning synthetic and biological data.

  • Takeaways & Limitations

    The approach supports combining prior knowledge about temporal evolution with data-driven methods for Schrödinger bridge estimation.

  • Takeaways & Limitations

    Gaussian processes may underperform on high-dimensional data, where a parametric estimator such as a neural network may be useful.

Abstract

from arXiv · show

The Schrödinger bridge problem (SBP) finds the most likely stochastic evolution between two probability distributions given a prior stochastic evolution. As well as applications in the natural sciences, problems of this kind have important applications in machine learning such as dataset alignment and hypothesis testing. Whilst the theory behind this problem is relatively mature, scalable numerical recipes to estimate the Schrödinger bridge remain an active area of research. We prove an equivalence between the SBP and maximum likelihood estimation enabling direct application of successful machine learning techniques. We propose a numerical procedure to estimate SBPs using Gaussian process and demonstrate the practical usage of our approach in numerical simulations and experiments.

1 Introduction

The paper generalizes the Schrödinger bridge problem to arbitrary diffusion priors and introduces IPML, a regression-based maximum-likelihood procedure implemented with Gaussian processes.

  • Motivation: The dynamic Schrödinger bridge seeks the most likely stochastic process evolving π0 to π1 under a specified diffusion prior.The generalized formulation can incorporate domain knowledge about temporal evolution.
  • Motivation: Existing Sinkhorn-Knopp approaches require complex adaptations for general diffusion priors and marginals.The paper identifies scalability challenges for related formulations in high-dimensional spaces.
  • Contributions: The algorithm is compared conceptually and empirically with density-estimation and optimal-transport approaches, including experiments on real-world embryoid-cell data.The paper reports both quantitative and qualitative comparisons with state-of-the-art deep-learning and optimal-transport methods.
  • Contributions: IPML recasts dynamic IPFP iterations as a regression-based maximum-likelihood objective, enabling scalable machine-learning regression methods.The formulation can use Gaussian processes or neural networks.
  • Contributions: Gaussian processes solve the regression objectives while incorporating functional prior information through their model specification.The paper motivates this choice through the connection between stochastic differential-equation drifts and Gaussian processes.

2 Technical Background

The technical formulation represents the bridge with forward and reversed diffusion processes, converts terminal constraints into initial constraints, and alternates KL minimizations through IPFP until convergence.

  • Stochastic Control Formulation: The solution reformulates the Schrödinger bridge as a stochastic differential equation with initial-value and final-value boundary constraints.The full boundary-value problem is approached through separate constrained diffusion processes.
  • Diffusion Prior: The prior is a drift-augmented Brownian diffusion with constant scalar volatility γ, generalized beyond the traditional Wiener-measure setting.The bridge must use the same diffusion coefficient for finite KL divergence.
  • Time Reversal: Forward and reversed drifts are related through time reversal, with the dual drift determined using the associated Fokker–Planck solution.The reversed process starts from the original process’s terminal state.
  • Drift Parametrization: The optimal drifts admit deterministic function parametrizations, providing the basis for estimating them with machine-learning models.The construction applies to the half-bridge objectives and their time-reversed forms.
  • Iterative Proportional Fitting: IPFP alternates forward and backward formulations, each enforcing one initial-value constraint, so the iterates converge toward the full boundary-value problem.The algorithm alternates KL minimization between marginal constraints.

3 Methodology

The methodology reformulates half-bridge solving as maximum-likelihood drift estimation, using time reversal to make terminal constraints easy to sample. Iterative proportional fitting combines these estimates into a general Schrödinger bridge solver, implemented with Gaussian processes.

  • 3.1 Approximate Half Bridge Solving as Optimal Drift Estimation: Half-bridge measures can be sampled by simulating their drift-parametrized stochastic differential equations with Euler–Maruyama.The resulting SDE solutions are distributed according to P∗− and P∗+.
  • 3.1 Approximate Half Bridge Solving as Optimal Drift Estimation: Theorem 1 establishes consistency of reverse maximum-likelihood formulations for estimating the half-bridge drift.
  • 3.2 On the Need for Time Reversal: Time reversal converts a terminal hitting condition into an initial-value problem that can be sampled without mesh-based PDE methods or additional density estimation.This avoids two high-dimensional difficulties associated with directly solving the forward terminal-value formulation.
  • 3. Methodology: The IPML routine fits dual drifts on sampled trajectories and generates new trajectories through stochastic differential equation solving.DriftFit can use any function-estimation procedure with consistency guarantees, while SDESolve uses Euler–Maruyama.
  • 3. Methodology: Gaussian-process DriftFit separately estimates each decoupled drift dimension and can encode the prior drift as the GP mean function.Using the prior as the mean makes fitted drifts fall back toward the prior when evaluating unlikely or unobserved paths.

4 Related Methodology

Existing Schrödinger bridge methods often solve static formulations and require costly approximations for dynamics or high-dimensional evaluation. IPML instead compares regression-based maximum likelihood with density-estimation and Sinkhorn approaches.

  • Sinkhorn-Knopp: Sinkhorn-Knopp solves discretized static bridges, but its cost matrix requires density estimation and prior simulation.Evaluating optimal drifts additionally requires high-dimensional integrals, and many enhancements assume Brownian priors.
  • Data Driven Schrödinger Bridge: DDSB iterates coupled density-estimation objectives for static bridges, then requires further approximation to estimate the optimal drift.Its normalizing term involves a multidimensional integral not taken with respect to a probability distribution.
  • Dynamic versus static bridges: Static bridge approaches do not directly provide an estimate of the optimal dynamics needed for dynamic interpolation.The paper emphasizes dynamic approaches because discrete observations omit information between sampled time points.

5 Numerical Experiments

Experiments evaluate IPML on synthetic alignment, double-well dynamics, finite-sample convergence, embryoid-body data, and human motion while incorporating different stochastic priors. Across these settings, IPML fits marginals, follows lower-energy paths, improves drift-error convergence with GP priors, and often outperforms comparison methods.

  • 5.1 Simple 1D and 2D Distribution Alignment: IPML obtains better marginals than DDSB at γ = 1, whereas DDSB required γ = 100 for sensible results when distributions were distant.The comparison uses EMD and KS statistics on unimodal and bimodal Gaussian alignment experiments.
  • 5.1 Simple 1D and 2D Distribution Alignment: IPML diffuses a unimodal distribution into a multimodal one and produces forward and backward trajectories that are mirror images.The learned bridge successfully splits in the reported 2D experiment.
  • 5.2 2D Double Well Experiments: IPML with the potential prior follows lower-energy passes around the double-well barrier, unlike OT geodesics or a Brownian-prior bridge.The prior enters through both the first drift and the Gaussian-process mean function.
  • 5.2 2D Double Well Experiments: 22.1, 19.9, and 18.7 are the reported mean squared errors for DDSB, IPML without a prior, and IPML with the prior, respectively.The paper reports that IPML considerably outperforms DDSB on this evaluation.
  • 5.3 Finite sample/iteration convergence: GP drift priors significantly improve the error and its convergence, although Theorem 1 does not cover finite-sample and discretization effects.The finite-effects analysis is empirical.
  • 5.4 Single Cell - Embryo Body (EB) Data set: On embryoid-body scRNA-seq data, IPML outperforms TrajectoryNet in most frames and performs similarly to OT.Evaluation fits endpoint models and predicts intermediate frames using EMD.
  • 5.4 Single Cell - Embryo Body (EB) Data set: Finite-data DriftFit permits nonlinear trajectories that improve fit, with better single-cell observation coverage and a reported improvement over OT for frame 4.The paper suggests domain-specific priors could further improve performance relative to OT.
  • 5.5 Human Motion: IPML models basketball sensor motion in a 4-dimensional space using either Brownian or second-order Langevin priors.The experiment models position and velocity for the right shoulder and elbow.

6 Limitations and Opportunities

The paper identifies extensions for broader volatility models and multiple data frames, while concluding that IPML solves Schrödinger bridges for arbitrary diffusion priors and supports science applications combining prior knowledge with data.

  • 6 Limitations and Opportunities: Gaussian processes estimate the drift, while IPML can also use parametric estimators such as neural networks for high-dimensional data.GPs additionally incorporate functional priors through their mean function.
  • 6 Limitations and Opportunities: Nonconstant volatility functions remain an extension needed for applications such as positivity-constrained processes modeled with geometric Brownian motion priors.Developing this extension requires broadening Schrödinger bridge theory to nonconstant volatility functions.
  • 6 Limitations and Opportunities: Multiple frames could replace the usual two boundary conditions, because IPFP adapts to multiple constraints.The paper proposes future experiments using multiple frames and leave-one-out evaluation.
  • 6.1 Conclusion: IPML solves Schrödinger bridges for arbitrary diffusion priors and is demonstrated on synthetic and biological data.The approach combines prior knowledge about temporal evolution with data-driven methods for modern scientific problems.

A Assumptions Across Proofs

The appendix states shared regularity, growth, compactness, differentiability, and boundary-moment assumptions, then introduces the Schrödinger system and its potential propagation from the static bridge.

  • A Assumptions Across Proofs: The proofs assume all SDEs have Lipschitz diffusion coefficients and satisfy linear growth.These are listed as assumptions shared across the appendix results.
  • A Assumptions Across Proofs: Optimal drifts are assumed to lie in a compact space and therefore satisfy the Hamilton–Jacobi–Bellman equations.The appendix notes that Theorem 1 can relax this using Γ-convergence.
  • A Assumptions Across Proofs: The SDEs are assumed to satisfy a Fokker–Planck equation, making p(x, t) = Law(X_t) differentiable with respect to x.
  • A Assumptions Across Proofs: Boundary distributions are assumed to have bounded second moments in L2(π_i).Specifically, the appendix requires a finite constant C satisfying E_πi[|X|^2] ≤ C.
  • A Assumptions Across Proofs: The Schrödinger system arises by marginalizing the dynamic bridge to its boundary distributions, with potentials propagated across time.The appendix introduces the static bridge, its Lagrangian derivation, and the resulting Schrödinger potentials φ and ˆφ.

C Disintegration Theorem - Product Rule for Measures

The appendix presents disintegration as a measure-theoretic analogue of the product rule, enabling marginal and conditional decompositions and a Radon–Nikodym derivative factorization.

  • C Disintegration Theorem - Product Rule for Measures: Disintegration decomposes probability measures on product spaces into a marginal measure and conditional measures.The construction uses measurable projections and a conditional measure P_y|x(·|x).
  • C Disintegration Theorem - Product Rule for Measures: For probability measures, disintegration plays a role analogous to the ordinary product rule.The appendix derives this by applying the theorem to indicator functions.
  • C Disintegration Theorem - Product Rule for Measures: The Radon–Nikodym derivative between two probability measures on the same product space can be decomposed into conditional and marginal derivatives.The proof assumes the relevant conditional absolute continuity relation.
  • C Disintegration Theorem - Product Rule for Measures: Replacing the reference measure Q with Lebesgue measure recovers the standard product rule for probability density functions.

D Proof Sketches For Half Bridges

The half-bridge proofs use KL-divergence disintegration to match boundary constraints while retaining the reference conditional dynamics, and characterize sampling through drift-based SDEs.

  • D Proof Sketches For Half Bridges: The reference measure can be parameterized by its drift through either of two stochastic differential equations.
  • D Proof Sketches For Half Bridges: Sampling the half bridges consists of simulating trajectories from the corresponding SDEs.The sampled paths are distributed according to P*− and P*+ respectively.
  • D Proof Sketches For Half Bridges: The forward half bridge is obtained by decomposing KL divergence into a boundary marginal term and a conditional KL term.
  • D Proof Sketches For Half Bridges: Setting the conditional path law equal to the reference conditional law makes the conditional KL term vanish while enforcing the boundary constraint.The resulting minimum coincides with the half-bridge minimum.

E Proof Sketch for Reverse-MLE Consistency

The proof connects reverse-time maximum likelihood with KL minimization between diffusions, showing consistency under increasingly fine discretization and large samples. It also identifies a limitation: the proof does not yet extend to approximate Euler–Maruyama samples.

  • Proof strategy: The reverse-time likelihood converges to a KL-divergence optimization whose minimum occurs when estimated and true drifts agree.The argument uses the equivalence between matching stochastic-process measures and matching their drifts.
  • Consistency result: Theorem 1 establishes consistency of the reverse-MLE formulation under the stated assumptions.The proof derives pointwise convergence of the negative normalized log-likelihood and then relates its maximizer to the dual drift.
  • Time reversal: The forward and reverse half-bridge drifts are linked through time reversal, with the optimal forward drift expressed using the score of the intermediate density.The stated relation assumes differentiability, compactness of the optimal drift space, and regularity of the prior drift.
  • Limitations: The proof applies to discretized samples from the original SDE but has not been extended to approximate Euler–Maruyama sampling schemes.The authors suggest that convergence properties of the Euler scheme may support such an extension, but leave the required stochastic-integral argument for future work.

F Towards a Finite Sample Analysis of Approximate IPFP Schemes

The paper sketches a finite-sample analysis for approximate IPFP schemes by bounding projection errors and iteration errors under Lipschitz and estimation-error assumptions. The analysis is explicitly presented as a strategy rather than a formal proof.

  • Setup: Approximate IPFP schemes replace exact IPFP projection steps with inexact approximations such as IPML.The section develops a preliminary framework for analyzing their finite-sample and discretization errors.
  • Error bounds: The iteration error can be bounded from above under Lipschitz approximate projection operators and bounded finite-sample projection error.The error is measured in the L2(Q*) norm for the half-bridge drift representation.
  • Scope: The section provides an analysis strategy rather than a proof, leaving formal finite-sample convergence rates for future work.The authors state that the required conditions and their effects on convergence rates remain to be established.
  • Required assumptions: The proposed analysis requires assumptions on exact IPFP errors, projection Lipschitz constants, and finite-sample projection error.The finite-sample error term is difficult to derive for GP-based drift estimation and is therefore assumed to be bounded.
  • Operator assumptions: Non-expansive or contractive projection operators prevent cumulative errors from growing exponentially with iteration count.For contractive operators, the section states that the error bound can remain constant with respect to the iteration number.

G Estimates Required by the Sinkhorn-Knop Algorithm

The section contrasts Sinkhorn-Knopp-style methods with IPML for general SDE priors, emphasizing the computational difficulty of unavailable transition densities and static-to-dynamic conversion. It then describes kernel-dependent behavior in potential-well experiments.

  • Sinkhorn-Knopp estimates: General SDE priors make Sinkhorn-Knopp costs difficult to compute when transition densities lack closed forms.Estimating these costs requires prior-SDE sampling and density estimation for each dataset point.
  • Static-to-dynamic conversion: Sinkhorn-Knopp produces discrete potentials, leaving an additional problem of converting static potentials into time-dependent trajectories.This is a distinct challenge from estimating the transition-density costs.
  • IPML comparison: IPML instead formulates its integrals as expectations over empirical distributions and fitted SDEs, avoiding the non-probability outer integral in the compared density-estimation approach.The alternative importance-sampling strategy is described as poorly scaling to higher dimensions.
  • DDSB limitation: The DDSB method can fail when the endpoint distributions are distant, including for suitable γ values such as 1, 2, and 3.The cited experiment uses the same marginals as an earlier experiment and reports the resulting failure in Figure 8.
  • Kernel choice: In the potential-well experiment, kernel choice changes whether trajectories split across both low-energy passes or select one.The exponential kernel with lengthscale 0.25 produced splitting, whereas the EQ kernel selected one pass; the authors relate this to smoothness preferences.

I.4 Motion experiment

The motion experiment uses a second-order Langevin prior encoded as a first-order SDE and compares its computational demands with Brownian-prior settings. The implementation relies on Gaussian-process drift fitting, whose cost dominates runtime and memory use.

  • Prior model: The motion experiment represents a second-order Langevin prior in companion form as a first-order SDE.The setup uses a dampening factor β of 0 and γ = 0.5, with both decoupled and coupled spring matrices explored.
  • Prior configuration: The prior frequency is set near ω ≈ π to represent approximately 0.5 Hz and half a revolution for the selected movement.This choice reflects motion that starts at rest, reaches a high position, and does not return or repeat.
  • Baseline: The Brownian baseline uses γ = 0.3, selected visually to compensate for the inherently smooth Langevin trajectories.The comparison uses a less noisy Brownian prior while retaining visually good results.
  • Computational cost: IPML’s main computational cost is Gaussian-process drift fitting, with cubic time complexity in the number of discretized samples.The section reports quadratic prediction and memory costs, and notes that GP approximations can reduce fitting cost.
  • Runtime: Cell-experiment runtime is between 1 and 2 hours, compared with around 5 minutes for the well experiment.The cell experiment also has significant RAM consumption because of Gaussian-process fitting.
Loading 2106.02081v9…