Supervised Training of Conditional Monge Maps
Optimal transport (OT) theory describes general principles to define and
select, among many possible choices, the most efficient 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 $(\mu,\nu)$, a
parameterized map $T_\theta$ that can efficiently map $\mu$ onto $\nu$. In many
applications, such as predicting cell responses to treatments, the data
measures $\mu,\nu$ (features of untreated/treated cells) that define optimal
transport problems do not arise in isolation but are associated with a context
$c$ (the treatment). To account for and incorporate that context in OT
estimation, we introduce CondOT, an approach to estimate OT maps conditioned on
a context variable, using several pairs of measures $(\mu_i, \nu_i)$ tagged
with a context label $c_i$. Our goal is to % extract from a dataset of labeled
pairs $\{(c_i, (\mu_i, \nu_i))\}$ learn a global map $\mathcal{T}_{\theta}$
which is not only expected to fit em all pairs in the dataset $\{(c_i, (\mu_i,
\nu_i))\}$, i.e., $\mathcal{T}_{\theta}(c_i) \sharp\mu_i \approx \nu_i$, but
should generalize to produce meaningful maps
$\mathcal{T}_{\theta}(c_{\text{new}})$ conditioned on unseen contexts
$c_{\text{new}}$. Our approach harnesses and provides a novel usage for
partially input convex neural networks, for which we introduce a robust and
efficient initialization strategy inspired by Gaussian approximations. We
demonstrate the ability of CondOT 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.