GENOT: Entropic (Gromov) Wasserstein Flow Matching with Applications to Single-Cell Genomics
Authors: Dominik Klein
Théo Uscidda
Fabian J Theis
marco cuturi
Venue: NeurIPS 2024 poster (2024)
Keywords: single-cell genomics
optimal transport
flow matching
Primary area: machine_learning_for_healthcare
TL;DR: Neural Optimal Transport Solvers based on Flow Matching for Gromov and Wasserstein OT motivated by tasks in single-cell genomics

Abstract: Single-cell genomics has significantly advanced our understanding of cellular behavior, catalyzing innovations in treatments and precision medicine. However,
single-cell sequencing technologies are inherently destructive and can only measure a limited array of data modalities simultaneously. This limitation underscores
the need for new methods capable of realigning cells. Optimal transport (OT)
has emerged as a potent solution, but traditional discrete solvers are hampered by
scalability, privacy, and out-of-sample estimation issues. These challenges have
spurred the development of neural network-based solvers, known as neural OT
solvers, that parameterize OT maps. Yet, these models often lack the flexibility
needed for broader life science applications. To address these deficiencies, our
approach learns stochastic maps (i.e. transport plans), allows for any cost function,
relaxes mass conservation constraints and integrates quadratic solvers to tackle the
complex challenges posed by the (Fused) Gromov-Wasserstein problem. Utilizing
flow matching as a backbone, our method offers a flexible and effective framework.
We demonstrate its versatility and robustness through applications in cell development studies, cellular drug response modeling, and cross-modality cell translation,
illustrating significant potential for enhancing therapeutic strategies.

GENOT: Entropic (Gromov) Wasserstein Flow
Matching with Applications to Single-Cell Genomics
Dominik Klein*t Helmholtz Munich
Théo Uscidda' CREST-ENSAE theo. www.ddaeensse........ fr
dominik. hww.nn/w..eiinennllllleze-----ee
Fabian Theis
Marco Cuturi Apple cuturiSapple. com
Helmholtz Munich
fabian. theishhelmmolleee munihhhhmmuchhhhh
Abstract
Single-cell genomics has significantly advanced our understanding of cellular be- havior, catalyzing innovations in treatments and precision medicine. However, single-cell sequencing technologies are inherently destructive and can only mea- sure limited array of data modalities simultaneously. This limitation underscores the need for new methods capable of realigning cells. Optimal transport (OT) has emerged potent solution, but traditional discrete solvers hampered by scalability, privacy, and out-of out-of-sample sampl estimation issues. These challenges have spurred the development of neural network- based solvers, known as neural OT solvers, that parameterize OT maps. Yet, these models often lack the flexibility needed for broader life science applications. To address these deficiencies, our approach learns stochastic maps (i.e. transport plans), allows for any cost function, relaxes mass conservation constraints and integrates quadratic solvers tackle the complex challenges posed by the (Fused) Ctroovv-Wasseree flow matching as a backbone, our method offers a ficxible and effective framework. We demonstrate its versatility and robustness through applications in cell develop- ment studies, cellular drug response modeling, and cross-modality cell translation, illustrating significant potential for enhancing therapeutic strategies.
problem. Utilizing
Introduction
Discrete Optimal Transport in Single-Cell Genomics. Due to the destructive nature of single-cell sequencing technologies, it is fundamental to realign distributions of sequenced cells. Discrete optimal transport (OT) has proven useful in this task. For example, it is standard practice to leverage the framework of OT to reconstruct cellular trajectories across developmental time points 772 and to study the response of cells to external perturbations like drugs or gene knockouts [18] In these settings, both distributions, i.e. earlier and later time point and before perturbation and after perturbation, respectively, live in the same space, setting which we refer to as linear OT. In contrast, most applications of OT in single-cell genomics require maps across (partially) incomparable spaces, the setting which we refer to as quadratic OT. For example, analysis of spatial transcriptomies datasets requires the quadratic formulation warped and rotated, thus not living in the same coordinate space [60]. Similarly, adding spatial information as a prior has been shown to improve the performance modeling the trajectory of cells in the discrete setting [39]. Analogously, adding information from lineage tracing readouts (89
OT as the coordinate systems of two different slices
Equual contribution
Work done during internship at Apple
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

recovers the evolution of cells more faithfully 42 Another prevalent task in single-cell genomics which is commonly approached with quadratic OT solvers is the alignment of cells across modalities 17 Most single-cell technologies capture only one modality, which is insufficient to obtain comprehensive representation of the cellular state. As single-cell datasets grow larger [65], traditional OT solvers become less applicable to tasks in single-cell genomics due to their high computational complexity. Thanks to recent advancements in low-rank OT [69 70 71, the aforementioned OT- based algorithms are more accessible to single-cell biologists 39 Yet, the non-parametric nature of discrete OT solvers entails issues with respect to data privacy and prevents their application to large scale single-cell atlases capturing millions of cells (65. Moreover, -ut-lessamples estimation is limited to very specific scenarios [63]
Current Limitations of Neural Optimal Transport in Single-Cell Genomics. To overcome these limitations, neural OT solvers have been leveraged to study cellular perturbations 7 86 and to model cellular trajectories [21 83 Yet, these methods estimate Monge maps, i.e., deterministic maps, contradicting the assumption that cells evolve stochastically [9]. Stochastic formulations are also favorable they can produce a conditional distribution that can be used to quantify uncertainty [25]. In the discrete setting, stochastic maps can be obtained from entropy-regularized OT (EOT) [I3. Recently, number of works have addressed learning EOT plans (29 76 84 85] in the neural setting. Yet, all of these methods are limited in the choice of the cost function, with most of them being restricted to the squared Euclidean cost. Single-cell genomics data is known to be non- non-Euclidean [57] and thus requires a flexible choice of the cost function in OT applications as demonstrated in the discrete OT setting in Demetci et al. [17]. Huguet et al. [35], Klein al. [39]. The third requirement for applying OT to single-cell genomics is the option to lift the mass conservation constraint, allowing for unbalanced optimal transport [11 24 74.. Unbalanced OT is crucial to model cellular growth and death as well as for automatically discarding outliers, which are prevalent in highly noisy measurements in single-cell genomics. Eyring et al. (20), Lübeec al. (50), Yang and Uhler (94) proposed ways to incorporate unbalancedness into deterministic linear OT maps, while unbalanced formulations for entropic neural OT have barely been explored. The most severe shortcoming of the existing plethora of neural OT estimators is their limitation to the linear OT scenario, i.e., learning maps within spaces. Yet, most applications in single-cell genomics require Wassersttnnn (FGW) [88] formulation. To
the Gromov- Wasserstenn (GW) [54] or the fused the best of our knowledge, the only neural formulation for GW proposed thus far learns deterministic. balanced maps for the inner product costs, using min- min-max-min optimization procedure, severely limiting its applications [58]. Hence, we arrive four necessities that need to be fulfilled and have to be flexibly combinable to make neural OT generally applicable to problems in single-cell genomics:
N1: modeling the evolution of cells stochastically as opposed deterministically,
N2: flexibly choosing cost functions due to non-Euclidean geometry of single-cell genomics data,
N3: allowing for unbalanced OT to model cellular growth and death or accounting for outliers,
N4: mapping across completely or partially incomparable spaces.
Contributions. Hence, propose GENOT (Generative Entropic Neural Optimal Transport),
powerful and flexible neural OT framework that satisfies all of the above requirements:
GENOT is the first method that parameterizes linear and quadratic EOT couplings for any cost by modeling their conditional distributions, using flow matching 45] as a backbone.
We extend GENOT to the unbalanced setting, resulting in U-GENOT, to flexibly allow for mass variations in both the linear and the quadratic setting.
We extend (U-)GENOT to address the fused problem and this way propose, to the best of our knowledge, the first neural OT solver for the FGW problem.
We showcase GENOT's s ability to handle common challenges in single-cell biology: we (i) quantify lineage branching events in the developing mouse pancreas, (ii) predict cellular responses to drug perturbations, along with a well-calibrated uncertainty estimation, and (iii) translate ATAC-seq to RNA-seq with GW and introduce a novel method to perform this translation task with FGW.
2 Background
Notations. Let X' C TRP, ) C IR2 compact sets, referred to as the source and the target domain, respectively. In general, P f q. The set of positive (resp. probability) measures on X' is denoted by MT)) (resp. Mi (X)). For e M (X x V), T == fT and T2 p2tT denote its marginals.
2

---

Then, for M (2), E ()), (( = f T == = p. x2 v}. du denotes the relative
de
density of μ w.r.t V, s.t. μ # v. For p.Y (1), - Jxx log( dp- dr+ fx dp.
Linear Entropic OT. Let X x y + R be continuous cost function, E and 2 O. The linear cetropy-reuulrizedd OT problem reads
e MT
min =ET)(a)))
Jxxy edw + (KL(
(LEOT)
solution = of (LEOT) always exists and is unique when E > (LETT) is also known as the static Schrödinger bridge (SB) problem [52). With E = O, we recover the Kantorovich (38] problem. For discrete ll and v, we can solve (LEOT) with the Sinkhorn algorithm [13, whose complexity for n points is (n²) in time, and (n) or On²)
(n²) in memory, depending c.
Quadratic Entropic OT. As opposed considering an inter-domain cost defined on x), quadratic entropic OT is concerned with seeking couplings that minimize the distortion of the geometries induced by intra-domain cost functions cx : X x X + R and ey
V x y - R:
min xell(u,e)
Dex.es d( @ T) + (L
(QEOT)
where Dex.cy (x,Y, x', y') == (xx(x, cyly."'" quantifies the pointwise cost distortion. solution .. to Prob. (QOT)) always exists (see B.I. With E = 0, we recover the Gromov- Wasserstein (54) problem, which is the standard OT formulation for transporting measures supported on incomparable spaces. In addition to statistical benefits (95), using & > 0 also offers computational benefits, since for discrete y and v, can solve (QEOT) with mirror descent scheme iterating the Sinkhorn algorithm [62]. For measures on n points, its time complexity is O(n²) or O(n²),
depending on ex and cy 770 Alg. 1&2], while its memory complexity always O(n²).
Unbalanced Extensions. The EOT formulations presented above can only handle measures with the same total mass. Unbalanced optimal transport (UOT) [11 [44 lifts this constraint by penalizing the deviation of p1 प to j and Patı to v with a divergence. Using the KL and introducing weightings X1 A2 0 the unbalanced extension of (LEOT) reads
Sasiis//n/ J edw + ( @ v) + A1KL(TT + X2KL(Tz)V).
min xEM+(Xx>))
(ULEOT)
This problem can be solved efficiently in the discrete setting using a variant of the Sinkhorn algorithm 24 80). Analogously, quadratic OT also admits an unbalanced generalization, which reads
min / TEM-(XXX)) (xxy)²
Dexiey d(T TT) + <KL @ (T| KL @ (T (m) + A (T2 (m))
(UQEOT)
where KL (phr) = KL 2 p 8 7). A solution -े. to Prob. (UEOTT always exists (see B.
In the discrete setting, Prob. (UQEOT) can be solved using an extension of Peyré et al. scheme introduced by Séjourné et al. 81 Each solver has the same time and memory complexity as its balanced counterpart. For both (ULEOT) and (UQOTT) instead of selecting A; we introduce
x s.t. we recover the marginal constraint for 7, when A + +oo. We write (71 72). T3 A
Flow Matching (FM). Given a prior Po E M (Rd) and a time-dependent vector field (vx (b)))e(o.1]- one can define probability path pt starting from Po using the flow o, solving
d
(x( = (z)( de
Φo(z) = z,
by setting pt drtPo. We then say that v; generates the path pt- Continuous Normalizing Flows 99 (CNFs) model 2t,0 with a neural network which is trained to match a terminal condition PI 01 MT (Rd). FM 45) is simulation-free technique to train CNFs by constructing individual paths between samples, and minimizing
E, Z..00o0.Zspp.lll e ([Zo, Zıkr - (Zi - 10)
(2)
where [Zo. Zil. := (1- -) + tZ₁ with t ~ ((0, 11). If this loss is 0, the flow maps Pu to P1. i.e.
$1 P1- This property is often referred to FM preserving the marginal distribution.
3

---

ATAC rep. of cells
Conditional x, N1, Us)
Source f
RNA rep. of cells @ RR9
Conditioning on X1y U1
Cell features R
11
Flow (r(1(x1111)) 11)
(y vi][( un) Samples
(Shared) features R
1 C
space X' C R
R2
RNA space y C R?
Generate RNA & features
c R
x; (+10,,
(y,, v.)
x
Zi
Y1
Noise p
Target
What do do? One task consider is generating RNA cell profiles from ATAC measurements and an additional cell feature. This is explained in $5.22 and demonstrated in Fig. 4 As the cells live on manifolds in two (partially) incomparable spaces, we rely on the Fused Gromov- -asoo---sssrttiin (FGW) formulation, as described in $ 3.3 Here, the incomparable structural information is contained in the ATAC and the RNA measurements, while the comparable information are the cell features. Right: How do we do it? For each (x, u) in the support of the source μ., we learn a flow Φ1 [x, u) from the noise p to the conditional .² (-([|x, u). whose support lies in that of the target v. The flow is multi-modal: It allows sampling structural informations y, as well as features simultaneously. We highlight this procedure for specific pair (x1, U1), with p 2 and q =
Figure 1: Left:
Generative Entropic Neural OT
We introduce GENOT, a method to learn EOT couplings (thus satisyfing NI) with any cost (N2) by learning their conditional distributions. In $3.1 we focus on balanced OT, and show that GENOT approximates linear or quadratic EOT couplings (N4), solutions to (LFOT) or (QEOT) respectively. Then, in $ 3.2 we extend GENOT to the unbalanced setting by loosening the conservation of mass constraint (N3). We define U-GENOT, which approximates solutions to (ULEOT) and (UQEOT) In 3.3 we highlight that GENOT also adresses a fused problem, combining (LEOT) and (QEOT) (N4). Finally, we demonstrate in $ the GENOT algorithm, that allows to flexibly adapt to different OT formulations, which is key for easy usability for single-cell genomics analysts.
3.1 Learning Entropic Couplings with GENOT
Let μ MM (X), M (1) and m3 be an EOT coupling between y and v, which can be solution of problem LEOT) or QQOT.. By the measure disintegration theorem, we get
drr yx du(x) dx (ylx).
Knowing μ4, we can hence fully describe =5 via the conditional distributions (x (-1x))xxxx The latter are of great practical interest, they provide way to transport point sampled from μ to the target domain ): either stochastically by sampling yI, Yn from के (-Ix), or deterministically by averaging over conditional samples: (x(x) := By [Y]. Moreover, we can assess the uncertainty of these predictions using any statistic of a (-1x), which is crucial in single-cell genomics.
Learning the Conditional Distributions. Let p = (0, 4g) the standard Gaussian on the target space R 3 y. From the noise outsourcing lemma (37), there exists a collection of conditional generators (T* (-xx)1 s.t. T (-Ix) R Rº and for each x in the support μ., x (-1x)
[+ This means that if Z ~ p. then Y T (Z x² (). Here, we seek to learn neural collection (To (-o(1xx) of such conditional generators, fitting the constraint (- n² (-Ix), for any x in
the support of μ. We employ the FM framework (see $2) and parameterize each To(- x) implicitly, as
the flow induced by a neural vector field ve v-,x(-1x) I89 RR9 Namely, T-(- Φi (- [x) R + R
where di (+x)) solves
dt (z (z)(((() do(z|x) z.
(3)

---

We stress that while x E X' C RR, the flow Φ1 (-1xx) from to x (-|x)) is defined on the target space F' V. Hence, we can map samples within the same space when p = q. but also across incomparable spaces when p f q. In particular, this allows us to consider the quadratic OT problem (QEOT). Thus, for each x, we optimize or,0 (-I) by minimizing the FM loss (2) to map p to T. (-1x), i.e.
E, -= lller,o ())x)) (Y - 2)
where [2,Y], (1 z)z + tY interpolates between noise and conditional vectors. Averaging for all the support of ir and using Fubini's Theorem, we derive the GENOT loss:
LGENOT (0) == (x,,=-(((X,))===]]bb Ill Z]|x) - (Y 2)1
(5)
In practice, we optimize a sample-based GENOT loss by estimating T. Dis P from XX, ~ j1 and Y,....,Y, Y,, ~ v, with a discrete EOT solver, see Alg. [
GENOT Captures the Signal from the Mini-Batch Couplings. Standard FM only preserves straight couplings [47] while bridge matching (its stochastic counterpart) only preserves the linear EOT coupling for e(x. y) = ||x y 122 [5]. In contrast, GENOT is conditional FM model, learning the conditional distributions of the coupling independently of each other: For each x, we leverage FM to learn a conditional flow d1 (- (x)x) that maps P to x (-Ix). Therefore, as FM preserves the marginal distributions (see 2 GENOT preserves all couplings. Thus, it always captures the signal carried out by the mini batch couplings, in both linear and quadratic settings, and for any cost. We highlight this property in Fig. 5 comparing GENOT to other OT- FM approaches 664
85].
GENOT Handles Any Cost. We can use GENOT to approximate linear or quadratic EOT couplings. In both cases, we make no assumptions on the cost functions: We only need evaluate them on samples to estimate म. using discrete solver, see GENOT Alg. []] can use implicitly defined costs, whose evaluation requires non-differentiable sub-routine. For instance, we can use the geodesic distance on the data manifold. which can be approximated from the shortest path distance on the k-nn graph induced by the Euclidean distance [22 77). This approach has demonstrated its effectiveness in a range of single-cell genomics tasks relying on discrete OT
line 5 In particular, we
17 35 39
We use with GENOT for both linear (see $ $5.1) and quadratic OT (see $5.22..
GENOT Approximates Conditional Densities. For each x, we build a CNF pe(- = de( {x) between Po = P and pi (-1x). From the instantaneous change of variables formula (9], we can then approximate the conditional density n. (y |x) at an arbitrary point y E ), proving useful to evaluate the likelihood of one-to-one matches between cells, see Fig 23
3.2 U-GENOT: Extension to the Unbalanced Setting
Re-Balancing the UOT Problems. In its standard form, GENOT respects marginal constraints, so cannot directly handle unbalanced formulations ULLET or (UQEOT) We show that unbalanced
EOT problems can be re-balanced. Eyring et 200 Lübeck et al. [00] Yang and Uhler [ introduced previously these ideas in the Monge map estimation setting, namely, in a static and deterministic setup. Our method stems from the fact that, for both linear and quadratic OT, the unbalanced EOT coupling "²,r between μ E M MM (X), E ()²) solves balanced EOT problem between its marginals, which are re- weighted versions jJ and i of ju and v, that have the same mass.
Proposition 3.1 (Re-Balancing the unbalanced problems.). Let गË.T be an unbalanced EOT coupling, solution of (ULEOT) or (UQEOT) between μ E -M and v E M ()). We note jı p1 ! and i = p its marginals. Then, in both cases, ia (resp. 8) has density w.r.t. μ (resp. v). That is, there exist two re-weighting functions, one on each space, 1 X + R+ and & : y - R s.I. ju 6== and i = E. v. Furthermore, ja and i have the same lotal mass and 1. (Linear) nビ, solves the balanced problem (EOTT) between it and i with the same e.
2. (Quadratic) Provided that cx, Cy (or -cx, -cy) are CPD kernels (see Def. B.2) Te. solves the balanced problem (QEOT) between jı and v with e' m( where m (T,))
(X x ))
Learning the Coupling and the Re- Re-Weightings Simultaneously. Thanks to Prop. 3.1 we aim to [i] learn a balanced HOT coupling between j and along with [ii] the seigeigtiigg functions m. E. Learning them is desirable since they model the creation and destruction of mass. We do both
5

---

we use simultaneously by adapting the GENOT procedure. Formally, we seek to optimize:
LU.GENOT (0) = -Fr,,=>,.XX.Y)00r Ille, ((Z, M](x) - (Y - [i]
+ +Ex~p no) By (X)
[ii
where Ti0 Eo are (non-negative) neural re-weighting functions. Crucially, similar (balanced) GENOT, we only need to estimate a discrete unbalanced EOT coupling ",r using samples XI,...,X, μ and Yi.... Yn v to compute the two components, [i] and [ii], of the U- GENOT loss. We build upon theoretical insights on the linear OT case and extend them to the quadratic OT case in practice.
Proposition 3.2 (Pointwise estimation of re-weighting functions.). Let T = Eis Pld d (X.)))) solution to (ULOTT) between empirical counterparts of ş and v. Let a Pe,rin and P 1n its marginals weights. Then, almost surely, na,; -> (X,) and n b, + E(V.).
Using Prop. 3.1. fg.r is a balanced EOT coupling between its marginals, which are empirical approximations j and iz. Hence, we estimate the first term [i] of the loss as we do in the balanced case by sampling from the discrete coupling. Furthermore, Prop 3.2 highlights that the estimation of # also provides a consistent estimate of the re-weighting function evaluations at each X; and y, This enables estimating the second term [ii]. Therefore, switching from GENOT to U-GENOT simply involves using an unbalanced solver instead of a balanced one, and regressing the neural re- re-weighting functions on the marginal weights of the estimated discrete coupling. We detail our showing the additional steps w.r.t. GENOT in teal. procedure in Alg. |
3.3 Combining Linear and Quadratic OT
We show in 3.1 and $ 3.2 how to use GENOT to solve OT problems within the same space across incomparable spaces. On the other hand, numerous real-world problems pose the challenge of the source and target domains being only partially incomparuble [8] Therefore. suppose that the source and target space can be decomposed as X' = Ω x X and y = Ω x D, respectively. Intuitively, a sample (u,x) Ω x X' can be interpreted as a structural information x equipped with feature u. Assume we are given cost e fl - R to compare features, along with the intra-domain costs cx. Cg. The entropic Fused Gromov- Wassersttinn (FGW) problem reads
I J Dc cg
min KEIl(A)))
d(T TT) EKL(T)
EFGW
((xx))
where De De..og ((u, x), (v,y), = (1-a) c(u, v) + (cx y' and a E [o, 1]. This loss combines the pointwise structural distortion and the feature information. We refer to the additional cost on the feature information as the fused term. The weighting a allows us to interpolate between purely linear OT on the feature (a 0), and purely quadratic OT (a 1) on the structural information. Problem (EFGW) also admits an unbalanced extension, derived similarly as (QEEOT) with the quadratic KL 1221. An (un)balanced fused EOT coupling always exists (see B.1 minimizes distortion along the structural information and displacement cost along the features.
(U-)G for the Fused Setting. Whether in the balanced or unbalanced setting, we can use our method to learn a specific coupling as soon as it can be estimated from samples. We stress that the discrete solvers introduced by Peyré et al. [62] and Séjourné et al. [8]] we use for (QEOT) and (UQEOT), respectively, are still applicable the fused setting. As result, we can approximate solutions of (FFGW) and its unbalanced counterpart with (U-)GENOT: To illustrate the learning outcome, take a solution 5 of (FFGW). Learning 5 with our method amounts to training vector fields that are conditioned pairs of modalities from the target domain via the flow: Z ~ p, Φ1 (Z)u, x) = (V, Y) ~ o ( u, x).
pairs of modalities from the source domain (t, x), to sample
The sampled modalities (V, YY (i) minimize transport cost quantified by along the feature (ii) while minimizing the distortion along the struetural information quantified by Cg and cy. In 5.2
5.2
fused couplings to enhance the translation between different modalities of cells.
Why Does Adding a Fused Term Provide Benefits? Generally, unlike the linear EOT prob- lem (LEOT. the solution "5 to the quadratic HOT problem (QEOT) is not unique. This is where adding a fused term helps, as introduces more constraints. Intuitively, the fused term on addi- tional features allows to 'select' an optimal GW coupling: trade some GW optimality

---

by choosing coupling that not only minimizes the distortion of the structural information x.y: y. but also reduces the displacement cost, quantified by c, along the features u. V. Empirically, this significantly mitigates the issue of non-uniqueness in (pure) GW couplings. In $ we demonstrate that it enhances the stability of our procedure, particularly for single-cell data, as illustrated in Fig. and Fig. 18 On related note, but independently of the addition of the fused term, we investigate in App. E.2 other strategies to mitigate the non-uniqueness of the GW coupling. Notably, we show that using the geodesic distance on the data manifold as costs, instead of the squared Euclidean distance, helps address this issue on single-cell data. As an alternative, we also introduce an initialization scheme to bias the discrete GW solver to specific solution based on the neural coupling obtained in the previous iteration (App. E.2) Yet, this approach renders GENOT non- non-simulation free.
3.4 One algorithm to flexibly switch between neural OT problems
Algorithm U-GENOT. Skip teal steps for GENOT.
Single-cell genomics data is inher- ently noisy, due to biases incurred by sequencing protocols and the high sparsity of the measurements [34.. Thus, it is indispensable to flexibly switch between configurations of the problem setup. For example, it is of- ten not clear whether the mass con- servation constraint should be actu- ally loosened (for example whether to allow for modeling cell death), or whether incorporating prior informa- tion for trajectory inference via the quadratic term (e.g. with spatial in- formation (39 lineage barcoding (42) is beneficial. Similarly, trying different costs is crucial in single-cell genomics [17 35 (3).. The GENOT formulation offers this flexibility by
1:
Require parameters: Batch size n;; entropic regulariza-
tion E; unbalancedness parameter т; linear, quadratic or fused discrete OT solver Solverg.r.
2: Require network:
Time-dependent conditional velocity field v, TRY x TRP + R2; weightntti functions 110 RP - R. Eo R2 R.
3: for 1,..., Titrr do
4: Sample XI....,X,
μ and Yi...., Yn
5: Pe,r f- Solver,, (ix (Xi)==)) E RT
6: a ++ Pe,,1 and b P
7: Sample (i1.Ji,, (In) Pe,t
8:
Sample Zg.... Zn ~ and t1 in ~2((0, 1))
9: C(0)= t Ex |1,,0 ([Zg, (x)) -Y
10:
+ + (Xx) - naz) 2 (Ex(Y) - -nb;)²
11: 0 - Update(0, VZ))).
12:
end for
introducing only minor changes into the algorithm. We detail our procedure in Alg. [ showing the additional steps switch from GENOT to U-GENOT in teal. Limitations are discussed in
4 Related work and Discussions
Static Neural EOT. While GENOT is the first model to consider the quadratic (and fused) EOT setting, various methods have been proposed in the linear EOT scenario. The first class of methods solves (LEOTY's 's dual. While some of them (2 73 do not allow direct sampling from to ~², Daniels
et al. [5] and Mokrov et al. (55] model n (- However, these methods might be costly and
unstable, as they rely on Langevin sampling during training or inference.
Dynamic Neural EOT. The second class of linear EOT solvers builds upon the link between (LOTT)
and the SB problem [10 16 29 87]. Although they operate primarily in the balanced setting, Gazdieva et al. (26 recently considered the unbalanced one. As simulations are costly in this setting,
76 84 85 96]. While OT, our method fundamentally different, as we do not build upon the link between EOT and SB. We only use flow matching as a powerful black box to learn a flow from the noise p to each conditional n (-1x). Therefore, our approach allows for more flexibility. First (i), the abovementioned methods assume that X, y C. Rdd since they learn velocity field (or drift) by directly bridging x E XX, E y. Then, they map to v with the induced (stochastic) flow. On the other hand, conditionally to each
recent works consider simulation- free training via bridge matching [41 46 61
Tong et al. 84 3
use
X gd we learn a velocity field (t,o(-ix) Rº + R in the target space, by building paths between noise E R and y e C RR. Then, we map p to each n (-Ix) with the flows Φ1 (-Ix), and recover V = oi (-1-)t(p & p). Second (ii), they can only approximate the SB for c(x,y) = X
-1113. This results from [S] Prop. 2]. In contrast, our method handles any cost, as shown in Fig 5 Third (iii), as they learn a single flow directly transporting to v, they do not approximate the conditional densities nद (y)).. Similarly to Tong et al. 185) Pooladian et al. 64 couple samples from μ and v. but they only model deterministic maps and assume μ = (o0,, La). Finally, Bortoli et al. [5] recently proposed augmented bridge matching
7

---

Geodesic Y n (-1x)
Euclidean E Y Fg (+|x)
Uncertainty score per cell
X ~ pb
0.08
Terminal states
Progenitors
& EAVVO
TSI ScOre (aaker betivl
0.778 (0.000)
0. 716 (0.009)
CT error Gueve beoser
0.128 (0.067)
0.227 (0.046)
0.0
UMAP
Figure 2: Left: Source cell from the early time points (top left) and samples of the conditional distributions of the EOT coupling learned with GENOT for the geodesic cost dı (middle) and the & cost (right) projected onto a UMAP [53]. along with biological assessment of the learnt dynamics (TSI score, CT eITor $
C.2
Fig. 9 Right: UMAP colored according to the uncertainty score
cos-var(" (adM (-IX)) of each source cell x. Target cells are colored in gray.
procedure that preserves couplings. However, it still requires X. C R and does not approximate conditional densities, limiting its applicability single-cell genomics tasks.
Mini-batches and Biases Quantifying non-asy as-maattiiccally the bias resulting from minimizing a sample-based GENOT loss, and not its population value, is challenging. The OT curse of dimension-
ality [91] has been discussed in generative models relying on mini-batch couplings [22 27 84 85
Yet, our goal is not model regulaauldd OT, such deterministic MMonge map, Benamou Brenier vector field. We explicitly target the entropy- rerropy-regularized OT coupling. Thus, using & >
helps mitigate the curse of dimensionality because of two qualitative factors:
(i) Statistical. For both linear and quadratic OT, all statistical recovery rates that relate to entropic
costs 28 95 maps 663 66), or couplings (90), have far more favorable regime, with parametric rate & > 0 that dodges the curse of dimensionality.
(ii) Computational. While the benefits of employing large enough in Sinkhorn's algorithm
widely known, Rioux et al. [67 have recently shown that as increases, the quadratic OT problem (EOTT) becomes convex, making discrete solvers faster and more reliable.
To demonstrate this aspect, we empirically study the influence of the batch size on GENOT, using a recent benchmark (30, with known true linear EOT coupling TT, see Fig. 6
5 Experiments
While there is evidence that cells exactly evolve according to an entropic OT plan, leveraging OT is an established way to realign cells. We also evaluate the performance on simulated data in settings which closed-form solutions of the EOT coupling available for. Metrics and datasets discussed in App.[C]as [[andd App. D respectively. Further experimental details or results are reported in App. [ Setups for baselines are listed in App- Implementation details can be found in App. G We denote by GENOT-L the GENOT model for solving the linear problem (LOT), GENOT-Q for the quadratic (QEOT) and GENOT-F for the fused (EFGW) one. When considering the unbalanced counterparts, we add the prefix U. Moreover, when using the conditional mean of a GENOT model, we add the suffix CM. We denote our learned neural couplings by wg.
5.1 GENOT-L for modeling cell trajectories
GENOT-L simulated data We show that GENOT accurately computes EOT couplings using recent benchmark [30]. Tab. [ shows that GENOT-L outperforms all baselines considered. These results are of particular significance since most of the baselines are tailored towards the specific task of solving (i) the linear EOT problem LEOT). (ii) in the balanced setting (iii) with the & cost. In contrast, the contributions of GENOT go beyond this setting, see N2, N3, and N4. Going beyond the balanced problem (N3), we show that U-GENOT learns meaningful unbalanced EOT couplings
between Gaussians (36] (TT), which visualize in Fig. 10
GENOT-L learns the evolution of cells. Trajectory inference is a prominent task in single-cell biology. It is key to understand the development of cells from less mature to more mature cell
8

---

Source ja
Source j
X; ~ 1
Dacinostal Civinostat Belinostat
sample
Accuracy
0.90
Hesperadin
Quisinostar
Alvespimycin
0.85
Tanespimycin
Tak 901 lavopizidol
0.80
1.00
01.99
0.955
0.90
0.50
Unbalancedness r = Via
Target x
Prod. Target publi
xo(x|X))
Figure 3: Left: Accuracy of cellular response predictions of U-GENOT-L for cancer drugs with varying unbalancedness parameter T Ti = T2. Smaller T implies more unbalancedness (3 runs per T). Right: Mapping a Swiss roll in I& (()) to a spiral in IR² (v) with GENOT-Q. Center: Color code tracks where samples from j (top) are mapped (bottom). Right column: samples from p (top) and the corresponding conditionals, along with conditional density estimates. The learned QEOT coupling minimizes the distortion: Points close support of µ generate points close in support of v.
states. In the following, we consider dataset capturing gene expression of the developing mouse pancreas at embryonic days 14.5 (source) and 15.5 (target) [2]. Using the geodesic distance on the data manifold d as cost ($ 3.1 improves the learnt vector field compared to using the L cost (TSI score in Fig $ C.2). The CT error measures to what extent cells follow the manifold ($ Fig. 9
which is visually confirmed by samples of the conditional distribution (Fig These results support the need for flexible choice of the cost function (N2). As cells are known to evolve stochastically, we leverage GENOT's capability to generate samples from the conditional distribution to model
non-deterministic trajectories (N1). We follow Gayoso et al. (25) for assessing the uncertainty of cell trajectories by computing cos-var(Te xx) Vary Sory>~ma--- (-(x)) ccos-sim(X (Y, (- [X]]) ($ C.Ip.
We expect high uncertainty in cell types in early developmental stages and low variance in mature cell types. Indeed, Fig. 2 and Fig. 8 show that the computed uncertainty is biologically meaningful.
U-GENOT-L predicts single-cell responses to perturbations. Neural OT has been successfully applied to model cellular responses to perturbations using deterministic maps 7 GENOT has the comparative advantage to model conditional distributions, allowing for uncertainty quantification. We consider single cell RNA--eq data measuring the response of cells to 163 cancer drugs (78.. Each drug has been applied to population of cells that can be partitioned into 3 different cell types. The source (resp. target) distribution consists of cells before (resp. after) drug application. While there is no ground truth in the matching between unperturbed and perturbed cells due to the destructive nature know which unperturbed subset of
of sequencing technologies (hence the need for realignment), cells is supposed to be mapped to which perturbed subset of cells. We use this to define an accuracy metric (see App. C.2) while assess the uncertainty (N1) of the prediction using cos-var. Fig. 12 shows that for 117 out of 163 drugs, the model is perfectly calibrated (see App. C.T). while it yields a negative correlation between error and uncertainty only for one drug. We can improve the generation of cells by accounting for class imbalances between different cell types using U-GENOT. These imbalances might occur due to biases in the experimental setup due to cell death (78). Indeed, Fig. 3 shows that allowing for mass variation increases accuracy for nine different cancer drugs, selected as they are known to have strong effect. Fig. 13 and 14 confirm the results visually (N3).
5.2 GENOT for Quadratic EOT Problems
As outlined in [ quadratic OT problems are of particular interest in single-cell genomics (N4). Here, we focus on the task of aligning different cellular modalities, which necessary as most sequencing technologies are limited to measuring one modality [34). Yet, there cxist few protocols measuring two modalities simultancously, thus allowing us to assess the alignment as the ground truth matching is known. We leverage this to assess GENOT-Q's performance on the task of translating cells from ATAC-seq to RNA-seq, and also propose a novel way to improve the alignment using GENOT-F. Before, we demonstrate the task of mapping across incomparable spaces on simulated data.
GENOT-Q and U-GENOT-Q on simulated data. To visualize the problem of quadratic OT, we transport Swiss role in 3 to spiral in R Fig. shows that GENOT-Q successfully mimics an isometric alignment. Here, we set e = 0.01 and investigat seettttgtee its influence more detail Fig. [5] To quantitatively assess the performance of GENOT-Q and U-GENOT-Q, compute (unbalanced)

---

cy 29
Cx = dxi cx = dx (Geodesic)
Re (GENOT-F) vs. 1g
0.5
FOSCTTM
0.4
UMAP 2
0.3-
Sinkhorn divergence
0.2
109.
cell Eype B
10¹
UMAP 2
pDC
10
10/10
30/10
30/30
98/50
10/10
30/10
20/30
18/50
10/30
30/30
38/30
10/30
30/30
38/30
UMAP 1
Dimensions
Figure 4: Left: Benchmark of GENOT-Q models against discrete GW (GW-LR, App. on translating cells between ATAC space of dim. d1 and RNA space of dim. d2, with performance measured with FOSCTTM score (Ap) and Sinkhorn divergence between target and predicted target distribution. (left) intra-domain costs cx Cy lz. (right) geodesic distances cx dx and cy dy. We show mean and std across 3 runs. Right: Top: UMAP of transported cells with GENOT-F (colored by cell type) and cells in the target distribution (gray). Cells of the same cell type generate cells which cluster together in RNA space. Bottom: UMAP of transported cells with a GENOT model trained on batch-wise independent couplings, thus not using OT, generating cells which are randomly mixed.
entropic OT plans between Gaussian distributions, which the closed form EOT plan is known for 43 Fig 19 and Fig 20 show that GENOT-Q is able to meaningfully learn (un-)balanced entropic EOT plans for lower dimensions, but its performance decreases with increasing dimensionality.
GENOT-Q translates modalities of single cells. We use GENOT-Q to translate ATAC measurements (source) to gene expression space (target) on a bone marrow dataset [ST] As both modalities were measured in the same cell, the true match of each cell is known for this specific dataset. We compare GENOT-Q with discrete OT extended out-of-sample with linear regression (GW-LR, see F.2). We assess the performance using (i) the FOSCTTM ("Fractions of Samples Closer to the True Match" see C that measures the optimality of the learned coupling, and (ii) the Sinkhorn divergence 23 between the predicted target and the target to assess distributional fitting. As in 5.11 we leverage GENOT's flexibility and use as intra-domains costs the geodesic distances on the source and target domain, namely Cx dx. сy dy (N2). We also estimate the EOT coupling for the Eg cost for comparison. Results are shown in Fig. 3 Regarding the FOSCTTM score, we see that (i) using geodesic costs is crucial in high dimensions and (ii) GW-LR is competitive in low dimensions but not for higher ones. Regarding distributional fitting, GENOT models outperform by orders of magnitude.
GENOT-F improves modality translation. To enhance the performances attained by purely quadratic OT-based models, we introduce a novel method for translating between ATAC and RNA. We extend the idea of translating between cell modalitiies proposed by Demetci et al. [17 to the fused setting: We approximate RNA data from ATAC measurements using gene activity (79]. and we further process the data using conditional VAE
[48) to reduce batch effects. This way,
construct joint space Ω. Following the notations in $ 3.3 RNA (source x E X') and ATAC (target y ") carry the structural information, while features E Ω and v S2 are obtained from the VAE embedding. Fig. 21 shows that the additional linear penalty on the feature term helps to obtain a better alignment compared to GENOT-Q. Fig. 4 visualizes the learned fused coupling compared to GENOT trained with batch- wise independent couplings. When aligning multiple cell modalities, the proportions of cell types in the two modalities often differ (N3). We simulate this setting by removing cells belonging to certain cell types Tab, shows that U-GENOT-F preserves high accuracy while learning meaningful rescaling functions.
Conclusion. Motivated by applications single-cell genomics, we introduce GENOT, a versatile neural OT framework to learn cost-efficient stochastic maps, within the same space or across incom- parable spaces. GENOT can be used for learning EOT couplings with any cost function, and allows for loosening mass constraints if desired.
10

---

References
[1] David Alvarez-Melis and Tommi S Jaakkola. Gromov-wasserseenn alignment word embedding 2018.
spaces. arXiv preprint
F.2
[2] Aimée Bastidas-Ponce Sophie Tritschler, Leander Dony, Katharina Scheibner, Marta Tarquis- Medina, Ciro Salinno, Silvia Schirge, Ingo Burtscher, Anika Böttcher, Fabian J Theis, et al. Comprehensive single cell mrna profiling reveals a detailed roadmap for pancreatic endocrino-
genesis. Development, 146(12):dev I 173849, 2019. 5.1 D.2
[3] Jean-David Benamou. Numerical resolution of an "unbalancedd" mass transport problem. ESAIM:
Mathematical Modelling and Numerical Analysis, 37(05):851-866 2003. +
[4] Jean-David Benamou, Guillaume Carlier, and Luca Nenna. A numerical method to solve optimal transport problems with coulomb cost, 2015.
[5] Valentin De Bortoli, Guan-Horng Liu, Tianrong Chen, Evangelos A. Theodorou, and Weilie
Nie. Augmented bridge matching, 2023. 3.1
[6] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milhe and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
com/ goto///oool//aaa G
[7] Charlotte Bunne, Stefan G Stark, Gabriele Gut, Jacobo Sarabia del Castillo, Kjong-Van Van Lehmann, Lucas Pelkmans, Andreas Krause, and Gunnar Ratsch. Learning Single-Cell Pertur- bation Responses using Neural Optimal Transport. bioRxiv, 2021. 5.1
[8] Charlotte Bunne, Andreas Krause, and Marco Cuturi. Supervised training of conditional monge maps. Advances in Neural Information Processing Systems (NeurIPS), 2022. !
[9] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018
[10] Tianrong Chen, Guan-Horng Liu, and Evangelos A Theodorou. Likelihood training of schr odinger bridge using bowaard-bakwwar sdes theory. arXiv preprint arXiv:21 2021.
E.
[11] Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François- Xavier Vialard. Unbalanced optimal transport: geometry and Kantorovich formulation. Journal of Functional Analysis, 274 (11):300031223 2018.
T
[12] Keenan Crane, Clarisse Weischedel, and Max Wardetzky, Geodesics in heat: new approach to computing distance based on heat flow. ACM Transactions Graphics (TOG), 1215) 2013. 3.1
[13] Marco Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. In Advances in Neural Information Processing Systems (NeurIPS), volume 26, 2013. 2
[14] Marco Cuturi, Laetitia Meng-Papaxxathon Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein. arXiv Preprint arXiv:2201. 12324, 2022. E.
G.4
[15] Max Danicls, Tyler Maunu, and Paul Hand. Score-based generative neural networks for large- scale optimal transport. Advances in neural information processing systems, 14:12955-12965. 2021.
[16] Valentin De Bortoli, James Thorton,, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications basre generative modeling. Advances in Neural Information Processing Systems, 34:17695-17709, 2021. E.I
[17] Pinar Demetci, Rebecca Santorella, Bjöm Sandstede, William Stafford Noble, and Ritambhara Singh. Scot: single-cell multi-omics alignment with optimal transport. Journal of Computational Biology, 29():3-188 2022.
I 3.1
3.4 5.2 D.4

---

[18] Mingze Dong, Bao Wang, Jessica Wei, Antonio H de O. Fonseca, Curtis J Perry, Alexander Frey, Feriel Ouerghi, Ellen F Foxman, Jeffrey J Ishizuka, Rahul M Dhodapkar, et Causal identification of single-cell experimental perturbation effects with cinema-ot. Nature Methods, 20(11):1769-177) 2023. |
[19] Michael B Elowitz, Arnold J Levine, Eric D Siggia, and Peter S Swain. Stochastic gene expression in a single cell. Science, 297(5584): 183-1186, 2002.
[20] Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, and Fabian Theis. Unbalancedness in neural monge maps improves unpaired domain translation. arXiv preprint arXiv:2311.15100 2023.
3.2
[21] Luca Vincent Eyring, Dominik Klein, Giovanni Palla, Soeren Becker, Philipp Weiler, Niki Kilbertus, and Fabian Theis, Modeling single-cell dynamics using unbalanced parameterized monge maps. bioRxiv, 2022. doi: 10.1 11111100222.1.04..... 510666. hRtpttts: //wtts::/wwww... biorxiv
10. 10.04. 510766
[22] Kilian Fatras, Younes Zine, Szymon Majewski, Rémi Flamary, Rémi Gribonval, and Nicolas Courty. Minibatch optimal transport distances; analysis and applications. arXiv preprint urXiv:2001.01792, 2021.
[23] Jean Feydy. Thibault Séjourné, François-Xavies Vialard, Shun-Ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between Optimal Transport and MMD using Sinkhorn Divergences. In International Conference on Artificial Intelligence and Statistics (AISTATS), volume 22, 2019.
5.2 B.1
[24] Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya, and Tomaso A Poggio. Learning with a wasserstein loss. Advances in neural information processing systems, 28, 2015.
2
[25] Adam Gayoso, Philipp Weiler, Mohammad Lotfollahi, Dominik Klein, Justin Hong, Aaron M Streets, Fabian J Theis, and Nir Yosef. Deep generative modeling of transcriptional dynamics for rna velocity analysis in single cells. bioRxiv, pages 2022-08, 2022. 5.1
[26] Milena Gazdieva, Arip Asadulaev, Alexander Korotin, and Evgeny Burnaev. Light unbalanced optimal transport, 2024. URL.htttss://rrav..orr/// 2230g//bb33333..oo 07988
[27] Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample Complexity of Sinkhorn Divergences. In International Conference on Artificial Intelligence and Statistics (AISTATS), volume 22, 2019. 4
[28] Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample com- plexity of sinkhorn divergences. In The 22nd international conference artificial intelligence and statistics, pages 1574-1583. PMLR, 2019.
[29] Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry Vetrov, and Evgeny Burnaev, Entropic neural optimal transport via diffusion processes. arXiv preprint arXiv:2211.0111 2022.
E.
[30] Nikita Gushchin, Alexander Kolesov, Petr Mokrov, Polina Karpikova, Andrey Spiridonov, Evgeny Burnaev, and Alexander Korotin. Building the bridge schr' " odinger: A continuous entropic optimal transport benchmark. arXiv preprint arXiv:2306. 10161, 2023. 5.
F.
[31] Muzlifah Haniffa, Deanne Taylor, Sten Linnarsson, Bruce J Aronow, Gary D Bader, Roger A Barker, Pablo G Camara, J Gray Camp, Alain Chédotal, Andrew Copp, et al. A roadmap for the human developmental cell atlas. Nature. 597(7875):196-205 2021. E.1
[32] Zhisong He, Leander Dony, Jonas Simon Fleck, Artur Szalata, Katelyn X Li, Irena Sliskovic, Hsiu-Chuan Lin, Malgorzata Santel, Alexander Atamian, Giorgia Quadrato, et An integrated transcriptomic cell atlas of human neural organoids. bioRxiv, pages 2023-10, 2023. E.I
12

---

[33] Leon Hetzel, Simon Boehm, Niki Kilbertus, Stephan Günnemann, Fabian Theis, et al. Predicting cellular responses to novel drug perturbations at single-cell resolution, Advances in Neural Information Processing Systems, 35:26711
2022. D.3
[34] Lukas Heumos, Anna C Schaar, Christopher Lance, Anastasia Litinetskaya, Felix Drost, Luke Zappia, Malte D Lücken, Daniel C Strobl, Juan Henao, Fabiola Curion, et al. Best practices for single-cell analysis across modalities. Nature Reviews Genetics, pages 1-23, 2023. 3.4 5.2 D.4
[35] Guillaume Huguet, Alexander Tong, María Ramos Zapatero. Guy Wolf, and Smita Krish- naswamy. Geodesic sinkhorm: optimal transport for high nigh-dimensional datasets. arXiv preprint vrXiv:22/1.00850 2022.
3.1 3.4
[36] Hicham Janati, Boris Muzellec, Gabriel Peyré. and Marco Cuturi. Entropic optimal transport between unbalanced gaussian measures has closed form. Advances in neural information processing systems, 11.10468-104777 2020. 5.1
[37] Olav Kallenberg. Foundations of Modern Probability. Springer, 2002. URL https:://inkk
wtt.///oo/ooo/..1000001000 1007///10..07077788-1 --00000-------------00000000000- 61871-1
3.1
[38] L. Kantorovich. On the transfer of masses (in russian). In Doklady Akademii Nauk, volume 37, page 227, 1942. 2
[39]
Dominik Klein, Giovanni Palla, Marius Lange, Michal Klein, Zoe Piran, Manuel Gander, Laetitia Meng-Papaxxntthss Michael Sterr, Aimee Bastidas-Ponee Ponce, Marta Tarquis-Medina, et al. Mapping cells through time and space with moscot. bioRxiv, pages 2023-05 2023.
3.1 3.4
D.4 G.4
[40] Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. 2022. doi: 10.48550/ARXIV:220) 12220. http/stts:://et org/ abs//2/s//2212 12220 E.
[41] Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev. Light schrödinger bridge, 2023.
[42] Marius Lange, Zoe Piran, Michal Klein, Bastiaan Spanjaard, Dominik Klein, Jan Philipp Junker, Fabian Theis, and Mor Nitzan. Mapping lineage-traced cells across time points with moslin. bioRxiv, pages 2023- 04, 2023. 3.4
C.2
[43] Khang Le, Dung Q Le, Huy Nguyen, Dat Do, Tung Pham, and Nhat Ho. Entropic gromov- wasserstein between gaussian distributions. In International Conference on Machine Learning, pages 12164-12203. PMLR, 2022.
5.2
E.2
19
20
[44] Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal entropy-transport problems and new hellinger-kanttrrviie kantorovich distance between positive measures. Inventiones Mathematicae, 21 211(())969-11 (110)999-111.,,. 2018. 2
[45] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling, 2023. 2
[46] Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A. Theodorou, Weili Nie, and Anima Anandkumar. 1ss:: Image-to-image schrödinger bridge, 2023.
[47] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow, 2022. 3.1
[48] Romain Lopez, Jeffrey Regier, Michael B Cole, Michael Jordan, and Nir Yosef. Deep generative modeling for single-cell transcriptomics. Nature methods, 15(12):1053-1058, 2018.
5.2
[49] Romain Lopez, Jeffrey Regier, Michael B Cole, Michael Jordan, and Nir Yosef. Deep generative modeling for single-cell transcriptomics. Nature methods, 15((2), 2018. D.4
[50] Frederike Lübeekk Charlotte Bunne, Gabriele Gut, Jacobo Sarabia del Castillo, Lucas Pelkmans, and David Alvarez-Melis. Neural unbalanced optimal transport via cycle-consistent semi- couplings. arXiv preprint arXiv:2209.15621, 2022.
3.2
13

---

[51] Malte D Luecken. Daniel Bernard Burkhardt, Robrecht Cannoodt, Christopher Lance, Aditi Agrawal, Hananeh Alice, Ann Chen, Louise Deconinck, Angela M Detweiler, Alejandro A Granados, et al. A sandbox for prediction and integration of dna, ma, and proteins in single cells. In Thirty-fitth conference on neural information processing systems datasets and benchmarks
track (Round 2), 2021. . 5.2 D.4
[52] Christian Léonard. A survey of the schrödinger problem and some of its connections with optimal transport, 2013. 2
[53] Leland McInnes, John Healy, and James Melville. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint /802.034000 2018.
E.2
[54] Facundo Mémoli. Gromov-wasserstin distances and the metric approach to object matching.
Foundations of computational mathematics, 11:417-487, 2011.
2
[55] Petr Mokrov, Alexander Korotin, and Evgeny Burnaev. Energy-guided entropic neural optimal transport. arXiv preprint arXiv:2304.0600 2023. E.1
[56] Gaspard Monge. Mémoire sur la théoric des déblais et des remblais. Histoire de I' l'Académie Royale des Sciences, pages 366-704, 1781.
[57] Kevin R Moon, Jay S Stanley III, Daniel Burkhardt, David van Dijk, Guy Wolf, and Smita Krishnaswamy, Manifold bearning-based methods for analyzing single-cell sequeenunnin data. Current Opinion in Systems Biology, 7:36-46, 2018.
[58] Maksim Nekrashevich, Alexander Korotin, and Evgeny Burnaev. Neural gromov-wasserstinn
optimal transport. arXiv preprint arXiv:2303.05976 2023, T
[59] Marcel Nutz. Introduction to entropic optimal transport.
B33
[60] Giovanni Palla, David S Fischer, Aviv Regev, and Fabian Theis. Spatial components of molecular tissue biology. Nature Biotechnology, 40(3):308-31 2022. [
[61] Stefano Peluchetti. Diffusion bridge mixture transports, schrödinger bridge problems and generative modeling, 2023.
[62] Gabriel Peyré, Marco Cuturi, and Justin Solomon. distance matrices. In International Conference on Machine Learning, pages 2664-26662 2672, 2016.
averaging of kernel and
3.3
[63]
Alam-Alexandre Pooladian and Jonathan Niles-Weed. Entropic estimation of optimal transport maps. arXiv preprint arXiv:2109. 12004, 2021.
[64] Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Ennich. Brandon Amos, Yaron Lipman, and Ricky Chen. Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint arXiv:2304.1477 2023. 3.1
E.
[65] Aviv Regev, Sarah Teichmann, Eric Lander, Ido Amit, Christophe Benoist, Ewan Birney, Bernd Bodenmiller, Peter Campbell, Piero Carninci, Menna Clatworthy, et al. The human cell atlas. elife, 6:e27041, 2017.
[66] Philippe Rigollet and Austin J. Stromme. On the sample complexity of entropic optimal transport, 2022.
[67] Gabriel Rioux, Ziv Goldfeld, and Kengo Kato. Entropic gromov- wasserstein distances: Stability, algorithms, and distributional limits. arXiv preprint 27Xiv:2306.00188 2023, 4
[68] Filippo Santambrogio. Optimal Transport for Applied Mathematicians. Birkhäuser, NY, 55 (58-63):94, 2015.
B.1
[69] Meyer Scethon, Marco Cuturi, and Gabriel Peyré. Low-rank sinkhorn factorization. In International Conference Machine Learning, pages 9344-9354, PMLR, 2021. !
14

---

[70] Meyer Scetbon, Gabriel Peyré, and Marco Cuturi. Linear-time gromov wasserstein distances using low rank couplings and costs, In International Conference on Machine Learning, pages 19347-19365. PMLR, 2022. 2
[71] Meyer Scetbon, Michal Klein, Giovanni Palla, and Marco Cuturi. Unbalanced low-rank optimal transport solvers. Advances in Neural Information Processing Systems, 36, 2024. |
[72] Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-Transport Analysis of Single-Cell Gene Expression Identifies Developmental Trajectories in Reprogramming. Cell, 176(4), 2019.
E.1
[73] Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large-scale optimal transport and mapping estimation. arXiv preprint arXiv: 2017.
[74] Thibault Séjoumé, François-Xan Xavier Vialard, and Gabriel Peyré. The unbalanced gromov wasserstein distance: Conic formulation and relaxation. Advances in Neural Information Processing Systems, 34:8766-8779, 2021. B.2
[75] Thibault Séjourné, Gabriel Peyré, and Trançois-Xavier Vialard. Unbalanced optimal transport, from theory to numerics. arXiv preprint arXiv:221 2022. A
[76] Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schr\" odinger bridge matching. arXiv preprint arXiv:2303. 2023.
[77] Justin Solomon, Fernando De Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional wasserstein distances: Efficient optimal transportation geometric domains. ACM Transactions on Graphics (ToG), 3414):1-11, 2015.
3.1
[78] Sanjay R Srivatsan, José L CeFFaline-Figuere Vijay Ramani, Lauren Saunders, Junyue Cao, Jonathan Packer, Hannah A Pliner, Dana L Jackson, Riza M Daza, Lena Christiansen, et al. Massively multiplex chemical transcriptomics at single-cell resolution. Science, 367(6473),
2020. 5. D.3
[79] Tim Stuart, Avi Srivastava, Shaista Madad, Caleb A Lareau, and Rahul Satija. Single-cell chromatin state analysis with signac. Nature methods, 18(11):1333-1341 2021. 5.2 D.4
[80] Thibault Séjourné, Gabriel Peyré, and François Xavier Vialard. Unbalanced optimal transport,
from theory to numerics, 2023. B.2 B.3 B.3
B.2
B.2
[81] Thibault Séjouréé, François-Xave Xavier Vialard, and Gabriel Peyré. The unbalanced gromov
wasserstein distance: Conic formulation and relaxation, 2023. 3.3
B.1
B.3
[82] Alexis Thual, Huy Tran, Tatiana Zemskova, Nicolas Courty, Rémi Flamary, Stanislas Dehaene, and Bertrand Thirion. Aligning individual brains with fused unbalanced gromov- wasserseen,,
2023. 3.3
[83] Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Tra- jectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics. In International Conference on Machine Learning (ICML), 2020. [
[84] Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation-fiee schr odinger bridges via score and flow matching. arXiv preprint arXiv:2307.03667. 2023.
[85] Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Conditional flow matching: Simulation dynamic optimal transport. arXiv preprint arXiv:2302.00482. 2023.
H.T
[86] Théo Uscidda and Marco Cuturi. The monge gap: A regularizer to learn all transport maps, 2023. |
15

---

[87] Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving schrödinger bridges via maximum likelihood. Entropy. 23(9):1134, 2021.
[88] Titouan Vayer, Laetitia Chapel, Rémi Flamary, Romain Tavenard, and Nicolas Courty. Optimal transport for structured data with application on graphs. arXiv preprint 4rXiv:1805.09114, 2018.
3.3
B.1 B.I
[89] Daniel E Wagner and Allon M Klein. Lineage tracing meets single-cell omics: opportunities and challenges. Nature Reviews Genetics, 21(7):410 427-427, 2020.
[90] Tao Wang and Ziv Goldfeld. Neural estimation entropic optimal transport, 2024.
[91] Jonathan Weed and Francis Bach. Sharp asymptotic and umite-sampll rates of convergence of empirical measures in wasserstein distance, 2017. URL pestps://arxivy
1b77
1707
00087
[92] Philipp Weiler, Marius Lange, Michal Klein, Dana Pe'er er, and Fabian Theis. Unified fate mapping in multiview single-cell data. bioRxiv, pages 2023 07, 2023.
[93] Karren D Yang and Caroline Uhler. Scalable Unbalanced Optimal Transport using Generative Adversarial Networks. International Conference on Learning Representations (ICLR), 2019.
3.2
E.1
[94] Karren D. Yang and Caroline Uhler. Scalable unbalanced optimal transport using generative adversarial networks, 2019. !
[95] Zhengxin Zhang. Ziv Goldfeld, Youssef Mroueh, and Bharath K. Sriperumbudur. Gromov- wasserstein distances: Entropic regularization, duality, and sample complexity, 2023. 2
[96] Linqi Zhou, Aaron Lou, Samar Khanna, and Stefano Ermon. Denoising diffusion bridge models, 2023. 4
16

---

Appendix
A Limitations
While our study opens up the use of neural optimal transport to a wide range of tasks in single-cell genomics, there are limitations of our work:
The solution to the entropic (Fused) Gromov-Wassersténn problem is not necessarily unique. As GENOT-Q and GENOT-F leverage discrete solvers in each iteration, the orientation of the solution across iterations is not maintained. To prevent learning mixture of solutions. we propose to use a fused term whenever possible, make use of different cost functions if helpful, or using GENOT with our proposed initialisation scheme ($ While the latter can always be applied, its major drawback is the costly training as a neural ODE has to be solved every iteration.
When learning unbalanced EOT plans, choosing value for the hyperparameters (Tas Tb) is not evident, As discussed Séjourné et al. (75,, the choice of these parameters is mostly empirical.
While GENOT is motivated by the extensive use of optimal transport in single-cell genomics, there is no evidence that cells evolve exactly according to entropic OT plan.
Evaluations of learnt transport plans for tasks in single-cell genomics experiments are task-specific, and do not necessarily reflect to what extent the learnt coupling is an entropic OT coupling. In effect, we rely on prior biological knowledge (CT error, Fig. (2) or on common downstream trajectory inference methods in single-cell genomics to evaluate the leannt velocity field (TSI score, Fig. ( when computing EOT plans for trajectory inference. Similarly, we can only assess learnt EOT plan on a cluster level, and not a single data point level when studying cellular perturbations (3). For the task of modality translation (Fig. 4. we know the ground-truth coupling. Yet, the ground-truth coupling deterministic and one-to-one, and hence is not valid solution of an entropy regglarriee Gromov- Wasserseenn OT plan with 0.
B Propositions and Proofs
B.1 Existence of EOT couplings
In this section, we point out that solution to the entropic fused GW problem (EFGW). D. and its unbalanced counterpart, always exists. Note that for a 0, we recover the existence for the pure entropic GW (EFGW) and its unbalanced counterpart (UQEOT) problems. The arguments are
adapted from 88 Prop. 2]
Balanced setting: We first recall the fused Gromov- Wasserstein problem below.
min rEll(ae) ( Oc Dop. d(τ @ m)) + =KI((
where Qeg.ey ((u, x), (v,)), = (1 - a)c(u, v) + xx) cy) We first need
to show that a solution x² to this problem always exists. We assume that the intra-domain cx, cy costs, and the inter-domain cost e are l.s.c, an assumption that matches all our experimental settings. From 88 Lemma 3), this implies that the functional:
[ ((0xx))x(xx)))>>>)))))))))))))))
k E M(p, v)
Veg.sy d(T @ TT)
weakly I.s.c, i.e. I.s.c for the convergence law. Moreover, from 23 Prop. 8), ग + KL(T|p @ v) is also weakly Ls.c. Therefore, the objective function of the fused
problem:
+ [ (((xx))x(X=)))
F F: E I(H,))
D Pigieg d(T 8x) + [KL(T
(8)
is weakly l.s.c, as sum of weakly l.s.c functionals. Then, since we assume that X' and y are compact, they are polish spaces, so from [88 Theorem 1.7], I(u, v) is weakly compact. Therefore, from Weierstrass Theorem, J attains its minimum 11( v), i.c. it exists गद solution to this problem.
17

---

Unbalanced setting: The proof is very similar. We first recall the unbalanced fused Gromov Wasserstein problem.
J (xxxxxxxxxxxxxx Q Cegi,
min (XM)(XX)))
+ (TIlK) + A2KL (T2
where KT, 8 (a|3) KT,a @ al s 9). From Prop. 8), KT. is weakly l.s.c, so - KI, @ (m (p) and æ + KL (T2l). Therefore, the objective function of the unbalanced fused Gromov- Wasserstein problem is also weakly lss.c. as a sum of weakly lss.c functions. Moreover, since X' and y are assumed be compact, X' y also compact and the Sanach-Alaope provides the compactness of M x )). Then, from Weierstrass Theorem, similarly to the balanced setting, we get the existence solution के. this problem.
B.2 Proofs of $ 3.2
Proposition B.1 (Re-Balancing the unbalanced problems.). Let Te,r be an unbalanced EOT coupling, solution of (ULFOT) (UQEOT) between E (X) and v E M ())) We note jú piim and D = p2t its marginals. Then, both cases, j (resp. 8) has a density w.r.t. l (resp. v). That is, there exist two re-weighting functions, one on each space, T] X + Rt and & : y - Rt s.t. 77 -- and i E. v. Furthermore, ji and ù have the same total mass and 1. (Linear) ,े,r solves the balanced problem (LFOT) between it and i with the same e.
2. (Quadratic) Provided that cx, ey (or -Cx, -ey) are CPD kernels (see Def. [.2), Wh.s solves the balanced problem (QEOT) between jı and with E' = -m( e, where m ( = - x V).
We first remind the definition of a conditionally positive kernel, which is involved in point 2. of Prop. 3.1
Definition B.2. A kernel k IRd Rd + conditionally positive if it is symmetric and for any X1y Xn E Rd and a E IR" s.t. a 1, = O, one has
a,a, k(x, x,)
Conditionally positive kernels include all positive kernels, such y(x yx or (x.
exp(- - but also the negative squared Euclidean distance k(x,y) -| y1
Proof of 3.7 Step Re-weightings. We first show that for पे. solution of (LEOT) or QEOT exists 17, E Rd +
s.t. j = pitte, = and v = pi f -.V.
We start with the linear case and consider -े, solution of (LEOT). The result follows from duality. Indeed, from [80 Prop. 2], one has the existence of the so-called entropic potentials f" E C(x) and g" E C(V) s.t.
d d(u Ø v) (x.y) =
[x(() + g gy( c(x,y)
exp exp
(10)
Therefore, j ==11 =) and i c v where 77 : X' + R and & y + B+ are defined by:
n(x) / exp
fx(x + g" gy() c(x,y) dr(y)
(11)
[x(x) + gr() c(x,y)
and {(y) exp
dp(x).
We now handle the quadratic case and consider The solution of (QQOOT) We remind that the (UQEOT) problem between μ and v reads:
min /
(cx(x,x)) eyly, dr do
nEM-(xx)) (2x) (xxx)2
+ =KL 2
(T + X1KL (FIlm) + 12LL (m22P), >>2P((TT))),
18

---

For x E M xX'))), the quadratic relative entropy KL 0 ( ( = ( (p finite i.f.f. & x has a density w.r.t. (μ 2 v) x
(μ @ v), which implies that x has density w.r.t.
0 v. Therefore, one can reformulate (UQFOT)
as:
min (xX)(
(xxO + +eKL(M 8 v) + T TiKL(hi (M) + (2KI((bll))
eyl, h(x, y))(', y') du(xx) du(x') dv(y) drly')
(12)
where we extend the KL divergence for densities: (T,(rl7) [(rlog(1) r 11) 1y and define
the marginal densities h1 : E X + Jy y) dy and h2 : 3 + fx (xx,)))))))))) dy(y). As result, it exists ht E LÍ (X x )) s.t. पे, = ht a @ v. It follows that f = 1) μ and i = E v with
7 hī and & h2.
Remark B.3. Note that in both cases, since dā(() = n7x) du((x) and di(yy) = E(y) dy(y). the cquality of mass of ji and 5 yields [xX)] Ey (y-))
Step 2: Optimality in the balanced problem for the linear case. We now prove point 1 stating
that if æt.r solves problem ULEOT between μ and v, then it solves problem (LOOT) between fú and
for the same entropic regularization strength e. It also follows from duality. We remind that that thanks to 80 Prop. 2]
d d(u @ v) (x,y) = exp
[x((x) + g* c(x,
(13)
with f E C(X) and g" E C(Y). Moreover, by 59 to the optimality problem (LEOT). Therefore, -े. is solves problem LEOT its marginals and i i.e.
Theorem 4.2], such a decomposition is equivalent
/
- arg min xEIl(ā.)) Jxxy y) dx(x,)) + =
(14)
Step Optimality in the balanced problem for the quadratic case. We now prove point 2 stating that, provided that the costs cx and Cy are conditionally positive (or conditionally negative), if n² solves problem (UQOOT) between μ and v. then it actually solves problem (QEÖT) between f and ū for the entropic regularization strength e' =m( We first define the functional:
/
F: (
E (X' x1) +
1(XXP) (xx(x, x') cyly, (x) y)
(15)
+ +eKL r )) + X (L(T @ n 1) A2 (T2 72/1 v);
s.t. THe E arg F T). Then, if we define the lincarized cost
/ JxXY ccx(x, - -/
: X' y
(16)
we get from 80 Proposition 9], that TË,r solves:
/
Thr E arg =EM(xy) Xxy
&e dT + JLL
(17)
Aimm( (LL (x) (KL(
Therefore, T2,r solves problem ULEOT between j and v for a new cost c and the regularization strength s' em(( ).) We seek to apply point [ to get that -द, solves problem LEOT between j == pi pm and v = prttEEr for the same entropic regularization strength e'. To that end, we first verify that -,, is continuous. Since cx and cy are continuous on X' X' and y x 37, the function:
(x,,,'''') E (X 2 +> (cx(xx,)) - cy(y y')
continuous. Therefore, it is bounded since (X' x 2 is compact as a product of compact sets. Then, since n has finite mass, Lebesgue's dominated convergence yields the continuity of ct,,. We then to apply point and get:
- E arg min xEll((.)) xxx
mig
(x, dT(( +elr))+KKKK((//)//
(18)
19

---

Since the costs are conditionally positive (or conditionally negative) kernels. (18) finally yields the desired result by applying [8]] Theorem 3]:
/ (xx x') yy dx y') xell(a.e) xxy
-े.r arg min min
( (x, + (TKT(
(19)
Remark B.4. In various experimental settings, μ and v have mass 1, and we impose one of two hard marginal constraints, for instance, je, by setting = 1. Then 7 has also mass I and m(= so = m(T) = e and we keep the same regularization strength by re-balancing (UQHOT)
Proposition B.5 (Pointwise estimation of re- se-weighting functions.). Let _e, = Ei, solution to (ULEOT) between empirical counterparts of μ and Let a Pe,r In and b P 2, its marginals weights. Then, almost surely, n n.a - rT().) and n b; E(Y:).
5,t (X.,))
Proof. As we saw in the proof of Prop 3.1 using Séjourné et al. 80 Prop. 2], one has the existence of f ((X) and g" C()) s.t.
d d(p & v)
(x) + gy e(x,)) -h(x,y)
(x.y) exp
E
and the relative densities 1, and & s.t. ā piamerr 17 and E =22tm = are given by
n: 1: + /,, yxx and E / h(x,y)
(20)
Now, let consider āe.r = Pi the solution of problem LEOT between An
E dx, and in = "E sy. Similarly, using Séjourné et al. 88 Prop. 2], one has the existence of fi E C(Y) and 9 (())) s.t.
dã
(x)() ()(()) c(x,,)
(x, = exp d(in 2 D,,)
1o (x,y)
(21)
8
and the relative densities th and En s.t. in Pit#t nn in and 8n 12 En in given by
Th x h (x,,,) and &n Gyyy Σ hn (x).y)
(22)
From Eq. (21). we get that Púj = 1 1(x, y,), so reminding that = Pe.r and b = P 1n,
n² one has n n.; = no (X,) and n nb (Y,). We now show that almost surely, nn + n and En +
pointwise, which implies the desired result that, almost surely, na, - rxx)) and nb, + <(Yi).
Almost surely, in - and in v. so using [74] Proposition 10), fí - I and g - g" in sup-norm. Since fū - on X and th -> g" sup-norm on X' x V. Indeed, for (x E (X' x )), one has:
in sup-norm,
show that hn + h in
1/4 (x, y) - (xx.))
(x) (v) - c(x,x) - exp
f*(x) + gy( c(x,y)
exp
8
c(x,y) E
f+(x) + (x)
fx((x)+ gy)
(23)
=exp
exp
-exp
8
fi((x) + gn(y))
fx( + 9" (y)
<.aa, exp
exp
8
20

---

with Me,s sup(x.y sup(x,)) exp the compact X' V, as e is continuous. Afterwards,
exp c(x,y)
too, since y) exp
c(x.y]
continuous on
[x) + (y)
((xx)+9'(1
exp
exp
E
(x)
In (x)
(x)
(x)
exp
exp
exp
exp
(24)
e
E
(x)
g" (x)
[x*()
(x)
exp
exp
exp
exp
8
E
For the first term, one has:
(x)
(x)
(x)
g g'(x) (x)
exp
exp
exp
exp
8
e
Jr (x)
9 (x)
9" (x)
exp
exp
exp
(25)
E
[| no
9h (x)
9 (x)
S exp
exp
exp
First, we can bound uniformly exp Ilffillo// since fa converges in sup-norm, so (llfft )000 bounded. Then, since Ih convergences in sup-norm, it is uniformly bounded, and since g is continuous on the compact X, it is bounded. Therefore, we can find a compact K C R s.t.
g" (X) C K and for each n, 9h (x) K. Then, applying the mean value theorem the C function x + exp(x/e) on K, we can bound:
(x)
g" (x(x))
exp
exp
S
sup (Pp(22)l9 (x(x zek
(26)
This yields the existence of a constant
M1 s.t.
(x)
I
(x)
g' g'(x)
exp
exp
exp
exp
- M1lll g01lo0 (27)
E
E
E
Using the same strategy, we get the existence of a constant M2 > O s.t.
(x)
9x())
(x)
gx(x)
exp
exp
exp
exp
< M2lls 10x
(28)
Combining (7 and 28 with (23 we get that:
(x) h(x)) < Mc. (M, g" 1lx + Mallft 1xx
(29)
from which we deduce that h, - h in sup-norm, from the convergence of fa I and yn. + g in sup-norm.
Now, we can show the pointwise convergence of îm. For any E X, one has:
(x) - 7xx)
Nnn (x. (D dD f (
Nn, (x,, y) dii / 1(xxy) di),x))) +1 / (xx,,)dū,( - ((x,)) -
(30)
/ ||h, - 1ollo den(y) / () / M(x,y) du(y))
=1|A,, -blloo + h(x,y) den(y) - / 1(xx, x) du(y)l
Therefore, it almost surely holds that in (x) - rx(x). Indeed, ||,, - -1lls - 0 since we have shown that hn - h in sup-norm. Then, is continuous the compact X' V, is bounded, since Nh - n. we get hhd - dd. Next, we show similarly that, almost surely, &n + & pointwise.
21

---

C
Metrics
We start with introducing general metrics in $ C.I. some of which will be used in the metrics introduced in the context of experiments on single-cell data in C.2
C.1 General metrics
In the following, we discuss way how to classify predictions in generative model. We start with the setting where each mapped sample to be assigned to a calegory based on labelled data in the target distribution. We then continue with the case where there are also labels for samples in the source distribution, and this way define classifier fetass mapping from the source distribution to labels in the target distribution with the possibility to compute accuracy. Building upon this, we assign the classifier felass an uncertainty score for each prediction. Finally, define a calibration score assessing the quality of given uncertainty score.
Turning a generative model into a classifier In the following, consider a finite set of samples in the target domain Y1 .I.....38 . E V. Assume {ym)m= allows for partition {ym) M = Like x Th. Hence, each sample belongs to exactly one class, which we interchangeably refer to as the sample being labelled. Let T X' y be map (deterministic stochastic), and let Ji.NN y {Tx) be the 1-nearest neighbor classifier. We obtain map g from X' to } K by the concatenation of fINN and T. This map proves useful in settings when mapped cells are to be categorized, e.g. to assign mapped cells a cell type.
A metric to assess the accuracy of generative model In the following, assume that the set of samples in the source domain xI:... xx allows for a partition [x,) ).N K We want to construct classifier felass assigning each category in the source distribution {Sa J T= probabilistically
to category in the target distribution [T Define felnss (Sa J - N via (Jelnss (Sx)),
Ex, (o(x.))=T;) where g X (Tx k=1 was defined above.
Assume that there exists known one-to-one (injective would be sufficient) match between elements in {Sx k= and elements in (Tx [ Then we can define a confusion matrix A with entries Aij Ex,ET. (((xx))SS))- In the context of entropic OT the confusion matrix is element-wise
defined as
Aig = Σ ((s)(((()))==
(31)
X.ET.
This way we obtain an accuracy score of the classifier fetuss mapping a partition of one set of samples to partition of another set of samples.
To assess the meaningfulnes, of an uncertainty score, we introduce the following calibration score. Assume we have classifier which yields predictions along with uncertainty estimations. Let E g be vector containing an uncertainty estimation for each element in {S*) Moreover, let a E k be a vector containing the accuracy for each element in {Sx k= We then define our calibration score to be the Spearman rank correlation coefficient between u and 1x - a, where 1g denotes the K -idimensional vector containing in every entry. In effect, the calibration score is close to 1 if the model assigns high uncertainty to wrong predictions and low uncertainty to true predictions, while the calibration score is close uncertainty to correct predictions and low uncertainty to wrong predictions.
Calibration score
if the model assigns high
In the following, consider stochastic map T. Let y1,..,, To(+1x) obtained from T. To obtain a calibration score for felass we estimate statistic ((To(.Ix) from the samples y1,-.. reflecting estimation of uncertainty. Then, let the uncertainty of the prediction of felass for category S, be the mean uncertainty statistic, i.e. Exes: felms (St) we get the uncertainty score
V(xe())))) In effect, for each prediction [S
V(To(-Ix))
04
Σ XES.
(32)
[S
22

---

Assessing the uncertainty with the cos-var metric Gayoso et al. [25] introduce statistic to assess the uncertainty of deep generative RNA velocity methods from samples of the posterior distribution, which we adapt to the OT paradigm to obtain
cos-var(me (- x)) = Vary~ao( ccos-sim( by [xD]
(33)
where cos-sim denotes the cosine similarity. We refer to this metric as cos-var, as it computes the variance of the cosine similarity of samples following the conditional distribution and the conditional mean. We use 30 samples from the conditional distribution to compute this metric.
C.2 Single-cell specific metries
TSI score for developmental systems In the following, we describe TSI metric used in Fig. 2 The terminal state identification score (TSI) was introduced in Weiler et al. [92] and assesses to what extent the learnt vector field of a developmental system points towards the correct directions. In effect, we assume that terminal states of the biological developmental process are known. CellRank 42 identifies macrostates using a Schur decomposition based on the transition matrix computed from the learnt vector field (using CellRank's VelocityKernel), From these macrostates, CellRank infers a set of putative terminal states, loosely speaking by looking for sinks in the transition matrix. To define the TSI score, consider system containing m terminal states and the function f that assigns each number of macrostates n the corresponding number of identified terminal states. In the case of a strategy that identifies terminal states optimally, fups describes the step function
(n) = [n, < m
m, 2
The TSI score of the learnt vector field @ is the area under the curve so relative to the area under the curve fopts i.e.,
TSI(0)= = (n(
Nuea m=
2
Naw
fe(n).
Nun n= fopt m1 + Nmax m))
==1
Thus, the TSI score lies between and 1. with being optimal. We compute the LSI score with httpps//cellrankk. readthedocs http:///ww...ooo/a//ao/////////e att.ap//aaattttuurrraaaao/aa estimatoors/celrrannk www.eettttss...... GPCCA. hww.hheeelaaa......m.emmml www.wanww....iiitaa..e.eoeoom.e. GPCCA. tsi
In the specific dataset considered here, the terminal cell states are Alpha, Beta, Delta, and Epsilon cells Hence, the TSI score measures what extent the underlying vector field captures the dynamics such that cells develop into either of these cell states, while these cell states being a sink. In the CellRank API, set n_ macrostates to 15, i.e., the algorithm looks for 15 macrostates.
Cell type scores/errors As in most single-cell tasks there is no ground truth of matches between cells, we rely labels of clusters of the data, i.e. cell types. We then assess the accuracy of generative model by considering the accuracy of the corresponding classifier felass as described above. The correct matches between classes have be considered task sss--specifically.
In the developmental pancreas example considered Figures 2 and 9 we have {S (Tr) = (Alpha, Beta, Delta, Epsilon, Fev+ Alpha, Fev+ Beta, Fev+ Delta, Fev+ Delta, Fev+ Pyy, Ngn3 High late, Ngn3 High early, Ngn3 low EP). Here, we do not look for one-to-one match, as from a developmental biology perspective, it reasonable (and even expected) that c.g. Fev+ Alpha cell evolves into Alpha cell. To measure to what extent the dynamics follow the manifold, we measure how many cells in one of the terminal cell states Alpha, Beta, Delta, Epsilon are derived directly from Ngn3 low EP cells. These transitions are biologically unlikely, cells are known to evolve via the Ngn3 high cluster and either of Fev+ Alpha, Fev+ Beta, Fev+ Delta or Fev+ Epsilon. of the terminal cell states is very unlikely.
Hence, the direct transition from Ngn3 low EP cells Thus, we denote this fraction as the cell transition error (CT error).
In the following, we discuss the choice of the labels (S&)k and [K&) for perturbation prediction in 3ig Fig T2 Each drug was applied to cells belonging to three different cell types/cell lines, namely A549, K562, and MCF7. Hence, can define {S {T
23

---

{A549, K562, as for each perturbed cell we know the cell type at the time of injecting the drug.
FOSCTTM score In the following, we consider a setting where the true match between samples is known. The FOSCTTM score ("Fraction of Samples Closer than True Match") measures the fraction of cells which are closer to the true match than the predicted cell. Hence, a random match has FOSCTTM score 0.5, while a perfect match has FOSCTTM score of 0.0. In the following only consider discrete distributions. To define the FOSCTTM score for a map T : X + V, let be samples from the target distribution, such that X3 and Yı form a true match. Moreover, let yı = T (x+).. Let
Xi.... жK E be samples from the source distribution and Y1... .... УУ
Pj ΣNEK llye 8, - [K]
(34)
and
9j EKEK 1 I Vt sallisly s
(35)
K
Then, the FOSCTTM score between the predicted target {5x FRek and the target {yx boek is obtained
FOSCTMM((ŠA)eeR (y=]kek) =
Σ
Pj +
4i
(36)
2
kEK
D Datasets
D.1 Gushchin et al. (30]s benchmark datasets
We follow the notebook in wtw.aaaaaaaaaaaaaaaaaaaaaaaaaaaa wwwpoottoooo//iittooooooeeeeeeee
visualization coam__ot..ppp ipynb which allows to down- load the data when instantiating the samplers by setting download=True in the function
D.2 Pancreas single-cell dataset
The dataset of the developing mouse pancreas was published in Bastidas-Ponce et al. 2 and can be downloaded following the guidelines on https://www ncbi nlm. nih, gov/ pnov/gee/queras hwtp/////////////eeeeeeeee goth.g/g////eeer//accc
cgi7acc=GSE1121 88 The full dataset contains measurements of embryonic days 12.5, 13.5, 14.5, and 15.5, while we only consider time points 14.5 and 15.5.
We filter the dataset such that we only keep cells belonging to the cell types of the endocrine branch ensure that learnt transitions are biologically plausible. Moreover, cells annotated Ngn3 high cycling were removed due to its unknown stage in the developmental process (2). The removal is justified by the small number of cells belonging to this cell type and its outlying position in gene expression space. Experiments were performed on 30-dimensional PCA space of log-transformed gene expression counts.
D.3 Drug perturbation single-cell dataset
The dataset was published in 78 We download the dataset following the instructions detailed on https://github. www.///wwwo///////oooooooeeeeeee httpa//eeceeeeooteeeeeeeaaaaee
For all analyses we computed PCA embeddings (50 dimensions) on the filtered dataset including the control cells and the corresponding drug only. This ensures the capturing of relevant distributional shifts and hence prevents the model from near-constant predictions as the effect of numerous drugs is weak.
The drugs in figure were chosen as in Hetzel et
D.4 Human bone marrow single-cell dataset for modality translation
This dataset contains paired measurements of single-cell RNA-seq readouts and single-nucleus ATAC-seq measurements [5]]. This means that we have ground truth one-to-one matching

---

(x,y) The Coulomb
(x T
GENOT (x,y) T
OT-CFM
Figure Fitting the EOT coupling between two Gaussian mixtures for the Coulomb cost c(x,y) = 1/||x - yyl2,, and 0.01, using GENOT and OT-CFM [5]. For both methods, we hence use mini-batch couplings computed with this cost. We connect paired samples with line. The EOT coupling pairs source samples and target samples diagonally (left). GENOT (middle) generates samples correctly, while OT-CFM (right) fails to preserve the signal from mini-batch couplings. The data-setting is inspired by Bortoli et al. [5, Fig. 1]
for each cell. We the processed data provided in moscot 39 which be downloaded following the instructions on htttss://moscooo www.oeahhctteeraccecccccc.ccce.e datasets. bone, marrow. html#moscott datasets. bone, marrow This version of the dataset additionally contains shared embedding for both the RNA and the ATAC data, which we use the fused term. This embedding was created using variational autoencoder (scVI [49] by integrating the RNA counts of the gene expression dataset and gene activity [79] derived from the ATAC data, commonly used approximation for gene expression estimation from ATAC data 34
readthedocs. http://n/aatesstgaapp///a///a/ep mmmps//
In RNA space we use the PCA embedding (the dimension of which is detailed in the corresponding experiments), while the embedding used in ATAC space is the given LSI (latent semantic indexing) embedding, followed by a feature- wise L2- ..-nommalization as proposed Demetci et al. [17]. The legend for the cell type labels is given in Fig. 25
E Additional information and results for experiments
not stated otherwise, the GENOT model configuration follows the setup described in appendix
E.1 (U)-GENOT-L
Entropic OT benchmark by Gushehin et al. [30] We assess the performance of GENOT-L following the benchmark in Gushchin al. [30]. While the results reported Gushchin et al. (30 are the best ones across different configurations of hyperparameters for certain models, we choose single hyperparameter configuration across all dimensions and entropy regularisation parameters. In particular, GENOT is trained for 100,000 iterations and the learning rate is set to 10-5, The dimension of the layers in the vector field is set 1024, and the entropy regularization parameter is set according to the task (i.e. 0.1, 1.0 or 10.0.). The batch size is set to 2048, and its influence is further studied in figure The number of layers per block is set 3, and the embedding dimension of the time, condition, and noise is set to 1, 1,024. All remaining parameters are chosen as described in appendix G
Following Gushchin et al. (30], Tab. reports the wasse-raasenntiin Unexplained Variance Percentage (BW2-UVD between the marginal distributions p2ime and defined
100%
BW3-UV UVP (patmo, v)
BW? (p2tTo,
(37)
1/ Var (v)
25

---

where BW3 (P2tTa, mτ) is the Warrs-Waasersttinn metric between the Gaussian approximations of the distributions PT and v:
BW2 (pedmo, = W (V (m,,Ey), , (mpztno Mpstse))
(38)
Here, W denotes the Wasserstein-2 distance, and for any measure 7, my denotes its mean and Σn its covariance. Tab. E displays the conditional Bures- Wasserstein Unexplained Variance Percentage (cBW-UVP), which is the integration of the Bures- Wasserstein distances of the estimated and the true conditional distributions:
cBW2- (To, m Var (v) / BW2 (TT( (+1x)) (u(x).
100%
(39)
Algorithm 2 U-GENOT with stratified sampling of the discrete conditionals. Skip teal steps for GENOT.
1: Require parameters:
Batch size n, number of (per conditional sample k entropic regulariza- tion E, unbalancedness parameter 7, discrete solver Solver,.r.
2: Require network: Time-dependent conditional velocity field vz,a(++)) R2 x Rº R4, re-
weighting functions Vle IRP + R, &a I29 - R.
3: for 1..., Ther do
4: XI...., Xn ~ μ and II...., Yn V.
5: Pe,r t Solver,, ((x;) (Yi)==) E RT
6: a + Pe In and b - In.
7:
Sample i,, . in ~ and set (I,,..,X,n) + (Xip...,,
8:
for 1..... do
9:
Sample .1.... ln P 5.r [.
10:
Sample noise vectors Z 2,.,..., Zy p and time-steps f1,1....,, ti, u(10, 1))
11:
end for
12:
L(O)
E 1 Ei= ( Y,,((/,,) (Yi:
Zr.1)113
+1 D - (no (X,) mak
2 + (Y, nbb)²²
13:
0 - Update(0, VL))).
end for
Hence, BW3 UVP (petmo, v) measures the fitting property of the push-forward distribution while cBWs- (T,, measures the fitting property of the joint distribution. Hence, Bw? UVP (patTe, v)) measures the fitting property of the push-forward distribution while cBW3-UVP (TO, π measures the fitting property of the joint distribution.
The values for the competing methods are taken from the original manuscript Gushchin et al. [30 Competing methods are Seguy et al. (73 (LSOT), Daniels et al. [15] (SCONES), Korotin et al. 40 (NOT), Mokrov et al. (55] (EgNOT), Gushchin et al. (29 (ENOT), Vargas et al. [87] (MLE-SB), De Bortoli et al. [6] (DiftSB), Chen et al. [10) (FB ADEAA, FB-SDE-J).
To make the results of the benchmark more interpretable, we rank the methods for each experimental setup. If multiple methods yield the same result, we assign the methods the mean of the ranks. Tab. reports the average rank across the twelve different experimental setups.
Ablation study of GENOT A natural way to extend the GENOT algorithm is to sample k > 1 instances from the conditional distribution. Alg. describes this procedure in lines & to II resulting only in an additional sum in the loss in line 12 As for each condition X; we then have multiple samples (Z,, k and (Y,, 1
} (in the unbalanced case depends on i), we can then couple the noise samples with the samples from the conditional distribution using discrete Sinkhorn solver (note that the noise and the conditional distribution always live in the same space). This corresponds to using the approach introduced in Pooladian et al. [64], Tong et al. [85] for each single conditional generator (xx).. We study this natural extension of the GENOT model in Fig 7 While in a few cases this extensions helps, in the majority of experiments using & decreases the performance.
GENOT-L for trajectory inference Discrete optimal transport is an established method for trajectory inference in single-cell genomics (72). Due to the ever increasing size of these datasets
26

---

Model
rank BW2 -W2-UVP UVP
rank cBW2 UVP
LSOT
11.33
11.33
SCONES
8.67
9.33
DifIS
8.92
9.50
FB-SDE-J
6.04
9.33
FB-SDE-A
8.08
8.66
EgNOT
3.04
7.08
DSBM
8.83
6.0
NOT
5.42
3,66
ENOT
6.92
3,58
CFM-SB
5,54
3,42
MLE-SB
3.33
3.25
GENOT-L
1.88
2.83
Table Average ranks of each method across experiments, on estimating the ground-truth EOT cou- pling क5 between each pair of measures μt, in dimension d E {2, 64, 128] and E {0.1, 10) of Gushchin et al. [30]ss benchmark. As proposed by the authors of the benchmark, we measure performances with both BW2 UVP (37) and cBW?-UW UVP detailed results.
39) metrics. See Table 2 and Table for
Model
0.1
J0
-(11.)
-(11.5)
((12)
-(12)
-(10.5)
(110.5
-(10.5)
SCONES
(11.5)
-(11.)
-(1..5)
-01.5)
1.06
2.08 (6)
0016:33
(7)
2.603
(4)
FeNOT
(6)
0.31 (5)
0.22
(7)
0.3.4)
0.85
0.12(0
(1)
0.00 (1.5)
0.23
0.2(7)
(.8(7)
04466
7.8(10)
(10)
1:3.85
(.(4)
ML.F-BB
(59
1.02
(5)
-00.5)
((0.5)
-00.5)
-10.5
RI-SDE-A
2.37.(8)
2.55
19 (9)
27.11 (9)
0.6 (8)
0.83 (7)
0.65
0.71
((10.5)
(10.5)
-(10.5
-01.5)
FR-SDE-J
(4.5)
0.05
0.25
1.95
0.07 (2)
(3)
-(1.5)
(10.5)
-11055
-00.5
CFM-SB
(4.5)
0.30
(355
0.99
(53
17.50.(7)
(185
54.23 (10)
(8)
2.12.111
11.38 (11)
19.27 (11)
39.19
(7)
632.58
GENOT
(100
0,015 (1)
0.13
0.33 (2)
0.085 (4)
0.06
0,08 (3)
BL18
0.10
0,02 (1.5)
0.50 (2)
0U47 (2)
Table 2: BW3-UVP (equation (37) between Tं and patrr between the distributions defined in Gushchin et al. [30]. with numbers in parantheses denoting the rank within the experimental configu- ration defined by the dimension of the space and the entropy regularization parameter
0.1
1.0
10.0
Dimenson
10.0
60
nar 2.0
10
2.0
aBW?-0 UVP
100
$4.0
80
20
10
214
batchh siza
batcch
size
batch
1520
Figure 6: cVW3-UVP (37) between n5 and ma between the distributions defined in Gushchin et al. 30 depending on the batch size. The dotted line corresponds to batch size 2048 and hence displays the values reported in Tab.
B1. neural OT solvers are of particular interest and deterministic Monge map estimators have been successfully applied millions of cells [32.
Obtaining samples from the conditional distribution allows for an assessment of the uncertainty of the trajectory of cell. We use the metric Vary~<o(-1xx (o0ssssooosssmm [oos-sim((: Byy [Y]) suggested in
Gayoso et al. (25) for generative RNA velocity models (appendix C.T). Therefore, we use 30 samples from the conditional distribution.
27

---

Mede
--0.1 0.1
10
2
15
64
125
16
64
126
2
16
64
125
LSOT
-(1.5)
-(11.5)
(11.5)
-01.5)
-(12)
-(12)
-(12)
-(12)
-(10.5)
-(10.5)
-110.5
-(30.;
SCONES
-(1.55
-(11.5)
(31.5)
-01.5
34.88 (10)
71.34 (10)
59.12 (2)
136. (9)
32.9 (8)
50.54:00 (7)
60.44.17
$2.31 (6)
NOTT
(34
13.67 (3)
11.74 (1)
114
4.77 (3)
23.27 (6)
41.75
26.56 (4)
2.86 (3)
4.57 (5)
3.41 (3)
6.56 (3)
EENOT
(29.5(10)
(5.2
(7)
4322 (7)
90.4 (31)
(1.4
63.8 (19)
53.2.6)
4.14 (6)
2.64 (4)
2,36 (2)
1,31 (2)
ENOT
3.65(8)
22.(5)
13.6
12.6133
(1)
9.464)
21.6.64
(3:55)
L.4
2.4
19.6
30
MLE SB
4.57(5)
16.12(4)
16.1
11,81 (4)
4,13.02
9.08 (2)
18.05 (31
15,226
1.0. (2)
1.27.(1
3.900)
129.0)
Dinsa
73.54
397.09
1386.4(10)
33.76
70.86 (9)
53.42.8)
(10)
-(10.5
-C10.5
-(105)
(9)
53.2
30.62
63.48.77)
131.72 (8)
-(10.5
-00.5)
-10.5)
-00.5)
FB-SDEJJ
(7)34(7)
59.16 (10)
(19.32
(73.96(8)
29.34
(9.2
135.14 (311)4-(1))
(77.52.
-(10.5)
(30.3)
-10.5)
-30.3)
(3)
(1)
7.17
(39
4.76
5.26:11
10.38 (7)
(6)94
(0
$7.39 00)
41.52.00
5.78
17.12
(5)
60.51
2.95 (4)
35.59
(42.1((8
Independent
176.
72.0
2.76
GENOT
3.24
1.97
Table 3: cBw? -VP 37 between मे and T. between the distributions defined in Gushchin et al.
80. with numbers in parantheses denoting the rank within the experimental configuration defined by the dimension of the space and the entropy regularization parameter E.
0.1
1.0
10.0
dAnl dAnn-Ma 1A0
16 64
128
35 64
126
2
15
54
123
Figure 7: Evaluation of the influence of the number of data points sampled from the conditional distribution in each step, and optimally coupling noise and samples from the target distribution, see Alg. k corresponds to the model performances reported in table
Fig. 8 visualizes the pancreatic endocrinogenesis dataset (.2 14.5) and target distribution (15.5) as well as according to cell type. To quantitatively confirm the visual results from Fig. 2 aggregate the uncertainty to cell type level. Indeed, we observe higher uncertainty in cells which have not committed to a certain lineage (Ngn3 low, Ngn3 High early, Ngn3 High late), while more mature cell types indicate lower variances.
according to source (embryonic day
Sourcs and target disteibution
Coll typos
Unoetainly sooee grodomic ooe
0.05
Must
CO0- TERT
Feri Cyniee
0.0
Aloba
Bors Dois
Base
Newd Ave
easu Mech
For
Figure 8: UMAP of the mouse pancreas development dataset colored by sample (left) and cell type (middle). We transport samples from embryonic day 14.5 to embryonic day 15.5. Uncertainty score as displayed in Fig 2 aggregated to cell type level.
cell is expected to have an uncertain trajectory when it can still develop into different cell lineages. In contrast, cells are expected to have less uncertain trajectory when their descending population is homogencous or they belong to terminal cell state, and hence have committed to a certain lineage.
GENOT-L produces meaningful uncertainty assessments Indeed, the Ngn3 low EP population and the Ngn3 high EP populations have a higher uncertainty than later cell types. The uncertainty is particularly high for Ngn3 high EP late cells, state where cells commit to a fate towards alpha, beta, delta, or epsilon cells (2).
be seen from Fig. and Fig. g
We now focus on the conditional distributions of cells from the Ngn3 low population, i.e. the least mature cell population in the dataset. Fig. [hsows more examples of selecting one Ngn3 low cell and generating 30 samples from its conditional distribution displayed in Fig 2 It is important to
28

---

X~ N
Y~ To (.)
dx (1+x Y ~ B
0.40
Ngn3 Low
Proportion originating
0.0
dM
Ngn3 High late
0.25
0.0
Alpla
Beta
Delts
Epsilon
Figure 9: Left: Samples from the Ngn3 Low EP population in the leftmost column, and next to 30 samples from their conditional distribution samples from GENOT with the geodesic cost (middle) and the squared Fuclidean cost (right). Right: Proportions of alpha, beta, delta, and epsilon cells derived from the Ngn3 Low population (top) and the Ngn3 High population (bottom). Mean and standard deviation across three different runs are reported.
note that the UMAP has to recomputed to include the generated samples and hence the UMAPs change slightly. The three instances Fig 9 support the hypothesis that the conditional distribution follows the data manifold when using the geodesic cost (here we assume that the UMAP represents the manifold well). In contrast, choosing the squared Euclidean distance results in samples of the conditional distribution being either very close on the UMAP or very far away, while the intermediate stage (i.e. Ngn3 High late) skipped. Biologically, this result is not plausible as all cells are known to evolve from Ngn3 Low EP cells to Ngn3 High EP cells, before they finally commit to lineage (towards alpha, beta, delta, or epsilon cells). important to note that both GENOT-L models are trained with the same entropy regularization parameter 0.01.
Based the examples, hence hypothesize that when using the squared Euclidean cost, generated cells in the terminal cell states (alpha, beta, delta, epsilon) are much more likely to be derived from Ngn3 low EP cells than Ngn3 High EP cells (in the following neglect Ngn3 High EP early cells as their number very low, see Fig 88. As outlined above, this is biologically not plausible. The impression conveyed by the UMAP plots confirmed numerically on the right hand side of Fig. 9 For each generated cell, we compute its nearest neighbor in the target distribution, and this way assign it label (cell type), App. C.2 We then obtain cell type cell (with columns containing the cell types of the
focus Ngn3 High Late cells and
type transition matrix, which we generated cells), such that for each cell type, we obtain the distribution of ancestor cell types. Figure 9 shows that for the Euclidean cost, around 30% of all delta cells are derived from Ngn3 Low EP cells, while only 5% are derived from Ngn3 High EP cells. In contrast, the proportion of progenitors in Ngn3 low EP cells and Ngn3 high EP cells when using the geodesic cost is much more comparable. Similar results can be observed for alpha, beta, and epsilon cells.
1d simulated data Figure 10 visualizes the influence of 7 simple setting between mixtures of Gaussians. While 7 1.0 corresponds to the fully balanced case, setting т = 0.97 results in a complete discardment of one mode in the target distribution. The ground truth is computed with discrete entropy-regularized OT solver [14), with kernel density estimation to create the plot. This figure is meant visualize the task, while in the following, assess the results quantitatively.
Benchmark of U-GENOT-L between Gaussian distributions Therefore, we leverage the closed- form solutions of the unbalanced EOT coupling between Gaussian distributions for the squared Euclidean cost [36]. We follow the data generation process described by Janati et al. (36), and evaluate the results for different entropy-regularisation parameters and different unbalancedness
parameters À I-T: and across different dimensions.
29

---

1.0
0.995
7-
j
vv v
Eu MM
T
T
0.99 1
T 0.97
Target
T
n.
Yp +
8-
Eo
Figure 10: Unbalanced entropic neural optimal transport plan with e = 0.05 and varying unbalanced-
ness parameter Ti T2.
8 1,0,= 1.0
1.0, À 0.5
U-GENOT-L SUOT
0.5
log BW 03 2/ X²*
BW (R, IT")
005
1.0
1.5
2.0
22
23
2
22
28
27
2
= 0.1, 1.0
0.1, À 0.5
log BW (1a & E4 nT"
BW RR 71")
-0.2
-0.4
-0.8
-1
-1.0
2
23
21
22
23
24
2
2
Dimension
Dimension
Figure 11: Benchmark U-GENOT-L on the task of learning an unbalanced entropic OT plan between Gaussian distributions across different dimensions, entropy regularisation parameters and unbalancedness parameters X, following Janati et (36). The dotted line represents the outer coupling, while SUOT denotes the method proposed in Yang and Uhler [93].
Yang Uhler 93
propose a method (scalable unbalanced optimal transport, SUOT) to learn unbalanced EOT plans, but their implementation only supports unbalanced deterministic Monge maps. We visualize the results in Fig TT Besides the simple switch learning unbalanced EOT plans, U-GENOT-L was with the same setup as in the balanced linear EOT benchmark.
30

---

Perturbation modeling with GENOT-L and U-GENOT-L For each drug, we project the single- cell RNA-seq readout of the unperturbed and perturbed cells to a 50-dimensional PCA embedding. Subsequently, we split the data randomly to obtain train and test set with ratio of 60%/40%. This preprocessing step holds for both the calibration score experiments and the experiments conducted with U-GENOT-L to assess the influence of unbalancedness to the accuracy score.
The uncertainty score for the calibration study computed based on 30 samples from the conditional distribution, see appendix
117
Number of drugs
-0.5
0.5
1.0
Calibration score
Figure 12: Calibration score for the predictions of GENOT-L for modeling cellular responses 163 cancer drugs.
In Figure 13 and Figure 14 we visualize the influence of unbalancedness for perturbation modeling with Dacinostat and Tanespimycin, respectively. These experiments were conducted on the full dataset for visualization reasons. While the fitting property seems to be little affected by incorporating unbalancedness (top rows), the cell type clusters are better separated for U-GENOT-L. transport plans than for GENOT-L transport plans.
1.0
0.8
Ungeterber cub and bosttd with Deciodie
BAD MA DA
MIT DALUM INPLRRN
Nodile
LROP
UNAr
Berisate sepertanet
UP3F2
A349
K362
C56O
MEF?
UMAP
Figure 13: Visual assessment of the influence of unbalancedness in modeling cellular predictions to the cancer drug Dacinostat. In the left column, the source and target distribution are jointly plotted with cells colored by whether they belong to the source (unperturbed) or the target (perturbed) distribution (top), and which cell type they belong to (bottom). In the center column, we plot a UMAP embedding of target and predicted target distribution. The top plot colors cells according to whether a cell belongs to the target distribution or the predicted target distribution. The bottom plot is colored by cell type. The cell type of the predicted target distribution the cell type of the pre-image of the predicted cell. The right column visualizes the same results, but this time obtained from U-GENOT-L with unbalancedness parameters 7 = Ti = 72 = 0.8.
31

---

== 1.0
0.9
UPAP
JPA
anperturted
DROP
Celtypee
UNAF2
2049
A249
AC4R
R362
LMAF1
DSS2P3
LPAAP
Figure 14: Visual assessment of the influence of unbalancedness in modeling cellular predictions to the cancer drug Tanespimycin. In the left column, the source and target distribution are jointly plotted with cells colored by whether they belong to the source (unperturbed) or the target (perturbed) distribution (top), and which cell type they belong to (bottom). In the center column, we plot a UMAP embedding of target and predicted target distribution. The top plot colors cells according to whether a cell belongs to the target distribution or the predicted target distribution. The bottom plot is colored by cell type. The cell type of the predicted target distribution is the cell type of the pre-image of the predicted cell. The right column visualizes the same results, but this time obtained from U-GENOT-L with unbalancedness parameters T = T1 = 0.9.
E.2 (U)-GENOT-Q & (U)-GENOT-F
GENOT-Q for mapping Swiss role to a spiral Fig. 15 visualizes the dependence of the con- ditional distribution on the entropy regularization parameter & with the same setup as in Figure 3 In particular, we display what happens when we consider the outer coupling. Fig. I6 visualizes the density in the target space, evolving from a 2-dimensional Gaussian distribution at 0 to conditional distributions at 1. Analogously to Fig 15 regularization parameter E.
we visualize the influence of the entropy
5 0.005
8 0.01
8 0.05
8 +
Figure 15: Samples from conditional distributions To(-Ix) for GENOT-Q models trained with different entropy regularization parameters 8. The setup is the same as in Figure 3 in effect we transport a three-dimensional Swiss roll to a two-dimensional spiral, which is colored in blue (with high transparency). The right plot shows the result when choosing the outer coupling as opposed to an EOT coupling. The source distribution as well as the data points which are conditioned on are visualized in Figure
How to handle non-uniqueness of quadratic EOT couplings In the following, we show how to empirically overcome the non-uniqueness of solutions to (QOT). The left hand side of Fig. 17 visualizes the solutions to an entropic GW coupling between two Gaussian distributions learnt by GENOT-Q without any modification. It becomes evident that the learned solution is not close to valid solution, but rather resembles a mixture of solutions.
32

---

1 To Ty
To To(y|Xg (++)
0.005
Card deraltyatt-u
deurts 234-
dearty
Searty
Card deribyatt-ouu
Card
Card att-644
0.05
Www.tyy Mi
Meailyy a1-0-11
Juaiyy #1110
Zard dealdy attutts
Dar Bearly di
Fand om aris
denity attaa
dually amnas
duaiy a0:0
pard deetyy at1-
card dealdy att-oox
Card Bealdy ath-44
8
T
Figure 16: Evolution of the density from t 0 corresponding to the latent distribution (i.e. mul- tivariate standard Gaussian distribution) to = corresponding to the density of the conditional distribution ( X,). The setup is the same in Figure in effect transport ttree-dimensional Swiss roll to two- dimensional spiral The left panel displays the densities of one conditional distribution, while the right panel visualizes mixture of densities over time. The top part of the figure corresponds to GENOT-Q models with entropy regularization parameter = 0.005, while 0.005. The bottom correspond to the outer coupling. The
the middle ones corresponds markers on the 3-dimensional plot on the very top correspond to data points in the source distribution.
We discuss three ways overcome the limitations non non- non-uniqueness of quadratic OT couplings:
Incorporation of a fused term: Whenever the data allows, using a fused term empirically helps to stabilize the solution.
The set of solutions to (QEOT) depends the cost function. Thus, finding a cost with a small set of solutions overcomes the limitation to certain extent. Empirically, the geodesic cost helps do for single-cell data.
Changing the cost function:
Initialisation of the discrete solver: We initialize the discrete GW solver at iteration + using the solution to the linear EOT problem between p2 and v, with #J denoting the learnt coupling step
Visually, the initialisation scheme helps for learning the GW-EOT coupling between Gaussian distributions (I7).
We now quantify what extent we can reduce the diversity the solutions by considering the variance of the conditional distribution of single data point. Therefore, we consider two datasets:
30 dimensional Gaussian distributions (20 dimensions in the quadratic term and dimen- sions in the additional fused term): For both the source (μ) and the target (v) distribution
33

---

Figure 17: Left: Learnt transport plan of GENOT-Q without initialization between two one- dimensional Gaussian distributions. Right: Learnt transport plan of GENOT-Q with initialization. The initialization scheme prevents GENOT from learning mixtures of valid transport plans.
sample mean vectors from U([-1, and construct the covariance matrices by sampling each entry of their square roots from (4(1 3])
2. 30 dimensional single-cell data (as in 4 20-dimensional PCA of the normalized ATAC data, 20-dimensional PCA projection the gene expression data, and 10-dimensional PCA of the VAE embedding (shared space).
To assess the stability of the solution of discrete EOT, we fix one sample from the source distribution, and compute the variance of R- [10 :)), i.e. of the incomparable space only, which is obtained from a discrete EOT solver with different samples in source and target. We do this for e = 0.01 and 0.001 with the squared Euclidean cost.
We consider the following couplings:
1. outer coupling, supposed to serve as a baseline.
2.
(GW): We X and 5, i... we use the last 20 dimensions of both the source and the target distribution.
(FGW): We Ω X and Ω Ý with 1 4. Gromov-Wasserster with geodesic cost (Geodesic); We use X' and 5 with the geodesic cost,
3. Fused
5. wasser-Wasserttein with initialization described above.
Fig 18 shows that both the fused term and the initialization scheme empirically help to preserve the variance of the conditional distributions. Moreover, the geodesic cost helps for small to reduce the conditional variance, suggesting less diverse set of solutions to (QEOT.
GENOT-Q und U-GENOT-Q for learning EOT plans between Gaussians To quantitatively assess the performance of GENOT-Q, we evaluate its performance by learning quadratic EOT couplings between Gaussian distributions [43. We follow the authors' data generation process and computations for the closed-form solution and the computation of the discrete GW solver. Note that we can only evaluate cost of the learnt transport plan as solution of the EOT plan is not unique. Thus, we use the initialisation scheme suggested above. The discrete solution calculated based on the discrete coupling (1000 samples), while we report results from GENOT-GW on both training data (30k samples) and test data (5k samples). GENOT-GW is trained with batch size 1024 and for 5k iterations, with the same network architecture as used for the benchmark for balanced, linear EOT for the squared Euclidean cost ($ F.T. As baseline, we report the cost induced by the outer coupling. To keep the results on a similar scale across different dimensions, we report the L1-distance between the predicted cost and the true cost, divided by the true cost.
continue with evaluation in the unbalanced setting in Fig. 20 Here, we compare only the costs induced by the distortion, i.e. not the one induced by the KI. penalisation, as for the latter one requires access to the densities. As demonstrated, GENOT can estimate the densities, but comes at the cost of an estimation emrrr As the purpose of this benchmark is to show the accuracy of the learnt
34

---

Gaussians, 6 0.001
Gaussians, 0.0001
Conditional Veriance
122
220
12
Single-cell data, 0.001
Single-cell data, - 0.0001
Conditional Variance
Conditional
oeter coepting
aaarr
areen derdeeic
GinD with enttaiiorios
Aasetine
ooter coupling
amennaa Desdeer
Geser with pnlitivation
Coupling
Discrete Coupling
Figure 18: Variance of conditional distributions of different couplings for = 0.001 (left) and 0.0001 (right) for both Gaussian data (top) and single-cell data (HSCP also considered in e.g.
Fig.
Error bars denote standard error across dimensions of the 20-dimensional space.) (bottom)
= 0.1
8 0.5
cOsterue
10
train GENOT-GY test
outer coupling
discrete solver
costtrue
COStextimatoo
6
5.0
1.0
[oStestimmtton costerue
1.0
true
0.8
0.6
2.1
4.2
8.4
16.6
32
64.32
16.8 Dimensions dx _dy
Dimensions d. -dy
Figure 19: Learning Gromov-Wasserstenn entropic optimal transport couplings between Gaussian distributions where the ground truth plan is known for 433. The quality of the learnt coupling is assessed with the induced cost.
couplings (and not the estimation of the density), compare based the distortion cost. As done by the authors Le et we set the total mass of both distributions to 1.0.
Modality translation with GENOT-Q For experiments, we perform a random 60/40 split for training and test data. All results reported on the test dataset. The cost matrices of all models
35

---

0.05, À 10.0
E 0.1, 10.0
costagtinated cOsterue
GENOT-GA GOW train
10a
GENOT- test
10a
outer coupling
costtruc
discrete solver
102
10
10a
10
10
10
4.2
8.4
0.5, À
10.0
1.0, X 10.0
rrottgstimnted costtrue
10a
0088 buture
10-
10a
10a
10a
10
10
10
10a2
2.1
4.2
0.4
16.0
32
64.32
16.0
32
Figure 20: Learning unbalanced -aasoosensss entropic optimal transport couplings between Gaussian distributions where the ground truth plan is known for 43 The quality of the learnt coupling assessed with the induced cost.
were scaled by its mean and the entropy regularization parameter & was set to 0.0001. Moreover, the models were trained for 5,000 iterations.
Modality translation with GENOT-F For all experiments, we perform a random 60-40 split for training and test data. All results are reported on the test dataset. The cost matrices of all models were scaled by its mean and the entropy regularization parameter e was set to 0.001. Moreover, the models were trained for 20,000 iterations.
Figure 21 21 reports results of the GENOT-F model with interpolation parameters a = 0.3 (left) and 0.7 (right). While the Sinkhorn divergences not comparable with results of the GENOT-Q model due to the respective target distributions living in different spaces, we can compare GENOT-Q with GENOT-F with respect to the FOSCTTM score. Figure 21 shows that GENOT-F strikingly outperforms GENOT-Q, hence the incorporation At the same time, it is important mention that the GW terms add valuable information to the problem setting, which can be derived from the results for GENOT-F with a = 0.3. Here, the higher influence of the fused term causes the model to perform overall worse.
the fused term is crucial for a good performance.
53 To demon- strate the robustness of our model, train GENOT-F model with 0.01, = 0.5 and the Euclidean distance on 60% of the dataset (38 dimensions for the ATAC LSI embedding. 50 dimen- sions for the RNA PCA embedding, and 28 dimensions for the VAE embedding in the fused term) and evaluate the learnt transport plan visually. Figure 22 shows joint UMAP embeddings of predicted target and target distributions, the full legend of cell types can be found in Figure 25 Qualitatively, a good mix between data points of the predicted target and the target distribution suggests a good fitting term. Optimality of the mapping can be visually assessed by considering what extent cell types are mixed (low optimality) or separated from other cluster (high optimality).
Moreover, we can visualize the optimality and fitting term in a UMAP embedding
We observed that taking the conditional mean improves results on the FOSCTTM score, but can impair the fitting property. Indeed, the mixing rate between data points belonging to the target and data points belonging to the predicted target seems to be slightly worse when considered in the joint embedding as well as when considering only the fused space and only considering the quadratic space
Modality translation with U-GENOT-F To simulate setting where there is not a match for certain cells in the target space (RNA), remove the cells labelled as Proerythroblasts, Erythroblasts,
36

---

0.3
Q: 0.7
Squared Euclidean
Geodesic
Squared Euclidean
Geodesic
Mathod
Mithoo
FOSCITM
GENOT-2 GENOT-P CM PGW-LR
GENOT GENOTL I' CM
PGW-LR
divorgesoss
10c
10%
10%
10/10/15
10/10/25
36/50/5
36/00/25
10/10,75
10/10/20
38/50/6
38 28/000/2
10/10/5
10/ 10/10/25
35/50/03
38/R0/25
10/10/5
38/50/5
34/000/15
Dirscuaions
Dimensloms
Figure 21: Mean and 95% confidence interval (across three runs) of the FOSCTTM score (top) and the Sinkhorn divergence (bottom) of GENOT-F and discrete FGW with linear regression for out-of-sample estimation. Experiments are categorized by the numbers /d2 /d.. where di the dimension of the space corresponding to the GW of the source distribution, d2 is the dimension of the space corresponding to the GW target distribution, and d3 the dimension Results are reported for the interpolation parameter 0.3 (left half) and a 0.7 (right half).
shared space.
UMAP based on x ==)=y'
UMAP on I CC only
Deferne by tngelpped terge.
Coloret by
UNNP
UNND
type
UNAP
UMAP based y y (GENOT-F CM)
UMAP based &2 C y only
Celered by la ha large.
ope
LMAP
Figure 22: UMAP embedding of target and predicted target based on subspaces of the target domain. Top left: full target space. Top right: incomparable space only. Bottom left: full target space (with predicted target computed by the mean of the conditional distributions). Bottom right: share space only. Cells are colored based on whether they belong to the target distribution or the predicted target distribution in the left plot of each subplot. Cells are colored based on their cell type in the right plot of each subplot. For cells which belong to the predicted target distribution, the cell type defined as the cell type of the preimage. Shown are results on the test data set, corresponding to 40% of the full dataset.
and Normoblasts these cells form lineage, developing into mature Reticulocytes (not present in the dataset). Thus, they are similar in their cellular profile while being clearly distinguishable from the remaining cells. Note that this simulates extreme setting of class imbalances, these cells make up for more than 20% the dataset.
While we keep the right marginals constant, as we have true match for each cell in the target distribution, we introduce unbalancedness in the source marginals. It is important to note that the influence of the unbalancedness parameters are affected by the number of samples, as well as the entropy regularization parameter e. demonstrate the robustness of GENOT-F with respect to hyperparameters. we still choose = 0.7, but this time set 5 10-3 We use 50-dimensional PCA-space the Gromov term in the RNA space, 38-dimensional LSI-space for the Gromov term in the ATAC space, and 30-dimensional VAE-embedding for the shared space.
37

---

Source: ATAC space
UNAP MNDC
UMAP ATAC
Target: RNA space
Cell
Conditional density
16095
RINA
UMAP RMA
28500
66035
UMAP
UMAP RNA
Figure 23: Top: A single Erythroblast cell is highlighted while all other cells the source distribution are depicted with increased transparency. Bottom: Target distribution colored by cell type (left) and colored by the conditional density (conditioned on the highlighted cell above) evaluated on all cells in the target distribution. All cells with relatively high density belong to the Eythroblast cluster, hence the mapping learnt by GENOT-F is meaningful.
model (T1)
Normoblast
Erythroblast
Proerythroblast
other
FOSCTTM
Discrete UFGW (0.8)
0.788
0.820
0.842
0.945
0.258
U-GENOT-F (0.8)
0.622
0.733
0.894
1.077
0.131
Discrete UFGW (0.3)
0.591
0.586
0.734
0.761
0.311
U-GENOT-F (0.3)
0.295
0.430
0.554
1.186
0.162
Table 4: Mean value of the rescaling function per cluster for U-GENOT-F and discrete unbalanced FGW together with the FOSCTTM scores three runs. Tab. reports the variances for the GENOT-Q models.
The computation of the growth rates for the discrete setting is described in appendix H2 a random 60-40 split to divide the data into training and test set. The FOSCTTM score only considers those cells which have true match, i.e. cells in the source distribution belonging to the Normoblast, Erythroblast, and Proerythroblast cell types are not taken into account as their true match was removed from the target distribution.
We perform
We assess the performance w.r.t. the FOSCTTM score to ensure that the model still learns meaningful results, and consider the average reweighting function per cell type (appendix E.2) We consider two values (r₁ = 0.8 and 7 = 0.3) of the left unbalancedness parameter, while 72 1.0 as for every cell in the target distribution there exists the true match in the source distribution. Tab. shows that U-GENOT-F learns more meaningful reweighting functions than discrete UFGW as the average rescaling function the left-out cell types is closer O, while the mean value of the rescaling function on all remaining cell types ("other") is closer to At the same time, U-GENOT-F yields lower FOSCTMM scores and hence learns
optimal couplings.
Tab. shows the variance across three runs, demonstrating the stability of both the learnt rescaling function well as the performance with respect to the FOSCTTM score.
38

---

model (71)
Normoblast
Erythroblast
Proerythroblast
other
FOSCTTM
U-GENOT-F (0.8)
3- 1010
10-5
1.10--
5 10-3 3
10-3
U-GENOT-F (0.3)
2- 10
9 10
5. 6.10-4
8 1010"
9 9-10-5
Table 5: Comparison of reweighting functions learnt by U-GENOT-F and discrete unbalanced FGW.
Fig. 24 shows the mean and the 95% confidence interval of the learnt growth rates per cell type. First, it is interesting to see that Normoblasts have the lowest mean of rescaling function evaluations (for both discrete UFGW and U-GENOT-F), which due to them being most mature among the left cell types and hence being furthest away in gene expression space / ATAC space from the common origin of all cells, the HSC cluster. Moreover, it is obvious that the 95% confidence interval of the reweighting function (across cells in one cell type) is much smaller for U-GENOT-F than for discrete UFGW. This is desirable cells within one cell type are very similar in their ATAC profile.
T1 0.8
0.3
U-GENOT-P
U-PGW
1/8 (X)
cone ca
Figure 24: Comparison of learnt growth rates of discrete UFGW and U-GENOT-F aggregated to cell type level for unbalancedness parameters 71 0.8, T2 = 1 (left) and T. 0.3, T2 = (right).
mypkod
Prserythenbliet
Figure 25: Complete legend of cell types for Fig.
Fig. 22 and Fig
23
39

---

Competing methods
F.1 EOT benchmark in the balanced, linear case
As results for competing methods were taken from Gushchin et al. [30]. we refer to the original publication for details the competing methods. For SB-CFM we follow the setup the au- thors' tutorial in http://ttt.....coooooooo/o//oood
single-ggggsee cell ongrrpee DSBM, we follow the setup in https://gt thub. com/yuyan shi/sshi/ir
examples stagssssingii
httecch//oommmm///aaaea DBLLL//DBMMMaaass Gaussian. py, but in- crease the number of inner iterations from 10 50, the number of outer iterations from 40 to 80, and the btach size from 128 to 2048 account for higher dimensionality.
F.2 Regression for out-of-sample data points
Out-of-sample prediction for GW has been considered [ Yet, their methods rely on orthogonal projection, which only works if both the sample size and the feature dimensions are the same in both spaces. Hence, we rely on barycentric projection for in-sample data points. For out-of-sample data points project data point onto the training set and apply the barycentric projection to the linear combination of points the in-sample distribution. Let X Rn be the matrix containing in-sample data points.
Then, for data point in the source distribution Rdd let
3 = arg min ||* x" 312
(40)
BER
where the sum is taken over the n in-sample data points. Moreover, let p. = Ex= Ily Then, the barycentric projection of point in the source distribution given as
11y5 e..
(41)
Similarly, we can apply this procedure estimate rescaling factors in the unbalanced setting. To ensure non-negativity of the rescaling function, we perform regression with non-negative weights:
à = arg min oER
| - xT
(42)
To estimate the rescaling function for data point 2, the estimated left rescaling function is given as
Σ aank ER
(43)
where {n is the set of reweighting function evaluations of in-sample data points.
40

---

G
Implementation
The GENOT framework is implemented in JAX 66. We use the discrete OT solvers provided by JAX- [4]
G.1 Parameterization of the vector field
The vector field is parameterized with feed-forward neural network which takes input the time (which is cyclically encoded with frequency 128), the condition (i.e. the samples from the source distribution) and the latent noise. Each of these input vectors are independently embedded by one block of layers before the embeddings are concatenated and applied to another block of layers, followed by one output layer.
stated otherwise, one block of layers consists of 8 layers of width 256 with silu activation functions.
G.2 Parameterization of the rescaling functions
Rescaling functions are parameterized feed-forward neural networks with 5 layers of width 128, followed by final sofiplus activation function to ensure non-negativity.
G.3 Training details
We report default values for the different parameters of the GENOT Alg I If not stated otherwise in the corresponding experiments section, we use these parameters. As A parameters related to U- U-GENOT provided in teal.
Batch size: = 1024.
Entropic regularization strength: E 10- By default, we do not scale the cost matrices
passed to discrete OT solvers.
Unbalancedness parameter: T = (1, 1). This means that by default, we impose the hard
marginal constraints.
Number of training iterations: The 10, 000.
Optimizer: AdamW with learning rate Ir = 10-4, and weight decay X = 10-10, This is used for both, the vector field vr,,, and the re-weighting functions 7Ja Es.
When using the graph distance, construct k-nearest neighbor graph with batch_size number of edges. Then we apply a negative exponential kernel to obtain connectivties from the distances in the knn-graph. For the approximation of the heat kernel, we use the default parameters provided by the implementation in OTT-JAX 14]
G.4 Code availability
The GENOT model along with the code reproduce the experiments can be found at https //github. combb/uDDM/CCCC /MCDD//geot, genot, while more modular implementation can be found OTT-JAX
14.
Additionally, we implement applications in moscot 339
H
Impact statement
This work introduces novel approaches in neural optimal transport and presents applications in single-cell biology. Leveraging single-cell genomics data holds promise for advancing personalized medicine but requires cautious handling due to its potential inclusion of sensitive information. Moreover, neural optimal transport techniques find applicability across diverse domains, the societal ramifications of this research extend beyond single-cell genomics. Our intention is to provide access to the GENOT code, enabling its utilization by both general machine learning practitioners and, in subsequent release, ensuring its accessibility to advance and accelerate research specifically tailored for single-cell biologists.
41

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Everything claim we do in the abstract and introduction we actually do in the following.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations in A
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
42

---

Justification:
We provide proofs in S which we refer to from the main paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We describe the model trainings
G andd and the dataset curation
and $ B
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
43

---

Answer: [Yes]
Justification:
We provide links to multiple implementations of GENOT |
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/
swt.iissssoooicyyoyoo for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We report training details and dataset curation details in $
D S
E and
G
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report results across multiple runs for all experiments in the main text. For experiments in the appendix we only report results on one run due to limitations in computational resources.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
44

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [No]
Justification: We only perform single-GPU training, thus we assume there is no limita- tion reproduce single experiments / there is no aavironmenta//socieaall effect due single experimental run. The computational burden lies in the large number of different experimens/contiiuuaaiooss conducted for this study.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
ctwteieeeeeitiiieeeeieeelllleel ines?
Answer: [Yes]
Justification: We are aware of the NeurIPS Code of Ethics, and to the best of our knowledge- adhere to it. Morcover, we include an impact statement in $ H
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. s if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We include an impact statement in H
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
45

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We don't release pretrained model, dataset.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [No]
Justification: We cite the papers publishing the original data, and provide links to datasets. We don't it mention the licenses explicitly, as they are provided by the original publications.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
46

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We don't release assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects other than the bone marrow dataset [51). The data is anonymized, and the original publication discusses questions related research with human individuals.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Only the bone marrow dataset [1] stems from human individuals, but the data is anonymized, and the original publication discusses these aspects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
47

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
48


Reviews:
[{"review_id": "BnEdiermUE", "reviewer": "Reviewer_GAfL", "summary": "The authors present Generative Entropic Neural Optimal Transport (GENOT), a flexible method for learning entropic couplings for linear and quadratic entropic optimal transport (EOT), unbalanced OT, and OT across incomparable spaces. At its core, their method uses conditional flow matching (CFM) to solve for the OT couplings. Through extensive empirical experiments, the authors show the use-case of their method for learning couplings of cell trajectories, predicting single-cell response to perturbations, and translating between single-cell data modalities.", "questions": "- For the experiment on single-cell response to perturbations, it appears there are no comparisons to baseline methods? Is the reason for this because existing baseline methods are deterministic? Is there intuition on why the stochastic approach is helpful (or necessary) for this task?", "limitations": "The authors discuss limitations and broader impacts in the appendix.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors present promising and comprehensive set of empirical results that support the applicability of their method for OT related problems in single-cell biology. They demonstrate a novel use case of learning OT couplings between different data modalities in singel-cell biology, i.e. from RNA-seq to ATAC-seq. Their proposed method is also supported with sound theoretical background and justification. In general, the paper is well motivated and addresses important and challenging problems in single-cell biology.", "weaknesses": "Overall, this work has all the elements that constructs a complete and thorough application-based contribution to the field of neural optimal transport and single-cell biology. However, there remain several items that I feel hinder the understanding and impact of this work:\n\n- It seems to me that a key novelty of this work is the use of quadratic OT for learning couplings between incomparable spaces -- i.e. from ATAC-seq to RNA-seq for the single-cell case considered in this work. The authors consider a large quantity of experiments for both linear and quadratic OT over various settings, with many of the experiment being reported in the appendix. This makes some of the results and discussion hard to follow, and overall diluting the presentation of the key items of this paper. Given the large quantity of experiments and dense content of this paper, maybe it would be beneficial to shift more focus on these aforementioned experiments and move other to the appendix. This may help the reader better understand some of the key contributions of this work.\n- A large amount of items that are placed in the appendix tare referenced in the main text and seem to be important. This in general disrupts the overall fluidity of the paper. For instance, propositions in section 3 are listed in the appendix. Moreover, some experiments discussed in the main text only reference figures that are in the appendix. For example, experiments in section 5.1 \"GENOT-L on simulated data\" and \"U-GENOT-L predicts single-cell responses to perturbations\". In general, there seem to be a large proportion of experimental results that are placed in the appendix and referenced in section 5, which seem to be given equal wait to the experiment included in the main text. This makes it difficult to decipher the significance/importance of the respective results. As mentioned in my previous comment, possibly focussing the paper slightly may be beneficial to the overall presentation of this work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> ***The authors present promising and comprehensive set of empirical results [...] and addresses important and challenging problems in single-cell biology.***\n\n➤ We thank the reviewer for the positive feedback.\n\n> ***It seems to me that a key novelty of this work is the use of quadratic OT for learning couplings between incomparable spaces [...] This may help the reader better understand some of the key contributions of this work..***\n\n➤ We are grateful to the reviewer for providing this constructive criticism. We acknowledge that the paper includes a large number of experiments, which could hinder the reading flow. We decided to focus on the motivation of this work, which is the need for widely applicable and flexible neural OT estimators in single-cell genomics. In particular, we tried to motivate each of the 4 necessities **N1**, **N2**, **N3**, and **N4** (see lines 67-70 of the submission) as much incrementally as possible. Yet, we will try to include more experiments of the fused Gromov-Wasserstein case in the final version of the manuscript, which is indeed the major novelty. \n\nAdditionally, it is important to note that GENOT satisfying **N2**—approximating EOT coupling for **any cost function**—is a significant novelty. As detailed in lines 156-161 and illustrated in Figure 5, previous simulation-free linear EOT solvers, such as [Shi+2023], [Liu+2023], [Pooladian+2023], and [Tong+2024;a,b], are limited to the squared Euclidean cost $c(\\mathbf{x},\\mathbf{y}) = |\\mathbf{x}-\\mathbf{y}|_2^2$. This flexibility is crucial. We have demonstrated that in both linear (Figure 2) and quadratic EOT settings (Figure 4), using data-driven cost functions that approximate the geodesic distance on the data manifold yields more meaningful results from a biological perspective.\n\n\n> ***A large amount of items that are placed in the appendix are referenced in the main text and seem to be important. [...] As mentioned in my previous comment, possibly focussing the paper slightly may be beneficial to the overall presentation of this work.***\n\n➤ We agree with the reviewer that the large number of references in the main text to the appendix might hinder readability. As mentioned above, we intended to highlight the need for flexible neural OT estimators in single-cell genomics, and thus prioritized the motivation from a single-cell point of view. We admit that the readability of the paper would benefit from having e.g. more experiments or propositions in the main body. Yet, as we hope this paper will motivate single-cell biologists to build upon these methods for analyzing their data, we prioritized the applied point of view. \n\n> ***For the experiment on single-cell response to perturbations, it appears there are no comparisons to baseline methods? Is the reason for this because existing baseline methods are deterministic?***\n\n➤ This experiment serves as a motivating example for the need of neural unbalanced EOT estimators. The reason for not having included baselines in the perturbation prediction task is the lack of unbalanced entropic OT estimators. In fact, [Yang+2019] proposes a neural OT estimator in both a stochastic and a deterministic version, but their implementation only allows for the deterministic version. We decided to still include it in the benchmark of learning an unbalanced EOT plan between Gaussians (Figure 11) to demonstrate that learning an unbalanced EOT coupling is not trivial and cannot be replaced by estimators learning an unbalanced deterministic Monge Map.\n\n> ***Is there intuition on why the stochastic approach is helpful (or necessary) for this task?***\n\n➤ As motivated in the introduction (line 45: “[...] cells evolve stochastically […]”), cells evolve stochastically rather than deterministically, and hence obtaining stochastic predictions is relevant to model trajectories of cells. For example, stochastic evolution is relevant in developmental single-cell data (measured across different time points) where a homogeneous progenitor population (in the most extreme case, a single fertilized egg cell) stochastically proliferates into more mature cells. This is the motivation for the use case we consider in Figure 2. \n\nAnother source of stochasticity can be introduced from technical/experimental errors and biases, as we state in line 325: “imbalances might occur due to biases in the experimental setup or due to cell death”. The topic of uncertainty estimation is prevalent in single-cell genomics. For example, [Laehnemann+2020], state in “Eleven grand challenges in single-cell data science”: “Optimally, sc-seq analysis tools would accurately quantify all uncertainties arising from experimental errors and biases.” In particular, the perturbation data we consider e.g. in Figure 3, LHS, measures cells before and after perturbation. While in such experimental setups, it is common to perturb the same number of cells with each drug, the final data rarely has the same number of cells for each condition due to experimental errors. This is what motivates the calibration score in Figure 13, where we showed that cells with a high variance are mapped incorrectly (see Appendix C.2 for the description of the metric), resulting in a good calibration score.\n\n [Laehnemann+2020] Laehnemann et al., Eleven grand challenges in single-cell data science, 2020", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JF3mLQS3tC", "reviewer": "Reviewer_NoBU", "summary": "The paper proposes a framework for realigning cells in single-cell genomics which lies in the field of neural Optimal Transport (OT) solvers. The method utilizes a generative flow-based model for computing entropic OT couplings and tackles several practical challenges, e.g., it allows for using arbitrary cost functions, learning stochastic transport plans, relaxing mass conservation constraint, and tackling challenges of the (Fused) Gromov-Wasserstein problem. Specifically, the method consists of training of the conditional flow matching model on top of estimated discrete entropic OT plans. The approach is tested in different tasks from the single-cell biology field.", "questions": "- In the related work section you cited several works on unbalanced EOT [5,6,7]. Why did you include only one of them [7]  in comparison in Gaussians experiment? \n\n- In Table 1-3 you use abbreviations DBSM and CFM-SM which are not introduced in the text. Please introduce them in the relevant section.\n\n- Could you explain in more details why the issues of your implementation in the GW case will not arise in biological tasks? I see your explanation in lines 1345-1346, however, it seems that the issue might arise not only in the case of highly symmetric distributions", "limitations": "The authors adequately addressed the limitations and potential negative societal impact of their work in a specific section of their work.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper proposes a flexible framework which allows for addressing different challenges of application of neural OT solvers in the single cell genomics field. It is shown that the approach provides good results in several tasks from the single cell genomics according to the metrics under consideration. The paper is well-written and structured.", "weaknesses": "My major concern is that the provided approach is not guaranteed to learn the ground-truth plans. Most of the experimental details and theoretical results related to this topic are moved to Appendix which hampers the understanding of this question.\n\nIndeed, GENOT is based on the idea of distilling discrete entropic OT (EOT) solutions using the flow matching (which is close, but as the authors explain - different, to the idea from the work [1]). It raises a question: how well does such a kind of distillation of discrete plans approximate the ground-truth continuous ones (assuming that the approach works with batches of finite size)? Some explanations regarding the mini-batch biases are provided in lines 267-280. However, as far as I see it, the provided discussion (and references) do not cover the statistical properties of the discrete *unbalanced* linear or *quadratic plans*. The situation with quadratic setup is the most tricky here, since it is known that GW problem admits multiple solutions and, thus, the optimization problem may differ at each step of training depending on the calculated discrete entropic GW plans.\n\nFrom the empirical side, the authors address the ‘biasedness’ issue by testing their approach using a benchmark with analytically known EOT [2], U-EOT [3], and GW/Unbalanced GW [4] plans (all in Appendix E). According to these experiments, I tend to agree that the proposed approach provides meaningful results for the balanced (or unbalanced) linear case. However, the situation with the GW case is not clarified – according to the experiment description (Appendix E.3), the authors introduce some heuristics to “preserve the orientation of GW solution” and claim that in the case of biological tasks it is not needed. (Still, as a non-expert in biology, I can not assess the validity of this claim.) \n\n**In summary**, there are some evident issues with the implementation of the proposed approach in the Gromov-Wasserstein case. I appreciate that the authors mentioned this issue in the limitations section of their work (Appendix A). However, it is not obvious for me to what extent this is a serious drawback hampering the applicability of the approach to biological tasks. Besides, I think that as an important aspect of the implemented approach, this issue should be mentioned in the main body of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> ***It is shown that the approach provides good results in several tasks from the single cell genomics according to the metrics under consideration. The paper is well-written and structured.***\n\n➤ We thank the reviewer for the positive feedback.\n\n> ***My major concern is that the provided approach is not guaranteed to learn the ground-truth plans.***\n\n➤ We agree with the reviewer that distilling mini-batch coupling introduces a bias, which means that we can, in theory, only learn the true coupling in the asymptotic setting of infinite batch size. We discussed this aspect in detail in lines 267-280. On the other hand, it's important to remember that we are not aiming to learn a deterministic Monge map or a Benamou-Brenier velocity field. Instead, we operate in the entropic regime with $\\varepsilon \\gg 0$. In this regime, the estimation rates are parametric in $\\varepsilon$, helping to avoid the curse of dimensionality and significantly reducing the bias introduced compared to estimating deterministic maps (see the discussion related to statistical estimation below.) \n\nMoreover, solving a continuous OT problem in practice is very challenging, and most neural OT techniques introduce either a bias or a significant computational challenge into the learning process. In this paper, we have developed a method that is easy to train, even though it introduces a bias. This approach proved effective, as GENOT outperforms 12 baselines on a recent benchmark [Gushchin+2023], where we estimate a ground truth linear EOT known in closed form (see Tables 1-2). \n\nFinally, we want to emphasize that, to the best our knowledge, there is only one other (unpublished) neural OT solver targeting quadratic OT, which comes with a number of limitations, especially for single-cell genomics. As stated in lines 62-64 of the manuscript: *\"To our knowledge, the only neural formulation for GW proposed so far by [57] learns deterministic and balanced maps for inner product costs, using a min-max-min optimization procedure, which severely limits its application in single-cell genomics.\"* In general, solving a quadratic OT problem is extremely challenging. The only case where we can satisfactorily approximate GW couplings is in the discrete setting, using [Peyré+2016]'s scheme. Therefore, with GENOT, we aimed to develop a method that builds on these discrete GW couplings.\n\nAll in all, we would also like to highlight that by no means we assume there is no better way to solve some of the tasks which does not rely on discrete OT. Instead, we would like to encourage the community to build upon our method, as we hope to clearly motivate the challenges which are faced in single-cell genomics while providing in some of these cases a first method which can be benchmarked against.\n\n> ***Most of the experimental details and theoretical results related to this topic are moved to the Appendix, which hampers the understanding of this question.***\n\n➤ We agree with the reviewer that it is not ideal that experimental details and theoretical results are moved to the appendix. As the development of GENOT arose from the need for flexible neural OT estimators in single-cell genomics, we decided to highlight the applied point of view. We admit that the readability of the paper would benefit from having e.g. more experiments or propositions in the main body. Yet, as we hope this paper will motivate single-cell biologists to build upon these methods for analyzing their data, we prioritised the application, but we still hope to include a few more experimental details and/or theoretical results in the final version of the manuscript.\n\n> ***How well does such a kind of distillation of discrete plans approximate the ground-truth continuous ones? [...] Some explanations regarding the mini-batch biases are provided in lines 267-280. However, as far as I see it, the provided discussion (and references) do not cover the statistical properties of the discrete unbalanced linear or quadratic plans.***\n\n➤ This is a very good question. Most papers dealing with the estimation of EOT objects provide estimation rates for (Gromov)-Wasserstein distances. As stated in lines 267-281:\n- In the linear EOT setting, both balanced [Genevay+2019] and unbalanced [Séjourné+2021] settings achieve a parametric rate in $\\varepsilon$ dodging the curse of dimensionality.\n- In the quadratic EOT setting, for the balanced case, [Zhang+2023] demonstrate that similar parametric rate in $\\varepsilon$ can also be achieved. They prove this result by linearization of the quadratic EOT problem. Specifically, this rate depends on the minimum dimension of the source and target domains, indicating that estimation is easier when one of the domains is low-dimensional. \n\nThere is only one very recent paper [89] that provides a rate estimate for the ground truth coupling $\\pi^\\star_\\varepsilon$ by the empirical coupling $\\hat{\\pi}^n_\\varepsilon$ in the balanced linear EOT setting. This paper shows the same type of parametric rate in $\\varepsilon$, mitigating the curse of dimensionality.\n\nExtending this result to (i) the unbalanced EOT setting and (ii) the quadratic EOT setting is a very exciting perspective. For (i), since we have shown that an unbalanced EOT coupling is essentially a balanced EOT coupling between re-scaled measures (see Prop B.1), this result seems intuitively extendable. For (ii), the extension could be achieved by linearizing the quadratic EOT problem, similar to the approach used by [Zhang+2023] for costs.\n\nFor the remaining comments, we would like to kindly ask the reviewer to read the \"Additional Answer to Reviewer NoBU 1/2\" and \"Additional Answer to Reviewer NoBU 2/2\" added.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2tqCSSyBzp", "reviewer": "Reviewer_GBWj", "summary": "The paper proposes a multi-variable OT-based framework for solving Single-Cell related problems. The main idea behind the approach is to use the solutions for different variants of discrete entropy-regularize OT problems to learn a continuous parametric plan (measure) given by its conditional distributions, which approximates the true solution to the corresponding continuous problem. The authors propose utilizing conditional Flow Matching to parameterize the learned plan. The methodology is validated on various synthetic and real (single-cell) use cases. The obtained results testify the competitiveness of the framework.", "questions": "* The main of my questions is about the utilization of (conditional) flow matching for learning the parametric plan out of discrete minibatches. Why flow matching? As I understand, the proposed framework could be used on par with other generative techniques like GANs/normalizing flows/diffusion models. Why do you restrict yourself by flow matching?\n* I found it interesting that GENOT achieves (on average) the best results on EOT benchmark [3] among the competitors. Could the authors provide the code for the benchmark experiment?\n* Unbalanced Entropic OT work (to be mentioned): [1]. Also, a recent preprint you may find to be interesting: [2].\n* Line 89: Do I understand correctly that $\\mu$ and $\\nu$ are not necessarily probability distributions, but could be positive measures with the same mass?\n* Why does in (UQEOT) one utilize $\\varepsilon \\text{KL}^{\\otimes}$, while in (QEOT) - just  $\\varepsilon \\text{KL}$ (without $\\otimes$)?\n* Line 289: What is the “conditional mean” regime of the GENOT model?\n* Why does the unfused methods work badly on the single cell modalities translation problem with $\\ell_2^2$ intra-costs (Figure 4, FOSCTTM metric is almost everywhere near 0.5)?\n\nMisprints:\n- (QEOT) formula: $Q_{c_X, c_Y}$ -> $D_{c_X, c_Y}$\n- line 101: complexity if -> complexity is\n- line 351: Fig. 3 -> Fig. 4\n\n[1] Pariset et. al., Unbalanced Diffusion Schrödinger Bridge, ICML’23 workshop. \n\n[2] Gazdieva et. al., Light Unbalanced Optimal Transport.\n\n[3] Gushchin et. al., Building the Bridge of Schrödinger: A Continuous Entropic Optimal Transport Benchmark\n\n------------------------\n\nPOST REBUTTAL\n\nI thank the authors for the answers provided. I am satisfied with them. I have decided to raise my grade.", "limitations": "Ok", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Despite the fact that the paper is full of specific technical details on single-cell stuff, the manuscript is clear and pleasant for reading.\n\nThe method is clearly introduced. It seems that the authors adequately and fairly exhibit both the advantages of the method - its universality (coverage of different variants of OT formulations), simplicity (both for understanding and I guess for implementation); and limitations of the methodology - good limitation section + discussed reliance on discrete OT. \n\nAlso, I would like to commend a good experimental section - a lot of interesting use cases are considered and interesting results (sometimes even negative) are presented. These experiment settings are worth considering regardless of the proposed method.", "weaknesses": "* It seems the method could not (or hardly could)  be adapted to high-dimensional (e.g., image-data) setups due to reliance on discrete EOT. At least, I expect that the learned plan in such setups will significantly deviate from the true EOT plan. Still, I understand that the proposed framework is targeted on single cell data, which is not that high-dimensional.\n* Reliance on Discrete OT is something I personally treat as a weakness. It introduces biases which hardly could be analyzed. Still, in practice, the minibatch OT seems to work well.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> ***The method is clearly introduced [...] These experiment settings are worth considering regardless of the proposed method.***\n\n➤ We thank the reviewer for appreciating our contributions.\n\n> ***It seems the method could not (or hardly could) [...] which is not that high-dimensional.***\n\n➤ We appreciate that the reviewer understands that our method is motivated by and targeted at single-cell data. As mentioned in the general response, we found this question very interesting and applied GENOT to image data. Again, we would like to highlight that we do not intend to include the results in the paper and leave further investigation and optimization of GENOT applied to image data to future work, as i) we believe it does not fit the story line, ii) reviewers appreciated the large number of experiments, but also reviewer GAFl also critized that such a large number of experiments is in the appendix, hence adding computer vision examples would hamper readability even more. In addition to what we wrote in the general response, we would like to add the following: We follow the same hyperparameter and architectural setup as the authors in [Eyring+24]. We only added the FiLM condition module and modified the number of iterations to 350k. Moreover, we consider the balanced scenario, because - as mentioned - we did not perform any hyperparameter optimization. \n\n> ***Reliance on Discrete OT is something I personally treat as a weakness. It introduces biases which hardly could be analyzed. Still, in practice, the minibatch OT seems to work well.***\n\n➤ We respect that the reviewer considers the reliance on discrete OT as a weakness. We would also like to highlight that by no means we assume there is no better way to solve some of the tasks which does not rely on discrete OT. Instead, we would like to encourage the community to build upon our method, as we hope to clearly motivate the challenges which are faced in single-cell genomics while providing in some of these cases a first method which can be benchmarked against. There are clear limitations of our method, which we explicitly discuss in the appendix A or can be directly seen from the results. Yet, as stated by the reviewer, we also see that in some cases the performance of GENOT is good. We thank the reviewer for appreciating that the results of the experiments justify the usage of discrete OT.\n\n> ***The main of my questions is about the utilization of (conditional) flow matching for [...]. Why do you restrict yourself by flow matching?***\n\n➤ We do agree that our proposed method could be built with any generative model. The choice of flow matching is mainly practical, but it also has some methodological advantages. First and most importantly, training is fast (simulation-free as opposed to normalizing flows) and stable due to its simple minimization objective (as opposed to GANs). Moreover, we experienced diffusion models to be harder to train on single-cell data (worse performance in terms of distributional matching). On the methodological side, flow matching allows for estimating densities (as opposed to GANs), which can be useful in single-cell genomics as demonstrated in Figure 21.\n\n> ***I found it interesting that GENOT achieves (on average) the best results on EOT benchmark [3] among the competitors. Could the authors provide the code for the benchmark experiment?***\n\n➤ We provide the code in an **anonymous** Google Drive: [https://drive.google.com/drive/folders/17DUZr_bnTY4gv4nd8szk1h0wJ5NdQjlH?usp=sharing](https://drive.google.com/drive/folders/17DUZr_bnTY4gv4nd8szk1h0wJ5NdQjlH?usp=sharing). The folder `genot_rebuttal` includes a self-contained notebook `benchmark_example_64_1.ipynb` to reproduce an example of the results for the benchmark pair in dimension $d=64$ and $\\varepsilon=1$. We also included a file `requirements.txt` containing the package requirements.\n\n\n> ***Unbalanced Entropic OT work (to be mentioned): [1]. Also, a recent preprint you may find to be interesting: [2].***\n\n➤ We apologize for having missed this work, and now explicitly mention the work by replacing the sentence in line 56-59: \"[Eyring+2024], [20], [Lübeck+2023], [Yang+2019] proposed a way to incorporate unbalancedness into deterministic linear OT maps, while unbalanced formulations for entropic OT in both the linear and the quadratic case have not been explored.” by [Eyring+2024], [20], [Lübeck+2023], [Yang+2019] proposed a way to incorporate unbalancedness into deterministic linear OT maps, while [Yang+2019] extended their method to the entropic setting (but do not provide code), while recent preprints [Pariset+2023], [Gazdieva+2024] also suggested methods to approximate unbalanced linear EOT solutions. To the best of our knowledge, unbalanced neural formulations in quadratic OT have not yet been explored.\"\n\n> ***Line 89: Do I understand correctly that \\mu and \\nu are not necessarily probability distributions, but could be positive measures with the same mass?***\n\n➤ In the balanced OT setting, $\\mu$ and $\\nu$ are taken as probability measures, i.e., their mass is 1. This extends to all measures with arbitrary but equal masses, as they can always be normalized, *with the same renormalization constant*, without losing generality. The unbalanced setting is used for positive $\\mu$ and $\\nu$ measures with different masses.\n\n> ***Why does in (UQEOT) one utilize,  while in (QEOT) - just  (without )?***\n\n➤ This is a very good question. For QEOT, we do not need to use $\\mathrm{KL}^\\otimes$ because, when $\\mu$ and $\\nu$ have the same total mass, as shown in [Séjourné+2024, Prop 8.], $\\mathrm{KL}^\\otimes(\\pi|\\mu\\otimes\\nu) = 2 \\mathrm{KL}(\\pi|\\mu\\otimes\\nu)$. Therefore, using $\\mathrm{KL}$ is equivalent to using $\\mathrm{KL}^\\otimes$, with the adjustment of replacing $\\varepsilon$ by $\\varepsilon/2$.\n\nFor the remaining comments, we would like to kindly ask the reviewer to read the \"Additional Answer to Reviewer GBWj\" added.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank all reviewers for their encouraging feedback, constructive criticism, and thoughtful comments, as well as for pointing out typos.

In response to questions about the uniqueness of quadratic OT solutions raised by reviewer NoBU, we included an analysis of the stability of discrete OT solutions (see pdf document). We perform experiments on both simulated Gaussian data and single-cell data (from Fig. 4). In both cases the data has 20 dimensions (+10 dim. for fused exp.). The empirical analysis (Alg. 1, pdf) suggests:
1. Including a fused term in GW makes the solution more stable in both datasets (Fig. 2, pdf). Intuitively, the addition of this fused penalty on extra features $\mathbf{u},\mathbf{v}$ enables the "selection" of an optimal GW coupling. Essentially, we "trade off a bit of GW optimality" by choosing a coupling that minimizes the distortion of structural information $|c_\mathcal{X}(\mathbf{x}, \mathbf{x}') - c_\mathcal{Y}(\mathbf{y}, \mathbf{y}')|^2$ **while also minimizing** the cost on features $c(\mathbf{u}, \mathbf{v})$. Empirically, this mitigates the problem of the non-uniqueness of (pure) GW coupling to a large extent, making our procedure more stable. 
2. The geodesic cost makes the solution more stable for suffiently small $\varepsilon$ in single-cell experiments, but barely in Gaussian experiments (Fig. 2, pdf) 
3. The initialization scheme (App. A/E) for quadratic OT solvers makes the solution more stable (Fig. 2, pdf) 
 
We intend to include these results in the final version of the manuscript.
For experimental details and a more detailed discussion, we refer the reader to the responses to reviewer NoBU. 

While we thank reviewer GBWj (as well as reviewers GAfL and NoBU) for appreciating that GENOT is built for single-cell data, reviewer GBWj expressed skepticism about the applicability of GENOT to high-dimensional data like image data. We agree with reviewer GBWj that GENOT is built for single-cell data specifically, but we found the question interesting and thus conducted experiments on image data. Importantly, we would like to highlight that we do **not intend to include computer vision experiments in the final manuscript**, as in accordance with all reviewers, GENOT is built for single-cell data. We set out to apply GENOT-L to the common task of image translation on the CelebA dataset, translating females to males. We leveraged the flexibility of using any cost function and used CLIP [Radford+21] embeddings for both the discrete matching and the conditioning (together with FiLM [Perez+17]). 

While FID scores and examples of translated images can be found in the provided pdf, we made the following observations:
1. In terms of FID, GENOT-L performs comparably, but slightly worse than other flow-matching based methods (Table 1, pdf)
2. Visually, the generated images look realistic, but the coupling information is not always well preserved (Fig. 3, 4, pdf), as predicted by reviewer GBWj.



[Liu+14] Liu et al. “Deep Learning Face Attributes in the Wild.”, 2014

[Perez+17] Perez et al. “FiLM: Visual Reasoning with a General Conditioning Layer.”, 2017

[Eyring+24] Eyring et al., "Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation", 2024

[Radford+21] Radford et al., “Learning Transferable Visual Models From Natural Language Supervision” 2021

Meta-review: This paper applies optimal transport/conditional flow matching to learn cell development dynamics from single cell transcriptomics. 

All the reviewers liked the paper and are in favour of acceptance. The reviewers appreciated the clarity of the paper which is an accomplishment given that it covers a lot of pretty complicated ground.
Decision: Accept (poster)