Unpaired Image-to-Image Translation via Neural Schrödinger Bridge
Authors: Beomsu Kim
Gihyun Kwon
Kwanyoung Kim
Jong Chul Ye
Venue: ICLR 2024 poster (2024)
Keywords: Style Transfer
Schrödinger Bridge Problem
Adversarial Training
Primary area: generative models
TL;DR: We apply Schrödinger Bridges to unpaired image-to-image translation.

Abstract: Diffusion models are a powerful class of generative models which simulate stochastic differential equations (SDEs) to generate data from noise. While diffusion models have achieved remarkable progress, they have limitations in unpaired image-to-image (I2I) translation tasks due to the Gaussian prior assumption. Schrödinger Bridge (SB), which learns an SDE to translate between two arbitrary distributions, have risen as an attractive solution to this problem. Yet, to our best knowledge, none of SB models so far have been successful at unpaired translation between high-resolution images. In this work, we propose Unpaired Neural Schrödinger Bridge (UNSB), which expresses the SB problem as a sequence of adversarial learning problems. This allows us to incorporate advanced discriminators and regularization to learn a SB between unpaired data. We show that UNSB is scalable and successfully solves various unpaired I2I translation tasks. Code: \url{https://github.com/cyclomon/UNSB}

Published as a conference paper at ICLR 2024
UNPAIRED IMAGE-TO-IMAE TRANSLATION VIA NEURAL SCHRÖDINGER BRIDGE
Beomsu Kim
Gibyun Kwon*
Kwanyoung Kim
Jong Chul Ye!
Equaal contribution
1 Kim Jaechul Graduate School of AI, KAIST
2 *Deaarmeen of Bio and Brain Engineering, KAIST beomsu.. kim, cyclomon, cubeyoung, jong ye yel@kkist..cc .c..sst.c.... kr
ABSTRACT
Diffusion models powerful class of generative models which simulate stochas- tic differential equations (SDEs) to generate data from noise. While diffusion models have achieved remarkable progress, they have limitations in unpaired image- to-image (I2I) translation tasks due to the Gaussian prior assumption. Schrödinger Bridge (SB), which learns SDE to translate between arbitrary distributions, have risen as an attractive solution to this problem. Yet, to our best knowledge, none of SB models so far have been successful unpaired translation between high-resolution images. In this work, we propose Unpaired Neural Schrödinger Bridge (UNSB), which expresses the SB problem as a sequence of adversarial learning problems. This allows us to incorporate advanced discriminators and regularization to learn a SB between unpaired data. We show that UNSB is scalable and successfully solves various unpaired 121 translation tasks. Code: htttps://ittuut com cttp:///wocoooo/ooooooooooo////o
INTRODUCTION
Diffusion models DSckl-Dickstein al.) 2015; Ho et al. 2020 Song et al., 2021a b). class of generative model which generate data by simulating stochastic differential equations (SDEs), have achieved remarkable progress over the past few years. They are capable of diverse and high-quality sample synthesis (Xiao et al. 2022), desiderata which difficult obtain for several widely acknowledged models such as Generative Adversarial Networks (GANs) (Goodfellow et al. 2014)). Variational Autoencoders (VAEs) (Kingma & Welling 2014), and flow-based models (Dinh et al. 2015). Furthermore, the iterative nature of diffusion models proved to be useful for a variety of downstream tasks, e.g., image restoration (Chung al. 2023) and conditional generation (Rombach) ct al. 2022).
However, unlike GANs and their family which are free to choose any prior distribution, diffusion models often assume a simple prior, such as the Gaussian distribution. In other words, the initial condition for diffusion SDEs is generally fixed be Gaussian noise. The Gaussian assumption prevents diffusion models from unlocking their full potential in unpaired image-to-image translation tasks such as domain transfer, style transfer, or unpaired image restoration.
Schrödinger bridges (SBs), a subset of SDEs, present promising solution to this issue. They solve the reguaaz-e optimal transport (OT) problem, enabling translation between two arbitrary distributions. Their flexibility have motivated several approaches to solving SBs via deep learning. For instance, [Bortoli al. (202T) extends the Iterative Proportional Fitting (IPF) procedure to the continuous setting, Chen et al. (2022) proposes likelihood training of SBs using Brwward-Baccward SDEs theory, and Tong et al. (202)) solves the SB problem with a conditional variant of flow matching.
Some methods have solved SBs assuming one side of the two distributions is simple. Wang et al. (2021a) developed two-stage unsupervised procedure for estimating the SB between Dirac delta and data. iSBB (Liu et al. 2023) and InDI (Delbracio Milanfar 2023) uses paired data to lear SBs between Dirac delta and data. Su et al. (023) discovered DDIMs (Song al. 2021a) as SBs between

---

Published as a conference paper at ICLR 2024
data and Gaussian distributions. Hence, they were able to perform image-to to--nnge translation by concatenating two DDIMs, i.e., by passing through an intermediate Gaussian distribution.
Yet, to the best of our knowledge, work so far has successfully trained SBs for direct translation between high-resolution images in the unpaired setting. Most methods demand excessive computation, and even if it is not the case, we observe poor results. In fact, all representative SB methods fail even on the simple task of translating points between two concentric spheres as dimension increases.
In this work, first identify the main culprit behind the failure of SB for unpaired image- to-image translation as the curse of dimensionality. As the dimension of the considered data increases, the samples become increasingly sparse, failing capture the shape of the underlying image manifold. This sampling error then biases the SB optimal transport map, leading to an inaccurate SB.
Based on the self-similarity of SB, referring to the intriguing property that SB restricted to sub- interval of its time domain also SB, we therefore propose the Unpaired Neural Schrödinger Bridge (UNSB), which formulates the SB problem as sequence of transport cost minimization problems under the constraint on the KL divergence between the true target distribution and the model distribution. We show that its Lagrangian formulation naturally expresses the SB a composition of generators learned via adversarial learning (Goodfellow et al. (014). One of the important advantages of the UNSB formulation of SB that it allows us to mitigate the curse of dimensionality on two levels: extend the discriminator in adversarial learning to more advanced one, and we can add regularization to enforce the generator to learn a mapping which aligns with our inductive bias. Furthermore, all components of UNSB are scalable, UNSB is naturally scalable as well to large scale image translation problems.
Experiments on toy and practical image-to torse-tttoooo-tt im-ima translation tasks demonstrate that UNSB opens up a new research direction for applying diffusion models to large scale unpaired translation tasks. Our contributions be summarized as follows.
We identify the cause behind the failure of previous SB methods for image-to-image transla- tion as the curse of dimensionality. We empirically verify this by using a toy task as sanity check for whether an OT-based method robust the curse of dimensionality.
We propose UNSB, which formulates SB as Lagrangian formulation under the constraint on the KL divergence between the true target distribution and the model distribution. This leads to the composition of generators learned via adversarial learning that overcome the curse of dimensionality with advanced discriminators.
UNSB improves upon the Denoising Diffusion GAN (Xiao et al. 2022) by enabling trans- lation between arbitrary distributions. Furthermore, based on the comparison with other existing unpaired translation methods, we demonstrate that UNSB is indeed a generalization of them by overcoming their shortcomings.
2 RELATED WORKS
Schrödinger Bridges. The Schrödinger bridge (SB) problem, commonly referred to as the entropy- regularized Optimal Transport (OT) problem (Schrödinger 1932 Léonard 2013). is the task of learning a stochastic process that transitions from an initial probability distribution to terminal distribution over time, while subject to a reference measure. It is closely connected to the field of stochastic control problems (Caluya & Halder 2 Chen al. 2021). Recently, the remarkable characteristic of the SB problem, which allows for the choice of arbitrary distributions as the initial and terminal distributions, has facilitated solutions to various generative modeling problems. In particular, novel algorithms leveraging Iterative Proportional Fitting (IPF) (Bortoli et al, 202T Vargas et al. (02T) have been proposed to approximate score-based diffusion. Building upon these algorithms, several variants have been introduced and successfully applied in diverse fields such as, inverse problems (Shi et al. 2022). Mean-Field Games (Liu et al.] (022), constrained transport problems (Tamir et al. 2023). ) Riemannian manifolds (Thornton et al. (022), and path samplers (Zhang & Chen 20e1 Ruzayqat et al. 2023). Some methods have solved SBs assuming one side of the two distributions simple. In the unsupervised setting, |Wang et al. (2021a) first learss a SB between a Dirac delta and noisy data, and then denoises the noisy samples. In the supervised setting, JSSB (Liu et al. 2023) and InDI (Delbracio & Milanfar 2023) used paired data to learn SBs
2

---

Published as a conference paper at ICLR 2024
Souree
NFB
NEB =2
NFB=3
N
Obearved Points
Vanilla UNSB
Unolserved Toints
Horse - Zebra
Map + Satellite
Manifold A
Schrödiager Bridge Didfasion
Figure 1: Left: Illustration of trajectories for Vanilla SB and UNSB. Due to the curse of dimensional- ity, observed data in high dimensions become sparse and fail describe image manifolds accurately. Vanilla SB learns optimal transport between observed data, leading to undesirable mappings. UNSB employs adversarial learning and regularization to learn an optimal transport mapping which success- fully generalizes beyond observed data. Right: UNSB can be interpreted as successively refining the predicted target domain image, enabling the model to modify fine details while preserving semantics. See Section Here, NFE stands for the number of function evaluations.
between Dirac delta and data. DDIB (Su et al. 2023) concatenates two SBs between data and the Gaussian distribution for image-to-image translation. Recent works (Gushehin et al. 2023a Shi et al.
2023) have achieved unpaired image-to-image translation for S 128 x 128 resolution images. But, they are computationally intensive, often taking several days train. To our best knowledge, our work represents the first endeavor to efficiently learn SBs between higher resolution unpaired images.
Unpaired image-to-image translation. The aim of image-to-image translation (I2I) is produce an image in the target domain that preserves the structural similarity to the source image. A seminal work in 12I pix2pix (Isola et al., 2017). which undertook the task with paired training images, utilizing straightforward pixel-wise regularization strategy. However, this approach is not applicable when dealing with unpaired training settings. the context of unpaired data settings, early methods like (Zhu et al. 2017 Huang et al. 201 8) maintained the consistency between the source and output images through a two-sided training strategy, namely sycle-consistency. However, these models are plagued by inefficiencies in training due to the necessity of additional generator training. In order to circumvent this issue, recent I2I models have shifted their focus to one-sided I2I translation. They aim preserve the correspondence between the input and output through various strategies, such as geometric consistency (Fu et al.) 2019) and mutual information regularization (Benaim & Wolf 2017). Recently, Contrastive Unpaired Translation (CUT) and its variants (Park et al.) 2020]Jungg ct al. 2022 Wang et al. 2021b Zheng et al.) [222]) have demonstrated improvements in 12I tasks by refining patch-wise regularization strategies. Despite the impressive performance demonstrated by previous GAN-based models in the 12I domain, this paper demonstrates that our SB SB-based approach paves the way for further improvements the 12I task by introducing the iterative refinement through SB that overcomes the potential mode collapsing issue from single-step generator from GAN. We believe that it represents a new paradigm the ongoing advancement image-to-image translation.
SCHRÖDINGER BRIDGES AND THE CURSE OF DIMENSIONALITY
Given distributions no. कı Rd the Schrödinger Bridge problem (SBP) seeks the most likely random process (xg E [0, }}} that interpolates To and m. Specifically, let be the path space on Rd i.c.. the space of continuous functions from [0, to Rdd and let P(2) be the space of probability measures on Ω. Then, the SBP solves
Q
SB
arg min DKL. (Q[/I) s.t. Qo = Tos
Q
(I)
QEP(R)
where WT is the Wiener measure with variance 7, and Qr denotes the marginal of Q at time t. We call QB the Schrödinger Bridge (SB) between To and m.

---

Published as a conference paper at ICLR 2024
The SBP admits multiple alternative formulations, and among those, the stochastic control formulation and the static formulation play crucial roles in our work. Both formulations will serve as theoretical bases for Unpaired Neural Schrödinger Bridge algorithm, and the static formulation will shed light on why previous SB methods have failed on unpaired image-to-image translation tasks.
Stochastic control formulation. The stochastic control formulation Pra 1991) shows that { QSB can be described by an Itō SDE
dx. u SB dt + VT dws
(2)
where the Line-arriing drift u hh is solution the stochastic control problem
arg min E J. | 2 dt 2
d Ue dt + VT dw,
SB
s.t.
Io To, æ1
u
assuming satisfies certain regularity conditions. Eq. (3) says that among the SDEs of the form Eq. (2) with boundaries and Ti the drift for the SDE describing the SB has minimum energy. This formulation also reveals
useful properties of SBs. First, {æ+) } is Markov chain, and second, {æ,} converges to the optimal transport ODE trajectory + 0. Intuitively, controls the amount of randommess in the trajectory {æe).
The static formulation of SBP shows that sampling from QSB is extremely simple, assuming we know the joint distribution of the SB at E {0,1), denoted QSB Concretely, for {x,} ~ Q conditioned upon initial and terminal points and 21. the density of æ can be described by Gaussian density (Tong et al.
Static formulation.
2023)
p( (o 21) = Nax [l:
(1 () ((1 l)r
Hence, simulate the SB given initial point æo. may sample æ, according to
pl (m)) p( æ1) d (a1 (xo)
(5)
where Q denotes the conditional distribution of given æa. Surprisingly, Q is shown to be solution to the entroyy-regularee optimal transport problem
Q
arg min ((=001))~~110>> x 11 27 (7)) >Ell(mo,a1)
(6)
where (T,, Tı) is the collection of joint distributions whose marginals are To and T1, and H denotes the entropy function. For discrete To and Ii it is possible to find QSP via the Sinkhorn-Knopp
algorithm, and this observation has inspired several algorithms (Tong al. 2023 Pooladian et al. 2023) for approximating the SB.
Curse of dimensionality. According to the static for- mulation of the SBP. any algorithm for solving the SBP can be interpreted ing to the entropy-regularized optimal transport Eq. (0) of the marginals No and Ti. However, in practice, we only have finite number of samples { /==1 and {at") m== ~ TI. This means, the SB is trained to transport samples between the empirical distributions 1 M + E dap and 37 Em= dem Due the curse of dimensionality, the samples fail to describe the image
learning an interpolation accord-
Dim.
Dim. 10
20
Dim. 100
Density
15
Dim. 1000
5-
manifolds correctly in high dimension. Ultimately, Q yield image pairs that do not meaningfully correspond to one another (see Figure I or the result for Neural Optimal Transport (NOT) in Figure 5).
cosine-sim(xe.e (xo. ug!!
Figure 2: Curse of dimensionality.
To illustrate this phenomenon in the simplest scenario, we consider the case where To and are supported uniformly on two concentric d-spheres of radii 1 and 2, respectively. Then, samples from QSB should have near-one cosine similarity, since QSB should transport samples from To radially outwards. However, when we draw M N 1000 i.i.d. samples from To and Ti, and calculate the cosine similarity between (æa, +1) ~ Q where @ is estimated using the Sinkhorn-Knoy increases (see Figure 2). Thus, high dimension, QB approximated by Sinkhorn BinkhornKnnopp will interpolate between nearly orthogonal points.
algorithm, we observe decreasing similarity dimension
4

---

min Published as a conference paper at ICLR 2024
4 UNPAIRED NEURAL SCHRÖDINGER BRIDGE (UNSB)
We now explain our novel UNSB algorithm which shows that SB can be expressed as a composition of generators learned via adversarial learning.
Specifically, given a partition {4:]--0 of the unit interval [0, 1] such that to - O, IN = 1, and < 1441, i, we can simulate SB according to the Markov chain decomposition
p( (r }) p( ) - (xx
p( Fio (p(())))
The decomposition Eq. 7 allows us to learn the SB inductively: we learn p( (xt.) assuming we able to sample from p) ( Having approximated p(ar, l. we can sample from p(#*; 1). so we may learn p( [t; Since the distribution of æto = 20 is already known as Tn. our main interest lies in learning the transition probabilities p( assuming we can sample from p(æ).). This procedure is applied recursively for i = .,...,N 1.
Let qo. (x) ( be conditional distribution parametrized by a DNN with parameter Φ.. Intuitively, qo, (x) (z,,) is generator which predicts the target domain image for Zt.- We define
gd, ( x1) = qo. (x1 xt, ),(;)), qo. (x1) = Ep(m.,) 19o, (x1 [x,,)]
(8)
The following theorem shows how to optimize o; and sample from p( xt.)) ) Theorem 1. For any tú consider the following constrained optimization problem
CSB(dr,ti) tt) = Egg, (m)) =1
H4) x1),
(9)
s.1. LAdv(di. (i) = DKL. (uo, (01))) =
(10)
and define the distributions
p(ar.4s Rt,) (Z |sst 13 + (1 - sit (a 8441 1
8,4 (11))(t-)))) I,)
(11)
where 8141 (1:41 t1)/(1 12) and
Jo. (Iezin (zt.) = Bgg; (as (mm.) [p( 1 (21, Ie, Ioh (Zt14.) ) = Eptrs, 190, (K 1 [2))). (12)
If oi solves Eq. 99 then have
4e. (x) ( = p( (at,,), qo. (We, (2x,,) p( 100e), (300,)
Lo. (2 .) = p(
(13)
Proof Sketch. Using the stochastic control formulation of SB, we show that the SB satisfies certain self-similarity property. which states that the SB restricted any sub-interval of [0, 1] is also SB.
Note that Eqs. (TI) 12 and (9) under the constraint Eq. (IO) are indeed counter-parts of the static formulation Eqs.
and (6) respectively, for the restricted domain [4,, 1]. The static formulation
of SB restricted to the interval t,, shows that 9a. (x1 xt, = p( (xx,) is the solution to Eq. (9).
If we assume Φ; solves Eq. (9) such that 4o. (01 (x.,) p(z1 (x Eq. (I3 follows by simple
calculation. The detailed proof can be found in the Appendix.
In practice, by incorporating the equality constraint in Eq. (10) into the loss with a Lagrange multiplier, we obtain the UNSB objective for a single time-step 7,
min LUNSB 44) == (di 42) + ASH,
(14)
Since is impractical to use separate parameters Φ; for each time-step t, and to learn p( (x,))
sequentially for O, O,..., N 1, we replace 4o: (ær (a,,)) which takes æt, as input, with time- conditional DNN 4: [m.,), which shares a parameter & for all time-steps 4, and takes the tuple (xx,, t4) input. Then, optimize the sum of LUNSB (0) t,) over i O..... NN -1.
Training, For the training of our UNSB, we first randomly choose a time-step &; to optimize. To calculate LUNSB(ó, f4), we sample 2t, and ~ Ri, where म. denotes the target distribution. The sampling procedure of xi, will be described soon. The sample V1, is then passed through go (0(x)) (x (xx.) to obtain 21 (z,,) that refers to the estimated target data sample given øt,- The pairs (æ), 21 (Z.
and (æ), x (xx,)) )) are then used to compute Lsn(o, t²) and C t.) in Eq. (9) and Eq. (10). respectively. Specifically, we estimate the entropy term in Een with a mutual information estimator,
5

---

Published as a conference paper at ICLR 2024
(a (21,)
generation for 0g....8--
Regulrrzztion Loss
Lt ),1)
x (x(())
OT & Entropy Loes
Sg
x(1 a(1)+1))
8(1++1ix
Ti~~ n0
+
0o(21 m4,)
22 (3t..)
Adv. Loss
\O, oj+t 1)
0t 1++1
<(t-1) = 18++1 -g)/(1
21
Øj+1 = +(}++)) a((+1))11 4)r
Generation Stage
Training Stage
Figure 3:
Generation and training process UNSB for time step i.. sg means stop gradient.
using the fact that for random variable X, [XX, (() where I denotes mutual information. We then estimate the divergence in LAdv with adversarial learning. Intuitively, æ1 and æ (æx,) are "real" and "fake" inputs to the discriminator. This process shown in the training stage of Figure
Generation of intermediate and final samples. We now describe the sampling procedure for the intermediate samples for training and inference. We simulate the Markov chain Eq. (7) using ds as follows: given Lt; 4o (zxx) (note that Lt 4o if 0), we predict the target domain image ær (3)) ~ go(2r ( ). We then sample 2tt5 ~ 1s ( (1) according to Eq. 11) by interpolating æo and xi (xx)) and adding Gaussian noise. Repeating this procedure for j =0...,??
get It. ~ (o(,),) With optimal é, by Theorem It, ~ (0(x),) = p(x,,). Thus, the trajectory (x1 (axt, : O,..., N }} can be viewed iterative refinement of the predicted target domain sample. This process is illustrated within the generation stage of Figure
4.1 COMBATING THE CURSE OF DIMENSIONALITY
Advanced discriminator in adversarial learning. One of the most important advantages of the formulation Eq. (14) is that we can replace the KL-divergence in LAdv by any divergence or metric which measures discrepancy between two distributions. Such divergence or metric can be estimated through adversarial learning, i.e., its Kantorovich dual. For example,
2 DIsD (go, (2))) (#1))l((z1)) - 1og(4) max Ep( [log (xx1)] + E (m.) (log) [(zz)) (15)
where D is discriminator. This allows to use various adversarial learning techniques to mitigate the curse of dimensionality. For instance, instead of using standard discriminator which distinguishes generated and real samples the instance level, distinguishes samples on the patch level. The Markovian discriminator is effective at capturing gegh-frequency characteristics (e.g., style) of the target domain data (Isola al. (017).
Markovian discriminator which
Regularization. Furthermore, we augment the UNSB objective with regularization, which enforces the generator network ds satisfy consistency between predicted æ and the initial point Zg:
L (s, (2) == Ep(2o.0,, Ey (z) (e.,) (R(wo a1)
(16)
Here, is a ecalar-valued differentiable function which quantifies an application-speciii measure of similarity between its inputs. In other words, R reflects inductive bias for similarity between two images. Thus, the regularized UNSB objective for time f; is
LuNSB(ó, t.) Laav (Ø, tx) + (SB,,,Cs((d AReg,e, CReg (b,tz);
(17)
which is the final objective in our UNSB algorithm.
4.2 SANITY CHECK ON TOY DATA
Two shells. With the two shells data of Section we show that representative SB methods suffer from the curse of dimensionality, whereas UNSB does not, given appropriate discriminator and regulariza- tion. For SB methods, consider Sinkhom-Knopp (SK), SB Conditional Flow Matching (SBCFM)

---

Published as a conference paper at ICLR 2024
(Tong et al. 2023 Diffusion SB (DSB) (Bortoli et al. 2021). and SB-FBSDE (Chen et al. 2022). For baselines, we use recommended settings, and for UNSB, we train the discriminator to distinguish real and fake samples by input norms, and choose negative cosine similarity as R in Eq. (16). Only 1/ samples from each To. тi are used throughout the training. In Figure 4 we see all baselines either fail to transport samples to
2.1
the target data manifold or fail maintain high
cosine similarity between input and outputs. On the other hand, UNSB is robust to dimension.
TBSCE LN59
04
500TM
1000
4001 Dinwnice
Two Gaussians. We also verify whether UNSB can learn the SB between two Gaussians. The SB between Gaussian distributions is also Gaus- sian, and its mean and covariance can be derived in closed form (Bunne al.) (023). Thus, we may measure the error between approximated and true solutions exactly. We consider learn- ing the SB between NV 11, I) and (1, in 50-dimensional space. In Table D we that UNSB recovers the mean and covariance of the ground-truth SB relatively accurately.
Dininiie
Figure 4: Results on two shells.
Method
MSE
Σ MSE
SK
1.5c-5
3.10-6
SBCFM
4.7e-5
2.3e-5
DSB
4.028
0.139
SB- SB-FBSDE
2.974
2.5c-9
UNSB
0.008
6.4e-7
Table 1: Results on two Gaussians.
4.3 RELATION TO OTHER UNPAIRED TRANSLATION METHODS
UNSB vs. GANs. N 1 version of UNSB is nearly equivalent to GAN-based translation methods such as CUT, i.e., UNSB may be interpreted as a multi-step generalization of GAN methods. Multi-step models are able to fit complex mappings that GANs are unable to by breaking down complex maps into a composition simple maps. That is why diffusion models often achieve better performance than GANs (Salmona al. 2022). Similarly, we believe the multi-step nature of UNSB arising from SB allow it to obtain better performance than GAN-based translation methods.
UNSB vs. previous diffusion methods. Previous diffusion- based methods such as SDEdit (Meng al. 2022) also translate images along SDEs. However, the SDEs used by such methods (e.g.. VP-SDE or VE-SDE) may be suboptimal in terms of transport cost, resulting in large NFEs. On the other hand, SB allows fast generation, as it learns the shortest path between domains. Furthermore, UNSB uses adversarial learning, it may generate feasible samples even with NFE 1. On the contrary, diffusion methods do not perform so well on complex tasks such as 2orse2Zebra.
RESULTS FOR UNPAIRED IMAGE-TO-IMAGE TRANSLATION
In this section, we provide experimental results of UNSB for large scale unpaired image-to-image translation tasks.
UNSB settings. We use the Markovian discriminator in Cadv and the patch-wise contrastive matching loss (Park et al.) (2200) in Reg.. We set the number of timesteps unit interval into to. f...... (5. We set ASB,t, = AReg,t, = 1 and r = 0.01. We measure the sample for each O...., --1. So, NFE i denotes sample evaluation or visualization of æ (. Other details are deferred to the Appendix.
N 5, so we discretize the
quality of xx (xx,.
Evaluation. We use four benchmark datasets: Horse2Zebra, Map2Cityscape, Summer2Winter, and Map2Satellite. All images are resized into 256x256, We use the FID score (Heusel et al. 2017) and the KID score (Chen et al. 2020) to measure sample quality.
Baselines. We note all SB methods SBCFM, DSB, and SB-FBSDE do not provide results on 2 256 resolution images in their paper due to scalability issues. While we were able to make SBCFM work on 256 resolution images, the results were poor, we deferred its results to the Appendix. Instead, we used Neural Optimal Transport (NOT) (Korotin et al. 2023) as a representative for OT. For GAN methods, used CycleGAN (Zhu al. 2017). MUNIT (Huang al.) 2018), 8), DistanceGAN (Benaim & Wolf, 2017). GeGAN (Fu et al. (2019). and CUT (Park et al. 2020). For diffusion, we used SDEdit (Meng et al.
2022 and P2P (Hertz et al. 2022) with L.DM (Rombach al. 2022).
7

---

Published as a conference paper at ICLR 2024
Horse2Zebra
Summer2Winter
Label2Cityscape
Map2Satellite
Method
NFE
Time
FID.
KIDI
FID.
KID,
FID
KIDI
FID
KID.
NOT
0.006
104,3
5.012
185.5
8.732
221.3
19.76
224.9
16.59
CycleGAN
1.957
84,9
1.022
76.3
MUNIT
0.011
133.8
3,790
115.4
4.901
91.4
6.401
181.7
12.03
Distance
0.009
72.0
1.856
97.2
2.843
81.8
4,410
98.1
5.789
GeGAN
0.0027
2.051
97.5
2.755
105.2
6.824
79.4
5.153
CUT
0.0033
45.5
0.541
84.3
1.207
56.4
1.611
56.1
3.301
SDEdit
30
1.98
97.3
4.082
118.6
3.218
P2P
50
120
60,9
1.093
99.1
2.626
Ours-bes 5
0.045
35.7
0.587
73.9
0.421
53.2
1.191
47.6
2.013
Table 2: NFE and generation time (sec.) per image and FID and XIDx 100 of generated samples.
Inpur
Ours
P2P
SDEdi
NOT
CUT
CyekGAN
DEAGAN
GeGAN
MUNIT
(a) Comparison natural domain data.
Inpus
CUT
CyeleGAN
DisIGAN
MUNIT
(b) Comparison artificial domain data.
Figure 5: Qualitative comparison of image-to-image translation results from our UNSB and bascline 12I methods. Compared to other one-step baseline methods, our model generates more realistic domain-changed outputs while preserving the structural information of the source images.
Since LDM is trained on natural images, we can use LDM on Horse2Zebra and Summer2Winter but not on Label2Cityscape and Map2Satellite. Hence, omit diffusion results those two tasks.
Comparison results. We show quantitative results in Table 2 where we observe our model out- performs baseline methods in all datasets. In particular, out model largely outperformed early GAN-based translation methods. When compared to recent models such as CUT, our model still shows better scores. NOT suffers from degraded performance in all of the datasets.
Qualitative comparison in Figure 5 provides insight into the superior performance of UNSB. Our model successfully translates source images to the target domain while preventing structural inconsis-
8

---

Published as a conference paper at ICLR 2024
tency. For other GAN-based methods, the model outputs do not properly reflect the target domain information, and some models fail to preserve source image structure. For NOT, we observe that the model failed to overcome the curse of dimensionality. Specifically, NOT hallucinates structures not present in the source image (for instance, see the first row of Figure 3. leading to poor samples. For diffusion-based methods,
they fail to fully reflect the target domain style.
NFE analysis. We investigate the relationship be- tween NFE and the translated sample quality. In Fig- ure top, we observe that using NFE 1 resulted in relatively poor performance, sometimes even worse than existing one-step GAN methods. However, we see that increasing the NFE consistently improved generation quality: best FIDs were achieved for NEE values between and for all datasets. This trend also evident in our qualitative comparisons. Looking at Figure]l [ we observe that as NFE increases, there gradual improvements in image quality. However, UNSB at times suffers from the problem where ar- tifacts occur with large NFE, as shown in Figure 6 bottom. We speculate this causes increasing FID for large NFEs on datasets such as Map2Satellite.
10
60-
50-
815-1
Lebe
Napobele Ba
tep4
Figure 6: NFE analysis and failure case.
Ablation study. Unlike GAN-based methods, UNSB able to generate samples in multiple steps. Also unlike SB methods, UNSB is able to use advanced discriminators and regularization. To understand the influence of those factors on the final performance, we performed an ablation study on the Horse2Zebra dataset. We report the results in Table 3 In the case where we use neither advanced discriminator nor reg- ularization (corresponding to the previous SB method SBCFM), we see poor results. Then, as we add multi- step generation, advanced discriminator (Markovian discriminator, denoted "Patch"), " and regularization, the result gradually improves until obtain the best performance given by UNSB. This implies that the three components of UNSB play orthogonal roles.
Dise.
Reg.
NFE
FI1
KID
x
x
5
230
18.8
Instance
1
104
5.0
66.3
Patch
5
58.9
1.5
Patch
5
35.7
0.557
Table 3: Ablation study.
Figure 7: UNSB output variation.
Stochasticity analysis. SB is stochastic and should return diverse outputs. In Figure 7 we see meaning- ful variation in the generated images. Hence, UNSB indeed learns a stochastic map.
200-
200
150-
100
Transport cost analysis. SB solves OT, so inputs and outputs must be close. In Figure we compare input-output distance for UNSB pairs and SK pairs computed between actual dataset images. we observe that input-output distance for UNSB is smaller than those of SK. This implies UNSB successfully generalizes beyond observed points to generate pairs which are close under L2 norm and are realistic.
102
525
5207
12C
Figure 8: Inputt distance.
CONCLUSION
In this work, we proposed Unpaired Neural Schrödinger Bridge (UNSB) which solves the Schrödinger Bridge problem (SBP) via adversarial learning. UNSB formulation of SB allowed us to combine SB with GAN training techniques for unpaired amaae-to-imaage translation. We demonstrated the scalability and effectiveness of UNSB through various data-to-data or image-to-image translation tasks. In particular, while all previous methods for SB or OT fail, UNSB achieved results that often surpass those of one-step models. Overall, our work opens up a previously unexplored research direction for applying diffusion models to unpaired image translation.

---

Published as a conference paper at ICLR 2024
ETHICS AND REPRODUCIBILITY STATEMENTS
Ethics statement. UNSB extends diffusion to translation between two arbitrary distributions, allowing us to explore wider range of applications. In particular, UNSB may be used in areas with beneficial impacts, such as medical image restoration. However, UNSB may also be used to create malicious content such as fake news, and this must be prevented through proper regulation.
Reproducibility statement. Pseudo-codes and hyper-parameters are described in the main paper and the Appendix.
ACKNOWLEDGMENTS
This research was supported by the National Research Foundation of Korea (NRF) (RS- 202300262527), Field-oriented Technology Development Project for Customs Administration funded by the Korean government (the Ministry of Science & ICT and the Korea Customs Service) through the National Research Foundation (NRF) of Korea under Grant NRF202 11097910 & NRF2021 M311 A 1097938, Korea Medical Device Development Fund grant funded by the Korea government (the Ministry of Science and ICT, the Ministry of Trade, Industry, and Energy, the Ministry of Health & Welfare, the Ministry of Food and Drug Safety) (Project Number: 1711 1137899, KMDF PR 20200901 0015), and Culture, Sports, and Tourism R&D Program through the Korea Creative Content Agency grant funded by the Ministry of Culture, Sports and Tourism in 2023.
REFERENCES
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and R Devon Hjelm. Mutual information neural estimation. In ICML, 2018.
Sagie Benaim and Lior Wolf. One-sided unsupervised domain mapping. In I. Guyon, U. V. Luxburg, S. Bengio. H. Wallach, R. Fergus, S. Vishwanathan, and R. Gamett (eds.), Advances in Neural Information Prucessing Systems, volume 30. Curran Asso- ciates, Inc., 2017. URL -appeeeebberrdbbbeepaaep...e.p. 995500100522022e3332223333330333
Valentin De Bortoli, James Thornton, Jeremy Heng, and Amaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. NeurIPS, 2021.
Charlotte Bunne, Ya- Et-Ping Hsich, Marco Cuturi, and Andreas Krause. The schrödinger bridge between gaussian measures has a closed form. In AISTATS, 2023.
Kenneth F Caluya and Abhishek Halder. Reflected schrödinger bridge: Density control with path constraints. In 2021 American Control Conference (ACC), pp. 1137-1142. IEEE, 2021.
Runfa Chen, Wenbing Huang, Binghui Huang, Fuchun Sun, and Bin Fang. Reusing discriminators for encoding: Towards unsupervised image-to-imas tomge translation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 8168-8177, 2020.
Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood training of schrödinger bridge using barkwrd sdes theory. In ICLR, 2022.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Stochastic control liaisons: Richard 2021 sinkhorn meets gaspard monge on schrodinger bridge. Siam Review, 63(2):249-313, 2021.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung- Woo Ha. Stargan v2: Diverse image synthesis for multiple domains. In CVPR, 2022.
Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In ICLR, 2023.
Mauricio Delbracio and Peyman Milanfar. Inversion by direct iteration: An alternative to denoising diffusion for image restoration. arxiv preprint arXiv:2303. 1435, 2023.
10

---

Published as a conference paper at ICLR 2024
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components estimation. arxiv preprint arXiv:1410.8516, 2015.
Huan Fu, Mingming Gong, Chaohui Wang, Kayhan Batmanghelich, Kun Zhang, and Dacheng Tao. Deeomerry-consisee generative adversarial networks for one-sided unsupervised domain mapping. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2019.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In NeurIPS, 2014.
Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry Vetrov, and Evgeny Burnaev. Entropic neural optimal transport via diffusion processes. In NeurlPS, 2023a.
Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry Vetrov, and Evgeny Burnaev. Entropic neural optimal transport via diffusion processes. In NeurlPS, 2023b.
Nikita Gushchin, Alexander Kolesov, Petr Mokrov, Polina Karpikova, Andrey Spiridonov, Evgeny Burnaev, and Alexander Korotin. Building the bridge of schrödinger: A continuous entropic optimal transport benchmark. In NeurlPS Track on Dulusets and Benchmarks, 2023c.
Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfir Aberman, Yael Pritch, and Daniel Cohen-Or. Prompt- to-prompt image editing with cross attention control. arXiv preprint arXiv:228.001626, 2022.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In NeurIPS, 2020.
Xun Huang, Ming- Yu Liu, Serge Belongie, and Jan Kautz. Multimodal unsupervised image-to-iman to--iaaee translation. In Proceedings of the European Conference on Computer Vision (ECCV), September 2018.
Phillip Isola, Lun-Yan Zhu, Tinghui Zhou, and Alexei A Efros. Image-to-image translation with conditional adversarial networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1125-1134, 2017.
Chanyong Jung, Gihyun Kwon, and Jong Chul Ye. Exploring patch-wis wise semantic relation for con- trastive learning in image-to-image translation tasks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp.
June 2022.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of GANs for improved quality, stability, and variation. In ICLR, 2018.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In ICLR, 2014.
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. In ICLR, 2023.
Christian Léonard. A survey of the schrödinger problem and some of its connections with optimal transport. arXiv preprint arX 2013.
Jae Hyun Lim, Aaron Courville, Christopher Pal, and Chin- Chin-Wei/l Huang. Ar-dae: Towards unbiased neural entropy estimation. In JCML, 2020.
Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos A Theodorou. Deep generalized schrödinger bridge. arXiv preprint arXiv:2209.09880 2022.
Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A. Theodorou, Weili Nie. and Anima Anandkumar. 12sb: Image-to-image schrödinger bridge. arxiv preprint arXiv:230..0587. 2023.
Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun- Yan Zhu, and Stefano Ermon. Sdedit: Guided image synthesis and editing with stochastic differential equations. In ICLR, 2022.
11

---

Published as a conference paper at ICLR 2024
Taesung Park, Alexei A. Efros, Richard Zhang, and Jun- Yan Zhu. Contrastive learning for unpaired image-to-image translation. Andrea Vedaldi, Horst Bischof, Thomas Brox, and Jan-Mion -inchel Frahm (eds.), Computer Vision ECCV 2020, pp- 319-345, Cham, 2020. Springer International Publishing. ISBN 978-5-030-58545-7.
Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enin -nnngc,,E Brandon Amos, Yaron Lipman, and Ricky T. Q. Chen. Multisample flow matching: Straightening flows with minibatch couplings. In ICML, 2023.
Paolo Dai Pra. A stochastic control appraoch to reciprocal diffusion processes. Applied Mathematics and Optimization, 23(1):313 329, 1991.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björ Ommer. High- resolution image synthesis with latent diffusion models. In CVPR, 2022.
Hamza Ruzayqat, Alexandros Beskos, Dan Crisan, Ajay Jasra, and Nikolas Kantas. Unbiased estimation using class of diffusion processes. Journal of Computational Physics, 472: 1 1643, 2023.
Antoine Salmona, Agnès Desolneux, Julie Delon, and Valentin De Bortoli. Can push-forward generative models fit multimodal distributions? In NeurIPS, 2022.
Erwin Schrödinger. Sur la théorie relativiste de l'électron I' I'Interprétation de la mécanique quantique. Annales de l'institut Henri Poincaré, 2(4))269-311 1932.
Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. Conditional simulation
using diffusion schrödinger bridges. In Uncertainty in Artificial Intelligence, pp. 1792-1802. PMLR, 2022.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. In NeurlPS, 2023.
Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learing using nonequilibrium thermodynamics. In /CML, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In ICLR, 2021a.
Yang Song, Jascha Dick--iccstenn, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. score-based generative modeling through stochastic differential equations. In ICLR, 2021b.
Xuan Su. Jiaming Song, Chenlin Meng, and Stefano Ermon. Dual diffusion implicit bridges for image-to-image translation. In ICLR, 2023.
Ella Tamir, Martin Trapp, and Amo Solin. Transport with support: Data-conditional diffusion bridges. arXiv preprint arXiv:2301.1363: 2023.
James Thornton, Michael Hutchinson, Emile Mathieu, Valentin De Bortoli, Yee Whye Teh, and Arnaud Doucet. Riemannian diffusion schrödinger bridge. arXiv preprint 0rXXiv2207.030024 2022.
Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Conditional flow matching: Simulation-free dynamic optimal transport. arxiv preprint arXiv: 2023.
Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving schrödinger bridges via maximum likelihood. Entropy, 23(9): 134, 2021.
Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via schrödinger bridge. In International Conference Machine Learning, pp. 10794-10804, PMLR. 2021a.
Weilun Wang, Wengang Zhou, Jianmin Bao, Dong Chen, and Houqiang Li. instance-wise hard negative example generation for contrastive learning in unpaired image-to-image translation. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 14020- 14029, October 2021b.
12

---

Published as a conference paper at ICLR 2024
Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. ICLR, 2022.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: a stochastic control approach for sampling. arXiv preprint arXiv:2111.15144 2021.
Min Zhao, Fan Bao, Chongxuan Li, and Jun Zhu. Egsde: Unpaired image-to-imane translation via energy-guide stochastic differential equations. In NeurlPS, 2022.
Chuanxia Zheng, Tat-Jen Cham, and Jianfei Cai. The spatially-correhae loss for various image translation tasks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 16407-16417, June 2021.
Yan Zhu, Taesung Park, Phillip Isola, and Alexei A. Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), Oct 2017.
13

---

Published as a conference paper at ICLR 2024
A PROOFS
Lemma 1 (Self-similarity). Let [aa to] _ [0, 1] and {aa} ~ QBB. The SB restricted to [&a to]. defined as the distribution {xt) [t.dd]] = (Z)()) s E [0, 1]} where t(s) == to + (ts to)s solves
min DKL (QI (QI s.t. Qo Qia QEP(R)
SB
Qiay
(18)
Proof. We claim up restricted to the interval [La to] is also SB from Q Q Suppose the claim false, there is another drift út [l, [b] such that
dde - -t dl + VT dw
[ell dt
SB u
and
@r ~ Qu.. It. ~ Qi
We extend ú; to the entire interval [0, 1] by defining
u
if S to:
i =
út
if S t < to
u
sB
th S t < 1.
We then have
%
dx = ún dl + /T dwi,
E
|| L. 112 dt
I 112
and
2o ~ To, 21 TT,
which contradicts our assumption that ui solves Eq.
We note that by change of time variable,
dx; = u dt + V dwi.
ta stsssa to
equivalent the SDE
dæ. (ts to Ju
SB ds + V T(t - ta) dw., 0 0 1
By comparing the stochastic control formulation Eq. with the original SBP Eq. @ we see that this means the reference Wiener measure for the SBP restricted to [ta, to] has variance Tto to). O
Lemma (Static formulation of restricted SBs). Let t [ta, tol S [0, 1] and {a.} Q.B Then p(x Rt,)) [(tx,sstt) + (1- (())aat) s(t) s(t)) t.)
(19)
where s(t) = - (t)/a)/(to - ta) is the inverse function of t(s). Moreover
arg min E( Pt QE.
x1, 11-] - 2T - -ta)
(20)
Proof. Translate the SBP Eq. (18 into the static formulation, taking into account the reduced variance Ttts - -ta) of the reference Weiner measure (7(tt-fa).
Proof of Theurem 1. Let 4, = 4; and 4, in Lemma 2 If the class of distributions expressed by 1o, (0) (,,) is sufficiently large, for any II( Q1).
d7(2,,, = 4o. (It,, Φ1)
(21)
for some o; which satisfies
DKL (go, (z) (l)))(p()))))
(22)
Thus, Eq. 20 and Eq. 2 under the constraint Eq. (10) are equivalent optimization problems, which yield the same solutions, namely p(æt, a.). Then, with optimal Φi.
go. (#1 (t, )p(x(,) = 4o, (w x1) = p(x,, x1).
(23)
Together with the observation that Eq. (11) identical to Eq. 19 Eq. (23 implies Eq.
14

---

Published as a conference paper at ICLR 2024
B OMITTED EXPERIMENT DETAILS
Training, All experiments are conducted on a single RTX3090 GPU. On each dataset, we train our UNSB network for 400 epochs with batch size and Adam optimizer with 31 0.5, B2 0.999, and initial learning rate 0.0002. Learning rate is decayed linearly to zero until the end of training. All images are resized into 256 x 256 and normalized into range [-1, 1]. For SB training and simulation, we discretize the unit interval [0, 1] into 5 intervals with uniform spacing. We used λSB = AReg = 1 and r 0.01. To estimate the entropy loss, we used mutual information neural estimation method Belghazi et al. ()) To incorporate timestep embedding and stochastic conditioning into out UNSB network, we used positional embedding and AdaIN layers, respectively, following the implementation of DDGAN Xiao et al. (2022). For I2I tasks, used CUT loss regularization. On Srrmmer2Wintee translation task, we used a pre-trained VGGI6 network as our feature selection source, following the strategy in the previous work by Zheng et al. (2021).
Entropy approximation. We first describe how to estimate the entropy of a general random variable. We observe that for a random variable X,
T(X, X) H(X) - - H(()
(24)
so mutual information may be used to estimate the entropy of a random variable. Works such as Belghazi et (2018) and[Lim et al. (2020) provide methods to estimate mutual information of two random variables by optimizing a neural network. For instance, Belghazi et al. (2018) tells us that
le(X, Z) == sup Ez x [To] - log (Epx e
(25)
where T' neural network parametrized by 0 < ④ is capable of approximating [XX, 2) up to arbitrary accuracy. Hence, we can estimate H (X) (XX, X)) by setting X = Z and optimizing Eq. (25). In the UNSB objective Eq. (9). we use Eq. 25 estimate the entropy of X
(x Lt,)) where (It,, 21) qa. (It,: x1 for qd. defined in Eq. (). In practice, we use the same neural net architecture for Ta and the discriminator. We use the same network for Tg for each time-step t, by incorporating time embedding into Ta. Ta is updated once with gradient ascent every qá; update, analogous to adversarial training.
Evaluation.
Metric calculation codes
FID KID
httt:://ggttab.cco//aaatttttttrt GAN MANN-pytorch
https:
Evaluation protocol: we follow standard procedure, as described in Park et al. (2020).
KID measurement
DistanceGAN, MUNIT, NOT, SBCFM trained from scratch using official code.
Baselines excluding above four methods KID measured using samples generated by models from official repository.
FID measurement
Horse2Zebra and Label2Cityscape: numbers taken from Table 1 in Park et al. (2020).
Summer2Winter and Map2Satellite: all basclines trained using official code.
15

---

Published as a conference paper at ICLR 2024
Input
Ours
CFM
Method
FID |
KIDJ
SBCFM
229.5
18.8
Ours
35.7
0.587
Table 4: Quantitative comparison result with another SB-based method on Horse2Zebra. The baseline model performance is severely degraded, while our generated results show superior perceptual quality.
Figure Visualization of our generated result and the result from conditional flow matching (CFM).
Figure 10: UNSB output pixel-wise standard deviation given input.
EGSDE
StrrGAN22
Method
Time
NFF
FIDI
KIDJ
EGSDE
60
500
45.2
3.62
StarGAN v2
0.006
48.55
3,20
0.006
Ours
0.045
37.87
1.54
Table 5: Male2Female results.
Figure 11: Fale2Female samples.
C ADDITIONAL EXPERIMENTS
C.
OTHER SB METHONS
In this part, we compare our proposed UNSB with other Schrödinger bridge based image translation methods. The recent work of ISSB and InDI require supervised training setting, therefore we cannot compare our proposed method with the baselines. For most of other SB-based methods, they focus on relatively casy problem, where one side of distribution is simple. Since we firstly proposed applying SB-based approach to high-resolution 12I problem, there is no direct baseline to compare. However, compared model with SB -B-sased approach of Schrödinger Bridge Conditional Flow Matching (SBCFM) (Shi et al. 2022) by adapting the SBCFM to high-resolution case.
that the baseline SBCFM model could not generate the proper outputs where the outputs are totally unrelated to the input image. The results show that the model could overcome the curse of dimensionality problem.
In Table 4 and Fig.
16

---

Published as a conference paper at ICLR 2024
C.2 MORE STOCHASTICITY ANALYSIS
In Figure 10 we visualize output pixel-wise standard deviation given input for more samples. We note that all pixel values are normalized into [0, 1]. A SB between two domains is stochastic, so does one-to-many translation. Thus, a necessary condition for a model to be SB is stochasticity. Figure shows that UNSB satisfies this necessary condition.
Greater std for foreground pixels tells us UNSB does one-to-many generation, and generated images lie in the target domain. For instance, model which simply adds Gaussian noise also does one-to- many generation, but is not meaningful. UNSB output shows high variation for locations which are relevant to the target domain (Zebra).
C.3 MALE2FEMALE TRANSLATION
We provide results the Fele2 translation task with the CelebA--Q--2 dataset (Karras eet al. 2018). Baseline methods are EGSDE (Zhao et al. (022), StarGAN v2 (Choi et al. 2022] and NOT. UNSB is trained according to Appendix B We used official pre-trained models for EGSDE and StarGAN v2 and trained NOT from scratch using the official code. We use the evaluation protocol
in CUT (Park al. 2020) to evaluate all models, i.e., we use test set images to measure FID. We note that IZhao et al. 2022) and (Choi et al. (2022) use different evaluation protocol in their papers, (they use train set images to measure FID) so the numbers in our paper and their papers may differ.
Results are shown in Table 3 and Figure I We note that UNSB produces outputs closer to the target image domain, as evidenced by FID and KID. Also, EGSDE, StarGAN v2, and NOT often fail to preserve input structure information. For instance, EGSDE fails to preserve eye direction or hair style. StarGAN v2 hallucinates earrings and fails to preserve hair style. NOT changes mouth shape (teeth are showing for outputs).
C.4 ENTROPIC OT BENCHMARK
We evaluate UNSB on the entropic OT benchmark pro- vided by Gushchin al. (2023c) for 64 64 resolution images. We use e = 0.1 (using the notation in (Gushchin 2023c)) and follow their training and evaluation protocol. Specifically, we trained UNSB to translate noisy CelebA images (i.e., y-samples) to clean CelebA images (i.e., a--saplle)). For comparison, we provide cFID num- bers for SB p(æ) y) learned by ENOT (Gushchin et al.
Method
ENOT
Ours
p(æ)
cHID
40.5
63.0
104
Table 6: EOT benchmark.
2023b) and unconditional p( p = p(æ) where we use 5k ground-truth samples from p(æ). cFID for ENOT is taken from Table in Gushchin et al. (2023c).
In Table 6 we see that UNSB performs much better than the unconditional map p(æly) = p(æ) but not quite as well as ENOT. We emphasize that this is only a preliminary result with ad hoc hyper-parameter choices due to time constraints, and we expect better results with further tuning. Some points of improvement are:
Using more samples to measure cFID. We only used 100 æ-samples for each y-sample. The benchmark provides 5k sa-mapllss for each y-sample, using more w-samples per y y-sample could improve the cFID.
Better architecture design. UNSB architecture is specialized to 2 256 x 256 resolution images, whereas the benchmark consists of 64 x 64 resolution images. Currently, we set UNSB networks for this task as resize(ær, 64) G(xx) o resize(ær, 256), where resize(ær, 8) is function which resizes images x resolution, and G(x)) the UNSB network used in our main experiments.
Better regularization. CUT regularization is specialized to translation between clean images, whereas the benchmark consists of translation between noisy and clean images.
Longer training. We only trained UNSB for 150k iterations. On, for instance, Felale, we needed to train about 1.7m iterations to get good results.
17

---

Published as a conference paper at ICLR 2024
Sourve
Ours
P2P
SDEdii
NOT
CUT
CydGGAN
GeGAN
DisheeGANN
Figure 13: Additional samples on Horse2Zebra and Summer2Winter.
C.5 ZEBRA2HORSE TRANSLATION
We show that UNSB is capable of reverse translation on tasks in our paper as well. Specifically, we com- pare CycleGAN and UNSB on Zebra2Horse trans- lation task. In Table 7 we see that UNSB beats CycleGAN in terms of both FID and KID. In Figure 12 we observe UNSB successfully preserves input structure as well.
CydeGAN
C.6 MORE ABLATION STUDY
Figure 12: Zebra2Horse samples.
In Table 8 we provide another ablation study result on the Summer2Winter translation task, where we observe similar trends as before. Adding Markovian discrminator, regularization, and multi-step sampling monotonically increases the performance. This tells us the components of UNSB play orthogonal roles in unpaired image-to-image translation.
Method
Time
NFE
FID
CycleGAN
0.004
135.75
2.998
Ours
0.045
134,09
2.581
Table Zebra2Horse results.
Dise.
Reg.
NFF
FID
KID
C.7 ADDITIONAL SAMPLES
X
254.2
23.91
Instance
185.5
8.732
256 256 images. In Fig. 13 generated samples from our proposed UNSB and baseline models for tasks in main paper.
14
show additional
Patch
155.5
7.697
Patch
75.6
0.491
Patch
73.9
0.421
512 x 512 images. In Fig. 15 UNSB for Cat2Dog translation with AFHQ 512 resolution images, NFE = This is preliminary result with UNSB training epoch 100.
we show samples from
Table 8: Summer2Winter ablation.
18

---

Published as a conference paper at ICLR 2024
Souree
Ours
NOT
CUT
CydeGAN
GeGAN
DistanceGAN
Figure 14: Additional samples on Label2Cityscapes and Map2Satellite dataset.
Cat
Dog
Figure 15: Cat2Dog on AFHQ 512 512 with UNSB. Preliminary result at training epoch 100.
19

---

Published as a conference paper at ICLR 2024
D DISCUSSION ON SCENE-LEVEL VS. OBJECT-LEVEL TRANSFER
We note that unpaired image-to-image translation tasks can roughly be categorized into scene- level transfer and object-level transfer. scene-level transfer (e.g., Horse2Zebra, Summer2Winter, etc.), mostly color and texture need to be altered during translation. In object-level transfer (e.g.. Male2Female, Cat2Dog, etc.), object identities need to be altered during translation. We discuss differences between UNSB and diffusion-ben based translation methods on scene-level and object-level transfer tasks, respectively.
While diffusion-based methods such SDEdit are capable of translating between high-resolution images, we note that they do not perform so well on scene-level translation tasks considered in our paper. For instance, see Table 2 for SDEdit and P2P results on Horse2Zebra, where they show 97.3 FID and 60.9 FID, respectively. UNSB achieves 35.7 FID.
Scene-level transfer.
Figure 5a provides some clue as to why diffusion-based methods do not perform well. P2P fails to preserve structural integrity (in the first row, the direction fails to properly translate images zebra domain. Similar phenomenon occurs for Summer2Winter task as well. On the other hand, UNSB does not suffer from this problem as it uses a Markovian discriminator along with CUT regularization.
right horse is reversed) and SDEdit
Object-level transfer. Diffusion-based methods perform reasonably on object-level translation tasks such Male2Female or Wild2Cat. We speculate this is because object-level translation tasks use images with cropped and centered subjects, so it relatively easy to preserve structural integrity. In fact, shown in Figure T6 semantically meaningful pairs on such dataset pairs be obtained even with Sinkhorn-Knopp (a computational optimal transport method), which does not rely on deep learning. However, Sinkhorn- Knopp is unable find meaningful pairs on Horse2Zebra, where subjects are not cropped and centered.
Wild Animal
Cat
Horse
Zebra
(a) Wild2Cat
(b) Horse2Zebra
Figure 16: Unpaired mmage-to-image translation results with Sinkhorn- Knopp.
20


Reviews:
[{"review_id": "hKYPteqqvC", "reviewer": "Reviewer_ek8v", "summary": "This paper aims to address the problem that the diffusion model based unpaired image translation tasks rely on the Gaussion prior assumption. It resorts to Schrödinger Bridge and proposes the Unpaired Neural Schrödinger Bridge (UNSB) to express SB problem as a sequence of adversarial learning problems. By this formulation, SB can regarded as Lagrangian formulation under the constraint on the KL divergence between the true target distribution and the model distribution. This formulation leads to the composition of generators learned via adversarial learning that overcome the curse of dimensionality with advanced discriminators. Experiments demonstrated the effectiveness of the proposed method compared to its competitors.", "questions": "no", "limitations": "", "rating": 8, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper experimentally identifies the cause behind the failure of previous SB methods for image-to-image translation\nas the curse of dimensionality, by a toy task.\n\nThe proposed method gives a new formulation of SB, which benefits on addressing the identified curse of dimensionality.\n\nThe proposed method shows better results than existing methods.", "weaknesses": "The paper claimed that \"none of SB models so far have been successful at unpaired translation between highresolution images\", however, the experiments in this paper do not include highresolution images yet (256x256 can not be regarded as highresolution). How will the proposed method perform on this setting is still unclear.\n\nIt is unclear that the findings by ablation study are still valid for other datasets?\n\nWhen presenting the stochasticity analysis, it only gives results for one image. It would be better if the paper can provide statistical numbers across many different images/cases.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1kS7igqarM", "reviewer": "Reviewer_f3t2", "summary": "The paper proposes a method for unpaired image-to-image translation based on Schrödinger bridge. The authors combine several ideas to obtain a practically-sounding method. At first, they propose to learn their model in a multistep manner, which is a natural and worth-considering given the recent success of diffusion models. Secondly, they propose to use the additional GAN loss to better match the target distribution. Additionally, the authors report the usefulness of the additional regularization applied at the training stage. Overall, the proposed approach demonstrates good practical results and has clear theoretical motivation.\n\n**Update**: My concerns seems to be addressed. I rise my score to 6", "questions": "- I have some doubts regarding the statement (end of page 1) that “no work has successfully trained SBs for direct translation between high-res images” and (end of page 2) “our work represents the first endeavor on this problem”. For example, [1], [2] deals with $64 \\times 64$ images and tackles unpaired image-to-image problems. I am not sure, if the methods from [1], [2] works with $256 \\times 256$, but, anyway, [1] and [2] should be covered as the related works and the competitive methods.\n- Regarding the competitive methods (GANs), the authors miss the StarGAN and StarGAN v2 [3]. The latter model demonstrates a good quality on the unpaired image2image and I think this model is a good baseline method. Please, add it to the comparison.\n- I have one concern about objective (14). The transition from constrained problem (9)-(10) to the combination of losses (14) is non-trivial and not fully theoretically justified. I am not sure, that the objective (14) indeed solves the SB problem. Some comments/clarifications will contribute to the clearness of the manuscript flow.\n- I ask the authors to add the clarifications how to estimate the entropy with samples, a technique introduced at the beginning of page 6. It will facilitate the comprehension and smoothness of the reading. \n- I strongly encourage the authors to evaluate how does their approach solve SB problem on a nontrivial setup beyond gaussian2gaussian case. In particular, there is a recent paper [4] which proposes benchmark pairs of distributions with known GT EOT/SB solutions, even for images data. The validation on such a benchmark will highly contribute to the quality of the paper and will probably answer the question (3).\n- I have one more question regarding the baseline methods. The authors choose the work [5] (called them as NOT) “as the representative of OT [methods]”. This choice really surprised me because [5] also deals with SB problem, and, as I understand, their method is not designed for image data (at least, the authors of [5] does not consider image data use cases). That is why comparing with [5] seems strange. I recommend authors take [6, 7, 8] as the representative of OT methods instead. These methods also has a possibility to generate the samples of controllable diversity, have GAN-resembling objective. Moreover, they are designed specifically for unpaired image2image problems.\n\n[1] Diffusion Schrödinger Bridge Matching, NeurIPS’2023\n\n[2] Entropic Neural Optimal Transport via Diffusion Processes, NeurIPS’2023\n\n[3] StarGAN v2: Diverse Image Synthesis for Multiple Domains, CVPR’2020\n\n[4] Building the Bridge of Schrödinger: A Continuous Entropic Optimal Transport Benchmark, NeurIPS’2023\n\n[5] Transport with support: Data-conditional diffusion bridges, arxiv\n\n[6] Neural Optimal Transport, ICLR’2023\n\n[7] Neural Monge Map estimation and its applications, TMLR\n\n[8] Kernel Neural Optimal Transport, ICLR’2023", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is overall-well written. It has clear motivation (unpaired image-to-image translation via entropic optimal transport (or Schrödinger bridge) problems. The idea of multistep approach, where the timeline of the Schrödinger process is discretized into several steps and a practitioner learn a kind of “denoising” model (an intuition from diffusion models) is fresh and worth to be considering. Moreover, it seems that the considering several steps indeed improves the qualitative performance, as shown by the practical experiments. The practical results are also encouraging.", "weaknesses": "My concerns are covered in the questions section (see below). I expect the authors to address them.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bSbYJlpTeT", "reviewer": "Reviewer_PUqf", "summary": "The paper proposes a method leveraging Schrödinger Bridge (SB) to bypass the limitation of applying Diffusion Models (which assume a Gaussian prior) to Image-to-Image (I2I) translation. Although previous approaches for utilizing SB to I2I translation have been explored, the proposed method, UNSB, is designed to be directly applied to I2I and scalable to high resolution images. The method is implemented as a composition of generators learned via adversarial learning that overcome the curse of dimensionality with advanced discriminators, showing success in several unpaired I2I translation tasks.", "questions": "1. The description of the Stochasticity analysis (Fig. 7) is not clear to me. Could the author describe this ablation experiment and what does the generated map tell us? I understand that high pixel-wise std implies greater diversity, but why is it surprising that the std is greater for the foreground pixels? What is the scale of the std?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The formulation of SB as a composition of generators learned via adversarial learning seems novel. Moreover, the ability to apply patch-level discriminators is crucial for data with diverse textures as the ones considered in the paper.\n2. The translation quality and faithfulness seem to be good.\n3. The paper is well written and easy to follow.", "weaknesses": "1. I think that some relevant baselines were not included in the experimental section. Could the author compare the proposed method to EGSDE [1] (Diffusion-based model)?\n\n[1] Zhao et al. EGSDE: Unpaired Image-to-Image Translation via Energy-Guided Stochastic Differential Equations. In NeurIPS, 2022.\n\n2. The authors attribute the failure of SB for unpaired image-to-image translation to the curse of dimensionality. However, methods such as EGSDE and SDEdit do achieve successful translations in translating high resolution images. Could the authors elaborate on the difference between the cases?\n\n3. Following my previous point, the datasets considered in the paper are all of “scene-level\" image to image translation (in which mostly color and texture need to be altered during translation). In most of other I2I translation works, many \"object-level\" datasets are evaluated e.g. CelebA-HQ Male-to-Female, Cat-to-dog, etc. Do the authors expect different performance for the latter case? I believe it would be helpful to include benchmarks for at least some of these additional datasets to understand the strengths and the limitations of the proposed method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UL7fA6GPsx", "reviewer": "Reviewer_279L", "summary": "The paper presents Unpaired Neural Schrödinger Bridge (UNSB), an innovative approach for unpaired image-to-image translation tasks, combining Schrödinger Bridge model with adversarial training. It expands the applicability of diffusion models by overcoming the Gaussian prior assumption limitation, which typically restricts them in unpaired image-to-image translations.The experiment results signal the potential of the proposed approach in dealing with high dimensional datasets, commonly referred to as the \"curse of dimensionality.\" Experiments have shown the superiority of the proposed method where comparisons are conducted with their method and multiple baseline I2I methods on multiple I2I datasets. An ablation study is conducted with clear analyses.", "questions": "See the weakness section.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1 The paper is aligned with recent advancements in stochastic generative modeling, image-to-image translations, and optimal transport, pushing the boundary of what's possible in these domains.\n\n2 The experimentation shows promising results in I2I tasks in terms of multiple metrics.\n\n3 The paper clearly explains the underlying mathematical conventions for Schrödinger Bridge and UNSB models and their connection with stochastic differential equations.\n\n4 The authors have taken additional steps to ensure reproducibility of their work by providing code and detailed descriptions of hyperparameters in the paper.", "weaknesses": "There are no significant weaknesses. It could be valuable to include the reverse translations (such as Zebra to Horse, Winter to Summer) in the supplementary material, providing more comprehensive experimental results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)