Technical Summary: Data-to-Energy Stochastic Dynamics

Problem Statement

The Schrödinger bridge (SB) problem seeks to find a stochastic dynamical system that bridges two marginal distributions, $p_0$ and $p_1$, while minimizing a transportation cost relative to a reference process (typically an Itô stochastic differential equation). This problem generalizes optimal transport to the stochastic setting and underpins modern generative modeling techniques like diffusion models and flow matching.

Existing algorithms for solving the SB problem, particularly those based on the Iterative Proportional Fitting (IPF) procedure, generally require access to samples from both marginal distributions. This limitation prevents the application of SB methods in scenarios where one or both distributions are defined only by unnormalized densities (energy functions) $p(x) = e^{-E(x)}/Z$, where the partition function $Z$ is unknown and no direct samples are available. This gap is critical for applications such as Bayesian posterior inference in latent spaces of generative models, where the posterior is often known only up to a normalizing constant via a reward or likelihood function.

Methodology

The paper proposes Data-to-Energy (and Energy-to-Energy) IPF, a general method for modeling Schrödinger bridges when distributions are given by unnormalized densities. The approach extends the standard IPF algorithm by integrating techniques from off-policy reinforcement learning and diffusion sampling.

1. Generalized IPF Framework

The standard IPF algorithm iteratively updates forward and reverse-time processes ($\overrightarrow{P}_t$ and $\overleftarrow{P}_t$) by solving half-bridge problems.

• Backward Step (Data-to-Data): When samples from $p_0$ are available, the backward process is trained via maximum likelihood (standard IPF step).
• Forward Step (Data-to-Energy): When $p_1$ is given by an unnormalized density $e^{-E_1(x)}$, the forward process cannot be trained via standard likelihood. Instead, the authors introduce a log-variance loss (conditional VarGrad). This loss minimizes the variance of the log-ratio between the forward trajectory probability and the target density weighted by the reverse trajectory probability:
  $$\text{LLV}(x_0, \theta) = \text{Var} \left( \log \frac{\overrightarrow{p}_\theta(\tau | x_0)}{p_1(x_1) \overleftarrow{p}_\phi(\tau | x_1)} \right)$$
  This objective enforces the proportionality required by the SB solution without needing samples from $p_1$.

2. Off-Policy Reinforcement Learning Strategies

To address the challenge of sampling from complex, high-dimensional unnormalized densities where on-policy training fails to explore all modes, the authors adapt off-policy techniques:

• Replay Buffer: A buffer stores final samples $x_1$ from the forward process. During training, $x_0$ is generated by sampling $x_1$ from the buffer and running a reverse trajectory. This focuses exploration on high-density regions of the target.
• Langevin Updates: The buffer is periodically updated using unadjusted Langevin dynamics on the target density $p_1$ to correct for imperfect fits.
• Trajectory Reuse: The algorithm reuses reverse trajectories to form batches of trajectories sharing an initial point $x_0$, improving the efficiency of the variance-based loss estimation.
• Mixing: The training policy mixes on-policy samples with off-policy samples from the buffer to balance exploration and exploitation.

3. Learnable Diffusion Coefficients

A secondary methodological contribution involves generalizing the IPF algorithm to learn not only the drift functions but also the diffusion coefficients ($\sigma_t$) of the dynamics. By replacing fixed variances in the discretized Euler-Maruyama steps with learnable functions, the method aims to correct for time-discretization errors, a concept inspired by recent work on diffusion samplers.

Key Contributions

1. First General Data-to-Energy SB Method: The paper introduces the first algorithm capable of inferring Schrödinger bridges when one or both distributions are provided as unnormalized densities without access to samples. It successfully learns stochastic bridges with low transport cost between synthetic datasets and densities, performing on par with data-to-data IPF that uses ground-truth samples.
2. Learnable Diffusion Coefficients: The authors demonstrate that existing data-to-data IPF algorithms can be substantially improved by learning the diffusion coefficient in addition to the drift. This generalizes previous findings for diffusion samplers to the broader SB setting, showing that learnable variance allows for more accurate modeling, particularly with coarse time discretizations.
3. Outsourced Posterior Sampling: The method is applied to the problem of sampling posterior distributions in the latent spaces of generative models (e.g., GANs, VAEs). By modeling a bridge between the prior latent distribution and a reward-weighted posterior, the authors create a data-free image-to-image translation method. This approach preserves semantic content (style, background) not constrained by the target condition, unlike standard diffusion samplers.

Experimental Results

The efficacy of the proposed method is demonstrated across several benchmarks:

• Synthetic 2D Benchmarks: On tasks like Gaussian to Gaussian Mixture Models (GMM) and "Two Moons" distributions, the Data-to-Energy IPF achieves performance comparable to Data-to-Data IPF. The inclusion of learnable variance significantly reduces the Wasserstein-2 distance and Path KL, especially when the number of discretization steps is small.
• Energy-to-Energy Bridging: Preliminary results show the algorithm can successfully bridge two distributions defined solely by energy functions, without access to samples from either.
• Image Translation (MNIST & CIFAR-10): In latent space experiments using SN-GAN and StyleGAN, the method translates prior samples to target class posteriors (e.g., generating "dogs" or "fives").
  • Metrics: The method achieves lower FID scores compared to rejection sampling and Langevin dynamics in several classes.
  • Qualitative: Visual inspection confirms that the bridge preserves style features (e.g., digit thickness, background orientation) while transforming the content to the target class, suggesting effective preservation of unconstrained semantic information.
• Ablation Studies: Experiments confirm that replay buffers, Langevin updates, and trajectory reuse significantly improve Path KL metrics. Reusing backward trajectories improves Path KL but may negatively impact mean log-reward, suggesting a trade-off between mode coverage and transport cost.

Significance and Claims

The paper claims to provide a general framework for data-to-energy Schrödinger bridges, filling a gap in the literature where SB methods were previously restricted to settings with sample access for both marginals.

The authors position their work as a significant step toward scalable, data-free generative modeling and posterior inference. By leveraging off-policy reinforcement learning techniques, the method enables the learning of optimal transport dynamics in high-dimensional latent spaces where only energy functions are available. The ability to learn diffusion coefficients further enhances the robustness of these dynamics against discretization errors.

The authors modestly note that while the method scales to high-dimensional image generation tasks, future work is needed to address mode collapse and further improve mode coverage. They also identify potential future directions, including text-conditional image reward fine-tuning and discrete Schrödinger bridge problems, but do not claim immediate success in these unexplored areas.