Neural Unbalanced Optimal Transport via Cycle-Consistent Semi-Couplings
Authors: Frederike Lübeck
Charlotte Bunne
Gabriele Gut
Jacobo Sarabia del Castillo
Lucas Pelkmans
David Alvarez-Melis
Venue: Submitted to ICLR 2023 (2023)

Abstract: Comparing unpaired samples of a distribution or population taken at different points in time is a fundamental task in many application domains where measuring populations is destructive and cannot be done repeatedly on the same sample, such as in single-cell biology. Optimal transport (OT) can solve this challenge by learning an optimal coupling of samples across distributions from unpaired data. However, the usual formulation of OT assumes conservation of mass, which is violated in unbalanced scenarios in which the population size changes (e.g., cell proliferation or death) between measurements. In this work, we introduce NubOT, a neural unbalanced OT formulation that relies on the formalism of semi-couplings to account for creation and destruction of mass. To estimate such semi-couplings and generalize out-of-sample, we derive an efficient parameterization based on neural optimal transport maps and propose a novel algorithmic scheme through a cycle-consistent training procedure. We apply our method to the challenging task of forecasting heterogeneous responses of multiple cancer cell lines to various drugs, where we observe that by accurately modeling cell proliferation and death, our method yields notable improvements over previous neural optimal transport methods.

Under review as a conference paper at ICLR 2023
NEURAL UNBALANCED OPTIMAL TRANSPORT VIA SYCLE-CONSISTENT SEMI-COUPLINGS
Anonymous authors Paper under double-blind review
ABSTRACT
Comparing umpaired samples of distribution or population taken at different points in time is a fundamental task in many application domains where measuring pop- ulations is destructive and cannot be done repeatedly the same sample, such in single-cell biology. Optimal transport (OT) can solve this challenge by learning an optimal coupling of samples across distributions from unpaired data. However, the usual formulation of OT assumes conservation of mass, which is violated in unbalanced scenarios in which the population size changes (e.g.. cell proliferation death) between measurements. In this work, we introduce NUBOT, neural un- balanced OT formulation that relies on the formalism of semi-couplings to account for creation and destruction of mass. To estimate such semi-couplings and general- ize samppe, we derive an efficient parameterization based on neural optimal transport maps and propose novel algorithmic scheme through cycle-consistent training procedure. We apply our method to the challenging task of forecasting heterogeneous responses of multiple cancer cell lines to various drugs, where we observe that by accurately modeling cell proliferation and death, our method yields notable improvements over previous neural optimal transport methods.
INTRODUCTION
Modeling change is at the core of various problems in the natural sciences, from dynamical processes driven by natural forces to population trends induced by interventions. In all these cases, the gold standard is to track particles or individuals across time, which allows for immediate estimation of individual (or aggregate) effects. But maintaining these pairwise correspondences across interventions or time is not always possible, for example, when the same sample cannot be measured more than once. This is typical in biomedical sciences, where the process of measuring is often altering or destructive. For example, single-cell biology profiling methods destroy the cells and thus cannot be used repeatedly on the same cell. In these situations, one must rely on comparing different replicas of population and, absent natural identification of elements across the populations, infer these correspondences from data in order to model evolution or intervention effects.
The problem of inferring correspondences across unpaired samples in biology has been traditionally tackled by relying on average and aggregate perturbation responses (Green & Pelkmans, 2016; Zhan al., 2019; Sheldon et al., 2007) by applying mechanistic or linear models (Yuan et al., 2021; Dixit et al., 2016) in, potentially, a learned latent space (Lotfollahi et al., 2019). Cellular responses to treatments are, however, highly complex and heterogeneous. To effectively predict the drug response of patient during treatment and capture such cellular heterogeneity, it is necessary to learn nonlinear maps describing such perturbation responses on the level of single cells. Assuming perturbations incrementally alter molecular profiles of cells, such as gene expression or signaling activities, recent approaches have utilized optimal transport to predict changes and alignments (Schiebinger et al., 2019; Bunne et al., 2022a; Tong et al., 2020). By returning coupling between control and perturbed cell states, which overall minimizes the cost of matching, optimal transport can solve that puzzle and reconstruct these incremental changes in cell states over time.
Despite the advantages mentioned above, the classic formulation of OT is ill-suited to model processes where the population changes in size, e.g., where elements might be created destroyed over time. This is the case, for example, in single-cell biology, where interventions of interest typically promote proliferation of certain cells and death of others. Such scenarios violate the assumption of conservation

---

Under review as a conference paper at ICLR 2023
a.
source #l
semicouplings
balanced Monge map
Y1
V
reweighting
71
B
reweighting
7/2
observed
subpopulation
cells
disappeared
Figure 1: A semi-coupling pair (7i, 72) consists of two couplings that together solve the unbal- anced OT problem. Intuitively, "/ describes where mass goes as it leaves from μ., and 722 where it comes from it arrives in v. b. NUBOT parameterizes the semi-couplings (71, 72) the composition of reweighting functions 17 and Ç and the dual potentials f and g between the then balanced problem.
of mass that the classic OT problem relies upon. Relaxing this assumption yields generalized formulation, known as the unbalanced OT (UBOT) problem, for which recent work has studied its properties (Liero et al., 2018; Chizat et al., 201 8a), numerical solution (Chapel et al., 2021), and has applied it successfully to problems in single cell biology (Yang & Uhler, 2019). Yet, these methods typically scale poorly with sample size, prone unstable solutions, or make limiting assumptions, e.g., only allowing for destruction but not creation of mass.
In this work, we address these shortcomings by introducing novel formulation of the unbalanced OT problem that relies on the formalism of semi-couplings introduced by Chizat et al. (2018b), while still obtaining explicit transport map that models the transformation between distributions. The advantage of the latter is that it allows mapping new out-of-sample points, and it provides an interpretable characterization of the underlying change in distribution. Since the unbalanced OT problem does not directly admit a Monge (i.c., mapping- based) formulation, we propose to learn to jointly 're-balance' the two distributions, thereby allowing us estimate map between their rescaled versions. To do so, we leverage prior work (Makkuva et al., 2020; Korotin et al., 2021) that learns the transport map as the gradient of convex dual potential (Brenier, 1987) parameterized via an input convex neural network (Amos et al., 2017). In addition, we derive simple update rule to learn the rescaling functions. Put together, these components yield a reversible, parameterized, and computationally feasible implementation of the semi-coupling unbalanced OT formulation (Fig. 1).
In short, the main contributions of this work are: (i) A novel formulation of the unbalanced optimal transport problem that weaves together the theoretical foundations of semi-couplings with the practical advantage of transport maps; (ii) A general, scalable, and efficient algorithmic implementation for this formulation based on dual potentials parameterized via convex neural network architectures; and (iii) An empirical validation on the challenging task of predicting perturbation responses of single cells multiple cancer drugs, where our method successfully predicts cell proliferation and death, in addition to faithfully modeling the perturbation responses on the level of single cells.
BACKGROUND
2.1
OPTIMAL TRANSPORT
For two probability measures v in PX)) with X = =dd and real-valued continuous cost function c E C(x²),, the optimal transport problem (Kantorovich, 1942) is defined as
/ c(x,v))(dz,du))
OT(H,)) == inf 7EL(a.))
(1)
where [(u, fr E M+ s.t. (Proj (10j))47=00 u4 (Proj2)r7 = n} is the set of couplings in the cone of nonnegative Radon measures M+ with respective marginals μu, v. When instantiated on finite discrete measures, such f Σi= u,ox, and Σド= Uy with u E En, Σm this problem translates to a linear program, which can be regularized using an entropy term (Peyré & Cuturi, 2019). For 2 0, set
(T, == min PeU((uv))
(P, - HP(),
(2)
2

---

Under review as a conference paper at ICLR 2023
where (P) == - Eij Ps (log P.j - 1) and the polytope (u, v) is the of matrices {P E IR+ 1R P1m = P, = v]. For clarity, we will sometimes write ( v, (x.). {y W. (1) when
Notice that the definition above reduces differentiable proxy to approximate OT and allows fast numerical approximation via the Sinkhorn algorithm (Cuturi, 2013), but introduces bias, since in general OT, n)) 0.
0. Setting > O yields faster and
Neural optimal transport. To parameterize (1) and allow to predict how a measure evolves from Ju v, introduce an alternative formulation known the Monge problem (1781) given by
OT(H, inf / c(x, (()))dp(x))
T:T,PEN
with pushforward operator and the optimal solution T known as the Monge map between jr and When cost is the quadratic Euclidean distance, i.c., c c= - 112. Brenier's 's theorem (1987) states that this Monge map is necessarily the gradient V of a convex potential
X' + E such that Velyu ie., (z) = (z). This connection has far-reaching impact and a central component recent neural optimal transport solvers (Makkuva et al., 2020; Bunne et al., 2022c; Alvarez-Melis et al., 2022; Korotin et al., 2020; Bunne et al., 2022b; Fan et al., 2021b). Instead of (indirectly) learning the Monge map 7 (Yang & Uhler, 2019; Fan et al., a021a), is sufficient to restrict the computational effort to learning a good convex potential Ve. parameterized via input convex neural networks (ICNN) (Amos et al., 2017), s.t. VoNitu v. Alternativly, parameterizations of such maps can be carried out via the dual formulation of (1) (Santambrogio, 2015, Proposition 1.11, Theorem 1.39), i.c.
/ fdw
f gdv.
OT(,,)) v) = sup
(4)
f.gEC(x))
f0g5c
where the dual potentials f. are continuous functions from X' to R, and f + f(z) + g(x).
Based on Brenier (1987), Makkuva et al. (2020) derive an approximate min-max optimization scheme parameterizing the duals J,g via two convex functions. The objective thereby reads
OT(,, v/) sup inf convex convex
E [
Exlf(=)) Ey Vo(v)) (( s(Vo(v)) (5)
V...(f.g)
Cy..
When parameterizing f and via pair of ICNNs with parameters 0y and 0g. this neural OT scheme then allows to predict v or (4 via g 1 or V Jo, respectively. We further discuss neural primal (Fan et al., 202 Yang & Uhler, 2019) and dual approaches (Makkuva et al., 2020; Korotin et al., 2020; Bunne et al., 2021) in $D.2.
2.2 UNBALANCED OPTIMAL TRANSPORT
A major constraint of problem (1) is its restriction to pair of probability distributions and v of equal mass. Unbalanced optimal transport (Benamou, 2003; Liero et al., 2018; Chizat et al., 2018b) lifts this requirement and allows comparison between unnormalized measures, i.c. via
inf f (x², (x*) c(x,,)) (dx, dy) n Dr ((Proji)s7) ()) + +2 ((Proj2))) 1);
(6)
with f-divergences Dj, and Dj induced by f1, fz, and parameters T2) controlling how much mass variations are penalized as opposed to transportation of the mass. When introducing an entropy regularization in (2), the unbalanced OT problem between discrete measures u and v, i.c.,
UBOT(u, v) := min (T, [c(#,, ,,)),)) 71 Dj (TI,, u) 72D,, (T 1m =H((T) (7) TER
can be efficiently solved via generalizations of the Sinkhorn algorithm (Chizat al., 201 Sa; Cuturi, 2013; Benamou al., 2015). We describe alternative formulations of the unbalanced OT problem in detail, review recent applications, and provide a broader literature review in the Appendix ($A.1).
A NEURAL UNBALANCED OPTIMAL TRANSPORT MODEL
The method propose weaves together rigorous formulation of the unbalanced optimal transport problem based on semi-couplings (introduced below) with a practical and scalable OT mapping estimation method based input convex neural network parameterization of the dual OT problem.
3

---

Under review as a conference paper at ICLR 2023
Semi-coupling formulation. Chizat et al. (2018b) introduced class of distances that generalize optimal transport for the unbalanced setting. They introduce equivalent dynamic and static formu- lations of the problem, the latter of which relies on semi-couplings to allow for variations of mass. formulation closely related (by a change of variables) to the notion of semi-couplings was inde- pendently introduced by Liero et al. (2016; 2018). These are generalizations of couplings whereby only one of the projections coincides with prescribed measure. Formally, the set of semi-couplings between measures μr and v defined as
F, (1, def. {170- 7) (M+ (23))² (Projn); 7o μ. (Proi2); n =
(8)
With this, the unbalanced Kantorovich OT problem can be written as Ce(H..) = inf (1om)))))))))) J c(2, 7. M/s 22 d),((z))) y). where y is any joint measure for which 70, 71 & 7.
Although this formulation lends itself to formal theoretical treatment, it has at least two limitations. First, it does not explicitly model mapping between measures. Indeed, no analogue of the celebrated Brenier's theorem is known for this setting. Second, deriving a computational implementation of this problem is challenging by the very nature of the semi-couplings: being undetermined along one marginal makes hard to model the space in (8).
Rebalancing with proxy measures. To turn the semi-coupling formulation of unbalanced OT into a computationally feasible method, we propose to conceptually break the problem into balanced and unbalanced subproblems, each tackling a different aspect of the difference between measures: feature transformation and mass rescaling. These in turn imply decomposition of the semi-couplings of (8), as we will show later. Specifically, we seek proxy measures jú and i with equal mass (i.e., (xX)= = )(x)) across which to solve a balanced OT problem through Monge/Brenier formulation. To decouple measure scaling from feature transformation, we propose to choose jı and i simply as rescaled versions of μ and v. Thus, formally, we seek Jı, E MM (X) and T, S : X - X' such that
i p- p. T 5.V. T- = D, S = jt,
(9)
where xx R+ are scalar fields, 17 denotes the measure with density n=(x) du(x) (analogously for (v)), and S are pair of forward/backwarrd optimal transport maps between jí and E (Fig. lb).
Devising an optimization scheme to find all relevant components in (9) is challenging. In particular, it involves solving an OT problem whose marginals not fixed, but that will change as the reweighting functionals 77, S are updated. We propose an alternating minimization approach, whereby alternative solve for v) (through an approximate scaling update) and T, S (through gradient updates on ICNN convex potentials, as described in Section 2.
Updating rescaling functions. Given current estimates of 1) and T, we consider the UBOT problem (6) between Ta(1) (4) = Ti and P. Although in general these two measures will not be balanced (hence why need use UBOT instead of usual OT), our goal is to eventually achieve this. To formalize this, let us use the shorthand notation YÜB (a, B) == argmin UBOT(y; a,, 8), where UBOT
is defined in (7). For a fixed 7', our goal is to find 17 such that (Projj):htt (Ta 14) v)] T: [i(n-pu p).
i... to rescale μ so that the unbalanced solution would in fact be 'balanced' along that marginal. For the discrete setting (finite samples), this corresponds to finding a vector e E R" satisfying:
I min = where argmin UBOT(e O u, T(x)).
(10)
For fixed T, the vector e" satisfying this system can be found via fixed-point iteration. In practice, we approximate it instead with a single-step update using the solution to the unscaled problem:
r + argmin UBOT(u, T(x;), yy);
c c Ø u
(11)
which empirically provides a good approximation on the optimal e but is significantly more efficient. Apart from requiring single update, whenever and uniform (as in most applications where the samples are assumed to be drawn i..d.) solving this problem between unscaled histograms will be faster and more stable than solving its scaled (and therefore likely non-uniform) counterpart in (10).
Analogously, for a given S, we choose ( that ensures (Projz) ( v).n)l = Sy(S v). For empirical measures, this yields the update:
r argmin UBOT(v,S(y)) u,
r
(12)
4

---

Under review as a conference paper at ICLR 2023
Algorithm 1 Neural Unbalanced Optimal Transport (NUBOT)
Input: g. ICNNs, initialized s.t. Vg(æ) se a and V f(y) R mynn, NNs
for t in epochs do
Sample batch (a:)i=1 ~ and (uyly= ~V
ÿ - Vg(z)
5 Vfly)
** Marginal fitting update Eq. (11)
1 Tı + unbalanced. sinoddrnnkhhnnn 1n m 1m) E,
/* Marginal fitting update Eq. (12)
*/
I2 unbalanced, sinkhorn(i, 1m, 2, 1nn)
Σ, Ty
Ex Po
J(Og. fy) - (f(VV(())))) (Ze, Vg(a:)) m Ee= - 2;f(y))
10
(,( = MSE(e, nT)))
(((0e) MSE(z, ((y))
12
Update fg to minimize J, 0g to minimize Lyp 0g to minimize I.c, and f to maximize
In order to be able to predict mass changes for samples, will the discrete e, to fit continuous versions of Ç via functions parameterized as neural networks trained to achieve (xx.) e; Vi n and ((y,) ss 2 Vj E fI,...,m,, m) through a mean squared error loss.
Updating mappings. Since Ç are tasked with modeling all mass rescalings, for fixed rescaled measures t and we can model the transformation between them with usual (balanced) OT formulation, whereby we seek and S to be pair of optimal (determinstic) OT maps between them. In particular, we use the formulation of (Makkuva et al., 2020) to fit them. That is, T Vg and S = f for convex potentials f and 9. parameterized as ICNNs with parameters 0, and og. The corresponding objective for these two potentials
((f.9) = E [f(Vg(z)) - (x, Vg(z))]] E [f(r)
xmp
= =f [((VV(()) (z, v(())))(())dd( / f())(())
In the finite sample setting, this objective becomes:
C(f.g)
(V(((()))) (xi, Vg(x:))]
Σ (gf(ys)
(13)
The optimization procedure is summarized in Algorithm
Transforming new samples. After learning fig. C. we use these functions transform (map and rescale) new samples, i.e., beyond those used for optimization. For given source datapoint with mass transform it as (x, u) + (Vg(x), n)(x) ((Vg(x))) -1). Analogously, target points can be mapped back to the source domain using (y, u) + (Vf(y). <(y) v (Vf(()) -).
Recovering semi-couplings. Let define r, def diag(e - and F del diag(z 2) T2, where T1, 22 are the solutions of the UBOT problems computed in Algorithm 1 (lines 7 and 9, respectively). It is easy to see that (T,, [T] is a valid pair of semi-couplings between μ and v (cf. Eq. 8).
4 EVALUATION
We illustrate the effectiveness of NUBOT different tasks, including a synthetic setup well an important but challenging task to predict single-cell perturbation responses to a diverse set of cancer drugs with different modes of actions.

---

Under review as a conference paper at ICLR 2023
HIOT GAN NuBOT
Pradlcticos
Weighs
Dusa
Predutirrs
Werghts
2.
Predetions
1
Woggtt
20
0.60
0.2
6.23
Thue Waiahta
Predietions
Weights
True Wekistt
Predictioss
Wooghs
Predicions
Weighs
0.2
133
-1.
0.8
0.00
-1.5
0.3
0.5
0.81
0.00
4.5
0.94
3.35
03
1.02
1.5
1:18
15
Figure 2: Unbalanced sample mapping. In all three scenarios (a, b, c), the source (gray) and target (blue) datasets share structure but have different shifts and per-cluster sampling proportions. The true growth factors of the clusters depicted in the left pane bottom for each setup (True Weights). Tasked with mapping from source to target, NUBOT and UBOT GAN predict the locations (middle pane, red) and weights (right pane) of the transported samples. The number next the weights denotes the mean weights per cluster. While both methods map the samples to the correct location, NUBOT more accurately predicts the weights needed match the target distribution, creating mass (dark blue) or destroying it (red) as needed.
Baselines. To put NUBOT'S performance into perspective, we compare it to several baselines: First, we consider a balanced neural optimal transport method CELLOT (Bunne et al., 2021), based on the neural dual formulation of Makkuva et al. (2020). Further, we benchmark NUBOT against the current state-of-the-ant UBOT GAN, unbalanced OT formulation proposed by Yang & Uhler (2019), which simultaneously learns a transport map and a scaling factor for each source point in order to account for variation of mass. Additionally, we consider the identity matching and modeling cell behavior in absence of a perturbation, and OBSERVED, random permutation of the observed target samples and thus lower bound when comparing predictions to observed cells on the distributional level. Further, we consider DISCRETE OT, the cauropy-regularizee Wasserstein mapping returned by the Sinkhorn algorithm on finite source and target samples. As this method does not parameterize the transport map, we need to include the test cells while computing the optimal coupling and can therefore not be considered out-of-sample. GAUSSIAN APPROX computes a Gaussian approximation of the source and target samples separately, and uses the flosed-form solution of the entropy-regulariae optimal transport problem on unbalanced Gaussians (Janati et al., 2020b) for mapping. More details can be found in the Appendix $D.2.
naive baselines: IDENTITY, simulating
4.1 SYNTHETIC DATA
Populations are often heterogeneous and consist of different subpopulations. To simulate such hetero- geneous intervention responses which exhibit changes in their particle counts. we generate a dataset containing a two-dimensional mixture of Gaussians with three clusters in the source distribution μ.. The target distribution consists of the same three clusters, but with different cluster proportions. Further, each particle has undergone a constant shift in space upon intervention. We consider three scenarios with increasing imbalance between the three clusters (see Fig. 2a-c). We evaluate NUBOT on the task of predicting the distributional shift from source target, while at the same time correctly rescaling the clusters such that no mass is transported across non-corresponding clusters.
Results. The results (setup, predicted Monge maps and weights) are displayed in Fig. 2. Both NUBOT and UBOT GAN correctly map the points to the corresponding target clusters without transporting mass across clusters. NUBOT also accurately models the change in cluster sizes by predicting the correct weights for each point. In contrast, UBOT GAN only captures the general trend of cluster growth and shrinkage, but does not learn the exact weights required re- weighitht the cluster proportions appropriately. The exact setup and calculation of weights can be found in $B (see Table 1), as well as an evaluation of the robustness of NUBOT w.r.t. several hyperparamaters (see Fig. 7) and a comparison to UBOT (see Fig. 8, 9).
4.2 HINGLE-CELL PERTURBATION RESPONSES
Through the measurement of genomic, transcriptomic, proteomic or phenotypic profiles of cells, and the identification of different cell types and cellular states based on such measurements, biologists

---

Under review as a conference paper at ICLR 2023
NUBOT
CELGOT
UBOT GAN
INENIITY
OBSEWED
DISCURTE OT
GAUSSIAN ArPROX
8h Timestep
Weighted MMD
10-1
10-2
10
Imatinih
Ertatinil
Midostaurin
Dasatinil
Regerafonil:
Panobinestal
Trametinih
Paditazel
Cisplaun
Lanalidanide
Tomozolamide
Dectabine
Everolimus
Crizatinll
Dezamethusone
Malphalau
Irazanili
Claparili
Vindesihu
Sarafonil
Dacarbazine
Staurosporine
Hydrogyuret
Palbocselit
Ulisertinbb
Duhrafonib
24h Timestep
Drug
10
Weighted MMD
10-2
10-3
Inittinb
Ertetinih
Midostaurin
Dasatinil
Regorafonib
Fanobinantsi
Trametinih
Paditasel
Ciaplatin
Lenalddanide
Thmozolonide
Deatabine
Evarolimua
Crozoinib
Desanathasnne
Malphalan
Irammii
Olsparil
Vindesine
Sarafinib
Dacarbazine
Staunsporine
Hydrozyuree
Pulbocielli
Uliserinib
Duhrafonib
Drug
Figure 3: Distributional fit of the predicted perturbed cell states to the observed perturbed cell states for each drug and timestep, measured by a weighted version of kernel MMD on a set of held-out control cells. For NUBOT and UBOT GAN, MMD is weighted by the predicted weights, while for the other baselines it is computed with uniform weights. OBSERVED corresponds to a random permutation of the observed control cells, i... its distribution approximately the same as the observed cells.
have revealed perturbation responses and response mechanism which would have remained obscured in bulk analysis approaches (Green & Pelkmans, 6016; Liberali et al., 2014; Kramer et al., 2022). How- ever, single-cell measurements typically require the destruction the cells in the course of recording.
Thus, each measurement provides us only with a snapshot of cell populations, i.e., samples of probability distribution that is evolving in the course of the perturbation, from control j (source) to perturbed cell states v (target). Using NUBOT and the considered baselines, we learn a map T that reconstructs how individual cells respond to a treatment. The effect of a single perturbation frequently varies depending on the cell type or cell state, and may include the induction of cell death or proliferation. In the following, we will evaluate if NUBOT able to capture and predict heterogeneous proliferation and cell death rates of two co-cultures melanoma cell liens through and ( in response to 25 drug treatments.
The single-cell measurements used for this task were generated using the imaging technology 4i (Gut et al., 201 8) over the course of 24 hours, resulting in three different unaligned snapshots (1 0h, 1= 8h and 24h) for each of the drug treatments. The control cells, i.e., the source distribution μt, consists of cells taken from a mixture of melanoma cell lines at = 0/ that are exposed to a dimethyl sulfoxide (DMSO) as a vehicle control. Futher, We consider two different target populations captur- ing the perturbed populations after t = 8h and t = 24h of treatment, respectively. As both cancer cell lines exhibit different sensitivities to the drugs (Raaijmakers et al., 2015). their proportion (Fig. 15) as well as the total cell counts (Fig. 17) vary over the time points. Both cell lines are characterized by the expression of mutually exclusive protein markers. i.e., one cell line strongly expresses a set of proteins detected by an antibody called MelA (MelA+ cell type), while the other is characterized by high levels of the protein Sox9 (Sox9* cell type). An evaluation of this cell line annotation can be found in Fig. 14 (8h) and Fig. 16 (24h). As no ground truth matching is available, cells after and 24 hours of treatment (Fig. 15, 17), as well as the cell type annotation for each cell to further evaluate NUBOT'S performance. A detailed description of the dataset can be found in $ C.2.
use insights from the number of
7

---

Under review as a conference paper at ICLR 2023
Results. We split the dataset into a train and test set and train NUBOT as well as the baselines on unaligned unperturbed (control) and perturbed cell populations for each drug. During evaluation, we then predict out-of-sample the perturbed cell state from held-out control cells. Details on the network architecture and hyperparameters can be found in $ D.3. NUBOT and UBOT GAN additionally predict the weight associated with the perturbed predicted cells, giving insights into which cells have proliferated or died in response to the drug treatments. First, we compare how well each method fits the observed perturbed cells on the level of the entire distribution. For this, we measure the weighted version of kernel maximum mean discrepancy (MMD) between predictions and observations. More details on the evalu- ation metrics can be found in $ D.I. The results are displayed in Fig. 3. We additionally report the weighted Wasserstein distance in Fig. 10 in Appendix BB. NUBOT outperforms all baselines in almost all drug perturbations, showing its effectiveness in predicting OT maps and local variation in mass.
a. 8 h
0.65
R- 0.95
p- 2.2e-25
0.60
Cell Type Fraction
0.35
0.50
0.45
0.40
Nicensott
MelA
0.35-
Sox9
800
900
1200
b.
24 h
0.8
Jinwtak
0.7
0.44
Cell Type Fraction
P- 0,004
0.6
In the absence of a ground truth and in particular, given our in- ability measure (i.e.. observe) cells which have died upon treat- ment, we are required to base further analysis of NUBOT's pre- dictions on changes in cell count for each subpopulation (MelA Sox9+). Fig. 15 clearly shows that drug treatments lead to sub- stantially different cell numbers for each of the subpopulations compared to control. For example, Ulixetinni leads to the pro- liferation of both subpopulations after 8h, but to pronounced cell death in Sox9 and strong proliferation in MelA + cells after 24h. We thus expect, that weights predicted by NUBOT for all drugs correlate with the change in cell counts for each cell type (here measured as population fractions). This is indeed the case, Fig. 4 shows high correlation between observed cell counts of the cell types and the sum of the predicted weights of the respective cell types after 8h of treatment for all drugs. After 24 hours, treatment-ing -nduueed-- cell death (in at least one cell type) by some drugs can be so severe at that the number of observed perturbed cells becomes too low for accurate predictions and the evaluation of the task (Fig. 17). Further, we find that drugs influence the abundance of the cell lines markers MelA and Sox9, complicat- ing cell type classification (see Fig. 14, 16). We ignore drugs falling into these categories and find that whilst the correlation between predicted weights and observed cell counts is reduced after 24h (see Fig. 4a), NUBOT still captures the overall trend.
0.5
0.4
0.3
MelA
Sox9
0.2
100
1000
1200
2000
1600
Summed Weights
Figure 4: Given the ground truth on the known subpopulation (MelA (red) and Sox9 (blue)) sizes for each drug, we analyze their level of correlation to our pre- dicted weights after a. 8h and b. 24h. With increasing difficulty of the task and certain drugs com- pletely removing both or one of the subpopulations, the level of correlation reduces from 8 to 24h. R denotes Pearson correlation co- efficient and P the p-value.
The data further provides insights into biological processes such as apoptosis, a form of programmed cell death induced by enzymes called Caspases (CICasp3). While dead cells become invisible in the cell state space (they cannot be measured), dying cells are still present in the observed perturbed sample and can be recognized by high levels of CICasp3 (the apopotosis markers). Conversely, the protein Ki67 marks proliferating cells. Analyzing the correlation of CICasp3 and Ki67 intensity with the predicted weights provides an additional assessment of the biological meaningfulness of our results. For example, upon Ulixertinib treatment, the absolute cell counts show an increase in S cells, and a decline of MelA + cells 24h (Fig. 15). Fig. 5 shows UMAP projections of the control cells at both time points, colored by the observed and predicted protein marker values and the predicted weights. At 8h, NUBOT predicts only little change in mass, but few proliferative cells with high weights in areas which are marked by high values of the proliferation marker Ki67. At 24h, our model predicts cell death in the S (MelA ) cell type, and proliferation in the MelA + cell type, which matches the observed changes in cell counts per cell type, seen in Fig. 15 in $ B. We identify similar results for Trametinib (Fig. 11), Ixazomib (Fig. 12), and Vindesine (Fig. 13) which can be found in B. These experiments thus demonstrate that NUBOT accurately predicts heterogeneous drug responses at the single-cell level, capturing both, cell proliferation and death.
8

---

Under review as a conference paper at ICLR 2023
a. K167 (Nuclei) 8 in Control
Ki67 (Nuclei)
b. MelA (Cell) 24 in Control
MelA (Cell) in Predicted
in Predicted
Weights
Weights
2
-1
-1
0
Figure 5: UMAP projections of the control cells for Ulixertinib 8h and b. 24h. Cells are colored by the observed and predicted protein marker values (Ki67, MelA), and predicted weights. NUBOT thereby correctly predicts weights 2 for proliferating cells in the MelA population (a. and a. right panel), and increased levels of cell death in the Sox9+ population after 24h via weights S 1 (b., right panel), confirmed by the experimental observations (see Fig. 15).
CONCLUSION
This work presents novel formulation of the unbalanced optimal transport problem that bridges two previously disjoint perspectives on the topic: a theoretical one based on semi-couplings and practical one based recent neural estimation of OT maps. The resulting algorithm, NUBOT, is scalable, efficient, and robust. Yet, it is effective at modeling processes that involve population growth death, as demonstrated through various experimental results on both synthetic and real data. On the challenging single-cell perturbation task, NUBOT is able to successfully predict perturbed cell states, while explicitly modeling death and proliferation. Explicitly modeling proliferation and death at the single-cell level as part of the drug response, allows to link cellular properties observed prior to drug treatment to therapy outcomes. Thus, the application of NUBOT in the fields of drug discovery and personalized medicine could be of great implications, as it allows to identify cellular properties predictive of drug efficacy.
REFERENCES
David Alvarez-Melis, Yair Schiff, and Youssef Mroueh. Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks. Transactions on Machine Learning Research (TMLR), 2022.
Brandon Amos, Lei Xu, and Zico Kolter. Input Convex Neural Networks. International Conference Machine Leurning (ICML), volume 34, 2017.
Martin Bauer, Emmanuel Hartman, and Eric Klassen. The Square Root Normal Field Distance and Unbalanced Optimal Transport. Applied Mathematics & Optimization, 85(3), 2022.
Florian Beier, Johannes von Lindheim, Sebastian Neumayer, and Gabriele Steidl. Unbalanced Multi-Marginal Optimal Transport. Journal of Mathematical Imaging and Vision, 2022.
Jean-David Benamou. Numerical resolution of an "unbalanced" mass transport problem. ESAIM: Mathematical Modeling and Numerical Analysis, 37((), 2003.
Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyré. Iterative Bregman Projections for Regularized Transportation Problems. SIAM Journal Scientific Computing, 37(2), 2015.
Nicolas Bonneel and David Cocurjolly. SPOT: Sliced Partial Optimal Transport. ACM Transactions on Graphics (TOG), 38(4), 2019.
Yann Brenier. Décomposition polaire et réarrangement monotone des champs de vecteurs. CR Acad. Sci. Paris Sér. I Math., 305, 1987.
Charlotte Bunne, Stefan G Stark, Gabriele Gut, Jacobo Sarabia del Castillo, Kjong- Van Lehmann, Lu- cas Pelkmans, Andreas Krause, and Gumnar Ratsch. Learning Single-Cell Perturbation Responses using Neural Optimal Transport. bioRxiv, 2021.

---

Under review as a conference paper at ICLR 2023
Charlotte Bunne, Ya-Ping Hsieh, Marci Cuturi, and Andreas Krause. Recovering Stochastic Dynamies via Gaussian Schrödinger Bridges. arXiv Preprint arXiv:2202.05722, 2022a.
Charlotte Bunne, Andreas Krause, and Marco Cuturi. Supervised Training of Conditional Monge Maps. In Advances in Neural Information Processing Systems (NeurIPS), volume 36, 2022b.
Charlotte Bunne, Lactitia Meng-Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal Optimal Transport Modeling of Population Dynamics. In International Conference on Artificial Intelligence and Statisties (AISTATS), volume 25, 2022c.
Luis A Caffarelli and Robert McCann. Free boundaries in optimal transport and Monge-Ampere obstacle problems. Annals of Mathematics, 2010.
Laetitia Chapel, Mokhtar Z Alaya, and Gilles Gasso. Partial Optimal Transport with Applications on Unlablee Learning. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, 2020.
Laetitia Chapel, Rémi Flamary, Haoran Wu, Cédric Févotte, and Gilles Gasso. Unbalanced Optimal Transport through Non-negative Penalized Linear Regression. Advances in Neural Information Processing Systems (NeurIPS), 34, 2021.
Yize Chen, Yuanyuan Shi, and Baosen Zhang. Optimal Control Via Neural Networks: Convex Approach. In International Conference Learning Representations (ICLR), 2019.
Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and Srançois-Xaviee Vialard. Scaling algorithms for Unbalanced Optimal Transport Problems. Mathematics of Computation, 87(314), 2018a.
Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François- -avier Vialard. Unbalanced Optimal Transport: Dynamic and Kantorovich Formulations. Journal of Functional Analysis, 274(11), 2018b.
Marco Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. In Advances in Neural Information Processing Systems (NeurIPS), volume 26, 2013.
Atray Dixit, Oren Parnas, Biyu Li, Jenny Chen, Charles P Fulco, Livnat Jerby-Arnon, Nemanja D Marjanovic, Danielle Dionne, Tyler Burks, Raktima Raychowdhury, et al. Perturb-Seq: Dissecting Molecular Circuits with Scalable Single-Cell RNA Profiling of Pooled Genetic Screens. Cell, 167 (7), 2016.
Debidatta Dwibedi, Yusuf Aytar, Jonathan Tompson, Pierre Sermanet, and Andrew Zisserman. Temporal Coccs-Consi Learning. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2019.
Jiaojiao Fan, Shu Liu, Shaojun Ma, Yongxin Chen, and Haomin Zhou, Scalable Computation of Monge Maps with General Costs. arXiv preprint arXiv:2106.03812, 2021a.
Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen. Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks. In International Conference on Machine Learning (ICML), 2021b.
Kilian Fatras, Thibault Séjourné, Rémi Flamary, and Nicolas Courty. Unbalanced minibatch Optimal Transport; applications to Domain Adaptation. In International Conference on Machine Learning (ICML), 2021.
Alessio Figalli. The Optimal Partial Transport Problem. Archive for Rational Mechanics and Analysis, 195(2). 2010.
Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya, and Tomaso A Poggio. Learning with a Wasserstein loss. In Advances Neural Information Processing Systems (NeurIPS), volume 28, 2015.
Alexandre Gramfort, Gabriel Peyré, and Marco Cuturi. Fast Optimal Transport Averaging of Neuroimaging Data. In International Conference on Information Processing in Medical Imaging (IPMI). Springer, 2015.
10

---

Under review as a conference paper at ICLR 2023
Victoria A Green and Lucas Pelkmans. A systems survey of progressive host-cell reorganization during rotavirus infection. Cell Host & Microbe, 20(1), 2016.
Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf. and Alexander Smola. A kernel two-sample test. Journal of Machine Learning Research, 13, 2012.
Qipeng Guo, Zhijing Jin, Ziyu Wang. Xipeng Qiu, Weinan Zhang, Jun Zhu, Zheng Zhang, and Wipf David. Fork or Fail: Sycle-Consistent Training with Many-to-Ont Mappings. In International Conference Artificial Intelligence and Statistics (AISTATS), 2021.
Gabriele Gut, Markus D. Herrmann, and Lucas Pelkmans. Multiplexed protein maps link subcellular organization to cellular states. Science, 361 2018.
Ernst Hellinger. Neue Begründung der Theorie quadratischer Formen von unendlichvielen Veränder- lichen. Journal für die reine und angewandte Mathematik, 1109( 1909.
Judy Hoffman, Eric Tzeng, Taesung Park, Jun- Yan Zhu, Phillip Isola, Kate Saenko, Alexci Efros, and Trevor Darrell. CyCADA: Cycle- Consistenttt Adversarial Domain Adaptation. International Conference Machine Learning (ICML), 2018.
Chin-Wer Huang, Ricky T. Q. Chen, Christos Tsirigotis, and Aaron Courville. Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization. In International Conference Learning Representations (ICLR), 2021.
YoonHaeng Hur, Wenxuan Guo, and Tengyuan Liang. Reversible Gromov-Monge Sampler for Simulation- Based Inference. arXiv preprint 0rXiv:2109.14099 2021.
Hicham Janati, Marco Cuturi, and Alexandre Gramfort. Spatio-Temporal Alignments: Optimal transport through space and time. In International Conference Artificial Intelligence and Statistics (AISTATS), volume 23, 2020a.
Hicham Janati, Boris Muzellee, Gabriel Peyré, and Marco Cuturi. Entropic Optimal Transport between Unbalanced Gaussian Measures has Closed Form. Advances Neural Information Processing Systems (NeurIPS), 33, 2020b.
Zdenek Kalal, Krystian Mikolajezyk, and Jiri Matas. of Tracking Failures. In IEEE Conference Computer Vision and Pattern Recognition (CVPR). IEEE, 2010.
Error: Automatic Detection
Kantorovich. On the transfer of masses (in Russian). In Doklady Akademii Nauk, volume 37, 1942.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference Learning Representations, 2020.
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? continuous wasserstein-2 benchmark. Advances in Neural Information Processing Systems, 34, 2021.
Bernhard A Kramer, Jacobo Sarabia del Castillo, and Lucas Pelkmans. Multimodal perception links cellular state to decision-making in single cells. Science, 377 648, 2022.
John Lee, Nicholas P Bertrand, and Christopher J Rozell. Parallel Unbalanced Optimal Transport Regularization for Large Scale Imaging Problems. arXiv preprint arXiv:/909.00149 2019.
Prisca Liberali, Berend Snijder, and Lucas Pelkmans. A hierarchical map of regulatory genetic interactions in membrane trafficking. Cell, (57(6))14733144 2014.
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal Transport in Competition with Re- action: The Hellinger- Kantorovich Distance and Geodesic Curves. SIAM Journal on Mathematical Analysis, 48(4), 2016.
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal Transsor problems and new Hellinger-Kantoro/eee Kantorovich distance between positive measures. Inventiones Mathematicae, 211 (3), 2018
11

---

Under review as a conference paper at ICLR 2023
Mohammad Lotfollahi, F Alexander Wolf, and Fabian Theis. scGen predicts single-cell perturbation responses. Nature Methods, 16(8). 2019.
Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning (ICML), volume 37, 2020.
Facundo Mémoli and Tom Needham. Distance distributions and inverse problems for metric measure spaces. Studies Applied Mathematics, 2022.
Gaspard Monge. Mémoire sur théorie des déblais et des remblais. Histoire de l'Académie Royale des Sciences, pp. 666-704, 1781.
Debarghya Mukherjee, Mikhail Yurochkin, Moulinath Banerjee, and Yuckai Sun. Two Simple Ways to Learn Individual Fairness Metrics from Data. In International Conference on Machine Learning (ICML). PMLR, 2020.
Ofir Pele and Michael Werman. Fast and robust Earth Mover's Distances. In International Conference on Computer Vision (ICCV). IEEE, 2009.
Gabriel Peyré and Marco Cuturi. Computational Optimal Transport. Foundations and Trends in Machine Learning, 1 11(5-6), 2019. ISSN 1935-8245.
Khiem Pham, Khang Le. Nhat Ho, Tung Pham, and Hung Bui. On Unbalanced Optimal Transport: An Analysis of Sinkhorn Algorithm. In International Conference on Machine Learning (ICML), 2020.
Marieke IG Raaijmakers, Daniel S Widmer, Melanie Maudrich, Tabea Koch, Alice Langer, Anna Flace, Claudia Schnyder, Reinhard Dummer, and Mitchell P Levesque. A new live-cell biobank workflow efficiently recovers heterogencous melanoma cells from native biopsies. Experimental Dermatology, 24(5):377-380, 2015.
Jack Richter-Powell, Jonathan Lorraine, and Brandon Amos. Input convex gradient networks. arXiv preprint arXiv:2 2021.
Filippo Santambrogio. Optimal Transport for Applied Mathematicians. Birkhäuser, NY, 55(58-63): 94, 2015.
Paul-Edouard Sarlin, Daniel DeTone, Tomasz Malisiewicz, and Andrew Rabinovich. SuperGlue: Learning Feature Matching with Graph Neural Networks. In IEEE Conference on Computer Vision and Puttern Recognition (CVPR), 2020.
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Tpaimal-Transport Analysis of Single- Gene Expression Identifies Developmental Trajectories in Reprogramming. Cell, 176(4), 2019.
Thibault Séjourné, Fravçois-Xavier Vialard, and Gabriel Peyré. The Unbalanced Gromov Wasserstein Distance: Conic Formulation and Relaxation, Advances in Neural Information Processing Systems (NeurIPS), 34, 2021.
Thibault Séjoumé, Gabriel Peyré, and François- François-Xavier Vialard. Unbalanced Optimal Transport, from Theory to Numerics. arXiv preprint arXiv:22 1.08775, 2022a
Thibault Séjouréé, François- Xavier Vialard, and Gabriel Peyré. Faster Unbalanced Optimal Transport: Translation invariant Sinkhorn and I-D Trank-Wolfe. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2022b.
Daniel Sheldon, MA Elmohamed, and Dexter Kozen. Collective Inference Markov Models for Modeling Bird Migration. In Advances Neural Information Processing Systems (NeurIPS), volume 20, 2007.
Tianxiao Shen, Tao Lei, Regina Barzilay, and Tommi Jaakkola. Style Transfer from Non-Parallel Text by Alignment. Advances in Neural Information Processing Systems (NeurIPS), 30, 2017.
12

---

Under review as a conference paper at ICLR 2023
Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics. International Conference Machine Learning (ICML), 2020.
Stefan Van der Walt, Johannes L Schönberger, Juan Nunez-Iglesias, François Boulogne, Joshua D Warner, Neil Yager, Emmanuelle Gouillart, and Tony Yu. scikit-image: image processing in python. PeerJ, 2te453, 2014.
Karren D Yang and Caroline Uhler. Scalable Unbalanced Optimal Transport using Generative Adversarial Networks. International Conference Learning Representations (ICLR), 2019.
Bo Yuan, Ciyue Shen, Augustin Luna, Anil Korkut, Debora S Marks, John Ingraham, and Chris Sander. CellBox: interpretable machine learning for perturbation biology with application to the design of cancer combination therapy. Cell Systems, 12(2), 2021.
Tianzuo Zhan, Giulia Ambrosi, Anna Maxi Wandmacher, Benedikt Rauscher, Johannes Betge, Niklas Rindtorff, Ragna S Häussler, Isabel Hinsenkamp. Leonhard Bamberg, Bernd Hessling, et al. MEK inhibitors activate Wnt signalling and induce stem cell plasticity in colorectal cancer. Nature Communications, 10(1). 2019.
Qiang Zhang, Tete Xiao, Alexei A Efros, Lerrel Pinto, and Xiaolong Wang. Learning Cross-Domain Correspondence for Control with Dynamics pooll-Consistency. International Conference on Learning Representations (ICLR), 2021.
Zhengxin Zhang, Youssef Mroueh, Ziv Goldfeld, and Bharath Sriperumbudur. Cycle Consistent Probability Divergences Across Different Spaces. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2022.
Jun-Yan Zhu, Taesung Park, Phillip Isola, and Alexei Efros. Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2017a.
Jun Yan Zhu, Richard Zhang, Deepak Pathak, Trevor Darrell, Alexei A Efros, Oliver Wang, and Eli Shechtman. Toward Multimodal Image-to-Image Translation. In 30, 2017b.
13

---

Under review as a conference paper at ICLR 2023
APPENDIX
A RELATED WORK
In the following, provide further information and review related literature concepts discussed throughout this work.
A.1 UNBALANCED OPTIMAL TRANSPORT
Unbalanced optimal transport is generalization of the classical OT formulation (I), and as such allows mass to be created and destroyed throughout the transport. This relaxation has found recent use cases in various domains ranging from biology (Schiebinger et al., 2019; Yang & Uhler, 2019), imaging (Lee et al., 2019), shape registration (Bonneel & Coeurjolly, 2019), domain adaption (Fatras et al., 2021), positive-unlabeled learning (Chapel et al., 2020), to general machine learning (Janati et al., 2020a; Frogner et al., 2015). Problem (6) provides a general framework of the unbalanced optimal transport problem, which can recover related notions introduced in the literature: Choosing for D, the Kullback-Leibler divergence, one recovers the so-called squared Hellinger distance, with D l2 norm, arrive at Benamou (2003), and an & norm retrieves concept often referred to robust OT (Mukherjee et al., 2020) with connections to a concept known as partial OT (Figalli, 2010). The latter comprises approaches which do not rely a relaxation of the marginal constraints as in (6). In particular, some strategies of partial OT expand the original problem by adding virtual mass the marginals (Pele & Werman, 2009; Caffarelli & McCann, 2010; Gramfort et al., 2015), or by extending the OT map by dummy rows and columns (Sarlin et al., 2020) onto which excess mass can be transported. The semi-coupling formulation (also known as Wasserstein- Seattein-Fisher-RRee Fisher-Rao distance) has been independently proposed by Liero et al. (2016; 2018), and recently been connected to shape analysis of surfaces (Bauer et al., 2022). Recent work has furthermore developed alternative computational schemes (Chapel et al., 2021; Séjourné et al., 2022b) as well as provided a computational complexity analysis (Pham et al., 2020) of the generalized Sinkhorn algorithm solving entropic regularized unbalanced OT (Chizat et al., 2018a). Besides Yang & Uhler (2019), 9), these approaches do not provide parameterizations of the unbalanced problem and allow for an out-of-sample generalization which consider in this work. The notion of unbalanced optimal transport has been further generalized to the multi-marginal setting Beier et al. (2022) as well as for Gromov-Wassen Wasserstein (Séjouréé et al., 2021). Janati et al. (2020b) further provide closed-form solution of entropic optimal transport between unbalanced Gaussian measures. For complete review, we refer the reader to Séjourné et al. (2022a) as well as (Peyré & Cuturi, 2019, Chapter 10.2).
It becomes evident from the discussion above that unbalanced optimal transport is better thought of as a family of problems, rather than a single specific one. Unbalanced OT arises in any setting where one seeks a notion of distance or correspondence across datasets/populattinn not normalized in the same way and/or (ii.) one does not want to rigidly enforce full correspondence between the datasets (e.g., to minimize the effect of outliers or because there is reason to believe some datapoints do not occur in both datasets). This work proposes a method to parameterize an UBOT problem which does not necessarily match the typical Kantorovich entropy-regularizee formulation. We seek a solution that learns which part of the distribution shrinks and which one exhibits overall growth, motivated by heterogeneous sobpopulation structures present in single-cell biology. To that end, NUBOT aims at modeling birth and death dynamics throughout the distribution shift recovered through optimal transport mapping.
and (i.) the populations are
A.2 CYCLE-CONSISTER LEARNING
The principle of cycle-consistency has been widely used for learning bi-directional transformations between two domains of interest. Cycle-consistency thereby assumes that both the forward and backward mapping are roughly inverses of each other. In particular, given unaligned points æ E X and y E y. well as maps g : X + y and f 3 + X. cycle-consistency reconstruction losses enforce ||- ff((x)) as well as | ((((v))l using some notion of distance || I]. assuming that there exists such ground truth bijection g = f-1 and f f= =g" The advantage of validating good matches by cycling between unpaired samples becomes evident through the numerous use cases to which sycle-consistency has been applied: Originally introduced within the field of computer vision
14

---

Under review as a conference paper at ICLR 2023
(Kalal et al., 2010) and applied to sange-to-image translation tasks (Zhu et al., 2017a), it has been quickly adapted to multi-modal problems (Zhu et al., 2017b), domain adaptation (Hoffman et al., 2018), and natural language processing (Shen et al., 2017). The original principle has been further generalized to settings requiring a many-to-one or surjective mapping between domains (Guo et al., 2021) via conditional variational autoencoders, dynamic notions of cycle-consistency (Zhang et al., 2021), or to time-varying applications (Dwibedi et al., 2019). These classical approaches enforce cycle-consistency by explicitly composing both maps and penalizing for any deviation from this bijection. In this work, we treat cycle-consisteney differently. It is enforced implicitly by coupling the two distributions of interest through sequence of reversible transformations: re-weighting. transforming, and re-weighting (Eq. (9) and Fig. 1).
Similarly to our work, Zhang et al. (2022) and Hur et al. (2021) establish a notion of cyclstcon (reversibility) for pair of pushforward operators to align two unpaired measures. Both methods rely on the Gromov- -Mygge distance (Mémoli & Needham, 2022), a divergence to compare probability distributions defined on different ambient spaces X' and y--a setting not considered in this work. They proceed by defining a reversible metric through replacing the single Monge map by a pair of two Monge maps, i.e., f X + y and g : y + X, minimizing the objective
inf GM(H, = A. (s:)))) A3((: + A g3
(14)
A) (B(lex(a,>')) (()
D5(g:v)- [E 3/) -x((60))000))))
A. v) (Ex - cy)(())))))
Problem (14) shows similarities the classical sycle-consistenc objective of Zhu et al. (2017a), where cycle-consistency is indirectly enforced through 2 Zhang et al. (2022) parameterize both Monge maps through neural networks in similar fashion as done in (Yang & Uhler, 2019; Fan et al., 2021a). Our approach differs from Zhang et al. (2022); Hur et al. (2021) as we model the problem through a single Monge map with duals f, g. allowing us to map back-and-foo and 1s, and using a different parametrization approach (ICNNs). More importantly, the approaches presented by Zhang et al. (2022); Hur et al. (2021) do not generalize to the unbalanced case. While Zhang et al. (2022) proposed an unbalanced version of (14) by relaxing the marginals as done in Chizat et al. (2018a) they require the unbalanced sample sizes to be known (i.e., and m need to be fixed). In our application of interest, particle counts of the target population are, however, not known a priori.
between measures
A.3 CONVEX NEURAL ARCHITECTURES
Input convex neural networks (Amos et al., 2017) are a class of neural networks that approximate the family of convex functions to with parameters 6, i.c., whose outputs the (z) are convex w.r.t. an input æ. This property is realized by placing certain constraints the networks parameters 0. More specifically, an ICNN is an I--layer feed-forward neural network, where each layer I = {0, {O...., 1} is given by
11+1 = o7 (Wfx + Wf + bx) and the (z) = 2L,
(15)
where an convex non-decreasing activation functions, and (W7, WF, be is the set of parameters, with all entries in W/2 being non-negative and the convention that zg and W3 are O. As mentioned above and through the connection established in $ 2, convex neural networks have been utilized to approximate Monge map T (3) via the convex Brenier potential ei connected to the primal and dual optimal transport problem. In particular, has been used to model convex dual functions (Makkuva et al., 2020) as well as normalizing flows derived from convex potentials (Huang et al., 202 The expressivity and universal approximation properties of ICNNS have been further studied by Chen et al. (2019), who show that any convex function over a compact convex domain can be approximated in sup norm by an ICNN. To improve convergence and robustness of ICNNs -kown to be notoriously difficult to train (Richter-Powell et al., 2021)- different initialization schemes have been proposed: Bunne et al. (2022b) derive two initialization schemes ensuring that upon initialization Vih mimics affine Monge map T mapping either the source measure onto itself
15

---

Under review as a conference paper at ICLR 2023
Table I: Setup of the synthetic mixture of Gaussians dataset, showing the proportions of the three clusters in source and target distribution in three different settings (a., b., c.) as well as the required scaling factor per cluster needed match the target without transporting points non-corresponding clusters. The last two columns show the mean weights obtained by NUBOT and UBOT GAN.
Setting
Cluster
True Scaling Factor
Mean Weights
Difference to Ground Truth
NubOT
ubOT GAN
NubOT
ubOT GAN
0.3
0.26
0.8
0.04
0.5
2
1,35
1,36
0,99
0.01
0.36
1.35
1.32
1.02
0.03
0.33
b.
0.3
0.29
0.81
0.01
0.51
0.6
0.64
0,88
0.04
0.28
2.1
2.08
1.18
0.02
0.92
c.
4.5
4.6
1,44
0.1
3.06
1.0
0.98
0,94
0.02
0.06
3
0.22
0.23
0.79
0.01
0.57
(identity initialization) or providing map between Gaussian approximations of measures μ and v (Gaussian initialization). Further, Korotin et al. (2020) proposed to quadratic layers as well pre-training pipeline to initialize ICNN parameters to encode an identity map.
B ADDITIONAL EXPERIMENTAL RESULTS
B. SYNTHETIC DATA
In our synthetic two-dimensional dataset, the source and target distribution are mixtures of Gaussians with varying proportions. Both source and target consist of three corresponding clusters, and by changing the proportions of each cluster, we illustrate scenario which subpopulations grow and shrink at different rates. The clusters are sampled according to the following probabilities: a. source (0.3, 0.3, 0.3) to target 10.1, 0.45,,444.. 0.45). b. source (0.3, 0.3, 0.3) target (0.1, 0.2, 0.7). c. source 10. 0.45, 0.3) to target (0.45, 0.45, 0.1]. Table shows the shares of the three clusters in the source and target distributions. In order to match the target distribution without transporting mass across non correspopoggg clusters, the clusters have be re-scaled with the factors presented in column "True Scaling Factor'. The last two columns show the mean weights per cluster obtained by NUBOT and UBOT GAN, respectively. UBOT GAN captures only the general trend in growth and shrinkage, the exact weights do not scale the cluster proportions appropriately. In contrast, the weights obtained NUBOT match the required scaling factors very closely. Fig. 6, shows the weighted MMD between the source distribution and the target distribution, confirming superior performance of NUBOT.
NabOT
ubOT GAN
Weighted MMD
B
A
Data Sctup
Figure Distributional of the predicted samples to the target samples on synthetic data, measured by weighted version of kernel MMD.
16

---

Under review as a conference paper at ICLR 2023
B.1.1 ROBUSTNESS NUBOT W.R.T. HYPERPARAMETERS
In the following, we evaluate NUBOT'S sensitivity to hyperparameter choices. For this we screen various parameter ranges of the batch size (bs), as well hyperparameters of the generalized Sinhorn algorithm (Chizat et al., 2018at 8a) to solve the UBOT (see Algorithm 1. 5 and 7), i.c., the entropy regularization parameters (reg), and relaxation penalties (regm). ). All considered hyperparameters can be found in the legend of Fig. 7. The base model thereby denotes the final hyperparameters chosen for all experiments conducted in this work. We conduct this analysis the synthetic data setup introduced in $ 4.1 and Figure 2a and b, i... mixture of three Gaussians exhibiting different changes in particle count with known ground truth. We compare the obtained weights of each hyperparameter configuration for each setups (a. and b.). The results displayed in Fig. 7 (left panels) demonstrate that the correct weights are robustly learned for all (in setting a.) and most (in setting b.) hyperparameter choices. In setup b., the weights deviate in particular for high values of regm. (2 0.5). Fig. (right panels) further show the correlation between the obtained weights the base model (i.e., the parameters chosen for all experiments) with all other hyperparameter configurations
Ryperparaneter
Setting a.
me 0,003
tx: 259
1.75
ng 0,005
ma.m
190
0.04
Whytt SyBitee OUIue
mg 1158
2:65
19 03654
5.00
6264
63071
m2 1:88
191 6301
0:25-
13,8 1/68
192 0300
Tx 399
0.03
19 6084
69646
26 250
Setting b.
3.0
1.0
2.0
0.55
0.50
Weights (snun)
6.25
Lio 1.55
2084
Wiiggna is aitineal SNEAND aetagg
"g 1,53
14:291
259
R3.7.m 1m
2.16
m. 130
26.59
8g.80 C104
my 0354
0.22
2:35
191
22 n. 1258
99 6344
14.39
92.5
mg 6/83
29
0.00
-32
0.5
1.0
2.0
2.5
Weagto
Weightt (8000)
6,054
199
Figure 7: Hyperparameter screen of the NUBOT model. We screen various different hyperpa- rameter configurations different to the base model (the one used for the experiments above) and evaluated them for different synthetic data settings (a. and b.). (left panels) Distribution of learned weights obtained for each hyperparameter setting and (right panels) correlation of the weights for all hyperparameter configurations compared to the base setup.
B.1.2 COMPARISON TO UBOT
NUBOT provides a parameterization of an unbalanced optimal transport problem. At each training step, NUBOT executes the generalized Sinkhorn algorithm (Chizat et al., 2018a; Cuturi, 2013; Benamou et al., 2015) as subroutine obtain estimates of the weights e and To compare the solution of NUBOT to the mapping computed by the UBOT problem in (7), we further analyze the obtained solutions and hyperparameter sensitivities of UBOT itself, i... not integrated within NUBOT. In particular, we compare different relaxation penalties = choices of the entropy regularization e.
(see (7)), and different
We consider the same synthetic data settings (a., b.) of Gaussian mixtures, as introduced in $ 4 and shown in Fig. 2, where the three clusters grow/shrink at different rates which are known. We show the couplings and weights obtained from UBOT in Fig. 8 (setting a.) and Fig. 9 (setting b.), for different regularization (c) and relaxation penalty (r) parameters. The weights per source point thereby computed by summing relative to all other weights, we additionally normalize these weights by the total sum the coupling matrix. In most cases UBOT couples points only between corresponding clusters. As increases
the columns in the coupling matrix. To get weight value
17

---

Under review as a conference paper at ICLR 2023
for 0.05, mass variation is penalized stronger, and therefore, mass has to be coupled between non-coreeponddng clusters as well in order to the marginals.
The non-normalized weights exhibit significantly higher variation between different values of especially for e = 0.05. Comparing the average weights per cluster to the true scaling, which are shown in Table 1 and Fig. 2, observe that both the non-normalized and normalized weights matching these values closely for r E {1,10) 10) and c = 0.005. For c = 0.05 and small values of т, the coupling is still only mapping points between corresponding clusters. The weights, however, not matching the true underlying scalings.
Setting a.
E 0.005
UBOT
E
0.05
Coupling
Weights
Norm. Weights
Coupling
Weights
Norm. Weights
1.5
1.5
0.97
2.25
0.37
0.85
521.01
1 1603..00 8288
0.38
1.12
T= 0.01
-1.0
-10
4.63
1.76
3.72.2
1.52
0.0
0.0
Coupling
Weights
Norm. Weights
Coupling
Weights
Norm. Weights
1.5
1.5
T= 0.05
0.64
1.3
0.55
1.12
110
14.06
33.77
0.49
1.17
1.0
1.55
1.33
38.57
1.34
0.0
0.0
Coupling
Weights
Weights
Coupling
Weights
Norm.
T= 0.1
1.23
1.16
1.5
3.98
8.76
0.54
1.19
1.5
-10
1.33
1.26
9.42
1.28
0.0
0.0
Coupling
Weights
Norm. Weights
2.0
Coupling
Weights
Norm. Weights
0.32
1.17
0.37
1.35
0.93
1.38
0.75
1.11
1.5
T=
1.0
1.0
-0.5
1.13
1.43
1.14
0.0
Coupling
Weights
Weights
Coupling
Weights
Norm. Weights
0.36
1.34
0.36
1.35
1.5
0.99
1.03
0.97
1.01
1.5
T= 10
1.0
1.29
1.04
1.02
0.0
0.0
Figure Discrete solution synthetic experiments in setting a. using unbalanced Sinkhom (UBOT), with different entropy regularization parameters and different penalization parameters T Ti (7). For each parameter setup, we show the coupling from source (red) target (blue), where the strength of each line is proportional to its value in the coupling matrix (left pane), the weights (middle pane) and the normalized weights (right pane). The weight per point is computed as the sum over the columns of the coupling matrix. We additionally show the average weight per cluster.
B.2 SINGLE-CELL PERTURBATION RESPONSES
In addition to the weighted MMD metric shown in Fig. 3, we evaluated our method on another distributional metric, the weighted Wasserstein distance between the predicted perturbed and observed perturbed cells. We compute it with the Sinkhorn algorithm, whereby for NUBOT and UBOT GAN, we pass the normalized predicted weights as source weights. Results are shown in Fig. 10.
As we lack ground truth for the correspondence of control and perturbed cells, we assess the biological meaningfulness our predictions by comparing the weights to CICasp3 and Ki67 intensity, the apoptosis and proliferation markers, respectively. Figures 11, 12 and 13 show UMAP projections computed on control cells for the drugs Trametinib, Ixazomib and Vindesine. and Figure 13 d.. regions of low predicted weights accurately correspond to regions of increased CICasp3 intensity. Additionally, we compare predicted weights between the two cell types, and contrast them with observed cell counts.
Figure 12 c., d.,
18

---

Under review as a conference paper at ICLR 2023
Setting b.
0.005
UBOT
0.05
Coupling
Weights
Norm. Weights
2.0
Coupling
Weights
Norm, Weights
2.0
1.5
1.5
458.93
0.0.01
0.5
0.5
5.33
0.0
Coupling
Weights
Weights
Coupling
Weights
Norm, Weights
*
1.5
1.5
0.64
0.75
0.57
0.68
11.67
14.55
0.49
0.62
0.05
1.0
110
1.9
1.7
43. 13.08
1.83
0.0
0.0
Coupling
Weights
Norm. Weights
2.0
Coupling
Weightss
Norm. Weights
1.5
0.1
3.57
4.36
0.95
0.68
1.0
1.65
1.62
0.5
11.01
1.71
00.
0.0
Coupling
Weights
Weights
Caupling
Weights
Norm. Weights
T=
0.34
0.48
0.37
0.52
1.0
1.5
0.96
0.99
0.8
0.83
1.5
1.
1.87
2.03
1.6
1.34
0.0
0.0
Coupling
Weights
Norm. Weights
2.0
Coupling
Weights
Norm. Weights
2.0
T=
-1.0
-0.5
0.5
1.05
Figure Discrete solution to synthetic experiments in setting using unbalanced Sinkhorn (UBOT), with different entropy regularization parameters and different penalization parameters
T2 (7). For each parameter setup, we show the coupling from source (red) to target (blue), where the strength of each line is proportional to its value in the coupling matrix (left pane), the weights (middle pane) and the normalized weights (right pane). The weight per point is computed as the sum over the columns of the coupling matrix. We additionally show the average weight per cluster.
C DATASETS
We evaluate NUBOT on several tasks including synthetic data as well as perturbation responses of single cells. both settings, we are provided with unpaired measures μ and y and aim to recover map T' which describes how source transforms into target v. While in the synthetic data setting provided with ground truth matching, this is not the case for the single-cell data as measuring a cell requires destroying it. In the following, describe generation and characteristics of both datasets, as well as introduce additional biological insights allowing us to shed light on the learned matching T.
C.I SYNTHETIC DATA
To evaluate NUBOT in a simple and low-dimensional setup with known ground-truth, we generate synthetic example: We model source population with clear subpopulation structure through mixture Gaussians. Next, we generate a second (target) population aligned to the source population. We then simulate intervention to which the subpopulations respond differently, including different levels of growth and death. Specifically, we generate batches of 400 samples with three clusters with different proportions before and after the intervention. Table 1 shows the proportions of the three clusters in the source and target distribution, as well as the required weight-facce and the obtained results from NUBOT and UBOT GAN.
19

---

Under review as a conference paper at ICLR 2023
NeMOT
CRLLOT
GAN
Ineveer
DISCRETE
GAUSSIAN
8h Timestep
Weighted Waywerteei Disaatoo
100
Imottnib
Erietteuh
Daastinn
Regoralenia
Panebinental
Trametinu
Fachtoonl
Caspatia
Lanalidamido
Tomamiamide
Destahbiee
Emnrllnus
Crisetinad
Melphalan
Itatrnul
Obapanzb
Vindoettoo
Berafonih
Staurepport trourspportee
Hnsrooyurea
Patlecieil
Ultxertinb
Dabratamu
Drug
Timestep
Dietanion
Weightted Waoweotein
Imoutnb
Brintiat
Milustaurr
Daviing
Reporafent
Panehitontall
Tramstinnl
Fachtaanl
Casplai
Aanatelomde
Temoniomide
Deatrahire
Eraroliuus
Crtentin
Dowmethreee
Molphalan
Isatomil
Olapard
Vitdintion
Soratonil
Decartantion
Skaurcspotter
Hndroryarea
Pallectota
Dabrafent
Drug
Figure 10: Distributional fit of the predicted perturbed cell states to the observed perturbed cell states for each drug and timestep, measured by the Wasserstein distance (2) on set of held-out control cells. Here, histogram in (2) of the predictions contains the learned weights (for those methods returning an unbalanced solution, i.e., NUBOT and UBOT GAN). For the other baselines u contains uniform weights.
MelA (Cell) Control
MelA (Cell) Predicted
Weights
2
Figure 11: UMAP projections computed on control cells for Trametinib at = 24h. High predicted weights in the MelA* + cell type suggest proliferation, while the + population shows higher levels of cell death. This prediction is confirmed by the relative cell counts, where MelA cell counts increase and Sox9+ counts decrease, demonstrating opposite response behaviors Trametinib for each subpopulation, i.c., MelA + cells show proliferation and S cells death.
20

---

Under review as a conference paper at ICLR 2023
Sox3 (Cell) 8 h
Sox9 (Cell) Predicted
ClCasp3 (Cell)
ClCesp3 (Cell) Predicted
in Control
Weights
Control 8 h
Weights
b.
MelA (Cell) Contral
MeLA (Cell) Predicted
d. ClCasp3 (Cell) Control
ClCaspá (Cell) Predicted
8
Weights
Weights
24
2
UMAP projections computed on control cells for Ixazomib for 8h, and 24h, colored by protein marker intensities a. MelA and b. Sox9, markers for the two subpopulations, as well CICasp3, marker for cell death, at 8h and d. 24h. The UMAPs confirm the measured relative cell counts of each subpopulation. After 8h a.-c, neither MelA + nor Sox9* cells are affected by the treatment, i.e., mainly predict weights around d. After 24h, we observe low weights in regions of high predicted apoptosis marker intensities (CICasp3), especially at t = 24h, where the observed cell counts suggest death predominantly the MelA cell cluster.
Figure 12:
(Cell)
ClCaspá (Cell) In Control
ClCasp3
h
Weights
Predicted
Weights
3
2
(Cell)
(Cell)
d. CICaap3 (Cell) in Control
CICasp3 (Cel)) in Predicted
24 h
Predicted
2
Weights
24
Weights
2
2
Figure 13: UMAP projections computed on control cells for Vindesine for 8h, and 24h. colored by protein marker intensities a. MelA and b. Sox9, markers for the two subpopulations, as well as CICasp3, a marker for cell death, 8h and d. 24h. The predicted weights (left) at e. and d. 24h match the observed effects each subpopulation, as initially only Sox9 cells are affected by treatment with Vindesine, and only after 24h MelA + cells show increased cell death.
21

---

Under review as a conference paper at ICLR 2023
Table 2: Overview of all treatments and their inhibition type considered in this work. Abbre- viations PROTi (Proteasome inhibitor), DNASYnthi (DNA synthesis inhibitor), panKi (pan kinase inhibitor), IomuneMod. (Immune modulatory compound), MTDisruptor (Microtubule disruptor), ApopInducer (Apoptosis inducer).
Drug Name
Inhibitor Type
Drug Name
Inhibitor Type
Ixazomib
PROTi
Olaparib
PARPi
Sorafenib
RAFI
Paclitaxel
MTDisruptor
Dabrafenib
BRAFI
Melphalan
Alkylator
Everolimus
mTORi
Regorafenib
panKi
Hydroxyurea
DNASynti
Vindesine
MTDisruptor
Midostaurin
panKi
Cisplatin
Alkalyting
Dexamethasone
ImmuneMod.
Ulixertniib
ERKi
Temozolomide
Alkylator
Staurosporine
ApopInducer
Decitabine
DNAMeti
Lenalidomide
ImmuneMod.
Dasatinib
SRCi-ABLi
Crizotinib
METi
Trametinib
MEKi
Imatinib
KITI-PDGFRi-ABLi
Erlotinib
EGFRi
Palbociclib
CDK4/6i
Dacarbazine
Alkylator
C.2 SINGLE- DATA
Biological experiment. The single-cell dataset used in this work was generated by the a multiplexed microscopy technology called Iterative Indirect Immunofluorescence Imaging (4i) (Gut et al., 2018), which is capable of measuring the abundance and localization of many proteins in cells. By iteratively adding, imaging and removing fluorescently tagged antibodies, a multitude of protein markers captured for each cell. Additionally, cellular and morphological characteristics are extracted from microscopical images, such as the cell and nucleus area and circularity. This spatially resolved phenotypic dataset is rich in molecular information and provides insights into heterogeneous responses of thousands of cells to various drugs. Mcasuring different morphological and signaling features captures pre-existing cell-to-cell variability which might influence perturbation effect, resulting in various different responses. Some of these markers of particular importance, they provide insights into the level of a cell's growth or death as well as subpopulation identity. We utilized mixture of two melanoma tumor cell lines (M130219 and M130429) at a ratio of 1:1. The cell lines can be differentiated by the mutually exclusive expression of marker proteins. The former is positive for Sox9, the latter for a of four proteins which are all recognized by and antibody called MelA (Raaijmakers et al., 2015). Cells were seeded in 384-well plate and incubated at 37C and 5% CO2 ovemiigtt. Next, the cells were exposed to multiple cancer drugs and Dimethyl sulfoxide (DMSO) as vehicle control for 8h and 24h after which the cells were fixed and six cycles of 4i were performed TissueMAPS and the scikit-image library (Van der Walt et al., 2014) were used to process and analyze the acquired images, perform feature extraction and quality control steps using semi-supervised random forest classifiers.
Data generation and processing. Our datasets contain high-dimensional single-cell data of control and drug-treated cells measured at two time points (8 and 24 hours). For both the 8h-dataset and the 24h-dataset, we normalized the extracted intensity and morphological features by dividing each feature by its 75th percentile, computed on the control cells. Additionally, values were transformed by a loglp function (r + log(æ + 1)). In total, our datasets consist of 48 features, of which 26 are protein marker intensities and the remaining 22 morphological features. For each treatment, have measured between 2000 and 3000 cells. For training the models, we perform a 80/20 train/test split. We trained all models on control and treated cells for each time step and each drug separately. The considered drugs as well as their inhibition type can be found in Table 2.
Cell type assignment. We assigned M130219 and MI M130429 cells to the Sox9 and MelA cell types, respectively, by first training two component Gaussian mixture model on the features 'intensity- MelA-mean' and intensity-nucled nuclei-Sox9 mean' of the control cells. Next, we used the aforementioned features and the labels provided by the mixture model to train a nearest neighbor classifier, which
22

---

Under review as a conference paper at ICLR 2023
dtip
skppr2
6:04
-6.3
Ino milphaxe
eug repeedees
drap Sytteey
dray villme
1.3
0.34
tng
3641
-8.3
dng ssstrapprtiv
Clug in decraber
drag - keakdeeise
Grag Gaviab
trap StoRtE
8.0
-8.5
kdeaaby
Figure 14: Classification of cells into cell types (MelA+, Sox9 based on protein marker intensities of MelA and Sox9, for all drugs, at t = 8h $ C.2. Each tile represents one drug. MelA cells colored in red, S in blue.
a. 8
Sox9 Cell Type
b. 24
Sox9 Cell Type
Cell Count
-2000
3000
MelA Coll Type
MelA Type
Coll Coll Connt
5030
2000
Ougten
Cdentin
Dabrofash
Decarbotte-
nocthaton
Pectoons
-3000
bechaMMes
snnting
Figure 15: Drug treatment-indde count of the respective cell types the control condition. a. Cell count change for cell types Sox9 + (top) and MelA (bottom) at t = 8h. b. Cell count change for cell types Sox9 (top) and MelA (bottom) at = 24h.
change in cell counts in the two cell types compared to the cell
23

---

Under review as a conference paper at ICLR 2023
Ine noueecl
Kng zlpprk
Ine paditca))
6.3
-3.0
axdptakn
Ing calootieee
drig diplasi
drip deanedaicee
4npp tezoodenall
6.5
allops
-1.0
siassepartve
3.5
ing tnuustil
dro InAtz42
drup siaah
drug in pabooidb
Grg decarbuie
8.3-
-30
ol Siitkrazas
Figure 16: Classification of cells into cell types (MelA+, Sox9-) based on protein marker intensities of MelA and SOX9, for all drugs, at t 24h $ C.2. MelA cells colored in red, Sox9+ in blue.
trapprist
1.0
X3
Cell coant
0.6
0.4
0.0
Cisplatio
Costrod
Caizeiaia
Dabrafeaib
Dacarbazine
Dasalaii
Dodtabine
Edottsib
Everolamss
Hydrosyurso
Imattelo
hazomis
Lenaldomdee
Malphalon
Midostaurta
Olapanio
Paditaxo)
Palrociell
Reporafenis
Sarafosib
Saurosporie
Temozolozide lonide
Tramntinl
Uixodiaii
Viadeesie
Figure 17: Observed cell counts of drug- cells normalized to control cell counts, per drug and time point. 8h treatment in light blue, 24h treatment in dark blue.
24

---

Under review as a conference paper at ICLR 2023
then used predict the cell type labels of the drug treated cells. The procedure performed separately for the 8h- and 24h dataset. Results of the classification can be found in Figure 14 and Figure 16 respectively.
D EXPERIMENTAL DETAILS
NUBOT consists of several modules and its performance is compared against several baselines. In the following, we provide additional background on experimental details, including description of the evaluation metrics and baselines considered, as well as further information the parameterization and hyperparameter choices made for NUBOT.
D.1 EVALUATION METRICS
We evaluate our model by analyzing the distributional similarity between the predicted and observed perturbed distribution. For this, compute the kernel maximum mean discrepancy (MMD) (Gretton al., 2012). We utilize the RBF kernel, and as is usually done, report the MMD as an average over several length scales, i.e., 0.5, 0,1, 0.01, 0.005. To take the mass variation into consideration, compute a weighted version of MMD, by weighting each predicted point by its associated normalized weight. Additionally, we compute the weighted Wasserstein distance between the predicted and observed perturbed cells (2).
D.2 BASELINES
We compare NUBOT against several baselines, comprising balanced OT based method (Bunne et al., 2021, CELLOT) and an unbalanced OT-based method (Yang & Uhler, 2019, NUBOT), i.e., current of-te-of- methods as well as ablations of our work. We further provide a comparison to the closed-for frrsed-formmo of eutropy-regularize optimal transport on unbalanced Gaussians (Janati et al., 2020b). In the following, we briefly motivate and introduce each baseline.
CELLOT. By introducing reweighting functions and Ç. NUBOT recovers balanced problem parameterized by dual potentials and g. An important ablation study to consider is thus to compare its performance its balanced counterpart. Ignoring the fact that the original problem includes cell death and growth, and thus varying cell numbers, we apply ideas developed in Makkuva et al. (2020); Bunne al. (2021) and learn balanced OT problem via duals and g. These duals are parameterized by two ICNNs and optimized in objective (5) via an alternating min-max scheme.
UBOT GAN. Using (6), Yang & Uhler (2019) propose to model mass variation in unbalanced OT via a relaxation of the marginals. Similar to Fan et al. (2021a), Yang & Uhler (2019) reformulate the constrained Monge problem (3) as saddle point problem with Lagrange multiplier h for the constraint Typ- i.e.,
T / c(x,T(x)) ((c,T(()))(2)dd / h(y) Tap) dy
sup inf h
f - h(T()))((()))) / h(u)v(uddy
parameterizing 7 and via neural networks. To allow mass to be created and destroyed, Yang Uhler (2019) introduce scaling factor ç X R allowing to scale mass of each source point 24. The optimal solution then nceds to balance the cost of mass and the cost of transport, potentially measured through different cost functions C₁ X y I (cost of mass transport) and 02 IR R+ (cost of mass variation). Parameterizing the transport map To. the scaling factors Eo, and the penalty h with neural networks, the resulting objective is
1 Σ (a Tolx:)) +o (E6(()))) + ( (T)))) (hwh()
1(0, 0, w) =
with V approximating the divergence term of the relaxed marginal constraints (see (6)), and is optimized via alternating gradient updates.
25

---

Under review as a conference paper at ICLR 2023
GAUSSIAN APPROX. Janati et al. (2020b) provide closed-fome solution of the entropy-regulariz regularized optimal transport problem on unbalanced Gaussians. They show that the unbalanced optimal transport plan is, minimizer of (7), is also a Gaussian distribution over R x R
In order to use this baseline on the single cell data, we first compute a Gaussian approximation the control and treated cells separately in the original cell data space. Then, we compute the closed-form joint coupling, æ = =N, Σ), where
fzz by
Σ Σ y/a
Exy Σ yyy
and
Σ
(16)
Given a test source sample Xt, we then compute the conditional expectation of the transported target y, L...,
Ex x = pi + Exy Dy (x+ -
(17)
DISCRETE OT. Additionally, we consider the entropy-regularized Wasserstein mapping returned by the Sinkhorn algorithm (Chizat et al., 2018a; Cuturi, 2013; Benamou et al., 2015) on finite sets. This algorithm does not return a parameterized solution, but rather a transport map between finite source and target sets. Thus, the solution is computed on the full dataset (inclduing train and test data) and cannot be considered in the out-of-sample setup. We have included a comparison anyway for completion.
IDENTITY. A trivial baseline is to compare the predictions to map which does not model any perturbation effect. The IDENTITY baseline thus models an identity map and provides an upper bound on the overall performance, also considered in Bunne et al. (2021).
OBSERVED. In a similar fashion we might ask for a lower bound on the performance. As a ground construct baseline for comparison on distributional level
truth matching is not available, by comprising different set of observed perturbed cells, which only vary from the true predictions up to experimental noise. The closer method can approach the OBSERVED baseline, the more accurate it fits the perturbed cell population.
D.3 HYPERPARAMETERS
We parameterize the duals f and y using ICNNs with 4 hidden layers, each of size 64, using ReLU as activation function between the layers. We choose the identity initialization scheme introduced by Bunne et al. (2022b) such that Vg and f resemble the identity function in the first training iteration. As suggested by Makkuva et al. (2020), relax the convexity constraint ICNN g and instead penalize its negative weights W?*
R (0) =
Σ ||max (-W7, WFEO
The convexity constraint on ICNN is enforced after each update by setting the negative weights of all W77 E O, to zero. Duals g and f are trained with an alternating min-max scheme where each model trained at the same frequency. Further, both reweighting functions 17 and Ç are represented by a multi-layer perceptron (MLP) with two hidden layers of size 64 for the single-cell and of size 32 for the synthetic dataset,, with ReLU activation functions. The final output further passed through softplus activation function we do not assume negative weights. For the unbalanced Sinkhorn algorithm, we choose an entropy regularization of 0.005 and a marginal relaxation penalty of 0.05. We use both Adam for pairs g and ff as well 77 and with learning rate 10-4 and 10 as well as 31 = 0.5 and 32- = 0.9, respectively. We parameterize both baselines with networks of similar size and follow the implementation proposed by Yang & Uhler (2019) and Bunne al. (2021).
E REPRODUCIBILITY
The code will be made public upon publication of this work.
26


Reviews:
[{"review_id": "X85itBrjbQj", "reviewer": "Reviewer_CvpT", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper extends CellOT to the case of unbalanced optimal transport, which is essential to modelling most cell systems. This extension is accomplished through the formalism of semi-couplings, learning reweighted distributions which are then transported with balanced transport, in this case using the ICNN framework. A novel algorithm is presented to solve this problem efficiently through gradient descent. This algorithm is evaluated on a toy dataset and a single-cell drug perturbation dataset captured with 4i imaging. NubOT fits the data better in that it is able to predict cell states for unseen cells (but in seen conditions) better than comparable methods.", "summary_of_the_review": "Overall a useful improvement concept with a clear and specific application. There are a few clarifications on the algorithm and the experiments. If these things are cleared up I have no problem raising my score. ", "strength_and_weaknesses": "Strengths:\n\n- Important extension to the CellOT formulation which accounts for relative size changes in cell populations, which is extremely important in heterogenous cell systems.\n- Well grounded theoretical formulation of unbalanced OT with semi-couplings.\n- Good dataset choice for the evaluation of the unbalanced transport with a ground truth measurement of the number of cells at a population level before and after treatment.\n\nWeaknesses:\n\n- The connection between Algorithm 1 and the presented theorems is not entirely clear. From Algorithm 1 I would have expected the transformation for new samples to be $(x, u) \\to (\\nabla g(x), \\eta(x) \\cdot u)$. Could the authors clear up what is correct, algorithm 1 or figure 1?\n- To me “forecasting heterogenous responses of multiple cancer cell lines to various drugs” would imply extrapolation to unseen drugs or cell line or at least combinations, instead of the forecasting the response of an unobserved cell in a seen drug and cancer cell line as is done here. Especially considering a separate model was trained for each treatment. Would it be possible to tone down the language on applications? In particular the \"This is an unprecedented achievement in the field of single-cell biology\" line? Modelling in distribution responses with proliferation has been tackled in prior work for time series, applying it to a seen perturbation response is in my view a relatively straight forward extension.\n- The application is extremely specific, it is unclear how interesting it is to a wider ICLR audience. It would be useful to discuss other applications for unbalanced OT such as domain alignment, transfer learning etc.\n\nComments / Questions:\n\n- In the “Updating Mappings” paragraph: “Since these are guaranteed to be balanced due to the argument above”. Could the authors clarify this please? Why is this guaranteed? How is guaranteed that $\\int \\tilde{\\mu} = \\int \\tilde{\\nu}$?\n- There is something that is not lining up for me between the optimization in algorithm 1 and how samples are transformed. To me it seems like you would want to weight $\\hat{y}$ by\n- With a batchsize less than the entire data this is similar to a mini-batch OT problem. The unbalanced minibatch OT problem is in some senses creates a similar (full) map to the (full) balanced OT problem (Fatras et al. ICML 2021). Further, it doesn’t have the same weights as the regular non-batched problem. Does batch size have any affect on your results? I think this should at least be mentioned as not solving the same problem as UBOT.\n- I’m curious how the learned weights compares to a discrete UBOT algorithm. I understand that there are advantages of learning a continuous map, however, I would still like to see a comparison to understand how accurately the proposed NubOT algorithm solves UBOT.\n- I don’t believe that there are any constraints on the total sum of the transported measure $Z = \\int_{(x,u)} \\eta(x) \\cdot u \\cdot \\zeta( \\nabla(g(x))^{-1})$. The weighted MMD is invariant to Z, thus there is no evaluation whether the total measure is accurate. Is this total mass at all important to your application?\n- Is there (albeit small) data leakage from data normalization in the single cell data to the 75% level?\n- How was the relaxation penalty chosen, and do the authors expect it to be tuned on new datasets? What was the batch size chosen?\n\nMinor Comments:\n\n- Line 7 of Algorithm 1 should have an $x$ instead of $y$. Line 1 might want to be batches instead of epochs?\n- I would note that $L^1$ norm unbalanced OT is not exactly the same as partial OT (As stated in A.1). Partial OT is slightly different. Instead this $L^1$ norm is sometimes “Robust optimal transport” (Mukherjee et al. ICML 2020).\n- Could the true weights in Table 1 be depicted in Figure 2? Its difficult to understand that these are the correct weights without scrolling to Table 1.\n- What is the weighted kernel MMD? What is the kernel? Gaussian with what bandwidth?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "QQiwxuI92X", "reviewer": "Reviewer_irGq", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes a method for finding optimal transport maps between unbalanced distributions, i.e., distributions with different total masses. The method is composed of two steps repeated iteratively. First, it seeks for scaling factors for source and target measures, such that the re-scaled measures are balanced. This problem is considered as a discrete unbalanced optimal transport problem and solved using the Sinkhorn algorithm. These factors are used to train neural networks to predict scaling factors for arbitrary points. At the second step,  the method finds OT maps between re-scaled measures using standard balanced OT formulation (Makkuva et al., 2020). \n\nThe method is tested on the synthetic problem of finding a map between unbalanced mixtures of Gaussians, and the real-world problem of testing cells' response to different drugs which may induce proliferation (quantitative growth) or death of cells, i.e., change in distribution mass. ", "summary_of_the_review": "The authors propose a method for finding an optimal transport map between unbalanced distributions. However, from the construction of the method it is unclear how well it approximates actual scaling factors and OT map between the distributions. The experimental section also does not clarify this issue. Due to this, I think that not enough evidence is given to convince the reader that the method computes what it is intended to do. Moreover, the entire focus of the method is a particular biological task, which is hard to follow for an unprepared reader. Because of all these issues, I am currently more on the negative side about the paper. If the authors solve these evaluation/clarity issues during the rebuttal phase, I will be happy to reconsider my assessment.", "strength_and_weaknesses": "**Strengths.** The paper proposes an optimal transport approach for an important unbalanced setup, which is not yet fully explored.\n\n**Weaknesses.**\n1. The proposed algorithm is based on the solutions of the discrete unbalanced OT problem. Thus, it is unclear to what extent the method approximates the actual underlying solution of the continuous problem. This might lead to bias in method solutions.\n2. The proposed approach lacks comprehensive quantitative evaluation. It is not shown that it learns actual scaling factors and OT map between distributions.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Ox4bbIBa14g", "reviewer": "Reviewer_rxRW", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "Optimal transport distances (a.k.a. Wasserstein distance) have recently drawn ample attention in statistics and machine learning communities as powerful discrepancy measures for probability distributions. The standard formulation of OT assumes equality of mass to be transported from the source to the target, which is referred as to the conversation mass constraint. In other applications, this constraint is not filled; hence application of standard OT is not done directly. Existing lines of research investigate this problem by handling the non-equality of mass transportation. They consist of unbalanced Optimal transport (UBOT). These approaches are given by the sum of an OT cost and two constraints involving the non-equality of mass transportation. The second term is often written using divergence, for example, a Kullback-Leibler divergence, etc.\n\nThis paper proposes NUBOT, a neural unbalanced OT that relies on the formalism of semi-couplings to account for creation and destruction of mass. NUBOT takes advantage of the neural optimal transport, where it was shown that the Monge plan can be learned via the Kantorovich potentials of the dual problem, which can be parametrized via input convex neural networks (ICNNs). Towards this end, the authors suggest learning a proxy measures $\\tilde{\\mu} = \\eta \\cdot \\mu$ and  $\\tilde{\\nu} = \\zeta \\cdot \\nu$ to rescale the mass. The functions $\\eta, \\zeta$ are also parametrized via NN. An algorithm of NUBOT is given with an extensive numerical study with an accent on cell-biology applications.", "summary_of_the_review": "\n\n### Typos ###\n- Page 3: Eq. (3): there is no dependency on $\\varepsilon$.\n- Page 3: Eq. (5): OT$_$n. I don't understand this dependence on n.\n- Page 4: The notation of the set of semi-coupling between measures $Gamma(\\mu, \\nu)$ is the same as the set of couplings in the nonnegative Radom measures (page 2). \n- Page 4: \n$\\Gamma \\gets \\text{UBOT}(\\boldsymbol{u}, T(\\boldsymbol{x}_i), \\boldsymbol{v}, \\boldsymbol{y}_j)$ --> $\\Gamma \\gets \\arg\\min\\text{UBOT}(\\boldsymbol{e}\\odot\\boldsymbol{u}, T(\\boldsymbol{x}_i), \\boldsymbol{v}, \\boldsymbol{y}_j)$\n- Page 4: \n$\\Gamma \\gets \\text{UBOT}(\\boldsymbol{v}, S(\\boldsymbol{y}_j), \\boldsymbol{u}, \\boldsymbol{x}_i)$ --> $\\Gamma \\gets \\text{UBOT}(\\boldsymbol{z}\\odot\\boldsymbol{v}, S(\\boldsymbol{y}_j), \\boldsymbol{u}, \\boldsymbol{x}_i)$\n- Page 6: \"intervention respnses\" --> \"intervention responses\"", "strength_and_weaknesses": "### Strength ###\n- Formulation UBOT with a neural approach based on semi-couplings proxy measures. The Kantorovich potentials are learned with ICNNs. \n- Extensive numerical study on synthetic and real data with an accent on cell-biology applications: prediction of cell proliferation and death, and the perturbation responses on the level of single cells.\n\n### Weaknesses ###\n- Steps 5 and 7 in Algorithm 1 deal with calculating an unbalanced Sinkhorn: it is a little be confused since the ultimate goal of this paper is to tackle UBOT. \n- The paper lacks comparisons with the SOTA approaches of UBOT: NUBOT is only compared to ubOT GAN? One can compare it to a vanilla generalized Sinkhorn algorithm (Chizat et al. 2018). \n- What is the computational complexity of NUBOT. I think this point should be clarified or at least the authors can give some insight bout it. Steps 5 and 7 in Algorithm 1 might be the most expensive ones.\n- In Algorithm 1, it is not clear to me the update of $\\eta$ and $\\zeta$.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "G9DttlRkZo", "reviewer": "Reviewer_zLtd", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "First, I am an emergent reviewer and not an expert on OT. Therefore, I'm not sure whether the literature is reviewed properly and whether the evaluation setup is fine.\n\nThis work presents a formulation of the unbalanced optimal transport problem. To solve the unbalance problem, it relies on the formalism of semi-couplings to account for the creation and destruction of mass. \n\nIt further presents a cycle-consistent training procedure upon neural OT to estimate such semi-couplings and generalize out-of-sample, named NUBOT. \n\nThe method is evaluated on a synthetic dataset and the single-cell perturbation task. It is compared to CellOT (a balanced neural OT method) and UBOT GAN (an imbalanced OT approach).  \n\nAccording to the reweighted MMD (and other metrics), NUBOT outperforms the baselines on both datasets. Moreover, on the challenging single-cell perturbation task, NUBOT is able to successfully predict perturbed cell states, while explicitly modeling death and proliferation. ", "summary_of_the_review": "This is a solid paper on addressing the unbalanced OT problem and I tend to accept it now.", "strength_and_weaknesses": "**Strength**\n\n1. The paper is well-organized and clearly written. \n\n2. The empirical results (both qualitative and quantitative ones) seem pretty good.\n\n**Question**\n\nI have several questions here.\n\n1. Why use MMD instead of the OT distance as the evaluation metric in Fig. 4? It is more natural to evaluate the objective on the validation dataset for machine learning algorithms.\n\n2. The method presented in Section 3 seems complex to me. It would be better to discuss the necessity for each part in depth and present the empirical contribution of each part in the experiments.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject