Supervised Training of Conditional Monge Maps
Abstract
Optimal transport (OT) theory describes general principles to deﬁne and select, among many possible choices, the most efﬁcient way to map a probability measure onto another.
That theory has been mostly used to estimate, given a pair of source and target probability measures ( µ, ⌫ ) , a parameterized map T ✓ that can efﬁciently map µ onto ⌫ .
In many applications, such as predicting cell responses to treatments, pairs of input/output data measures ( µ, ⌫ ) that deﬁne optimal transport problems do not arise in isolation but are associated with a context c , as for instance a treatment when comparing populations of untreated and treated cells.
To account for that context in OT estimation, we introduce C OND OT, a multi-task approach to estimate a family of OT maps conditioned on a context variable, using several pairs of measures ( µ i , ⌫ i ) tagged with a context label
c i .
C OND OT learns a global map T ✓ conditioned on context that is not only expected to ﬁt all labeled pairs in the dataset { ( c i , ( µ i , ⌫ i )) } , i.e., T ✓ ( c i ) ]
µ i ⇡ ⌫ i , but should also generalize to produce meaningful maps T ✓ ( c new ) when conditioned on unseen contexts c new .
Our approach harnesses and provides a novel usage for partially input convex neural networks , for which we introduce a robust and efﬁcient initialization strategy inspired by Gaussian approximations.
We demonstrate the ability of C OND OT to infer the effect of an arbitrary combination of genetic or therapeutic perturbations on single cells, using only observations of the effects of said perturbations separately.
Introduction
A key challenge in the treatment of cancer is to predict the effect of drugs, or a combination thereof, on cells of a particular patient.
To achieve that goal, single-cell sequencing can now provide measurements for individual cells, in treated and untreated conditions, but these are, however, not in correspondence.
Given such examples of untreated and treated cells under different drugs, can we predict the effect of new drug combinations?
We develop a general approach motivated by this and related problems, through the lens of optimal transport (OT) theory , and, in that process, develop tools that might be of interest for other application domains of OT.
Given a collection of N pairs of measures ( µ i , ⌫ i ) over R d (cell measurements), tagged with a context c
i (encoding the treatment), we seek to learn a context-dependent, parameterized transport map T ✓ such that, on training data, that map T ✓ ( c i ) :
R d ! R d ﬁts the dataset, in the sense that T ✓ ( c i ) ]
µ i ⇡ ⌫ i .
Additionally, we expect that this parameterized map can generalize to unseen contexts and patients, to predict, given a patient’s cells described in µ new , the effect of applying context c new on these cells as T ✓ ( c new ) ] µ .
Learning Mappings Between Measures From generative adversarial networks, to normalizing ﬂows and diffusion models, the problem of learning maps that move points from a source to a target distribution is central to machine learning.
OT theory (Santambrogio, 2015) has emerged as a principled approach to carry out that task:
For a pair of measures µ, ⌫ supported on R d , OT suggests that, among all maps T such that ⌫ can be reconstructed by applying T to every point in the support
[Table] [Table]
e.g
, metadata or identifiers .. by action
[Table]
[Figure]
and Cuturi (2019), addresses computational challenges of estimating such maps, with impactful applications to various areas of science (cf., Hashimoto et al., 2016;
Schmitz et al., 2018; Schiebinger et al., 2019;
Yang et al., 2020; Janati et al., 2020; Bunne et al., 2022a).
Neural OT
We focus in this work on neural approaches that parameterize the optimal maps T as neural networks.
An early approach is the work on Wasserstein GANs (Arjovsky et al., 2017), albeit the transport map is not explicitly estimated.
Several recent results have exploited a more explicit connection between OT and NNs, derived from the celebrated Brenier theorem (1987), which states that Monge maps are necessarily gradients of convex functions.
Such convex functions can be represented using input convex neural networks (ICNN) (Amos et al., 2017), to parameterize either the Monge map (Jacob et al., 2018; Yang and Uhler, 2019; Bunne et al., 2021, 2022b) or a dual potential (Makkuva et al., 2020;
Korotin et al., 2020) as, respectively, the gradient of an ICNN or an ICNN itself.
In this paper, we build on this line of work, but substantially generalize it, to learn a parametric family of context-aware transport maps, using a collection of labeled pairs of measures.
Contributions We propose a framework that can leverage labeled pairs of measures { ( c i , ( µ i , ⌫ i )) }
i to infer a global parameterized map T ✓ .
Hereby, the context c i belongs to an arbitrary set C .
We construct T ✓ so that it should be able, given a possibly unseen context label c 2 C , to output a map
T ✓ ( c ) : R d ! R d , that is itself the gradient of a convex function.
To that end, we propose to learn these parameterized Monge maps T ✓ as the gradients of partially input convex neural networks (PICNN), which we borrow from the foundational work of Amos et al.
Our framework can be also interpreted as a hypernetwork (Ha et al., 2016):
The PICNN architecture can be seen as an
ICNN whose weights and biases are modulated by the context vector c , which parameterizes a family of convex potentials in R d .
Because both ICNN —and to a greater extent PICNN— are notoriously difﬁcult to train (Richter-Powell et al., 2021; Korotin et al., 2020, 2021), we use closed-form solutions between Gaussian approximations to derive relevant parameter initializations for (P)ICNNs:
These choices ensure that, upon initialization , the gradient of the (P)ICNNs mimics the afﬁne Monge map obtained in closed form between Gaussian approximations of measures µ i , ⌫ i (Gelbrich, 1990).
Our framework is applied to three scenarios: Parameterization of transport through a real variable (time or drug dosage), through an auxiliary informative variable (cell covariates) and through action variables (genetic perturbations in combination)
(see Fig. 1).
Our results demonstrate the ability of our architectures to better capture on out-of-sample observations the effects of these variables in various settings, even when considering never-seen, composite context labels.
These results suggest potential applications of conditional OT to model personalized medicine outcomes, or to guide novel experiments, where OT could serve as a predictor for never tested context labels.
Background on Neural Solvers for the 2-Wasserstein Problem
Optimal Transport The Monge problem between two measures µ, ⌫ 2 P ( R d ) , here restricted to measures supported on R d and compared with the squared Euclidean metric, reads
[Equation]
The existence of T ? is guaranteed under fairly general conditions (Santambrogio, 2015, Theorem 1.22), which require that µ and ⌫ have ﬁnite L 2 norm, and that µ puts no mass on ( d  1) surfaces of class C 2 .
This can be proved with the celebrated Brenier theorem (1987), which states that there must exist a unique (up to the addition of a constant) potential f ? : R d ! R such that T ? =
r f ? .
This theorem has far-reaching implications: It is sufﬁcient, when seeking optimal transport maps, to restrict the computational effort to seek a “good” convex potential, such that its gradient pushes µ towards ⌫ .
This result has been exploited to propose OT solvers that rely on input convex neural networks (ICNNs)
(Amos et al., 2017), introduced below
[Equation]
In practice, Monge maps can be estimated using a dual formulation (Makkuva et al., 2020; Korotin et al., 2020; Bunne et al., 2022b; Alvarez-Melis et al., 2021;
Mokrov et al., 2021).
Indeed, T ? in (1) is recovered as r f ? , where f ? is deﬁned in (2), writing f ⇤ for the Legendre transform of f .
Convex Neural Architectures Input convex neural networks (ICNN) are neural networks ✓ that admit certain constraints on their architecture and parameters ✓ , such that their output ✓ ( x ) is a convex function of their input x (Amos et al., 2017).
As a result, they have been increasingly used as drop-in replacements to the set of admissible functions in (2).
Practically speaking, an ICNN is a
K -layer, fully connected network such that, at each layer index k from 0 to K  1 , a hidden state vector z k is deﬁned recursively as in (3),
[Equation]
and ✓ ( x ) =
z K , where, by convention, z 0 and W z 0 are 0 ;  k are convex non-decreasing activation functions; ✓ = { b k , W zk , W xk } K  1 k =0 are the weights and biases of the neural network.
While ample ﬂexibility is provided to choose dimensions for intermediate hidden states z k , the last layer must necessarily produce a scalar, hence W xK  1 and W zk  1 are line vectors and b K  1 2 R .
ICNNs are characterized by the fact that all weight matrices W zk associated to latent representations z must have non-negative entries.
This, along with the speciﬁc activation functions, ensures the convexity of ✓ .
We encode this constraint by identifying these matrices as the elementwise softplus or ReLU of other matrices of the same size, or, alternatively, using a regularizer that penalizes the negative entries of these matrices.
Since the work by Amos et al.
(2017), convex neural architectures have been used within the context of OT to model convex dual functions (Makkuva et al., 2020), or normalizing ﬂows derived from convex potentials (Huang et al., 2021).
Their expressivity and universal approximation properties have been studied by Chen et al.
(2019), who show that any convex function over a compact convex domain can be approximated in sup norm by an ICNN.
Supervised Training of Conditional Monge Maps
We are given a dataset of N pairs of measures, each endowed with a label, ( c i , ( µ i , ⌫ i ))
2 C ⇥ P ( R d ) 2 .
Our framework builds upon two pillars: (i.) we formulate the hypothesis that an optimal transport T ? i (or, equivalently, the gradient of a convex potential f ? i ) explains how measure µ i was mapped to ⌫ i , given context c i ; (ii.)
we build on the multi-task hypothesis (Caruana, 1997) that all of the N maps
T ? i between µ i and ⌫ i share a common set of parameters, that are modulated by context informations
c i .
These ideas are summarized in an abstract regression model described below.
3.1 A Regression Formulation for Conditional OT Estimation
✓ 2 ⇥ ⇢ R r , T ✓ describes a function that takes an input vector c 2 C , and outputs a function
T ✓ ( c ) : R d ! R d , as a hypernetwork would (Ha et al., 2016).
Assume momentarily that we are given ground truth maps T i , that describe the effect of context c i on any measure, rather only pairs of measures ( µ i , ⌫ i ) .
This is of course a major leap of faith, since even recovering an OT map T ? from two measures is in itself very challenging (Hütter and Rigollet, 2021;
Rigollet and Stromme, 2022;
Pooladian and Niles-Weed, 2021).
If such maps were available, a direct supervised approach to learn a unique ✓ could hypothetically involve minimizing a ﬁt function composed of losses between maps
[Equation]
Unfortunately, such maps T i are not given, since we are only provided unpaired samples before µ i and after ⌫ i that map’s application.
By Brenier’s theorem, we know, however, that such an OT map
T ? i exists, and that it would be necessarily the gradient of a convex potential function that maximizes (2).
As a result, we propose to modify (4) to (i.) parameterize, for any c , the map T ✓ ( c ) as the gradient w.r.t.
x of a function f ✓ ( x, c ) : R d ⇥ C ! R that is convex w.r.t.
x , namely T ✓ ( c ) := x 7!
r 1 f ✓ ( x, c ) ; (ii.) estimate ✓ by maximizing jointly the dual objectives (2) simultaneously for all N pairs of measures, in order to ensure that the maps are close to optimal, to form the aggregate problem
[Equation]
We detail in App.
B how the Legendre transforms that appear in the energy terms E µ i , ⌫ i are handled with an auxiliary function.
3.2 Integrating Context in Convex Architectures
We propose to incorporate context variables, in order to modulate a family of convex functions f ✓ ( x, c ) using partially input convex neural networks (PICNN).
PICNNs are neural networks that can be evaluated over a pair of inputs ( x, c ) , but which are only required to be convex w.r.t.
x .
Given an input vector x and context vector c , a K -layer PICNN is deﬁned as ✓ ( x, c ) =
z K , where, recursively for 0  k  K  1 one has u
[Equation]
where the PICNN is initialized as u 0 =
c, z [Equation] ,  denotes the Hadamard elementwise product, and ⌧ k is any activation function.
The parameters of the PICNN are then given by
[Equation]
Similar to ICNNs, the convexity w.r.t.
input variable x is guaranteed as long as activation functions  i are convex and non-decreasing, and the weight matrices
W zk have non-negative entries.
We parameterize this by storing them as elementwise applications of softplus operations on precursor matrices of the same size, or, alternatively, by regularizing their negative part.
Finally, much like
ICNNs, all matrices at the K  1 layer are line vectors, and their biases scalars.
Such networks were proposed by Amos et al.
(2017, Eq. 3) to address a problem that is somewhat symmetric to ours: Their inputs were labeled as ( y, x ) , where y is a label vector, typically much smaller than that of vector x .
Their PICNN is convex w.r.t.
y , in order to easily recover, given a datapoint x
(e.g., an image) the best label y that corresponds to x using gradient descent as a subroutine, i.e. y ? ( x ) =
arg min y PICNN ✓ ( x, y ) .
PICNN were therefore originally proposed to learn a parameterized, implicit classiﬁcation layer, amortized over samples, whose motivation rests on the property that it is convex w.r.t.
label variable y .
By contrast, we use PICNNs that are convex w.r.t.
data points x .
In addition to that swap, we do not use the convexity of the PICNN to deﬁne an implicit layer (or to carry out gradient descent).
Indeed, it does not make sense in our setting to minimize ✓ ( x, c ) as a function of x , since x is an observation.
Instead, our work rests on the property that r 1 ✓ ( x, c ) describes a parameterized family of OT maps.
We note that PICNNs were considered within the context of OT in (Fan et al., 2021, Appendix B).
In that work, PICNN provide an elegant reformulation for neural Wasserstein barycenters.
Fan et al.
(2021) considered a context vector c that was restricted to be a small vector of probabilities.
3.3 Conditional Monge Map Architecture
Using PICNNs as a base module, the C OND OT architecture integrates operations on the contexts C .
As seen in Figure 1, context values c may take various forms:
A scalar t denoting a strength or a temporal effect.
For instance, McCann’s interpolation and its time parameterization, ↵ t =
((1  t ) Id + tT ) ] ↵ 0 (McCann, 1997) can be interpreted as a trivial conditional OT model that creates, from an OT map ✓ , a set of maps parameterized by t , T ✓
[Equation]
A covariate vector inﬂuencing the nature of the effect that led µ i to ⌫ i , (capturing, e.g., patient feature vectors).
One or multiple actions, possibly discrete, representing decisions or perturbations applied onto µ i .
To provide a ﬂexible architecture capable of modeling different types of conditions as well as conditions appearing in combinations, the more general C OND OT architecture consists of the hypernetwork T ✓ that is fed a context vector through embedding and combinator modules.
This generic architecture provides a one-size ﬁts all approach to integrate all types of contexts c .
Embedding Module To give greater ﬂexibility when setting the context variable c , C OND OT contains an embedding module E that translates arbitrary contexts into real-valued vectors.
Besides simple scalars t (Fig.
1a) for which no embedding is required, discrete contexts can be handled with an embedding module E  .
When the set C is small, this can be done effectively using one-hot embeddings
E ohe .
For more complicated actions a such as treatments, there is no simple way to vectorize a context c .
Similarly to action embeddings in reinforcement learning (Chandak et al., 2019;
Tennenholtz and
Mannor, 2019), we can learn embeddings for discrete actions into a learned continuous representation.
This often requires domain-knowledge on the context values.
For molecular drugs, for example, we can learn molecular representations E mol such as chemical, motif-based (Rogers and Hahn, 2010) or neural ﬁngerprints (Rong et al., 2020; Schwaller et al., 2022).
However, often this domain knowledge is not available.
In this work, we thus construct so-called mode-of-action embeddings, by computing an embedding E moa that encourages actions a with similar effect on target population ⌫ to have a similar representation.
In § 5, we analyze several embedding types for different use-cases.
[Equation]
[Table]
⌫ embeddingcombinator map i npu t s [Table]
_base64="F/KxlTcdBmmuH9nzvCI8RkzkvWQ=">AAAB6nicbVDLSgNBEOyJj8T4inoUZTAInsKuiHoMeFFPEc0DkiXMTmaTIbOzy8ysEpZ8ghcPinj1i7z5DX6ETh4HTSxoKKq66e7yY8G1cZxPlFlYXFrO5lbyq2vrG5uFre2ajhJFWZVGIlINn2gmuGRVw41gjVgxEvqC1f3+xciv3zOleSTvzCBmXki6kgecEmOlW9rutwtFp+SMgeeJOyXF8v73V/b6Ya/SLny0OhFNQiYNFUTrpuvExkuJMpwKNsy3Es1iQvuky5qWShIy7aXjU4f40CodHETKljR4rP6eSEmo9SD0bWdITE/PeiPxP6+ZmODcS7mME8MknSwKEoFNhEd/4w5XjBoxsIRQxe2tmPaIItTYdPI2BHf25XlSOy65p6WTG5vGFUyQg104gCNw4QzKcAkVqAKFLjzCM7wggZ7QK3qbtGbQdGYH/gC9/wCOy5D+</latexit
[Figure]
Combinator Module While we often have access to contexts c in isolation, it is crucial to infer the effect of contexts applied in combination.
A prominent example are cancer combination therapies, in which multiple treatment modalities are administered in combination to enhance treatment efﬁcacy (Kummar et al., 2010).
In these settings, the mode of operation between individual contexts c is often not known, and can thus not be directly modeled via simple arithmetic operations such as min, max, sum, mean .
While we test as a baseline the case, applicable to one-hot-embeddings, where simple additions are used to model these combinations, we propose to augment the
C OND OT architecture with a parameterized combinator module C  .
If the order in which the actions are applied is irrelevant or unknown, the corresponding network C  needs to be permutation-invariant, which can be achieved by using a deep set architecture (Zaheer et al., 2017).
Receiving a ﬂexible number of inputs from the embedding module E  , C OND OT allows for a joint training of the PICNN parameters ✓ , embedding parameters  , and combinator parameters  in a single, end-to-end differentiable architecture.
Training Procedure Given a dataset D
= { c i , ( µ i , ⌫ i ) } Ni =0 of N pairs of populations before µ i and after transport ⌫ i connected to a context c
i , we detail in Algorithm 1 provided in § B, a training loop that incorporates all of the architecture proposals described above.
The training loss aims at making sure the map T ✓ ( c i ) is an OT map from µ i to ⌫ i , where c i may either be the original label itself or its embedded/combined formulation in more advanced tasks.
To handle the Legendre transform in (2), we use the proxy dual objective deﬁned in (Makkuva et al., 2020, Eq. 6) (15)-(16) in place of (2) in our overall loss (5).
This involves training the C OND OT architecture using two PICNNs, i.e., PICNN ✓ f and PICNN ✓ g , that share the same embedding/combinator module, with a regularization (14) promoting that for any c , the PICNN ✓ g ( · , c ) resembles the Legendre transform of the other, PICNN ⇤ ✓ f ( · , c ) .
Initialization Strategies for Neural Convex Architectures
We address the problem of initializing the parameters of (P)ICNNs to ensure their gradient evaluated at every point is (initially) meaningful in the context of OT, namely that it is able to map the ﬁrst and second moments of a measure µ into those of a target measure ⌫ .
The initializers we propose build heavily on the quadratic layers proposed in the seminal reference (Korotin et al., 2020,
Appendix B.2), notably the “DenseQuad” layer, as well as on closed-form solutions available for
Gaussian approximations of measures (Gelbrich, 1990).
Closed-Form Potentials for Gaussians Given two Gaussian distributions N 1 , N 2 with means respectively m 1 , m 2 and covariance matrices ⌃ 1 , ⌃ 2 (where ⌃ 1 is assumed to be full rank), the
[Table] [Table]
+DBGZTgEspQAQIdeIRneHG48+S8Om/j1jlnMrMNf+C8/wA565DG</latexit> [Figure] <latexit sha1_base64="L5qrL0YA6qFN6oe0FDcO/nW4LxQ=">AAAB7XicbVDLSgNBEOyNrxhfUY9eBoPgKeyKqDcDXvQWwTwgWcPsZJKMmd1ZZnqFsOQfvHhQxKu/I978GyebHDSxoKGo6qa7K4ilMOi6305uaXlldS2/XtjY3NreKe7u1Y1KNOM1pqTSzYAaLkXEayhQ8masOQ0DyRvB8GriNx65NkJFdziKuR/SfiR6glG0Up3dt1HFnWLJLbsZyCLxZqR0+QkZqp3iV7urWBLyCJmkxrQ8N0Y/pRoFk3xcaCeGx5QNaZ+3LI1oyI2fZteOyZFVuqSntK0ISab+nkhpaMwoDGxnSHFg5r2J+J/XSrB34aciihPkEZsu6iWSoCKT10lXaM5QjiyhTAt7K2EDqilDG1DBhuDNv7xI6idl76x8euuWKjfTNCAPB3AIx+DBOVTgGqpQAwYP8AQv8Ooo59l5c96nrTlnNrMPf+B8/AAUKpBB
[Figure]
known labels to modulate the magnitude of each of the three potentials.
Brenier potential solving the OT problem from the ﬁrst to the second Gaussian reads:
[Equation]
deﬁne both quadratic and linear terms and t ( A, b ) can be any constant.
Importantly, note that we write the quadratic term in factorized form AA T to enforce psd-ness, as done by Korotin et al.
(2020), not as usually done with a single psd matrix (Peyré and Cuturi, 2019, Remark 2.31).
Our quadratic potentials are only injected in the ﬁrst state of hidden vector z 0 , to populate it with a collection of relevant full-rank quadratic convex functions, with the goal of recovering an afﬁne OT map from the start, as illustrated in the experiments from § C.1.
Quadratic Potentials Lower Bounded by
Naturally, for any choice of t ( A, b ) one recovers the property that r f ? µ, ⌫ ] N 1 = N 2 .
When used in deep architectures, the level of that constant does, however, play a role, since convex functions in ICNN are typically thresholded or modulated using rectifying functions.
To remove this ambiguity, we settle on a choice for t ( A, b ) that is such that the lowest value reached by f ?
N 1 , N 2 is 0.
This can be obtained by setting
[Equation]
which results in the following choice, writing !
= m 1  ( A T A [Equation]  1 m 2 ,
[Equation]
To mimic these potential functions, we introduce a quadratic layer parameterized by a weight matrix
M and a “bias” vector m , deﬁned as q M,m ( x ) = 12 k M ( x  m ) k 22 .
By design, q M,m ( x ) is a convex quadratic, non-negative layer.
Finally, one has the following relationships,
[Equation]
ICNN Initialization We explore two possible ICNN (3) initializers for OT.
Identity Initialization
The ﬁrst approach ensures that upon initialization the ICNN’s gradient mimics the identity map, i.e., r ✓ ( x ) =
x for any x .
We do so by injecting in the initial hidden state zthe norm of the input vector 12 k x k 2 , cast as a trainable layer q M,m initialized with M
=
I and m =
d , see (10).
The remaining parameters are chosen to propagate that norm throughout layers using averages.
This amounts to the following choices:
Set all  i to be activations such that {+ 
+}
i ( u ) ⇡ 1 for u large enough, e.g., (leaky) ReLU or softplus.
Introduce an initialization layer, z 0 = q M,m ( x ) 1 , itself initialized with M =
I and m =
0 d . 3.
Initialize all matrices W zi to ⇡ 1 d 2 ,d 1 /d 1 , where d 1 , d 2 are the dimensions of these matrices.
Initialize all matrices W xi to ⇡ 0 .
Initialize biases b i to s 1 , where s is a large enough value s so that {+ 
+}
i ( s ) ⇡ 1 .
Gaussian Initialization
The second approach can be used to initialize an ICNN so that its gradient mimics the afﬁne transport between the Gaussian approximations of µ and ⌫ .
To this end, we follow all of the steps outlined above, except for step 2 where the quadratic layer q M,m is initialized instead with M =
A and m =
m 1  ( A T A )  1 m 2 using notations in (7), (8), (9), where m 1 , m 2 , ⌃ 1 , ⌃
[Table]
(cid:19) (cid:21) (cid:23) (cid:19)(cid:17)(cid:19) (cid:19)(cid:17)(cid:21) (cid:19)(cid:17)(cid:23) (cid:19)(cid:17)(cid:25) (cid:19)(cid:17)(cid:27) (cid:20)(cid:17)(cid:19)
[Table]
(cid:19) (cid:20) (cid:21) (cid:22) (cid:19)(cid:17)(cid:19) (cid:19)(cid:17)(cid:24) (cid:20)(cid:17)(cid:19) (cid:20)(cid:17)(cid:24) (cid:21)(cid:17)(cid:19) (cid:21)(cid:17)(cid:24) (cid:22)(cid:17)(cid:19) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:20)(cid:23)(cid:24)(cid:28)(cid:21)(cid:19)(cid:17)(cid:20)(cid:23) (cid:19) (cid:20) (cid:21) (cid:22) (cid:23) (cid:19)(cid:17)(cid:19) (cid:19)(cid:17)(cid:24) (cid:20)(cid:17)(cid:19) (cid:20)(cid:17)(cid:24) (cid:21)(cid:17)(cid:19) (cid:21)(cid:17)(cid:24) (cid:22)(cid:17)(cid:19) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:20)(cid:20)(cid:26)(cid:26)(cid:21)(cid:23)(cid:17)(cid:20)(cid:21) (cid:19)(cid:17)(cid:19) (cid:19)(cid:17)(cid:24) (cid:20)(cid:17)(cid:19) (cid:20)(cid:17)(cid:24) (cid:19) (cid:21)(cid:19) (cid:23)(cid:19) (cid:25)(cid:19) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:21)(cid:26)(cid:22)(cid:26)(cid:26)(cid:26)(cid:17)(cid:23) (cid:2255)(cid:21) (cid:2255)(cid:20) (cid:19) (cid:20) (cid:21) (cid:22) (cid:23) (cid:19)(cid:17)(cid:19) (cid:19)(cid:17)(cid:21) (cid:19)(cid:17)(cid:23) (cid:19)(cid:17)(cid:25) (cid:19)(cid:17)(cid:27) (cid:20)(cid:17)(cid:19) (cid:20)(cid:17)(cid:21) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:20)(cid:26)(cid:22)(cid:23)(cid:27)(cid:21)(cid:17)(cid:20)(cid:25) (cid:2255)(cid:20) (cid:19) (cid:20) (cid:21) (cid:22) (cid:23) (cid:19)(cid:17)(cid:19) (cid:19)(cid:17)(cid:24) (cid:20)(cid:17)(cid:19) (cid:20)(cid:17)(cid:24) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:20)(cid:23)(cid:26)(cid:23)(cid:19)(cid:27)(cid:17)(cid:20)(cid:23) [Table] (cid:19)(cid:17)(cid:19)
[Table]
(cid:19) (cid:21) (cid:23) (cid:25) (cid:27) (cid:20)(cid:19)
[Table]
(cid:19) (cid:20)(cid:19) (cid:21)(cid:19) (cid:22)(cid:19) (cid:23)(cid:19) (cid:24)(cid:19) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:20)(cid:23)(cid:24)(cid:28)(cid:21)(cid:19)(cid:17)(cid:20)(cid:23) (cid:19) (cid:20) (cid:21)
[Table]
(cid:19)(cid:17)(cid:24) (cid:20)(cid:17)(cid:19) (cid:20)(cid:17)(cid:24) (cid:21)(cid:17)(cid:19) (cid:21)(cid:17)(cid:24) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:20)(cid:20)(cid:26)(cid:26)(cid:21)(cid:23)(cid:17)(cid:20)(cid:21) (cid:19)(cid:17)(cid:19) (cid:19)(cid:17)(cid:24) (cid:20)(cid:17)(cid:19) (cid:20)(cid:17)(cid:24) (cid:19) (cid:21)(cid:24) (cid:24)(cid:19) (cid:26)(cid:24) (cid:20)(cid:19)(cid:19) (cid:20)(cid:21)(cid:24) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:21)(cid:26)(cid:22)(cid:26)(cid:26)(cid:26)(cid:17)(cid:23) (cid:2255)(cid:19)(cid:17)(cid:21)(cid:24) (cid:19)(cid:17)(cid:19)(cid:19) (cid:19)(cid:17)(cid:21)(cid:24) (cid:19)(cid:17)(cid:24)(cid:19) (cid:19)(cid:17)(cid:26)(cid:24) (cid:20)(cid:17)(cid:19)(cid:19) (cid:20)(cid:17)(cid:21)(cid:24) (cid:19) (cid:24) (cid:20)(cid:19) (cid:20)(cid:24) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:20)(cid:26)(cid:22)(cid:23)(cid:27)(cid:21)(cid:17)(cid:20)(cid:25)
[Table]
(cid:2255)(cid:19)(cid:17)(cid:24) (cid:19)(cid:17)(cid:19) (cid:19)(cid:17)(cid:24) (cid:20)(cid:17)(cid:19) (cid:20)(cid:17)(cid:24) (cid:21)(cid:17)(cid:19) (cid:21)(cid:17)(cid:24) (cid:19) (cid:20) (cid:21) (cid:22) (cid:23) (cid:24) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:20)(cid:19)(cid:27)(cid:25)(cid:19)(cid:21)(cid:17)(cid:20)(cid:26) (cid:2255)(cid:19)(cid:17)(cid:21) (cid:2255)(cid:19)(cid:17)(cid:20) (cid:19)(cid:17)(cid:19)
[Table]
(cid:19)(cid:17)(cid:22) (cid:19)(cid:17)(cid:23) (cid:19) (cid:21)(cid:19) (cid:23)(cid:19) (cid:25)(cid:19) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:20)(cid:23)(cid:24)(cid:28)(cid:21)(cid:19)(cid:17)(cid:20)(cid:23) (cid:2255)(cid:20) (cid:19) (cid:20) (cid:21) (cid:22) (cid:23) (cid:19)(cid:17)(cid:19) (cid:19)(cid:17)(cid:24) (cid:20)(cid:17)(cid:19) (cid:20)(cid:17)(cid:24) (cid:21)(cid:17)(cid:19) (cid:21)(cid:17)(cid:24)
[Table]
(cid:20)(cid:19)(cid:19) (cid:21)(cid:19)(cid:19) (cid:22)(cid:19)(cid:19) (cid:23)(cid:19)(cid:19)
[Table]
(cid:2255)(cid:19)(cid:17)(cid:24)
[Table]
(cid:21) (cid:23) (cid:25) (cid:27) (cid:20)(cid:19)
[Table]
K562 ENSG00000165092.
[Table]
(cid:19)
[Table]
(cid:23)
[Table]
(cid:19)(cid:17)(cid:19) (cid:21)(cid:17)(cid:24) (cid:24)(cid:17)(cid:19)
[Table]
(cid:20)(cid:19)(cid:17)(cid:19) (cid:20)(cid:21)(cid:17)(cid:24) [Table] (cid:19)(cid:17)(cid:19) (cid:19)(cid:17)(cid:24) (cid:20)(cid:17)(cid:19) (cid:20)(cid:17)(cid:24) (cid:19) (cid:21) (cid:23) (cid:25) (cid:27) (cid:20)(cid:19) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:20)(cid:19)(cid:27)(cid:25)(cid:19)(cid:21)(cid:17)(cid:20)(cid:26) (cid:19)(cid:17)(cid:19) (cid:19)(cid:17)(cid:21) (cid:19)(cid:17)(cid:23) (cid:19)(cid:17)(cid:25) (cid:19) (cid:20)(cid:19) (cid:21)(cid:19) (cid:22)(cid:19) (cid:23)(cid:19) (cid:24)(cid:19) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:20)(cid:23)(cid:24)(cid:28)(cid:21)(cid:19)(cid:17)(cid:20)(cid:23) (cid:19) (cid:20) (cid:21) (cid:22) (cid:19)(cid:17)(cid:19) (cid:19)(cid:17)(cid:24) (cid:20)(cid:17)(cid:19) (cid:20)(cid:17)(cid:24) (cid:21)(cid:17)(cid:19) (cid:21)(cid:17)(cid:24) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:20)(cid:20)(cid:26)(cid:26)(cid:21)(cid:23)(cid:17)(cid:20)(cid:21) (cid:19)(cid:17)(cid:19) (cid:19)(cid:17)(cid:24) (cid:20)(cid:17)(cid:19) (cid:20)(cid:17)(cid:24) (cid:19) (cid:21)(cid:24) (cid:24)(cid:19) (cid:26)(cid:24) (cid:20)(cid:19)(cid:19) (cid:20)(cid:21)(cid:24) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:21)(cid:26)(cid:22)(cid:26)(cid:26)(cid:26)(cid:17)(cid:23) (cid:2255)(cid:19)(cid:17)(cid:21)(cid:24)
[Table]
(cid:24) (cid:20)(cid:19) (cid:20)(cid:24)
[Table]
(cid:19)(cid:17)(cid:19) (cid:19)(cid:17)(cid:24) (cid:20)(cid:17)(cid:19) (cid:20)(cid:17)(cid:24) (cid:19) (cid:20)(cid:19) (cid:21)(cid:19) (cid:22)(cid:19) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:21)(cid:24)(cid:27)(cid:26)(cid:22)(cid:22)(cid:17)(cid:24) (cid:2255)(cid:19)(cid:17)(cid:20) (cid:19)(cid:17)(cid:19) (cid:19)(cid:17)(cid:20) (cid:19)(cid:17)(cid:21) (cid:19) (cid:21)(cid:24) (cid:24)(cid:19) (cid:26)(cid:24) (cid:20)(cid:19)(cid:19) (cid:20)(cid:21)(cid:24) (cid:20)(cid:24)(cid:19) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:20)(cid:19)(cid:25)(cid:26)(cid:20)(cid:23)(cid:17)(cid:20)(cid:26) (cid:19) (cid:20) (cid:21) (cid:19) (cid:20) (cid:21) (cid:22) (cid:40)(cid:49)(cid:54)(cid:42)(cid:19)(cid:19)(cid:19)(cid:19)(cid:19)(cid:20)(cid:27)(cid:22)(cid:19)(cid:22)(cid:21)(cid:17)(cid:20)(cid:20)
[Table]
[Figure]
Lotfollahiet al
[Figure]
are replaced by the empirical mean and covariances of µ and ⌫ .
Throughout the experiments, we use the Gaussian and identity initialization.
Further comparisons between the vanilla initialization and those introduced in this work can be found in § C.1 (Fig.
PICNN Initialization Recall for convenience that a K -layer PICNN architecture reads:
[Equation]
In their original form (Amos et al., 2017, Eq. 3), PICNNs are initialized by setting u
= c and z 0 = 0 to a zero vector of suitable size.
Intuitively, the hidden states u k act as context-dependent modulators, whereas vectors z k propagate, layer after layer, a collection of convex functions in x that are iteratively reﬁned, while retaining the property that they are each convex in x .
A reasonable initialization for a PICNN that is provided a context vector c is that if c ⇡ c j (where j is in the training set), one has that r 1 ✓ 0 ( · , c ) maps approximately µ j to ⌫ j , which can be obtained by having ✓ 0 ( · , c ) mimic the closed-form Brenier potential between the Gaussian approximations of µ j , ⌫ j .
Alternatively, one may also default to an identity initialization as discusses above.
To obtain either behavior, we make the following modiﬁcations, and refer to the illustration in Fig.
The modulator u 0 ( c ) = softmax ( c T M ) , where C is initialized as M =
[ c j ] j , and V i
=
I, v i = 0 .
z
=
[ q M j ,m j ( x )] j , where weight matrices and bias ( M j , m j ) are either initialized to ( I, 0 ) or as ( A j , ! j ) recovered by solving the Gaussian afﬁne map from µ j to ⌫ j using (9).
Modulator u 0 is passed directly to hidden state upon ﬁrst iteration W zu
=
I, b z 0 = 0 .
All subsequent matrices W zk are initialized to ⇡ 1 d 2 ,d 1 /d 1 , where d 1 , d 2 are their dimensions, 5.
W x k and W xu k are ⇡ 0 , the biases b z k ⇡ 1 , b u k ⇡ 0 .
Evaluation
Biological cells undergo changes in their molecular proﬁles upon chemical, genetic, or mechanical perturbations.
These changes can be measured using recent technological advancements in high-
W a ss e r s t e i n L o ss o f I
CNN
O T
W a ss e r s t e i n L o ss o f I
CNN
O T
W a ss e r s t e i n L o ss o f C
P A
Wasserstein Loss of CondOT
Wasserstein Loss of CondOT Wasserstein Loss of CondOT
Known Perturbations Unknown Perturbations a. b. c. MoA Embedding
One-Hot Embedding
One-Hot Embedding
[Figure]
resolution multivariate single-cell biology.
Measuring single cells in their unperturbed or perturbed state requires, however, to destroy them, resulting in populations µ and ⌫ that are unpaired.
The relevance of OT to that comes from its ability to resolve such ambiguities through OT maps, holding promises of a better understanding of health and disease.
We consider various high-dimensional problems arising from this scenario to evaluate the performance of C OND OT (§ 3) versus other baselines.
5.1 Population Dynamics Conditioned on Scalars
Upon application of a molecular drug, the state of each cell x i of the unperturbed population is altered, and observed in population ⌫ .
Molecular drugs are often applied at different dosage levels t , and the magnitude of changes in the gene expression proﬁles of single cells highly correlates with that dosage.
We seek to learn a global, parameterized transport map T ✓ sensitive to that dosage.
We evaluate our method on the task of inferring single-cell perturbation responses to the cancer drug Givinostat, a histone deacetylase inhibitor with potential anti-inﬂammatory, anti-angiogenic, and antineoplastic activities (Srivatsan et al., 2020), applied at different dosage levels, i.e., t 2 { 10 nM, 100 nM, 1 , 000 nM, 10 , 000 nM } .
The dataset contains 3 , 541 cells described with the gene expression levels of 1 , 000 highly-variable genes.
In a ﬁrst experiment, we measure how well C OND OT captures the drug effects at different dosage levels via distributional distances such as MMD (Gretton et al., 2012) and the ` 2 -norm between the corresponding perturbation signatures (PS), as well as the entropy-regularized Wasserstein distance (Cuturi, 2013).
We compute the metrics on 50 marker genes, i.e., genes mostly affected upon perturbation.
For more details on evaluation metrics, see § E.2.
To put C OND OT’s performance into perspective, we compare it to current state-of-the-art baselines (Lotfollahi et al., 2021) as well as parameterized Monge maps without context variables (Bunne et al., 2021; Makkuva et al., 2020, ICNN OT), see § E.1.
As visible in Table 1 and Fig.
4a, C OND OT achieves consistently more accurate predictions of the target cell populations at different dosage levels than OT approaches that cannot utilize context information, demonstrated through a lower average loss and a smaller variance.
This becomes even more evident when moving to the setting where the population has been trained only on a subset of dosages and we test C OND OT on out-of-sample dosages.
Table 1 and Fig.
4a demonstrate that C OND OT is able to generalize to previously unknown dosages, thus learning to interpolate the perturbation effects from dosages seen during training.
For further analysis, we refer the reader to § E (see Fig.
9 and 10).
We further provide an additional comparison of C OND OT, operating in the multi-task setting, to the single-task performance of optimal transport-based methods § C.4.
While the single-task setting of course fails to generalize to new contexts and requires all contexts to be distinctly known, it provides us with a pseudo lower bound, which C OND OT is able to reach (see Table 2).
5.2 Population Dynamics Conditioned on Covariates
Molecular processes are often highly dependent on additional covariates that steer experimental conditions, and which are not present in the features measures in population µ or ⌫ .
This can be, for instance, factors such as different cell types clustered within the populations.
When the model can only
[Table]
n
[Table]
a. b. c.
[Table]
[Figure]
be conditioned w.r.t.
a small and ﬁxed set of metadata information, such as cell types, it is sufﬁcient to encode these contexts using a one-hot embedding module E ohe .
To illustrate this problem, we consider cell populations comprising three different cell lines (A549, MCF7, and K562).
As visible in Table 1,
C OND OT outperforms current baselines which equally condition on covariate information such as
CPA (Lotfollahi et al., 2021), assessed through various evaluation metrics.
Figure 4b displays a gene showing highly various responses towards the drug Givinostat dependent on the cell line.
C OND OT captures the distribution shift from control to target populations consistently across different cell lines.
5.3 Population Dynamics Conditioned on Actions
To recommend personalized medical procedures for patients, or to improve our understanding of genetic circuits, it is key to be able to predict the outcomes of novel perturbations, arising from combinations of drugs or of genetic perturbations.
Rather than learning individual maps T a ✓ predicting the effect of individual treatments, we aim at learning a global map T ✓ which, given as input the unperturbed population µ as well as the action a of interest, predicts the cell state perturbed by a .
Thanks to its modularity, C OND OT can not only learn a map T ✓ for all actions known during training, but also to generalize to unknown actions, as well as potential combinations of actions.
We will discuss all three scenarios below.
5.3.1 Known Actions
In the following, we analyze C OND OT’s ability to accurately predict phenotypes of genetic perturbations based on single-cell RNA-sequencing pooled CRISPR screens (Norman et al., 2019; Dixit et al., 2016), comprising 98 , 419 single-cell gene expression proﬁles with 92 different genetic perturbations, each cell measured via a 1 , 500 highly-variable gene expression vector.
As, in a ﬁrst step, we do not aim at generalizing beyond perturbations encountered during training, we utilize again a one-hot embedding E ohe to condition T ✓ on each perturbation a .
We compare our method to other baselines capable of modeling effects of a large set of perturbations such as CPA (Lotfollahi et al., 2021).
Often, the effect of genetic perturbations are subtle in the high-dimensional gene expression proﬁle of single cells.
Using ICNN-parameterized OT maps without context information, we can thus assess the gain in accuracy of predicting the perturbed target population by incorporating context-awareness over simply predicting an average perturbation effect.
Figure 5a and b demonstrate that compared to OT ablation studies, Fig. 5c and Fig.
6a for the current state-of-the-art method CPA (Lotfollahi et al., 2021).
Compared to both, C OND OT captures the perturbation responses more accurately w.r.t.
the Wasserstein distance.
5.3.2 Unknown Actions
With the emergence of new perturbations or drugs, we aim at inferring cellular responses to settings not explored during training.
One-hot embeddings, however, do not allow us to model unknown perturbations.
This requires us to use an embedding E , which can provide us with a representation of an unknown action a 0 .
As genetic perturbations further have no meaningful embeddings as, for example,
CondOT
[Figure]
UMAP
[Figure]
A P
[Figure]
ICNN OT
Identity initialization
UMAP
[Figure]
A P
CPA (Lotfollahi et al., 2021)
P er t u
[Figure]
b a t i o n
K L F+ M A
P 2 k
Data
Predictions molecular ﬁngerprints for drugs, we resort to mode-of-action embeddings introduced in § 3.3.
Assuming marginal sample access to all individual perturbations, we compute a multidimensional scaling (MDS)-based embedding from pairwise Wasserstein distances between individual target populations, such that perturbations with similar effects are closely represented.
For details, see § E.
As current state-of-the-art methods are restricted to modeling perturbations via one-hot encodings, we compare our method to ICNN OT only.
As displayed in Fig.
5a, C OND OT accurately captures the response of unknown actions (BAK1, FOXF1, MAP2K6, MAP4K3), which were not seen during training, at a similar Wasserstein loss as perturbation effects seen during training.
For more details, see § E.
5.3.3 Actions in Combination
While experimental studies can often measure perturbation effects in biological systems in isolation, the combinatorial space of perturbations in composition is too large to capture experimentally.
Often, however, combination therapies are cornerstones of cancer therapy (Mokhtari et al., 2017).
In the following, we test different combinator architectures to predict genetic perturbations in combination from single targets.
Similarly to Lotfollahi et al.
(2021), we can embed combinations by adding individual one-hot encodings of single perturbations (i.e., C ohe + ).
In addition, we parameterize a combinator via a permutation-invariant deep set, as introduced in § 3.3, based on mode-of-action embeddings of individual perturbations (i.e., C moa  ).
We split the dataset into train / test splits of increasing difﬁculty (details on the dataset splits in §D.2).
Initially containing all individual perturbations as well as some combinations, the number of perturbations seen in combination during training decreases over each split.
For more details, see § E. We compare different combinators to ICNN OT
(Fig. 5b) and CPA (Lotfollahi et al., 2021)
(Fig. 5c, Fig.
6b, c).
While the performance drops compared to inference on known perturbations (Fig.
6a) and decreases with increasing difﬁculty of the train / test split, C OND OT outperforms all baselines.
When embedding these high-dimensional populations in a low-dimensional UMAP space (McInnes et al., 2018), one can see that C OND OT captures the entire perturbed population, while
ICNN OT and CPA fail in capturing certain subpopulations in the perturbed state (see Fig.
7 and 11).
Conclusion
We have developed the C OND OT framework that is able to infer OT maps from not only one pair of measures, but many pairs that come labeled with a context value.
To ensure that C OND OT encodes optimal transports, we parameterize it as a PICNN, an input-convex NN that modulates the values of its weights matrices according to a sequence of feature representations of that context vector.
We showcased the generalization abilities of C OND OT in the extremely challenging task of predicting outcomes for unseen combinations of treatments.
These abilities and PICNN more generally hold several promises, both as an augmentation of the OTT toolbox (Cuturi et al., 2022), and for future applications of OT to single-cell genomics.
Acknowledgments and Disclosure of Funding
This publication was supported by the NCCR Catalysis (grant number 180544), a National Centre of
Competence in Research funded by the Swiss National Science Foundation.
We thank Stefan Stark and Gabriele Gut for helpful discussions and the reviewers for their thoughtful comments and efforts towards improving our manuscript.