Neural Optimal Transport with General Cost Functionals
Authors: Arip Asadulaev
Alexander Korotin
Vage Egiazarian
Petr Mokrov
Evgeny Burnaev
Venue: ICLR 2024 poster (2024)
Keywords: Optimal Transport
Neural Networks
Generative Modelling
Primary area: generative models

Abstract: We introduce a novel neural network-based algorithm to compute optimal transport (OT) plans for general cost functionals. In contrast to common Euclidean costs, i.e., $\ell^1$ or $\ell^2$, such functionals provide more flexibility and allow using auxiliary information, such as class labels, to construct the required transport map. Existing methods for general cost functionals are discrete and do not provide an out-of-sample estimation. We address the challenge of designing a continuous OT approach for general cost functionals in high-dimensional spaces, such as images. We construct two example functionals: one to map distributions while preserving the class-wise structure and the other one to preserve the given data pairs. Additionally, we provide the theoretical error analysis for our recovered transport plans. Our implementation is available at \url{https://github.com/machinestein/gnot}

Published as a conference paper at ICLR 2024
NEURAL OPTIMAL TRANSPORT WITH GENERAL COST FUNCTIONALS
Arip Asadulaev* Alexander Korotin 2,1 Vage Egiazarian 4.5 Petr Mokrov Evgeny Burnaev Artificial Intelligence Research Institute Skolkovo Institute of Science and Technology Moscow Institute of Physics and Technology *HSE University "Yandex ar wiipsadulaeeeaii.... net, a wwk.koorrnnssse. .okottinskolteehh.ccm
ABSTRACT
We introduce novel neural network-basee algorithm compute optimal transport (OT) plans for general cost functionals. contrast to common Euclidean costs, i.e., f or f2, such functionals provide more flexibility and allow using auxiliary informa- tion, such as class labels, to construct the required transport map. Existing methods for general cost functionals are discrete and do not provide an out-of-sample estima- tion. We address the challenge of designing a continuous OT approach for general cost functionals in high-dimensional spaces, such as images. We construct two example functionals: one to map distributions while preserving the class-wise struc- ture and the other one to preserve the given data pairs. Additionally, we provide the theoretical error analysis for our recovered transport plans. Our implementation is available ht tps: //gi//g///t //ggttttung withub.... com/ www.///anhinnnnnnnnnoo..nm
Figure 1: Results of method with the pair-guided cost functional ($ applied to the supervised amage-to-image translation task (Celeba-MaskHQ dataset, 256 x 256 images).
INTRODUCTION
Optimal transport (OT) a powerful framework to solve mass-moving problems for data distribu- tions which finds many applications in machine learning and computer vision (Bonneel & Digne 2023). Most existing methods to compute OT
Transport map
Q
Z
plans designed for discrete distributions (Fla-
mary et al., 2021| Cuturi 2013). These methods have good flexibility and allow control the
Unlabeled
IP
properties of the plan (Peyré et al. 2019). How- ever, discrete methods find optimal matching between two given (train) sets which does not generalize new (test) data points. This limits the use of discrete OT plan methods in scenar- ios where new data needs be generated, e.g., image-to-image transfer (Zhu al.) 2017).
Figure 2: Dataset transfer problem. Input P= Σn on, F n. target Q Σn, 3,Q, distributions are mixtures of N classes. The task is to learn a transport map T preserving the class. The learner has the access to labeled input data F. and only partially labeled target data Q
"qqual contribution.

---

Published as a conference paper at ICLR 2024
Recent works (Rout et al.) 2022 Korotin al. 2023b
2021b
Fan et al. 2023 Daniels al.
2021] propose continuous methods to compute OT plans. Thanks to employing neural networks to parameterize OT solutions, the learned transport plan can be used directly as the generative model in data synthesis (Rout et al., 2022) and unpaired learning (Korotin et al. 2023b Rout et al. 2022
Daniels et al. 2021 Gazdieva et al. 2022).
Existing continuous OT methods mostly focus on classic cost functions such as e² (Korotin et al.
2021b 2023b, [an al. 2023 Gazdieva al. 2022) which estimate the closeness of input and
output points. However, choosing such costs for problems where a specific optimality of the mapping required may be challenging. For example, when one needs to preserve the object class during
the transport (Figure common e² cost may be suboptimal (Su et al. 2022 Appendix C). (Daniels) al. [02T Figure 3). This limitation could be fixed by considering general cost functionals (Paty
& Cuturi 2020) which may take into account additional information, e.g., class labels.
Despite the large popularity of OT, the approach for continuous OT with general cost function- als (general OT) is still missing. We address this limitation. The main contributions of our paper are:
1. We show that the general OT problem (#2) can be reformulated as a saddle point optimization problem, which allows to implicitly recover the OT plan ($4.1) in the continuous setting. The (Algorithm 2
problem can be solved with neural networks and stochastic gradient methods (
We provide the error analysis of solving the proposed saddle point optimization problem via the duality gaps, i.e., errors for solving inner and outer optimization problems ($4.
3. We construct and test examples of general cost functionals for mapping data distributions with the preservation of the class-wise (5 5.1 Algorithm
| and paired data structure ($5.2 Algorithm B)
From the theoretical perspective, our max- min reformulation is generic and subsumes previously known reformulations for classic (Rout et al. 2022 Fan et al. (023) and weak (Korotin et al.] (022b) OT. Furthermore, existing error analysis works exclusively with the classic OT and operate only under certain restrictive assumptions such as the the convexity of the dual potential. Satisfying these assumptions in practice leads severe performance drop (Korotin et al. 2021c| Figure 5a). In contrast, our error analysis is free from assumptions on the dual variable and, besides general OT, it applicable to weak OT for which there is currently no existing error analysis.
From the practical perspective, we apply our method to the dataset transfer problem (Figure 2 previously not solved using continuous optimal transport. This problem arises when is necessary to repurpose fixed or black-box models to classify previously unseen partially labelled target datasets with high accuracy by mapping the data into the dataset on which the classifier was trained (Alvarez. Melis & Fusil 2021 ). Our method achieves notable improvements in accuracy over existing algorithms. Also, we show the performance of our method on the supervised image-to-image translation task.
Notations. For a compact Hausdorff space S, we use P(S) to denote the set of Borel probability distributions on S. We denote the space of continuous R-valued functions on endowed with the supremum norm by C(S). Its dual space is the space M(S) O P(S) of finite signed Borel measures over S. Let X, be compact Hausdorff spaces and TP E P(X). Q E P())). We use X' x Y, which projection onto the first marginal is P. We use (I,,, Q) C [I(P) to denote the subset of probability distributions
I(P) C P(X y) denote the subset of probability distributions
(transport plans) X' x y with marginals IP, Q. For a measurable map T X x Z + V, we denote the associated push-forward operator by Tч.
2 BACKGROUND
In this section, we provide key concepts of the optimal transport theory. Throughout the paper, we consider compact X' = y C D and IP, Q E P(X),,
Classic and weak OT. For a cost function e C(X )), the OT cost between PP, Q is
Cost(P, def inf
rEII(P.Q) x c(x, (x)
(1)
see Villani 2008 @)). We call T the classic OT. Problem [ admits a minimizer x² E [IP, Q) which is called an OT plan (Santambrogio) 2015 Theorem 1.4). It may be not unique (Peyré et al. 2019 Remark 2.3). Intuitively, the cost function c(æ, y) measures how hard it is to move mass piece between points I E X and y E V That is, π shows how to optimally distribute the mass of IP Q, i.e, with minimal effort. For cost functions c(x, || yll2 and c(x; y) 2 | v
2

---

Published as a conference paper at ICLR 2024
the OT cost ① is called the Wasserstein- 1 (W1) and the (square of) Wasserstein-2 (W22 distance, respectively, see Villani 2008 $1) or (Santabrogio 2015 81, 2).
Recently, classic OT obtained the weak OT extension (Gozlan et al. 2017 Backhoff- Veraguua et al.
2019). Consider C : X x P()) + IR, i.c., a weak cost function whose imputs are a point x E and distribution of y y. The weak OT cost
inf / C(1, =Ell(P.Q)
Cost(IP,
where πt (-I1a) denotes the conditional distribution. Weak formulation (2) is reduced to classic formu- lation ①T when C(x, p1) = fy c(z, y)du(y). Another example of a weak cost function is the y-weak quadratic cost C( ()) Sy 2 o y (d)()) 3 Var(()), where 7 2 0 and Var(μ) is the variance of μr, see (Korotin et al. 2023b] Eq. 5). (Alibert et al. 2019, 25. (Gozlan & Juillet 2020) $5.2) for details. For this cost, we denote the optimal value ( by W and call 7-weak Wasserstein-2.
Regularized and general OT. The expression inside ① is linear functional. It is common to add lower cemi-continuous convex regularizer .M x) + R U {oo} with weight y > O:
{// c(2, y)dT + *EI((..))
Cost(P,
Q
def
inf
(3)
Regularized OT formulation typically provides several advantages over original formulation (T) For example, if (TT)) is strictly convex, the expression inside (3) a strictly convex functional in IT and yields the unique OT plan x². Besides, regularized OT typically has better sample complexity
(Genevay 2019. Mena Niles- Weed, 2019) Genevay et al., 2019). Common regularizers are the
entropic (Cuturi 2013). quadratic (Essid & Solomon, 2018), lasso (Courty et al. 2016), etc.
To consider a general OT formulation, let F MMXX y) R {+oo] be a convex lower semi-continuous functional. Assume that there exists [[( Q) for which F(r) < co. Let
(P, def inf =EIT(P,))
F(r).
This problem a generalization of classic OT weak OT 2 and regularized OT 3 Following
(Paty & Cuturi 2020), we call problem ( general OT problem. It admits minimizer (OT plan) π
(aty & Cuturi 2020 Lemma 1). One may note that regularized OT 33) represents a similar problem:
is enough to put c( y) = 0, 7 = 1 and (R(T)) F(x) to obtain (4) from 33. i.e., regularized
and general OT (4) can be viewed as equivalent formulations.
RELATED WORK: DISCRETE AND CONTINUOUS OT SOLVERS
Solving OT problems usually implies either finding an OT plan म' or the OT cost. Many approaches in generative learning use OT cost as the loss function to update generative models, such as WGANs
(Arjovsky& & Bottou| 2017 Petzka et al., 2018 Liu et al.) 2019). see (Korotin et al., 20226 for related to our work as they do not compute OT plans or maps. Existing computational OT plan methods can be roughly split into two groups: discrete and continuous.
survey. These
0y. Discrete OT considers discrete distributions PNN EN Pn. and QN = -m and aims find the OT plan (]. (2). (4), 33 directly between IP PN and Q QM. In this case, the OT plan n can be represented as a doubly stochastic N x M matrix. For survey of computational methods for discrete OT, we refer to (Peyré al., 2019). In short, one of the most popular the Sinkhorn algorithm (Cuturi 2013) which is designed to solve formulation (3) with the entropic regularization.
General discrete OT is extensively studied Nash 2000 Courty et al. 2016 Flamary et al. 202
Ferradans et al. 2014 Rakotomamonjy et al., 2015); these methods are often employed in domain adaptation problems (Courty et al.] al.] (016). Additionally, the available labels can used to reconstruct the classic cost function, to capture the underlying data structure (Courty et al. 2016; Stuart &
Wolfram 2020 Liu et al. 2020 Li et al. 2019).
The major drawback of discrete OT methods is that they only perform a (stochastic) matching between the given empirical samples and usually do not provide out-of- estimates. This limits their application real-world scenarios where new (test) samples frequently appear. Recent works (Hütter & Rigollet, [202]] Pooladian & Niles- Weed [202 Manole et al. 2021 Deb et al.) (02T) consider the OT problem with the quadratic cost and develop out-of-sample estimators by wavelet/kernel-basee
3

---

Published as a conference paper at ICLR 2024
plugin estimators by the barycentric projection of the discrete entropic OT plan. In spite of tractable theoretical properties, the performance of such methods in high dimensions is questionable.
Continuous OT usually considers Pn d and qn. Mr and assumes that the given discrete distri-
QM= butions PN 1. _n=1 6z,, QM = M Σm= 6ym are the empirical counterparts of the underlying N.
distributions IP, QQ That the goal of continuous OT is recover the OT plan between P, Q which are accessible only by their (finite) empirical samples { .2....., ,NN J ~ P and (V1. : U2..-.... MM ~ Q- In this case, represent the plan one has to employ parametric approximations of the OT plan x² or dual potentials u', which, in turn, provide straightforward out- out-of -am-ol-sample estimates.
A notable development is the use of neural networks compute OT maps for solving weak and classic (1) functionals (Korotin et al., 20233 2022a 2021b Rout et al. [022]] Fan et al. 2023
Henry-Labordes 2019). Previous OT methods were based on formulations restricted to convex
potentials (Makkuva et al.) 2020 Korotin et al. 2021a с Mokrov et al.
2021 Fan et al. 2023 Bunne
et al. | Alvarez- et al. 2022). and used Input Convex Neural Networks (Amos et al. 2017
ICNN) approximate them, which limited the application of OT in large-scale tasks (Korotin et al.
2021b Fan et al. 2022: Korotin et al. 2022a). In (Genevay et al.) 2016; Seguy et al. 2018 Daniels et al.
2021 Fan et al.] 2022), the authors propose methods for -diieegeence regularized costs
While the discrete version of the general OT problem 4 is well studied the literature, its continuous counterpart is not yet analyzed. In our work, we fill this gap by proposing the algorithm to solve the (continuous) general OT problem ($4), provide error bounds ($4.2). As an illustration, we construct examples of general cost functionals which can take into account the available task-specific information as labels ($5.1) or pairs ($5.2).
4
MAXIMIN REFORMULATION OF THE GENERAL OT
In this section, we derive saddle point formulation for the general OT problem which we later solve with neural networks. All the proofs of the statements are given in Appendix
4.1 GENERAL OT MAXIMIN REFORMULATION VIA STOCHASTIC MAPS
In this subsection, we derive the dual form for (4. which can be used to get the OT plan x".
Our formulation utilizes the implicit representation for plans [I(P) via stochastic maps, idea inspired by (Korotin et al., 2023b, $4.1). We introduce a latent space 2 = R² and an atomless distribution S E P(2)
it, e.g., $ = -O Iz). For every function T = Tx X' x Z + y which implicitly represents it. Such TT satisfies (z, :)ES = . (-12x) for all X. That is, given x E X and random latent vector z S, the function T produces sample Tx 2) ~ π (r). In particular, if ~ IP, the random vector [r. Tx, =)] is distributed as т. Thus, every π E II(F) can be implicitly represented (non-uniquely) as a function Tx : & x + Y. And vice-versa, every measurable function T X x Z + ) is an implicit representation of the distribution कy which the joint distribution of random vector [ar, T(x, =)] with P, = S.
P(X' x )), there exists measurable
Our two following theorems constitute the main theoretical idea of our approach. They are proven for separably * *-increasing functionals F (see the Definition in Appendix [A]. Note that one can eliminate this restiction by taking the advantage of the minimax theorems (Terkelsen) 1972).
Theorem 1 (Maximin reformulation of the general OT). For separably *-increasing convex and lower semi-continuous functional F .M(XXX) IR{++00 it holds (we identify F(I) def F(TT)):
Cost(P, Q) -sup inf L(u, T) def supinf inf F(T) - -xx(( (T)(x,)) ((0))dQ(1)) (5)
T
T
where the sup is taken over potentials v E C())) and inf - over measurable measurabllffunciionn T X'x y.
From 55 we also see that it is enough to consider values of F in कт E [[(P). For convention, in
further derivations we always consider F(T,)) F(x) +oo for M(x x 3) IT(P).
We say that T is a stochastic OT map if it represents some OT plan x² solving i.e., T (x. ) = 7" (-|ær) holds PPalmost surely for all x
Theorem 2 (Optimal saddle points provide stochastic OT maps). Let v- E Eargsup, inf7 L(v, , T') be any optimal potential. Then for every stochastic OT map T holds:
T E arg inf L(v", ,T).
(6)
T
4

---

Published as a conference paper at ICLR 2024
Furthermore, if F is strictly convex in n, then 4 permits the unique OT plan x². In this case, T E arg inf, L(v", , ,T)
)) @ T is a stochastic OT map.
From our results it follows that by solving (5) and obtaining an optimal saddle point (v", )") one gets stochastic OT map T*, To ensure that all the solutions are OT maps, one may consider adding strictly convex regularizers to F with small weight, e.g., conditional interaction energy, see Appendix which is also known as the conditional kernel variance (Korotin et 2023a).
Practical considerations. Every term in (5) can be estimated with Monte Carlo by using random empirical samples from P, Q, allowing us to approach the general OT problem ( the continuous setting (33) To solve the problem (5) in practice, one may use neural nets To R x IR - R and tw R + R to parametrize T and v, respectively. To train them, one may employ stochas- tic gradient ascent-descent (SGAD) by using random batches from PP Q.S. We summarize the optimization procedure for general cost functionals F Algorithm 2 of Appendix [ In the text below ($5.1). we focus on the special cases of the class-guided functional FG. which is targeted to be used in the dataset transfer task (Figure 2 and pair-guided functional Fs for supervised image-to-image style transfer (55.2).
Relation to prior works. Maximin reformulations analogous to our (5) appear in the continuous OT literature (Krottin et al. 2021c]] 2023b Rout et al. 2022; Fan et al. 2023) yet they are designed only for classic (1) and weak (2) OT. Our formulation is generic and automatically subsumes all of them. allows using general cost functionals which, e.g., may easily take into account side information.
4.2 ERROR BOUNDS FOR APPROXIMATE SOLUTIONS FOR GENERAL OT
For pair (0, 2 approximately solving (5). it is natural to ask how close is TH. to the OT plan n². Based on the duality gaps, i.e., errors for solving outer and inner optimization problems with (6, 1') in we give an upper bound on the difference between Ty. and x². Our analysis holds for functionals F which are strongly convex in some metric p(.,)). see Definition 2 in Appendix A Recall that the strong convexity of F also implies the strict convexity, i.e., the OT plan T²" is unique.
5),
BURRR RUUUOOS be convex cost functional. Let p(:--) be metric IT(P) C MM(X' x y). Assume that F is B-strongly convex in p on (I). Consider the duality gaps for approximate solution (0,T) of (5) E2(6) def supinf -inf L(;,T))
Theorem 3 (Error analysis via duality gaps for stochastic maps). Let F MMX
E1 T def C(o, f) inf ((),T),
which are the errors of solving the outer sup, and inner infT problems in 5 respectively. Then for OT plan 't in (4) between IP and Q the following inequality holds:
p(T S 12 (VELCO.T) VE2(0)
i.e. the sum of the roots of duality gaps upper bounds the error of the plan Te w.r.l. n' in p(:)))
The significance of our Theorem Bis manifested when moving from the theoretical objective (5) to its numerical counterpart. In practice, the dual potential in (5) is parameterized by NNs (a subset of continuous functions) and may reach the optimizer v" Our duality gap analysis shows that we can still find good approximation of the OT plan. It suffices to find pair (b, I') that achieves nearly optimal objective values in the inner infr and outer sup,, problems of (5) In such pair, nn
close to the OT plan n². To apply our duality gap analysis, the strong convexity of F required. We give example of a strongly convex regularizer and general recipe for using it in Appendix D In turn, Appendix D.I demonstrates the application of this regularization technique in practice.
Makkuva al. 2020) carried out error analysis via duality gaps resembling our Theorem Their error analysis works only for classic OT (T) and requires the potential to satisfy certain convexity properties. Our error analysis is free from assumptions on é and works for general OT with strongly convex F.
Relation to prior works. The authors of (Fan al. 2023). (Rout al.
2022
LEARNING WITH GENERAL COST FUNCTIONALS
In this section, we show class-guided general cost functional $ for dataset transfer problem 6. and mair-guided cost functional $5.2 for supervised image-to-image translation 6.2
5

---

Published as a conference paper at ICLR 2024
Algorithm 1: Neural optimal transport with the class-guided cost functional Fo.
Input DDistributions P = >n. Fn Q = En, 3n Qn. S accessible by samples (unlabeled); weights On are known and samples from each I,,,, Qn are accessible (labeled); mapping network To IR x R + Rº; potential network vi RR + R:
number of inner iterations Kr;
Output : stochastic OT map To representing an OT plan between distributions P, Q:
repeat
Sample (unlabeled) batches Y Q, X ~ IP and for each r E X sample batch 2[x] ~ S;
C, Σ Σ
(Tatz.))
Σ
(v)
1112
#EX =E(a)
NEY
Update w by using ac.. 8w
for ky = 1, 1,2...................... K7 do
Pick 21,22..... N at random with probabilities (ai,.... .N.):
Sample (labeled) batches Xn ~ Pn, Yn Qni for each x E sample batch Zn [x] ~
Cr + A T(X,, kn)) Σ Σ
(Ta(,))))
xE,, sE [x11
Update by using 0Cr. 36
until not converged;
5.1 CLASS-GUIDED COST FUNCTIONAL.
To begin with, we theoretically formalize the problem setup. Let each input P and output Q distributions be a mixture of N distributions (classes) {IP } and (Qn (Qn). respectively. That is P P= P= P, and Q > Qn where On; Pn 2 0 the respective weights (class prior probabilities) satisfying Σ = On = 1 and Σ B = 1. In this general setup, we aim to find the transport plan (xx, y) E [I(,, Q) for which the classes of pairs (ar, y) ~ и as possible. That is, its respective stochastic map T should map each component P,, (class) of IP to the respective component Qn (class) of Q
X and ) are the same for many
The task above is related to domain adaptation or transfer learning problems. does not always have a solution with each Fr, exactly mapped to Qn due to possible prior/posterior shift (Kouw & Loog 2018). We aim to find a stochastic map T between P and Q satisfying (P,,, x S) z Q. for all N. To solve the above-discussed problem, we propose the following functional:
Fo(T) Fo(T,) def
an E2 (T (Ix Qn).
(7)
where & denotes the energy distance
(8 For distributions Q Q' P()) with y R.,, the (square of) energy distance & (Rizzo & Székely 2016 between them is:
(2 Q/) B Y2ll2
2 BIIY Yill E Yglls.
(8)
where Y,
Q Q² Q' are independent random vectors. Energy distance (8) is particular case of the Maximum Mean Discrepancy (Sejdinovic et al. (013)). It equals zero only when value when the components of
Q1 Q2. Hence, functional 7 is non- non-negative and attains are correctly mapped to the respective components of Q (if this is possible).
Theorem 4 (Properties of the class- class-guided guided cost functional Fo). Functional Fo(a) is convex in E (T(P), lower semi-continuous and +-separubly increasing.
In practice each of the terms 8 (Taa(P,,X)) x)). Qn) in 7 admits estimation from samples from m. Proposition 1 (Estimator for E²), Lel Xn P. be a butch of x samples from class n. For each E X,, let 2n [x]] ~ & be latent batch of size Kz. Consider a batch Yn Qn of size Ky. Then
11
T(x, (xx,2)) 2)112
A Xn,, T(Xn: Yn)
del
ΣΣ
Ky Kx Kz
yEY, x 207.
[[ z) T(x',=')]1
Σ Σ
Σ
(9)
2. (K* Kx) K2
ZEX Z'EX x (x) d'Z

---

Published as a conference paper at ICLR 2024
is an estimator of 8 (Ti(P.X8), Qa) up to a constant constant T-independent shift.
To estimate FT(T), one may separately estimate terms E (S)) Qa) for each n and sum them up with weights Onn We only estimate n-th term with probability On at each iteration.
We highlight the two key details of the estimation of which are significantly different from the estimation of classic (1 and weak OT costs 2 appearing in related works (Korotin et al. 2023bb 2021b; Fan et al. 2023). First, one has to sample not just from the input distribution IP, but separately from each its component (class) ,.. Moreover, one also has to be able to separately sample from the target distribution's s Q components Qn- This is the part where the guidance (semi-supervision) happens. We note that to estimate costs such classic weak (2) target samples from Q are needed at all, i.e., they can be viewed as unsupervised.
In practice, we assume that the learner is given labelled empirical sample from IP for training. In contrast, we assume that the available samples from Q are only partially labelled (with 2 labelled data point per class). That is, we know the class label only for limited amount of data (Figure 2 In this case, all n cost terms (9) can still be stochastically estimated. These cost terms are used to learn the transport map Te in Algorithm 2 The remaining (unlabeled) samples will be used when training the potential Uw., as labels are not needed to update the potential in (5). We provide the detailed procedure for learning with the functional FG (7) in Algorithm T
COST FUNCTIONAL In this section, we demonstrate general OT formulation with another peeticcall--apeallin specification. In particular, we define the pair-guided general OT cost functional. For given paired data set (21, (x1)),........ (IN,, Nee,,y" y"" (æx)) with samples XI:N = f ... 2:n) and y" (X {y "....)))))N)) which are assumed to follow the source I and target Q distributions, respectively, we introduce:
5.2
Jx
Fs(T))
xxxy E(y, v² (x)))d((a,))
(10)
The function é X' x IR is appropriate loss measuring the difference between samples. In the majority of our experiments we choose E(y, 3/) = 11 2/1|22- 2- In practice, to estimate (T)) we assume that the learner is given labelled empirical sample of P and Q for training. Using the cost (10). which can handle such information, we can train optimal mapping in a supervised manner. In 6.6.2 show how our method, together with the pair-guided functional Fs, is directly applicable the paired image-to-image translation problem.
6 EXPERIMENTAL ILLUSTRATIONS
Our Algorithm D is capable of learning both stochastic (one-to-many) T(x, 2) and deterministic
(one-to-one) T 2) = T(x) transport maps. For the latter, no random noise is added to input.
First, we implement stochastic and deterministic maps along with our class glass- cost function Fa.
to address the dataset transfer problem ($6.1). Second, using the deterministic transport map T(x).
we apply our pair-guided functional Fs to solve the image-to-image supervised translation ($ 2). In
the Appendix we provide experiments with toy data C.2 biological batch effect problem C.I3 and
various paired image-to-image datasets E
The code for the experiments can be found at
https: /ttp://gittbb http://hhhhuu..ooooccccheeeeeeee
6.1 CLASS-GUIDES EXPERIMENTS
Datasets. We use MNIST LeCun & Cortes 2010 FashionMNIST (Xiao et al. 2017) and MNIST-
M (Ganin & Lempitsky 2015) datasets P. Q Each dataset has 10 (balanced) classes and the pre-defined train-test split. In this experiment, the goal is to find class-wise map between unrelated domains: FMNIST - MNIST and MNIST - MNIST-M. We use the default class correspondence between the datasets. For completeness, in Appendices we provide additional results with imbalanced classes C.8), non-default correspondence (C.IT). and other datasets (C).
Baselines. We compare our method to the pixel-level adaptation methods such as (one-to-many AugCycleGAN (Almahairi al. 2018| Zhu et al. 2017 Hoffman et al., 2018 Almahairi et al. 2018 and MUNIT (Huang et al. 2018] Liu et al. (017)). We use the official implementations with the hyperparameters from the respective papers. We test Neural OT (Korotin et al., 2023b Fan et al.) 2023) with Euclidean cost functions: the quadratic cost 1l|z - -21ī (W2) and the w-weak (one-to-
many) quadratic cost (W2,, 7 1/)) For semi-supervised mapping, considered (one-to-one)
7

---

Published as a conference paper at ICLR 2024
2
8
3
4
/
2
3
4
5
6
7
8
9
DEF OTDO CAN 9INT y Ce
(
2
3
4
5
6
7
8
9
0
\
3
3
4
6
7
1
5
7
7
4
9
4
7
O
0
3
5
6
7
8
9
7
I
3
1
3
2
8
6
5
6
MA DAN
/
2
3
4
5
8
9
6
4
/
2
3
4
5
7
8
9
8
3
3
4
5
4
7
8
9
7
8
1
0
6
2
4
2
16
3
4
5
6
7
S
(ouns) louns) 1A.
6
6
2
8
7
9
2
8
(curs) (ourz) 19,
9
3
6
3
7
92
To 15.
0
1
2
3
4
5
6
9
8 8
9
3
6
7
8
0
1
3
3
7
5
6
7
9
3
4
5
4
8
(a) FMNIST + MNIST
(b) MNIST + MNIST-M
Figure 3: The results of class- preserving mapping between two unrelated (left) and related (right)
datasets. Each column shows the transfer result of a random (test) input x ~ Pr (first row) from a
particular class (n 0, 1,....9)) 9). Each row shows the results of transfer via particular method. For
methods which learn a stochastic map T T(r, 2), we show their output T(x z) for a random noise s.
Emage-to-Imagg Translation
Flows
Discrete OT
Neural Optimal Transport
Datasets 32)
MUNIT
CycleGAN
Aug
OTDD
SinkhornL.pL1
W2
Wa
Fo. no
[Ours]
[Ours]
FMIIST - MNIST
12.03
10.28
10:67
10.96
8.02
82.79
83,22
MNIST -+ .NIST-M M
97.95
98.2
83.26
38.77
37.0
95.27
94.62
Table 1:
Accuracyt of the maps learned by the translation methods in view.
Datasets (32 x 32)
MUNIT
Aug CycleGAN
OTDD
SinkhornLpLII
W2
Wa.v
Fa. no a [Ours]
[Ours]
FMNIST - MNIST
7.91 11.68
26.35
> 100
> 100
7.51
7.02
7.14
5.26
MNIST - MNIST.M
26.87
>
19.43
17.48
18.56
Table 2: FIDI of the samples generated by the translation methods in view.
OTDD flow Melvas & Fusi 2021 2020). This method employs gradient flows to perform the transfer preserving the class label. We also examine a General Discrete OT (DOT) which use
labels. In particular, we adopted the solver from ot.. da (Flamary et al.) 2021 with its default
out-of-sample estimation procedure. The solver utilizes the Sinkhorn (Cuturi (2033) with Laplacian
cost regularization (Courty al. 2016). We show the results of ICNN-based OT method (Makkuval
et al. 2020 Korotin et al. 2021a) in Appendix C.9
Metrics. All the models are fitted on the train parts of datasets; all the provided qualitative and quantitative results are exclusively for test (unseen) data. To evaluate the visual quality, we compute FIID (Heusel et al. 2017) of the entire mapped source test set w.r.t. the entire target test set. To estimate the accuracy of the mapping we use a pre-trained ResNet18 8 (He et al. 2016) classifier (with 95+ accuracy) on the target data. We consider the mapping T' to be correct if the predicted label for the mapped sample I T'x,, 2) matches the corresponding label of x.
Results. Qualitative results are shown in Figure 3 FID, accuracies - in Tables aadd respectively. To keep the figures simple, for all the models (one-to-one, one-to-many), we plot single output per input. For completeness, in Appendices C.5 C.6 we show multiple outputs per each input for our method, and in Appendix C.7 we provide ablation study on 2 size. Our method, general discrete OT and OTDD, use 10 labeled samples for each class in the target. Other baselines lack the capability to use label information. As seen in Figure 3 and Table [ our approach preserves the class-wise structure accurately with just 10 labelled samples per class. The accuracy of other neural OT methods is around 10%, equivalent to a random guess. Both the general discrete OT and OTDD methods do not preserve the class structure in high dimensions, resulting in samples with poor FID, see table 2 Visually, the OTDD results are comparable to those in Figure 3 of (Alvarez-Menl) Melis & Fusi 2021).
8

---

Published as a conference paper at ICLR 2024
X
y'()
RMSE
Wi
(ours) Pix2Pix
F's
Figure 4: Results of our method with the pair- guided cost functional in comparison to other methods, applied to the supervised image-to-image translation task (Comic-Faces- V1, 256 x 256).
6.2 FATR-GUIDEL EXPERIMENTS
Datasets and metrics. We utilize three popular datasets for our evaluation: Comic-Faces-VI Edges- to-Shoes (Isola al. 2017), and CelebAMask-HO (Lee al. 2020). These datasets all contain pairs of images (handmade or formulated synthetically) and are commonly employed in benchmarking supervised image translation methods. For all experiments, we operated the resolution of 256 x 256 pixels. All the results (qualitative, quantitative) are on the test sets in accordance with the default train-test splits for the datasets. As the metric, we use the test FID (as in 36.1
Baselines. As the baselines, we consider the basic (unsupervised) NOT algorithm Korotin et al. 20236), RMSE regression, and well-celebrated supervised Pix2Pix method (Isola et al. 2017).
Details. In our method, we use U2Net the transport map (z) and WGAN-QC discriminator's ResNet architecture (He et al. 2016) for potential tw. In Comic-Faces-VV V1 and Edges-to-Shoes
experiments, we use RMSE as the function E in our method. In CelebAMask-HO case, we use VGG-based perceptual loss. Other details are given in Appendix $ [E.3
Results. The evaluation results and comparisons with the baselines for the Comic-Faces-VV V1 and Edges-to-Shoes datasets are presented in Figure and Figure 2T respectively. Additionally, the computed FID scores for these datasets detailed Table 7 in the appendix. For CelebAMask-HO dataset, we achieve the FID score of 21.1. The examples of generated images are provided in Figures and [2 Further qualitative experimentation is conducted the Comic-Faces: V| dataset higher resolution of 512 x 512, see Figure Overall, the obtained results show that our method achieves competitive quality and can be further applied to high quality generation and editing tasks.
7 DISCUSSION
Our method is generic tool to learn transport maps between data distributions with a task-specific cost functional F. In general, the potential impact of work society depends on the scope of its application in digital content creation. As limitation, we can consider the fact that to apply our method, one has to provide estimator F(T) for the functional F which may be non-trivial. Besides, the construction of a cost functional F for a particular downstream task may be not straightforward. This should be taken into account when using the method in practice. Constructing task-specific functionals F and estimators F is promising future research avenue.

---

Published as a conference paper at ICLR 2024
8 ACKNOWLEDGEMENT
The work was supported by the Analytical center under the RF Government (subsidy agreement 000000DD3003211500000, Grant No. 70-2021-00145 02.11.2021).
REFERENCES
J-J Alibert, Guy Bouchitté, and Thierry Champion. A new class costs for optimal transport planning. European Journal of Applied Mathematics, 30(6):1229-12655 2019.
Amjad Almahairi, Sai Rajeshwar, Alessandro Sordoni, Philip Bachman, and Aaron Courville. Augmented cyclegan: Learning many-to-many mappings from unpaired data. In International Conference on Machine Learning, pp. 195-204. PMLR, 2018.
David Alvarez-Melis and Nicolo Fusi. Geometric dataset distances via optimal transport. Advances in Neural Information Processing Systems, 33:21428- 2-1-211441 21439, 2020.
David Alvarez-Melis and Nicolò Fusi. Dataset dynamics via gradient flows in probability space. In International Conference on Machine Learning, pp. 219-230 PMLR, 2021.
David Alvarez-Melis, Yair Schiff, and Youssef Mrouch. Optimizing functionals on the space of probabilities with input convex neural networks. Transactions on Machine Learning Research, 2022.
Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 146-155. JMLR. org, 2017.
-aalunn-VVolume
Martin Arjovsky and Léon Bottou. Towards principled methods for training generative adversarial networks. arXiv preprint
2017.
Julio Backhoff- Backhoff-Verrguas, Mathias Beiglböck, and Gudmun Pammer. Existence, duality, and cyclical monotonicity for weak transport costs. Calculus of Variations and Partial Differential Equations, 58(6):1-28, 2019.
Lukas Biewald. Experiment tracking with weights and biases, 2020. URL httph:tt//////t wandb Softwareevaiile available from wandb.com.
com/ Software
Nicolas Bonneel and Julie Digne. A survey of optimal transport for computer graphics and computer vision. In Computer Graphics Forum, 2023.
Victor Bouniakowsky. Sur quelques inégalités concernant les intégrales ordinaires les intégrales aux différences finies, volume 1. Mem. Acad. St. Petersburg, 1859.
Charlotte Bunne, Lactitia optimal transport modeling of population dynamics, 2021.
Andreas Krause, and Marco Cuturi. Jkonet: Proximal
Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy. Optimal transport for domain adaptation. IEEE transactions on pattern analysis and machine intelligence, 1 2016.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in neural information processing systems, pp. 2292-2300, 2013.
Grady Daniels, Tyler Maunu, and Paul Hand. Score-based generative neural networks for large-scale optimal transport. Advances in Neural Information Processing Systems, 34, 2021.
Nabarun Deb, Promit Ghosal, and Bodhisattva Sen. Rates of estimation of optimal transport maps using plug-in estimators via barycentric projections. Advances in Neural Information Processing Systems, 34:29736-29753. 2021.
Montacer Essid and Justin Solomon. Quadratically regularized optimal transport on graphs. SIAM Journal Scientific Computing, 1961):A 2018.
10

---

Published as a conference paper at ICLR 2024
Jiaojiao Fan, Qinsheng Zhang, Amirhossein Taghvaei, and Yongxin Chen. Variational wasserstein gradient flow. In International Conference on Machine Learning, pp. 6185-6215 PMLR, 2022.
Jiaojiao Fan, Shu Liu, Shaojun Ma, Hao-Min Zhou, and Yongxin Chen. Neural monge map estimation and its applications. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL Featured Certification.
/tttp::://penreveew net www.moodooruaaaiaa
Sira Ferradans, Nicolas Papadakis, Gabriel Peyré, and Jean-François Aujol. Regularized discrete optimal transport. SIAM Journal Imaging Sciences, 7(3):1853-1882, 2014.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, et al. Pot: Python optimal transport. The Journal of Machine Learning Research, 22(1))357)-3578, 2021.
Yaroslav Ganin and Victor S. Lempitsky. Unsupervised domain adaptation by backpropagation. In Francis R. Bach and David M. Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop
and Conference Proceedings, pp. 1180-1 1880-1189, 189. JML. JML..org. 2015. URL http:: httt:::/pproeeeinng mlr peress/v37.m m15 ht.l
Milena Gazdieva, Litu Rout, Alexander Korotin, Alexander Filippov, and Evgeny Burnaev. Unpaired
image super-resolution with optimal transport maps. arXiv preprint arXiv:2202.01116, 2022.
Aude Genevay. reguoiy-r optimal transport for machine learning. PhD thesis, Paris Sciences et Lettres (ComUE), 2019.
Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large- scale optimal transport. In Advances in neural information processing systems, pp. 3440.3448, 2016.
Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample complexity of sinkhorn divergences. The 22nd international conference on artificial intelligence and statisties, pp. 1574-1583. PMLR, 2019.
Nathael Gozlan and Nicolas Juillet. On a mixture of brenier and strassen theorems. Proceedings of the London Mathematical Society, 120(3):4344463 2020.
Nathael Gozlan, Cyril Roberto, Paul-Marie Samson, and Prasad Tetali. Kantorovich duality for general transport costs and applications. Journal of Functional Analysis, 273(11):3377-3405 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Pierre Henry-Laborden (martingale) optimal transport and anomaly detection with neural networks: A primal-dual algorithm. Available at SSRN 3370910, 2019.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by two time-scale update rule converge to a local nash equilibrium. In Advances in neural information processing systems, pp. 6626-6637, 2017.
Judy Hoffiman, Eric Tzeng. Taesung Park, Jun-Ye Yan Zhu, Phillip Isola, Kate Saenko, Alexei A. Efros, and Trevor Darrell. Cycada: Cycle-consistent adversarial domain adaptation. Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference Machine Learning, ICML 2018, Stockholmsmässsan Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 1994-2003. PMLR, 2018. URL http:/7
proceedi ngs m r. prwss/essvvv /ho fman 8a. html
Xun Huang, Ming Yu Liu, Serge Belongie, and Jan Kautz. Multimodal unsupervised image-to- tongge image translation. In Proceedings of the European conference on computer vision (ECCV), pp. 172-189, 2018.
11

---

Published as a conference paper at ICLR 2024
San-Christian Hütter and Philippe Rigollet. Minimax estimation of smooth optimal transport maps. 2021.
Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A Efros. Image-to-image translation with conditional adversarial networks. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 1125-1134, 2017.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6998 2014.
Lev Borisovich Klebanov, Viktor Beneš, and Ivan Saxl. N-distances and their applications. Charles University in Prague, the Karolinum Press Prague, Czech Republic, 2005.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safn, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations. 2021a. URL wttpt:ttt///peeeeeeeettttts
net/ forum? wwtpmmidd-bie d=bboox EXsa
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. Advances in Neural Information Processing Systems, 34: 15533-165646 2021b.
Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous wasserstein-2 barycenter estimation without minimax optimization. In International Conference on Learning Representations, 2021c. URL httpps://ppeereeeeeettt net/ hat.:eeenneettromm22ddd33tttaa-
Alexander Korotin, Vage Egiazarian, Lingxiao Li, and Evgeny Burnaev. Wasserstein iterative networks for barycenter estimation. In Thirty-Sixth Conference on Neural Information Processing Systems, 2022a. URL htttps://penreviee..... wttp::/.eeeefforeeeeeeeeee wiw.:rrummmniciiiiiieeeeeeeeeeee
Alexander Korotin, Alexander Kolesov, and Evgeny Burnaev. Kantorovich strikes back! wasser- stein GANs are not optimal transport? In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022b. URL /ttl:ttts////peeeeee// net/Eoourn wttp:rrmmmmmdctEep---eeo
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Kernel neural optimal transport. In International Conference on Learning Representations, 2023a. URL /tttps:/
n Zuc MHt
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. In International Conference on Learning Representations, 2023b. URL BRRLatpps///ppeeeeiier net fnet/Eorung wtw.:oorrnnndddcnnnnneeeeee
Wouter M Kouw and Marco Loog. An introduction to domain adaptation and transfer learning. arXiv preprint arXiv: 112. 1806, 2018
Cosmin Lazar, Stijn Meganck, Jonatan Taminau, David Steenhoff, Alain Coletta, Colin Molter, David Y Weiss-Solís, Robin Duque Hugues Bersini, and Ann Nowé. Batch effect removal methods for microarray gene expression data integration: survey. Briefings in bioinformatics, 14 (4):469-49)) 490, 2013.
Yann LeCun and Corinna Cortes. MNIST handwritten digit database. MNIST, 2010. URL http: 7/yann. lecun. com/ wwtm// / miist/
Cheng-Han Lee, Ziwei Liu, Lingyun Wu, and Ping Luo. Maskgan: Towards diverse and interactive facial image manipulation. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
Jeffrey T Leek, Robert B Scharpf, Héctor Corrada Bravo, David Simcha, Benjamin Langmead, W Evan Johnson, Donald Geman, Keith Baggerly, and Rafael A Irizarry. Tackling the widespread and critical impact of batch effects in tigh-throughhut data. Nature Reviews Genetics, 11(10): 733-739, 2010.
Ruilin Li, Xiaojing Ye, Haomin Zhou, and Hongyuan Zha. Learning to match via inverse optimal transport. Journal of machine learning research, 20, 2019.
12

---

Published as a conference paper at ICLR 2024
Huidong Liu, Xianfeng Gu, and Dimitris Samaras. Wasserstein GAN with quadratic transport cost. In Proceedings of the IEEE International Conference on Computer Vision, pp. 1882-48413 2019.
Ming-Yu, Yu Liu, Thomas Breuel, and Jan Kautz. Unsupervised image-to-image translation networks. In Advances in neural information processing systems, pp. 700-708, 2017.
Ruishan Liu, Akshay Balsubramani, and James Zou. Learning transport cost from subset cor-
respondence. International Conference on Learning Representations, 2020. URL https //openreview. net/ htt.://ww.eet/toommoio-oooaoaaa
Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pp. 6672-6681.
PMLR, 2020.
Tudor Manole, Sivaraman Balakrishnan, Jonathan Niles-Weed, and Larry Wasserman. Plugin 2021.
estimation of smooth optimal transport maps. arXiv preprint
Mark W Meckes. Positive definite metric spaces. Positivity, 17(3):733-7577 2013.
2013.
Gonzalo Mena and Jonathan Weed. Statistical bounds for entropic optimal transport: sample complexity and the central limit theorem. Advances in Neural Information Processing Systems, 32, 2019.
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, and Evgeny Burnaev. Large-scale wasserstein gradient flows. Advances in Neural Information Processing Systems, 34:15243-15256, 2021.
John C Nash. The (dantzig) simplex method for linear programming. Computing in Science & Engineering, 2(1):29-31, 2000.
François-Pierr Paty and Marco Cuturi. Regularized optimal transport is ground cost adversarial. In International Conference Machine Learning, pp. 7532-7542. PMLR, 2020.
Michaël Perrot, Nicolas Courty, Rémi Flamary, and Amaury Habrard. Mapping estimation for discrete optimal transport. Advances in Neural Information Processing Systems, 29, 2016.
Henning Petzka, Asja Fischer, and Denis Lukovnikov. On the regularization of wasserstein gans. In International Conference on Learning Representations, 2018.
Gabriel Peyré, Marco Cuturi, al. Computational optimal transport. Foundations and Trends@ in Machine Learning, 1 15556))355-607, 2019.
Aram-Alexandre Pooladian and Jonathan Niles- Weed. Entropic estimation of optimal transport maps. arXiv preprint arXiv:2109. 2021.
Alain Rakotomamonjy, Rémi Flamary, and Nicolas Courty. Generalized conditional gradient: analysis of convergence and applications. arXiv preprint arXiv:1510.0655 2015.
Maria L Rizzo and Gábor Székely. Energy distance. wiley interdisciplinary reviews: Computational
statistics, 8(1):27-35 2016.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer-assisted intervention, pp. 234-241. Springer, 2015.
Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. In International Conference on Learning Representations, 2022. URL TRRl:tttpa://ppeeeeviee
net /odt
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63):94, 2015.
Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Math- ieu Blondel. Large scale optimal transport and mapping estimation. In International Conference Learning Representations, 2018.
13

---

Published as a conference paper at ICLR 2024
Dino Sejdinovic, Bharath Sriperumbudur, Arthur Gretton, and Kenji Fukumizu. Equivalence of distance-buse based and rkhs-based statistics hypothesis testing. The Annals of Statistics, pp. 2263- 2291, 2013.
Andrew M Stuart and Marie-Therese Wolfram. Inverse optimal transport. SIAM Journal Applied Mathematics, 80(1)599-611, 2020.
Xuan Su, Jiaming Song, Chenlin Meng, and Stefano Ermon. Dual diffusion implicit bridges for image- to-image translation. In The Eleventh International Conference on Learning Representations, 2022.
Frode Terkelsen. Some minimax theorems. Mathematica M Sendnnavia,, 413, 1972. ISSN 00255521, 19031807. URL httthtttp//ww// http:///ttor..ogggaaaaa//////220
Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist a novel image dataset for benchmarking machine learning algorithms. arXiv preprint 1eXXiv1708.077440 2017.
Luke Zappia, Belinda Phipson, and Alicia Oshlack. Splatter: simulation of single-cell ma sequencing data. Genome biology, 18( 8(1):174, 2017.
Jun- Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. In Proceedings of the IEEE international conference on computer vision, pp. 2223-2232, 2017.
14

---

sup sup sup sup Published as a conference paper at ICLR 2024
A PROOFS
A.1 PROOFS OF RESULTS OF $4
To begin with, we recall some basics which play an important role in the further derivations.
Conjugate functional. Let F MM x ))
R U {oo} be functional. The convex conjugate
functional of F is F C(K )) - {oo}:
(h) def
[ h(s)dT(s)) -F(m)
sup
xcM(Xx))
Let u, o E C()). We write u @ v E C(X x }) to denote the function u 4 v (2, y) H u(x)) v(y). The next definition is borrowed from (Paty & Cuturi 2020 Definition 2).
Definition (Separably * --increasing functional). For F M(X }) + I say that it is separably *-increasing if for all functions u. v E C(X),C()) and any function c E C(X x y) from u v S e (point-wise) it follows F (u +) v) (c)...
Proof of Theorem D Since F is convex, Isc and separably * i-increasing functional, general OT prob- lem 4 permits the following dual representation (Paty & Cuturi 2020 Theorem 2):
Cost(P: (P, Q sup U u(z)dP(s) I, oly)dQ(u) -F)
(11)
where optimization is performed over u, v E ((x))C()) which are called potentials. We use the dual form (1) to derive
Cost(P, Q) = sup sup [
/ (())d(()))))
F -F(
(12)
y
+ / v())d(w))
u(x)dP(x)) sup X (u@v)dT(I,)) F(T))
(13)
| (u (udddd((r v))
uu (dP inf (F(T)
+ v())Q(w))
(14)
xxx
u
sup
sup inf
F(T)- u(x)d(= PP) - u(v)dT(y)
(v)
(15)
x
sup
sup inf =EIT(P)
F(T)) u(z)d(T F)())
v(v)dT(1)
vv))d(())
(16)
x
sup sup inf =EIl(P)
F(T))
(())d(())
(v))(()))
(17)
1/ v(w)dQ())
sup
inf TεT((P)
(((T)
v(y)dT(1)
(18)
sup F(T")
e(y) d= v())d(()))
=F(T") Cost(I, Q). (19)
(Q(D)
In line 12). we group the terms involving the potential u. In line 13). functional F by using its definition. In the transition to line (14). we replace inf, operator with the equivalent sup, operator with the changed sign. In transition to (15). we put the term Jx u( under the inf, operator; we use definition (utv)(ax v) = u(x)+((() to split the integral over π y) into two separate integrals over 7 (I:) and T (yy) respectively. In transition to (16). we restrict the inner inf, to probability distributions π E [I)) which have P as the first marginal, i.e. dT(x) = dI.(). This provides result, we remove the sup,, operator in line (18. In transition to line (19) we substitute an optimal plan x² E (TP,, Q) c TT(Q) upper bound (8). Since Cost(IP Q) turns to be both an upper bound
we express the conjugate
upper bound (5). in particular, all u-dependent terms vanish, see (17) As
15

---

Published as a conference paper at ICLR 2024
19
and lower bound 12 for (1). conclude that:
TEIT()) F(m) J ((w))a(()> + / v(o)d((v) def
(F Q) sup inf
sup inf Cp(u, (20) *T)((?)
The desired equation 5 is obtained by replacing plans E (PPP) in (20 with their stochastic map representations T. see the first paragraph of $4.1
Proof Theorem Assume that T arg inf L(v", T). This yields
L(v", T*) > inf L(v", T) Cost(P, Q).
T
On the other hand, by substituting 7 with the optimal OT plan TT. we write:
C(v", = F(TT-) - S o² v"(() drr- / e² (v)dd(()) = F(Tr.) = Cost(P. Q).
aQ())
which is contradiction. Thus, the assumption is wrong and (6) holds.
Let F be strictly convex, i.e. (4) permits unique OT plan n² From the first part of the theorem, it holds that the corresponding stochastic map T.- solves (6), i.e., T.- E arg inf, L(v", ,T) Consequently, x² E inf=en(r)) (n'"", .), see (20). Let T E arg inf L(v", .) be another stochastic map solving (6). Similarly, Tr E inf (P)()) Lp("", π). Since F is strictly convex, then f,("", л) is also strictly convex as functional of æ. The latter immediately yields that n" = nr, ie. T is stochastic OT map. Combining this conclusion with the first part of the theorem, we derive:
T' inf L(v", T) TT- is OT plan,
T
which finishes the proof.
Definition 2 (Strongly convex functional w.r.t. metric p(., 1). Let Let :M x)) - RU be
a convex lower semi-continuous functional. Let U C C(X x y) C CMX x y) be a convex subset such that 3m U: F(T) < +00. Functional F is called B-strongly convex on U w.r.t metric p(.,-)
VT 22 E U, Va O0, 1] holds:
F(am + (1 - a)T2) S (F(())) + (1 -a)F(F2)
2 o( - (001-a)p'(60 (T1 T2)-
(21)
Lemma 1 (Property of minimizers of strongly convex cost functionals). Consider a lower- semicontinuous B-strongly convex in metric p(-.-) on U c P(X x Y) functional F. Assume that x" U satisfies F(m²)) inf F(T). Then Va E U it holds: REMl
8
0000
(m (m²,))) .7).
(22)
Proof of Lemmal
7 We substitute क = n², :22 ग to formula (2 and fix a E [0, 1]. We obtain
F(ax* + (1 - (1-
( (1) (1- (1-o)
a(1 (1l-a)pp2 a)) m) =
F(an" + (1- -aF = (1-
o0 a) (='
2
inf F(r")=F(x*) e'eu
(1 ( (1-a)F(m²) S (1- 8
aa(1 a
B
F( < F(T)
app( T))
(23)
Taking the limit >> 1- in inequality 23 obtain
22
16

---

Published as a conference paper at ICLR 2024
Theorem 5 (Error analysis via duality gaps). In the conditions of Theorem 3 consider the duality
gaps for an approximate solution (0, #t) E C()) x [I) of 20
ei (t, def (ū, #) inf (p((, R), (24) (0) sup inf C,(u, T) -inf C, T), (25) =EI((P) -
v =EI(??)
=EI((P)
which are the enrors of solving the outer sup,, and inner inf problems in (0 respectively. Then for the OT plan x² in (4) between P and Q the following inequality holds
x() S 2 (Venle,tt) () + vea(i))
(26)
i.e. the sum of the roots of duality gaps upper bounds the error of the plan &t w.r.t. nt p(:s ))
Proof of Theorem 5 Given a potential v E C(V), we define functional V. II(P) + +RU {+00}:
(Tr def F(m) / e(y)dr(y)
(27)
Since the term fo v(y)d((())))) is linear w.r.t. т, the 3-strong convexity of F implies B-strong convexity of Vv. Moreover, since V. is lower semi-continuous and (T(PP) is compact (w.r.t weak-* topology), it follows from the Weierstrass theorem (Santambrogio [0035] Box 1.1) that
3T" LI(IP): (e(")) inf (.));;; TEIT(P)
(28)
i.e the infimum of (x) is attained. Note that n² minimizes the functional x + C(u, m) as well since L(v, T² = (T) + Const(v). Therefore, the duality gaps (24 (25) permit the following reformulation:
(b = Lp(0, 4) (o( )- where T is minimizer 28 for v 6. Consider expression
(29)
( = Lp("", (²)) L0 7));
(30)
29
Lemma 1 3
el ( +) = Co R) Cp(0,
2 2
2
(0, +) 2 T")
(31)
B
Consider expression 330
(2 L, xT) - (,(
F(T') - / ('(()((T'' -0) + / o(v)d(T' (v)
F(x²) - / b(y)d(T -2) + / (6(v) - -v(()dd(=" Q)( -v) F(T))+ + / ((y)d(T' - -V) = (32)
F(TT)) / ((y)do ((y))d(())++++
(v) + / (6(v) - -w')))((( -Q) -F(=0) + / i(y) ((y)dr
=Vo(=²)
=0, since di (v)=)Q((v)
=-Vo(=E)
Lemma
B
Ve(TT) - Vo(
2
2 p= (²(x",TT) m²)
2
1 B
c(( 2 ;,);;;));=)>> (33)
where in line 32 we add and subtract ʃy ((y)d(7* -Q)(v)
The triangle inequality p(T=", f) S p(TT" ,²)"" + p(ft, = = (Va(t,)) + +V2()) for norm
p(.))) finishes the proof.
Proof of Theorem 3 The statement directly follows from Theorem 5 by substituting plans л E II(PP with their stochastic map representations Tz, see the first paragraph of 4.1
17

---

Published as a conference paper at ICLR 2024
A.2 PROOFS OF RESULTS OF 5..1
First, we prove that F = Fa it is *-separately increasing. For ग E M(X x )) ((P)) it holds that F(x) +oc. Consequently,
Proof of Theorem +
J ()d,() F(T) =
(aam)((xx xxxy (u(z) + u())+ F(T) = -00 (34)
When (T)(P) it holds that is a probability distribution. We integrate u(x) + v(v) S ctz, v) w.r.t.
T, subtract F(r) and obtain
xxy ((I.))dT(.V)) F(T) 2 / xxy
(u(z) + (())++(())d((1) F(x).
(35)
By taking the sup of and 35 w.I.t. x E M(X V). we obtain F 2 e)
Next, we prove that F is convex. We prove that every term (T S) Q.) is convex in n.
First, we show that + (TT def [F x 3) is linear in (I).
Pick any TT, E II(P) which lie on the same line. Without loss of generatity we assume that Ty E [11, #2], i.e., my = am1 (1 22 for some a E [0, 1.. We need to show that
(TS) a (T) + (1 fm
(36)
In what follows, for random variable U we denote its distribution by LUW(
The first marginal distribution of each n, is P. From the glueing lemma (Villani 2008 that there exists a triplet of (dependent) random variables (X, Y1, Y2) such that Law X, Y, = T for = 2, We define Y3- Y,, where r an independent random variable which takes values in {1, with probabilities {aa, 1 a). From the construction of Y3 it follows that Law(X, Y2) is mixture of Law(X, Y1) = π and Law(X: X, Y2) T2 with weights and a. Thus, Law (X, Ys) = aτ1 + (1 7T = We conclude that Law(Yg) X=x) = Ts(-(ar) for P-almost all x E X (recall that (aw(()=PP) On the other hand, again by the construction, the conditional Law(Y3) X = is mixture of (Yi
51) it follows
=x)
x) Ti (-1x)) and Law(Y,) X =x) π2(-11² with weights a and 1 Thus, T (2) = am1 (.1x) + (1 (2)2((11x) holds true for P-almostt all a X.
Consider independent random variables X, ~ Pr and Z S. From the definition of Tx we conclude that Law Law(T=, (z, Z)) (-(x)) for P-almost all x E X' and, since Pn is a component of P. for P. almost all x E X' as well. As result, we define T; = Tn, (Xn, Z) and derive
Law = x) T2(+22 = am (+(x) (1 at T2 (-(zz)) eLaw((T = r) + (1 Law (T, X, =
for Pllalmost all x E X. Thus, (Xw, is also a mixture of Law(X,., Ti) and Law(Xn, Th) with weights a and LTw(T,) = (T;) by the definition of fn and obtain (36).
In particular, (aw(T aLaw(Ti) + (1 (T2)). We note that
Second,
we highlight that for every P()"), the functional P(Y) μ E2(,, v) is convex in Indeed, 8² is a particular case of (the square of) Maximum Mean Discrepancy (MMD, (Sejdinovic al. 2013)). Therefore, there exists Hilbert space H and function & y - H (feature map), such that
E² (² - d(v)du(v) - / o (v
1
Since the kernel mean embedding μ1 ++ Jyz (())d(())) linear in μ and - convex, we conclude that (1, v) is convex in μ. To finish this part of the proof, it remains to combine the fact that + (IP, x E) is linear and (, Qx) is convex in the first argument.
Third, we note that the lower semi comi-continuity F(T)) ) in (IP) follows from the lower semi cemi-continuity of the Energy distance (E2) terms in 7 That is, it suffices to show that 82 defined in equation
'The proof is generic and works for any functional which equals outside P(X y).
18

---

Published as a conference paper at ICLR 2024
is indeed lower semi-continuout in the first argument Q1 In (8. there two terms depending [| - 22l 2d x2)) dQu (21) is linear in Q1. It is on Q1. The term X1 - Xi = Ix [Jx
just the expectation of continuous function w.r.t. Q1. Hence it is lower semi-continuous by the definition of the lower semi-continuity. Here we also use the fact that y is compact. The other term fxxx ||ا| 2/2 22112d(2 (Q1 x Q1) (21, .222 is quadratic term in Q1. This term can be viewed as the interaction energy (Santambrogio 2015 57) between particles in Q1 with the
- Xill2
interaction function W d -|| x 2/2 Thanks to the compactness of ) it is also lower
semi-continuous in Q1. see (Santabrrgiio 2015 Proposition 7.2) for the proof.
Proof of Proposition
Direct calculation of the expectation of yields the value
1
E T(X, 2) ||2
T(X, 2) T(X', 122
2
1 2 _ Y'
E - T(X, -))
イ B Y'
82 (T x 3), Q.) + Elly - -Y'l2, (37)
where Y.Y" ~ On and (X, 2X', (P,, x 3) are independent random variables. remains to note that E 12 is a T -iddepeendent constant.
B ALGORITHM FOR GENERAL COST FUNCTIONALS
In this section, we present the procedure to optimize 5 for general cost functionals F. In practice, one may utilize neural networks To TDD x RRS - TRDD and ow TR - to parameterize 7 and Bas correspondingly, solve the problem 5) One may train them with stochastic gradient ascent- desseetee (SGAD) using random batches from P. Q.S. The procedure is summarized in Algorithm 2
Algorithm 2: Neural optimal transport for general cost functionals
Input Distributions P. Q, 8 accessible by samples; mapping network T. R.P. x R$ - RP:
potential network vw E + IR; number inner iterations Kr; empirical estimator F(x, T(X, 2)) for cost F(T);
Output:Learned Learned stochastic OT map To representing an OT plan between distributions PP, Q; repeat
Sample batches Y ~ Q, X ~ P and for each x E X sample batch Z ~ S;
C. + Σ Σ
V. (Ta(z.))
Σ
NEY
Update w by using ac.,
dw
for ky 1,2,... Kr do
Sample batch X ~ P and for each x E X sample batch Z[x] ~ S;
LT +- (xX,To(X,22) Update by using SCr.
Σ Σ (T(Ta(()))),
ZEX (EZ(u) [x|12]]]
00
until not converged;
Algorithm requires an empirical estimator F for F(T). Providing such an estimator might be
non-trivial for general F. If F(T) = C(z, T (-122) (x)). i.c., the cost is weak (2). one may use the
following unbiased Monte-Carle CCrlo estimator: F(x, T(X, 2)) def -.EE (s., T(x Z(r))), where Ō is the respective estimator for the weak cost C and Z(ær denotes random batch of latent vectors for given x E X. For classic costs and the y-weak quadratic cost, the estimator Ō
is given by (Korotin et al. 2023b Eq. 18 and 19) and Algorithm 2 for general 4 to the
neural OT algorithm (Korotin et al. 2023b, Algorithm 1) for weak (2) or classic (I) OT. Unlike the
predecessor, the algorithm is suitable for general OT formulation (4 In $5.1
propose cost
functional Fa to solve the class class-guided dataset transfer task (Algorithm)
5..2 we propose the
functional Fs for the paired domain translataion (Algorithm)
19

---

Published as a conference paper at ICLR 2024
C CLASS-GUIDER EXPERIMENTS
C.I TRAINING AND COMPARISON DETAILS
The code is written in PyTorch framework and publicly available at htttts://ggthhui com/ machinesteine gnot On the image data, our method converges in 5-15 hours on Tesla VTO0
(16 GB). We use WandB for babysitting the experiments (Biewald 2020).
Algorithm details. In our Algorithm 1 for both T3 and Ew- The number of inner iterations for Te is Kr = 10. Doing preliminary experiments,
we use Adam (Kingma & Ba 2014) optimizer with Ir 10- 4
we noted that it sufficient to use small mini-batch sizes x, Ky, K2 in ()) Therefore, we decided to average loss values over K3 small independent mini-batches (each from class n with probability ar)) ) rather than use a single large batch from one class. This is done parallel with tensor operations.
Dataset an hyperparameters. We rescale the images to size 32x32 and normalize their channels to [-1, I]. For the grayscale images, we repeat their channel times and work with 3-channel images. We do not apply any augmentations data. We use the default train-test splits for all the datasets.
We use WGAN-QC discriminator's 's ResNet architecture (He et al. 2016) for potential dw We use UNe Ronneberger et al. 2015) as the stochastic transport map Te 2). To condition it on we insert conditional instance normalization (CondIN) layers after each UNet's upscaling block We use CondIN from AugCycleGAN (Almahairi et al. 2018). In experiments, is the 28-dimensional standard Gaussian noise.
The batch size is KB 32, Kx = Ky = 2. K2 2 for training with a. When training without 2, we use the original UNet without conditioning; the batch parameters are the same (Kz does not matter). Our method converges in R 60k iterations of Uw-
For comparison in the image domain, we use the official implementations with the hyperparameters from the respective papers: AugCycleGAN" Almahhaii al. 2018). MUNITS Huang et al. 2018). For comparison with neural OT 2W W2,1). we use their publicly available code о For the stochastic maps T(x, 2). we only sampled one a per x during computing the metric, no averaging over & was applied. The accuracy of the ResNet18 classifiers is 99.17 for the MNIST and 95.56 for the USPS datasets. For the KMNIST and MNITST, the accuracy's are 99.39 and 97.19, respectively.
OTDD flow details. As in our method, the number of labelled samples in each class is 10. We learn the OTDD flow between the labelled source datase and labelled target samples. Note the OTDD method does not use the unlabeled target samples. As the OTDD method does not produce out-of-sample estimates, we train UNet to map the source data to the data produced by the OTDD flow via regression. Then we compute the metrics on the test (FID, accuracy) for this mapping network.
the same in method. We tested variety of discrete OT solvers from Python Optimal Transport (POT) package (Flamary et al., 202T), including EMD, Trappingrransorttt (Perrot et al.) 2016 and Sinkhorn'Trnnspan
DOT details. Input pre- pre-processing
with Laplacian and I.2 regularization (Courty et al. 2016) from ot da (Flamary et al. 2021). These methods are semi-supervised and can receive labels to construct a task-specific plan. As in our method, the number of labelled samples in each class is 10. For most of these methods, two tunable hyper-parameters are available: entropic and class regularization values. We evaluated a range of these values (1, 2, 5, 10, and 100). To assess the accuracy of the DOT solvers, we used the same oracle classifiers as in all the other cases. Empirically, we found that the Sinkhorn with Laplacian regularization and both regularization values equal to 5 achieves the best performance in most cases. Thus, to keep Table [ simple, we report the test accuracy results only for this DOT approach. Additionally, we calculated its test FID (Table 2
@tthhbb.ccm/alleeaaa////ooooohbp
citthb..com////aas//WWV
hittps://githin cwmp:ccomaaaaaa exxorrcti/Neeut 10t mma Transport
We use only 15k source samples since OTDD is computationally heavy (the authors 2k samples).
20

---

Published as a conference paper at ICLR 2024
C.2 MOONS
The task is to map two balanced classes of moons (red and green) between F and Q (circles and crosses Figure respectively). The target distribution Q is rotated by degrees. The number of randomly picked labelled samples in each target moon is 10. The maps learned by neural OT algorithm with the quadratic cost 2. (Fan et al., 2023 Korotin et al. 2023bb) and our Algorithm
I with functional FG given in Figures 5(c) and ((d) respectively. In Figure 5( we show the
matching performed by a discrete OT-SI algorithm which learns the transport cost with a neural net from known classes' correspondence (Liu al. 2020). As expected, the map for W2 does not preserve the classes (Figure
while our map solves the task (Figure ((d)). We use 500 train and 150 test samples for each moon. We the fully-connected net with 2 ReLU hidden layers size of 128 for both Tp and uu. We train the model for 10k iterations of ew with K = 32, Kx Ky = zKz plays no role as we do use here).
-10-
-12
-10 463
65
-15 -10 -55 do
as io 15
-15 -1 -3 06 26 15
-1 -10 -63 60
(a) Pa. (circles), y Qn. (crosses).
(b) Discrete OT-SI.
(c) Neural OT (E2)
(d) Ours (Fi))
Figure 5: The results of mapping two moons using OT with different cost functionals.
C.3 GAUSSIANS MIXTURES.
(a) Input 27 ~ P.
(b) Target Q
(c) Ours (Fa)
Figure 6: Illustration of the mapping between two Gaussian mixtures learned by our Algorithm
15
10
0.5
-05
0.0
1.5
-0.5
0.0
0.5
3.5
(a) Input 27 ~ P.
(b) Target
Q
(c) Ours (Fa)
Figure 7: Illustration of the mapping between two Gaussian mixtures learned by our Algorithm when the classes are overlapping.
21

---

Published as a conference paper at ICLR 2024
In this additional experiment both P. Q are balanced mixtures of 16 Gaussians, and each color denotes a unique class. The goal is to map Gaussians in (Figure ( to respective Gaussians in Q which have the same color, Figure [6(b] The result of our method (7O known target labels per class) is given in Figure It correctly maps the classes. Neural OT for the quadratic cost is not shown as results in the identity map (the same image as Figure (()))))((a))))))))())) which is completely mistaken in classes. We use the fully connected network with 2 ReLU hidden layers size of 256 for both Z' and Uw There are 10000 train and 500 samples in each Gaussian. We train the model for 10k iterations of ew with Kg = 32, Kx = Ky = 2 (Kz plays no role here as well).
Using the same settings as in the previous experiment, we conducted additional tests involving overlapping classes. In this scenario, specific samples within one class are identical to samples from another class. To execute this experiment, we adjusted the Gaussian modes of each class to be closer to each other, as illustrated in Figure [[a] The target classes are depicted in Figure 7(b) Our method handles this scenario, demonstrating the robustness of our model in handling overlapping classes.
The visual representation can be observed in Figure 7ce)
C.4
EXPERIMENTS WITH OTHER IMAGE DATASETS
3
6
2
4
7
W5 DOT DID0 CON NUNT 7-0 1-h
0
/
2
3
4
5
6
7
8
9
0
1
2
3
4
5
7
7
8
0
2
3
4
5
6
7
8
6
7
8
1
01
1
4
4456
9
NONTT
3
1
01238
3
4
5
7
9 9 2 7 7 8
9
3
3
2 2 7 2
O
2
4
3
6
8
9
To Fouer We. We DOT 0180
2
H
55
6
8 7 2 7 1
4 3 4 2 4
3
4
6
3
3
2
7 8
4 9
&
5
7 3
4
x
6
tours! (ours) 12
2
9
9
8
3
3
3
3
2 2
Tc Founs
2
4 4
C
6
7
8
4
1
2
3
156
6
7
8
8
(oues)
4 1
MNIST USPS
(b) MNIST KMNIST
Figure 8: The results of mapping between two datasets.
Image-to-Imaee Translation
Flows
Discrete OT
Neural Optimal Transport
Datasets (32 x 32)
MUNIT
Aug CycleGAN
OTDD
Sinkhornl.pI.I
W2
Wa.v
no
[Ours]
[Ours]
MNIST - USPS
97.95
98.2
83.26
38.77
37.0
95.27
94.62
MNIST KMNIST
12.27
8.99
4.46
4.27
6.13
6.82
79.20
61.91
Table 3: Accuracyt of the maps learned by the translation methods in view.
Datasets (32 32)
MUNIT
Aug CycleGAN
OTDD
SinkhornI.pI.I
W2
W2.v
Fa. no [Ours]
[Ours]
Fg
MNIST - USPS
6.86
22.74
> 100
51.18
4,60
3.05
5.40
2,87
MNIST
>
9.46
17.26
9.69
Table 4: FID of the samples generated by the translation methods in view.
Here we test the case with different datasets, again in one when the source and target domains related, in the second case when they are not. We consider MNIST- and -MMNSTK As in the main text, unlabeled. The results are shown in Table 3 and Figure X
given only 10 labeled samples from the target dataset; the rest
In this case (Figure 8 GaN-based methods and our approach with our guided cost Fu show high accuracy 2 90%. However, neural OT with classic and weak quadratic costs provides low accuracy (35-50%). We presume that this is because for these dataset pairs the ground truth OT map for
22

---

Published as a conference paper at ICLR 2024
the (pixel-wise) quadratic cost simply does not preserve the class. This agrees with (Daniels et al. 2021 Figure 3) which tests an rerroy--rgullriiedd quadratic cost in a similar MNIST-USSPS setup. For our method with cost FG. The OTDD gradient flows method provides reasonable accuracy -NIST~uUSPS. However, OTDD has much higher FID than the other methods.
C.5
ADDITIONAL VISUALIZATION
0
O
O
2
2
2
3
Y
4
4
5
5
5
6
7
7
7
8
8
9
9 0 0
8
8
/
5
9
5
8
7
5
7
/
7
4
6
5
7
6
4
6
9
8
4
7
7
9
/
2
3
3
3
7
1
8
4
4
O
3
9
0
4
8
5
6
4
9
C
4
0
0
0
0
8
4
8
4
V
7
5
?
7
y
2
9
3
6
7
0
7
4
7
4
4
3
2
/
7
7
J
6
9
9
3
9
9
3
2
1
4
O
2
4
3
9
T
8
0
6
5
2
2
4
8
4
5
9
5
6
7
7
7
8
8
9 9 9
0
2
2
2
4
4
5
5
5
6
6
2
7
7
9
9
Figure 9: FMNIST -MIIST mapping results. Three input images per class are presented.
To further qualitative demonstrate that our model preserves the classes well, we provide an additional visualization of the learned maps. The same models were used in Figure The only difference is that in Figure 6. we show a single input and target per class, while here (Figure 9) we show three inputs and three outputs per class for the different methods.
C.6 ADDITIONAL. EXAMPLES OF STOCHASTIC MAPS
In this subsection, we provide additional examples of the learned stochastic map for Fo (with 2). We consider all the image datasets from the main experiments (. The results are shown in Figure [] and demonstrate that for fixed z and different 2, our model generates diverse samples.
T(xi Z)
X1
n(y(x,)
P X
my)(ss
n(ylxa))
X
T(Xm. 2)
S
Noise
Figure 10: Implicit representation of π E II(P)) via function T Ta X +V.
C.7 ABLATION STUDY OF THE LATENT SPACE DIMENSION
In this subsection, we study the structure of the learned stochastic map for FG with different latent space dimensions Z. We consider MNIST - USPS transfer task (10 classes). The results are shown
in Figures 22 13 and Table 5 As be seen, our model performs comparably for different Z.
Metrics
Z -=1
Z = 4
Z = 8
Z = 16
Z 32
Z 64
Accuracy
86.96
93.48
91.82
92.08
92.25
92.95
FID
4.90
5.88
4.63
3.80
4.34
4.61
Table 5: Accuracyf and FIDJ of the stochastic maps MNIST - USPS learned by our translation method with different noise size Z.
23

---

Published as a conference paper at ICLR 2024
/
2
3
4
S
6
7
8
9
0
1
2
3
4
5
6
7
8
9
y-d DRZzE TRUSI 70U,0) TI..1 TOC
0
23
4
44567
9
011
3
5
7
9
1
223467724
0
1
2
3
5
8
9 9
0
413456781
26
7
7
6
7
O
I
13316778
2
5
6
8
2
3
7/4
3
8
7
0
23
7 5 6 7
5
6
8
9
(a) MNIST > USPS
(b) MNIST
MNIST-M
DA DEES TOK TIARI T3030 X-P
2
3
4
5 6
0
2
3
4
5
6
7
8
9
11
35
0
2
3
4
5
6
7
8
9
2
/
2
3
4
5
6
7
8
4
6
1
2
3
4
5
6
7
8
9
8
/
2
3
4
6
6
7
8
9
3
8
&
2
3
/
2
3
4
5
6
7
8
9
(c) MNIST + KMNIST
(d) FMNIST + MNIST
Figure 11:
Stochastic transport maps To (x(x 2) learned by our Algorithm |
Additional examples.
TIx.sITTx TI. 23 TIx. TIx 2. Tix. 7)),, X-P
3
4
5
3
4
5
3
4
5
3
4
5
3
4
5
3
4
5
3
4
5
3
4
5
3
4
5
3
4
5
3
4
5
3
4
5
3
4
5
3
4
5
3
4 4
5
3
4
6
3
4
5
3
1 1 7 4 4 4
6 5 S 5 5
3
4
5
3
4
5
3
5
3
1
6
3
1
5
3
3
4
5
3
9
5
3
4
5
3
4
6
8
4
5
3
4
5
3
4
5
3
9
6
3
4
8
3
4
S
3
4
5
3
4
5
3
4
5
3
4
5
3
4
5
3
3
5
3
4
5
3
4
5
3
4
5
3
4
5
3
5
3
5
3
C
5
3
4
S
3
1
5
3
1
5
3
4
5
(a) =
(b) 2 = 4
(c)
=
(d) Z = 16
(c) 2 = 32
(f) 2 = 64
Figure 12: MNIST + USPS translation with functional Fu and varying Z 1, 4, 8, 16, 32, 64.
C.8 IMBALANCED CLASSES
In this subsection, study the behaviour of the optimal map for FG when the classes are imbalanced in input and target domains. Since our method learns a transport map from IP to Q it should capture the class balance of the Q regardless of the class balance in P. We check this below.
24

---

Published as a conference paper at ICLR 2024
ZeM Ze.32 Ze16 Ze8 Ze4 2e1 x-Pp
0
1
2
3
4
5
6
7
8
9
0
2
3
4
5
6
7
8
9
0
1123
7
8
9
0123
OO
1123
9
O
3
3
4
8
9
100
2
3
4
5
6
7
8 9
2
3
4
5
6
7
8
9
Figure 13:
Stochastic transport maps To x) learned by our Algorithm with different sizes of Z.
n=0
n=
n=
0
n=0
n=0
n=
n=
n=
1
n= 1
0= 1
/
1
3.1
0
1.6
On T()
2
0.33
3.2
1.7
1
0
0
(a) Examples of transported digits x + To()).
(b) Confusion matrix of (a
Figure 14: Imbalanced MNIST - USPS translation with functional Fu (deterministic, no 2).
We consider MNIST + USPS datasets with n = 3 classes in MNIST and n = 3 classes in USPS. We assume that the class probabilities are 01 Qr 1/2 03 0 and 31 32 = .3 1 That is, there is no class 3 in the source dataset and it is not used anyw during training. In turn, the target class 3 is not used when training To but used when training fu. All the hyperparameters are the same as in the previous MNIST USPS experiments with 10 known labels in target classes. The results are shown in Figures 14(a) and [5(a) We show deterministic (no 2) and stochastic (with 2) maps.
Our cost functional Fa stimulates the map to maximally preserve the input class. However, to transport Q, the model must change the class balance. We show the confusion matrix learned maps T3 in Figures [4(b] 15(b) illustrattes that the model maximally preserves the input classes 0, and uniformly distributes the input classes 0 and 1 into class 2, as suggested by our cost functional.
C.9 ICNN-BASED BASED DATASET TRANSFER
For completeness, show the performance of ICNN-based method for the classic | quadratic transport cost e(x, y) x y on the dataset transfer task. We use the non-minimax version (Korotin et al. 202 the ICNN-base method by (Makkuva et al. 2020). We employ the publicly available code and dense ICNN architectures from the Wasserstein-2 benchmark repository 3 The batch size is KB 32, the total number of iterations is 100k, /r 3- 10 3 and the Adam optimizer is used. The datasets are preprocessed as in the other experiments, see Appendix C.1
The qualitative results for MNIST- and Fashion 16. The results are reasonable in the first case (related domains). However, they are visually unpleasant the second case (unrelated domains). This is expected as the second case is notably harder. More generally, as derived in the Wasserstein-2 benchmark (Korotin al.) (021b), the ICNN models do not work well in the pixel space due to the poor expressiveness ICNN architectures. The ICNN method achieved 18.8% accuracy and >>> 100 FID in the and 13.9 FID in the MNIST case. All the metrics are much worse than those achieved by our general OT method with the class-guided functional Fe. see Table 2 for comparison.
transfer are given in Figure
transfer, and 35.6% and accuracy
i hhu.. com ama www.oelkkooooooo wtc.m//maalexoorooooooooo h
25

---

Published as a conference paper at ICLR 2024
R-0
n-1
TIx. Zol
2
2
2
y- Q T(x. 24 T
2
3
1.7
0.16
3.3
1.9
0
(a)
Examples of transported digits x + (2 =), ~ S.
(b) Confusion matrix of (Tr, =).
Figure 15: Imbalanced MNIST + USPS translation with functional Fu (stochastic, with 2).
0
1
2
3
4
5
6
7
8
9
KCNN yoG
0
2
3
4
5
6
7
8
IONN y-:0,
2
3
4
5
6
7
8 9
2
3
4
$
6
UNIST-+USP transfer.
(b) FMNIST- transfer.
Figure 16: Results of ICNN-based method applied to the dataset transfer task.
C.10 CLASSIC COST OT FOR DATASET TRANSFER
Our general cost functional-bes based algorithm training, which can be useful for the data transfer tasks. Existing continuous OT approaches do not handle such type of training. Indeed, suppose we have additional information (labels) in the dataset and try to solve the class class-guid guided mapping using the ICNN-basee (Amos et al.) 2017) W2 algorithms (Korotin al. 2023b] Fan et al.) (023). In this scenario. we can train OT using only the labeled samples (10 separate maps in case of MNIST). The unlabeled data immediately becomes useless. Indeed, using unlabeled data for a class during training for that class implies that we know the labels
both labeled and unlabeled target samples for
for that class, which is contradiction.
O
(a) FMNIST -NISST transfer on '0' class.
Figure 17: Results of classic OT (W2 method applied the dataset transfer task.
For illustrative purposes, we performed these experiments using the We algorithm for one of the clases the FMNIST to MNITS mapping problem. We clearly see that the qualitative results not competitive with our algorithm. This is because W, is forced to train with only 10 target samples, the only labeled target samples in the problem setup.
26

---

Published as a conference paper at ICLR 2024
C.I1 NON-DEFAULT CLASS CORRESPONDENCE
0
2
3
5
6
8
9
9
\
2 3
4
5
6
7
8
Jours!
0
/
2
3
4
5
6
7
8
Figure 18: FMNIST -NIST mapping with FG no cost, classes are permuted.
To show that our method can work with any arbitrary correspondence between datasets, we also consider FMNIST->MNIST dataset transfer with the following non non-default correspondence between the dataset classes:
0 - -
7 6, 7,9+
In this experiment, we use the same architectures and data preprocessing as in dataset transfer tasks; Appendix C.T We (no 2). In this setting, our method produces comparable results to the previously reported Section 6 accuracy equal 83.1, and FID 6.69. The qualitative results are given in Figure
Fa 7 the cost functional and learn deterministic transport map T
C.12 IN DOMAIN CLASS-PRESERVING
6
1
2
3
4
5
6
7
8
9
O
lourshyy-n On
0
1
2
3
4
5
6
7
8
9
Pc
0
/
2
3
4
5
6
7
8
9
Figure 19: MNIST mapping with Fa. classes are permuted.
To provide an additional illustration, we also consider MNIST-MNNE dataset transfer with the following non-default correspondence between the dataset classes:
0 9, 0, 1,3 2, 3,5 4, 6,6 5,7 6, 8 7,9 8.
In this experiment, we use the same architectures, data preprocessing and metrics as in dataset transfer tasks (C.I). We our Fo (7) the cost functional and learn transport map T. The resulted accuracy is equal to 95.1, and FID 3.35. The qualitative results are given in Figure 19
C.13 SOLVING BATCH EFFECT
The batch effect is well-known issue biology, particularly in high-throughhut genomic studies such as gene expression microarrays, RNA-seq, and proteomics (Leek et al.) (210)). It occurs when non- biological factors, such as different processing times or laboratory conditions, introduce systematic variations in the data. Addressing batch effects is crucial for ensuring robust and reproducible findings in biological research (Lazar et al. 2013). By solving this problem using our method, directly in the input space, we can preserve the samples' intrinsic structure, minimizing artifacts, and ensuring biological validation.
In our experiments, we map classes across two domains: TM- WM-baron-mane Mharon-mouse-for monouse-oor-segesttlll and segerstolpe-humen consisting of 3,329 and 2, 2,108 samples, respectively. The data was generated by the Splatter package (Zappia et al., 2017). Each domain consists of eight classes. The source domain IP is fully-labelled, and the target Q contains 10 labelled samples per class, Each sample is 657-sized vector pre-pot pre-procesee with default Splatter settings. (Zappia et al. (017).
We employed feed-forward networks with one hidden layer of size 512 for the map T3 and hidden layer of size 1024 for the potential network ew To evaluate the accuracy, we trained single-layer neural network classifiers with soft-man output activation, using the available target data. Our method improved accuracy from 63.0 + 92.5. Meanwhile, the best DOT solver (EMD) identified through search, described in Appendix C.I reduced accuracy from 63.0 - 50.4.
27

---

Published as a conference paper at ICLR 2024
D GENERAL FUNCTIONALS WITH CONDITIONAL INTERACTION ENERGY REGULARIZER
Generally speaking, for practically useful general cost functionals F M x }) - - [+oo}
it may be difficult or even impossible to establish their strict or strong convexity. For instance. our considered class-guided functional FG 7 is not necessarily strictly convex. In such cases, the maps T* which solve (5) are not necessarily stochastic OT maps, and our duality gap analysis (Theorem 3) is not directly applicable.
In this section, we propose generic way to overcome this problem by means of strongly convex regularizers. Let F, R: M(X x )) + R U {+00) be convex, lower semi-continuous functionals, which are equal to too μ M(X y) P(X )). Additionally, we will assume that R is 3-strongly convex on (I) in some metric p(..). ). For y > O, one may consider the following R-regularized general OT problem:
inf xEM((.))
(F(T) + +R(T))
Note that π + F(x) + +(((()) is convex, lower semi-continuous, separately * *-increasing (since it equals +00 outside π E P(X J), see the proof of Theorem in Appendix [A] and f7-strongly convex on (IP)) in p(s In the considered setup, functional } corresponds to real problem a practitioner may want solve, and functional
the regularizer which slightly shifts the resulting solution but induces nice theoretical properties. Our proposed technique resembles the idea of the Neural Optimal Transport with Kernel Variance (Korotin et al. 2023a). In this section, we generalize their approach and make it applicable to our duality gap analysis (Theorem 3) Below we introduce an example of a strongly convex regularizer. Corresponding practical demonstrations are left to Appendix D.I
Conditional interaction energy functional. Let (V,1) be a semimetric space of negative type (Sejdinovic et al.) 2013 $2.1), L: y>y' + IR is the semimetric and V 2 2, 1/, 1/2...................... ....... MN E y and Vo1, -...-.. ON IR such that Σ Og it holds Σ Σ 'n'= =1 Onp Fn') S 0. The (square of) energy distance & w.r.t. semimetric / between probability distributions Q1, Q2 E P())
(Sejdinovic et al. 2013 $2.2)):
E (Q Q2) = 21E Y2) - BI(Y, Y1 EI(Y,,
(38)
where Yi,Y; ~ Q1: Y2, Y ~ Q2. Note that for I(y, y/)= = 1 ly 1/12 formula 38 reduces to (8)
The energy distance is known be metric on P()) (Klebanov al. 2005) (note that y is compact).
The examples of semimetrics of negative type include I(y, v/) = || vll mln for 0 < p S 2
(Meckes) 2013 Th. 3.6).
Consider the following generalization of energy distance & on space (IP). Let T2 E (TP).
(T T2) del / (+1)) T2 (a())))
(39)
Proposition 2. It holds that pi(.s) :) is metric on II((P).
Proof of Proposition 2 Obviously, Vi E IT(P): : p(( T) = 0 and VTi, 22 E IT(P): pr (m) π2) =
(T(TT, 11) 2 O. We are left check the triangle inequality. Consider T1, T2, T3 [I()). In what
follows, for π E (I())) and x E X, we denote the conditional distribution TT (-|2)) as -":
Pi(TII, m2) + (T(; T3) > pt (TI, (T) +
(p((TT, T2) (1 mz))² 2 > p [T T3) +
[EP(TR, mZ) xz)pP(x)+ ( (355 xT)))())+ +2 87 (T.. m3 (aP())) E7 (Tž T dF) of T5 dTP = (40)
x
I (( E7 (T)d 2 Ei Tz) ( dP))) >
E7 dP(c))
x
(E((T), + ( - 87 m5) dT 2 0,
20 due triangle incquality
28

---

Published as a conference paper at ICLR 2024
where in line 40 we apply the Baluchy-Bunyakova inequality (Bouniakowsky 1859):
J m mdPP())))
md)((()) (/x E (3)
Now we are ready to introduce our proposed strongly convex (w.I.t pt) regularizer. Let т E (I))
and be semimetric y of negative type. We define
1 R₂(x)= - / f Uy /dd///dd" (y)da" (y)dd ()))d"(/))P(())))) (f)d)(z))
(41)
We call R₂ be conditional interaction energy functional. In the context of solving the OT problem, it was first introduced in (Korotin et al. 2023a) from the perspectives of RKHS and kernel embeddings (Sejdinovic al. 2013) $3). The authors of (Korotin al., (023a) establish the conditions under which the semi-dual (max-min) formulation of weak OT problem regularized with R2 yields the unique solution, i.e., they deal with the strict convexity of R2, contrast, our paper exploits strong convexity and provides the additional error analysis (Theorem 3 which helps with tailoring theoretical guarantees to actual practical procedures for arbitrary strongly convex functionals. Below, we prove that R, is strongly convex w.r.t. pr and, under additional assumptions on I, is lower semi-continuous.
Proposition 3. R2 is 1-strongly convex II(P) w.r.t. pt.
Proof of Proposition
Let TT, T2 E I(P). S a S 1. Consider the left-hand side of 21
Re + (1 a)T2) =
J / 1(V, Y'lo (1 ())=)=))))))))))))
(1 T2)= (())d((x))
x yxx
1
a²
((,,,')ddT ddP(2)
2
x Nxyy
-a(1 a) / [(V, (lv.)")dT (F)
yxy
1
(1 a)² /
/ 1( (1 (P))
2
yxxy
a a(1 11-))))) 01) 1 / 1(v, ((v.v")dT (u) dF
1
- 2 2
vx3
-a(1 a) / /
(P)
(
1 (1 (1-
2 aa - a)
1
/
I(v, "( dP(())
2 0/ Se Uy v/ ddd (v)dd (v/ )dP())
1
1 a)
1(v, 1/ (1 (dP(x) +
2
x vxx
=aR((A1)
=(1-a)R((x2)
a(l-a)
[ Mly (yy.v'ddii s()drr(()d (v')+ My., vf )dzz (v) 2/ /
1(y, v j')dr( (1) dP(xx)
2
Jyxy
"yxy
=-E((af,mz)
(R((TI) + (1 R _a(1 - pf (m,, m2):
i.e., ( a(1-a) pp (MT,,T)) which finishes the proof.
Propusition 4. Assume that is continuous (it is the case for all reasonable seminetrics 1). Then R₁ lower semi-continuous on (PP).
29

---

Published as a conference paper at ICLR 2024
Proof of Proposition Consider the functional W/ & P()) + -RUffOO
W1 Lox Uy.. (gg-))ddy(y)dull)
then the conditional interaction energy functional could be expressed as follows: ((( (x (x)). We are check that Wi satisfies Condition (A+) in (Backhoff- Veraguas et 2019 Definition 2.7). Note that W/2 actually does not depend on X.
The lower--seiicottinnity of Wi follows from (Santambrogio 2015 Proposition 7.2) and the equivalence of the weak convergence and the convergence w.r.t. Wasserstein metric on P())) where y is compact, see (Villani 2008 Theorem 6.8).
Since (y. 1/) + -,y, y') a lower-semicontinuous, achieves its minimum on the compact Y y which lower-bounds the functional Wi
The convexity (even 1-strong convexity w.r.t metric E((:.)) of functional Wi was de facto
established Proposition 3
In particular, given Pl, μ2 E P()). a E (0, 1), E X it holds: a(1 - aa) &
Wi amn (ott::/wwwtHtoooooooooaa/aaaat--
The application of (Backhoff-Veraguas et al.
2019 Proposition 2.8, Eq. (2.16)) finishes the proof.
D.1 EXPERIMENTS WITH CONDITIONAL INTERACTION ENERGY REGULARIZER
In the previous Section D we introduce example of the strongly convex regularizer. In this section, we present experiments to investigate the impact strongly convex regularization on our general cost functional FG- In particular, we conduct experiments on the FMNIST-MNIST dataset transfer problem using the proposed conditional interaction energy regularizer with (y,, y') = Jy--l 1ly-y'l Il2. To empirically estimate the impact of the regularization, we test different coefficients [0.001, 0.01, 0.1]. The results are shown in the following Figure 20 and Table
0
1
2
4
5
6
7
9
0.001 00 Y= C.c
0
2
3
4
5
6
7
8
9
0
1
2
3
4
5
6
7
8
9
0
1
2
3
4
4
6
7
8
9
0.01
/
21
3
4
5
6
7
8
9
2
3
0
S
6
7
8
9
Figure 20: Qualitative results of the FMNIST -MNITT mapping with Fo cost and using different values y of the conditional interaction energy regularization.
0.001
0.01
0.1
Accuracyt 83.33
81.87
79.47
65.11
FID 4
5.27
7.67
3.95
7.33
Table 6: Accuracyt and FID J. of the map learned on FMNIST with Fo cost and different values y of the conditional interaction energy regularization.
can be seen that the small amount of regularization (7 = 0.001) does not affect the results. But high values decrease the accuracy, which is expected because the regularization contradicts the dataset transfer problem. Increasing the value of y shifts the solution to be more diverse instead matching the classes.
30

---

In our experiments, we compare several methods. As the baselines we consider (unsupervised) NOT, and RMSE regression. Here use U2Net ||| as the transport map To (=) (generator). For comparison employ the (unsupervised) RMSE the Published as a conference paper at ICLR 2024
E PAIR-GUIDED COST FUNCTIONAL
E
ALGORITHM
Recall that in our main manuscript we parameterize the learned plan via stochastic map [z, (z, z)]. ~ IP, 2 ~ 8. In practice, we found that substitution T(x, 2) with deterministic map T(x) generally improves the results in the case of pair-guite pair-guided cost functional. This is possibly due to the paired nature of the considered problem. Therefore, we adapt our proposed Algorithm map T (2) and report all metrics and demonstrations exactly for this setup.
3 for deterministic
Algorithm 3: Neural optimal transport with pair-guided cost functional and deterministic map Input Distributions P, Q accessible by samples; paired data set (æ1, (z1)).... (Ir N.N.21 y (IN)), where Ti:N ~ PP and y (FI:N) ~ Q; mapping network
To R
R potential network % IR + IR; number of inner iterations Kr;
Output: Learned OT map Ta representing (pair-guided) OT plan between distributions Q
repeat
Sample batches Y ~ Q ~ P: (v(,) IYT
C
(Ta(2)) IXI
Σ
VEY
Update by using 8Cu: for kr 1.2,.... KT do
Sample batch Xa from the paired data set; A(Talma). (xa))
C Σ
[Xal
2aEXa
#tEXp
Update 0 by using DCr. 00
until not converged;
E.2 DATASETS
Comic-Faces- V 7 This dataset contains paired samples which are useful for real-to-comic convertion. The original resolution is 512x512, 10000 pairs (total 20k images)
Edges-to-Shoes: This dataset consists of 50,025 shoe images and their corresponding edges split into train and test subsets.
CelebA Ma This is large-scale face image dataset that has 30,000 high-resolution face images selected from CelebA-HQ dataset. Each image has segmentation mask of facial attributes corresponding to CelebA. The masks of CelebAMask-HH were manually-annotace with the size of 512 x 512 images.
E.3 TRAINING DETAILS
with NOT, we use their publicly available code 12 There cost function (the method is denoted by W the table and figures). Note that unsupervised NOT (Figures 4 and 21 fails to perform the translation in both the cases. For the comparison with Pix2Pix, we use the official implementations with the default hyperparameters
To train model, use Adam Kingma& & Ba 2014 optimizer with Ir 10 for both Ta and Uu- The number of inner iterations for To is K 10. The batch size of KB was set for the
comic-faces and shoes experiments. The batch size of KB 32 was used for Mas- Our method converges in z 60k iterations of ey-
synthetic
httter//itttinnoorrrrtttaaaoooos H0
31

---

Published as a conference paper at ICLR 2024
F,
Datasets (256 256)
W1
RMSE
Pix2Pix
[Ours]
Komic-Faces-11
100
79.16
37.02
35.42
Edges to io-Shoes
100
61,55
49.53
Table 7: FID 4 of the maps learned by the translation methods in view.
RMSE (x) X-
W1
Ps (ours)
Figure 21: Results for Edges-to-Shoce Shoes with the Pair-guided cost, images resolution is 256 x 256.
bu
Figure 22: Results for CelebAMask with the VGG-based perceptual Pair-guided cost, images resolution is 256 x 256.
X
T(x)
Figure 23: Results for Comic-Faces-Vi V1 with the Pair-guided cost, images resolution is 512 x 512.
32


Reviews:
[{"review_id": "Z4FmwKaBLV", "reviewer": "Reviewer_nKa4", "summary": "This paper proposes to study the so-called general cost OT problem, that is $\\min_{\\pi\\in\\Pi(P, Q)}F(\\pi)$ where $F$ is a general functional and $|\\Pi(P, Q)$ is the set of couplings between $P$ and $Q$. This is called general cost OT, because classical OT is a special case, taking $F(\\pi) = \\int c(x, y) d\\pi$.\nThe authors derive a max-min reformulation for the approach as $\\sup_{v}\\inf_{\\pi\\in \\Pi(P)} F(\\Pi) - \\int v d\\pi(y) + \\int v dQ(y)$, and then propose to parametrize the coupling $\\pi$ as samples from $[x, T(x, z)]$ where $x\\sim P$ , $z$ are samples from a Latent distribution like a Gaussian, and $T$ is a map that can itself be parameterized by a neural net. The potential $v$ is also parameterized by a NN.\nIn cases where the functional $F$ has a nice structure, the above max-min formulation can be estimated from random samples of $P, Q$ and $z$. \n\nThe authors then provide an error analysis for the method where $F$ is strongly convex.\n\nThe main application of the method discussed in the paper is to do optimal transport that is faithful to labels. Given two mixtures $P = \\sum \\alpha_n P_n$ and $Q = \\sum \\beta_n Q_n$, the authors want to estimate a transport plan between P and Q that should also map, as well as possible, each $P_n$ to $Q_n$. The corresponding cost function, a sum of energy distances between $T\\\\#P_n$ and $Q_n$, fits nicely into the proposed general cost framework, as it can be estimated from samples using a (costly) U-statistics, and is different from classical OT costs.\n\nThe experiments are on toy MNIST, KMNIST and fashion MNIST datasets, where the authors try to match each class. They compute the corresponding FID between mapped source and target dataset, and accuracy on the mapped set of a resnet trained on the target set.", "questions": "See above.\n\n\nMisc. minor remarks:\n- in the abstract, the use of transport \"map\" vs \"plan\" can be confusing\n- Why write $\\mathcal{X} = \\mathcal{Y}$ in the notation? why bother with two spaces if they are the same.\n- $\\gamma$ is used for two different things on top of page 3\n- \"The performance of such methods in high dimensions is questionable\": a reference to elaborate on this would be welcome.\n- In eq.6, the second $\\sup \\inf$ can be removed to make it clearer what $\\mathcal{L}$ is.\n- middle of page 5: \"it follows by solving (2)\"  I think this refers to the unlabeled equation in corollary (2), not to eq.(2). Same with \"Overall, Problem (2)\" shortly after. It would probably be best to number the equation in corollary 7.\n- \"it does not always have a solution with each P_n exactly mapped to Q_n\" a few more words about why this is the case would be nice.\n- \"we employ the Sinkhorn\"\n- what is the baseline accuracy of each resnet used in the experiments?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is very well written and easy to follow.\nThe algorithm developed in the paper is sound, and is an interesting method to estimate transport plans for general costs.\nThe theoretical results of the paper are interesting. \nThe problem of optimal transport with label faithfulness is also interesting for ML applications, and the proposed method is an elegant solution.", "weaknesses": "The main weakness of this paper is the experimental validation.\n- the setup is very toyish: the datasets are toy datasets, and the labels are entirely unrelated: why match the digit '1' to 'trouser'? This is fine as a first toy experiment but there must be some more interesting ML applications where the labels from P and from Q have a relationship and are not paired randomly: having only this artificial experiment is underwhelming, and does not convince the reader that it is actually an interesting problem for machine learning.\n- what does fig.3 show ? the description is far too short. Some methods are stochastic (i.e. T(x, z) with z random), how is the sample chosen? \n- Same question for the metrics: how it is computed for stochastic outputs could be clearer. is it averaged over z?\n- is FID computed per class or on the whole dataset? this should be clarified.\n- all fonts are too small in the figures\n\nAnother point is that the cost function proposed in Prop.1 contains a quadruple sum over samples: a discussion about its variance would be welcome. Also, the proposition mentions that it is an estimator: in which sense?", "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": "mGZ9tp2Myw", "reviewer": "Reviewer_LT3k", "summary": "The paper discusses the problem of computing optimal transportation plans for a general cost functional. In particular, the paper proposes a max-min formulation of the problem, provides theoretical consistency results, and a stochastic optimization algorithm to numerically solve it. The problem is motivated by a \"dataset transfer problem\" where class labels are required to be preserved with the transportation plan. The proposed algorithm is illustrated on an a toy example with moon dataset and an example involving image dataset.", "questions": "- It is interesting to see and discuss how the algorithm performs when the data from classes overlap. \n- Regarding the discussion at the beginning of Sec 3.2., why is there a \"measurable\" map for each coupling? I understand existence for each x, but not sure how to argue existence of a measurable map as a function of x and z.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The topic of the paper is interesting and valuable to the researchers\n- The paper is written with mathematical rigor\n- The proposed algorithm is supported by theoretical arguments and numerical experiments \n- The theoretical results are important and useful", "weaknesses": "1- The novelty of the paper, in comparison to Ref[1], is weak. \n- The theoretical novelty, in comparison to the existing theoretical result in Ref[1] is not explained well (what is the new approach or new tool that is being used here). What are the challenges of considering a general cost functional that the previous approach could not handle.   \n- The computational algorithm is also very similar. The NOT algorithm block in Ref[1] can be simply extended to general cost functional. The proposed algorithm block is very similar, with the only difference that it is written for a class-guided functional.       \n\n2 - As I was reading the paper, I found the motivation, the theoretical discussion, and the numerical examples a bit disconnected. If the main contribution of the paper is the OT with general cost functional, it is necessary to provide several examples of a general cost functional, rather than just focusing on class guided cases. If the main goal of the paper is to do the class guided transportation, it should be reflected in the title, there should be more motivation why this is useful in practice, and what are the existing approaches for this particular problem. \n\n3 - The comparison with the existing OT approaches does not seem fair as they do not optimize the same cost function as your apporach. A possible comparison is to use the existing OT algorithms to do to transportation for each class separately, resulting in 10 different maps (for the MNIST case) and discuss how your proposed method, which only trains one map, is computationally more efficient, while the loss in accuracy is not significant.\n\n[1] Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. In International Conference on Learning Representations, 2023", "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": "X5GIVGEh4V", "reviewer": "Reviewer_iaG1", "summary": "This paper introduces a novel neural network-based algorithm for computing optimal transport plans (OT) for cost functionals that go beyond typical Euclidean costs like $\\ell^1$ or $\\ell^2$. These functionals offer greater flexibility and allow the incorporation of auxiliary information, such as class labels, in constructing the transport map. \n\nExisting methods for general costs are discrete and lack out-of-sample estimation capabilities. The paper addresses the challenge of designing a continuous OT approach for general costs that can generalize to new data points in high-dimensional spaces like images. Additionally, it provides theoretical error analysis for the recovered transport plans. \n\nAs an application, the paper demonstrates how to construct a cost functional that maps data distributions while preserving class-wise structures.", "questions": "- A direct comparison of the theoretical findings with those presented in (Korotin et al. 2023a) is essential. Readers are likely seeking a consolidated presentation rather than having to review two separate papers with overlapping theoretical content.\n- From my perspective, the paper's novelty appears to lie more in its practical applications, particularly in utilizing a cost functional to map data distributions while maintaining class-wise structures. In light of this, it would be beneficial for the paper to allocate more of its main text to discussing practical implementation aspects.\n- From my own interests, I would find it valuable if the paper could delve further into potential applications involving general cost functionals.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper is well-structured with a clear motivation, thorough literature review, rigorous theoretical analysis, comprehensive numerical experiments, detailed implementation, and insightful discussions.\n- The paper expands upon existing neural OT techniques to accommodate general cost functionals, offering potential applications in mapping data distributions while maintaining class-wise structures.", "weaknesses": "- The paper heavily draws upon the prior work of (Korotin et al. 2023a) for its theoretical foundations. Approximately 5 out of 9 pages are dedicated to presenting these theoretical results, which may not constitute the primary novelty of the paper. Based on my personal reading, it seems that the authors might be overselling their theoretical contributions, potentially leading to a less reader-friendly introduction.\n- The paper exceeds the strict upper limit of 9 pages for the main text of the submission by including a section on reproducibility (Section 7) on Page 10. As a reviewer, this doesn't pose an issue for me, but I would advise adhering to the prescribed page limits as a matter of following the rules.\n\n\n[Korotin et al. 2023a] Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Kernel neural optimal transport. In\nInternational Conference on Learning Representations, 2023a.", "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": "I3hhIkd7xD", "reviewer": "Reviewer_BH6f", "summary": "This paper proposes a general OT formulation, which uses a functional \\mathcal{F} to encompass common objectives and regularizers as special cases. The authors introduced a method for addressing the continuous general OT problem and illustrated how a general functional \\mathcal{F} can incorporate information, such as the presence of class labels in the data. To validate the method, the authors use synthetic datasets and various MNIST datasets for testing.", "questions": "See the \"Weaknesses\" section above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The closest prior work appears to be the study by Korotin et al. (2023b), as cited in the paper. Building upon this previous research, the authors demonstrate how one can preserve the class-label structure in OT. This contribution is novel to my knowledge.", "weaknesses": "As I understand it, one benefit of employing a general functional F is to account for the class-label structure. Are there any other intended applications of a general F? If the sole purpose is to consider the class-label structure, perhaps some proofs (e.g., the proof to Theorem 3) could be simplified.\n\nI do not fully understand the image data experiments depicted in Fig 3(a) and (b) in section 5.2, and I would appreciate it if the authors could provide further explanation. From my understanding, the goal here is to identify an optimal transport (OT) map between two data distributions (e.g., the distribution of MNIST and KMNIST images) while preserving the class correspondence. To achieve this, one could visualize several source images from the same class in the source dataset and check if the corresponding target images are from the same class. However, Fig 3 appears to display only a single source image per class (top rows) along with a target image per class (2nd rows), making it unclear whether the class correspondence is preserved overall. \n\nAnother reason why Figures 3(a) and 3(b) are challenging to understand is the absence of a natural correspondence between the classes of MNIST -> KMNIST and FMNIST -> MNIST images. It is thus difficult to check the correspondence visually. To enhance visualization, would it be reasonable to use only a single dataset, such as MNIST, and establish the correspondence of images from class 0, 1, 2, ..., 9 to a permuted class order, such as 1, 2, 3, ..., 9, 0? This approach would help to visualize the class correspondence.", "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)