Adversarial Schrödinger Bridge Matching

TL;DR

A novel Discrete-time IMF (D-IMF) procedure is proposed in which learning of stochastic processes is replaced by learning just a few transition probabilities in discrete time, and it is shown that this procedure can provide the same quality of unpaired domain translation as the IMF.

Abstract

The Schrödinger Bridge (SB) problem offers a powerful framework for combining optimal transport and diffusion models. A promising recent approach to solve the SB problem is the Iterative Markovian Fitting (IMF) procedure, which alternates between Markovian and reciprocal projections of continuous-time stochastic processes. However, the model built by the IMF procedure has a long inference time due to using many steps of numerical solvers for stochastic differential equations. To address this limitation, we propose a novel Discrete-time IMF (D-IMF) procedure in which learning of stochastic processes is replaced by learning just a few transition probabilities in discrete time. Its great advantage is that in practice it can be naturally implemented using the Denoising Diffusion GAN (DD-GAN), an already well-established adversarial generative modeling technique. We show that our D-IMF procedure can provide the same quality of unpaired domain translation as the IMF, using only several generation steps instead of hundreds. We provide the code at https://github.com/Daniil-Selikhanovych/ASBM.

Figures (18)

• Figure 1: Our D-IMF approach performs unpaired image-to-image translation in just a few steps, achieving results comparable to the hundred-step IMF shi2023diffusion. Celeba liu2015faceattributes, male$\rightarrow$female ($128\times 128$).
• Figure 2: Markovian projection of a reciprocal stochastic process $T_q$.
• Figure 3: Reciprocal projection of a stochastic process $T$, i.e., $\text{proj}_{\mathcal{R}}(T) = \int W^{\epsilon}_{|x_0,x_1} dp^T(x_0, x_1)$.
• Figure 4: Reciprocal projection of a discrete stochastic process $q$, i.e., $r(x_0, x_{t_1}, ..., x_{t_N}, x_1) = p^{W^\epsilon}(x_{t_1}, ..., x_{t_N}|x_0, x_1)q(x_0, x_1)$.
• Figure 5: Markovian projection of a reciprocal discrete stochastic process $q$.

Theorems & Definitions (18)

• Theorem 3.1: Discrete Markovian and reciprocal process is the solution of static SB
• Definition 3.2: Discrete Reciprocal Projection
• Proposition 3.3: Discrete Reciprocal projection minimizes KL divergence with reciprocal processes
• Definition 3.4: Discrete Markovian Projection
• Proposition 3.5: Discrete Markovian projection minimizes KL divergence with Markovian processes
• Theorem 3.6: D-IMF procedure converges to the the Schrödinger Bridge
• Theorem 3.7: Reciprocal projection of a process whose joint marginal distribution is Gaussian
• Theorem 3.8: Markovian projection of a discrete Gaussian process
• proof : Proof of Theorem \ref{['thm:discrete-recioprocal-and-ma']}
• proof : Proof of Proposition \ref{['thm:reicprocal-minimizes-KL']}