Data-to-Energy Stochastic Dynamics

Kirill Tamogashev, Nikolay Malkin

Keywords: Schrödinger bridge, Bayesian posterior inference, stochastic differential equations, Iterative Proportional Fitting

TL;DR: The paper presents a novel algorithm for modelling data-to-energy Schrödinger bridges.

Abstract:

The Schrödinger bridge problem is concerned with finding a
stochastic dynamical system bridging two marginal distributions
that minimises a certain transportation cost.
This problem, which represents a generalisation of optimal
transport to the stochastic case, has received attention due to its
connections to diffusion models and flow matching, as well as its
applications in the natural sciences.
However, all existing algorithms enable the inference of such
dynamics only for cases where samples from both distributions are available.
In this paper, we propose the first general method for modelling
Schrödinger bridges when one (or both) distributions are given by
their unnormalised densities, with no access to data samples.
Our algorithm relies on a generalisation of the iterative
proportional fitting (IPF) procedure to the data-free case,
inspired by recent developments in off-policy reinforcement
learning for training of diffusion samplers.
We demonstrate the efficacy of the proposed data-to-energy
IPF on synthetic problems, finding that it can successfully learn
transports between multimodal distributions.
As a secondary consequence of our reinforcement learning
formulation, which assumes a fixed time discretisation scheme for
the dynamics, we find that existing data-to-data Schrödinger bridge
algorithms can be substantially improved by learning the diffusion
coefficient of the dynamics.
Finally, we apply the newly developed algorithm to the problem of
sampling posterior distributions in latent spaces of generative
models, thus creating a data-free image-to-image translation method.

Primary Area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Submission Number: 9096