72 research outputs found
Likelihood Training of Schr\"odinger Bridge using Forward-Backward SDEs Theory
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
Deep Generalized Schr\"odinger Bridge
Mean-Field Game (MFG) serves as a crucial mathematical framework in modeling
the collective behavior of individual agents interacting stochastically with a
large population. In this work, we aim at solving a challenging class of MFGs
in which the differentiability of these interacting preferences may not be
available to the solver, and the population is urged to converge exactly to
some desired distribution. These setups are, despite being well-motivated for
practical purposes, complicated enough to paralyze most (deep) numerical
solvers. Nevertheless, we show that Schr\"odinger Bridge - as an
entropy-regularized optimal transport model - can be generalized to accepting
mean-field structures, hence solving these MFGs. This is achieved via the
application of Forward-Backward Stochastic Differential Equations theory,
which, intriguingly, leads to a computational framework with a similar
structure to Temporal Difference learning. As such, it opens up novel
algorithmic connections to Deep Reinforcement Learning that we leverage to
facilitate practical training. We show that our proposed objective function
provides necessary and sufficient conditions to the mean-field problem. Our
method, named Deep Generalized Schr\"odinger Bridge (DeepGSB), not only
outperforms prior methods in solving classical population navigation MFGs, but
is also capable of solving 1000-dimensional opinion depolarization, setting a
new state-of-the-art numerical solver for high-dimensional MFGs. Our code will
be made available at https://github.com/ghliu/DeepGSB.Comment: NeurIPS 202
Generative Modeling with Phase Stochastic Bridges
Diffusion models (DMs) represent state-of-the-art generative models for
continuous inputs. DMs work by constructing a Stochastic Differential Equation
(SDE) in the input space (ie, position space), and using a neural network to
reverse it. In this work, we introduce a novel generative modeling framework
grounded in \textbf{phase space dynamics}, where a phase space is defined as
{an augmented space encompassing both position and velocity.} Leveraging
insights from Stochastic Optimal Control, we construct a path measure in the
phase space that enables efficient sampling. {In contrast to DMs, our framework
demonstrates the capability to generate realistic data points at an early stage
of dynamics propagation.} This early prediction sets the stage for efficient
data generation by leveraging additional velocity information along the
trajectory. On standard image generation benchmarks, our model yields favorable
performance over baselines in the regime of small Number of Function
Evaluations (NFEs). Furthermore, our approach rivals the performance of
diffusion models equipped with efficient sampling techniques, underscoring its
potential as a new tool generative modeling
Augmented Bridge Matching
Flow and bridge matching are a novel class of processes which encompass
diffusion models. One of the main aspect of their increased flexibility is that
these models can interpolate between arbitrary data distributions i.e. they
generalize beyond generative modeling and can be applied to learning stochastic
(and deterministic) processes of arbitrary transfer tasks between two given
distributions. In this paper, we highlight that while flow and bridge matching
processes preserve the information of the marginal distributions, they do
\emph{not} necessarily preserve the coupling information unless additional,
stronger optimality conditions are met. This can be problematic if one aims at
preserving the original empirical pairing. We show that a simple modification
of the matching process recovers this coupling by augmenting the velocity field
(or drift) with the information of the initial sample point. Doing so, we lose
the Markovian property of the process but preserve the coupling information
between distributions. We illustrate the efficiency of our augmentation in
learning mixture of image translation tasks
- …