Neural Unbalanced Optimal Transport via Cycle-Consistent Semi-Couplings
Abstract
Comparing unpaired samples of a distribution or population taken at different pointsin time is a fundamental task in many application domains where measuring populations is destructive and cannot be done repeatedly on the same sample, such asin single-cell biology.
Optimal transport (OT) can solve this challenge by learningan optimal coupling of samples across distributions from unpaired data.
However,the usual formulation of OT assumes conservation of mass, which is violated inunbalanced scenarios in which the population size changes (e.g., cell proliferationor death) between measurements.
In this work, we introduce N UB OT, a neural unbalanced OT formulation that relies on the formalism of semi-couplings to accountfor creation and destruction of mass.
To estimate such semi-couplings and generalize out-of-sample, we derive an efficient parameterization based on neural optimaltransport maps and propose a novel algorithmic scheme through a cycle-consistenttraining procedure.
We apply our method to the challenging task of forecastingheterogeneous responses of multiple cancer cell lines to various drugs, where weobserve that by accurately modeling cell proliferation and death, our method yieldsnotable improvements over previous neural optimal transport methods.
Introduction
Modeling change is at the core of various problems in the natural sciences, from dynamical processesdriven by natural forces to population trends induced by interventions.
In all these cases, the goldstandard is to track particles or individuals across time, which allows for immediate estimation ofindividual (or aggregate) effects.
But maintaining these pairwise correspondences across interventionsor time is not always possible, for example, when the same sample cannot be measured more thanonce.
This is typical in biomedical sciences, where the process of measuring is often altering ordestructive.
For example, single-cell biology profiling methods destroy the cells and thus cannot beused repeatedly on the same cell.
In these situations, one must rely on comparing different replicasof a population and, absent a natural identification of elements across the populations, infer thesecorrespondences from data in order to model evolution or intervention effects.
The problem of inferring correspondences across unpaired samples in biology has been traditionallytackled by relying on average and aggregate perturbation responses (Green & Pelkmans, 2016; Zhanet al., 2019; Sheldon et al., 2007) or by applying mechanistic or linear models (Yuan et al., 2021;
Dixit et al., 2016) in, potentially, a learned latent space (Lotfollahi et al., 2019).
Cellular responses totreatments are, however, highly complex and heterogeneous.
To effectively predict the drug responseof a patient during treatment and capture such cellular heterogeneity, it is necessary to learn nonlinearmaps describing such perturbation responses on the level of single cells.
Assuming perturbationsincrementally alter molecular profiles of cells, such as gene expression or signaling activities, recentapproaches have utilized optimal transport to predict changes and alignments (Schiebinger et al.
,2019; et al., 2022a; Tong et al., 2020).
By returning a coupling between control and perturbed
[Table] [Table]
base64="VCM7gqRpqVIucX0+cOnR6uebBJc=">AAAB73icdVDLSgMxFM3UV62vqks3wSK4GjJ9d1d047KCtYV2KJk004YmmTHJCGXoT7hxoYhbf8edf2P6EFT0wIXDOfdy7z1BzJk2CH04mbX1jc2t7HZuZ3dv/yB/eHSro0QR2iYRj1Q3wJpyJmnbMMNpN1YUi4DTTjC5nPude6o0i+SNmcbUF3gkWcgINlbq9kdYCDwoDvIF5KJGtVaqQuR6jTKqIEtK9SKqlKDnogUKYIXWIP/eH0YkEVQawrHWPQ/Fxk+xMoxwOsv1E01jTCZ4RHuWSiyo9tPFvTN4ZpUhDCNlSxq4UL9PpFhoPRWB7RTYjPVvby7+5fUSE9b9lMk4MVSS5aIw4dBEcP48HDJFieFTSzBRzN4KyRgrTIyNKGdD+PoU/k9ui65XdcvX5ULzYhVHFpyAU3AOPFADTXAFWqANCODgATyBZ+fOeXRenNdla8ZZzRyDH3DePgEb25AK</latexit
[Figure]
cell states, which overall minimizes the cost of matching, optimal transport can solve that puzzleand reconstruct these incremental changes in cell states over
Despite the advantages mentioned above, the classic formulation of OT is ill-suited to model processeswhere the population changes in size , e.g., where elements might be created or destroyed over time.
This is the case, for example, in single-cell biology, where interventions of interest typically promoteproliferation of certain cells and death of others.
Such scenarios violate the assumption of conservationof mass that the classic OT problem relies upon.
Relaxing this assumption yields a generalizedformulation, known as the unbalanced OT (UBOT) problem, for which recent work has studied itsproperties (Liero et al., 2018; Chizat et al., 2018a), numerical solution (Chapel et al., 2021), and hasapplied it successfully to problems in single cell biology (Yang & Uhler, 2019).
Yet, these methodstypically scale poorly with sample size, are prone to unstable solutions, or make limiting assumptions,e.g., only allowing for destruction but not creation of mass.
In this work, we address these shortcomings by introducing a novel formulation of the unbalanced
OT problem that relies on the formalism of semi-couplings introduced by Chizat et al.
(2018b),while still obtaining an explicit transport map that models the transformation between distributions.
The advantage of the latter is that it allows mapping new out-of-sample points, and it provides aninterpretable characterization of the underlying change in distribution.
Since the unbalanced OTproblem does not directly admit a Monge (i.e., mapping-based) formulation, we propose to learnto jointly ‘re-balance’ the two distributions, thereby allowing us to estimate a map between theirrescaled versions.
To do so, we leverage prior work (Makkuva et al., 2020; Korotin et al., 2021) thatlearns the transport map as the gradient of a convex dual potential (Brenier, 1987) parameterized viaan input convex neural network (Amos et al., 2017).
In addition, we derive a simple update rule tolearn the rescaling functions.
Put together, these components yield a reversible, parameterized, andcomputationally feasible implementation of the semi-coupling unbalanced OT formulation (Fig. 1).
In short, the main contributions of this work are: (i) A novel formulation of the unbalanced optimaltransport problem that weaves together the theoretical foundations of semi-couplings with the practicaladvantage of transport maps; (ii) A general, scalable, and efficient algorithmic implementation forthis formulation based on dual potentials parameterized via convex neural network architectures; and(iii)
An empirical validation on the challenging task of predicting perturbation responses of singlecells to multiple cancer drugs, where our method successfully predicts cell proliferation and death, inaddition to faithfully modeling the perturbation responses on the level of single cells.
Background2.1 Optimal Transport
For two probability measures µ, ν in P ( X ) with X = R d and a real-valued continuous cost functionc ∈ C ( X 2 ) , the optimal transport problem (Kantorovich, 1942) is defined as
[Equation] [Equation]
where Γ( µ, ν ) = { γ ∈ M + ( X 2 ) , s.t.
(Proj 1 )
♯ γ = µ, (Proj 2 )
♯ γ = ν } is the set of couplings in thecone of nonnegative Radon measures M + ( X 2 ) with respective marginals µ, ν .
When instantiatedon finite discrete measures, such as µ = (cid:80) ni =1 u i δ x i and ν = (cid:80) mj
=1 v j δ y j , with u ∈ Σ n , v ∈ Σ mthis problem translates to a linear program, which can be regularized using an entropy term (Peyré
Cuturi, 2019).
For ε ≥ 0 , set
[Equation]
where H ( P ) :
= − (cid:80) ij P ij (log P ij − 1) and the polytope U ( u , v ) is the set of matrices{ P ∈ R n × m + , P1 m = u , P ⊤ 1 n = v } .
For clarity, we will sometimes write OT ε ( u , v , { x i } , { y j } ) .
79 Notice that the definition above reduces to (1) when ε = 0 .
Setting ε > 0 yields a faster anddifferentiable proxy to approximate OT and allows fast numerical approximation via the Sinkhornalgorithm (Cuturi, 2013), but introduces a bias, since in general OT ε ( µ, µ ) ̸ = 0 .
Neural optimal transport.
To parameterize (1) and allow to predict how a measure evolves fromµ to ν , we introduce an alternative formulation known as the Monge problem (1781) given by
[Equation]
with pushforward operator ♯ and the optimal solution T ∗ known as the Monge map between µ and ν .
When cost c is the quadratic Euclidean distance, i.e., c = ∥ · ∥ 22 , Brenier’s theorem (1987) states thatthis Monge map is necessarily the gradient ∇ ψ of a convex potential ψ : X (cid:55)→ R such that ∇ ψ ♯ µ = ν , 87 i.e., T ∗ ( x ) =
∇ ψ ( x ) .
This connection has far-reaching impact and is a central component of recentneural optimal transport solvers (Makkuva et al., 2020; Bunne et al., 2022c; Alvarez-Melis et al.
,2022; Korotin et al., 2020; Bunne et al., 2022b; Fan et al., 2021b).
Instead of (indirectly) learning the
Monge map T (Yang & Uhler, 2019; Fan et al., 2021a), it is sufficient to restrict the computationaleffort to learning a good convex potential ∇ θ , parameterized via input convex neural networks(ICNN)
(Amos et al., 2017), s.t. ∇ θ ψ♯µ = ν .
Alternatively, parameterizations of such maps can becarried out via the dual formulation of (1) (Santambrogio, 2015, Proposition 1.11, Theorem 1.39), i.e.,
[Equation]
where the dual potentials f, g are continuous functions from X to R , and f ⊕ g (cid:55)→ f
( x ) + g ( x ) .
Based on Brenier (1987), Makkuva et al.
(2020) derive an approximate min-max optimizationscheme parameterizing the duals f, g via two convex functions.
The objective thereby reads
[Equation]
When paramterizing f and g via a pair of ICNNs with parameters θ f and θ g , this neural OT schemethen allows to predict ν or µ via ∇ g θ g ♯ µ or ∇ f θ f ♯ ν , respectively.
We further discuss neural primal 99 (Fan et al., 2021a; Yang & Uhler, 2019) and dual approaches (Makkuva et al., 2020; Korotin et al.
,2020; Bunne et al., 2021) in §D.2.
2.2 Unbalanced Optimal Transport
A major constraint of problem (1) is its restriction to a pair of probability distributions µ and ν ofequal mass.
Unbalanced optimal transport (Benamou, 2003; Liero et al., 2018; Chizat et al.
, 2018b)lifts this requirement and allows a comparison between unnormalized measures, i.e., via
[Equation]
with f -divergences D f 1 and D f 2 induced by f 1 , f 2 , and parameters ( τ 1 , τ 2 ) controlling how muchmass variations are penalized as opposed to transportation of the mass.
When introducing an entropyregularization as in (2), the unbalanced OT problem between discrete measures u and v , i.e.,
[Equation]
( u , v
[Equation]
can be efficiently solved via generalizations of the Sinkhorn algorithm (Chizat et al., 2018a; Cuturi,2013; Benamou et al., 2015) .
We describe alternative formulations of the unbalanced OT problem indetail, review recent applications, and provide a broader literature review in the Appendix (§A.1).
A Neural Unbalanced Optimal Transport Model
The method we propose weaves together a rigorous formulation of the unbalanced optimal transportproblem based on semi-couplings (introduced below) with a practical and scalable OT mappingestimation method based on input convex neural network parameterization of the dual OT problem.
Semi-coupling formulation.
Chizat et al.
(2018b) introduced a class of distances that generalize optimal transport for the unbalanced setting.
They introduce equivalent dynamic and static formulationsof the problem, the latter of which relies on semi-couplings to allow for variations of mass.
Theseare generalizations of couplings whereby only one of the projections coincides with a prescribedmeasure.
Formally, the set of semi-couplings between measures µ and ν is defined as
[Equation]
With this, the unbalanced Kantorovich OT problem can be written as C k ( µ, ν )
=inf
[Equation]
( x, y ) , where γ is any joint measure for which γ 0 , γ 1 ≪ γ .
Although this formulation lends itself to formal theoretical treatment, it has at least two limitations.
First, it does not explicitly model a mapping between measures.
Indeed, no analogue of the celebrated
Brenier’s theorem is known for this setting.
Second, deriving a computational implementation ofthis problem is challenging by the very nature of the semi-couplings: being undetermined along onemarginal makes it hard to model the space in (8).
Rebalancing with proxy measures.
To turn the semi-coupling formulation of unbalanced OT intoa computationally feasible method, we propose to conceptually break the problem into balancedand unbalanced subproblems, each tackling a different aspect of the difference between measures:feature transformation and mass rescaling.
These in turn imply a decomposition of the semi-couplingsof (8), as we will show later.
Specifically, we seek proxy measures ˜ µ and ˜ ν with equal mass (i.e.,µ ( X ) =
˜
ν
( X ) ) across which to solve a balanced OT problem through a Monge/Brenier formulation.
To decouple measure scaling from feature transformation, we propose to choose ˜ µ and ˜ ν simply asrescaled versions of µ and ν .
Thus, formally, we seek ˜ µ, ˜ ν ∈ M + ( X ) and T, S : X → X such that
[Equation]
9) where η, ζ : X → R + are scalar fields, η · µ denotes the measure with density η ( x ) d µ ( x ) (analogouslyfor ζ · ν ) , and T, S are a pair of forward/backward optimal transport maps between ˜ µ and ˜ ν
(Fig. 1b).
Devising an optimization scheme to find all relevant components in (9) is challenging.
In particular,it involves solving an OT problem whose marginals are not fixed, but that will change as thereweighting functionals η, ζ are updated.
We propose an alternating minimization approach, wherebywe alternative solve for η, ζ (through an approximate scaling update) and T, S (through gradientupdates on ICNN convex potentials, as described in Section 2.1).
Updating rescaling functions.
Given current estimates of η and T , we consider the UBOT problem(6) between T ♯ ( η · µ ) =
T ♯ ˜ µ and ν .
Although in general these two measures will not be balanced(hence why we need to use UBOT instead of usual OT), our goal is to eventually achieve this.
Toformalize this, let us use the shorthand notation π ∗ UB ( α, β ) def.
= argmin π UBOT ( π ; α, β ) , where UBOTis defined in (7).
For a fixed T , our goal is to find η such that (
[Equation]
For the discrete setting (finite samples), this corresponds to finding a vector e satisfying:
[Equation]
where Γ = argmin UBOT ( e ⊙ u , T ( x i ) , v , y j ) .
[Equation]
For a fixed T , the vector e ∗ satisfying this system can be found via a fixed-point iteration.
In practice,we approximate it instead with a single-step update using the solution to the unscaled problem:
[Equation]
which empirically provides a good approximation on the optimal e ∗ but is significantly more efficient.
Apart from requiring a single update, whenever u and v are uniform (as in most applications wherethe samples are assumed to be drawn i.i.d.) solving this problem between unscaled histograms will befaster and more stable than solving its scaled (and therefore likely non-uniform) counterpart in (10).
Input: f, g : ICNNs, initialized s.t. ∇ g ( x ) ≈ x and ∇ f ( y )
≈ y ; η, ζ : NNs 1 for t in epochs do
Sample batch { x i } ni =1 ∼ µ and { y j } mj =1 ∼ ν ˆ y ← ∇ g ( x )ˆ x ← ∇ f ( y ) Γ 1 ← unbalanced.sinkhorn (ˆ y, 1 n 1 n , y, 1 m 1 m ) e i ← (cid:80) j Γ ij (cid:80) ij Γ ij · n Γ 2 ← unbalanced.sinkhorn (ˆ x, 1 m 1 m , y, 1 n 1 n ) z i ← (cid:80)
j Γ ij (cid:80) ij Γ ij · m
J ( θ g , θ f ) = 1 n (cid:80) ni
=1 e i [ f ( ∇ g ( x i )) − ⟨ x i , ∇ g ( x i ) ⟩ ] − 1 m (cid:80) mj
=
z j f ( y j )
L η ( θ η ) =
MSE ( e , η ( x ))
L ζ ( θ ζ ) =
MSE ( z , ζ ( y ))
Update θ g to minimize J , θ η to minimize L η , θ ζ to minimize L ζ , and θ f to maximize J
Analogously, for a given S , we choose ζ that ensures (Proj 2 )
♯
[ π ∗ UB ( S ♯ ( ζ · ν ) , µ )] =
S ♯ ( ζ · ν ) .
Forempirical measures, this yields the update:
[Equation]
In order to be able to predict mass changes for new samples, we will use the discrete e , z tofit continuous versions of η, ζ via functions parameterized as neural networks trained to achieveη ( x i )
≈ e i ∀ i ∈ { 1 , . . . , n } and ζ (
[Equation]
Updating mappings.
For a fixed pair of η, ζ , we want T and S to be a pair of optimal OT mapsbetween ˜ µ and ˜ ν .
Since these are guaranteed to be balanced due to the argument above, we can use ausual (balanced) OT formulation to find them.
In particular, we use the formulation of (Makkuvaet al., 2020) to fit them.
That is, T = ∇ g and S = ∇ f for convex potentials f and g , parameterizedas ICNNs with parameters θ f
and θ g .
The corresponding objective for these two potentials is:
[Equation]
In the finite sample setting, this objective becomes:
[Equation]
The optimization procedure is summarized in Algorithm 1.
Transforming new samples.
After learning f, g, η, ζ , we can use these functions to transform(map and rescale) new samples, i.e., beyond those used for optimization.
For a given source datapointx with mass u , we transform it as ( x , u ) (cid:55)→ ( ∇ g ( x ) , η ( x ) · u · ζ ( ∇ g ( x )) − 1 ) .
Analogously, targetpoints can be mapped backed to the source domain using ( y , v ) (cid:55)→ ( ∇ f ( y ) , ζ ( y ) · v · η ( ∇ f ( y )) − 1 ) .
Recovering semi-couplings.
Let us define ˜Γ 1 def.
=
diag ( e − 1 ) ⊤ Γ 1 and ˜Γ 2 def.
=
diag ( z − 1 )
⊤ Γ 2 , whereΓ 1 , Γ 2 are the solutions of the UBOT problems computed in Algorithm 1 (lines 7 and 9, respectively).
It is easy to see that (˜Γ 1 , ˜Γ ⊤ 2 ) is a valid pair of semi-couplings between µ and ν (cf.
Eq.
Evaluation
We illustrate the effectiveness of N UB OT on different tasks, including a synthetic setup for which aground truth matching is available, as well as an important but challenging task to predict single-cellperturbation responses to a diverse set of cancer drugs with different modes of actions.
Baselines.
To put N UB OT’s performance into perspective, we compare it to several baselines:
First,we consider a balanced neural optimal transport method C ELL OT
(Bunne et al., 2021), based on the
N ub O
T ub O T
G AN a.
N ub O
T ub O T
G AN b.
N ub O
T ub O T
G AN c. neural dual formulation of Makkuva et al.
Further, we benchmark N UB OT against the currentstate-of-the-art UB OT GAN, an unbalanced OT formulation proposed by Yang & Uhler (2019), whichsimultaneously learns a transport map and a scaling factor for each source point in order to accountfor variation of mass.
Additionally, we consider two naive baselines: I DENTITY , simulating theidentity matching and modeling cell behavior in absence of a perturbation, and O BSERVED , a randompermutation of the observed target samples and thus a lower bound when comparing predictions toobserved cells on the distributional level.
More details can be found in the Appendix §D.2.
4.1 Synthetic Data
Populations are often heterogeneous and consist of different subpopulations.
Upon intervention,these subpopulations might exhibit heterogeneous responses.
To simulate such heterogeneousintervention responses, we generate a dataset containing a two-dimensional mixture of Gaussianswith three clusters in the source distribution µ .
The target distribution ν consists of the same threeclusters, but with different cluster proportions.
Further, each particle has undergone a constant shiftin space upon intervention.
We consider three scenarios with increasing imbalance between thethree clusters (see Fig.
2a-c).
We evaluate N UB OT on the task of predicting the distributional shiftfrom source to target, while at the same time correctly rescaling the clusters such that no mass istransported across non-corresponding clusters.
Results.
The results (setup, predicted Monge maps and weights) are displayed in Fig.
Both
N UB OT and UB OT GAN correctly map the points to the corresponding target clusters without transporting mass across clusters.
N UB OT also accurately models the change in cluster sizes by predictingthe correct weights for each point.
In contrast, UB OT GAN only captures the general trend of clustergrowth and shrinkage, but does not learn the exact weights required to re-weight the cluster proportions appropriately.
The exact setup and calculation of weights can be found in the §B (see Table 1).
4.2 Single-cell Perturbation Responses
Through the measurement of genomic, transcriptomic, proteomic or phenotypic profiles of cells, andthe identification of different cell types and cellular states based on such measurements, biologistshave revealed perturbation responses and response mechanism which would have remained obscuredin bulk analysis approaches (Green & Pelkmans, 2016; Liberali et al., 2014; Kramer et al., 2022).
However, single-cell measurements typically require the destruction of cells in the course ofrecording.
Thus, each measurement provides us only with a snapshot of cell populations, i.e.,samples of probability distribution that is evolving in the course of the perturbation, from controlµ (source) to perturbed cell states ν (target).
Using N UB OT and the considered baselines, welearn a map T that reconstructs how individual cells respond to a treatment.
The effect of a singleperturbation frequently varies depending on the cell type or cell state, and may include the inductionof cell death or proliferation.
In the following, we will evaluate if N UB OT is able to capture andpredict heterogeneous proliferation and cell death rates of two co-cultures melanoma cell liensthrough η and ζ in response to 25 drug treatments.
8h 24h
[Table]
<
[Table]
latexit
<latexit
[Table]
/ExkJtU5D30yGBMZ62ZuJ/3m9BIZ1L+MyToBJulg0TASGCM96wAFXjIJIDSFUcXMrpmOiCAXTVsmU4C5/eZW0z6vuRbV2W6s06nkdRXSMTtAZctElaqBr1EQtRNEDekav6M16sl6sd+tjMVqw8swR+gPr8wcjTpSN</latexit> Identity
[Figure]
t
[Figure]
MelA for each drug, we analyze their level of correlation to our predicted weights after a. 8h and b.
24h.
With increasing difficulty of the task and certain drugs completely removing both or one of the subpopulations, the level of correlation reduces from 8 to 24h.
The single-cell measurements used for this task were generatedusing the imaging technology 4i (Gut et al., 2018) over the courseof 24 hours, resulting in three different unaligned snapshots( t = 0 h , t = 8 h and t = 24 h ) for each of the drug treatments.
The control cells, i.e., the source distribution µ , consists ofcells taken from a mixture of melanoma cell lines at t =
hthat are exposed to a dimethyl sulfoxide (DMSO) as a vehiclecontrol.
Futher, We consider two different target populations νcapturing the perturbed populations after t = 8 h and t = 24 hof treatment, respectively.
As both cancer cell lines exhibitdifferent sensitivities to the drugs (Raaijmakers et al.
, 2015),their proportion (Fig. 11) as well as the total cell counts (Fig.
13)vary over the time points.
Both cell lines are characterizedby the expression of mutually exclusive protein markers, i.e.,one cell line strongly expresses a set of proteins detected byan antibody called MelA (MelA + cell type), while the other ischaracterized by high levels of the protein Sox9 (Sox9 + celltype).
An evaluation of this cell line annotation can be foundin Fig.
10 (8h) and Fig.
12 (24h).
Contrary to the synthetic taskin § 4.2, the nature of these measurements destroys the groundtruth matching.
We thus use insights from the number of cellsafter 8 and 24 hours of treatment (Fig. 11, 13), as well as thecell type annotation for each cell to further evaluate N UB OT’sperformance.
A description of the data can be found in § C.2.
Results.
We predict perturbed cell states for each drug out-ofsample on a hold-out test set.
Details on the network architectureand hyperparameters can be found in § D.3.
N UB OT and UB OT
GAN additionally predict the weight associated with the perturbed predicted cells, giving insights into which cells haveproliferated or died in response to the drug treatments.
First, wecompare how well each method fits the observed perturbed cellson the level of the entire distribution.
For this, we measure theweighted version of kernel maximum mean discrepancy (MMD)between predictions and observations.
More details on the evaluation metrics can be found in § D.1.
The results are displayedin Fig.
N UB OT outperforms all baselines in almost all drugperturbations, showing its effectiveness in predicting OT mapsand local variation in mass.
In the absence of a ground truth and in particular, given ourinability to measure (i.e., observe) cells which have died upontreatment, we are required to base further analysis of N UB OT’s predictions on changes in cell count for each subpopulation (MelA + , Sox9 + ).
Fig.
11 clearly shows that drug treatments lead tosubstantially different cell numbers for each of the subpopulations compared to control.
For example,
Ulixertinib leads to the proliferation of both subpopulations after 8h, but to pronounced cell deathin Sox9 + and strong proliferation in MelA + cells after 24h.
We thus expect, that weights predictedby N UB OT for all drugs correlate with the change in cell counts for each cell type (here measuredas population fractions).
This is indeed the case, Fig.
4 shows a high correlation between observedcell counts of the two cell types and the sum of the predicted weights of the respective cell typesafter 8h of treatment for all drugs.
After 24 hours, treatment-induced cell death (in at least one celltype) by some drugs can be so severe at that the number of observed perturbed cells becomes too lowfor accurate predictions and the evaluation of the task (Fig. 13).
Further, we find that drugs influencethe abundance of the cell lines markers MelA and Sox9, complicating cell type classification (see
Fig.
10, 12, as well as Fig.
We ignore drugs falling into these categories and find that whilstthe correlation between predicted weights and observed cell counts is reduced after 24h (see Fig. 4a),
N UB OT still captures the overall trend.
The data further provides insights into biological processes such as apoptosis, a form of programmedcell death induced by enzymes called Caspases (ClCasp3).
While dead cells become invisible in thecell state space (they cannot be measured), dying cells are still present in the observed perturbedsample and can be recognized by high levels of ClCasp3 (the apopotosis markers).
Conversely, theprotein Ki67 marks proliferating cells.
Analyzing the correlation of ClCasp
and Ki67 intensitywith the predicted weights provides an additional assessment of the biological meaningfulness ofour results.
For example, upon Ulixertinib treatment, the absolute cell counts show an increase in
Sox9 + cells, and a decline of MelA + cells at 24h (Fig. 11).
Fig.
5 shows UMAP projections of thecontrol cells at both time points, colored by the observed and predicted protein marker values and thepredicted weights.
At 8h, N UB OT predicts only little change in mass, but a few proliferative cellswith high weights in areas which are marked by high values of the proliferation marker Ki67.
At 24h,our model predicts cell death in the Sox9 + (MelA − ) cell type, and proliferation in the MelA + celltype, which matches the observed changes in cell counts per cell type, seen in Fig.
11 in § B. Weidentify similar results for Trametinib (Fig. 7), Ixazomib (Fig.
8), and Vindesine (Fig. 9) which canbe found in § B. These experiments thus demonstrate that N UB OT accurately predicts heterogeneousdrug responses at the single-cell level, capturing both, cell proliferation and death.
Conclusion
This work presents a novel formulation of the unbalanced optimal transport problem that bridges twopreviously disjoint perspectives on the topic: a theoretical one based on semi-couplings and a practicalone based on recent neural estimation of OT maps.
The resulting algorithm, N UB OT, is scalable, efficient, and robust.
Yet, it is effective at modeling processes that involve population growth or death, asdemonstrated through various experimental results on both synthetic and real data.
On the challengingsingle-cell perturbation task, N UB OT is able to successfully predict perturbed cell states, while explicitly modeling death and proliferation.
This is an unprecedented achievement in the field of single-cellbiology, which currently relies on the use of markers to approximate the survival state of cell population upon drug treatment.
Explicitly modelling proliferation and death at the single-cell level as partof the drug response, allows to link cellular properties observed prior to drug treatment to therapyoutcomes.
Thus, the application of N UB OT in the fields of drug discovery and personalized medicinecould be of great implications, as it allows to identify cellular properties predictive of drug efficacy.
A Background
In the following, we provide further information and review related literature on concepts discussed throughout 454 this work.
A.1 Unbalanced Optimal Transport
Unbalanced optimal transport is a generalization of the classical OT formulation (1), and as such allows mass to 457 be created and destroyed throughout the transport.
This relaxation has found recent use cases in various domainsranging from biology (Schiebinger et al., 2019; Yang & Uhler, 2019), imaging (Lee et al., 2019), shape regis- 459 tration (Bonneel & Coeurjolly, 2019), domain adaption (Fatras et al., 2021), positive-unlabeled learning (Chapelet al., 2020), to general machine learning (Janati et al., 2020; Frogner et al., 2015).
Problem (6) provides a general 461 framework of the unbalanced optimal transport problem, which can recover related notions introduced in the literature: Choosing for D f the Kullback-Leibler divergence, one recovers the so-called squared Hellinger distance.
Alternatively, with D f = ℓ 2 norm, we arrive at Benamou (2003), while an ℓ 1 norm retrieves a concept often 464 referred to as partial OT (Figalli, 2010).
The latter comprises approaches which do not rely on a relaxation of themarginal constraints as in (6).
In particular, some strategies of partial OT expand the original problem by adding 466 virtual mass to the marginals (Pele & Werman, 2009; Caffarelli & McCann, 2010; Gramfort et al., 2015), or by extending the OT map by dummy rows and columns (Sarlin et al., 2020) onto which excess mass can be transported.
A further review is provided in (Peyré & Cuturi, 2019, Chapter 10.2).
Recent work has furthermore developedalternative computational schemes (Chapel et al., 2021; Séjourné et al., 2022) as well as provided a computational 470 complexity analysis (Pham et al., 2020) of the generalized Sinkhorn algorithm solving entropic regularized unbalanced OT (Chizat et al., 2018a).
Besides Yang & Uhler (2019), these approaches do not provide parameterizations 472 of the unbalanced problem and allow for an out-of-sample generalization which we consider in this work.
A.2 Cycle-Consistent Learning
The principle of cycle-consistency has been widely used for learning bi-directional transformations between 475 two domains of interest.
Cycle-consistency thereby assumes that both the forward and backward mappingare roughly inverses of each other.
In particular, given unaligned points x ∈ X and y ∈ Y , as well as maps 477 g : X (cid:55)→ Y and f : Y (cid:55)→ X , cycle-consistency reconstruction losses enforce ∥ x − f
( g ( x ))
∥ as well as∥ y − g
( f ( y ))
∥ using some notion of distance ∥ · ∥ , assuming that there exists such a ground truth bijection 479 g = f − 1 and f = g − 1 .
The advantage of validating good matches by cycling between unpaired samplesbecomes evident through the numerous use cases to which cycle-consistency has been applied: Originally 481 introduced within the field of computer vision (Kalal et al., 2010) and applied to image-to-image translationtasks (Zhu et al., 2017a), it has been quickly adapted to multi-modal problems (Zhu et al., 2017b), domain 483 adaptation (Hoffman et al., 2018), and natural language processing (Shen et al., 2017).
The original principlehas been further generalized to settings requiring a many-to-one or surjective mapping between domains (Guo 485 et al., 2021) via conditional variational autoencoders, dynamic notions of cycle-consistency (Zhang et al., 2021),or to time-varying applications (Dwibedi et al., 2019).
These classical approaches enforce cycle-consistency 487 by explicitly composing both maps and penalizing for any deviation from this bijection.
In this work, we treatcycle-consistency differently.
It is enforced implicitly by coupling the two distributions of interest through a 489 sequence of reversible transformations: re-weighting, transforming, and re-weighting (Eq.
(9) and Fig. 1).
Similarly to our work, Zhang et al.
and Hur et al.
(2021) establish a notion of cycle-consistency 491 (reversibility) for a pair of pushforward operators to align two unpaired measures.
Both methods rely on the
Gromov-Monge distance (Mémoli & Needham, 2022), a divergence to compare probability distributions defined 493 on different ambient spaces X and Y —a setting not considered in this work.
They proceed by defining areversible metric through replacing the single Monge map by a pair of two Monge maps, i.e., f : X (cid:55)→ Y and 495 g : Y (cid:55)→ X , minimizing the objective
[Equation] [Equation]
Problem (12) shows similarities to the classical cycle-consistency objective of Zhu et al.
(2017a), where cycle- 498 consistency is indirectly enforced through ∆ p
X , Y .
Zhang et al.
(2022) parameterize both Monge maps through neural networks in a similar fashion as done in (Yang & Uhler, 2019; Fan et al., 2021a).
Our approach differs 500 from Zhang et al.
; Hur et al.
(2021) as we model the problem through a single Monge map with duals f, g ,allowing us to map back-and-forth between measures µ and ν , and using a different parametrization approach 502 (ICNNs).
More importantly, the approaches presented by Zhang et al.
; Hur et al.
(2021) do not generalizeto the unbalanced case.
While Zhang et al.
(2022) proposed an unbalanced version of (12) by relaxing the 504 marginals as done in Chizat et al.
(2018a)
, they require the unbalanced sample sizes to be known (i.e., n and mneed to be fixed).
In our application of interest, particle counts of the target population are, however, not known 506 a priori .
A.3 Convex Neural Architectures
Input convex neural networks (Amos et al., 2017) are a class of neural networks that approximate the family of 509 convex functions ψ with parameters θ , i.e., whose outputs ψ θ ( x ) are convex w.r.t.
an input x .
This property isrealized by placing certain constraints on the networks parameters θ .
More specifically, an ICNN is an L -layer 511 feed-forward neural network, where each layer l
= { 0 , ..., L − 1 } is given by
[Equation]
and
[Equation]
(13) where σ k are convex non-decreasing activation functions, and θ = { W xl , W zl , b l } L − 1 l =0 is the set of parameters, 513 with all entries in W zl being non-negative and the convention that z 0 and W z 0 are 0 .
As mentioned above andthrough the connection established in § 2, convex neural networks have been utilized to approximate Monge 515 map T (3) via the convex Brenier potential ψ connected to the primal and dual optimal transport problem.
Inparticular, it has been used to model convex dual functions (Makkuva et al., 2020) as well as normalizing flows 517 derived from convex potentials (Huang et al., 2021).
The expressivity and universal approximation propertiesof ICNNs have been further studied by Chen et al.
(2019), who show that any convex function over a compact 519 convex domain can be approximated in sup norm by an ICNN.
To improve convergence and robustness of ICNNs—known to be notoriously difficult to train (Richter-Powell et al., 2021)— different initialization schemes have 521 been proposed: Bunne et al.
(2022b) derive two initialization schemes ensuring that upon initialization ∇ ψmimics an affine Monge map T mapping either the source measure onto itself (identity initialization) or providing 523 a map between Gaussian approximations of measures µ and ν (Gaussian initialization).
Further, Korotin et al.(2020) proposed to use quadratic layers as well as a pre-training pipeline to initialize ICNN parameters to encode 525 an identity map.
B Additional Experimental ResultsB.1 Synthetic Data
In our synthetic two-dimensional dataset, the source and target distribution are mixtures of Gaussians with 529 varying proportions (see Fig. 2).
Both source and target consist of three corresponding clusters, and by changingthe proportions of each cluster, we illustrate a scenario in which subpopulations grow and shrink at different 531 rates.
Table 1 shows the shares of the three clusters in the source and target distributions.
In order to match thetarget distribution without transporting mass across non-corresponding clusters, the clusters have to be re-scaled 533 with the factors presented in column ’True Scaling Factor’.
The last two columns show the mean weights percluster obtained by N UB OT and UB OT GAN, respectively.
UB OT GAN captures only the general trend in 535 growth and shrinkage, the exact weights do not scale the cluster proportions appropriately.
In contrast, theweights obtained by N UB OT match the required scaling factors very closely.
Fig.
6, shows the weighted MMD 537 between the source distribution and the target distribution, confirming superior performance of N UB OT.
fit of the predicted samples to the target samples on synthetic data, measured by a weighted version of kernel MMD
B.2 Single-Cell Perturbation Responses
[Table]
Trametinib proliferation markers, respectively.
Figures 7, 8 and
UMAP projections computed on control cells for 542 the drugs Trametinib, Ixazomib and Vindesine.
In Figure 8 c., d., and Figure 9 c., d., regions of low predictedweights accurately correspond to regions of increased ClCasp3 intensity.
Additionally, we compare predicted 544 weights between the two cell types, and contrast them with observed cell counts.
C Datasets
We evaluate N UB OT on several tasks including synthetic data as well as perturbation responses of single cells.
In both settings, we are provided with unpaired measures µ and ν and aim to recover map T which describeshow source µ transforms into target ν .
While in the synthetic data setting we are provided with a ground truth 549 matching, this is not the case for the single-cell data as measuring a cell requires destroying it.
In the following,we describe generation and characteristics of both datasets, as well as introduce additional biological insights 551 allowing us to shed light on the learned matching T .
C.1 Synthetic Data
To evaluate N UB OT in a simple and low-dimensional setup with known ground-truth, we generate synthetic 554 example: We model a source population with clear subpopulation structure through a mixture of Gaussians.
Next,we generate a second (target) population aligned to the source population.
We then simulate an intervention to 556 which the subpopulations respond differently, including different levels of growth and death.
Specifically, wegenerate batches of 400 samples with three clusters with different proportions before and after the intervention.
558 Table 1 shows the proportions of the three clusters in the source and target distribution, as well as the requiredweight-factor and the obtained results from N UB OT and UB OT GAN.
[Figure]
h
[Figure] [Figure]
C.2 Single-Cell Data
Biological experiment.
The single-cell dataset used in this work was generated by the a multiplexed 562 microscopy technology called Iterative Indirect Immunofluorescence Imaging (4i) (Gut et al., 2018), which iscapable of measuring the abundance and localization of many proteins in cells.
By iteratively adding, imaging and 564 removing fluorescently tagged antibodies, a multitude of protein markers is captured for each cell.
Additionally,cellular and morphological characteristics are extracted from microscopical images, such as the cell and nucleus 566 area and circularity.
This spatially resolved phenotypic dataset is rich in molecular information and providesinsights into heterogeneous responses of thousands of cells to various drugs.
Measuring different morphological 568 and signaling features captures pre-existing cell-to-cell variability which might influence perturbation effect,resulting in various different responses.
Some of these markers are of particular importance, as they provide 570 insights into the level of a cell’s growth or death as well as subpopulation identity.
We utilized a mixture oftwo melanoma tumor cell lines (M130219 and M130429) at a ratio of 1:1.
The cell lines can be differentiated 572 by the mutually exclusive expression of marker proteins.
The former is positive for Sox9, the latter for a setof four proteins which are all recognized by and antibody called MelA (Raaijmakers et al., 2015).
Cells were 574 seeded in a 384-well plate and incubated at 37C and 5% CO2 overnight.
Next, the cells were exposed to multiplecancer drugs and Dimethyl sulfoxide (DMSO) as a vehicle control for 8h and 24h after which the cells were 576 fixed and six cycles of 4i were performed TissueMAPS and the scikit-image library (Van der Walt et al., 2014)were used to process and analyze the acquired images, perform feature extraction and quality control steps using 578 semi-supervised random forest classifiers.
[Figure]
i n t e n s i t y -
[Figure]
c l e i
[Figure]
o x
[Figure]
n
[Figure]
i n t e n s i t y -
[Figure]
c l e i - S o x
[Figure]
m e a n
[Figure]
i
[Figure]
t
[Figure]
t
[Figure]
i n t e n s i t ynu c l e i - S
o x- m
e a n
[Figure]
Data generation and processing.
Our datasets contain high-dimensional single-cell data of control anddrug-treated cells measured at two time points (8 and 24 hours).
For both the 8h-dataset and the 24h-dataset, we 581 normalized the extracted intensity and morphological features by dividing each feature by its 75th percentile,computed on the control cells.
Additionally, values were transformed by a log 1 p function ( x ← log ( x + 1) ).
Intotal, our datasets consist of 48 features, of which 26 are protein marker intensities and the remaining 22 are 584 morphological features.
For each treatment, we have measured between 2000 and 3000 cells.
For training themodels, we perform a 80/20 train/test split.
We trained all models on control and treated cells for each time step 586 and each drug separately.
The considered drugs as well as their inhibition type can be found in Table 2.
Cell type assignment.
We assigned M130219 and M130429 cells to the Sox9 and MelA cell types, 588 respectively, by first training a two component Gaussian mixture model on the features ’intensity-cell-MelAmean’ and ’intensity-nuclei-Sox9-mean’ of the control cells.
Next, we used the aforementioned features and the 590 labels provided by the mixture model to train a nearest neighbor classifier, which we then used to predict thecell type labels of the drug treated cells.
The procedure was performed separately for the 8h- and 24h dataset.
592 Results of the classification can be found in Figure 10 and Figure 12 respectively.
D Experimental Details
N UB OT consists of several modules and its performance is compared against several baselines.
In the following, 595 we provide additional background on experimental details, including a description of the evaluation metrics and
[Figure]
C o un t
C e ll C o un t a. 8 h b h
[Figure]
i n t
[Figure]
n
[Figure]
i t y
[Figure]
t
[Figure]
s i t y
[Figure]
i
[Figure]
n
[Figure]
i n t e n s i t ynu c l e i -
[Figure]
o x- m
e a n
[Figure]
i n t e n s i t ynu c l e i -
[Figure]
x-
[Figure]
e a n
[Figure]
C i s p l a t i n
C o n t r o l
C r i z
o t i n i b
D a b r a f e n i b D a c a r b a z i n e
D a s a t
i n i b
D e c i t a b i n e
D e x a m e t
h a s o n e
E r l o t i n i b
E v e r o li m u s
H y d r o x y u r e a
I m a t i n i b
I x a z o m i b
L e n a li d o m i d e
M e l p h a l a n
M i d o s t a u r i n
O l a p a r i b
P a c li t a x e l
P a l b o c i c li b
R e g
o r a f e n i b
S
o r a f e n i b
S t a u r o s p
o r i n e
T e m o z
o l o m i d e
T r a m e t i n i b
U li x e r t i n i b
V i n d e s i n e 0.
[Figure]
timepoint8h24h
[Figure]
[Table]
baselines considered, as well as further information on the parameterization and hyperparameter choices made 597 for N UB OT.
D.1 Evaluation Metrics
We evaluate our model by analyzing the distributional similarity between the predicted and observed perturbed 600 distribution.
For this, we compute the kernel maximum mean discrepancy (MMD) (Gretton et al., 2012).
To takethe mass variation into consideration, we compute a weighted version of MMD, by weighting each predicted 602 point by its associated normalized weight.
D.2 Baselines
We compare N UB OT against several baselines, comprising a balanced OT-based method (Bunne et al., 2021,
C ELL OT) and an unbalanced OT-based method (Yang & Uhler, 2019, N UB OT), i.e., current state-of-the-art 606 methods as well as ablations of our work.
In the following, we briefly motivate and introduce each baseline.
C ELL OT.
By introducing reweighting functions η and ζ , N UB OT recovers a balanced problem parameterized 608 by dual potentials f and g .
An important ablation study to consider is thus to compare its performance to itsbalanced counterpart.
Ignoring the fact that
UB OT GAN.
Using (6), Yang & Uhler (2019) propose to model mass variation in unbalanced OT via a 614 relaxation of the marginals.
Similar to Fan et al.
(2021a)
, Yang & Uhler (2019) reformulate the constrained
Monge problem (3) as a saddle point problem with Lagrange multiplier h for the constraint T ♯ µ = ν , i.e.,
[Equation]
parameterizing T and h via neural networks.
To allow mass to be created and destroyed, Yang & Uhler (2019) 617 introduce scaling factor ξ : X → R + , allowing to scale mass of each source point x i .
The optimal solution thenneeds to balance the cost of mass and the cost of transport, potentially measured through different cost functions 619 c 1 : X × Y → R + (cost of mass transport) and c 2 : R + → R + (cost of mass variation).
Parameterizing thetransport map T θ , the scaling factors ξ ϕ , and the penalty h ω with neural networks, the resulting objective isl ( θ, ϕ, ω ) :
=
[Equation]
n
[Equation]
(cid:88)
[Equation]
with Ψ ∗ approximating the divergence term of the relaxed marginal constraints (see (6)), and is optimized viaalternating gradient updates.
I DENTITY .
A trivial baseline is to compare the predictions to a map which does not model any perturbation 624 effect.
The I DENTITY baseline thus models an identity map and provides an upper bound on the overallperformance, also considered in Bunne et al.
O BSERVED .
In a similar fashion we might ask for a lower bound on the performance.
As a ground truth 627 matching is not available, we can construct a baseline for a comparison on a distributional level by comprising adifferent set of observed perturbed cells, which only vary from the true predictions up to experimental noise.
The 629 closer a method can approach the O BSERVED baseline, the more accurate it fits the perturbed cell population.
D.3 Hyperparameters
We parameterize the duals f and g using ICNNs with 4 hidden layers, each of size 64, using ReLU as activation 632 function between the layers.
We choose the identity initialization scheme introduced by Bunne et al.
(2022b)such that ∇ g and ∇ f resemble the identity function in the first training iteration.
As suggested by Makkuva 634 et al.
(2020), we relax the convexity constraint on ICNN g and instead penalize its negative weights W zl
[Equation]
The convexity constraint on ICNN f is enforced after each update by setting the negative weights of all W zl ∈ θ f 636 to zero.
Duals g and f are trained with an alternating min-max scheme where each model is trained at the samefrequency.
Further, both reweighting functions η and ζ are represented by a multi-layer perceptron (MLP) with 638 two hidden layers of size 64 for the single-cell and of size 32 for the synthetic dataset, with ReLU activationfunctions.
The final output is further passed through a softplus activation function as we do not assume negative 640 weights.
For the unbalanced Sinkhorn algorithm, we choose an entropy regularization of ε = 0 .
005 and amarginal relaxation penalty of 0 .
We use both Adam for pairs g and f as well as η and ζ with learning rate 642 10 − 4 and 10 − 3 as well as β 1 = 0 .
5 and β 2 = 0 .
9 , respectively.
We parameterize both baselines with networksof similar size and follow the implementation proposed by Yang &