Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Soft-constrained Schrodinger Bridge: a Stochastic Control Approach

Jhanvi Garg, Xianyang Zhang, Quan Zhou

TL;DR

This work introduces the soft-constrained Schrödinger bridge (SSB), which replaces the hard terminal constraint of the classical Schrödinger bridge with a penalized KL term controlled by $\beta$, enabling robust diffusion-based sampling when the target is uncertain or data-limited. The authors derive the optimal drift as $u_t^* = \sigma^2 \nabla \log h(X_t^{u^*}, t)$ with $h$ defined via conditional expectations, and show that the optimal terminal law is a geometric mixture between the target $\mu_T$ and the uncontrolled terminal distribution. They extend the framework to time-series targets, provide an existence/uniqueness theory via a generalized Schrödinger system, and develop a score-matching with importance sampling algorithm to learn geometric mixtures. Experiments on MNIST illustrate how moderate values of $\beta$ yield high-quality samples by leveraging large, high-fidelity reference data while maintaining fidelity to noisy targets. Overall, SSB offers a flexible, theoretically grounded route to robust generative diffusion and time-series data synthesis with controllable regularization through $\beta$.

Abstract

Schrödinger bridge can be viewed as a continuous-time stochastic control problem where the goal is to find an optimally controlled diffusion process whose terminal distribution coincides with a pre-specified target distribution. We propose to generalize this problem by allowing the terminal distribution to differ from the target but penalizing the Kullback-Leibler divergence between the two distributions. We call this new control problem soft-constrained Schrödinger bridge (SSB). The main contribution of this work is a theoretical derivation of the solution to SSB, which shows that the terminal distribution of the optimally controlled process is a geometric mixture of the target and some other distribution. This result is further extended to a time series setting. One application is the development of robust generative diffusion models. We propose a score matching-based algorithm for sampling from geometric mixtures and showcase its use via a numerical example for the MNIST data set.

Soft-constrained Schrodinger Bridge: a Stochastic Control Approach

TL;DR

This work introduces the soft-constrained Schrödinger bridge (SSB), which replaces the hard terminal constraint of the classical Schrödinger bridge with a penalized KL term controlled by , enabling robust diffusion-based sampling when the target is uncertain or data-limited. The authors derive the optimal drift as with defined via conditional expectations, and show that the optimal terminal law is a geometric mixture between the target and the uncontrolled terminal distribution. They extend the framework to time-series targets, provide an existence/uniqueness theory via a generalized Schrödinger system, and develop a score-matching with importance sampling algorithm to learn geometric mixtures. Experiments on MNIST illustrate how moderate values of yield high-quality samples by leveraging large, high-fidelity reference data while maintaining fidelity to noisy targets. Overall, SSB offers a flexible, theoretically grounded route to robust generative diffusion and time-series data synthesis with controllable regularization through .

Abstract

Schrödinger bridge can be viewed as a continuous-time stochastic control problem where the goal is to find an optimally controlled diffusion process whose terminal distribution coincides with a pre-specified target distribution. We propose to generalize this problem by allowing the terminal distribution to differ from the target but penalizing the Kullback-Leibler divergence between the two distributions. We call this new control problem soft-constrained Schrödinger bridge (SSB). The main contribution of this work is a theoretical derivation of the solution to SSB, which shows that the terminal distribution of the optimally controlled process is a geometric mixture of the target and some other distribution. This result is further extended to a time series setting. One application is the development of robust generative diffusion models. We propose a score matching-based algorithm for sampling from geometric mixtures and showcase its use via a numerical example for the MNIST data set.
Paper Structure (22 sections, 12 theorems, 127 equations, 8 figures, 3 tables)

This paper contains 22 sections, 12 theorems, 127 equations, 8 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Schrödinger Bridge and Its Applications
  3. Overview of This Work
  4. Related Literature
  5. Problem Formulation
  6. Solution to Soft-constrained Schrödinger Bridge
  7. Extension to Time Series
  8. Experiments
  9. Problem Setup
  10. A Score Matching Algorithm for Learning Geometric Mixtures
  11. An Example Using MNIST
  12. Concluding Remarks
  13. Acknowledgements
  14. Simulation with Known Density Functions
  15. Monte Carlo Simulation with Densities Known up to a Normalization Constant
  16. ...and 7 more sections

Key Result

Theorem 1

Let $\mu_0$ be the Dirac measure such that $\mu_0(\{x_0\}) = 1$ for some $x_0 \in {\mathbb R}^d$, and let $X$ be a weak solution to eq:udp. Assume $\mathcal{D}_{\mathrm{KL}}(\mu_T, \mathcal{L}aw(X_T)) < \infty$. For Problem problem0, the optimal control is given by $u^*_t = \sigma^2 \nabla \log h(X_ Moreover, $J(u^*) = \mathcal{D}_{\mathrm{KL}}(\mu_T, \mathcal{L}aw(X_T))$.

Figures (8)

  • Figure 1: SSB Samples for MNIST Experiment
  • Figure A1: Densities of Geometric Mixtures of Normal Cauchy Distributions.
  • Figure A2: Q-Q plots for the SSB Samples with $N_{\rm{mc}} = 1,000$.
  • Figure A3: SSB Trajectories for Normal Mixture Targets.
  • Figure F4: Samples in $\mathcal{D}_{\mathrm{obj}}$.
  • ...and 3 more figures

Theorems & Definitions (40)

  • Remark 1
  • Theorem 1: dai1991stochastic
  • Remark 2
  • Theorem 2
  • proof
  • Remark 3
  • Example 1
  • Lemma 3
  • proof
  • Remark 4
  • ...and 30 more