Source-linked AI summary

Deep Generative Learning via Schrödinger Bridge

Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, Can Yang

arXiv:2106.10410v2cs.LGcs.CV

TL;DR

Generative modeling needs theoretically justified ways to transport a simple reference distribution to a target without restrictive assumptions. The paper uses finite-time Schrödinger Bridge interpolation with estimated SDE drifts and a two-stage Euler–Maruyama algorithm, proving convergence to the target and reporting competitive benchmark performance.

  • Problem

    Existing generative-model consistency results can require restrictive assumptions such as target log-concavity or a log-Sobolev inequality.

  • Method

    The paper learns generative models by interpolating reference and target distributions with a Schrödinger Bridge, using a two-stage drift-estimation and Euler–Maruyama procedure.

  • Results

    The learned distribution converges to the target under mild smoothness assumptions, while experiments show multimodal stability and performance comparable with state-of-the-art GANs.

  • Takeaways & Limitations

    Schrödinger Bridges provide a new formulation of generative modeling with theoretically guaranteed target-distribution learning and competitive empirical performance.

  • Takeaways & Limitations

    The theoretical guarantees assume smoothness of the target distribution, while existing comparison methods may require stronger assumptions such as log-concavity or a log-Sobolev inequality.

Abstract

from arXiv · show

We propose to learn a generative model via entropy interpolation with a Schrödinger Bridge. The generative learning task can be formulated as interpolating between a reference distribution and a target distribution based on the Kullback-Leibler divergence. At the population level, this entropy interpolation is characterized via an SDE on $[0,1]$ with a time-varying drift term. At the sample level, we derive our Schrödinger Bridge algorithm by plugging the drift term estimated by a deep score estimator and a deep density ratio estimator into the Euler-Maruyama method. Under some mild smoothness assumptions of the target distribution, we prove the consistency of both the score estimator and the density ratio estimator, and then establish the consistency of the proposed Schrödinger Bridge approach. Our theoretical results guarantee that the distribution learned by our approach converges to the target distribution. Experimental results on multimodal synthetic data and benchmark data support our theoretical findings and indicate that the generative model via Schrödinger Bridge is comparable with state-of-the-art GANs, suggesting a new formulation of generative learning. We demonstrate its usefulness in image interpolation and image inpainting.

1. Introduction

The paper proposes Schrödinger Bridges as a generative-learning formulation that interpolates a reference distribution to a target over finite time, addressing theoretical limitations of existing methods. Its two-stage algorithm has consistency guarantees and performs competitively with state-of-the-art GANs.

  • Motivation: Existing generative models have a gap between empirical success and theoretical justification, with some consistency results requiring restrictive target-distribution assumptions.In particular, certain SDE-based approaches require log-concavity or a log-Sobolev inequality.
  • Approach: The paper proposes Schrödinger Bridge generative learning, formulating interpolation between reference and target distributions through KL divergence and a finite-time SDE with time-varying drift.The sample algorithm estimates the drift and applies Euler–Maruyama.
  • Theory: The theory proves consistency under mild smoothness assumptions, guaranteeing that the learned distribution converges to the target.This contrasts with existing theories that rely on stronger assumptions such as log-concavity or log-Sobolev-type conditions.
  • Algorithm: The novel two-stage algorithm first learns a smoothed target distribution, then drives it to the underlying target distribution using two Schrödinger Bridges.Figure 1 presents this progression from initial samples through a smoothed distribution to the target.
  • Experiments: Experiments show stable learning of multimodal distributions and performance comparable with state-of-the-art GANs on benchmark data.The paper positions this result as support for Schrödinger Bridges as a new generative-model formulation.

2. Background

The Schrödinger Bridge problem seeks a path-space law connecting two endpoint distributions while remaining close to Brownian motion in relative entropy. Its optimal time-varying drift transports the reference distribution to the target over the unit interval with minimum quadratic energy.

  • Schrödinger Bridge Problem: The Schrödinger Bridge problem finds the most likely evolution between endpoint distributions by minimizing relative entropy relative to a Brownian-diffusion prior.The solution is a probability law on a continuous path space with prescribed initial and terminal marginals.
  • Characterization: The optimal bridge law has a factorized form involving endpoint functions f* and g*, which characterize the solution under the Schrödinger system.The time-marginal density is represented through corresponding forward and backward factors.
  • Dynamics: The resulting stochastic control solution has minimum energy under a quadratic cost among admissible Markov controls.The dynamic formulation connects the bridge drift to a finite-energy stochastic-control problem.
  • Dynamics: The bridge dynamics are represented by an SDE with a time-varying drift that evolves particles from the initial distribution μ to the target distribution ν over [0,1].This finite-time transport property motivates its use for generative learning.

3. Generative Learning via Schr¨odinger Bridge

The approach learns generative distributions by first smoothing the target and then using two finite-time Schrödinger Bridge SDEs with estimated, time-varying drifts. Its estimators and resulting algorithm are theoretically consistent under mild regularity assumptions.

  • Generative Learning via Schrödinger Bridge: The method first learns a smoothed target distribution, then transports it to the original target through two SDEs over the unit interval.Gaussian noise controls the smoothing level, while the two-stage construction starts from a simple distribution and proceeds through the smoothed target.
  • Estimation of the drift terms: The drift terms are estimated with a deep density-ratio estimator and a deep score estimator before Euler-Maruyama simulates both stages.The density ratio is learned using logistic regression, while denoising score matching estimates the score across noise levels.
  • Consistency of estimators: The deep density-ratio and score estimators are consistent as their sample sizes increase under the stated network and smoothness conditions.The score-estimation error in the stated L2(q̃σ) sense converges to zero as m and n tend to infinity.
  • Estimation of the drift terms: The two-stage algorithm runs Euler-Maruyama first from zero toward the smoothed distribution and then from that distribution toward the target.The first stage uses the estimated density ratio and score; the second uses an estimated time-varying score drift.
  • Consistency of Schrödinger Bridge Algorithm: The second stage recovers the reverse-time Variance Exploding SDE algorithm when its annealing schedule is linear.This identifies the Schrödinger Bridge construction with a particular annealing score-based sampling procedure.

4. Related Work

The paper contrasts Schrödinger Bridge with neural transform-map generative models and particle or SDE methods. Its distinguishing feature is a time-varying drift that reaches the target on the finite interval [0,1].

  • Related Work: Schrödinger Bridge differs from VAEs, GANs, and flow-based methods, which parameterize transform maps with neural networks and optimize an integral probability metric.
  • Related Work: Unlike related Langevin and particle methods with time-invariant drifts or asymptotic evolution, Schrödinger Bridge uses a time-varying drift to recover the target at terminal time.The cited comparison emphasizes the finite unit-time interval [0,1] as the distinguishing property.
  • Related Work: Reverse-time VE SDE methods retain a smoothed end-time distribution and theoretically require a smoothed initialization, whereas Schrödinger Bridge uses a first stage to learn that smoothed distribution.

5. Experiments

Experiments evaluate the two-stage Schrödinger Bridge on multimodal synthetic data, CIFAR-10, and CelebA. The method learns multimodal distributions, produces competitive CIFAR-10 scores, and supports denoising, interpolation, and inpainting.

  • 5.2. Learning 2D Multimodal Distributions: The method generates multimodal samples without missing modes, unlike vanilla GAN in the six-component Gaussian-mixture experiment.Estimated drift fields closely match ground truth except where probability density is nearly zero.
  • 5.3. Effectiveness of Two Stages for Image Generation: CIFAR-10 particle trajectories show a valid path from the initial distribution to the target distribution across the two stages.The center column marks particles obtained after stage 1.
  • 5.3. Effectiveness of Two Stages for Image Generation: Stage 1 provides a better initial reference for stage 2, while stage 2 denoises perturbed images toward the data distribution and retains some diversity.Using only stage 2 from non-informative Gaussian noise performs worse than the two-stage approach.
  • 5.4. Results: Performance is stable for prior Wiener variance τ ≥ 2.0, while reasonable τ choices generally yield relatively good generation results.The experiments vary τ from 1.0 to 4.0 on CIFAR-10.
  • 5.4. Results: FID 12.32 and Inception Score 8.14 on CIFAR-10 are competitive with WGAN-GP, SN-SMMDGAN, SNGAN, NCSN, and NCSNv2.Table 2 reports evaluation using 50,000 generated samples.
  • 5.5. Image Interpolation and Inpainting with Stage 2: On CelebA, stage 2 produces smooth facial-attribute interpolation and is adapted to generate image-inpainting samples from masked images.The interpolation experiment uses noise variance σ^2 = 0.4; the inpainting procedure perturbs observed pixels before stage-2 sampling.

6. Conclusion

The paper proposes Schrödinger Bridge generative learning through entropy interpolation and establishes consistency under smoothness assumptions, with learned distributions converging to the target.

  • The proposed Schrödinger Bridge learns generative models by interpolating a reference distribution and target distribution through Kullback-Leibler-based entropy interpolation.The population formulation uses an SDE on [0, 1] with a time-varying drift term.
  • A two-stage algorithm estimates the drift with deep score and density ratio estimators before applying Euler-Maruyama discretization.
  • Under smoothness assumptions, the score estimator, density ratio estimator, and resulting Schrödinger Bridge approach are consistent, so the learned distribution converges to the target.
  • The theoretical results are supported by synthetic and benchmark experiments, where the approach is comparable with state-of-the-art GANs.

A.3. Proof of Theorem 9

The proof characterizes Schrödinger Bridge dynamics through density ratios and shows that the resulting SDEs transport the initial distributions to their intended marginals.

  • Defining f(x) as the density ratio qσ(x) / Φ√τ(x) yields an SDE whose solution at t = 1 follows qσ(x).
  • The Schrödinger system provides the functions needed to represent the SDE construction under the Brownian reference measure.
  • The first-stage dynamics transport the point-mass initial distribution δ0 to qσ from t = 0 to t = 1.
  • With τ = σ2, the second-stage dynamics transport qσ to the target distribution pdata over the same interval.

A.4. Proof of Theorem 10

Theorem 10 establishes consistency of the neural density-ratio estimator using approximation and empirical-process bounds under compact-support and regularity assumptions.

  • Under compact support and bounded Lipschitz density-ratio assumptions, the neural estimator is analyzed through approximation and statistical errors.
  • The estimator minimizes an empirical logistic loss comparing samples from qσ(x) with samples from Φ√τ(x).
  • The network-size construction uses logarithmic depth and specified width and size growth to obtain the required approximation and estimation bounds.
  • The proof controls the empirical-process term using Lipschitzness, Rademacher variables, Gaussian complexity, covering numbers, and ReLU-network pseudo-dimension.

A.5. Proof of Theorem 11

Theorem 11 proves score-estimator consistency and uses it with density-ratio consistency and Euler-Maruyama error control to establish consistency of the two-stage algorithm.

  • Under differentiability, bounded support, and Lipschitz-bounded score assumptions, the neural score estimator is consistent as sample sizes increase.
  • The proof decomposes score-learning error into approximation and statistical components and bounds their empirical-process contribution.
  • The stated statistical error bound decreases as O(n^−2/(d+2)).
  • The first-stage consistency proof combines uniform estimator convergence, Monte Carlo drift approximation, and Euler-Maruyama discretization control.
  • The second-stage consistency follows by repeating the first-stage argument while using the established first-stage consistency results.

C. Network Architectures

The paper uses fully connected networks for the 2D example and image-specific architectures for its estimators. For images, the density ratio estimator uses a residual network, while the score estimator follows a U-Net based on a Wide ResNet.

  • C. Network Architectures: For the 2D example, the score and density ratio estimators use fully connected networks.
  • C. Network Architectures: The 2D score-estimator architecture includes a sinusoidal embedding of time t.
  • C. Network Architectures: For image datasets, the density ratio estimator is parameterized with a residual network, while the score estimator follows a U-Net based on a Wide ResNet.

D. More Implementation Details

The implementation centers data differently for density-ratio and score estimation, adjusts score inputs and signs, and adds a final denoising step for image generation. Experiments use CIFAR-10 and CelebA on one Tesla V100 GPU.

  • D. More Implementation Details: Density-ratio training centers pdata samples by subtracting an estimated image mean, whereas score training rescales samples to [−0.5, 0.5].
  • D. More Implementation Details: Score-estimator inputs are shifted by adding the estimated mean minus 0.5, with the output sign adjusted accordingly.
  • D. More Implementation Details: Image generation adds one final denoising step by repeating the last step without injecting noise.This step is intended to eliminate small noises in generated samples.
  • D. More Implementation Details: Experiments on CIFAR-10 and CelebA use one Tesla V100 GPU.

E. Additional Experiment Results

Additional CIFAR-10 experiments vary the noise parameter σ2 and compare different τ values. The reported metrics are FID and Inception Score across three result tables.

  • E. Additional Experiment Results: Additional CIFAR-10 results report FID and Inception Score for σ2 values of 0.5, 2.0, and 5.0.The σ2 = 1.0 results were already presented in the paper.
  • E. Additional Experiment Results: The additional experiments compare different τ values starting from τmin = σ2.
Loading 2106.10410v2…