Schrödinger Bridge Samplers

TL;DR

This work introduces Schrödinger bridge samplers, a regression-augmented IPF approach to approximate target distributions by iteratively reshaping Markov kernels so their T-step marginals converge to the target. It reframes Schrödinger bridges as forward/backward half-bridge KL projections and leverages approximate dynamic programming to learn policy refinements, enabling feasible sampling in continuous spaces without requiring conjugacy. The authors extend the method to sequential, multi-marginal settings (SSB), apply it to discretized Langevin dynamics, and demonstrate substantial variance reductions in normalizing-constant estimation and improved marginal fits in high dimensions and Bayesian regression. The framework connects to optimal transport, flow transport, and shortcuts to adiabaticity, offering a versatile toolkit for diffusion-based sampling and inference in challenging settings.

Abstract

Consider a reference Markov process with initial distribution $π_{0}$ and transition kernels $\{M_{t}\}_{t\in[1:T]}$, for some $T\in\mathbb{N}$. Assume that you are given distribution $π_{T}$, which is not equal to the marginal distribution of the reference process at time $T$. In this scenario, Schrödinger addressed the problem of identifying the Markov process with initial distribution $π_{0}$ and terminal distribution equal to $π_{T}$ which is the closest to the reference process in terms of Kullback--Leibler divergence. This special case of the so-called Schrödinger bridge problem can be solved using iterative proportional fitting, also known as the Sinkhorn algorithm. We leverage these ideas to develop novel Monte Carlo schemes, termed Schrödinger bridge samplers, to approximate a target distribution $π$ on $\mathbb{R}^{d}$ and to estimate its normalizing constant. This is achieved by iteratively modifying the transition kernels of the reference Markov chain to obtain a process whose marginal distribution at time $T$ becomes closer to $π_T = π$, via regression-based approximations of the corresponding iterative proportional fitting recursion. We report preliminary experiments and make connections with other problems arising in the optimal transport, optimal control and physics literatures.

Figures (11)

• Figure 1: Illustration of the iterative proportional fitting procedure. The blue line represents $\mathcal{P}_0(\pi_0) \subset \mathcal{P}(\mathsf{E}^{T+1})$, denoted $\mathcal{P}_0$ in the figure, while the red line represents $\mathcal{P}_T(\pi_T) \subset \mathcal{P}(\mathsf{E}^{T+1})$, denoted $\mathcal{P}_T$. The black line illustrates that the alternating KL projections $\mathbb{Q}^{(i)} \in \mathcal{P}_0(\pi_0)$ and $\mathbb{P}^{(i)} \in \mathcal{P}_T(\pi_T)$ converge towards the Schrödinger bridge $\mathbb{S}$.
• Figure 2: Illustration of the evolutions of the marginal distributions $q_t$ and $s_t$ of the reference process $\mathbb{Q}(\mathrm{d}x_{0:T})$ and the associated Schrödinger bridge $\mathbb{S}(\mathrm{d}x_{0:T})$, respectively. The colors transition from red to green to blue as $t$ increases from $0$ to $T = 40$. The Schrödinger bridge corresponds to the minimal modification of the reference process in terms of KL that interpolates between the initial distribution $\pi_0$ and the target distribution $\pi_T = \pi$, here given in grey.
• Figure 3: Illustration of the marginals of a continuous-time reference process $(q_s)_{s\in[0,\tau]}$ initialized at $\pi_0$ and its discretization $\{q_t\}_{t\in[0:T]}$, shown as solid and dashed red lines respectively. Due to the discretization error, the two paths do not overlap exactly. The solid and dashed black lines correspond to the marginal distributions of the continuous-time and discrete-time Schrödinger bridges $(s_s)_{s\in[0,\tau]}$ and $\{s_t\}_{t\in[0:T]}$ associated with the aforementioned reference processes. Respectively, they form continuous and discrete interpolations between the marginal distributions $\pi_0$ and $\pi_T = \pi_\tau$. The blue dashed line corresponds to the discrete-time sequential Schrödinger bridge, based on solving Schrödinger bridge problems between adjacent distributions in the discrete interpolation $\{\pi_t\}_{t\in[0:T]}$. The blue solid line corresponds to the smooth interpolation $(\pi_s)_{s\in[0,\tau]}$, which is also equal to the marginals of the continuous-time version of the sequential Schrödinger bridge in Section \ref{['sec:ssb_samplers']}.
• Figure 4: Distances between marginals of the Schrödinger bridge $s_t$ and the marginals of the IPF iterates $q_t^{(i)}$, measured as $\log \mathcal{W}_2(s_t,q_t^{(i)})$, for the LQG setting of Section \ref{['sec:lqg']} with discretized Brownian diffusion reference dynamics. Figure \ref{['fig:brownian_IPF_exact']} corresponds to the exact computation of the IPF iterates for $i \in [0:5]$. The remaining plots correspond to the proposed particle-based approximation of the IPF iterations using $N =1,000$ particles and different values of conditional SMC iterations $M$ (using $P = 128$ CSMC particles). The solid lines correspond to the median value of the log-distance calculated over $100$ independent simulations, and corresponding confidence bands represent the $5\%$ and $95\%$ quantiles. Note that the vertical axis of Figure \ref{['fig:brownian_IPF_exact']} is on a different scale than those of the other figures, as the exact IPF iterations yield smaller distances in general.
• Figure 5: Distances between marginals of the Schrödinger bridge $s_t$ and the marginals of the IPF iterates $q_t^{(i)}$, measured as $\log \mathcal{W}_2(s_t,q_t^{(i)})$, for the LQG setting of Section \ref{['sec:lqg']} with discretized Langevin diffusion reference dynamics. Figure \ref{['fig:langevin_IPF_exact']} corresponds to the exact computation of the IPF iterates for $i \in [0:5]$. The remaining plots correspond to the proposed particle-based approximation of the IPF iterations using $N =1,000$ particles and different values of conditional SMC iterations $M$ (using $P = 128$ CSMC particles). The solid lines correspond to the median value of the log-distance calculated over $100$ independent simulations, and corresponding confidence bands represent the $5\%$ and $95\%$ quantiles. Note that the vertical axis of Figure \ref{['fig:langevin_IPF_exact']} is on a different scale than those of the other figures, as the exact IPF iterations yield smaller distances in general.

Theorems & Definitions (2)

• Proposition 2.1
• proof : Proof of Proposition \ref{['prop:IPF_convergence']}