Diffusion & Adversarial Schrödinger Bridges via Iterative Proportional Markovian Fitting

TL;DR

This work addresses Schrödinger Bridge problems by introducing IPMF, a bidirectional IMF-based procedure that unifies IMF and IPF to stabilize training and enable flexible control over the similarity-quality trade-off in translations. The authors prove exponential convergence for Gaussian marginals and demonstrate the method across high-dimensional Gaussian, illustrative 2D, SB benchmark, and unpaired image-to-image translation tasks, highlighting robustness to starting couplings. They show that the starting coupling can be designed to tailor results for specific tasks, offering a common framework across discrete and continuous-time SB formulations. The practical impact lies in stabilizing SB solvers, mitigating error accumulation, and enabling task-aware translation with controllable fidelity and realism in applications like unpaired domain translation and diffusion-based modeling.

Abstract

The Iterative Markovian Fitting (IMF) procedure, which iteratively projects onto the space of Markov processes and the reciprocal class, successfully solves the Schrödinger Bridge (SB) problem. However, an efficient practical implementation requires a heuristic modification -- alternating between fitting forward and backward time diffusion at each iteration. This modification is crucial for stabilizing training and achieving reliable results in applications such as unpaired domain translation. Our work reveals a close connection between the modified version of IMF and the Iterative Proportional Fitting (IPF) procedure -- a foundational method for the SB problem, also known as Sinkhorn's algorithm. Specifically, we demonstrate that the heuristic modification of the IMF effectively integrates both IMF and IPF procedures. We refer to this combined approach as the Iterative Proportional Markovian Fitting (IPMF) procedure. Through theoretical and empirical analysis, we establish the convergence of the IPMF procedure under various settings, contributing to developing a unified framework for solving SB problems. Moreover, from a practical standpoint, the IPMF procedure enables a flexible trade-off between image similarity and generation quality, offering a new mechanism for tailoring models to specific tasks.

Figures (7)

• Figure 1: Diagrams of IPF, IMF, and unified IPMF procedure. All procedures aim to converge to the Schrödinger Bridge, i.e., a Markovian process in the reciprocal class, with marginals $p_0$ and $p_1$.
• Figure 2: Convergence of IPMF procedure with different starting process $q^0$.
• Figure 3: Samples from DSBM and ASBM learned with IPMF using IMF and $q^{\text{inv} 7}$ starting processes on Colored MNIST 3$\rightarrow$2 ($32\times 32$) translation for $\epsilon=10$.
• Figure 4: Results of CelebA at 64$\times$64 size for male$\rightarrow$female translation learned with ASBM and DSBM using various starting processes for $\epsilon=1$.
• Figure 5: Visualization of learned processes with DSBM and ASBM solvers for Gaussian$\!\rightarrow\!$Swiss roll translation using IMF, IPF, Identity starting processes for $\epsilon=0.1$.

Theorems & Definitions (28)

• Theorem 3.1
• Theorem 3.2: Convergence of IPMF for Gaussians
• proof : Proof idea
• Theorem 3.3: Convergence of IPMF under boundness assumption
• proof : Proof of Theorem \ref{['lemma-eot-gaussian']}
• Lemma D.1: Improvement after IPF steps
• proof
• Lemma D.2: Marginals norm bound during IPMF procedure
• proof
• Lemma D.3: IPF step does not change optimality matrix $A$