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

Abstract: Schr\"odinger 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.

• The paper presents a novel FBSDEs framework that converts Schrödinger Bridge optimality into log-likelihood objectives for deep generative models.
• It demonstrates that the SB-FBSDE approach attains a competitive 2.98 bits/dim NLL and a 3.18 FID score on datasets like CIFAR10.
• The study implies that integrating SGM techniques such as Langevin sampling into SB models enhances training dynamics and sample quality.

Likelihood Training of Schrödinger Bridge Using Forward-Backward SDEs Theory: An Analysis

The paper investigates the use of Schrödinger Bridge (SB) for deep generative modeling. SB presents an alternative to Score-based Generative Models (SGMs), focusing on an entropy-regularized optimal transport problem. While SGMs require pre-determined data-to-noise diffusion, SBs inherently learn the process, offering more flexibility. The study probes into whether SB's optimization principles align with deep generative model training, particularly concerning the construction of log-likelihood objectives.

The researchers present a novel computational framework leveraging Forward-Backward Stochastic Differential Equations (FBSDEs) theory. This approach translates the SB optimality conditions into a set of SDEs, forming a bridge to log-likelihood objectives synonymous with modern generative models. The result is a new optimization strategy retaining SB's optimality while incorporating SGM techniques, achieving effective image generation on datasets like MNIST, CelebA, and CIFAR10. The proposed method, SB-FBSDE, demonstrates that SB models can inherit generative techniques of SGMs while potentially improving training dynamics and sample quality.

Key Contributions

1. Novel Computational Framework: The paper introduces a framework grounded on FBSDEs theory, exposing theoretical connections between SB and SGM concerning log-likelihood training.
2. Generative Training Enhancements: The framework facilitates combining SB's flexibility and SGM-inspired correction techniques like Langevin sampling, promising improvements in training outcomes.
3. Empirical Results: On image generation tasks, the SB-FBSDE method shows superior sample quality over prior optimal transport models, proving to be on par with leading generative model classes.

Methodology and Results

The methodology centers around solving SB optimality through a set of coupled SDEs expressed via FBSDEs. The nonlinear Feynman-Kac formula is employed, allowing the translation of PDE optimality conditions into stochastic processes. The paper articulates how the likelihood objective can be calculated through FBSDEs, aligning with continuous-time and highlighting remarkable scalability in higher dimensions.

SB-FBSDE is trained using a divergence-based loss function, differing from previous approaches relying on regression to mean-matching. Empirical studies on CIFAR10 reveal that SB-FBSDE achieves a competitive negative log-likelihood (NLL) of 2.98 bits/dim and a Fréchet Inception Distance (FID) score of 3.18, positioning itself strongly against both traditional and contemporary models.

Implications and Future Directions

This work provides a robust alternative framework for generative modeling, especially in scenarios where pre-specifying diffusion processes is non-ideal. The ability to compute likelihoods under SB models expands their feasibility in applications requiring precise probabilistic interpretations. The theoretical connections established between SB and SGM further affirm SB's relevance in modern generative training practices.

Prospective research could explore optimizing FBSDEs implementation in tandem with larger-scale models, focusing on reducing computational overheads while maximizing convergence rates. Additionally, the impact of adaptive Langevin correction techniques could be further analyzed to enhance generative quality.

In conclusion, the study advances SB models as sophisticated counterparts to well-established generative frameworks, his work opens avenues for deeper exploration of entropy-regularized optimal transport's role in AI developments.