Improving Neural Optimal Transport via Displacement Interpolation
Authors: Jaemoo Choi
Yongxin Chen
Jaewoong Choi
Venue: ICLR 2025 Poster (2025)
Keywords: Optimal Transport
Displacement Interpolation
Image-to-image Translation
Primary area: generative models
TL;DR: This paper proposes new algorithm to improve the neural optimal transport using displacement interpolation, achieving superior results in image-to-image translation tasks.

Abstract: Optimal Transport (OT) theory investigates the cost-minimizing transport map that moves a source distribution to a target distribution. Recently, several approaches have emerged for learning the optimal transport map for a given cost function using neural networks. We refer to these approaches as the OT Map. OT Map provides a powerful tool for diverse machine learning tasks, such as generative modeling and unpaired image-to-image translation. However, existing methods that utilize max-min optimization often experience training instability and sensitivity to hyperparameters. In this paper, we propose a novel method to improve stability and achieve a better approximation of the OT Map by exploiting displacement interpolation, dubbed Displacement Interpolation Optimal Transport Model (DIOTM). We derive the dual formulation of displacement interpolation at specific time $t$ and prove how these dual problems are related across time. This result allows us to utilize the entire trajectory of displacement interpolation in learning the OT Map. Our method improves the training stability and achieves superior results in estimating optimal transport maps. We demonstrate that DIOTM outperforms existing OT-based models on image-to-image translation tasks.

Published as a conference paper at ICLR 2025
IMPROVING NEURAL OPTIMAL TRANSPORT VIA DIS- PLACEMENT INTERPOLATION
Jaemoo Choi
Yongxin Chen
Georgia Institute of Technology jchoi 8438 gatech. edu
Georgia Institute of Technology ychen31 1 4 @8@gatech.ede edu
Jaewoong Choi"
Sungkyunkwan University jaewoongchoi@@skku..
ABSTRACT
Optimal Transport (OT) theory investigates the cost-minimizing transport map that moves a source distribution to target distribution. Recently, several approaches have emerged for learning the optimal transport map for given cost function using neural networks. We refer to these approaches the OT Map. OT Map provides a powerful tool for diverse machine learning tasks, such as generative modeling and umpaired image-to-image translation. However, existing methods that utilize max-min optimization often experience training instability and sensitivity to hyperparameters. In this paper, we propose novel method to improve stability and achieve better approximation of the OT Map by exploiting displacement inter- polation, dubbed Displacement Interpolation Optimal Transport Model (DIOTM). We derive the dual formulation of displacement interpolation at specific time t and prove how these dual problems are related across time. This result allows us to utilize the entire trajectory of displacement interpolation in learning the OT Map. Our method improves the training stability and achieves superior results in estimating optimal transport maps. We demonstrate that DIOTM outperforms existing OT-based models on image-to-image translation tasks.
INTRODUCTION
Optimal Transport (OT) problem (Villani et al. 2009 Peyré et al. 2019) explores the problem of finding the cost-optimal transport map that transforms one probability distribution (source distribution) into another (target distribution). Recently, there has been growing interest in directly learning the optimal transport map using neural networks. Throughout this paper, we call these approaches as the OT Map. OT Map has been widely applied across various machine learning tasks by appropriately defining the source and target distributions, such as generative modeling (Rout et al. 2022 Choi et al.
2023 2024a Liu al. 2022 [Limman et al. 2023), image-to-image translation (Korotin al. [023 Fan et al. (022), and domain adaptation (Flamary et al. 2016). OT Map is particularly well-suited for unsupervised (unpaired) distribution transport problems, it enables the transport of one distribution into another using only a predefined cost function, without requiring data pairs.
Despite their potential, existing OT Map methods often encounter significant challenges in training stability. In particular, the OT models utilizing max- max-min objectives exhibit unstable training dynamics and sensitivity to hyperparameters (Makkuva al. 2020 Fan et al. 2022: Rout et al.) 2022 Flamary et al. 2016). These challenges limit their applicability to high digh-dimensional data. To address these instability issues, various approaches have been explored, such as introducing additional regularization
terms in the learning objective (Rout et al. 2022 Roth et al.] (017) and generalizing the standard OT problem to the unbalanced optimal transport problem (Choi et al. 2023: 2024b).
In this paper, we propose novel approach for learning the optimal transport map based on displace- ment interpolation. We refer to our model as the Displacement Interpolation Optimal Transport
Corresponding Author

---

Published as a conference paper at ICLR 2025
Model (DIOTM). We identify fundamental connection between the (static) optimal transport map and displacement interpolation. Motivated by this relationship. we derive a max-mii formulation of displacement interpolation, involving the optimal transport map and the time-dependent Kantorovich potential. Our experimental results demonstrate that DIOTM achieves more stable convergence and superior accuracy in approximating OT maps compared existing methods. In particular, DIOTM achieves competitive FID scores in image-to-image translation tasks, such as 5.27 for Malc-+Female (64 64), 7.40 for Male-Feemaee (128 128), and 10.72 for Wild-Cat (64 x 64), comparable to the state-of-ther the-art results. Our contributions can be summarized as follows:
We propose method to learn the optimal transport map based displacement interpolation.
We derive the dual formulation of displacement interpolation and utilize this formulate max-min optimization problem for the transport map and potential function.
We introduce a novel regularizer, called the HJB regularizer, derived from the optimality condition of the potential function.
Our model significantly improves the training stability and accuracy of existing OT Map models that leverage min-max objectives.
Notations and Assumptions Let X' be connected bounded convex open subspace of Rd and let and v be absolutely continuous probability distributions with respect to Lebesgue measure. We regard and the source and target distributions. For measurable map T, Th 14 represents the pushforward distribution of μr. I1(,, v) denote the set of joint probability distributions on X' x whose marginals are μ and respectively. c(2,, y) refers the transport cost function defined X' y. Throughout this paper, we consider X' = y C Rd with the quadratic cost, c(x, = alI - 2 where d indicates the dimension of data. Here, a given positive constant. Moreover, we denote W2(.,)) as the 2- Wassersteii distance of two distributions.
2 BACKGROUND
In this section, provide brief overview of Optimal Transport theory. These results, especially the dual formulation and displacement interpolation, will play a crucial role in our proposed method.
Optimal Transport The Optimal Transport (OT) problem investigates the optimal way to transport the source distribution μ to the target distribution y (iillani et al.] 2009). The optimality of the transport plan is defined as the minimization of given cost function. Initially, Monge (1781) formulated the OT problem with a deterministic transport map T where T V. However, the Monge OT problem is non-convex and the optimal transport map T* may not exist depending on μ and v. To overcome this problem. Kantorovich (1948) introduced the following convex formulation:
[ c(x, (dd
C(u, = inf xEll(a.)) Jxxx
where the minimization is conducted over the set joint probability distribution æ E I[p, v). We refer to this π as the transport plan or coupling between and v. When the optimal transport map T* from the Monge OT exists, the optimal coupling "* satisfies л (Id x ( For general cost function c(-.) -) that is lower semicontinuous and lower bounded, the Kantorovich OT problem (Eq. can be reformulated as follows (Villani al, (2009). Chapter 5):
('()dd(()) + f v(w)dv(u))
C(H,v) v) sup VELi(w)
(2)
where the potential function E (v) is an integrable function with respect to v and the c-transform of V is defined as
V²(x) = inf (c(æ, v) (v))
yey
This formulation (Eq.) is called the semi-dual formulation of OT.
Rout al. 2202 and Fan al. (2022) proposed method for learning the optimal transport map T by leveraging this semi-dual formulation (Eq. [ of OT and applied this T for generative modeling. In generative modeling, the source distribution f and the target distribution correspond to the
2

---

Published as a conference paper at ICLR 2025
Gaussian prior and the target data distribution, respectively. Specifically, these models parametrize the potential = Ve in Eq. 2 and the transport map Tg X + 3 as follows:
Tg + arg min [c(x, y) V (v) vEX
+
(z(2) = c Ta(x)) V (T)())).
Note that the parametrization of Te the left-hand side is equivalent to the representation of V on the right-hand side, by the definition of the c-transform (Eq. 3 Based on this, we arrive at the following optimization problem:
J inf (n(( - V (To(z)) du(x)) Ty
/ Vo(w)drrv))
LV = =sup Ve
(5)
Intuitively, Ta and Va serve similar roles as the generator and the discriminator of GAN (Goodfellww ct al. 2020). A key difference is that Tg trained minimize the cost с (z, To(x)), since Lg learns the optimal transport map. For convenience, the OT-based generative model (OTM) (Fan et al. 2022).
denote the optimization problem LV To (Eq. 5 as
Dynamic Optimal Transport and Displacement Interpolation In this paragraph, we provide a
close connection between the dynamic optimal transport problem and displacement interpolation. While the (static) optimal transport (Eq. (T) focuses solely on how each ~ μ is transported to y the dynamic optimal transport problem tracks the continuous evolution of to v. Formally, the dynamic formulation of the Kantorovich OT problem (Eq. D for the quadratic cost c(x, y)
alla a|ı - y/l can be expressed follows:
Opt JoJ (z)))d x:
inf
+ V. = Po = P1
(6)
Öt
This dynamic formulation of OT (Eq. 6 called the Benamou-Brente formulation (Benamou & Brenierr 2000). Note that the dynamic transport plan {pr Justgu evolves from μ to and this evolution is governed by the ODE dr /t/dt= = uy (x)) through the continuity equation.
Interestingly, the optimal solution of this dynamic OT problem has simple form. Along each ODE trajectory {x(t) dx/dt = (x), (0) = }. the velocity field v remains constant. Moreover, when the deterministic optimal transport map T exists, the following holds:
Pt = [1 (1-t) Id
T"]
and
p
p for Osts
where Td denotes the identity map and p² denotes the optimal dynamic OT plan of Eq. 6 die is called McCann's 's Displacement Interpolation (DI) (MeCann) 1997). Eq. shows that the pushforward of linear interpolation between Id and the static OT map 7 equivalent the dynamic OT plan pt. Hence, throughout this paper, we simply denote the displacement interpolation as p² Furthermore, it is well known that the displacement interpolants satisfy the following property (Theorem 7.21 in
Villani et al. )2009)):
pf = arg
arginf inf (Dr(t, where Cpi(t, p) = t) + (P,
(8)
Note that Lnr corresponds the Wasserstein-2 barycenter problem between the two probability distributions (t, y (Agueh & Carlier 201T Kolesov et al. 2024b). In other words, Eq. represents the equivalence between the displacement interpolants and the Wasserstein barycenter. This equivalence will be utilized in Se derive our approach to neural optimal transport, i.c., learning the optimal transport map T with neural network. We establish how the optimal potential and transport maps for each P are related and use this relationship to improve neural optimal transport.
3 METHOD
In this section, propose our method, called the Displacement Interpolation Optimal Transport Model (DIOTM). Our model leverages displacement interpolation to improve the stable estimation of the optimal transport map using neural networks. In See 3.]] we derive our two main theorems for deriving our learning objective. 3.2 we describe how we implement our DIOTM model based on these theoretical results.
3

---

Published as a conference paper at ICLR 2025
3.1 DUAL FORMULATION OF DI AND THE RELATIONSHIP BETWEEN INTERPOLATION POTENTIAL FUNCTIONS
In this subsection, we provide two theoretical results: the Dual Formulation of Displacement Interpolation (Thm (.T) and the Relationship between Interpolation Potential Functions (Thm 3.3 Thm 3.1 will be utilized derive our main learning objective (Eq. 17 and Thm 3.3 will be employed to formulate our regularizer (Eq. 19p.
Dual Problem of Displacement Interpolant Here, we begin with the dual formulation of the minimization characterization of the interpolant Pi (Eq. [8. The dual problem can be expressed as Eq, 9 in the following Theorem [.T Note that while the primal formulation (Eq. ( optimizes over the set of probability distribution p. the dual formulation optimizes over two potential functions fL,t, f2.1-
Theorem 3.1. Given the assumptions in Appendix A for a given t E (0, 1), the minimization problem
inl, (lr p) (Eq. 8) is equivalent to the following dual problem:
sup EC with t (1 (1- 1) S (.,())du(()+ + I/ 12..w))de(o)
(9)
where the supremum is taken over two continuous potential functions f1.t . R. and f2,1 X - R. which satisfy (1 f).t),i.t tt2, = O. Note that the assumptions in Appendix Aegenddxxxuuaaann guarantee the existence
and uniqueness of displacement interpolants pt, the forward optimal transport map T from μ to pt. and the backward transport map T from to Pi. Based on this, we have the following:
Pi(x)) E arginf [c(x, u) [t), T7(w) p.a.s. for x and v-a.s. for y. By applying the same c-transfrom parametrization in Eq. + derive the following max-min formulation of the dual problem:
E arginf (c(x,v) - f2,(()),
(10)
sup Ea with (1- 1 +1f2 (1 - 1) f inf (c(x,T S1, (t.)))) du(x) 7.
T, + / inf (at -ax du(y). (1I)
Sax(T
Note that the max-min formulation above requires optimization over four functions: the two potentials fi,e, f2,t and the two transport maps Te, T.. Here, provide an intuitive interpretation of this
max-min optimization (Eq. 11). Two transport maps T t. TT. I act as generators of the interpolant p from μ and v, respectively. Two potential functions f1, f2 serve similar roles to discriminators for the generated samples (fake samples) from 7., T The discriminator values for the true samples pt cancels out because of the potential condition (1 - t), + ff2. = O. Formally,
(1 (1-t) f fLe(3) (2)) +t
faa(2) dpt = / [11- t). + tfa.l (z) (p) = 0. (12)
In fact, using this potential condition, we can combine these two potentials into a single value function (or potential function) V,, This simplified formulation of Eq. 9 will be used to derive our regularizer in Theorem 3.3.
Corollary 3.2. For given I 10, let J1,(ly) (V((y) and f2,(()) -(1 (x) for some value function Ve : X = + R. Then, the dual formulation of displacement interpolation (Eq. 9 can be rewritten as follows:
V v. (x))((()())))))))))))))))))))))) + / (-V:)aua (-V))oo(())dv())) Vi
sup
(13)
where (x, = w/1 for every S 1
Relationship between Interpolation Potential Functions Here, based on Corollary 3.2 we derive the optimality condition, that each interpolant value function Vú2 satisfies, the time-dependent value function zt, = Vz(z). From now on. we will denote the value function in its time-dependent form: V(I,a)) (0, 1) X TR.
4

---

Published as a conference paper at ICLR 2025
Algorithm 1 Training algorithm of DIOTM
Require: The source distribution and the target distribution v. Transport networks 7' 01 T 02 and the discriminator network Vá- Total iteration number K, and regularization hyperparameter X.
1: for k 0, 2......... K do
2: Sample a batch x ~ pa v, v.t U,0,0
3: Te - (1 t)x + 0, (I), y1 + (1 - ()) bty
4: Update Ve increase the Ls
Ls = -V + Vo vo) - XR(Vo(t,zz)) AR(Vo(t,
5: Update To, 01 to decrease the loss: co. xt) - Vo((,
6: Update T 02 decrease the loss: (3, v) (o y/).
7:
end for
Theorem 3.3. Given the assumptions in Appendix A the optimal V in Eq. 13 satisfies the following:
V V. (x)dp(x) + J V+(x)dpp(xx)
V srg V.
(14)
up to constant p² -a.... Moreover, there exists (Vi that satisfies Hamilton-lacobi-Belle Belllobbbb-llllan (HJB) equation, i.e.
0,V," + 1 |V, = 0, p² aasa.s
(15)
4a
In Sec 3.2 we use this HJB optimality condition (Eq. [5 regularizer for the value function V in our model. In Sec we demonstrate that, when combined with our displacement interpolation, this regularizer significantly improves the stability of optimal transport map estimation.
3.2 DISPLACEMENT INTERPOLATION OPTIMAL TRANSPORT MODEL
In this subsection, introduce model, called Displacement Interpolation Optimal Transport Model (DIOTM). The goal of our model is to learn the optimal transport map between the source distribution yu. and the target distribution v. DIOTM is trained by utilizing the displacement interpo- lation D (Eq. 7 between these two distributions. In DIOTM, the forward and backward transport maps T, T trained match all intermediate distributions. As result, each transport map directly trained to generate the boundary distributions and v, but instead exploits the matching of intermediate distributions. This approach enables our model achieve a more stable estimation of the optimal transport map.
Parametrization of the DI Dual Problem Our DIOTM learns the static optimal transport map between y and by exploiting the displacement interpolation (DI) pt. which is the solution of the dynamic optimal transport problem (Villani al. 2009). However, there are some challenges when using DI to learn the static transport map. While our goal is to learn the optimal transport maps between i and v, the max-min dual formulation the dual problem (Eq, IT] applies to each specific time / (0, 1). In other words, the intermediate transport maps 7,, T, and the potential V. are defined separately for each Therefore, we represent these interpolant generators T., T through the boundary generators T. T by incorporating the optimality condition of DI (Eq. Specifically, we parametrize the interpolant generators T. and T, follows:
7((x)= - t)x + + (z), F.(0) = (1 t)y + + (v) for 10.
(16)
where T', o, and Te, parametrize the optimal transport maps from the source μ the target and from v to μ., respectively. The optimality condition of DI clarifies how the intermediate optimal transport maps 7; T; related to each other. We utilize this condition parametrize the entire interpolant pf for E (0, 1) using just two networks T To. : X + y and T o₂ y - X. Note that this optimality condition is satisfied under the optimal value function Va is given. Thus, this parametrization is
5

---

Published as a conference paper at ICLR 2025
introduced for better efficiency. Additionally, already investigated how the potentials. specifically the value function (tt, I), are related in Theorem 3.3 Therefore, by combining Eq 13 and 16 we arrive at our main max-min learning objective:
a T T((x)] - Volt, T ((())) du(x))
C
0,6
sup V
inf
E
|| -
X To.
T (1)) y112 + Volt, T ((lil) du(y). (17)
Er inf
Tez
Here, a indicates the cost intensity hyperparameter, i.e., cx.y) a a|I yl12- Note that the expectation with respect to / is for aggregating the interpolant pi for i & (0, 1). When represented with our transport map parametrizations Te, and Te, (Eq. 16), main learning objective 23/ can be expressed follows:
DI = sup / inf Ee |dl|l 7 T Vol,, 7, (x)) du(x) 0,8
Vi To,
in E. [((1 - 1)I F (v) y Voll, T Tx(i)) doly). (18) To
Moreover, we introduce the HJB regularizer R(Va). which is derived from the HJB optimality condition of the value function, proved in Theorem 3.3
R(Vo) = Ee 2a 0 (ot, + vvo(t
(19)
As result, the learning objective can be summarized as follows:
L..0 L3.0 xR(Vv((,x)).
(20)
where A > 0 denotes the HJB regularizer intensity hyperparameter.
Algorithm We present our training algorithm for DIOTM (Algorithm D.. Our adversarial training objective Ls.e updates alternatively between the value function Vo and the two transport maps To. To. similar to GAN framework (Goodfellow et al. 2020 Note that we simplified Algorithm 1 by omitting the non-dependent terms for each neural network. Additionally, apply the HJB regularizer to both generated distributions, i.e., the forward generated distribution from To and the backward generated distribution from To. Finally, throughout all experiments, distribution for time sampling. i.c., U[0, 1). However, the time sampling distribution can be freely modified. We leave the investigation of the optimal time sampling distribution to future work. In particular, we set ~ 51, i.e., t = 1, the training algorithm becomes similar to OTM (Rout et al.) 2022). In this case, we do not use any displacement interpolation information for & (0, 1). The only difference is our HJB regularizer.
the uniform
4 RELATED WORK
OT problem addresses a transport map between two distributions that minimizes the predefined
cost function. Starting from the dual formulations (Kantorovich) 1948 Vacher & Vialard
2022a b
Gallouët et al. [222)) diverse methods for estimating OT Map have been developed based on minimax
problem (Liu al. 2019 Makkuva et al. 2020
al. 2020 Fan et al.) 2022 Rout al. 2022
Choi et al. 2023 2024b Korotin et al.) 2023). In particular, Fan et al. (2022): Rout et al. 2022):
Korotin et al. (2023) derived adversarial algorithm from the semi-dual formulation of OT problem and
properly recovered OT maps compared other previous works (Makkuva et al.) 2020 Moreover, they provided moderate performance in image generation and image translation tasks for large-scale datasets. Another line of research focuses the dynamical formulation of OT problems (Chen et al.)
2021 Shi et al.| 2024b Liu et al. 2024 2022 Neklyudov et al. 2023 Gushchin et al., 20224). Several
works (Zhang & Chen 2021 Chen al. 2021 Shi et al. 2024b Liu et al. 2024) sampling-based

---

Published as a conference paper at ICLR 2025
Table Evaluation of optimal transport map on the synthetic datasets between OTM (Rout et al.
2022
and ours, based on the 2- Z-aasersteinn distance W2 and the L.2 distance between transport maps.
G-s8G
G G-+25G
Moon-+Spiral
G-XCrclls
Metric
OTM
DIOTM
OTM
DIOTM
OTM
DIOTM
OTM
DIOTM
22 (4)
4.93
3.72
10.09
6.49
0.40
0.55
3.96
2.34
L2 (4)
6,48
4.38
13.09
10.00
1.77
1.67
6.23
5.44
Source Dala
Target Date
Generated Data
GT:
OTM:
DI-OTM
Figure 1: Visualization of transport maps T on synthetic datasets. The transport map is visualized as black line connecting each source sample x to its corresponding generated data T (a:).
approaches, which require numerically simulating ODEs or SDEs during training. More recently,
some methods (Neklyudov et al. 2023 Gushchin et al. 2024 have introduced simulation-fiee techniques, often incorporating adversarial learning strategies. Since our approach also leverages the dynamical properties of OT, we compare method with Shi et al. (2024b); Gushchin al. (2024)
both of which have demonstrated scalability in image translation tasks.
5 EXPERIMENTS
In this section, we conduct experiments on various datasets to evaluate our model from following perspectives.
In Sec
compare model with the ground-truth solution from POT Flamary et al. 2021 on synthetic datasets to assess how well our model approximates OT maps.
In Sec 5.2
we compare our model with various OT models on image-to-image translation
tasks evaluate the scalability of our model.
we evaluate the training stability of our DIOTM model and investigate the effectiveness of HJB regularizer compared to other approaches, such as OTM regularizer and RI regularizer.
In See 5.3
For implementation details of experiments, please refer to Appendix B
5.1 OPTIMAL TRANSPORT MAP EVALUATION ON SYNTHETIC DATASETS
First of all, we evaluate whether our model can accurately learn the optimal transport map T* from the source distribution and the target distribution v. We assess our model against OTM (Rout et al. 2022 Fan et al. 2022) by comparing them to the discrete OT solution from the POT library (Flamary et al. 2021 As discussed in Sec 3.2 when the time sampling distribution is set to 81=1, our DIOTM presents a similar framework with OTM. Hence, we consider OTM an appropriate baseline for demonstrating the advantage of using displacement interpolation. Note that OTM demonstrated the most competitive performance as neural optimal transport map in Korotin al. 2022b). Also, the discrete OT solution indicates the optimal transport map between empirical
7

---

Published as a conference paper at ICLR 2025
Figure 2: Image-to-Inng Image translation results of DIOTM for Wild Cat (64 x 64) AFHQ CChoi ct al. 2020). The left figure shows the source images and the right figure shows the corresponding translated images by DIOTM.
Figure 3: Image translation results of DIOTM for Male + Female (128 x 128) on CelebA (Liu al. 2015).
distributions, i.e., / E, v= Σj 6g. computed via convex optimization. Hence, this discrete OT solution serves as the proxy ground-truth solution of the continuous OT map. We tested model on four synthetic datasets: Gaussian-ter Soaiintttt-GGssiiin (G-8G) Gaussian-to-25Gansiid (G-25G), Moon-to-Spiral, and Ctussian-to-Ciree (G (G-Circles).
Fig. visualizes the transport maps and Tab. presents the quantitative evaluation metric results. In Tab. we evaluate each transport map through two metrics. First, we calculate the 2- Wasseerstei
distance WT 1/) between the generated distribution To [1 and the target distribution v. Second, we evaluate whether the neural optimal transport correctly recovers the optimal pairings (s, T(x)), Specifically, we compute the discrete optimal transport Tässe on test datasets using POT and measure
the L. distance between transport maps, i.e., fx 7 (x) T (1> [13d (etttess((t (x:). Fig. 1 shows that our DIOTM more accurately approximates the target distribution, indicated by the smaller distribution error between the orange target data and the green generated data. This is further supported by quantitative results. In Tab. [ DIOTM achieves smaller distribution error (W2) three out of four datasets and consistently better recovers the optimal coupling (L2) across all datasets. In summary, DIOTM provides better approximation of the optimal transport map T compared to OTM.
5.2 SCALABILITY EVALUATION IN IMAGE-TO-IMAGE TRANSLATION TASKS
We assessed our model on several Image-to-Image (12I) translation benchmarks: Male + Female (Liu al. 2013) (64 x 64), Wild Cat (Choi et al. 2020) (64 x 64), and Male Female (Liu et al. 2075) (128 x 128). Intuitively, the optimal transport map serves as generator for the target distribution, which maps each input its cost-minimizing counterpart y. For instance, if consider Male > Female task with the quadratic cost | y v1l2.. the optimal transport map translates each male image into a female image while minimizing the pixel-level difference. Consequently, various (entropic) optimal transport approaches are widely used for the I2I translation tasks. Therefore, we compared 2022) and OTM (Rout et al. 2022)) and the entropic optimal transport models (DSBM Shi et al. 2024a) and ASBM (Gushchin et al. 2024 on image-to-image (12I) translation tasks.
our model with the optimal transport models (NOT (Fan et al.
Fig. and present the I2I translation results. In each figure, the left subfigure shows source image samples. The right subfigure displays the translated images generated by our transport map. DIOTM successfully generates the target distributions (Cat and Female images) while preserving the identity
8

---

Published as a conference paper at ICLR 2025
Table 2: Image-to-Image translation benchmark results compared to existing neural (entropic) opti- mal transport models. indicates the results con-
OTH (Wild Cal. 64)
OTM (Male Female,, 128)
(Wild 64)
(Male
ducted by ourselves. DSBM scores are taken from (Gushchin et al. 2024
70
De Bortoli et al. 2024).
Data
Model
FID ())
40
CycleGAN ZZu al. 2017
12.94
Male-=Female (64x64)
NOT (Komiun
11.96
30
OTM (Fan at 2022
6.42
DIOTM (Ours)
5.27
20
DSBM Shi et al
2024a
20+
10
Wild->Cat (64x64)
OTM (Fan et 2022
12.42
DIOTM (Ours)
1
10
50
100
DSBM
2024a
Male-FFemale (228x128)
ASBM (Gushchun ctail all 2024
16.08
Figure 4: Ablation study on the regularizer hyperparameter
OTM (Fan et all 2022
7.55
DIOTM (Ours)
7.40
V2 Loss (OTM)
Ve Loss (DIOTM)
Ta Loss (OTM)
Ts Loss (DIOTM)
15
Log10 Scale
Log10 Scale
1.
10
20
30
40
10
20
30
50
60
Num Iterations (K)
Num Iterations (K)
Figure 5: Visualization of the stable training dynamics of DIOTM on Wild-+Cat (64 64). The loss values are visualized on log scale.
Table 3: Comparison of our HJB regularizer with the OTM and R1 regularizers on the DIOTM model. Our HJB regularizer exhibits superior performance and stability to X.
Model
G-+8
G-25G
Moon Moon-Spiine Spiral
A
0.1
0.2
1,0
10
0.1
0.2
1.0
10
0.1
0.2
1.0
10
OTM
22.08
22.90
DIV
31.35
68.01
89.62
DIV
81.02
19.99
14.19
15.66
33.80
W, (1)
R1
3.59
5.01
3.29
4.42
9.20
9.94
11.78
DIV
1.91
2.08
1.05
2.74
HJB
1,93
2.69
2.92
3.21
7.19
14.64
7.99
12.38
0.54
0.59
0.30
1.31
OTM
27.41
28.21
DIV
34.47
96.89
97.98
DIV
87.05
20.96
15.01
34.31
33.80
L2 (1)
R1
4.49
5.39
3.87
5.14
86.05
17.64
19.52
DIV
2.88
3.56
2.36
3.74
HJB
3.05
3.44
3.36
3,98
16.51
15.82
11.11
15.64
1.42
2.25
1.13
2.27
of the input source images. In practice, our DIOTM model trains two transport maps in both directions
To and To. The results for the reverse image-to-ing to--tiiiimge translation are included in Appendix C
Tab. [lrrovidss the quantitative evaluation results. We adopted the FID score (Heusel et al., 2017)
for quantitative comparison. The FID score assesses whether each model accurately generates the target distribution. As shown in Tab. 2] the DIOTM model demonstrates state-of-the-art results among existing (entropic) OT-based methods on 12I translation benchmarks. Specifically, our model significantly outperforms other entropic OT-based models with FID score of 7.40 in the higher resolution case of Male + Female (128 x 128). While OTM achieved comparable results
specific hyperparameter of = 10, OTM diverges for all other À E {50, 100} with FID > 60 (Fig. 4). In contrast, our DIOTM consistently maintains stable performance of FID 10 for À (500100)
(Fig. 4. In summary, DIOTM exhibits superior scalability in handling higher-resolution image datasets compared to previous OT Map models such as NOT and OTM.
5.3 FURTHER ANALYSIS
In this subsection, we provide an in-depth analysis of our DIOTM model. Specifically, we demon- strate the stable training dynamics of DIOTM, compare the HJB regularizer RVV,)) with various

---

Published as a conference paper at ICLR 2025
regularizers introduced in other neural optimal transport models, and conduct ablation study the regularization hyperparameter X.
Stable Training Dynamics The previous approaches to neural optimal transport with adversarial learning often suffer from training instability (Choi et al.) 2024b). These models tend diverge after long training and are sensitive to hyperparameters. In contrast, dynamics. Fig. 5 visualizes the loss values for the transport map Te and the value function Vo throughout the training process. Note that we visualized these loss values on log scale. In OTM (Rout et al.] 2022), the To loss gradually explodes as training progresses. This unstable training dynamics has been major challenge for the OT Map models based on minimax objectives. On the contrary, DIOTM exhibits stable loss dynamics without the abrupt divergence phenomenon, observed in OTM.
DIOTM offers stable training
Comparison to Various Regularization Methods We investigate the effect of our HJB regularizer R(Vs) (Eq. 19 compared to other regularization methods. Specifically, we compare alternatives: OTM regularizer ROTM (Fan et al. 2022] Rout et al.] 2022) and R regularizer RRI (Roth et al.) (017).
These regularizers are also introduced to stabilize the training value function Vg. We incorporate these regularizers into our algorithm by modifying the regularization term R in line 4 of Alg. aa
follows:
(t, Te, Me) = V (co,(z, 24) vt,)),))) || + VV (C. Me) + Vo(()
(Ri Ft, y+) = VV Vo +1) Vot 1.
We assessed these regularizers on synthetic datasets to measure the accuracy of each neural optimal transport map, in Tab. [ Tab. provides the results of each regularization method. Note that because we are comparing different regularizers, a direct comparison under the same is not meaningful. Instead, need to focus on the best results of each regularizer and their robustness to the regularization intensity parameter A. In Tab. 3 our HJB regularizer attains the best and stable results on all three synthetic datasets. We interpret this result by focusing that the HJB regularizer is the only regularizer that incorporates the time derivative 8, V- Our time-dependent value function (t, æ) is trained to distinguish the displacement interpolation for each f. Therefore, regularizing the behavior of our value function across time t is beneficial for DIOTM.
regulaat
Effect of the Regularization Hyperparameter Finally, we conducted an ablation study on the regularization hyperparameter À (Eq. 20 in the image-to-image translation tasks of Wild - Cat (64 64) and Male Female (128 128). Note that, unlike Tab. 3 we compare our model with OTM (Rout et al.] 2022), not the DIOTM model with the OTM regularizer. Fig. demonstrates that our model exhibits significantly greater stability regarding the regularization hyperparameter X, in comparison to OTM. Specifically, in Wild + Cat (64 x 64), our model maintains decent performance for À E (1, 50, 100), showing the FID score around 20 even the worst case. In contrast, the FID) scores of OTM fluctuate severely from below 20 to around 60 depending on X.
6 CONCLUSION
In this paper, we introduced the Displacement Interpolation Optimal Transport Model (DIOTM), a neural optimal transport method based on displacement interpolation. Our method is motivated by the equivalence between displacement interpolation and dynamic optimal transport. We derived the dual formulation of displacement interpolant and developed a method that utilizes the entire trajectory of displacement interpolation improve neural optimal transport learning. Our experiments demonstrated that DIOTM achieves more accurate and stable optimal transport map estimation compared to previous method. A major limitation of this work the requirement to train both bidirectional transport maps T' and 'T. In image-to-image translation tasks, both transport maps are meaningful, e.g. Male <+ Female. However, in generative modeling, the reverse transport map T from the data to the prior noise distribution is not always necessary. In these cases, training T can be an unnecessary cost. Another limitation of this work is that our approach is limited to the quadratic cost. This is because our displacement interpolation parametrization in Eq. [6 is only valid under the quadratic cost assumption.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
Jaemoo was supported by the National Research Foundation of Korea(NRF) grant funded by the Korea governmen((KSIT) RSS-2024-00410661 and [RS- 2024-00342244) YC is supported by the National Science Foundation under Award No. by the National Research Foundation of Korea(NRI)) grant funded by the Korea government(MSST) (RS-2024-00349646]. We thank the Center for Advanced Computation in KIAS for providing computing resources.
Jaewoong was partially supported
Ethics Statement this paper, we aim advance the field machine learning by developing relatively stable and high-pette perhhreeformig algorithm. We hope that further investigation will enable our methodology to applied as a stabilized algorithm across various machine learning applications. By providing robust framework, we anticipate that our approach will contribute to improving the reliability and effectiveness of machine learning solutions in diverse contexts.
Reproducibility Statement To ensure the reproducibility our work, we submitted the anonymized source in the supplementary material, provided complete proofs of our theoretical results in Appendix A and included the implementation and experiment details in Appendix B
REFERENCES
Martial Aguch and Guillaume Carlier. Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis, 43(2):904 924, 2011.
Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In International Conference on Machine Learning, pp. 146-155 155. PMLR, 2017.
Dongsheng An, Yang Guo, Na Lei, Zhongxuan Luo, Shing- Tung Yau, and Xianfeng Gu. Ac-ot: A new generative model based on extended semi-discrete optimal transport. ICLR, 2020.
Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the monge- kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000.
Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood training of schrödinger bridge using tions, 2021.
sdes theory. In International Conference on Learning Representa-
Jaemoo Choi, Jaewoong Choi, and Myungjoo Kang. Generative modeling through the semi-dual formulation of unbalanced optimal transport. In Thirty-sevnnthh Conference on Neural Information Processing Systems, 2023.
Jaemoo Choi, Jaewoong Choi, and Myungjoo Kang. Scalable Wasserstein gradient flow for gener- ative modeling through unbalanced optimal transport. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research. PMLR, 21-27 Jul 2024a.
Jaemoo Choi, Jaewoong Choi, and Myungjoo Kang. Analyzing and improving optimal-transport- based adversarial networks. In The Twelfth International Conference on Learning Representations, 2024b.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung- Woo Ha. Stargan v2: multiple domains. In Proceedings the IEEE/CVE conference computer vision and pattern recognition, pp. 8188-8197, 2020.
Diverse image synthesis for
Valentin De Bortoli, Iryna Korshunova, Andriy Mnih, and Arnaud Doucet. Schr\," odinger bridge flow for unpaired data translation. arXiv preprint a7Xiv:2409.09347, 2024.
Jiaojiao Fan, Shu Liu, Shaojun Ma. Yongxin Chen, and Hao-Min Zhou. Scalable computation of monge maps with general costs. ICLR Workshop on Deep Generative Models for Highly Structured Data, 2022.
11

---

Published as a conference paper at ICLR 2025
R Flamary, N Courty, D Tuia, and A Rakotomamoniy. Optimal transport for domain adaptation. IEEE Trans. Pattern Anal. Mach. Intell, 1, 2016.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Lactitia Chapel. Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021. URL
html
Thomas Gallouët, Roberta Ghezzi, and François-Xaviee Vialard. Regularity theory and geometry of unbalanced optimal transport. arXiv preprint arXiv:2112.11056, 2021.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(1 1):119-44,, 2020.
Nikita Gushchin, Daniil Selikhanovych, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin. Adversarial schr\." odinger bridge matching. arXiv preprint arXiv:2405. 14449, 2024.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge local nash equílibrium. Advances in neural information processing systems, 30, 2017.
Leonid Vitalevich Kantorovich. On problem of monge. Uspekhi Mat. Nauk, pp. 225-226, 1948.
Alexander Kolesov, Petr Mokrov, Igor Udovichenko, Milena Gazdieva, Gudmund Pammer, Evgeny Burnaev, and Alexander Korotin. Estimating barycenters of distributions with neural optimal transport. In Forty-first International Conference on Machine Learning, 2024a.
Alexander Kolesov, Petr Mokrov, Igor Udovichenko, Milena Gazdieva, Gudmund Pammer, Evgeny Burnaev, and Alexander Korotin. Estimating barycenters of distributions with neural optimal transport. In Proceedings of the 41 International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 25016-25041. PMLR, 21-27 Jul 2024b.
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? continuous wasserstein-2 benchmark. Advances in neural information processing systems, 34: 1559-14655, 2021.
Alexander Korotin, Vage Egiazarian, Lingxiao Li, and Evgeny Burnaev. Wasserstein iterative networks for barycenter estimation. Advances in Neural Information Processing Systems, 35: 15672-15686, 2022a.
Alexander Korotin, Alexander Kolesov, and Evgeny Burnaev. Kantorovich strikes back! wasserstein gans are not optimal transport? In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022b.
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. In The Eleventh International Conference on Learning Representations, 2023.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling, In The Eleventh International Conference on Learning Repre- sentations, 2023. URL
Guan HuannHorng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos Theodorou, and Ricky TQ Chen. Generalized schrödinger bridge matching. In The Twelfth International Conference on Learning Representations, 2024.
Huidong Liu, Xianfeng Gu, and Dimitris Samaras. Wasserstein gan with quadratic transport cost. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 4832 9832-4841, 2019.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint 2tXXi:2209.030033 2022.
12

---

Published as a conference paper at ICLR 2025
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference on computer vision, pp- 3730-3738, 2015.
Ashok Makkuva, Amirhossein Taghvaci, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pp. 6672-6681.
PMLR, 2020.
Robert J McCann. A convexity principle for interacting gases. Advances in mathematics, 128(1): 153-179, 1997.
Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Mem. Math. Phys. Acad. Royale Sci., pp. 666-704, 1781.
Kirill Neklyudov, Rob Brekelmans, Alexander Tong, Lazar Atanackovic, Qiang Liu, and Alireza Makhzani. A computational framework for solving wasserstein lagrangian flows. arXiv preprint arXiv:2310.10649 2023.
Gabriel Peyré, Marco Cuturi, et Computational optimal transport: With applications to data science. Foundations and Trends in Muchine Learning, 1165-6):355-677, 2019.
Kevin Roth, Aurelien Lucchi, Sebastian Nowozin, and Thomas Hofmann. Stabilizing training of generative adversarial networks through regularization. Advances in neural information processing systems, 30, 2017.
Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. In International Conference on Learning Representations, 2022.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrüdinger bridge matching. Advances in Neural Information Processing Systems, 36, 2024a.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. Advances in Neural Information Processing Systems, 36, 2024b.
Maurice Sion. General Minimax Theorems. United States Air Force, Office of Scientific Research, 1957.
Thomas Staudt, Shayan Hundrieser, and Axel Munk. On the uniqueness of kantorovich potentials. Contributions to the Theory of Statistical Optimal Transport, 49:101, 2022.
Adrien Vacher and François-Xaun Xavier Vialard. Stability semi-dual unbalanced optimal transport: fast statistical rates and convergent algorithm. 2022a.
Adrien Vacher and François- Xavier Vialard. Stability and upper bounds for statistical estimation of unbalanced transport potentials. 2022b.
Cédric Villani et al. Optimal transport: old and new, volume 338. Springer, 2009.
Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. arXiv preprint arXiv:2112.07884 2021.
Qinsheng Zhang and Yongxin Chen. Diffusion normalizing flow. Advances in neural information processing systems, 34: 11116220-122911. 2021.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 586-595. 2018.
Jun-Yan Zhu, Taesung Park, Phillip Isola, and Alexei Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. In Proceedings of the IEEE international conference on computer vision, pp. 2223-2232, 2017.
13

---

Published as a conference paper at ICLR 2025
A PROOFS
Assumptions Let E IR be closure of connected bounded convex open subspace of IR Let dim(OX) < d. Let P2.0 be collection of probability distributions on X' that are absolutely continuous and have finite second moments. Let μ,, V E P2.00 Moreover, let c(x, y) = -|| y Here, 7 is given positive constant.
A.1 PROOF OF THEOREM
Theorem A.l. Given the assumptions in Appendix(A] for given t E (0, 1), the minimization problem inf, (pr(t, (Eq. is equivalent to the following dual problem:
(1- [ (idl))d(())) + (E.(i)dduw))
sup f1.e.22. with (I-t)f..ttfa..-
(21)
where the supremum taken over two potential functions f..: y R and f2,t X - R. which satisfy t1, + tf2,t = O. Note that the assumptions in Appendix A guarantee the existence and uniqueness of displacement interpolants pł. the forward optimal transport map T't from μ to pi, and the backward transport map T from pt. Based on this, have the following:
(x) E arginf yey [e(2, v) - [i,((0)]
p.a.s..
(22)
(i))) E arginf ZE [c(x, y) - [22c(x))] v-a.s.
(23)
Proof. We prove the theorem by two steps. First, we prove Eq. 21 following the duality theorem in Korotin et al. (2022a) and Kolesov et al. (222a)). Then, we discuss the existence and uniqueness of the potential f1. f2 and the optimal transport maps. Using these uniquesness, we prove Eq. 22
Step 1. We first prove Eq. 21 Suppose 0 < t < 1 is given. By applying the dual form of OT problem Kantorovich 1948). for every p E P2, the following equation satisfies:
(1 (1-1) (1-1)W2 (t, p) (p,
sup [(1 (f,, + (Eynn
[f2, (EymollE.
()/,(;)))))))))))))))))))))))( (u.) (24)
f.FxEC(x)
Then, by applying Theorem 3.4 in Sion 1957). we obtain the following:
L' =nf(1 - t)W p) + (p,
innf
sup [[1 - E (fj [x)] + tByy [v) + Esep (ffa))) (C())
sup fL.NdaEE(())
inf [(1 1) F [/ (x)
+ L.F.yayyllll UL (v) Ezmpl(1 -1 -1 ((.
(25)
Note that we can swap minimax to max-min problem due to the compactness assumption of the space X. Now, suppose m m(fi.,, f2.1) == i1 (ff2 Let f,, = (m(f1,, f2.)) tf2,,)(11 - t). Then, m = (1 - (x) + (=) s - t)fi, + tf2,(()). which implies fi. S fi.e- Then, we can easily obtain ff,, S ff,,- With this inequality, we can obtain the following equality:
sup inf [(1 E) C [f /f, (x) + B (g (w)) Ex-p((1 ()fia(z) (f2a())]] fL(.)2.E((x))
[[1 - -))))EE (f, (x)) + 1E
f2 (3)) + mf fax))
= sup
[1 - 1) (x)] + Lh y 112
ml (h f2.1 (2)) (1)
sup fa,EC(x))
(- (x) + th [[
sup 12,E0(x)
(1 - t) Exea
(26)
14

---

Published as a conference paper at ICLR 2025
Note that the last equation is obtained by using the following property: (fi, a) =fi, for any constant a. By letting fi,
ffa.s T-f" we finally obtain Eq. 21
Step 2. In this step, we prove the uniqueness of the optimal potential pair (fi, f2,1) of Eq. 21 Let
p arg in ( p) + (p,
(27)
First, we should show the well-definednes of above equation. In other words, we should show the existence and uniqueness p² Due to the absolutely continuity assumption of μ., there exists unique deterministic optimal transport map between and (See Chapter Theorem 9.2 of Villani et al.
(2009). Thus, by applying Corollary 7.23 in Villani et al. 2009),, there exist unique solution p².
Now, consider the following optimization problem:
sup (1 (1-1) 1) (Sx (idl))dull))+ / (alo)aoow)) f...2..
(L, + J (as w))]
(28)
Trivially, the first and the second term of Eq. 28 break down into two independent optimization problems. By using the Kantorovich duality theorem (See Kantorovich (1948) or Theorem 5.10
Villani et al. (2009)), the solution value of Eq. 28 is obviously L. Moreover, since the solution pair
(f f2,))
Eq. 21 satisfies (1 - t)ff,, + t 0, is included in the optimal potential pair of Eq.
28
Note that we assumed that is absolutely continuous and the space X' is convex, and a closure of a connected open set. Then, on the support of μ. the Kantorovich potentials fú, and f2, is unique up constant the connected support of and v, respectively (See (Staudt al. (022) ). Therefore, the optimal potentials fī, f2,, of Eq. 21 are a Kantorovich potentials of W2 (p, and (v,,)))) respectively.
Finally, since μ E P2.c. there exists a measurable deterministic optimal transport map T which transport j to p² (See Chapter or Corollary 9.4 in Villani al. (2009)). Then, by Theorem 5.10
and Remark 5.13 in Villani et al. 2209),, we can easily obtain Eq. 22
Let f1,(()) (V, x:) and fL,(2) = (1 - (t, for every t E (0, 1). Then, we can rewrite Eq. 9 as follows:
sup
V. (x)d + (-V,) (V)))(()))vv)))
(29)
Vi
xx
where (a,( 3) = la- vll 1-8
for every 0 < < 1.
The following theorem shows the connection between the optimal potential of transport problem between μ to p and our potential function V.
Theorem A.2. The optimal V," in Eq. [3 satisfies the following:
/ v, (x)dp(x) + J v(())dpt(a)) x
V arg sup V.
(30)
up constant p² -a.... Moreover, there exists (Vi Joseei that satisfies Hamiilonn-lccai--Bellmnn (HJB) equation, i.e.
V.V, + 47
|VV]
pa-s.s
(31)
Proof. By the Step of the proof of the Theorem A.1 have shown that the optimal (ff, f2]
which solves Eq. 22 Kantorovich dual function of W (H,, p²) and W3 p²). respectively. By directly applying this fact, we obtain Eq. 30
Now, we prove Eq. 31 the HJB equation. Let p" be defined as the Step 2 of the proof Theorem
A. Then, discussed in Step 2, there exists unique p" that satisfies Eq. [7 Since y is absolutely
15

---

Published as a conference paper at ICLR 2025
continuous, it is well known that p² is also absolutely continuous. Due to the compactness of the space X. the optimal Vú is differentiable with respect to z E A' p*-a.s.. Now, by applying Theorem 7.36 and Remark 7.37 of Villani et al. (200)), there exists V (0,1) X + R such that
(a) = inf (c,, y) - Vu()),,
(32)
for all 0 <<<<t<> 1. By applying Hopf-Lax formula, we obtain,
[. T||e>|| du v((x,)
- Xs, =:(0.1]xx'--
inf
(33)
where X, vx))... Thus, by organizing the Eq. 33
obtain Eq. 31 follows:
1 s-8 (Vi(X,) - vv(X)))
lim t+ v: v: x - +x
inf
=||0|12
(34)
(T 1|2 0,V, VV, v.) = - |VV.11
inf
(35)
47
Note that law(X)) Pt. Thus, V.Ve +
|V
-a... 4r
B IMPLEMENTATION DETAILS
For all experiments, parametrize T T x R X where X' is data space, and IR an auxiliary variable. As reported in [Choi et al. (2023 2024b), incorporating the auxiliary variable
slightly improves the performance. For all OTM experiments have implemented, use the same network and the same parameter to DIOTM, unless otherwise stated.
B. 2D EXPERIMENTS
Data Description In this paragraph, we describe our synthetic datasets:
8-Gaussian: For m, 12 (cos ii,, sin ( where i 0, ., 7 and a 0.4, the distribution is defined as the mixture of N (m,,,,;
sin (m;;0 o²) with equal probability.
25-Gaussian: For m,j (8i, 8j) where -2 -1.... 2 and = -2, -1, -,,.... 2, the distribution is defined as the mixture of MM o²) where 0.01.
Moon to Spiral: We follow Choi et al. 2024b).
Two Circles: We first uniformly sample from the circles of radius and 16 with the center at origin. Then, we add Gaussian noise with standard deviation of 0.2.
Network Architectures We first descrihe the value function network Vo,,,, ar). The input x E X' embedded using two-layer MLP with hidden dimension of 128. The time variable t is embedded using a positional embedding of dimension 128, followed by a two-layer MLP, also with a hidden dimension of 128. These two embeddings SiLU activation function applied to all MLP layers. Besides, for the transport map networks, we employ the same network Choi et al. (2023) with hidden dimension of 128.
then added and passed through three-layer MLP. The
Training Hyperparameters We set Adam optimizer with (81, = (0, 0.9), learning rate of 10-4 and the number of iteration of 120K. We set 0.1 and X =
Discrete OT Solver We used the POT library Flamary al. (021) to obtain an accurate transport plan Toot We used 1000 training samples for each dataset in estimating Tpot to sufficiently reduce the gap between the true continuous measure and the empirical measure.
B.2 IMAGE TRANSLATION
Training Hy Hyperparaeeters We follow the large neural network architecture introduced in Xiao
et al. (2021). We use Adam optimizer with (B1, (2) (0, 0.9), learning rate of 10-4, and trained
16

---

Published as a conference paper at ICLR 2025
for 60K iterations. We use cosine scheduler to gradually decrease the learning rate from 10- to 10 100$. The batch size 64 and 32 is employed for 64 64 and 128 x 128 image datasets, respectively. We use 0.001 for CelebA image dataset, and 0.0005 for AFHQ dataset. We use ema rate of 0.9999 for 64 x 64 image datasets and 0.999 for 128x 128 image datasets.
Evaluation Metric Regarding the FID computation, followed the evaluation scheme of De Bor toli et al., (024) for the Wild-Cat experiments and (Korotin et al. [023 Gushchin et al. 2024) for the CelebA experiments for a fair comparison. Specifically, in the Wild Cat experiments, we generated ten samples for each source test image. Since the source test dataset consists of approxi- mately 500 samples, we generated 5000 generated samples. Then, we computed the FID score with the training target dataset, which also contains 5000 samples. Also, in the CelebA experiment, computed the FID score using the test target dataset, which includes 12247 samples. We generated the same 12247 samples and compared them with the test target dataset.
C ADDITIONAL. QUALITATIVE RESULTS
In this section, include some qualitative results image-to-imaln to-gage translation tasks. We visualize the source images and its transported samples.
Figure 6: Unpaired Male + Female translation for 64 64 CelebA image.
17

---

Published as a conference paper at ICLR 2025
Figure 7: Unpaired Female
Male translation for 64 64 CelebA image.
DoH
da
CIROC
CIROC
Figure 8: Unpaired Female Male translation for 128 128 CelebA image.
18

---

Published as a conference paper at ICLR 2025
Figure 9: Unpaired Male Female translation for 128 x 128 AFHQ image.
Figure 10: Unpaired Wild Cat translation for 64 x 64 AFHQ image.
19

---

Published as a conference paper at ICLR 2025
Figure 11: Unpaired Cat - Wild translation for 64 64 AFHQ image.
20

---

Published as a conference paper at ICLR 2025
1
50
Source
OTM
DIOTM
OTM
DIOTM
Figure 12: Qualitative comparison of unpaired Male -> Female translation between OTM and DIOTM
on 128 x 128 CelebA images.
21

---

Published as a conference paper at ICLR 2025
2 = 10
7 = 50
Source
OTM
DIOTM
OTM
DIOTM
Figure 13: Qualitative comparison of unpaired Wild Cat translation between OTM and DIOTM
for x 64 AFHQ images.
22

---

Published as a conference paper at ICLR 2025
D ADDITIONAL QUANTITATIVE RESULTS
D.1 ADDITIONAL QUANTITATIVE RESULTS FOR IMAGE-TO-IMAGE TRANSLATION
Evaluation of Image Alignment To evaluate whether our transported image Te (x) preserves the properties of the original source image z, we compute the LPIPS (Zhang et al. 2018 metric between the source datapoint z and its generation Tz(z). The reported results the mean LPIPS values calculated over a test dataset. As shown in Tab. our model demonstrates comparable results to other image-to-image (I2I) benchmarks.
Table 4: Comparison of LPIPS (L) scores for perceptual similarity on Image-to tooeetto-iio to-Image translation benchmarks.
Model
Wild +Cat (64x64)
Male :Female (128x128)
DSBM
0.59
0.25
ОТM
0.47
0.21
DIOTM
0.45
0,25
FID Curves along Training Process As discussed in Sec. 5.3 our model demonstrates robustness to the regularization hyperparameter (Fig. 4 and maintains stable training dynamics (Fig. [5}. In this paragraph, we provide the FID curves throughout the training process for various values of hyperparameter X. Fig. 14 shows that our model exhibits stable convergence across various hyperparameters In contrast, OTM converges only under specific hyperparameter settings, i.e.,
= 10 for Wild-sCat (64 x 64) and X 50 for Male->Female (128 x 128).
150
OTM A=10
125
OTM A
OTM A=100
100
DIOTM A=10
75
DIOTM A=50
DIOTM 1=100
50
25
10
20
30
Num Iterations (K)
(a) Visualization of FID curve on Wild-Cat (64 x 64).
70
OTM 1=10
60
OTM A=50
DIOTM A=10
DIOTM A=50
40
20
10
10
20
30
40
50
60
Num Iterations (K)
(b) Visualization curve on Male->Female (128 x 128).
Figure 14: Visualization of FID (L) curves for various hyperparameter À on Image-to-Image transla-
tion tasks.
23

---

Published as a conference paper at ICLR 2025
Ablation Study Time Sampling Distribution We investigate alternative time samplers for line of Algorithm beyond the uniform distribution U 1]. Specifically, we test two time sampling distributions: Beta(2, 2) and (0ta 0.5). The Beta(2, 2) distribution samples intermediate time values (t R 0.5) more frequently, while Beta(0.5, 0.5) tends to sample near the boundaries of the interval (t 0, 1) more often. This ablation study time sampling distribution evaluated on the Wild->Cat dataset. As shown in Table the uniform sampler achieves slightly better performance compared to the other sampling methods. However, the differences in performance are relatively minor, indicating that the alternative sampling methods yield results comparable to the uniform sampler.
Table 5: Ablation study on the time sampling distribution. FID scores are evaluated on the Wild->Cat (64 64) dataset.
Sampling Distribution
Uniform
Bcta(0.5,0.5)
Bcta(2,2)
FID (1)
10.72
11.85
12.45
D.2 ADDITIONAL QUANTITATIVE RESULTS FOR 2D-TOY EXPERIMENTS
Cyclical Property Our DIOTM model learns the bidirectional transport maps To. To simultaneously. Here, To and To denote the optimal transport maps for - and v + μ., respectively. Theoretically, these optimal transport maps satisfy T2 Ti (x) x: and Ti T2(y) y almost surely under our assumptions. We refer to this property as a cyclical property. To evaluate whether our model satisfies this cyclical property, we measured the MSE reconstruction error between x and T2 ( (z). i.e. E>>~-l]?s o Ti (x) x x=|²]. For comparison, we trained two OTM models for + and + and measured the MSE reconstruction error using these models.
As shown in Tab. 6 our DIOTM achieves better reconstruction eπоо on four out of six experiments. Our model shows a larger reconstruction eтrоr in the (25G + G - 25G). However, this reconstruction is meaningful when the generating distribution errors also considered, i.e., To μ. V and To o v μu (Table 1.. We interpret this result being due to the larger distribution еπхог of OTM in the G + 25G case.
case
Table 6: Evaluation of the cyclical property synthetic datasets.
Model 8G-G G-8GG-x G
25G->G- -25G
G G->25G-4 G
M >S M
S~M->
OTM
1.06
0.040
4.13
0.56
4.01
1.12
DIOTM
0.22
0.015
12.43
0.68
1.11
0.46
D.3 COMPARISON TO OT BENCHMARKS ON HIGH DIMENSIONAL DATA
Table Continuous neural optimal transport benchmark Korotin et al. 2021 l.. The transport maps
evaluated based the L2-UVP (%, !) metric and Cosine Similarities (1). The baseline scores
are taken from Korotin et al. (202T) and indicates the results by ourselves.
Dimension
D=16
D=64
Metric
2 UVV (L)
cos (t)
L CV-UVP (L)
cos (t)
L
41.6
0.73
63.9
0.75
QC
47.2
0.70
75.2
0.70
MM
2.2
0.99
3.2
0.99
MMv1
1.4
0.99
8.1
0.97
MMv2
3.1
0.99
10.1
0.96
MMv2:R
7.7
0.96
6.8
0.97
MM-B
6.4
0.96
13.9
0.94
DIOTMI
2.5
0.98
10.4
0.96
24

---

Published as a conference paper at ICLR 2025
High-dimensional Experiment In this section, assess model using high high-dimensional Gaussian mixture experiment, following the benchmark proposed in Korotin et al. (2021). Table presents the results. Our model demonstrated reasonable performance, effectively learning the optimal transport map. Still, our model does not achieve state-of-the-an results. We believe this is because our primary goal
achieve better scalability and stable estimation high-dimensional image dataset. To achieve this, our DIOTM does not impose specific structure constraints on two neural networks in our min-max optimization algorithm, i.c., the transport map and the discriminator, such as the input convex neural networks ICNN (Amos et al. [2017). Note that most of the baselines in Table 7 use ICNNs to parameterize their potential functions (discriminator), However, this approach sacrifices scalability, particularly for image datasets.
Experimental Settings Unless otherwise stated, we follow the hyperparameter and the imple- mentation of Korotin al. (2021). We use the number of iterations of 100K, Adam optimizer of (81,82 = (0, 0.9) with the learning rate of 10-3.
For both forward and backward transport maps, use 3-MLP layers with SiLU activation function. For the discriminator, we embed time variable by using sinusoidal embedding of dimension 128, and then pass it through the two MI.P layers with SiLU activation. For the state variable æ, embed it through 2-MLP layers. Then, we add time and state embeddings and pass them through a 3-layer MLP with SiLU activation. We the hidden dimension of 2048 for all MLP layers.
25


Reviews:
[{"review_id": "288nwkQQnc", "reviewer": "Reviewer_xVqq", "summary": "The authors propose a novel method (DIOTM) to solve the optimal transport mapping problem for the quadratic transport cost (Wasserstein-2 OT) with neural networks. The approach is ideologically inspired by the previous works in the field which solve the dual (semi-dual) optimal transport problem by approximating an OT map and the dual potential (a.k.a. discriminator) with neural networks and optimizing them in the GAN-style adversarial manner (max-min).\n\nThe key innovative thing in the current paper lies in exploiting the properties of the W2 OT maps. They are related to the displacement interpolation linear interpolation from the input distribution to the target using the OT map). More precisely, the authors formulate the (semi-)dual problem for finding the displacement interpolation for a given time moment t in (0,1) which requires optimizing a particular t-dependent dual potential. Then they group all these problems together and obtain a dual problem when they have to optimize over one t-conditioned dual potential (and also additional t-dependent transport maps). In principle, each problem for different t can be viewed as independent, but\n\n1) The authors note that after some reparameterization, the t-dependent dual potentials should satisfy the Hamilton-Jacobi-Bellman (HJB) condition. At this point, the authors propose to incorporate the HJB-inspired regularization into the optimization, which helps connect optimization problems for each t together.\n\n2) The authors note that the optimal transport maps at each time moment t are connected with each other. In fact, they all can be expressed through each other and through the main transport map (from source to target). As a result, the authors use restricted parameterization where all these t-dependent maps are parameterized through a single map.\n\nThe resulting algorithm is a (simulation free) bi-directional max-min adversarial training scheme. The authors demonstrate the superiority of the proposed technique compared to previous dual form neural optimal transport solvers & their regularization techniques through a series of experiments (toy 2D data + image-to-image translation).", "questions": "I think the ideas in this paper are very interesting and should be presented to the community. My current score is based on the current condition of the paper but I may adjust it if the authors carefully reply to the weaknesses which I raised and revise the paper accordingly. Also, I have some additional questions:\n\n1) What is the point of introducing alpha? The OT map/displacement maps should be the same for all alpha, right?\n\n2) Could you please provide some analysis of the time sampling schemes (line 294)? In diffusion models, this is an important aspect, so I believe it may be important here as well and at least some analysis should be provided. For example, you can consider a scheme where t is mostly samples closer to 0/1 and the other scheme where t is concentrated around 0.5 and show the results.\n\n3) It looks like the training curves (figure 5) present the losses which are generally not very representative in adversarial learning. Could you please provide FID(epoch) plots to see how stably your method converges compared to the baselines? This would be much more convincing.\n\n4) Most comparisons are quantitative through FID which does not measure optimality but only measures matching the target. Could you please provide a side-by-side qualitative comparison with the baseline in I2I tasks? It would be nice to see how your trained generator preserves the content compared to the baselines.\n\n5) Could you please run your method in some I2I experiment several times. Does it converge to roughly the same solutions (qualitatively), i.e., recovers (nearly) the same map (which should be optimal)?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) The idea of exploiting the displacement interpolation overall looks interesting and fresh. To my knowledge, it has not been actively studied in the field, so I believe that further developing it may be interesting and fruitful for the community of adversarial/dual-based OT methods. Overall, the contribution of this paper looks as significant for the neural OT field, as WGAN-GP improved WGAN.\n\n2) The HJB based-regularization proposed here seems to be very natural and unbiased in the sense that it looks theoretically justified and does not bias the resulting solution. This is not the case for other GAN-based regularizing techniques which appear in related works (like R1 or other gradient penalty regularizers). However, for me it is still not clear from the main text if the authors in their method use only HJB or HJB+R1. This should be clarified.\n\n3) The experimental comparison on unpaired Image-2-Image looks rather convincing and supports the main claim that HJB regularizer is useful for stability and works (I deduce this from the results of comparison with various dual OT methods).\n\n4) The text is overall readable and the clarity is ok (although sometimes the amount of the bolded text is too annoying).", "weaknesses": "1) I believe that there might be a theoretical gap in the proposed DI-OTM approach which lies in the restricted parameterization of the t-dependent transport maps. Specifically, each transport map (for a particular t) should be parameterized the way that it should solve the corresponding inner conjugation (c-transform) minimization for a particular corresponding dual potential (for time t). However, when the authors tighten all the transport maps together via a single function, this may not hold and may spoil the theoretical validity of the proposed semi-dual form. This aspect should be discussed in more detail.\n\n2) I think that some of the results presented here are not completely novel and the authors miss a large set of related work. The key problem which is exploited in the current work is the displacement interpolation optimization (equation 8). In essence, this is the Wasserstein-2 barycenter problem and, to my understanding, it has already been well studied both in theory and in practice. For example, the W2 dual barycenter problem (equation 9 in theorem 3.1 in the current paper) has been derived in the founding work [1], see their derivations around proposition 2.2. The semi-dual version (which is the second part of theorem 3.1 in the current paper) seems to directly follow from the general semi-dual for barycenters which has been recently introduced in [2] (theorem 4.1). I think these relations to the barycenter literature (theoretical and computational) should be clearly clarified and the related literature should be included.\n\n3) The DIOTM approach proposed here seems to work only for the quadratic cost optimal transport (and may be for some lp-based OT as well) due to reliance on the displacement interpolation properties. It looks like it can not be generalized to more general OT formulations, e.g., formulations with non-lp transport costs. This point is more a limitation than a weakness as the authors specifically target the quadratic cost OT. Nevertheless, it should be mentioned in the paper and the background considers the general cost OT.\n\n4) While the authors claim that they significantly improve the accuracy of solving OT, they omit detailed evaluation of this aspect in high dimensions. The experiments in 2D are good but do not convincingly support the claim, more advanced and high-dimensional evaluation should be considered [3] and some recent baselines should be included like [4].\n\n5) Some of the theoretical statements are not very mathematically rigorous. For example, the authors prove some results regarding the optimal dual potentials (like eq. 10/11), but do not explain to which functional spaces they belong. If I correctly get it from the proof, they should be continuous functions. Does the supremum among the continuous functions is achieved, i.e., are f* also continuous functions?\n\nReferences\n\n[1] Agueh, M., & Carlier, G. (2011). Barycenters in the Wasserstein space. SIAM Journal on Mathematical Analysis, 43(2), 904-924.\n\n[2] Kolesov, A., Mokrov, P., Udovichenko, I., Gazdieva, M., Pammer, G., Burnaev, E., & Korotin, A. Estimating Barycenters of Distributions with Neural Optimal Transport. In Forty-first International Conference on Machine Learning.\n\n[3] Korotin, A., Li, L., Genevay, A., Solomon, J. M., Filippov, A., & Burnaev, E. (2021). Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. Advances in neural information processing systems, 34, 14593-14605.\n\n[4] Amos, B. On amortizing convex conjugates for optimal transport. In The Eleventh International Conference on Learning Representations.", "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": "BW9SBQsrZ0", "reviewer": "Reviewer_uCi4", "summary": "The paper presents a theoretically justified method for the computation of dynamic optimal transport using the theory of Displacement Interpolation. The authors derive a dual formulation of Displacement Interpolation. They show that the optimal potential for solving the dual problem satisfies the HJB equation and incorporates the HJB equation as a regularizer for the training of the proposed method. The method is validated on synthetic datasets - G->8G, G->25G, Moon->Spiral, G->Circles - and several image-to-image translation problems - Celeba and Wild->Cat. The proposed method achieves the best FID among considered competitors CycleGAN, OTM, NOT, DSBM, and ASBM for image-to-image translation problems and outperforms closely related OTM method on most synthetic datasets in W^2 distance.", "questions": "Questions:\n1. It looks like in Eq. 6 the integration should be over $x$ instead of $d\\rho_ {t}$, and $\\rho_ {t}(x)$ should be under the integral. Can you comment on this?  \n2. Can you clarify how long it took you to train your methods for image-to-image translation problems compared to competing methods? \n3. What is the number of parameters used by all the methods for image-to-image translation problems? Are they comparable?\n4. Have you experienced failures of your method, and if so, can you provide them?\n5. Can you provide an evaluation of your method on the Wasserstein-2 benchmark to show that the method is capable of solving optimal transport in high dimensions?\nTyping errors:\n1. Line 49 - double \"the\"; one should be deleted.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The method has a derivation of the dual problem for displacement interpolation, which opens the possibility of numerical optimal transport computation from the perspective of the Benamou-Brenier dynamic transport formulation.\n2. Experiments on toy examples and image-to-image translation problems show that the proposed method achieves good numerical results over competing methods for optimal transport computation and is scalable to image problems. \n3. The paper proves numerically that the HJB regularizer improves the training procedure and is better than the OTM and R1 regularizers. This regularizer seems to be novel in the literature of numerical optimal transport computation.", "weaknesses": "1. The method doesn't compare to closely related flow-based optimal transport methods, such as Rectified Flow (Flow straight and fast: Learning to generate and transfer data with rectified flow, ICLR-2023) and Flow Matching (Flow Matching for Generative Modeling, ICLR-2023). I suggest the authors compare with these methods as well.\n2. The paper lacks a visual comparison for image-to-image translation problems between different methods and a discussion of why competing methods perform worse. It is not clear why the proposed method achieves better numerical results when it has similar visual results to competitors.\n3. It is not clear how well the method computes optimal transport in high dimensions. I suggest that the authors evaluate their method on the Wasserstein-2 benchmark (Do neural optimal transport solvers work? A continuous Wasserstein-2 benchmark, NeurIPS-2021).", "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": "xqdoozFHTl", "reviewer": "Reviewer_WWpL", "summary": "This paper builds on displacement interpolation in Optimal Transport (OT) and introduces a time-derivative HJB regularizer, enhancing training stability. The training of the model is based on min-max optimization similar to GAN. It achieves state-of-the-art results on both synthetic data and image-to-image translation tasks w.r.t $W_2$, $L_2$ and FID score.", "questions": "Refer to weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The paper presents comprehensive and detailed theoretical derivations, with notable innovations within the OT framework. It leverages the dual formulation of displacement interpolation to derive a new min-max optimization function.\n\n2. In terms of experimental performance, the proposed HJB regularizer is effectively insensitive to the hyperparameter $\\lambda$, performing better than other regularizers such as R1 and OTM. And DIOTM outperforms other benchmarks and exhibits more stable training.", "weaknesses": "1. The motivation behind the theoretical innovation is unclear. There is no analysis explaining why decomposing the optimization of $T_\\theta$ in OTM into optimizations for forward $\\overrightarrow{T_\\theta}$ and backward $\\overleftarrow{T_\\theta}$ improves training stability.\n\n2. The experimental results in Table 2 appear unusual. I couldn't find related experimental setups for the benchmarks, and some references don’t report similar experiments or use different resolution datasets. Since the FID scores for these benchmarks couldn’t be directly cited, how were these results obtained? Were all models trained for the same number of steps? It would be beneficial to add an ablation study of FID vs. training steps.\n\n3. The paper argues that DIOTM is more stable than OTM, but Fig. 5 shows that OTM remains stable for the first 40K steps before experiencing a sudden spike in loss. What caused this increase? If the loss curve does not decrease further, why train for 60K steps rather than 40K?\n\n4. The paper only provides visualizations for DIOTM, making it hard to compare visually with baselines. The DSBM paper’s wild-to-cat results at 512x512 resolution look much better than those in Fig. 2, yet its FID score in Table 2 is much higher. Could the authors clarify this discrepancy?", "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": "PHdIsSKNMf", "reviewer": "Reviewer_Zhh7", "summary": "This paper proposes a new method to estimate the optimal transport map between two distributions - a source and a target. The proposed method, called Displacement Interpolation Optimal Transport (DIOTM), leverages displacement interpolation which is the optimal solution of a particular dynamic formulation of OT with quadratic cost. The core component of the training algorithm for DIOTM involves a min-max loss objective, similar to GAN framework. This min-max objective is derived from the dual problem of the original minimization problem of displacement interplant. The expression involves a supremum over two potential value functions which can be combined into a single potential value function. These potential functions play a role similar to discriminators in GANs and the transport maps are similar to generators. In addition, the regularization term of the loss objective is derived from Hamilton-Jacobi-Bellman (HJB) optimality condition of the value function. The training algorithm involves alternately updating the potential value function and  the two transport maps. The paper shows applications of the proposed approach on image-to-image translation on datasets such as Male $\\rightarrow$ Female (64$\\times$64, 128$\\times$128), Wild $\\rightarrow$ Cat (64$\\times$64), etc.", "questions": "1. Training stability: Can we have multiple curves to understand how frequently the training diverges for OTM? Also, how sensitive is training of OTM to various hyperparameters?\n2. What are some practical constraints on the source and target distributions when trying to learn an OT map with DIOTM? Can it learn OT map in the cases where the distance between the source and target distribution might be large? For instance, prior works in this space consider more complex datasets/tasks for image-to-image translation such as mask-to-image synthesis (COCO / ADE-20K), sketch-to-image synthesis, day-to-night, summer-to-winter, colorization etc. \n3. The results of Figure 11 seem much more suboptimal than other cases (with multiple faces) etc. What could be the reason for more failures for this  pair of distribution?\nMinor:\n- Line 243: typo - parametrization\n- Line 253, 256, 258: Consider using different parameter notation e.g. $\\overrightarrow{T}_\\theta$ and $\\overleftarrow{T}_\\tilde{\\theta}$ for the two transport maps, as these are parametrized with two different networks with different parameters. This would make it clear that these two networks are trained separately, as opposed to using a shared network. \n- Repeated citation for Diffusion Schrodinger bridge matching paper.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Writing**: The paper is well-written. It provides sufficient background on major concepts involved in DIOTM such as displacement interpolation. The core algorithm has been explained well and the underlying theoretically motivation has been explained well.  \n\n**Quality and significance**: Improving stability of Optimal transport is an important problem and this paper proposes a method to addresses it. \n- The experimental results on simple 2D toy datasets seem to indicate improved performance compared to prior methods as indicated in Table 1.  \n- DIOTM seems to outperform other optimal transport based models on image-to-image translation task in terms of metrics such as FID (Table 2).\n- The proposed HJB regularizer seems to help with improved training dynamics (Figure 5). Further, HJB regularizer seems to be less sensitivity to the choice of regularization hyperparameter  (Table 3) which is a desirable property.", "weaknesses": "1. This method trains two optimal transport maps from source to target distribution and vice versa which is a bit inefficient. Further, there are no experiments which demonstrate that the two independently trained transport maps are invertible, which they should be theoretically. How does source -> target -> source reconstruction perform on various datasets in the paper? Similarly, target -> source -> target reconstruction on images should be reported with a metric such as l2 error/reconstruction error. \n2. Qualitative Results: The paper should Include qualitative comparison with other methods on Image-to-Image translation baseline. FID doesn’t necessarily capture lot of semantic and perceptual information of images. A better comparison would be side-by-side comparison of images obtained from DIOTM and previous OT benchmarks.\n3. Quantitative results: Table 2 compares DIOTM with existing neural optimal transport models. For the sake of completeness, the paper should include another table that includes other state-of-the-art methods (e.g. GANs[1], flows as well as diffusion-based methods (e.g. Wang et al. [2]) for image-to-image translation task so that reader gets an overall picture of the landscape and the gap of DIOTM from SOTA method. I would like to reiterate that it is completely alright if DIOTM is not SOTA overall, compared to other methods for I2I task, but such a table should be included, as it is a standard practice. \n4. Implementation details: The paper is missing some of the implementation details, specifically architecture details of networks for image-to-image  translation task. Further, the number of images used to calculate FID is unclear.\n5. The largest image resolution considered in this work is 128X128 which is not very large. In order to reliably evaluate scalability, larger resolutions such as 256X256 or 512X512 should be considered. See Isola et al. [1] for a list of potential datasets for image-to-image translation tasks on larger resolution. \n\n[1] Isola, Phillip, et al. \"Image-to-image translation with conditional adversarial networks.\" Proceedings of the IEEE conference on computer vision and pattern recognition. 2017.\n[2] Wang, Tengfei, et al. \"Pretraining is all you need for image-to-image translation.\" arXiv preprint arXiv:2205.12952 (2022).", "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)