Improving Neural Optimal Transport via Displacement Interpolation

## Abstract

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 𝑡 t italic_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.

## 1 Introduction

Optimal Transport (OT) problem (Villani et al., ; Peyré et al., ) 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 a 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., ; Choi et al., ; ), image-to-image translation (Korotin et al., ; Gushchin et al., ; Fan et al., ), and domain adaptation (Flamary et al., ). OT Map is particularly well-suited for unsupervised (unpaired) distribution transport problems, as 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-min objectives exhibit unstable training dynamics and sensitivity to hyperparameters (Makkuva et al., ; Fan et al., ; Rout et al., ; Flamary et al., ).
These challenges limit their applicability to high-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., ; Roth et al., ) and generalizing the standard OT problem to the unbalanced optimal transport problem (Choi et al., ; ).

In this paper, we propose a novel approach for learning the optimal transport map based on displacement interpolation. We refer to our model as the Displacement Interpolation Optimal Transport Model (DIOTM). We identify a fundamental connection between the the (static) optimal transport map and displacement interpolation. Motivated by this relationship, we derive a max-min 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 to existing methods.
In particular, DIOTM achieves competitive FID scores in image-to-image translation tasks, such as 5.27 for Male$\rightarrow$Female ($64\times 64$), 7.40 for Male$\rightarrow$Female ($128\times 128$), and 10.72 for Wild$\rightarrow$Cat ($64\times 64$), comparable to the state-of-the-art results. Our contributions can be summarized as follows:

- •
We propose a method for learning the optimal transport map based on displacement interpolation.
- •
We derive the dual formulation of displacement interpolation and utilize this to formulate a 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 $\mathcal{X}$ be a connected bounded convex open subspace of $\mathbb{R}^{d}$, and let $\mu$ and $\nu$ be absolutely continuous probability distributions with respect to Lebesgue measure.
We regard $\mu$ and $\nu$ as the source and target distributions.
For a measurable map $T$, $T_{\#}\mu$ represents the pushforward distribution of $\mu$.
$\Pi(\mu,\nu)$ denote the set of joint probability distributions on $\mathcal{X}\times\mathcal{Y}$ whose marginals are $\mu$ and $\nu$, respectively.
$c(x,y)$ refers to the transport cost function defined on $\mathcal{X}\times\mathcal{Y}$.
Throughout this paper, we consider $\mathcal{X}=\mathcal{Y}\subset\mathbb{R}^{d}$ with the quadratic cost, $c(x,y)=\alpha\lVert x-y\rVert^{2}$, where $d$ indicates the dimension of data.
Here, $\alpha$ is a given positive constant.
Moreover, we denote $W_{2}(\cdot,\cdot)$ as the 2-Wasserstein distance of two distributions.

## 2 Background

In this section, we provide a brief overview of key concepts in 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 $\mu$ to the target distribution $\nu$ (Villani et al., ). The optimality of the transport plan is defined as the minimization of a given cost function. Initially, Monge () formulated the OT problem with a deterministic transport map $T$ where $T_{\#}\mu=\nu$. However, the Monge OT problem is non-convex and the optimal transport map $T^{\star}$ may not exist depending on $\mu$ and $\nu$. To overcome this problem, Kantorovich () introduced the following convex formulation of the OT problem:

$$ $C(\mu,\nu):=\inf_{\pi\in\Pi(\mu,\nu)}\left[\int_{\mathcal{X}\times\mathcal{X}} c(x,y)d\pi(x,y)\right],$ (1) $$

where the minimization is conducted over the set of joint probability distribution $\pi\in\Pi(\mu,\nu)$. We refer to this $\pi$ as the transport plan or coupling between $\mu$ and $\nu$. When the optimal transport map $T^{\star}$ from the Monge OT exists, the optimal coupling $\pi^{\star}$ satisfies $\pi^{\star}=(Id\times T^{\star})_{\#}\mu$. For a general cost function $c(\cdot,\cdot)$ that is lower semicontinuous and lower bounded, the Kantorovich OT problem (Eq. 1) can be reformulated as follows (Villani et al. (), Chapter 5):

$$ $C(\mu,\nu)=\sup_{V\in L^{1}(\nu)}\left[\int_{\mathcal{X}}V^{c}(x)d\mu(x)+\int_ {\mathcal{X}}V(y)d\nu(y)\right],$ (2) $$

where the potential function $V\in L^{1}(\nu)$ is an integrable function with respect to $\nu$ and the $c$-transform of $V$ is defined as

$$ $V^{c}(x)=\underset{y\in\mathcal{Y}}{\inf}\left(c(x,y)-V(y)\right).$ (3) $$

This formulation (Eq. 2) is called the semi-dual formulation of OT.

Rout et al. () and Fan et al. () proposed a method for learning the optimal transport map $T^{\star}$ by leveraging this semi-dual formulation (Eq. 2) of OT and applied this $T^{\star}$ for generative modeling. In generative modeling, the source distribution $\mu$ and the target distribution $\nu$ correspond to the Gaussian prior and the target data distribution, respectively.
Specifically, these models parametrize the potential $V=V_{\phi}$ in Eq. 2 and the transport map $T_{\theta}:\mathcal{X}\rightarrow\mathcal{Y}$ as follows:

$$ $T_{\theta}:x\mapsto\arg\min_{y\in\mathcal{X}}\left[c(x,y)-V_{\phi}\left(y \right)\right]\quad\Leftrightarrow\quad V_{\phi}^{c}(x)=c\left(x,T_{\theta}(x) \right)-V_{\phi}\left(T_{\theta}(x)\right).$ (4) $$

Note that the parametrization of $T_{\theta}$ on the left-hand side is equivalent to the representation of $V_{\phi}^{c}$ on the right-hand side, by the definition of the $c$-transform (Eq. 3). Based on this, we arrive at the following optimization problem:

$$ $\mathcal{L}_{V_{\phi},T_{\theta}}=\sup_{V_{\phi}}\left[\int_{\mathcal{X}}\inf_ {T_{\theta}}\left[c\left(x,T_{\theta}(x)\right)-V_{\phi}\left(T_{\theta}(x) \right)\right]d\mu(x)+\int_{\mathcal{X}}V_{\phi}(y)d\nu(y)\right].$ (5) $$

Intuitively, $T_{\theta}$ and $V_{\phi}$ serve similar roles as the generator and the discriminator of a GAN (Goodfellow et al., ). A key difference is that $T_{\theta}$ is trained to minimize the cost $c\left(x,T_{\theta}(x)\right)$, since $T_{\theta}$ learns the optimal transport map. For convenience, we denote the optimization problem $\mathcal{L}_{V_{\phi},T_{\theta}}$ (Eq. 5) as the OT-based generative model (OTM) (Fan et al., ).

#### 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. (1)) focuses solely on how each $x\sim\mu$ is transported to $y\sim\nu$, the dynamic optimal transport problem tracks the continuous evolution of $\mu$ to $\nu$. Formally, the dynamic formulation of the Kantorovich OT problem (Eq. 1) for the quadratic cost $c(x,y)=\alpha\lVert x-y\rVert^{2}$ can be expressed as follows:

$$ $\inf_{v:[0,1]\times\mathcal{X}\rightarrow\mathcal{X}}\left[\int^{1}_{0}\int_{ \mathcal{X}}\alpha\lVert v_{t}(x)\rVert^{2}d\rho_{t}dt;\ \ \frac{\partial\rho_ {t}}{\partial t}+\nabla\cdot(v_{t}\rho_{t})=0,\ \rho_{0}=\mu,\ \rho_{1}=\nu \right].$ (6) $$

This dynamic formulation of OT (Eq. 6) is called the Benamou-Brenier formulation (Benamou & Brenier, ). Note that the dynamic transport plan $\{\rho_{t}\}_{0\leq t\leq 1}$ evolves from $\mu$ to $\nu$, and this evolution is governed by the ODE $dx/dt=v_{t}(x)$ through the continuity equation.

Interestingly, the optimal solution of this dynamic OT problem has a simple form. Along each ODE trajectory $\{x(t)\mid dx/dt=v_{t}(x),x(0)=x_{0}\}$, the velocity field $v$ remains constant. Moreover, when the deterministic optimal transport map $T^{\star}$ exists, the following holds:

$$ $\rho_{t}^{dis}:=\left[(1-t)\cdot Id+t\cdot T^{\star}\right]_{\#}\mu\qquad\text {and}\qquad\rho_{t}^{dis}=\rho_{t}^{\star}\qquad\text{ for }\,\,0\leq t\leq 1.$ (7) $$

where $Id$ denotes the identity map and $\rho^{\star}$ denotes the optimal dynamic OT plan of Eq. 6. $\rho_{t}^{dis}$ is called McCann’s Displacement Interpolation (DI) (McCann, ). Eq. 7 shows that the pushforward of linear interpolation between $Id$ and the static OT map $T^{\star}$ is equivalent to the dynamic OT plan $\rho_{t}^{\star}$. Hence, throughout this paper, we simply denote the displacement interpolation as $\rho_{t}^{\star}$. Furthermore, it is well known that the displacement interpolants satisfy the following property (Theorem 7.21 in Villani et al. ()):

$$ $\rho^{\star}_{t}=\arg\inf_{\rho}\mathcal{L}_{DI}(t,\rho)\quad\textrm{where} \quad\mathcal{L}_{DI}(t,\rho)=(1-t)W^{2}_{2}(\mu,\rho)+tW^{2}_{2}(\rho,\nu).$ (8) $$

This property will be utilized in Sec 3 to derive our approach to neural optimal transport, i.e., learning the optimal transport map $T^{\star}$ with a neural network.

## 3 Method

In this section, we 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 Sec 3.1, we derive our two main theorems for deriving our learning objective. In Sec 3.2, we describe how we implement our DIOTM model based on these theoretical results.

### 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 3.1) and the Relationship between Interpolation Potential Functions (Thm 3.3). Thm 3.1 will be utilized to derive our main learning objective (Eq. 18) and Thm 3.3 will be employed to formulate our regularizer (Eq. 20).

#### Dual Problem of Displacement Interpolant

Here, we begin with the dual formulation of the minimization characterization of the interpolant $\rho_{t}^{\star}$ (Eq. 8). The dual problem can be expressed as Eq. 9 in the following Theorem 3.1. Note that while the primal formulation (Eq. 8) optimizes over the set of probability distribution $\rho$, the dual formulation optimizes over two potential functions $f_{1,t},f_{2,t}$.

###### Theorem 3.1 .

Given the assumptions in Appendix A, for a given $t\in(0,1)$, the minimization problem $\inf_{\rho}\mathcal{L}_{DI}(t,\rho)$ (Eq. 8) is equivalent to the following dual problem:

$$ $\sup_{f_{1,t},f_{2,t}\,\,\textrm{with}\,\,(1-t)f_{1,t}+tf_{2,t}=0}\left[(1-t) \int_{\mathcal{X}}{f^{c}_{1,t}(x)d\mu(x)}+t\int_{\mathcal{Y}}{f^{c}_{2,t}(y)d \nu(y)}\right].$ (9) $$

where the supremum is taken over two potential functions $f_{1,t}:\mathcal{Y}\rightarrow\mathbb{R}$ and $f_{2,t}:\mathcal{X}\rightarrow\mathbb{R}$, which satisfy $(1-t)f_{1,t}+tf_{2,t}=0$.
Note that the assumptions in Appendix A guarantee the existence and uniqueness of displacement interpolants $\rho_{t}^{\star}$, the forward optimal transport map $\overrightarrow{T}^{\star}_{t}$ from $\mu$ to $\rho^{\star}_{t}$, and the backward transport map $\overleftarrow{T}^{\star}_{t}$ from $\nu$ to $\rho^{\star}_{t}$. Based on this, we have the following:

$$ $\displaystyle\overrightarrow{T}^{\star}_{t}(x)\in{\rm{arginf}}_{y\in\mathcal{Y }}\left[c(x,y)-f_{1,t}^{\star}(y)\right]\quad\mu\text{-a.s.}.$ (10) $\displaystyle\overleftarrow{T}^{\star}_{t}(y)\in{\rm{arginf}}_{x\in\mathcal{X} }\left[c(x,y)-f_{2,t}^{\star}(x)\right]\quad\nu\text{-a.s.}.$ (11) $$

By applying the same $c$-transfrom parametrization as in Eq. 4, we derive the following max-min formulation of the dual problem:

$$ $\sup_{f_{1,t},f_{2,t}\,\,\textrm{with}\,\,(1-t)f_{1,t}+tf_{2,t}=0}(1-t)\int_{ \mathcal{X}}\inf_{\overrightarrow{T}_{t}}\left(c(x,\overrightarrow{T}_{t}(x))- f_{1,t}(\overrightarrow{T}_{t}(x))\right)d\mu(x)\\ +t\int_{\mathcal{Y}}\inf_{\overleftarrow{T}_{t}}\left(c(\overleftarrow{T}_{t}( y),y)-f_{2,t}(\overleftarrow{T}_{t}(y))\right)d\nu(y).$ (12) $$

Note that the max-min formulation above requires optimization over four functions: the two potentials $f_{1,t},f_{2,t}$ and the two transport maps $\overrightarrow{T}_{t},\overleftarrow{T}_{t}$. Here, we provide an intuitive interpretation of this max-min optimization (Eq. 12).
Two transport maps $\overrightarrow{T}_{t},\overleftarrow{T}_{t}$ act as generators of the interpolant $\rho_{t}^{\star}$ from $\mu$ and $\nu$, respectively. Two potential functions $f_{1},f_{2}$ serve similar roles to discriminators for the generated samples (fake samples) from $\overrightarrow{T}_{t},\overleftarrow{T}_{t}$. The discriminator values for the true samples $\rho_{t}^{\star}$ cancels out because of the potential condition $(1-t)f_{1,t}+tf_{2,t}=0$. Formally,

$$ $(1-t)\int f_{1,t}(z)\,d\rho_{t}^{\star}(z)+t\int f_{2,t}(z)\,d\rho_{t}^{\star} (z)=\int\left[(1-t)f_{1,t}+tf_{2,t}\right](z)\,d\rho_{t}^{\star}(z)=0.$ (13) $$

In fact, using this potential condition, we can combine these two potentials into a single value function (or potential function) $V_{t}$. This simplified formulation of Eq. 9 will be used to derive our regularizer in Theorem 3.3.

###### Corollary 3.2 .

For a given $t\in(0,1)$, let $f_{1,t}(y)=tV_{t}(y)$ and $f_{2,t}(x)=-(1-t)V_{t}(x)$ for some value function $V_{t}:\mathcal{X}=\mathcal{Y}\rightarrow\mathbb{R}$. Then, the dual formulation of displacement interpolation (Eq. 9) can be rewritten as follows:

$$ $\sup_{V_{t}}\left[\int_{\mathcal{X}}{V^{c_{0,t}}_{t}(x)d\mu(x)}+\int_{\mathcal {Y}}{(-V_{t})^{c_{t,1}}(y)d\nu(y)}\right],$ (14) $$

where $c_{s,t}(x,y)=\frac{\alpha\lVert x-y\rVert^{2}}{t-s}$ for every $0\leq s<t\leq 1$.

###### Theorem 3.1 .

###### Corollary 3.2 .

#### Relationship between Interpolation Potential Functions

Here, based on Corollary 3.2, we derive the optimality condition, that each interpolant value function $V_{t}$ satisfies, as the time-dependent value function $V(t,x)=V_{t}(x)$. From now on, we will denote the value function in its time-dependent form: $V(t,x):(0,1)\times\mathcal{X}=\mathcal{Y}\rightarrow\mathbb{R}$.

###### Theorem 3.3 .

Given the assumptions in Appendix A, the optimal $V^{\star}_{t}$ in Eq. 14 satisfies the following:

$$ $V^{\star}_{t}=\arg\sup_{V_{t}}\left[\int_{\mathcal{X}}{V^{c_{0,t}}_{t}(x)d\mu( x)}+\int_{\mathcal{Y}}{V_{t}(x)d\rho^{\star}_{t}(x)}\right],$ (15) $$

up to constant $\rho^{\star}$-a.s.. Moreover, there exists $\{V^{\star}_{t}\}_{0\leq t\leq 1}$ that satisfies Hamilton-Jacobi-Bellman (HJB) equation, i.e.

$$ $\partial_{t}V^{\star}_{t}+\frac{1}{4\alpha}\|\nabla V^{\star}_{t}\|^{2}=0, \qquad\rho^{\star}\text{-a.s.}$ (16) $$

In Sec 3.2, we use this HJB optimality condition (Eq. 16) as a regularizer for the value function $V$ in our model.
In Sec 5, we demonstrate that, when combined with our displacement interpolation, this regularizer significantly improves the stability of optimal transport map estimation.

Figure: Algorithm 1 Training algorithm of DIOTM

###### Theorem 3.3 .

### 3.2 Displacement Interpolation Optimal Transport Model

In this subsection, we introduce our model, called Displacement Interpolation Optimal Transport Model (DIOTM). The goal of our model is to learn the optimal transport map between the source distribution $\mu$ and the target distribution $\nu$. DIOTM is trained by utilizing the displacement interpolation $\rho_{t}^{\star}$ (Eq. 7) between these two distributions. By leveraging the optimal intermediate dynamics connecting two distributions, our model achieves a stable estimation of the optimal transport map.

#### Paramterization of the DI Dual Problem

Our DIOTM learns the static optimal transport map between $\mu$ and $\nu$ by exploiting the displacement interpolation (DI) $\rho_{t}^{\star}$, which is the solution of the dynamic optimal transport problem (Villani et al., ).
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 $\mu$ and $\nu$, the max-min dual formulation of the DI dual problem (Eq. 12) applies to each specific time $t\in(0,1)$. In other words, the intermediate transport maps $\overrightarrow{T}_{t},\overleftarrow{T}_{t}$ and the potential $V_{t}$ are defined separately for each $t$.
Therefore, we represent these interpolant generators $\overrightarrow{T}_{t},\overleftarrow{T}_{t}$ through the boundary generators $\overrightarrow{T},\overleftarrow{T}$ by incorporating the optimality condition of DI (Eq. 7).
Specifically, we parametrize the interpolant generators $\overrightarrow{T}_{t}$ and $\overleftarrow{T}_{t}$ as follows:

$$ $\overrightarrow{T}_{t}(x)=(1-t)x+t\overrightarrow{T}_{\theta}(x),\quad \overleftarrow{T}_{t}(y)=(1-t)y+t\overleftarrow{T}_{\theta}(y)\quad\textrm{ for }t\in(0,1).$ (17) $$

where $\overrightarrow{T}_{\theta}$ and $\overleftarrow{T}_{\theta}$ parametrize the optimal transport maps from the source $\mu$ to the target $\nu$ and from $\nu$ to $\mu$, respectively. The optimality condition of DI clarifies how the intermediate transport maps are related to each other and allows us to parametrize the entire interpolant $\rho_{t}^{\star}$ for $t\in(0,1)$ using just two networks $\overrightarrow{T}_{\theta}:\mathcal{X}\rightarrow\mathcal{Y}$ and $\overleftarrow{T}_{\theta}:\mathcal{Y}\rightarrow\mathcal{X}$. Note that we already investigated how the potentials, specifically the value function $V(t,x)$, are related in Theorem 3.3. Therefore, by combining Eq 14 and 17, we arrive at our main max-min learning objective:

$$ $\mathcal{L}_{\phi,\theta}^{DI}=\sup_{V_{\phi}}\int_{\mathcal{X}}\,\,\inf_{ \overrightarrow{T}_{\theta}}\mathbb{E}_{t}\left[\frac{\alpha}{t}\|x- \overrightarrow{T}_{t}(x)\|^{2}-V_{\phi}(t,\overrightarrow{T}_{t}(x))\right]d \mu(x)\\ +\int_{\mathcal{Y}}\,\,\inf_{\overleftarrow{T}_{\theta}}\mathbb{E}_{t}\left[ \frac{\alpha}{1-t}\|\overleftarrow{T}_{t}(y)-y\|^{2}+V_{\phi}(t,\overleftarrow {T}_{t}(y))\right]d\nu(y).$ (18) $$

Here, $\alpha$ indicates the cost intensity hyperparameter, i.e., $c(x,y)=\alpha\lVert x-y\rVert_{2}^{2}$.
Note that the expectation with respect to $t$ is for aggregating the interpolant $\rho_{t}^{\star}$ for $t\in(0,1)$.
When represented with our transport map parametrizations $\overrightarrow{T}_{\theta}$ and $\overleftarrow{T}_{\theta}$ (Eq. 17), our main learning objective $\mathcal{L}_{\phi,\theta}^{DI}$ can be expressed as follows:

$$ $\mathcal{L}_{\phi,\theta}^{DI}=\sup_{V_{\phi}}\int_{\mathcal{X}}\,\,\inf_{ \overrightarrow{T}_{\theta}}\mathbb{E}_{t}\left[\alpha t\|x-\overrightarrow{T} _{\theta}(x)\|^{2}-V_{\phi}(t,\overrightarrow{T}_{t}(x))\right]d\mu(x)\\ +\int_{\mathcal{Y}}\,\,\inf_{\overleftarrow{T}_{\theta}}\mathbb{E}_{t}\left[ \alpha(1-t)\|\overleftarrow{T}_{\theta}(y)-y\|^{2}+V_{\phi}(t,\overleftarrow{T }_{t}(y))\right]d\nu(y).$ (19) $$

Moreover, we introduce the HJB regularizer $\mathcal{R}(V_{\phi})$, which is derived from the HJB optimality condition of the value function, proved in Theorem 3.3:

$$ $\mathcal{R}(V_{\phi})=\mathbb{E}_{t,x\sim\rho_{t}}\left|2\alpha\ \partial_{t}V _{\phi}(t,x)+\frac{1}{2}\|\nabla V_{\phi}(t,x)\|^{2}\right|.$ (20) $$

As a result, the learning objective can be summarized as follows:

$$ $\mathcal{L}_{\phi,\theta}=\mathcal{L}_{\phi,\theta}^{DI}+\lambda\mathcal{R}(V_ {\phi}(t,x)).$ (21) $$

where $\lambda>0$ denotes the HJB regularizer intensity hyperparameter.

#### Algorithm

We present our training algorithm for DIOTM (Algorithm 1). Our adversarial training objective $\mathcal{L}_{\phi,\theta}$ updates alternatively between the value function $V_{\phi}$ and the two transport maps $\overrightarrow{T}_{\theta},\overleftarrow{T}_{\theta}$, similar to GAN framework (Goodfellow et al., ). Note that we simplified Algorithm 1 by omitting the non-dependent terms for each neural network. Additionally, we apply the HJB regularizer to both generated distributions, i.e., the forward generated distribution from $\overrightarrow{T}_{\theta}$ and the backward generated distribution from $\overleftarrow{T}_{\theta}$.
Finally, throughout all experiments, we use the uniform distribution for time sampling, i.e., $t\sim 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, if we set $t\sim\delta_{1}$, i.e., $t=1$, the training algorithm becomes similar to OTM (Rout et al., ). In this case, we do not use any displacement interpolation information for $t\in(0,1)$. The only difference is our HJB regularizer.

## 4 Related Work

OT problem addresses a transport map between two distributions that minimizes the predefined cost function.
Starting from the dual formulations (Kantorovich, ; Vacher & Vialard, ; ; Gallouët et al., ), diverse methods for estimating OT Map have been developed based on minimax problem (Liu et al., ; Makkuva et al., ; An et al., ; Fan et al., ; Rout et al., ; Choi et al., ; ; Korotin et al., ).
In particular, Fan et al. (); Rout et al. () and Korotin et al. () derived adversarial algorithm from the semi-dual formulation of OT problem and properly recovered OT maps compared to other previous works (Makkuva et al., ).
Moreover, they provided moderate performance in image generation and image translation tasks for large-scale datasets.
Another line of research focuses on the dynamical formulation of OT problems (Chen et al., ; Shi et al., ; Liu et al., ; ; Neklyudov et al., ; Gushchin et al., ).
Several works (Zhang & Chen, ; Chen et al., ; Shi et al., ; Liu et al., ) use sampling-based approaches, which require numerically simulating ODEs or SDEs during training.
More recently, some methods (Neklyudov et al., ; Gushchin et al., ) have introduced simulation-free techniques, often incorporating adversarial learning strategies.
Since our approach also leverages the dynamical properties of OT, we compare our method with Shi et al. (); Gushchin et al. (), 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 the following perspectives:

- •
In Sec 5.1, we compare our model with the ground-truth solution from POT Flamary et al. () 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 to evaluate the scalability of our model.
- •
In Sec 5.3, 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 R1 regularizer.

For implementation details of experiments, please refer to Appendix B.

**Table 1: Evaluation of optimal transport map on the synthetic datasets between OTM (Rout et al., ) and ours, based on the 2-Wasserstein distance $W_{2}$ and the L2 distance between transport maps.**
| Metric | G$\rightarrow$8G | G$\rightarrow$25G | Moon$\rightarrow$Spiral | G$\rightarrow$Circles |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| OTM | DIOTM | OTM | DIOTM | OTM | DIOTM | OTM | DIOTM |  |
| $W_{2}$ $(\downarrow)$ | 4.93 | 3.72 | 10.09 | 6.49 | 0.40 | 0.55 | 3.96 | 2.34 |
| L2 $(\downarrow)$ | 6.48 | 4.38 | 13.09 | 10.00 | 1.77 | 1.67 | 6.23 | 5.44 |

Figure: Figure 1: Visualization of transport maps $T$ on synthetic datasets. The transport map is visualized as a black line connecting each source sample $x$ to its corresponding generated data $T(x)$.
Refer to caption: extracted/5944077/figures/toy.png

### 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^{\star}$ from the source distribution $\mu$ and the target distribution $\nu$. We assess our model against OTM (Rout et al., ; Fan et al., ) by comparing them to the discrete OT solution from the POT library (Flamary et al., ).
As discussed in Sec 3.2, when the time sampling distribution is set to $\delta_{t=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 a neural optimal transport map in Korotin et al. ().
Also, the discrete OT solution indicates the optimal transport map between empirical distributions, i.e., $\mu=\sum_{i}\delta_{x_{i}},\nu=\sum_{j}\delta_{y_{j}}$, computed via convex optimization. Hence, this discrete OT solution serves as the proxy ground-truth solution of the continuous OT map. We tested our model on four synthetic datasets: Gaussian-to-8Gaussian (G$\rightarrow$8G), Gaussian-to-25Gaussian (G$\rightarrow$25G), Moon-to-Spiral, and Gaussian-to-Circles (G$\rightarrow$Circles).

Fig. 1 visualizes the transport maps and Tab. 1 presents the quantitative evaluation metric results. In Tab. 1, we evaluate each transport map through two metrics. First, we calculate the 2-Wasserstein distance $W_{2}(\overrightarrow{T_{\theta}}_{\#}\mu,\nu)$ between the generated distribution $\overrightarrow{T_{\theta}}_{\#}\mu$ and the target distribution $\nu$.
Second, we evaluate whether the neural optimal transport correctly recovers the optimal pairings $(x,T(x))$. Specifically, we compute the discrete optimal transport $T_{disc}^{\star}$ on test datasets using POT and measure the L2 distance between transport maps, i.e., $\int_{\mathcal{X}}\lVert\overrightarrow{T}(x)-T^{\star}(x)\lVert^{2}_{2}d\mu_{
test}(x)$.
Fig. 1 shows that our DIOTM more accurately approximates the target distribution, as indicated by the smaller distribution error between the orange target data and the green generated data. This is further supported by the quantitative results. In Tab. 1, DIOTM achieves a smaller distribution error ($W_{2}$) on three out of four datasets and consistently better recovers the optimal coupling ($L2$) across all datasets. In summary, our DIOTM provides a better approximation of the optimal transport map $T^{\star}$ compared to OTM.

Figure: Figure 2: Image-to-Image translation results of DIOTM for Wild $\rightarrow$ Cat ($64\times 64$) on AFHQ (Choi et al., ). The left figure shows the source images and the right figure shows the corresponding translated images by DIOTM.
Refer to caption: extracted/5944077/figures/wild2cat_source_main2.png

Figure: Figure 3: Image-to-Image translation results of DIOTM for Male $\rightarrow$ Female ($128\times 128$) on CelebA (Liu et al., ).
Refer to caption: extracted/5944077/figures/male2female_source_main.png

### 5.2 Scalability Evaluation in Image-to-Image Translation Tasks

We assessed our model on several Image-to-Image (I2I) translation benchmarks: Male $\rightarrow$ Female (Liu et al., ) ($64\times 64$), Wild $\rightarrow$ Cat (Choi et al., ) ($64\times 64$), and Male $\rightarrow$ Female (Liu et al., ) ($128\times 128$). Intuitively, the optimal transport map serves as a generator for the target distribution, which maps each input $x$ to its cost-minimizing counterpart $y$. For instance, if we consider Male $\rightarrow$ Female task with the quadratic cost $\|x-y\|_{2}^{2}$, 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 our model with the optimal transport models (NOT (Fan et al., ) and OTM (Rout et al., )) and the entropic optimal transport models (DSBM (Shi et al., ) and ASBM (Gushchin et al., )) on image-to-image (I2I) translation tasks.

Fig. 2 and 3 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 of the input source images. In practice, our DIOTM model trains two transport maps in both directions $\overrightarrow{T}_{\theta}$ and $\overleftarrow{T}_{\theta}$. The results for the reverse image-to-image translation are included in Appendix C.
Tab. 2 provides the quantitative evaluation results. We adopted the FID score (Heusel et al., ) 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 I2I translation benchmarks.
Specifically, our model significantly outperforms other entropic OT-based models with a FID score of 7.40 in the higher resolution case of Male $\rightarrow$ Female ($128\times 128$). While OTM achieved comparable results at a specific hyperparameter of $\lambda=10$, OTM diverges for all other $\lambda\in\{50,100\}$ with $\text{FID}>60$ (Fig. 5). In contrast, our DIOTM consistently maintains a stable performance of $\text{FID}<10$ for $\lambda\in\{50,100\}$ (Fig. 5).
In summary, DIOTM exhibits superior scalability in handling higher-resolution image datasets compared to previous OT Map models such as NOT and OTM.

Figure: Table 2: Image-to-Image translation benchmark results compared to existing neural (entropic) optimal transport models.
Refer to caption: extracted/5944077/figures/lambda.png

### 5.3 Further Analysis

In this subsection, we provide an in-depth analysis of our DIOTM model. Specifically, we demonstrate the stable training dynamics of DIOTM, compare the HJB regularizer $\mathcal{R}(V_{\phi})$ with various regularizers introduced in other neural optimal transport models, and conduct an ablation study on the regularization hyperparameter $\lambda$.

#### Stable Training Dynamics

The previous approaches to neural optimal transport with adversarial learning often suffer from training instability (Choi et al., ). These models tend to diverge after long training and are sensitive to hyperparameters. In contrast, our DIOTM offers stable training dynamics. Fig. 5 visualizes the loss values for the transport map $T_{\theta}$ and the value function $V_{\phi}$ throughout the training process. Note that we visualized these loss values on a $\log_{10}$ scale. In OTM (Rout et al., ), the $T_{\theta}$ loss gradually explodes as training progresses.
This unstable training dynamics has been a 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.

#### Comparison to Various Regularization Methods

We investigate the effect of our HJB regularizer $\mathcal{R}(V_{\phi})$ (Eq. 20) compared to other regularization methods. Specifically, we compare two alternatives: OTM regularizer $\mathcal{R}_{\text{OTM}}$ (Fan et al., ; Rout et al., ) and R1 regularizer $\mathcal{R}_{\text{R1}}$ (Roth et al., ). These regularizers are also introduced to stabilize the training value function $V_{\phi}$. We incorporate these regularizers into our algorithm by modifying the regularization term $\mathcal{R}$ in line 4 of Alg. 1 as follows:

- •
$\mathcal{R}_{\text{OTM}}(t,x_{t},y_{t})=\|\nabla_{y}\left(c_{0,t}(x,x_{t})-V_{
\phi}(t,x_{t})\right)\|+\|\nabla_{y}\left(c_{t,1}(y,y_{t})+V_{\phi}(t,y_{t})
\right)\|$.
- •
$\mathcal{R}_{\text{R1}}(t,x_{t},y_{t})=\|\nabla_{y}V_{\phi}(t,x_{t})\|^{2}+\|
\nabla_{y}V_{\phi}(t,y_{t})\|^{2}$.

We assessed these regularizers on synthetic datasets to measure the accuracy of each neural optimal transport map, as in Tab. 1. Tab. 3 provides the results of each regularization method. Note that because we are comparing different regularizers, a direct comparison under the same $\lambda$ is not meaningful. Instead, we need to focus on the best results of each regularizer and their robustness to the regularization intensity parameter $\lambda$. 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 $\partial_{t}V_{\phi}$. Our time-dependent value function $V_{\phi}(t,x)$ is trained to distinguish the displacement interpolation for each $t$. Therefore, regularizing the behavior of our value function across time $t$ is beneficial for DIOTM.

#### Effect of the Regularization Hyperparameter λ 𝜆 \lambda italic_λ

Finally, we conducted an ablation study on the regularization hyperparameter $\lambda$ (Eq. 21) in the image-to-image translation tasks of Wild $\rightarrow$ Cat ($64\times 64$) and Male $\rightarrow$ Female ($128\times 128$). Note that, unlike Tab. 3, we compare our model with OTM (Rout et al., ), not the DIOTM model with the OTM regularizer. Fig. 5 demonstrates that our model exhibits significantly greater stability regarding the regularization hyperparameter $\lambda$, in comparison to OTM. Specifically, in Wild $\rightarrow$ Cat ($64\times 64$), our model maintains decent performance for $\lambda\in\{1,50,100\}$, showing the FID score around 20 even in the worst case. In contrast, the FID scores of OTM fluctuate severely from below $20$ to around $60$ depending on $\lambda$.

## 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 to 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 is the requirement to train both bidirectional transport maps $\overrightarrow{T}$ and $\overleftarrow{T}$. In image-to-image translation tasks, both transport maps are meaningful, e.g. Male $\leftrightarrow$ Female. However, in generative modeling, the reverse transport map $\overleftarrow{T}$ from the data to the prior noise distribution is not always necessary. In these cases, training $\overleftarrow{T}$ can be an unnecessary cost.

#### Ethics Statement

In this paper, we aim to advance the field of machine learning by developing a relatively stable and high-performing algorithm. We hope that further investigation will enable our methodology to be applied as a stabilized algorithm across various machine learning applications. By providing a 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 of 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.

## Appendix A Proofs

#### Assumptions

Let $\mathcal{X}\in\mathbb{R}^{d}$ be a closure of a connected bounded convex open subspace of $\mathbb{R}^{d}$. Let $\text{dim}(\partial\mathcal{X})<d$. Let $\mathcal{P}_{2,ac}$ be a collection of probability distributions on $\mathcal{X}$ that are absolutely continuous and have finite second moments. Let $\mu,\nu\in\mathcal{P}_{2,ac}$.
Moreover, let $c(x,y)=\tau\lVert x-y\rVert^{2}$.
Here, $\tau$ is a given positive constant.

### A.1 Proof of Theorem 3.1

###### Theorem A.1 .

Given the assumptions in Appendix A, for a given $t\in(0,1)$, the minimization problem $\inf_{\rho}\mathcal{L}_{DI}(t,\rho)$ (Eq. 8) is equivalent to the following dual problem:

$$ $\sup_{f_{1,t},f_{2,t}\,\,\textrm{with}\,\,(1-t)f_{1,t}+tf_{2,t}=0}\left[(1-t) \int_{\mathcal{X}}{f^{c}_{1,t}(x)d\mu(x)}+t\int_{\mathcal{X}}{f^{c}_{2,t}(y)d \nu(y)}\right].$ (22) $$

where the supremum is taken over two potential functions $f_{1,t}:\mathcal{Y}\rightarrow\mathbb{R}$ and $f_{2,t}:\mathcal{X}\rightarrow\mathbb{R}$, which satisfy $(1-t)f_{1,t}+tf_{2,t}=0$.
Note that the assumptions in Appendix A guarantee the existence and uniqueness of displacement interpolants $\rho_{t}^{\star}$, the forward optimal transport map $\overrightarrow{T}^{\star}_{t}$ from $\mu$ to $\rho^{\star}_{t}$, and the backward transport map $\overleftarrow{T}^{\star}_{t}$ from $\nu$ to $\rho^{\star}_{t}$. Based on this, we have the following:

$$ $\displaystyle\overrightarrow{T}^{\star}_{t}(x)\in{\rm{arginf}}_{y\in\mathcal{Y }}\left[c(x,y)-f_{1,t}^{\star}(y)\right]\quad\mu\text{-a.s.}.$ (23) $\displaystyle\overleftarrow{T}^{\star}_{t}(x)\in{\rm{arginf}}_{x\in\mathcal{X} }\left[c(x,y)-f_{2,t}^{\star}(x)\right]\quad\nu\text{-a.s.}.$ (24) $$

###### Proof.

We prove the theorem by two steps. First, we prove Eq. 22 following the duality theorem in Korotin et al. () and Kolesov et al. (). Then, we discuss the existence and uniqueness of the potential $f_{1},f_{2}$ and the optimal transport maps. Using these uniquesness, we prove Eq. 23.

###### Theorem A.1 .

###### Proof.

#### Step 1.

We first prove Eq. 22. Suppose $0<t<1$ is given.
By applying the dual form of OT problem Kantorovich (), for every $\rho\in\mathcal{P}_{2}$, the following equation satisfies:

$$ $(1-t)W^{2}_{2}(\mu,\rho)+tW^{2}_{2}(\rho,\nu)\\ =\sup_{f_{1},f_{2}\in\mathcal{C}(\mathcal{X})}\left[(1-t)\mathbb{E}_{x\sim\mu} [f^{c}_{1,t}(x)]+t\mathbb{E}_{y\sim\nu}[f^{c}_{2,t}(y)]+\mathbb{E}_{z\in\rho}[ (1-t)f_{1,t}(z)+tf_{2,t}(z)]\right].$ (25) $$

Then, by applying Theorem 3.4 in Sion (), we obtain the following:

$$ $\displaystyle\begin{split}&\mathcal{L}^{*}=\inf_{\rho}(1-t)W^{2}_{2}(\mu,\rho) +tW^{2}_{2}(\rho,\nu),\\ =&\inf_{\rho}\sup_{f_{1,t},f_{2,t}\in\mathcal{C}(\mathcal{X})}\left[(1-t) \mathbb{E}_{x\sim\mu}[f^{c}_{1,t}(x)]+t\mathbb{E}_{y\sim\nu}[f^{c}_{2,t}(y)]+ \mathbb{E}_{z\in\rho}[(1-t)f_{1,t}(z)+tf_{2,t}(z)]\right],\\ =&\sup_{f_{1,t},f_{2,t}\in\mathcal{C}(\mathcal{X})}\inf_{\rho}\left[(1-t) \mathbb{E}_{x\sim\mu}[f^{c}_{1,t}(x)]+t\mathbb{E}_{y\sim\nu}[f^{c}_{2,t}(y)]+ \mathbb{E}_{z\sim\rho}[(1-t)f_{1,t}(z)+tf_{2,t}(z)]\right].\end{split}$ (26) $$

Note that we can swap minimax to max-min problem due to the compactness assumption of the space $\mathcal{X}$. Now, suppose $m=m(f_{1,t},f_{2,t}):=\inf_{z\in\mathcal{X}}(1-t)f_{1,t}(z)+tf_{2,t}(z)$. Let $\tilde{f}_{1,t}=(m(f_{1,t},f_{2,t})-tf_{2,t})/(1-t)$. Then, $m=(1-t)\tilde{f}_{1,t}(z)+tf_{2,t}(z)\leq(1-t)f_{1,t}(z)+tf_{2,t}(z)$, which implies $\tilde{f}_{1,t}\leq f_{1,t}$. Then, we can easily obtain $f^{c}_{1,t}\leq\tilde{f}^{c}_{1,t}$. With this inequality, we can obtain the following equality:

$$ $\displaystyle\begin{split}&\sup_{f_{1,t},f_{2,t}\in\mathcal{C}(\mathcal{X})} \inf_{\rho}\left[(1-t)\mathbb{E}_{x\sim\mu}[f^{c}_{1,t}(x)]+t\mathbb{E}_{y\sim \nu}[f^{c}_{2,t}(y)]+\mathbb{E}_{z\sim\rho}[(1-t)f_{1,t}(z)+tf_{2,t}(z)]\right ],\\ =&\sup_{f_{1,t},f_{2,t}\in\mathcal{C}(\mathcal{X})}\left[(1-t)\mathbb{E}_{x \sim\mu}[f^{c}_{1,t}(x)]+t\mathbb{E}_{y\sim\nu}[f^{c}_{2,t}(y)]+m(f_{1,t},f_{2 ,t})\right],\\ =&\sup_{f_{2,t}\in\mathcal{C}(\mathcal{X})}\left[(1-t)\mathbb{E}_{x\sim\mu}[ \tilde{f}^{c}_{1,t}(x)]+t\mathbb{E}_{y\sim\nu}[f^{c}_{2,t}(y)]+m(\tilde{f}_{1, t},f_{2,t})\right],\\ =&\sup_{f_{2,t}\in\mathcal{C}(\mathcal{X})}\left[(1-t)\mathbb{E}_{x\sim\mu} \left[\left(-\frac{tf_{2,t}}{1-t}\right)^{c}(x)\right]+t\mathbb{E}_{y\sim\nu}[ f^{c}_{2,t}(y)]\right].\end{split}$ (27) $$

Note that the last equation is obtained by using the following property: $(f_{1,t}+a)^{c}=f^{c}_{1,t}-a$ for any constant $a$.
By letting $f_{1,t}=-\frac{tf_{2,t}}{1-t}$, we finally obtain Eq. 22.

#### Step 2.

In this step, we prove the uniqueness of the optimal potential pair $(f^{\star}_{1,t},f^{\star}_{2,t})$ of Eq. 22.
Let

$$ $\rho^{\star}=\arg\inf_{\rho}(1-t)W^{2}_{2}(\mu,\rho)+tW^{2}_{2}(\rho,\nu).$ (28) $$

First, we should show the well-definedness of above equation. In other words, we should show the existence and uniqueness of $\rho^{\star}$.
Due to the absolutely continuity assumption of $\mu$, there exists unique deterministic optimal transport map between $\mu$ and $\nu$ (See Chapter 1 or Theorem 9.2 of Villani et al. ()). Thus, by applying Corollary 7.23 in Villani et al. (), there exist unique solution $\rho^{\star}$.

Now, consider the following optimization problem:

$$ $\sup_{\hat{f}_{1,t},\hat{f}_{2,t}}\bigg{[}(1-t)\left(\int_{\mathcal{X}}\hat{f} ^{c}_{1,t}(x)d\mu(x)+\int_{\mathcal{X}}\hat{f}_{1,t}(x)d\rho^{\star}(x)\right) \\ +t\left(\int_{\mathcal{X}}\hat{f}^{c}_{2,t}(y)d\nu(y)+\int_{\mathcal{X}}\hat{f }_{2,t}(y)d\rho^{\star}(y)\right)\bigg{]}.$ (29) $$

Trivially, the first and the second term of Eq. 29 break down into two independent optimization problems. By using the Kantorovich duality theorem (See Kantorovich () or Theorem 5.10 Villani et al. ()), the solution value of Eq. 29 is obviously $\mathcal{L}^{*}$.
Moreover, since the solution pair $(f^{\star}_{1,t},f^{\star}_{2,t})$ of Eq. 22 satisfies $(1-t)f^{\star}_{1,t}+tf^{\star}_{2,t}=0$, it is included in the optimal potential pair of Eq. 29.

Note that we assumed that $\mu$ is absolutely continuous and the space $\mathcal{X}$ is convex, and a closure of a connected open set. Then, on the support of $\mu$, the Kantorovich potentials $\hat{f}_{1,t}$ and $\hat{f}_{2,t}$ is unique up to constant on the connected support of $\mu$ and $\nu$, respectively (See Staudt et al. ()).
Therefore, the optimal potentials $f^{\star}_{1,t},f^{\star}_{2,t}$ of Eq. 22 are a Kantorovich potentials of $W^{2}_{2}(\mu,\rho^{\star})$ and $W^{2}_{2}(\nu,\rho^{\star})$, respectively.

Finally, since $\mu\in\mathcal{P}_{2,ac}$, there exists a measurable deterministic optimal transport map $T^{\star}$ which transport $\mu$ to $\rho^{\star}$ (See Chapter 1 or Corollary 9.4 in Villani et al. ()). Then, by Theorem 5.10 and Remark 5.13 in Villani et al. (), we can easily obtain Eq. 23.
∎

Let $f_{1,t}(x)=tV(t,x)$ and $f_{1,t}(x)=-(1-t)V(t,x)$ for every $t\in(0,1)$. Then, we can rewrite Eq. 9 as follows:

$$ $\sup_{V_{t}}\left[\int_{\mathcal{X}}{V^{c_{0,t}}_{t}(x)d\mu(x)}+\int_{\mathcal {X}}{(-V_{t})^{c_{t,1}}(y)d\nu(y)}\right],$ (30) $$

where $c_{s,t}(x,y)=\frac{\tau\lVert x-y\rVert^{2}}{t-s}$ for every $0\leq s<t\leq 1$.

The following theorem shows the connection between the optimal potential of transport problem between $\mu$ to $\rho^{\star}_{t}$, and our potential function $V$.

###### Theorem A.2 .

The optimal $V^{\star}_{t}$ in Eq. 14 satisfies the following:

$$ $V^{\star}_{t}=\arg\sup_{V_{t}}\left[\int_{\mathcal{X}}{V^{c_{0,t}}_{t}(x)d\mu( x)}+\int_{\mathcal{X}}{V_{t}(x)d\rho^{\star}_{t}(x)}\right],$ (31) $$

up to constant $\rho^{\star}$-a.s.. Moreover, there exists $\{V^{\star}_{t}\}_{0\leq t\leq 1}$ that satisfies Hamilton-Jacobi-Bellman (HJB) equation, i.e.

$$ $\partial_{t}V_{t}+\frac{1}{4\tau}\|\nabla V_{t}\|^{2}=0,\ \ \rho^{\star}\text{ -a.s.}$ (32) $$

###### Proof.

By the Step 2 of the proof of the Theorem A.1, we have shown that the optimal $(f^{\star}_{1},f^{\star}_{2})$ which solves Eq. 22 is a Kantorovich dual function of $W^{2}_{2}(\mu,\rho^{\star})$ and $W^{2}_{2}(\nu,\rho^{\star})$, respectively.
By directly applying this fact, we obtain Eq. 31.

Now, we prove Eq. 32, the HJB equation.
Let $\rho^{\star}$ be defined as the Step 2 of the proof of Theorem A.1. Then, as discussed in Step 2, there exists unique $\rho^{\star}$ that satisfies Eq. 28. Since $\mu$ is absolutely continuous, it is well known that $\rho^{\star}$ is also absolutely continuous.
Due to the compactness of the space $\mathcal{X}$, the optimal $V_{t}$ is differentiable with respect to $x\in\mathcal{X}$ $\rho^{\star}$-a.s..
Now, by applying Theorem 7.36 and Remark 7.37 of Villani et al. (), there exists $V:(0,1)\times\mathcal{X}\rightarrow\mathbb{R}$ such that

$$ $-V_{s}(x)=\inf_{y}\left(c_{s,t}(x,y)-V_{t}(y)\right),$ (33) $$

for all $0<s<t<1$. By applying Hopf-Lax formula, we obtain,

$$ $-V(s,X_{s})=\inf_{v:[0,1]\times\mathcal{X}\rightarrow\mathcal{X}}\left[\int^{t }_{s}\tau\lVert v_{t}\rVert^{2}du-V_{t}(X_{t})\right],$ (34) $$

where $\dot{X}_{t}=v_{t}(X_{t})$. Thus, by organizing the Eq. 34, we obtain Eq. 32 as follows:

$$ $\displaystyle 0$ $\displaystyle=\lim_{t\rightarrow s}\ \inf_{v:[0,1]\times\mathcal{X}\rightarrow \mathcal{X}}\left[\frac{1}{t-s}\int^{t}_{s}\tau\lVert v\rVert^{2}-\frac{1}{t-s }\left(V_{t}(X_{t})-V_{s}(X_{s})\right)\right],$ (35) $\displaystyle=\inf_{v}\left(\tau\lVert v_{s}\rVert^{2}-\partial_{s}V_{s}- \nabla V_{s}\cdot v_{s}\right)=-\left(\partial_{s}V_{s}+\frac{\lVert\nabla V_{ s}\rVert^{2}}{4\tau}\right).$ (36) $$

Note that $\text{law}(X_{t})=\rho^{\star}_{t}$. Thus, $\partial_{t}V_{t}+\frac{\lVert\nabla V_{t}\rVert^{2}}{4\tau}=0$ $\rho^{\star}_{t}$-a.s..
∎

###### Theorem A.2 .

###### Proof.

## Appendix B Implementation Details

For all experiments, we parametrize $\overrightarrow{T}_{\theta},\overleftarrow{T}_{\theta}:\mathcal{X}\times
\mathbb{R}^{n}\rightarrow\mathcal{X}$ where $x\in\mathcal{X}$ is the samples from given datasets, and $z\in\mathbb{R}^{n}$ is an auxiliary variable. As reported in Choi et al. (; ), the auxiliary slightly improves the performance. For all OTM Fan et al. (); Rout et al. () experiments, we use the same network and the same parameter unless otherwise stated.

### B.1 2D Experiments

#### Data Description

In this paragraph, we describe our synthetic datasets:

- •
8-Gaussian: For $m_{i}=12\left(\cos{\frac{i}{4}\pi},\sin{\frac{i}{4}\pi}\right)$ where $i=0,1,\dots,7$ and $\sigma=0.4$, the distribution is defined as the mixture of $\mathcal{N}(m_{i},\sigma^{2})$ with an equal probability.
- •
25-Gaussian: For $m_{ij}=\left(8i,8j\right)$ where $i=-2,-1,\dots,2$ and $j=-2,-1,\dots,2$, the distribution is defined as the mixture of $\mathcal{N}(m_{ij},\sigma^{2})$ where $\sigma=0.01$.
- •
Moon to Spiral: We follow Choi et al. ().
- •
Two Circles: We first uniformly sample from the circles of radius $8$ and $16$ with the center at origin. Then, we add Gaussian noise with standard deviation of $0.2$.

#### Network Architectures

We first describe the value function network $V_{\phi}(t,x)$. The input $x\in\mathcal{X}$ is embedded using a two-layer MLP with a 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 are then added and passed through a three-layer MLP. The SiLU activation function is applied to all MLP layers.
Besides, for the transport map networks, we employ the same network to Choi et al. () with hidden dimension of 128.

#### Training Hyperparameters

We set Adam optimizer with $(\beta_{1},\beta_{2})=(0,0.9)$, learning rate of $10^{-4}$ and the number of iteration of 120K.
We set $\alpha=0.1$ and $\lambda=1$.

#### Discrete OT Solver

We used the POT library Flamary et al. () to obtain an accurate transport plan $\pi_{pot}$.
We used 1000 training samples for each dataset in estimating $\pi_{pot}$ to sufficiently reduce the gap between the true continuous measure and the empirical measure.

### B.2 Image Translation

#### Training Hyperparameters

We follow the large neural network architecture introduced in Xiao et al. ().
We use Adam optimizer with $(\beta_{1},\beta_{2})=(0,0.9)$, learning rate of $10^{-4}$, and trained for 60K iterations. We use a cosine scheduler to gradually decrease the learning rate from $10^{-4}$ to $5\times 10^{-5}$. The batch size of 64 and 32 is employed for 64x64 and 128x128 image datasets, respectively. We use $\alpha=0.001$ for CelebA image dataset, and $\alpha=0.0005$ for AFHQ dataset. We use ema rate of 0.9999 for 64x64 image datasets and 0.999 for 128x128 image datasets.

#### Evaluation Metric

For the Wild$\rightarrow$Cat experiments, we compute FID based on train datasets of the target dataset and generate samples from each source samples. This follows the implementation of De Bortoli et al. (). We sampled ten generated samples from each source.
For CelebA experiments, we follow the metric used in Korotin et al. (); Gushchin et al. (). Specifically, we compute FID through test datasets.

## Appendix C Additional Qualitative Results

In this section, we include some qualitative results on image-to-image translation tasks. We visualize the source images and its transported samples.

Figure: Figure 6: Unpaired male $\rightarrow$ female translation for 64 × 64 CelebA image.
Refer to caption: extracted/5944077/figures/appendix/male2female64_source.png

Figure: Figure 7: Unpaired female $\rightarrow$ male translation for 64 × 64 CelebA image.
Refer to caption: extracted/5944077/figures/appendix/male2female64_target.png

Figure: Figure 8: Unpaired female $\rightarrow$ male translation for 128 × 128 CelebA image.
Refer to caption: extracted/5944077/figures/appendix/male2female128_source.png

Figure: Figure 9: Unpaired male $\rightarrow$ female translation for 128 × 128 AFHQ image.
Refer to caption: extracted/5944077/figures/appendix/male2female128_target.png

Figure: Figure 10: Unpaired wild $\rightarrow$ cat translation for 64 × 64 AFHQ image.
Refer to caption: extracted/5944077/figures/appendix/wild2cat64_source.png

Figure: Figure 11: Unpaired cat $\rightarrow$ wild translation for 64 × 64 AFHQ image.
Refer to caption: extracted/5944077/figures/appendix/wild2cat64_target.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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:

1) What is the point of introducing alpha? The OT map/displacement maps should be the same for all alpha, right?

2) 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.

3) 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.

4) 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.

5) 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)?
soundness: 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.

2) 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.

3) 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).

4) The text is overall readable and the clarity is ok (although sometimes the amount of the bolded text is too annoying).
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).

The 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

1) 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.

2) 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.

The 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).
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.

2) 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.

3) 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.

4) 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].

5) 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?

References

[1] Agueh, M., & Carlier, G. (2011). Barycenters in the Wasserstein space. SIAM Journal on Mathematical Analysis, 43(2), 904-924.

[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.

[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.

[4] Amos, B. On amortizing convex conjugates for optimal transport. In The Eleventh International Conference on Learning Representations.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: Questions:
1. 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?  
2. Can you clarify how long it took you to train your methods for image-to-image translation problems compared to competing methods? 
3. What is the number of parameters used by all the methods for image-to-image translation problems? Are they comparable?
4. Have you experienced failures of your method, and if so, can you provide them?
5. 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?
Typing errors:
1. Line 49 - double "the"; one should be deleted.
soundness: 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.
2. 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. 
3. 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.
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.
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.
2. 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.
3. 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).

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: Refer to weaknesses.
soundness: 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.

2. 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.
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.
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.

2. 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.

3. 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?

4. 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?

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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?
2. 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. 
3. 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?
Minor:
- Line 243: typo - parametrization
- 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. 
- Repeated citation for Diffusion Schrodinger bridge matching paper.
soundness: 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.  

**Quality and significance**: Improving stability of Optimal transport is an important problem and this paper proposes a method to addresses it. 
- The experimental results on simple 2D toy datasets seem to indicate improved performance compared to prior methods as indicated in Table 1.  
- DIOTM seems to outperform other optimal transport based models on image-to-image translation task in terms of metrics such as FID (Table 2).
- 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.
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.
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. 
2. 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.
3. 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. 
4. 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.
5. 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. 

[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.
[2] Wang, Tengfei, et al. "Pretraining is all you need for image-to-image translation." arXiv preprint arXiv:2205.12952 (2022).

Meta-review: The reviewers agree that the paper presents a novel and interesting method for solving the optimal transport mapping problem using neural networks.  They appreciate the idea of exploiting displacement interpolation and the theoretically justified HJB-based regularization.  The experimental results on image-to-image translation are also considered convincing.

However, there are some concerns about the theoretical validity of the proposed semi-dual form and the limited scope of the method.  The reviewers also point out that some of the results are not completely novel and miss related work.

Decision: Accept (Poster)