Data-to-Energy Stochastic Dynamics

TL;DR

This work addresses learning Schrödinger bridges between distributions when one or both marginals are unnormalised densities, i.e., a data-free setting. It generalises iterative proportional fitting (IPF) to data-to-energy and energy-to-energy SB via time-discretised SDEs and off-policy reinforcement learning techniques, optionally learning the diffusion coefficient to mitigate time discretisation error. The method yields transports between multimodal densities, improves diffusion-sampler training, and enables data-free image-to-image translation by operating in latent spaces, with outsourced sampling to posterior-like tasks. The approach is demonstrated on synthetic benchmarks and latent-space outsourcing problems (e.g., MNIST and CIFAR-10) using metrics such as ELBO, path KL, and Wasserstein distances, and code is released for reproducibility.

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 allow to infer 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. Code: https://github.com/mmacosha/d2e-stochastic-dynamics

Figures (11)

• Figure 1: Data-to-Data IPF
• Figure 2: $\mathcal{W}_2^2$ and Path KL depending on the number of discretization steps for Gauss $\leftrightarrow$ GMM.
• Figure 3: Comparison of learnt processes at various IPF iterations for data-to-data (top), data-to-energy (middle) and energy-to-energy (bottom) settings. For the energy-to-energy setting we use two different mixtures of Gaussian distributions.
• Figure 4: Outsourced Schrödinger bridge on MNIST and CIFAR-10. The bridge preserves style features (thickness, background colour, orientation) while tranforming digits to the target class.
• Figure 5: Curated examples of outsoursed SB with SN-GAN and StyleGAN generators.