Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation
Authors: Luca Eyring
Dominik Klein
Théo Uscidda
Giovanni Palla
Niki Kilbertus
Zeynep Akata
Fabian J Theis
Venue: ICLR 2024 poster (2024)
Keywords: optimal transport
domain translation
image translation
flow matching
Primary area: generative models
TL;DR: We propose a theoretically grounded method to incorporate unbalancedness into any Monge map estimator and show how unbalancedness yields enhanced results across three distinct tasks employing three estimators.

Abstract: In optimal transport (OT), a Monge map is known as a mapping that transports a source distribution to a target distribution in the most cost-efficient way. Recently, multiple neural estimators for Monge maps have been developed and applied in diverse unpaired domain translation tasks, e.g. in single-cell biology and computer vision. However, the classic OT framework enforces mass conservation, which
makes it prone to outliers and limits its applicability in real-world scenarios. The latter can be particularly harmful in OT domain translation tasks, where the relative position of a sample within a distribution is explicitly taken into account. While unbalanced OT tackles this challenge in the discrete setting, its integration into neural Monge map estimators has received limited attention. We propose a theoretically
grounded method to incorporate unbalancedness into any Monge map estimator. We improve existing estimators to model cell trajectories over time and to predict cellular responses to perturbations. Moreover, our approach seamlessly integrates with the OT flow matching (OT-FM) framework. While we show that OT-FM performs competitively in image translation, we further improve performance by
incorporating unbalancedness (UOT-FM), which better preserves relevant features. We hence establish UOT-FM as a principled method for unpaired image translation.

Published as a conference paper at ICLR 2024
UNBALANCEDNESS IN NEURAL MONGE MAPS IMPROVES UNPAIRED DOMAIN TRANSLATION
Luca Eyring!.200 1,2,3+ Dominik 2lein Théo Uscidda Giovanni Palla
Niki Kilbertus2.3.5 Zeynep Akatal.2.3
Fabian Theis
TTbbineen AI Center
Heimmolttz Munich
TTU Munich
'CREST-ENSAE MMunich Center for Machine Learning (MCML)
luca. eyring aunituubinnenn {dominik. klein, f fabian. theis whweeseeeeeeeeeeeeeeeeeeeeeeaele mew.hhhhhttmmmmcenn... de
theo. seseiiddaeessee... fr
ABSTRACT
In optimal transport (OT), Monge map is known mapping that transports source distribution to a target distribution in the most cost-efficient way. Recently, multiple neural estimators for Monge maps have been developed and applied in diverse unpaired domain translation tasks, c.g. in single-cell biology and computer vision. However, the classic OT framework enforces mass conservation. which makes it prone to outliers and limits its applicability in real-world scenarios. The latter can be particularly harmful OT domain translation tasks, where the relative position of a sample within distribution is explicitly taken into account. While un- balanced OT tackles this challenge in the discrete setting, its integration into neural Monge map estimators has received limited attention. We propose theoretically grounded method to incorporate unbalancedness into any Monge map estimator. We improve existing estimators model cell trajectories over time and to predict cellular responses to perturbations. Moreover, our approach seamlessly integrates with the OT flow matching (OT-FM) framework. While show that OT-FM performs competitively in image translation, we further improve performance by incorporating unbalancedness (UOT-FM), which better preserves relevant features. We hence establish UOT-FM as a principled method for unpaired image translation.
INTRODUCTION
Unpaired domain translation aims to transform data from source to target distribution without access paired training samples. This setting poses the significant challenge of achieving meaning- ful translation between distributions while retaining relevant input features. Although there many ways to define the desired properties of such a transformation, optimal transport (OT) offers a natural framework by matching samples across distributions in the most cost-efficient way. If this optimal correspondence can be formulated as a map, such a map is known as a Monge map.
Recently, considerable number of neural parameterizations to estimate Monge maps have been proposed. While earlier estimators were limited to the squared Fuclidean distance (Makkuva et al., 2020; Korotin al., 2020; Amos, 2022), more flexible approaches have been proposed recently (Uscidda & Cuturi, 2023; Tong et al., 2023b; Pooladian et al., 2023a;b). Neural Monge maps have been successfully applied to a variety of domain translation tasks including applications in computer vision (Korotin et al., 2022; Tong et al., 2023b; Pooladian et al., 2023a; Mokrov et al., 2023) and the modeling of cellular responses to perturbations (Bunne et al., 2021).
In its original formulation, optimal transport assumes static marginal distributions. This can limit its applications as it cannot account for [i] outliers and [ii] undesired distribution shifts, e.g. class imbal- ance between distributions overcomes these limitations by replacing the conservation of mass constraint with a penalization on mass deviations. The practical significance of unbalancedness in discrete OT has been demonstrated
visualized in Figure 1. Unbalanced OT (UOT) (Chizat et al., 201 8a)
"qqual contribution

---

Published as a conference paper at ICLR 2024
Target V
Rescaled Target v v -== P2tTUOT
B
O
17
Unbalanced OT
8
7
a a
Tuor = UOT(H
0 O 0
0
3
Z 7
88
Source μ
Rescaled Source u. 1A = P1 PitTUOT
Figure 1: Comparison of balanced and unbalanced Monge map computed on the EMNIST dataset translating digits - letters. Source and target distribution are rescaled leveraging the unbalanced OT coupling. The computed balanced mapping includes - {0, B), and - {O, because of the distribution shift between digits and letters. With unbalancedness 8 + B, and 1 I are recovered.
in various applications. e.g. in video registration (Lee et al., 2020), computer vision (Plaen et al., 2023), or domain adaptation (Fatras et al., 202 aa). Existing methods for estimating neural Monge maps with unbalancedness (Yang & Uhler, 2019; Lübeek et al., 2022) are limited to specific Monge map estimators and rely adversarial training, Section 4 for detailed discussion.
In light of these limitations, we introduce a new framework for incorporating unbalancedness into any Monge map estimator based re-scaling scheme. We motivate our approach theoretically by proving that we can incorporate unbalancedness into neural Monge maps by rescaling source and target measures accordingly. To validate the versatility of our approach, we showease its applicability on both synthetic and real-world data, utilizing different neural Monge map estimators. We highlight the critical role of incorporating unbalancedness to infer trajectories in developmental single-cell data using ICNN-bare based estimators (OT-ICNN) (Makkuva et al., 2020) and to predict cellular responses to perturbations with the Monge gap (Uscidda & Cuturi, 2023).
Monge maps can also be approximated with OT flow matching (OT-FM) (Lipman et al., 2023; Liu al., 2022: Albergo & Vanden-
2023), simulation-free technique for training continuous normalizing flows (Chen et al., 20 relying on mini-bated batih OT couplings (Pooladian et al., 2023a; Tong al., 2023b). The universal applicability of our method allows us to extend OT-FM to the unbalanced setting (UOT-FM). We demonstrate that unbalancedness. is crucial for obtaining meaningful matches when translating digits to letters in the EMNIST dataset (Cohen et al., 2017). Additionally, we benchmark OT-FM unpaired natural image translation and show that it achieves competitive results compared to established methods. Moreover, UOT-FM clucidates the advantages of unbalancedness in image translation as it [i] improves overall performance upon OT-FM while additionally [ii] helping to preserve relevant input features and lowering the learned transport cost. new principled method for unpaired image translation. To summarize:
This establishes UOT-FM
1. We propose an efficient algorithm to integrate any balanced Monge map estimator into the unbalanced OT framework.
2. We theoretically verify our approach by proving that computing Monge map between measures of unequal mass can be reformulated computing Monge map between two rescaled measures.
3. We demonstrate that incorporating unbalancedness yields enhanced results across three distinct tasks employing three different Monge map estimators. We find that our proposed approach enables the recovery of more biologically plausible cell trajectories and improves the prediction of cellular responses to cancer drugs. Furthermore, our method helps to preserve relevant input features on unpaired natural image translation.
2 BACKGROUND
Notation. For C. RRd, MM) and (f) are respectively the set of positive measures and probability measures on Q. For E (M), we write μ < La if μ is absolutely continuous w.r.t. the Lebesgue measure. For a Lebesgue measurable map T Ω + Ω and μ E M MM2).. Tipe denotes the pushforward of μ by T, namely, for all Borel sets A C Ω, Tit(A)
p(T (A)). For
μ, V E M (2), I1(, v) := { pT1T = μ. P2t = v) C P(2 x 12) where p1 (x,y) + x and P2 (x, y) y are the canonical projectors, pat and p2tT are the marginals of m.
2

---

Published as a conference paper at ICLR 2024
Monge and Kantorovich Formulations. Let Ω C. R be compact set and Ω f2 - R continuous cost function. The Monge problem (MP) between Ju, V E MI (S2) consists of finding a map T 92 + 52 that mush-forwards j onto v, while minimizing the average displacement cost
inf / c(x, T(x)) du(x)).
(MP)
T:TH
We call any solution T Problem (MP) Monge map between j and v. Solving this problem is difficult, especially for discrete μt, v, for which the constraint set can even be empty. Instead of transport maps, the Kantorovich problem (KP) of OT considers couplings л E V);
Welu,)) = min [ c(x, dx( #Ell(u.) Rx
(KP)
An OT coupling #oT. a solution of (KP), always exists. When Problem (MP) admits a solution, these formulations coincide with тOT (La, Hu,, i.c. for is deterministie, which means that (x. y) ~ т, then y = Tx).. Problem (KP) admits a dual formulation. Denoting by Po(u, v) = ((fg))lf.g Ω R. f ① g S c, μ @ v-a.e.} where f @ g (x,)) + f(x) + g(y) is the tensor sum, it reads
( f(x) du(x)+ + / gly)dv()) dy(y).
($.9 E arg min (1.))
(DKP)
Extension to the Unbalanced Setting. The Kantorovich formulation imposes mass conservation, so it cannot handle arbitrary positive measures v, (M))... Unbalanced optimal transport (UOT) (Benamou, 2003; Chizat et al., 201 lifts this constraint by penalizing instead the deviation of pı tor to μ and patir
via é-divergence Dø. These divergences are defined trhough entropy function ø : (0, +o0) + too], which is convex, positive, lower-semi-contimuuss and s.t. F(1) Denoting dx lim d(x)/x its recession constant, a, B M(22), we have
Do(a] n. / Φ
do
da + d'co ( da-
(1)
dB
da 3 + at, with 33 the relative density of where we write the Lebesgue decomposition a de a w.r.t. 3. this work, we consider strictly convex and differentiable entropy functions o, with d'e = +oo. This includes, for instance, the KL, the Senen-Shhaoon or the x² divergence. Introducing A2 > 0 controlling how much mass variations are penalized as opposed to transportation, the unbalanced Kantorovich problem (UKP) then seeks measure GM (X' )):
/ Jaxa y) + A (ri + A2D (p2tTlu). (UKP)
min UWolu, u) (ExM)
An UOT coupling TUOT always exists. In practice, instead of directly selecting As we introduce Ti = A, Xr+5: where & the entropy regularization parameter as described Appendix B. Then, we
recover balanced OT for 7; 1, equivalent to X; -> +00, and increase unbalancedness by decreasing Ts. We write (TI, 72) accordingly. Finally, the (UKP) also admits dual formulation which reads
(f.g) E arg arg min (1.))44.(p.))
- (-)((((())
/
I - (o((()))
(UDKP)
UNBALANCED NEURAL MONGE MAPS
Unbalanced Monge Maps. Although the Kantorovich formulation has an unbalanced extension, deriving an analogous Monge formulation that relaxes the mass conservation constraint is more challenging. Indeed, by adopting the same strategy of replacing the marginal constraint Tép with penalty D& (T' v), we would allow the distribution of transported mass Tip to shift away from v, but the total amount transported mass would still remain the same. Instead, we can follow a different approach and consider a two-step procedure where (i) we re-weight μ and v into measures Ā R= u μ and v- = v y having the same total mass using u, V Rd + R+ and (ii) then compute a (balanced) Monge map between á and v. Hence, it remains to define the re-weighted measures ÿú, i. Therefore, we seek to achieve two goals: [i] create or destroy mass to minimize the cost of transporting measure á to D, while [ii] being able control the deviation of the latter from the input measures μ and v. To that end, we can show that computing Tuor between j and v implicitly amounts to following this e-baalnniing procedure. Additionally, j and can be characterized precisely.
3

---

Published as a conference paper at ICLR 2024
target
source
transported source
(a) "Unbalanced" data
(b) Discrete OT (r = 1.0) (c) Disc. UOT (r = 0.99) (d) Disc. UOT (r = 0.9)
(e) Flow Matching
(f) OT-FM (r = 1.0)
(g) UOT-FM (r = 0.99) (h) UOT-EM (r 0.9)
Figure 2: Different maps data drawn from mixture of uniform distribution, where the density in the bottom left and the top right ( is higher than in the top left and bottom right (3) (Appendix F.6). Besides the data Figure 2a, the first row shows results of discrete balanced OT (2b), and discrete unbalanced OT with two different degrees of unbalancedness T (2c, 2d). The second row shows the maps obtained FM with independent coupling (2e), OT-FM (2f), and UOT-FM (2g. 2h).
Proposition 3.1 (Re-balancing the UOT problem). Let TUOT be the solution of problem (UKP) between p. E M
M- (S2). for T1, > O. Then, the following holds:
1. FUOT solves the balanced problem (KP) between its marginal jı pi TUOT and i = PLITUOT. which re-weighted versions of μı an that have the total mass. Indeed, === 6( f+) f*) and = --) where ò = ( and f",g" S2 - are solution of (UDKP).
2. If we additionally assume that μ < La and that the cost c(x,y) = h(x y) with h strictly convex, then TUOT is unique and deterministic: TUOT = (Ia, T", where T = la- Vii is the Monge map between i and for cost
Remark 3.2. All the cost functions c(x, y) = -1x->>lllll with p > can be dealt with via point 2 of Prop. 3.1. For p = 2, we recover an unbalanced counterpart of Brenier ( (1987) Theorem stating that TUOT is supported on the graph of Vip" where x* Rd + convex.
Prop. 3.1 highlights that the optimal re-weighting procedure relies setting PitTUOT and = POTTUTT where we directly control the deviation to the input measures μ and v with and 72. This enables define unbalanced Monge maps formally. We illustrate this concept in Figure 1.
Definition 3.3 (Unbalanced Monge maps). Provided that it exists, we define the unbalanced Monge map between v E M (Q)) as the Monge map T between the marginal jı and is of the TUOT between μ and From point of Prop. 3.1, it satisfies TUOT (Id, T )t...
Proofs are provided in Appendix A. We stress that even when μ and v have the same mass, the unbalanced Monge map does not coincide with the classical Monge map as defined in (MP), Indeed, a and ū + y/ unless we impose Ti 72 1 to recover each marginal constraint.
3.1 ESTIMATION OF UNBALANCED MONGE MAPS
Learning an unbalanced Monge map between μi, v remains to learn a balanced Monge map between μ4 E = и. Provided that map estimator. practice, we can produce approximate samples from the re-weighted measures using samples from the input measure X1 ...,, μ and .1...Y. Yn (Fatras al., 2021b). First, [i] we compute #UOT x the solution of problem (UKP) between A., = + - 6x, and D, E 6y,, then [ii] we sample V1). (X,, ~ #UoT. Indeed, if we set a = Avor and b = An. and define jin := Σn -i=1 a, PIUTT and V, : -n b:dy, PutRUoT. these are empirical approximation of ç and V, and X1 ... X,, ~ and y1... Yn ~ D. Such couplings fuor be estimated efficiently using entropic regularization (Cuturi, 2013; Séjourné et al., 2022) (Appendix B.l). Moreover, we can learn the re- re-weighting functions as we have access to estimates of their pointwise evaluation: a(x;) n. (4; (dîn/ (x;))xx) and v(y.)
sample from and V. this can be done with any Monge
by
4

---

Published as a conference paper at ICLR 2024
(di,/dūn)(y:) dén))y;))))))))) We provide our general training procedure in Appendix D and note that UOT introduces the hyperparameter т, which controls the level of unbalancedness as visualized in Figure 2. The main limitation of our proposed framework revolves around the choice of T. which as discussed in Séjourné et al. (2023) is not always evident and thereby facilitates a grid search (Appendix B.2).
Merits of Unbalancedness. In addition to removing outliers and addressing class imbalances across the entire dataset, as mentioned in Section 1, the batch-wise training approach employed by many neural Monge map estimators makes unbalancedness a particularly favorable characteristic. Specifically, the distribution within a batch is likely not reflective of the complete source and target distributions, leading to a class imbalance at the batch level. Similarly, data point considered normal in the overall distribution may be treated as an outlier within a batch. Consequently, the use of unbalanced Monge maps serves to mitigate the risk of suboptimal pairings within individual batches, which can help stabilize training and speed up convergence (Fatras et al., 2021a; a;; Choi et al., 2023). Next, we introduce the three estimators that we employ to compute unbalanced Monge maps.
||x y12.. Brenier (1987)'s Theorem states that T Vyp", where 5p* Ω + IB is convex and solves a reformulation of the dual problem (DKP)
OT-ICNN. When y)
80* E arg inf / +p convex Jn p(x) du((x) + / (y) duly),
where yp" denotes the convex conjugate of yp Makkuva et al. (2020) propose to solve Eq. 2 using Input Convex Neural Networks (ICNNs) (Amos et al., 2017), which are parameterized convex functions x (x). To avoid the explicit computation of the convex conjugate 1g. they approximate it by parameterizing an additional ICNN no and solve
sup inf LOTHICN == [
Jaxo (ne (Vip(()) - (x, Vipa - (o(y)) d dv(y),
(3)
no ICNN Ye ICNN
where estimate of the Monge map is recovered through TV Vig. We refer to this estimation procedure as OT-ICNN and use it to learn unbalanced Monge maps (UOT-ICNN) in Section 5.1.
Monge Gap. ICNN-based methods leverage Brenier (1987)'s Theorem, that can only be applied when c( y) ) = x II.. Recently, Uscidda & Cuturi (2023) proposed a method to learn Monge maps for any differentiable cost. They define regularizer M called the Monge gap, that quantifies the lack of Monge optimality of a map T': : Rd + Rd. More precisely,
(T) = J. c(x,y) du(x) Wo(u, (tp).
From the definition, M (T) 2 0 with equality i.f.f. T is Monge map between a and Thy for the cost c. Then, the Monge gap can be used with any divergence A to define the loss
LOTMa(O) = ATooiu, v) MT
(5)
which is 0 i.f.f. Te is a Monge map between and for cost c. We refer to the estimation procedure minimizing this loss as OT-MG and use it to learn unbalanced Monge maps (UOT-MG) in Section 5.2.
Flow Matching. Continuous Normalizing Flows (CNFs) (Chen et al., 2018) are a family of continuous time deep generative models that construct a probability path between μ and v using the flow (ét)teju,1) induced by neural velocity field (Gr, (cr.a)egjo.;; [e Flow Matching (FM) (Lipman et al., 2023; Liu et al., 2022; Albergo Vanden- Sanden-Eijnden, 2023) is a simulation-free technique to train CNFs by constructing probability paths between individual data samples x μ., X1 V, and minimizing
(FM (xx. Il
[1, xx) (x: xo)
(6)
where Tind f 2 P. When this loss is zero, Lipman et al. (2023, Theorem 1) states that di is push-forward map between and 1/, namely orEP = P. While here FM yields individual straight paths between source and target pairs, the map Φ1 not a Monge map. OT-FM (Pooladian al., 2023a; Tong et al., 2023b) suggests to use an OT coupling Ror instead of Tind in the FM loss, which in practice is approximated batch-wise. It can then be shown that 1/1 approximates a Monge map asymptotically (Pooladian et al., 2023a, Theorem 4.2). Thanks to Proposition 3.1, we could extend OT-FM to the unbalanced Monge map setting (UOT-FM) by following Algorithm and rescaling the measures batch-wir before applying OT-FM. In this special case however, both the unbalanced rescaling and the coupling computation can be done in one step by leveraging the unbalanced coupling TUOT, which again is approximated batch-wise. The full algorithm is described in Appendix D)
5

---

Published as a conference paper at ICLR 2024
4 RELATED WORK
While the majority of literature in the realm of neural OT considers the balanced setting (Makkuva al., 2020; Korotin al., 2021; Uscidda & Cuturi, 2023; Tong et al., 2023b; Korotin al., 2022), only a few approaches have been proposed to loosen the strict marginal constraints. Gazdieva et al. (2023) consider partial OT, less general problem setting than unbalanced OT. Yang & Uhler (2019) propose a GAN-based approach for unbalanced OT plans (as opposed to deterministic maps) and very recently, Choi et al. (2023) employ a methodology based on the semi-dual formulation for generative modeling of images. Finally, Lübbeck et al. (2022) introduced way to incorporate unbalancedness in ICNN-ÖT. In the following, we highlight the main differences to our approach.
First, our proposed method can be applied to any neural Monge Map estimator, while the method proposed Lübeck et al. (2022) requires the explicit parameterization of the inverse Monge map. While this requirement is satisfied in ICNN-OT, more flexible and performant neural Monge map estimators like MG-OT (Uscidda & Cuturi, 2023) and OT-FM (Pooladian et al., 2023a; Tong et al., 2023b) do not model the inverse Monge map. Second, Lübeck et al. (2022) suggest rescaling the distributions based on the solution of a discrete unbalanced OT coupling between the push-forward of the source distribution and the target distribution, while we consider unbalancedness between the source and the target distribution. This entails that their rescaled marginals are not optimal (Proposition 3.1). Third, our approach is computationally more efficient, as our method requires the computation of at most one discrete unbalanced OT plan, while the algorithm proposed in Lübeck al. (2022) includes the computation of two discrete unbalanced OT plans. Related work for each of the specific domain translation tasks is discussed in Appendix E due to space restrictions.
5 EXPERIMENTS
We demonstrate the importance of learning unbalanced Monge maps on three different domain translation tasks leveraging three different (balanced) Monge map estimators, which showcases the flexibility of our proposed method. To start with, we demonstrate the necessity of incorporating unbalancedness when performing trajectory inference on single-cell RNAseq data, which we use ICNN-OT for: Subsequently, we improve the predictive performance of modeling cellular responses to perturbations with an unbalanced Monge map building upon MG-OT. Both of these tasks can be formulated as unpaired single-cell to single-cell translation task. Finally, show in Section 5.3 that OT-FM performs competitively on unpaired image-to-image translation tasks, while its unbal- anced formulation (UOT-FM) further improves performance and helps to preserve characteristic input features. We utilize cetropy-regularized UOT with E = 0.01 (Appendix B. and peform a small grid search over т. Note that while we train estimators to learn an unbalanced Monge map. i.c. Monge maps between the rescaled distribution jú and v, we evaluate them on balanced distributions i and v.
5.1 SINGLE--CELI TRAJECTORY INFERENCE
OT has been successfully applied to model cell trajectories in the discrete set- ting (Schiebinger et al., 2019). The need for scalability motivates the use of neu- ral Monge maps. Hence, we leverage OT-ICNN to model the evolution of cells on single-cell RNA-seq dataset com- prising measurements of the developing mouse pancreas at embryonic days 14.5 and 15.5 (Appendix 1.5. To evaluate how accurately OT-ICNN recovers bio- logical ground truth, we rely on the es- tablished single-cell trajectory inference tool CellRank (Lange et al., 2022).
OT-ICNN
(b) UOT-ICNN
Figure 3: Velocity stream embedding plots (Ap-
pendix F.5.6). The orange box highlights the direction of Ngn3 EP cells. With OT-ICNN these move to the "right", which contradicts biological ground truth. This due to the distribution shift shown in Appendix F.5.1 and demon- strates the need to incorporate unbalancedness.
The developing pancreas can be divided into two major cell lineages, the endocrine branch (EB), and the non-endocrine branch (NEB) (Appendix F.5.4). It is known that once a cell has committed to either of these lineages, it won't develop into cell belonging to the other lineage (Bastidas-Poncee

---

Published as a conference paper at ICLR 2024
et al, 2019). The Ngn3 EP cluster is population of early development in the EB, and hence cells belonging to this population are similar in their gene expression profile to cells in the NEB branch. Thus, when measuring the performance of OT-ICNN, we consider this cell population isolatedly. We evaluate OT-ICNN and UOT-ICNN by considering the predicted evolution of a cell.
We report the mean percentage of correct cell transi- tions for each lineage in Table 1, with full results in Ap- pendix C.5. While UOT-ICNN consistently outperforms its balanced counterpart, the effect of the unbalanced neu- ral Monge map estimator is particularly significant in the Ngn3 EP population. This improvement is visually con- firmed Figure To demonstrate the biological rele- vance of our proposed method, we compare the results attained with OT-ICNN and UOT-ICNN with established trajectory inference methods in single-cell biology (Appendix C.5).
Table 1: Evaluation of unbalancedness in OT-ICNN based correct cell type transitions.
Correct transitions
Model
EB
Ngn3 EP
NEB
OT-ICNN
54%
7%
67%
UOT-ICNN
57%
69 %
85 %
5.2 MODELING PERTURBATION RESPONSES
In-silico prediction of cellular responses drug perturbations is a promising approach to accelerate the development of drugs. Neural Monge maps have shown promising results in predicting the reaction of cells different drugs (Bunne et al., 2021). Since Uscidda & Cuturi (2023) show that OT-MG outperforms OT-ICNN in this task, we con- tinue this line of research by applying un- balanced OT-MG (UOT-MG) to predict cel- lular responses to 35 drugs from data pro- filed with 4i technology (Gut et al., 2018). We leverage distribution-level metrics to measure the performance of map T mod- cling the effect of a drug. More precisely, given metric A, we keep for each drug a batch of unseen control cells Meesa and unseen treated cells Viest and compute (T Mees))
Sinkhom Divergence: Sy Fpt,
Drug Signature: (s[DS) (2%p,
Unbalaxed: UOT-G
Balancod: OT-MG
Balanced: OT-MG
Figure 4: Fitting of a transport map T' to predict the responses cell populations to cancer treatments on 4i (upper plot), using balanced (OT-MG) and unbalanced Monge maps (UOT-MG) fitted with the Monge gap. A point below the diagonal indicates that unbalancedeness improves performance.
We measure the predictive performances using the Sinkhorn divergence (Feydy et al., 2019) and the 1.2 drug signature (Bunne et al., 2021, Sec. 5. Figure shows that adding unbalancedness through UOT-MG improves the performances for almost all 35 drugs, w.r.t. to both metrics. Fach scatter plot displays points 2, = (24,,) where y/ the performance obtained with UOT-MG and x, that of OT-MG, a given treatment and for a given metric. A point below the diagonal y refers to an experiment in which using an unbalanced Monge map improves performance. We assign a color to each treatment and plot five runs, along with their mean (the brighter point). For given color, higher variability of points along the x axis means that UOT-MG is more stable than OT-MG, and vice versa.
5.3 UNPAIRED IMAGE TRANSLATION
Metrics. To evaluate the generative performance in the image translation settings, we employ the Fréchet inception distance (FID) (Heusel al., 2018). In image translation, it is generally desirable to preserve certain attributes across the learned mapping. How- ever, FID computed on the whole dataset does not take this into account. Hence, we also compute FID attribute-wise to evaluate how well relevant input features are preserved.
Moreover, we consider the transport cost (T-Cost) Table 2: FID EMNIST digits - letters (xx) induced by the learned flow 1/1. Full details are described in Appendix F.2.2.
Method
0 -
- I
8 + B
Average
EMNIST. To illustrate why unbalancedness be crucial in unpaired image translation, we compare the performance of FM, OT-FM, and UOT-FM the
FM
30.1
21.2
122.8
58.0
OT
9.2
18.6
83.1
36.9
UOT-FM
12.5
13.9
42.6
23.0
7

---

Published as a conference paper at ICLR 2024
Male Female
Female
Malc
Input
FM
OT-FM
LOI-FM
Add Glasses
Remove Glasses
Input
FM
OT-FM
UOT-FM
Figure 5: CelebA 256x256 translated test samples with FM, OT-FM, and UOT-FM.
EMNIST (Cohen et al., 2017) dataset translating the digits {0, 1,8) to letters {0. We select this subset to obtain a desired class-wise mapping {0- 1-,0,,11,8--11 As seen in Appendix F.2.1 samples of class B and I are underrepresented in the target distribution. Because of this class imbalance, not all digits can be mapped to the correct corresponding class of letters. Figure sketches how unbalanced Monge maps can alleviate this restriction. Indeed, Table 5 confirms this behavior numerically, while Figure 6 visually corroborates the results. The difference is especially striking for mapping class 8 B, where the largest distribution shift between source and target occurs. While this class imbalance is particularly noticeable in the considered task, these distribution-leve discrepancies can be present in any image translation task, making unbalancedness crucial to obtain meaningful mappings with OT-FM.
Digits Letters
CelebA. To demonstrate the applicability of our method on more complex tasks in computer vi- sion, we benchmark on four translation tasks in the CelebA dataset (Liu et al., 2015), namely Male - Female, Female > M Male, Remove Glasses, and Add Glasses. Here, the distribution shift occurs be- tween image attributes. For example, when translat- ing from Male to Female, samples of males with hats (10.1%) significantly outnumber female ones (3.1%). Roughly speaking, this implies that when satisfying the mass conservation constraint, over half of the males with hats will be mapped to females without one. To compare to established methods, we use a common setup in high-dimensional image translation (Torbunov et al., 2022; Nizan & Tal, 2020; Zhao et al., 2021) where images are cropped and reshaped to 256x256. Due to computational reasons, we employ latent FM (Dao al., 2023) leveraging a pretrained Stable Diffusion (Rom- bach et al., 2022) variational auto-encoder ( (VAE), which compresses images from H x W C' to x 4. We run FM, OT-FM, and UOT-FM in latent space, based on which we also compute the minibatch OT couplings. We compare UOT-FM against FM, OT-FM, CycleGAN (Zhu et al., 2017), and UVCGAN (Torbunov et al., 2022), which achieves current state-of-the-a results. While UOT-FM only requires training one network and the selection of one hyperparameter, established methods require more complex model and network choices as described in Appendix E.3.
Input
8
8 8 O
0
0
FM
1
O
0
O
0
0
OT-FM
O
B
D O
0
0
/
O
UOT-FM
B B B O 0
0
/
I
Figure 6: Samples from EMNIST translating
digits + letters with different FM methods.
H. W. g 8
We report generative performance as well as transport cost results in Table 3. UOT-FM outperforms CycleGAN and OT-FM across all tasks with respect to the FID score while lowering the learned transport cost. Table 4 reports the attribute- wise FID between different FM methods. As expected,
8

---

Published as a conference paper at ICLR 2024
Table 3: Comparison results on CelebA 256x256 measured by FID and transport cost (T-Cost) over different tasks. Results marked with * are taken from Torbunov et al. (2022).
Method
FID
T-Cost
FID
T-Cost
FID
T-Cost
FID T-Cost
Male
Female
Female + Male
Remove Glasses
Add Glasses
UVCGAN
9.6'
13.9
14.4
13.6"
CycleGAN
15.2
22.2*
24.2*
19.8
FM
17.5
89.2
18.6
91.2
25.6
74.0
30.5
73.9
OT-FM
14.6
59.3
14.8
59.7
20.9
47.1
20.1
45.7
UOT-FM
13.9
56.8
14.3
56.4
18.5
44.3
18.3
42.0
Table 4: CelebA 256x256 attribute-wise FID on different tasks for different FM methods.
Male
Female
Female
Male
Method
Glasses
Hat
Gray hair
Average
Glasses
Hat
Gray hair
Average
FM
66.6
117.0
71.7
85.1
40.2
86.2
60.5
62.3
OT-FM
63.1
111.5
64.6
79.7
38.8
74.1
54.0
55.6
UOT-FM
62.6
108.5
62.7
77.9
38.2
72.4
51.1
53.9
Remove Glasses
Add Glasses
Male
Female
Hat
Gray hair
Average
Male
Female
Hat
Gray hair
Average
FM
46.9
47.7
110.6
89.1
74.3
19.6
43.2
46.6
39.7
37.3
OT-FM
42.3
40.3
99.5
83.1
66.3
12.0
36.9
44.4
27.4
30.2
UOT-FM
40.5
37.3
89.6
77.0
61.1
11.4
35.0
39.7
25.5
27.9
unbalancedness improves performance substantially. To confirm the superior performance of UOT-FM qualitatively, Figure 5 visualizes translated samples for FM, OT-FM, and UOT-FM. To demonstrate the strong performance of UOT holding in pixel space, follow Korotin et al. (2022) and downsample images to 64x64 to evaluate FM, OT-FM, and UOT-FM translating Male + Female. We benchmark UOT against two principled approaches for unpaired image translation, namely NOT (Korotin et al., 2022) and CycleGAN. Appendix C.1 confirms the superiority of UOT-FM over OT-FM. These convincing results put UOT-FM forward new principled approach for unpaired image translation. Lastly, UOT-FM also improves performance with a low-cost solver and fewer function evaluations, which allows for higher translation quality on compute budget (Appendix C.2).
CIFAR-16 10 image generation. The experiments above show that unbalanced neural Monge maps are mostly favorable over their balanced counterpart when translating between two data distributions. When translating from a source noise, i.e. parameterized distribution that is easy sample from, this task referred to as generutive modeling. Hence, we investigate whether UOT-FM also improves upon OT-FM in image generation task. We benchmark on the CIFAR-10 (Krizhevsky al.) dataset using the hyperparameters reported in Tong et al. (2023b) and show that also here UOT-FM improves upon OT-FM w.r.t. FID score. Moreover, we plot FID and transport cost convergence over training in Appendix C.7 and generated samples (C.8).
Table 5: OT-FM compared UOT- FM with different on CIFAR-IU 10 measure by FID and transport cost.
Method
FID
T-Cost
OT-FM
3.59
104.53
UOT-FM (7 0.99)
3.47
103.99
UOT-FM (7 0.97)
3.42
102.98
UOT-FM (7 0.95)
3.79
102.17
CONCLUSION
In this work, we present a new approach incorporate unbalancedness into any neural Monge map estimator. We provide theoretical results that allow us to learn an unbalanced Monge map by showing that this equivalent to learning the balanced Monge map between two rescaled distributions. This approach can be incorporated into any neural Monge map estimator. We demonstrate that unbal- ancedness enhances the performance in various unpaired domain translation tasks. In particular, improves trajectory inference time-resolved single-cell datasets and the prediction of perturbations on the cellular level. Moreover, we demonstrate that while OT-FM performs competitively natural image translation. unbalancedness (UOT-FM) further elevates these results. While use the squared Euclidean distance for all experiments, an interesting future direction lies in the exploration of more meaningful costs tailored specifically towards the geometry of the underlying data space.

---

Published as a conference paper at ICLR 2024
7 REPRODUCIBILITY
For our new findings, we supply complete proofs in Appendix A. We detail the proposed algorithms in Appendix D and provide all implementation details for reproducing the results of all three tasks reported in this work in Appendix F. Additionally, the code to reproduce our experiments can be found at https: //t::///iiii httpt/tttbtt.tt/t.///////ha.bhtt com Eann/xxplee 1a eetoooee//oo--. fm.
ACKNOWLEDGEMENT
The authors would like to thank Alexey Dosovitskiy, Sören Becker, Alessandro Palma, Alejandro Tejada, and Phillip Weiler for insightful and encouraging discussions as well as their valuable feedback. Luca Eyring and Dominik Klein thank the European Laboratory for Learning and Intelligent Systems (ELLIS) PhD program for support. Zeynep Akata acknowledges partial funding by the ERC (853489- DEXIM) and DFG (2064/1 Project number 390727645) under Germany's Excellence Strategy, Fabian Theis acknowledges partial founding by the ERC (101054957- DeepCell), consulting for Immunai Inc., Singularity Bio B.V., CytoReason Ltd, Cellarity, and an ownership interest in Dermagnostix GmbH and Cellarity. Views and opinions expressed are however those of the authors only and do not necessarily reflect those of the European Union or the European Research Council. Neither the European Union nor the granting authority can be held responsible for them.
REFERENCES
Michael S. Albergo and Eric Vanden-Eijnden Building normalizing flows with stochastic interpolants,
2023.
Brandon Amos. On amortizing convex conjugates for optimal transport. arXiv preprint arXiv:2210.12153. 2022.
Brandon Amos, Lei Xu, and J Zico Kolter. Input Convex Neural Networks. In International Conference on Machine Learning (ICML), volume 34, 2017.
Brandon Amos, Samuel Cohen, Giulia Luise, and levgen Redko. Meta optimal transport, 2023.
Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, John Quan, George Papanakarios, Roman Ring. Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Luyu Wang, Wojeiech Stokowiec, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http:/ httgtgtthuhss com/ deepmind.
Aimée Bastidas-Ponee Sophie Tritschler, Leander Dony, Katharina Scheibner, Marta Tarquis-Medina, Ciro Salinno, Silvia Schirge, Ingo Burtscher, Anika Böttcher, Fabian Theis, et al. Comprehensive single cell mrna profiling reveals a detailed roadmap for pancreatic endocrinogenesis. Development, (46(12)dev173384) 2019.
Jean-David Benamou. Numerical resolution of an "unbalanced" mass transport problem. ESAIM: Mathematical Modelling and Numerical Analysis, 37(05):851-866 2003.
Volker Bergen, Marius Lange, Stefan Peidli, F. Alexander Wolf, and Fabian J. Theis. Generalizing velocity to transient cell states through dynamical modeling. bioRxiv, 2019. doi: 10.1 101//2098. URL hwwph:tpps... /wwww//wwwww/..c/ biorxi org/ content weattt/eerrry//aa /t//////99///// /0/////29996....////////
Valentin De Bortoli, James Thornton, Jeremy Heng, and Amnad Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. 2023.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPP programs, 2018. URL
http:://ithie com/ com/googll /eooll///o jax.
10

---

Published as a conference paper at ICLR 2024
Yann Brenier. Décomposition polaire et réarrangement monotone des champs de vecteurs. CR Acad. Sci. Paris Sér. Math., 305, 1987.
Charlotte Bunne, Stefan G Stark, Gabriele Gut, Jacobo Sarabia del Castillo, Kjong- Van Lehmann, Lu- cas Pelkmans, Andreas Krause, and Gunnar Ratsch. Learning Single-Cell Perturbation Responses using Neural Optimal Transport. bioRxiv, 2021.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and Prançois-Xavier Vialard. Unbalanced optimal transport: geometry and Kantorovich formulation. Journal of Functional Analysis, 274(11): 3090-3123, 20
Lenaic Chizat, Gabriel Peyré. Bernhard Schmitzer, and François- Franveis-Xavier Xavier Vialard. Scaling algorithms for unbalanced optimal transport problems. Mathematics of Computation, 87(314):2563-22 -609,, 2018b.
Jaemoo Choi, Jaewoong Choi, and Myungjoo Kang. Generative modeling through the semi-dual formulation of unbalanced optimal transport, 2023.
Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune Gwon, and Sungroh Yoon. Ilvr: Conditioning method for denoising diffusion probabilistic models, 2021.
Yunjey Choi, Minje Choi, Munyoung Kim, Jung- Woo Ha. Sunghun Kim, and Jaegul Choo. Stargan: Unified generative adversarial networks for multi-domain mmage-to-image translation, 2018.
Gregory Cohen, Saeed Afshar, Jonathan Tapson, and André van Schaik. Emnist: an extension of
mnist to handwritten letters, 2017.
Marco Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. Advances in Neural Information Processing Systems (NeurIPS), volume 26, 2013.
Marco Cuturi, Lactitia PaangPapaxanthos, Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein. arXiv Preprint 22X1::2201.,2324,, 2022.
Quan Dao, Hao Phung, Binh Nguyen, and Anh Tran. Flow matching in latent space, 2023.
Prafulla Dhariwal and Alex Nichol. Diffusion models beat gans on image synthesis, 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x 16 words: Transformers for image recognition at scale, 2021.
Kilian Fatras, Thibault Séjourné, Nicolas Courty, and Rémi Flamary. Unbalanced minibatch optimal transport; applications to domain adaptation, 2021a.
Kilian Fatras, Younes Zine, Rémi Flamary, Rémi Gribonval, and Nicolas Courty. Learning with minibatch wasserstein: asymptotic and gradient properties, 2021b.
Jean Feydy, Thibault Séjourné, François-Xavier Xavier Vialard, Shun-Ichi Amari, Alain Trouvé, and Gabriel
Peyré. Interpolating between Optimal Transport and MMD using Sinkhorn Divergences. In International Conference Artificial Intelligence and Statistics (AISTATS), volume 22, 2019.
Milena Gazdieva, Alexander Korotin, and Evgeny Burnaev. Partial neural optimal transport. arXiv preprint S1VV:2333300988 2023.
Gabriele Gut, Markus Herrmann, and Lucas Pelkmans. Multiplexed protein maps link subcellular organization to cellular state. Science (New York, N. Y.). 361, 08 2018. doi: 10.1126/scieeee 224.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge local nash equilibrium, 2018.

---

Published as a conference paper at ICLR 2024
Yuge Ji, Mohammad Lotfollahi, F Alexander Wolf, and Fabian J Theis. Machine learning for perturbational single-cell omics. Cell Systems, 12(6))522-533) 507,,37,,22 2021.
Kantorovich. On the transfer of masses (in russian). In Doklady Akademii Nauk, volume 37, pp. 227, 1942.
Patrick Kidger. On Neural Differential Equations. PhD thesis, University of Oxford, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In International Conference on Learning Representations (ICLR), 2014.
Dominik Klein, Giovanni Palla, Marius Lange, Michal Klein, Zoe Piran, Manuel Gander, Laetitia Mang--apaxxanths, Michael Stem, Aimee Bastidas-Ponce Marta Tarquis-Medine al. Mapping cells through time and space with moscot. bioRxiv, pp. 2023-05, 2023.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations, 2020.
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? continuous wasserstein-2 benchmark. 2021. doi: 10448550/RRXIVV210660125
URL https::tp/// www.a..xvvvo....... abs /2106 01 954.
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. 2022. doi: 10.48550/ARXIV.220..1
URL https: //tp//ixvvv...///.. wwt./v.vorrrvooooooooo abss/a////// 22001 22...12220...2..........2.2.21.1
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 (canadian institute for advanced research). URL http:
heett://www/ /www. cs www.ooooooooooooooooooo cii ftr.
Marius Lange, Volker Bergen, Michal Klein, Manu Setty, Bernhard Reuter, Mostafa Bakhti, Heiko Lickert, Meshal Ansari, Janine Schniering, Herbert B Schiller, Dana Pe'er er, and Fabian Theis. CellRank for directed single-cell fate mapping. Nat. Methodds 19(2):199-11 -100 February 2022.
John Lee, Nicholas P. Bertrand, and Christopher J. Rozell. Parallel unbalanced optimal transport regularization for large scale imaging problems, 2020.
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal entropy-transport problems and a hellinger-kantorvvich distance between positive measures. Inventiones Mathematicae, 211 (3):969-11177 2018.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. 2023.
Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos Theodorou, Weili Nie, and Anima Anandkumar. Psb: Image-to-image schrödinger bridge, 2023.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow, 2022.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild, 2015.
Frederike Lübeck, Charlotte Bunne, Gabriele Gut, Jacobo Sarabia del Castillo, Lucas Pelkmans, and David Alvarez-Melis. Neural unbalanced optimal transport via cycle-consistent semi-couplings. arXiv preprini arXiv:2209. /6211, 2022.
Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. International Conference on Machine Learning (ICML), volume 37, 2020.
Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon. Sdedit: Guided image synthesis and editing with stochastic differential equations, 2022.
Petr Mokrov, Alexander Korotin, and Evgeny Burnaev. Energy-guides entropic neural optimal transport. arXiv preprint 20t4:2300006094 2023.
12

---

Published as a conference paper at ICLR 2024
Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Histoire de l'Académie Royale des Sciences, pp. 666-704, 1781.
Ori Nizan and Ayellet Tal. Breaking the cycle colleagues are all you need, 2020.
Taesung Park, Alexei A. Efros, Richard Zhang, and Jun-Yan Zhu. Contrastive learning for unpaired image-to-image translation, 2020.
Gabriel Peyré and Marco Cuturi. Computational Optimal Transport. Foundations and Trends in Machine Learning, 11(5-)), 2019. ISSN 1935-8245.
Henri De Plaen, Pierre-Françond De Plaen, Johan A. K. Suykens, Mare Proesmans, Tinne Tuytelaars, and Luc Van Gool. Unbalanced optimal transport: unified framework for object detection, 2023.
Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky Chen. Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint wwXivv2334.147770 2023a.
Aram-Alexandre Pooladian, Carles Domingo-Enrich, Ricky T. Q. Chen, and Brandon Amos. Neural optimal transport with lagrangian costs. In JCML Workshop on New Frontiers Learning, Control, and Dynamical Systems, 2023b. URL https: /opp///eeeeveeeeeeeeeeeeeeeeeeee net forum? dd== myb0FKB8C9.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and BjAmm Ommer. High- resolution image synthesis with latent diffusson models. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition (CVPR), 2022. URL https: httpg/gttttt/oo///o//ggeee com/ CompVis/ http//aatetttaafaiiaaaaaa
aarxiv..vvv. org/ abs 2tt//2222..111
Filippo Santambrogio. Optimal Transport for Applied Mathematicians. Birkhüusee, NY, 55(58-63):
94, 2015.
Hiroshi Sasaki, Chris G. Willcocks, and Toby P. Breckon. Unit-ddpm: Unpaired image translation with denoising diffusion probabilistic models, 2021.
Geoffrey Schicbinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Arych Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Tpaimsa-Transport Analysis of Single-Cel Gene Expression Identifies Developmental Trajectories in Reprogramming. Cell, 176(4), 2019.
Thibault Séjourné, Gabriel Peyré, and François-Xavier Xavier Vialard. Unbalanced optimal transport, from theory to numerics. arXiv preprint 0tt::2211.087755 2022.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. baore-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13453 2020.
Sanjay R Srivatsan, José L McFaline- Figueroa, Vijay Ramani, Lauren Saunders, Junyue Cao, Jonathan Packer, Hannah Pliner, Dana I. Jackson, Riza M Daza, Lena Christiansen, et al. Massively multiplex chemical transcriptomics at single-cell resolution. Science, 367(6473), 2020.
Thibault Séjourné, Gabriel Peyré, and François-Xavier Vialard. Unbalanced optimal transport, from theory to numerics, 2023.
Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lueas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, and Alexey Dosovitskiy. Mlp-mixer: all-mlp architecture for vision, 2021.
Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajectory A Dynamic Optimal Transport Network for Modeling Cellular Dynamies. In International Conference on Machine Learning (ICML), 2020.
Alexander Tong, Nikolay Malkin. Kilian Fatras, Lazar Atanackovie, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation-free schr\ odinger bridges via score and flow matching. arXiv preprint arXiv: 036672 2023a.
13

---

Published as a conference paper at ICLR 2024
Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport, 2023b.
Dmitrii Torbunov, Yi Huang, Haiwang Yu, Jin Huang, Shinjae Yoo, Meifeng Lin, Brett Viren, and Yihui Ren. Uvegan: Unet vision transformer cycle-consistent gan for unpaired image-to-image translation, 2022.
Théo Uscidda and Mareo Cuturi. The monge gap: A regularizer to learn all transport maps, 2023.
Karren D. Yang and Caroline Uhler. Scalable unbalanced optimal transport using generative adver- sarial networks, 2019.
Min Zhao, Fan Bao, Chongxuan Li, and Jun Zhu. Egsde: Unpaired image-to-image translation via energy-guided stochastic differential equations, 2022.
Yihao Zhao, Ruihai Wu, and Hao Dong. Unpaired Longe-to-image translation using adversarial consistency loss, 2021.
Jun-Yan Yan Zhu, Taesung Park, Phillip Isola, and Alexei Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. In Computer Vision (ICCV), 2017 IEEE International Conference on, 2017.
14

---

Published as a conference paper at ICLR 2024
A PROOFS
In this section, we prove Propositon 3.1, which we repeat here.
Proposition 3.1 (Re-balancing the UOT problem). Let aUoT be the solution of problem (UKP)
between μ. // E M T2, > 0. Then, the following holds:
1. AUOT solves the balanced problem (KP) between its marginal which are re-weighted versions of μ. an y that have the same total mass. Indeed, jı = - and - V, where & = Y and f*.g" - are solution of (UDKP).
p PITOT and i = P2tTUOT.
2. If we additionally assume that & La and that the cost c(x.y) = h(x y) with h strictly convex, then TUOT is unique and deterministic: TUOT = (Ld, T )iZ,, where T la Vh o is the Monge map between it and for cost
Proof. Proof of point 1. We remember that
/
NUOT E arg arg min xEM(X.)) Axn
+ (xxy))dt((xx))+DAAd + A2 (p2tmlu))
so since we assume that 6' +00, the terms (p())t)()) and Dø(p2tTlu) are finite i.f.f. pi Prr has density w.r.t j and p2tz has density w.r.t. v. Therefore, it exists u, Rd + R s.t. j = pi UTUOT = u y and jí = P2tTTUT = v. v. Moreover, and have the same total mass, since by applying Fubini's Theorem twice, one has
J divoo(x.y)= 1 djī(x) = f(att) dō(y) Jn
Moreover, from (Liero et al., 2018, Corollary 4. 4.16), exists (f*,g*) E Lo(4, v) solution of the unbalanced Kantorovich dual (UDKP) between μ and s.t. @ g Tuor-ae. This implies that (f*.g" solves the balanced Kantorovich dual problem (DKP) between fi and D. Indeed, since tUoT has and & marginals, this equality yields
J fx(x)) dī(x)) + [ gy) dū(y) = c(x, y) dTuoT(x,y)
(9)
On the other hand, for any (f.g)= E le 27., one has
J f(x)
/ o(0) dū(y)= / f @ g(x. dl
/ s
c(x, e)x dxuor(x,
(10)
which provides
( [(x) ((y) S J d(x,y) divoo(xxy)) (11)
sup (1.))t.((PP)))
Additionally, since (f*,g") V)(,,v), one has (f*,g") E L,( D) because we have shown that j < f and < v. Therefore, they optimal dual potentials. Then, since is continuous and Ω is compact, strong duality holds (Santambrogio, 2015, Theorem 1.46) and thus
( [(x) + [ dBl)) inf / =Ell((.e) Jaxa c(x,y)
sup (1.))
(12)
-f Jaxn c(x, y) dzuor(x,y).
which yields the optimality of TUOT in the balanced problem (KP) between j and i.
To conclude the proof of the first point, we show that = da = Ø(-f) and u de = ~ 6(-9). dp
dv
(Liero et al., 2018, Corollary 4.16) also states that f - w-a.e. and g" -o ou v-a.e Since ó is strictly convex, o' is invertible and (6')-1 = ( = (Santambrogio, 2015, Box 1.12), so the result follows.
Proof of point 2. First, given that μ & Cas one has jı < La since ÿ < μ. Then, since c(x, y) = hixx y) where h is strictly convex, we apply (Santambrogio, 2015, Theorem 1.17) to state that
15

---

Published as a conference paper at ICLR 2024
the Monge map between Ji and & exists and is unique, and T L1- Vh f since we have shown that gf** are optimal balanced dual potential between f and i. Then, since c is continuous, 12 is compact and f is atomless, the Kantorovich and the Monge formulation coincide (Santambrogio, 2015, Theorem 1.33) so TUOT = (Ta, T )tfi and it is unique.
B COUPLING COMPUTATION
In this section, we lay out details w.r.t the mini-batch coupling computation we leverage in our framework. In B.1 we detail entropy regularized OT, which we utilize as it offers more efficient way of estimating couplings as opposed to computing the non-regularized one, which we lay out in B.3. Moreover, we discuss limitations that might arise given the coupling computation and in general with our proposed framework (B.2).
B. ENTROPIC REGULARIZATION
When the measures l and are instantiated samples, as usual in machine learning context, the Kantorovich problem (UKP) translates to a convex program whose objective function can be smoothed out using entropic regularization (Cuturi, 2013). For empirical measures A. 1 in ンド= by, and E > 0, we form the cost matrix C [e(x,, y, ), and consider:
n -i= 8x.
min (P. C) D. (PI, In) + X2 1 1m) (H)) P20
(13)
where H(P) - - - P.j (og(Pig) (1,)). With e + O, we recover (UKP). Using e > 0, the above Eq. (13) admits tractable dual representation that can be leveraged to derive fast computational pro- cedure which generalizes the Sinkhorn algorithm and is commonly used in computational OT (Chizat et al., 2018b; Peyré & Cuturi, 2019; Séjouréé et al., 2023). In practice, we follow Cuturi et al. (2022) and define Ti x.+e* This facilitates hyper-pase parameter selection: т₃ E (0. 1] and we recover the i-th hard marginal constraint with т, 1, when X; + +oo. In this work, we use entropic regularization with small enough regularization strength e to approximate UOT couplings. In all the experiments, we use 0.01 C, where C denotes the mean of the cost matrix C.
B.2 LIMITATIONS
We consider the choice of the additional hyperparameter (T1, 12) as the main challenge of our proposed framework. A way to make the choice of these hyperparameters independent of the scale of the data, and hence to some degree comparable across different tasks, is scale the cost matrix of the discrete OT problem by its mean, which we leverage in all experiments as mentioned in Appendix B.1.
A second limitation is the implicit removal of outliers by UOT. This property is largely seen as a strength of UOT, e.g. Choi et al. (2023) show UOT's robustness to outliers. However, what considered an outlier with respect to UOT is again dependent on the hyperparameter r. As one lowers т, the data pairs with lower distance attain more relative mass, and as T - O sampling from the coupling approaches almost solely sampling the lowest distance pair. Depending there might be data points in the dataset that are "far" away from the remaining data points but are considered to be relevant the model training. When choosing a lower т, these points might get removed by UOT. Hence, the choice of r can be very important, but as discussed in Séjourné et al. (2023) the optimal choice is not evident and thereby usually facilitates least small grid search. Lastly, our framework also adds computational complexity, although negligible in most cases, as discussed in B.3.
the application,
B.3 COMPUTATIONAL COMPLEXITY
We compute all mini-batch UOT couplings using entropic regularization, which allows the of generalization of Sinkhorn's 's algorithm (Chizat et al., 2018b), having the same (nn2) time complexity, where n depicts the batch size. The added computational overhead of our framework compared to
16

---

Published as a conference paper at ICLR 2024
using the balanced alternative depends on the estimator used. We distinguish between two types of estimators in the following.
Estimators that don't tt leverage OT couplings (e.g. OT-MG OT-ICNNs). These estimators require samples of the marginal distributions between which the Monge map is to be calculated. Thus, each of the training iterations for these models involves first calculating unbalanced coupling #UoT using samples xI...., Xn l and yI..... .nn ~ V, then evaluate their respective training losses LOT-MG or LOTICNN samples x1y..., Xn fnn pi fuVOT and yiy.... in PIINUOT. Thus, for both these estimators, the additional computational cost lies in calculating #UoT-
Estimators that already leverage OT couplings (like OT-FM), Each iteration of OT-FM consists of sampling and calculating a coupling for from samples X...... X, ~ and y1...., v. then calculating the Flow Matching loss on samples of for. For UOT-FM, we directly replace the calculation of a balanced plan for by an unbalanced plan øor from samples x1,.... Xn and y1.....Ye yn v, then calculate the Flow Matching loss on for samples the same way. Since #oT and #UOT can be calculated with the same (n²)) runtime complexity, using Sinkhorn's algorithm (Cuturi, 2013) and its unbalanced generalization (Chizat et al., 2018b), UOT-FM has exactly the same runtime complexity as OT-FM. Therefore, since FM operates at coupling level, incorporating unbalancedness via UOT-FM does not add any additional computational cost.
In the first case, the relative computational overhead depends the cost of gradient step. In Pooladian et al. (2023a); Tong et al. (2023b) it was shown that even coupling computations of cost (nn²) are negligible compared to the gradient step of the Flow Matching loss. In this work, which involves computations scaling quadratically O(n") rather than cubically, this observation applies even more strongly. Hence, in most settings, our framework will significantly impact training time.
C ADDITIONAL EMPIRICAL RESULTS
In this section, we report additional results: In C.I results pixel space for 64x64 CelebA translating Male Female. Then for 256x256 CelebA translating Male - Female evaluate the performance of UOT-FM compared to OT-FM with few function evaluations and a low-cost Euler solver (C.2) and in C.3 report results for different levels of unbalancedness in UOT-FM. Moreover, we compare the performance of OT-ICNN and UOT-ICNN to established methods in single-cell trajectory inference (C.5). Furthermore, we benchmark UOT-ICNN against competing unbalancedness methods proposed in Lübeck et al. (2022) and Yang & Uhler (2019) on sciPlex perturbation dataset (Srivatsan al., 2020) in C.6. Lastly, plot FID and transport cost over training (C.7) for the generative modeling experiment on CIFAR-10 10 as well as some randomly generated samples in C.8.
C.I CELEBA IN PPXCE-SPACE
Table CelebA 64x64 Male Female. Results denoted with are taken from Korotin al. (2022).
attribute-wise FID
Method
FID
T-Cost
Glasses
Hat
Gray hair
Average
NOT
13.23*
CycleGAN
17.74*
OT-FM
11.52
40.37
47.40
85.80
49.53 46.73
60.91
UOT-FM
11.09
37.38
47.13
84.29
59.38
17

---

Published as a conference paper at ICLR 2024
C.2 CELEBA WITH LOW-COST SOLVER
Table CelebA 256x256 Male -> Female with low-cost Euler solver and varying number of
function evaluations.
FID
T-Cost
NFE
OT-FM
UOT-FM
OT-FM
UOT-FM
Adaptive
14.58
13.94
59.31
56.78
40
14.53
14.02
57.00
54.84
20
14.91
14.66
54,44
53,40
12
16.40
16.08
53.69
51.82
8
19.56
19.32
52.46
50.49
24.14
24.08
51.77
46.66
4
38.02
37.38
51.25
49.02
88.47
87.85
55.08
53.22
C.3 CELEBA OVER DIFFERENT LEVELS UNBALANCHDNESS
Table 8: 256x256 CclebA male > female results for UOT-FM with different = T1 = T2.
FID
T-Cost
FID Average
OT-FM
14.58
59.31
79.7
T = 0.99
14.75
58.94
80.3
0.98
14.23
57.63
79.8
T -= 0.95
13.94
56.78
77.9
0.90
14.30
55.24
78.8
C.4 EFFECT OF BATCH-SIZA AND EPSILON ON EMNIST
Here, we empirically evaluate the effect of changing the batch size and in UOT-FM applied the EMNIST dataset translating digits - letters. We report the learned Transport Cost and attribute- wise average FID score as detailed in Appendix F.2.2. The results reported in Section 5.3 are obtained with batch size of 256 and 0.01. Figure 7a shows the effect of changing the batch size. Our results align with previous studies (Pooladian et al., 2023a; Tong et al., 2023b), which show that already fairly low batch size yields competitive results with OT mini-batch couplings, and increasing over 128 does not change results significantly. Only when we decrease the batch size below 32, do observe significant drop in performance. For all batch size results, keep the number of samples seen during training the same set to 256 500k.
Secondly, Figure 7b confirms empirically how the performance of UOT-FM approaches FM increase E. Additionally, we find that lowering further than 0.01 does not improve the learned transport cost and average FID significantly while increasing the runtime of the coupling computation.
18

---

Published as a conference paper at ICLR 2024
Average
Avenge
Traqpunt Coe
50-
Transport Cost
40
20
20
10
16
64
256
512
10
10a
10-3
10
10
Boach
(a) Effect different batch sizes
(b) Effect different
Figure 7: Results EMNIST with UOT-FM translating digits + letters.
C.5 SINGLE-CELL TRAJECTORY INFERENCE COMPARISON TO ESTABLISHED METHODS
We compare OT-ICNN and UOT-ICNN against established methods in single-cell trajectory in- ference. We benchmark against sele (Bergen et al., 2019), TrajectoryNet (Tong et al., 2020), and Waddington OT (WOT) (Schiebinger et al., 2019). Table summarizes results, while the full transition probabilities on the EB and NEB branches while performing second best on the Ngn3 EP cells. Additionally, thirteen cell transitions.
reported in Table 10. UOT-ICNN outperforms all competing methods
UOT-ICNN improves keeps equal performance upon OT-ICNN in Implementation details of each competing method are described in Appendix F.5.5.
Table 9: Evaluation of different trajectory inference methods based on correct cell type transitions.
Correct transitions
Model
EB
Ngn3 EP
NEB
TrajectoryNet
0.33
0.01
0.71
scVelo
0.44
0.99
0.39
WOT
0.45
0.50
0.72
OT-ICNN
0.53
0.07
0.67
UOT-ICNN
0.57
0.69
0.85
Table 10: Cell type transition probabilities between cell types A and B such that cell type A maps exclusively cell type B. For each column, we underline the best, second best, and third best methods.
Model
FA
A
FB
B
FD
D
FE
E
NE
NL
DU
T
AC
+
+
4
->
A
B
D
E
ED
ED
DU
AC
AC
TrajectoryNet
0.07
0.46
0.07
0.39
0.11
0.75
0.10
0.52
0.00
0.01
0.32
0.79
0.99
SeVe
0.79
0.80
0.30
0.61
0.03
0.52
0.04
0.56
0.98
1.00
0.32
0.04
0.90
WOT
0.37
0.62
0.18
0.44
0.19
0.74
0.55
0.49
0.50
0.50
0.82
0.48
0.84
OT-ICNN
0.51
0.54
0.11
0.70
0.39
0.94
0.48
0.62
0.07
0.08
0.01
1.00
1.00
UOT-ICNN
0.25
0.60
0.29
0.75
0.78
0.99
0.36
0.55
0.65
0.73
0.59
1.00
1.00
19

---

Published as a conference paper at ICLR 2024
C.6 BENCHMARK AGAINST COMPETING UNBALANCEDNESS ESTIMATORS
To assess the performance our proposed approach UOT-ICNN, we benchmark it against competing methods proposed in Lübeck al. (2022) and Yang & Uhler (2019). While proposed approach for unbalanced neural Monge maps is significantly different from Lübeck et al. (2022), both approaches build upon the solver suggested in Makkuva al. (2020). We the exact same architecture and training procedure for both approaches and only perform the resampling in a different way. In detail, UOT-ICNN solves one discrete unbalanced OT problem between batch of the source distribution and batch of the target distribution. In contrast, Lübeck et al. (2022) compute an unbalanced discrete OT solution between the push-forward of the source distribution and the target distribution to estimate the left rescaling factor. Moreover, they compute second discrete OT solution between the pull-back of the target batch and the source batch to obtain an estimate of the right rescaling factor. We also compare Yang & Uhler (2019) but would like to highlight that we can only approximate a similar setup, e.g. by also training for 1000 iterations.
We evaluate the models on the sciPlex pertubation data (Srivatsan et al., 2020). We chose drugs that were reported to have a strong signal. The data was downloaded from https:/ / github. com/ bunnech/ cellot/ wwweeeeaiii We compute 30-dimensional PCA embedding and evaluate the performance with the Sinkhorn divergence (Feydy et al., 2019) (with entropy regularization parameter 0..1) (Feydy et al., 2019), analogously to the experimental setup for comparing OT-MG with UOT-MG in Section 5.2. Then we report the mean Sinkhorn divergence across all drugs for different levels of unbalancedness in Table 11. UOT-ICNN outperforms competing methods across all т. In particular, Lübeck et al. (2022) achieves close performance to UOT-ICNN with 0.95 but seems diverge for lower values of r.
Table 11: Sinkhorn divergence between predicted target and target distribution on sciPlex data per model. We highlight best and second best performance.
Sinkhorn Divergence
Method
T 0.95
== 0.9
T 0.85
== 0.8
T 0.75
= 0.7
UOT-ICNN
25.95
28.05
29.60
30.68
31.42
32.15
Lübecc et al. (2022)
26.39
1286.48
1841.31
3758.46
12913.06
27031.94
Yang & Uhler (2019)
37.40
37.79
37.84
38.03
38.09
37.83
C.7 CIFAR-10 CONVERGENCE OVER TRAINING
Here, we plot FID score and transport cost over training. As expected introducing more unbalanced- by lowering reduces the learned transport cost. Generative performance increases with the introduction of unbalancedness offering faster convergence up until a certain level.
5.00
FT LOOT-EM 0.99)
4.50
LOT--PM 0.95)
UOT-FM 0.97)
103
4.25
002
4.0
CT-FM
UOT-PM (T- 0.99)
UOT--M (T- 0.97)
3.50
(M 0.95)
99
100k
150k
2014:
50k
100k
200K
(a) score over training iterations.
(b) Learned transport cost over training iterations.
Figure 8: Convergence of training CIFAR-10 plotted for OT-FM and UOT-FM with different
20

---

Published as a conference paper at ICLR 2024
C.8 CIFAR- 10 GENERATED SAMPLES
Figure 9: Randomly generated images with UOT-FM (T 0.97) trained on CIFAR-10.
D ALGORITHMS
In this section, we describe the full algorithms for [i] estimating unbalanced Monge maps with any Monge map estimator in Algorithm [ii] computing unbalanced optimal transport Flow Matching in Algorithm 2.
Algorithm 1: Neural Unbalanced Monge maps
Require: Source and target distribution fu, v, unbalancedness parameters T = (Ti, T2). loss function of a balanced Monge map estimator Lor, solver Solver,, batch size n, number of iterations Piter, boolean flag learn rescaling whether to learn the re-weightings functions, neural re-weighting functions uo, 17g
1: for (= 1.... Tliter do
2: Sample batches Xi ... X, p and y ... Yn
3: Compute coupling tuor t Solver, ( (x 1 (y,)%m)
4: Sample new batches ([x;) {5, In (==1 RUor
5: Compute loss LUOT (0) < Lor (0; | (8,)2=)
6:
if learn rescaling then
7:
#uor and b x- * uorr
8:
(o) < (uor(0) + 1 n
(uo n a.)² + 1, Ex=((vo(3)) - n b,)
9: Update 9 with (eLvor(
10:
return 0
Algorithm 2: Unbalanced Optimal Transport Flow Matching (UOT-FM)
Require: Source and target distributions v, unbalancedness parameters 71, 72, parameterized vector field ni.
1: while Training do
2: Sample batches 3h ~ μ4, 32 ~ M, 1,,,~~U00 U(0,1)
3: Compute coupling tuor UOT Ti, T1: T2)
4:
Sample new batches (Øo. 21) ~ fuor
5:
Compute flow á;, (1 ()2a (.
6: Compute loss (F) = [1o 2x) ue (i. (1)
Update ! with (OCFM(
8: return Ug
21

---

Published as a conference paper at ICLR 2024
E RELATED WORK FOR TRANSLATION TASKS
Here, we describe related work specific to the three different domain translation tasks.
E. RELATED WORK IN SINGLE-CELL TRAJECTORY INFERENCE
Optimal transport has been established trajectory inference method in single-cell genomics, spearheaded by Schiebinger et al. (2019). The high computational burden of discrete OT has led to the development of low-rank solvers, which have been studied in the context of single-cell trajectory inference methods in Klein et al. (2023). Proofs of concept have also been conducted by Tong et al. (2020; 2023b;a). For developmental single-cell data, OT-based methods have been established as competitive methods in the field (Lange et al., 2022).
E.2 RELATED WORK IN PERTURBATIONAL. PREDICTIONS
Studying the effect of different perturbations a single-cell level is recent field in single-cell biology enabled by the progress of machine learning (Ji et al., 2021). 1). First models were based on VAEs (Ji al., 2021), while recently, the community has focused on OT-based methods (Bunne et al., 2021; Lübeck et al., 2022; Uscidda & Cuturi, 2023).
E.3 RELATED WORK IN UNPAIRED IMAGE TRANSLATION
One of the seminal works in unpaired image translation was CycleGAN (Zhu et al., 2017), which introduced a coneissonst loss to encourage translations that can reconstruct the original image when run through since built upon CycleGAN's foundational approach, seeking to improve the realism and variability of generated images. Notable examples include STARGAN (Choi et al., 2018), which extended CycleGAN to handle multiple attributes simultancously, as well as CUT (Park et al., 2020), and most recently UVCGAN (Torbunov et al., 2022), which employs self-supervised pre-training and an optimized architecture, a UNet Vision Transformer (Dosovitskiy et al., 2021). ) Note, that UOT-FM trains one network with one loss function and one hyperparameter 7 while in contrast, CycleGAN- based approaches train at least four networks with four loss functions and hyperparameters.
hewaadd-bacwwatt cycle between domains. Several state-of- of-the--tttttttt methods have
Diffusion models (DMs) (Song et al., 2020) have achieved ground-breaking results in image gen- eration. However, DMs were designed with only Gaussian noise source distribution in mind and thus are not directly applicable to image translation. This sparked the development of the extension of DMs to the unpaired image translation task. UNIT-DDPM (Sasaki et al., 2021) extends the CycleGAN concept to DMs by leveraging two DMs between domains with a cycle-consistency loss. ILVR (Choi al., 2021) and SDEdit (Meng et al., 2022) rely test source image during inference without leveraging information from the training source distribution. EDGSE (Zhao et al., 2022) suggests alleviating this constraint by pre-training two energy functions on the respective source and target domain. Very recently, a different approach was proposed based on the framework of Diffusion Schrödinger Bridges (Bortoli et al., 2023) applied in Liu al. (2023).
In contrast to DMs, OT-FM offers a more natural framework for unpaired image translation out of the box. It has been applied CelebA encoded in 128-dimensional VAE embedding (Tong et al., 2023b), but without comparison to any established methods and without reporting common metrics in image translation like FID. However, they show that OT-FM improves performance upon FM with the independent coupling, which we confirm with our results also holds in higher dimensional embeddings (Section 5.3) and pixel-space (Appendix C. Rectified flow (Liu et al., 2022) can be interpreted version of FM that approximates OT-FM. It has been applied to high-quality image translation and reports results visually, but again without comparison to established methods and no numerical results.
F EXPERIMENTAL AND EVALUATION DETAILS
Our code is based on Jax (Bradbury et al., 2018) while utilizing parts of the DeepMind Jax ecosys- tem (Babuschkin et al., 2020).
22

---

Published as a conference paper at ICLR 2024
F.1 UNBALANCED COUPLING ALGORITHM
To compute the entropic unbalanced coupling we leverage the ott-oatttto jax library (Cuturi et al., 2022) with the Sinkhorn algorithm (Cuturi, 2013). We choose the entropic coupling with small due its computational benefits.
Table 12: Hyperparameters for unpaired image translation on CelebA
CelebA-256 gender
CelebA-256 glasses
CelebA-64
Channels
128
128
192
ResNet blocks
4
3
Channels multiple
2, 2, 2
2, 2,
1, 2, 3, 4 4
Heads
1
Heads channels
64
64
64
Attention resolution
16
32, 16, 8
Dropout
0.1
0.1
0.1
GPU batch size
64
32
Effective Batch size
256
256
256
Iterations
400k
100k
400k
Learning rate
lc-4
lc-4
lc-4
Scheduler
constant
constant
constant
EMA-decay
0.9999
0.9999
0.9999
F.2 UNPAIRED IMAGE TRANSLATION
For each of the CelebA image translation tasks, we use a very similar hyperparameter setup based upon the UNet architecture used in Dhariwal & Nichol (2021), where all FM models were trained with the exact same setup. We report these in Table 12. For EMNIST we leverage the MLPMixer architecture (Tolstikhin et al., 2021) with hyperparameters detailed Table 13. Additionally, use the Adam optimizer (Kingma & Ba, 2014) with B1 = 0.9, $2 = 0.999, e = le 8 and no weight decay for all image translation experiments.
During inference, we solve for Pt at t 1 using the adaptive step-size solver Tsit5 with atol=rtol=le-5 implemented in the diffrax library (Kidger, 2021). For T we employ small grid search on the Male - Female task, where results are reported in Appendix C.3. 0.95 achieves the best performance and subsequently, we choose T 0.95 for all other tasks, including in pixel-space. For the EMNIST experiment, we choose = 0.9, 1.0.
F.2.1 DATASETS
For all image translation experiments, we rescale images from [0, 255] to [-1.0, 1.0). In practice, has been shown that data augmentation and regularization, e.g. random rotations or added noise, can be beneficial to achieve better generalization. We note that in our case, this entails learning a Monge map between adjusted distributions z and instead of μ and v. In pratice, we regularize source and
Table 13: Hyperparameters for the MLPMixer used for the EMNIST experiments.
Channels
Patch size
4
Token Mixer channels
512
Channel Mixer channels
512
Depth
Batch size
256
Iterations
500k
Learning rate
3e-4
Scheduler
constant
EMA-decay
0.9999
23

---

Published as a conference paper at ICLR 2024
Iaah Ldan
1330
39000
8
0
1
Figure 10: EMNIST data distribution.
target distribution with small amount of Gaussian noise such that we learn the Monge map between I = μ * (O, o² I) and i = * (0, o² where denotes the convolution operator. the FM framework, this is equivalent considering Gaussian probability paths px(x((x, N N(xxxx x2, o² I), which is implemented by adding a small amount of noise to the computation of æ; such that 2c ()ar læ ae, where (),
EMNIST. We leverage a subset of EMNIST, where we take digits {0, 1,8) as source and letters I, B} target distribution. The grayscale images of shape 28x28. The result- ing translation task contains a large distribution shift as visualized in Figure 10.
CelebA. CelebA contains images of size 218x178. Following Torbunov al. (2022); Nizan & Tal (2020); Zhao et al. (2021) we do not use the validation dataset for training, but instead add it to the test dataset. Then, the gender swap task contains about 68k males and 95k females while the glasses task entails 1kk samples with and 152k without glasses. For CelebA-256 we again follow a similar setup Torbunov al. (2022) and upsize images 313x256. Then instead of random cropping we center crop all images to 256x256. This is done to pre-compute all embeddings for the Stable Diffusion VAE. Note, that this is a disadvantage for the benchmarked FM methods as they cannot utilize the benefit of random cropping during training like Torbunov et al. (2022), For Celeba-64 we resize all images to 64x64.
F.2.2 METRIC COMPUTATION DETAILS
For general FID computation, we follow Torbunov et al. (2022); Korotin al. (2022) and compare translated test samples from the source distribution against the statistic from the test target distribution. For the attribute-wise computation, the number of test samples can be very low depending on the task and attribute. Thus, we compute the attribute-wise FID) comparing translated test samples to the statistics of the whole dataset w.r.t. the given attribute. For all FID computations, we use the jax-fi fii library. The transport cost || (xo) xoll2 is always computed and averaged over the hold-out test set. Note that we compute transport cost pixel space, scaled back to [0, 255]. The transport cost between the rescaled measures (f D) is always less than equal to the non-rescaled one We v). With this metric, we aim to measure whether this also generalizes to unseen samples from the balanced source distribution 26 /a when training UOT-FM.
F.3 IMAGE GENERATION
For the CIFAR- 10 experiments, we follow almost the exact setup reported in Tong et al. (2023b) except that we compute mini-batch couplings with entropy regularized OT as detailed in Appendix B. We report results at iteration 250k reproducing the OT-FM results from Tong et al. (2023b).
F.4 SINGLE-CELL PERTURBATION RESPONSES
For each drug, we model the neural map To using an MLP with hidden layer size [max(128, 2d), max 128, 2d), max(64, d), max(64, d)] where d is the dimension of the data. We train it for Riter 10, 000 iterations with a batch size n = 1, 022 and the Adam optimizer (Kingma & Ba, 2014) using learning rate y 0.001, along with a polynomial schedule of power p 1.5 to
24

---

Published as a conference paper at ICLR 2024
14.5 (carly time point)
15.5 (late time point)
Relative froquency %
EB
Ngn3 EP
NEB
Figure 11: Distribution over the selected lineages and time points. Cells of the non-endocrine branch much more abundant at the later time point due to very high proliferation rates of Acinar and Ductal cells.
decrease it to 10-5 Finally, we select the unbalancedness parameters using a grid search, imposing 72 and selecting among three values T; E (0.99, 0.95, 0.9).
F.5 SINGLE-CELL TRAJECTORY INFERENCE
F.5. PANCREATIC ENDOCRINOGENESH, DATA
The pancreatic endocrinogenesis data includes samples of embryonic days 14.5 and 15.5 (Bastidas- Ponce et al., 2019). After standard preprocessing 16, 206 genes remained. The OT-ICNN-basee algorithms and Waddington OT were distribution shift between the two time points. Note the increase in the NEB branch, which causes Ngn3 EP cells to be mapped to NEB without accounting for unbalancedness in OT-ICNN.
the 50 principal components. Figure 11 visualizes the
F.5.2 MODEL ARCHITECTURE
An Input Convex Neural Network (ICNN) parameterizes function f such that is convex with respect to its input by imposing certain constraints (Amos et al., 2017). Following Makkuva et al. we train two ICNNs, denoted by and g. with the following architecture (Makkuva al., 2020):
K dense layers consisting of weights Ao,..., Ax applied to the raw input 2,
K - positive dense layers consisting of non-negative weights W,,... Wx applied to intermediate outputs - as defined below.
Then, layer k is defined as
==<((WWW)))) + (Aga + bz))
(14)
where & is convex non-decreasing activation function, by the bias term and A₃ the weight matrix. In the last layer, we apply no activation function. Additionally, we use quadratic first layer:
zu (Ø(Agz + bu))
(15)
We enforce the non-negatyvity constraint of the weights W by weight clipping, while we only a penalization term for negative weights of
R(W9)
Σ
max(( mmax(0.-w)113
(16)
wEW's
where W denotes the set of weight matrices in the positive dense layers of the ICNN parameterizing g- We train with the following hyperparameters:
25

---

Published as a conference paper at ICLR 2024
learning rate: 0.001
optimizer: Adam| 31 0.5, 92 0.9)
hidden layers: [64, 64, 64, 64]
inner loop iterations: 10
outer loop iterations: 25000
batch size: 1024
activation function: Leaky ReLU(B = 0.01)
gradient clipping to norm: 1.0
Additionally, employ grid search over different levels of unbalancedness Results reported in Figure 12. Introducing unbalancedness gradually improves performance up to a certain level. The results reported in Section 5.1 use 0.85.
0.8
0.6
Percentage
0.4
EB
0.2
NEB
Ngn3 EP
0.95
0.90
0.85
Tau
Figure 12: Effect of unbalancedness parameter on aggregated transition probabilities. The trend shows that introducing unbalancedness to a certain extent by lowering T improves performance for all lineages. Metrics are computed the same way as in Table 1.
F.5.3 PRETRAINING
We pretrain the ICNN parameterizing / on the identity map suggested in Korotin et al. (2020); Amos et al. (2023) such that ff(z) = and then copy the weights to for them to be mutually and Vg(Vf(()) z. Therefore, we train on j (N(0,))) 3) for 15, 000 inverse (z(V iterations.
Additionally, for the balanced settings, perform best model selection based upon the lowest forward Sinkhorn divergence which is a debiased estimate of the Wasserstein distance. We also empirically observe better results using this stopping criterion for the unbalanced setting.
F.5.4 CELL TYPE TRANSITION METRIC
We follow Bastidas-Ponee al. (2019) to obtain the ground truth of cell type transitions. We only consider cell type transitions where the target cell state is a terminal cell state I. E T = (Acinar, Ductal, Alpha, Beta, Delta, Epsilon] types (Alpha, Beta, Delta, Epsilon). We assume the following cell type transitions are exclusively сorrect (denoted by +), there no descending cell type (or of cell types) other than the given one. We partition all considered cell type transitions into three categories.
union thereof. Let ED be the set of endocrine cell
The first set of considered transitions are endocrine branch (ED) transitions:
Fev+ Alpha (FA) Alpha (A)
Fev+ Beta (FB) Beta (B)
Fev+ Delta (FD) -> Delta (D)
Fev+ Epsilon (FE) -> Epsilon (E)
A A
26

---

Published as a conference paper at ICLR 2024
BB B
D
E
The second set of transitions are Ngn3 EP transitions:
Ngn3 high early (NE) - ED
Ngn3 high late (NL) - ED
The third set of transitions is the non-endocrine branch (NEB)
Ductal (DU) DU
Tip (T) - Acinar (AC)
AC + AC
In Table 10 we report the transition probabilities for all mentineet cell type transitions for OT-ICNN, UOT-ICN,, and established trajectory inference methods.
F.5.5 EVALUATION COMPETING METHODS
For evaluating cell type transitions we use CellRank kernels. Here, kernels quantify transition probabilities based vector fields. For all methods yielding velocity vectors (OT-ICNN, UOT-ICNN, Velo,, TrajectoryNet), we use the VelocityKernll. For evaluating Waddington OT (WOT), we use the WOTKernel.
Velo. To infer RNA velocity with Velo,, we first selected genes measured in at least 20 cells both unspliced and spliced transcripts. Next, cells are normalized by the median cell size, and the 2000 highly variable genes are selected. For these preprocessing steps, we used Velo's filter.and no normal the scvelo. pp .mmmentss function with the settings n_pcs=50 principal components, and n.neighbors=-3 nearest neighbors. RNA velocity was inferred using the recover.dynamie function implemented in seVelo.
function. Following, moments were calculated by
TrajectoryNet. We run trajectory net with default parameters as suggested by the author's Jupyter notebooks. Specifically, embeddi PCCCA max_dim=10 10, vecint=le-t We computed velocities by subtracting the inferred coordinates from the original coordinates in the embedding space. Since were able to retrieve the inferred coordinates only for one time point, we set the velocities of the other time point to
000 and
Waddington OT. The Waddington OT results For the corresponding transition matrix, we considered both inter and intra timepoint tran- sitions. The intra timepoint transitions by the cell-cell nearest neighbor graph, and assigned a weight of 0.2. Summarizing, WoTKernel's compute.transitic growth rate key=="roww rate conn.weight=0.2.
calculated with CellRank's WoTKernel.
quantified for each time point independently
method was run with growthiters= and
rate init",
F.5.6 VELOCITY STREAM EMBEDDING
Velocity vectors were projected onto the two-dimensional UMAP embedding using seVelo's velocity embeddin ng st ream function. To project the high-dimensional vectors, we consider the empirical displacement given by the difference of a cellular representation in the low-dimensional embedding. The displacement vector UMAP coordinates is then defined as the expected empirical displacement under a transition matrix and corrected by the expected shift under uniform distribu- tion. To define the entry (j. k) of the transition matrix, consider the reference cell and a neighbor k. The probability that cell transitions into cell & is defined as the normalized Pearson correlation between the empirical displacement in the high dimensional of the two cells, and the velocity vector of the reference cell.
27

---

Published as a conference paper at ICLR 2024
F.6 SIMULATED DATA
The simulated data consists of the union of draws of uniform distributions on 00[[-0...03 x
[-1.5, )0.5) and μl2 [4[[4.5, 5.5) x -..5).. Similarly, V1 sim2/([-0.5, 0.5] [0.5, 1.5] and = U1( 4.55 5.5] x [0.5, 1.5)). The source distribution is obtained by drawing 180 samples
from Yr and 120 samples from μ42. Similarly, the target distribution v is obtained by 180 samples from 1/2 and 120 samples from V1. To compute couplings, we use e = 0.1 in this case.
28


Reviews:
[{"review_id": "2y5BxV3gNg", "reviewer": "Reviewer_LDyN", "summary": "The paper introduces unbalancedness into any balanced Monge map estimator based on optimal transport (OT). Moreover, it demonstrates, both theoretically and experimentally, that this can enhance the performance of unpaired domain translation tasks.", "questions": "Please see the above weakness part.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper is well writing and is well organized, and proposed method is well-motivated and effective.\n2. This work innovatively integrates unbalanced optimal transport into the existing Monge map estimator, achieving competitive results in various domain translation generative tasks.\n3. Sufficient theoretical and experimental evidence is provided in the main paper and appendix.", "weaknesses": "1. In the single-cell trajectory inference and image translation experiments in Sections 5.1 and 5.3 (Table 1, 2, and 4), the article compared the performance of OT and unbalanced OT. Were there comparisons with other state-of-the-art methods?\n2. The experiments in Table 3 on the CelebA dataset show that the proposed method performs suboptimally FID score in comparison to UVCGAN [1], which was introduced in 2022. Further explanation is required to address this disparity.\n3. OT-FM [2] has conducted numerous experiments in image generation tasks and outperforms many GAN-based and diffusion-based methods. Have you considered experimenting with UOT-FM on the same tasks for further validation of the method's effectiveness?\n\n[1] Lipman Y, Chen R T Q, Ben-Hamu H, et al. Flow matching for generative modeling.\n[2] Torbunov D, Huang Y, Yu H, et al. Uvcgan: Unet vision transformer cycle-consistent gan for unpaired image-to-image translation.", "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": "mqkcakyZ5y", "reviewer": "Reviewer_qjsW", "summary": "The paper explores the role of introducing unbalancedness into neural Monge maps within the framework of optimal transport (OT). The authors investigate whether this unbalancedness can enhance the performance of unpaired domain translation tasks. Through a mix of theoretical analysis and empirical validation, the paper shows that incorporating unbalancedness leads to significant improvements in terms of cost-efficiency and generalization across domains.\n\nContributions:\n\nA modified OT framework that accommodates unbalancedness.\nNew neural estimators designed for unbalanced Monge maps.\nEmpirical evidence demonstrating the advantages of the modified approach in real-world applications like single-cell biology and computer vision.\nThe paper thus offers a novel approach to improving unpaired domain translation by tweaking the conventional OT framework and validating its effectiveness", "questions": "What are the limitations of introducing unbalancedness into Monge maps, and are there scenarios where this approach may not be beneficial?\n\nHow does the proposed unbalanced neural Monge map approach compare with existing methods in terms of efficiency and accuracy?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Novelty: The paper introduces the novel concept of \"unbalancedness\" into neural Monge maps, which is a fresh angle in the well-studied field of optimal transport.\n- Theoretical and Empirical Validation: The work combines both theoretical reasoning and empirical results, strengthening the validity of its claims.\n- Practical Impact: The paper demonstrates the utility of its approach in real-world applications, such as single-cell biology and computer vision, indicating its relevance beyond theoretical considerations.\n- Methodological Rigor: The research methodology appears to be sound, involving both the development of a new framework and neural estimators, as well as their validation on synthetic and real-world data.\n- Broad Applicability: The issue of unpaired domain translation is relevant in multiple fields, and the paper's contributions could be generalized to other domains, increasing its impact.", "weaknesses": "- Lack of External Benchmarks: Without a comparison to existing methods or frameworks, it's difficult to assess how much of an improvement the proposed approach offers.\n- Complexity: Introducing unbalancedness into neural Monge maps could add computational or conceptual complexity, which may not be fully addressed in the paper.\n- Dependency on Empirical Validation: While the paper does include empirical validation, its contributions could be further strengthened with more diverse datasets or a broader set of real-world applications.", "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": "1VzTFFqyWn", "reviewer": "Reviewer_Rczx", "summary": "This paper studies the methodological aspect of optimal transport (OT) under the practical application scenarios for domain translation, e.g., the mass proportion is imbalanced for different domains. The main pursuit is to approximately estimate the dual OT distance via the neural networks with significant capacity and provide an explicit retrieval form for the OT plan/map. To overcome the limitation in marginal distribution conservation, this paper introduces unbalanced OT to obtain a relaxed plan and then recovers the balanced OT from a reweighting view. Theoretical results are provided to connect the re-weighted OT and UOT and ensure the feasibility of recovering the Monge map from neural UOT. Experiments are conducted from several perspectives, which validates the effectiveness of the proposed method.", "questions": "1. As far as I understand this paper, the basic goal is to address the potential mismatch induced by the imbalanced weights of different domains; more precisely, the imbalance means the mass proportion of the ideal transport pairs is different. Such a scenario is also analogous to the label shift problem, which is usually considered in OT methodology. To address this problem, there are two natural and common solutions: relaxation (i.e., UOT) or reweighting (i.e., adjusting the marginal distribution) [r1, r2]. Since this work adopts relaxation as a solution (while also noting that the idea of reweighting is also implicitly shown in Sec. 3.1), I think more justification and discussion on the related fields are highly expected.\n\n2. Based on the reweighting solution mentioned above, can the imbalanced mass problem be addressed by detecting the degree of shifting mass (i.e., estimating the weights in Sec. 3.1) and solving the reweighting (balanced) OT? If it is feasible, what will be the advantages and weaknesses of the proposed methods?\n\n3. As the basic problem is similar to label shift, the experiments on related problems are highly appreciated. For example, the generalized label shift scenario [r1] and partial domain adaptation [r2] (which can be taken as an extreme scenario for label shift), where the comparison methods for corresponding tasks should also be carefully considered.\n\n4. Minor points:\n\n4.1) Page 3, typo in the reweighting formulation of $\\tilde{\\nu}$ in Sec. 3.1.\n\n4.2) Point 2 in Prop. 3.1, the definition of $\\mathcal{L}_d$ is not provided.\n\n[r1] Rakotomamonjy, Alain, et al. \"Optimal transport for conditional domain matching and label shift.\" Machine Learning (2022): 1-20.\n\n[r2] Gu, Xiang, et al. \"Adversarial reweighting for partial domain adaptation.\" Advances in Neural Information Processing Systems 34 (2021): 14860-14872.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ This paper is well-written and easy to follow; necessary justifications and discussions are provided around the technical parts.\n+ The proposed method is well-motivated and reasonable; the main difficulty stated in introduction is properly addressed with the proposed method and theoretical analysis.\n+ Diverse experiments are conducted to validate the proposed method from different perspectives and tasks, where the results are generally convincible.", "weaknesses": "- Some related concepts and fields are omitted, which should also be discussed appropriately.\n- The limitations of the proposed method should be discussed.\n- The experiments can be enhanced by considering more practical and challenging problems.", "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": "cz449P3cBZ", "reviewer": "Reviewer_1iBx", "summary": "This paper proposes a general framework for incorporating unbalancedness into any Monge map estimator. This paper shows that the unbalanced Monge map corresponds to the Monge map (from OT) between the reweighted source and target distributions. The proposed framework is to introduce existing Monge map estimator into these reweighted distributions. These reweighted distributions are estimated through the minibatch UOT estimator.", "questions": "-\t**Q1.** What is the definition of $\\tau$ in Prop 3.1? Is it the same $\\tau$ in Appendix B?\n-\t**Q2.** What is the reason for better performance of UOT compared to OT? This paper suggest that this improvement is primarily attributed to the discrepancy in the number of samples for each corresponding cluster, e.g. $8 \\leftrightarrow B$ in Fig 1. If the number of samples were similar, would the performance of UOT be comparable to OT?\n-\t**Q3.** The proposed framework is to optimize the OT Monge map on the estimated reweighted distribution. How does the performance of the proposed framework change according to the minibatch UOT coupling estimates, such as minibatch size and $\\epsilon$ of entropic regularization, in terms of quantitative results?\n- **Typo**  $\\phi^{\\*} (x) d \\mu (x) \\rightarrow \\phi^{\\*}(y) d \\nu (y)$ in Eq (2)\n\n**Reference**\n\n[1] Choi, Jaemoo, Jaewoong Choi, and Myungjoo Kang. \"Generative Modeling through the Semi-dual Formulation of Unbalanced Optimal Transport.\" arXiv preprint arXiv:2305.14777 (2023).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "-\tThis work is overall well-written.\n-\tThis paper proposes a general framework for extending existing OT Monge maps estimators into unbalanced cases.\n-\tThis paper shows that the introduction of unbalancedness improves the performance of existing OT Monge map estimators in unpaired domain translation tasks.", "weaknesses": "-\t**W1.** The main Prop 3.1-1, which provides the justification for the proposed general framework, was proved in [1] (The proof is different).\n-\t**W2.** The performance is only compared with the OT-counterpart and not with other UOT Monge maps estimators.", "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)