Supervised Training of Conditional Monge Maps
Authors: Charlotte Bunne
Andreas Krause
marco cuturi
Venue: NeurIPS 2022 Accept (2022)
Keywords: Optimal Transport
Single-Cell Dynamics
Convex Neural Networks

Abstract: 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, pairs of input/output data measures $(\mu,\nu)$ that define 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 CondOT, a multi-task approach to estimate a family of OT maps conditioned on a context variable, using several pairs of measures $(\mu_i, \nu_i)$ tagged with a context label $c_i$. CondOT learns a global map $\mathcal{T}_{\theta}$ conditioned on context that is not only expected to fit all labeled 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 also generalize to produce meaningful maps $\mathcal{T}_{\theta}(c_{\text{new}})$ when 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.

Supervised Training of Conditional Monge Maps
Charlotte Bunne" ETH Zurich bunnec@ethz.com ch
Andreas Krause ETH Zurich krausea@ethz.ccm ch
Marco Cuturi Apple cuturi@apple.com com
Abstract
Optimal transport (OT) theory describes general principles to define and select, among many possible choices, the most efficient way to map probability measure onto another. That theory has been mostly used to estimate, given pair of source and target probability measures (J4, v). parameterized map To that can efficiently map la onto In many applications, such as predicting cell responses to treatments, pairs of input/output data measures (μ, v) that define optimal transport problems do not arise in isolation but are associated with context c, as for instance treatment when comparing populations of untreated and treated cells. To account for that context in OT estimation, we introduce CONDOT, multi-task approach to estimate a family of OT maps conditioned on a context variable, using several pairs of measures (μ4, 14) tagged with context label c. CONDOT learns global map Ta conditioned on context that is not only expected to fit all labeled pairs in the dataset {(c, (Mi, v;))], i... To(e;) s Vi, but should also generalize produce meaningful maps (Caww) when conditioned on unseen contexts Cnew- Our approach harnesses and provides novel usage for partially input convex neural net- works, 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.
1 Introduction
A key challenge in the treatment of cancer is to predict the effect of drugs, or 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 (14, v3) over Rdd (cell measurements), tagged with a context e; (encoding the treatment), we seek to learn context-dependent, parameterized transport map To such that, on training data, that map To (e(c)) R - Rd fits the dataset, in the sense that Te Te(c))tp) R Vi- Additionally, we expect that this parameterized map can generalize unseen contexts and patients, predict, given patient's 's cells described (new, the effect of applying context Caew on these cells as (g(Cuew)
Learning Mappings Between Measures From generative adversarial networks, to normalizing flows and diffusion models, the problem of learning maps that move points from source target distribution is central to machine learning. OT theory (Santambrogio, 2015) has emerged as principled approach to carry out that task: For pair of measures μ4, v supported on Idd, OT suggests that, among all maps 7 such that can be reconstructed by applying T to every point in the support
Work done during internship at Apple.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

a.
by scalar
by covariate
c.
by action
e.g,, time-course or dosage levels
e.g. metadata identifiers
e.g.. perturbations decisions
target measures
Pi
CL.2
V
action
source measure
Vt
1
VT
C222'
V2
Voi d5
Vaij
A
Vy
covariate CL.
VI
combination as
learm To
learn To (C(P)))))
Iearm Ta( az) 1z
Figure 1: The evolution from a source μ to target measure v can depend on context variables c of various nature. This comprises scalars such as time or dosage which determine the magnitude of an optimal transport, b. flow of measures into another one based on additional information (possibly different between μ and P) stored in vectors Cr,, or discrete and complex actions a, possibly in combination aej. We seek a unified framework to produce a map Ta(e) from any type of condition c.
of µe (abbreviated with the push-forward notation as Tij - v), one should favor so-called Monge maps, which minimize the average -enrred-lengths of displacements || T(x)]|? ]|2 A rich literature, covered in Peyré 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 neural approaches that parameterize the optimal maps T as neural networks. An early approach is the work on Wasserstein GANs (Arjovsky et al., 201 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 parametric family of context-aware transport maps, using collection of labeled pairs of measures.
Contributions We propose framework that can leverage labeled pairs of measures {(a, (Mi 2/2 v4))) ) to infer a global parameterized map Ta. Hereby, the context e; belongs to an arbitrary set C. We construct To so that it should be able, given possibly unseen context label e C, to output map Tø(c)= : _dd + IR, that is itself the gradient of a convex function. To that end, we propose to learn these parameterized Monge maps Tg as the gradients of partially input convex neural networks (PICNN), which we borrow from the foundational work of Amos et al. (2017). Our framework can be also interpreted as hypernetwork (Ha et al., 2016): The PICNN architecture can be seen as an ICNN whose weights and biases are modulated by the context vector which parameterizes a family of convex potentials in R.d. Because both ICNN and to greater extent PICNN are notoriously difficult to train (Richter-Powell et al., 2021; Korotin et al., 2020, 2021). we use closed-forn 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 affine Monge map obtained in closed form between Gaussian approximations of measures μú; Ms (Gelbrich, 1990). Our framework is applied to three scenarios: Parameterization of transport through 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 oft-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 a predictor for never tested context labels.
2 Background on Neural Solvers for the 2- Wasserstein Problem
Optimal Transport
The Monge problem between two measures 14, E P(R))), here restricted to measures supported on I and compared with the squared Euclidean metric, reads
T == arg
inf
1d ||
(1)
T Thy=

---

The existence of T* is guaranteed under fairly general conditions (Santambrogio, 2015, Theorem 1.22). which require that fl and have finite L2 norm, and that μt puts no mass on (d 1) surfaces of class C2. This can be proved with the celebrated Brenier theorem (1987), which states that there must exist unique (up to the addition of a constant) potential f TR. - IR such that 7 = V This theorem has far-reaching implications: It is sufficient, 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
f := arg sup (f) == [ f'du+ [ fdv. convex
Jga
In practice, Monge maps can be estimated using a dual formulation (Makkuva et al., 2020; Korotin al., 2020; Bunne et al., 2022b; Alvarez-Melis et al., 2021: Mokrov al., 2021). Indeed, T in (I)
is recovered as Vf*" where f is defined in (2), writing f* for the Legendre transform of f.
Convex Neural Architectures Input convex neural networks (ICNN) are neural networks the that admit certain constraints their architecture and parameters 0, such that their output vv(x)) is convex function of their input x (Amos et al., 2017). As a result, they have been increasingly used as drop-in replacements the set of admissible functions in (2). Practically speaking, ICNN is
K- K-layer, fully connected network such that, at each layer index & from 0 to 1, a hidden state vector zę is defined recursively in (3),
-k+1 0he (W Wi bx)
(3)
and (x) = 2x. where, by convention, 20 and W are 0; o are convex non-decreasing activation functions; 0- (be, W k- are the weights and biases of the neural network. While ample flexibility is provided to choose dimensions for intermediate hidden states zg. the last layer must necessarily produce a scalar, hence - and W are line vectors and bk-1 R. ICNNs are characterized by the fact that all weight matrices W associated to latent representations z must have non-negative entries. This, along with the specific activation functions, ensures the convexity of tho- 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 flows 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.
3 Supervised Training of Conditional Monge Maps
We are given a dataset of N pairs of measures, each endowed with a label, (c;, (Hi, 14)) E C CRR())) Our framework builds upon two pillars: (i.) we formulate the hypothesis that an optimal transport TV (or, equivalently, the gradient of a convex potential fī² explains how measure 14 was mapped to Vi. given context ez (ii.) we build on the multi-task hypothesis (Caruana, 1997) that all of the N maps T, between (4; and in share a common set of parameters, that are modulated by context informations
G. These ideas summarized in an abstract regression model described below.
3.1 A Regression Formulation for Conditional OT Estimation
0 E O C R.,"""""""""""" To describes a function that takes an input vector e E C, and outputs a function Tolc): Rd + -Rd ; as a hypernetwork would (Ha et al., 2016). Assume momentarily that we are given ground truth maps T., that describe the effect of context c; on any measure, rather only pairs of measures (Ni, vi). This is of course 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, direct supervised approach to learn a unique 0 could hypothetically involve minimizing a fit function composed of losses between maps
min
[T(c - T=(x))1 dx.( dddx).))).())
(4)
Rd
3

---

Unfortunately, such maps T, are not given, since we are only provided unpaired samples before pa and after V4 that map's upplication. By Brenier's is theorem, we know, however, that such an OT map T 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.I.t. of a function fø( c) E x R that convex w.r.t. z, namely Tø(c) == + V Vife(æ, c); (ii.) estimate 8 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, form the aggregate problem
maX, Ei Eng, ii (fo(., .4)).
(S)
We detail in App. B how the Legendre transforms that appear in the energy terms Epg, are handled with an auxiliary function.
3.2 Integrating Context in Convex Architectures
We propose to incorporate context variables, order to modulate a family of convex functions fo(2, c) using partially input convex neural networks (PICNN). PICNNs are neural networks that can be evaluated over pair of inputs (xr, c), but which are only required to be convex w.r.t. æ. Given an input vector z and context vector a K K-llayer PICNN is defined as ( c) ZK, where, recursively for 0 S k I
1 one has
Uk+11 Tk (Viug Uk),
(6)
2k+1 = 0he (Wr (2x o [W 2u + b11+) + WT (x o (WT" uq + b2)) + W un + bb).
where the PICNN is initialized as ug = c, 20 = 0, o denotes the Hadamard elementwise product, and The is any activation function. The parameters of the PICNN are then given by
θ (Vk. W W W W W ; 0k bū. bg. 6%]-
Similar to ICNNs, the convexity w.r.t. input variable is guaranteed as long as activation functions 0; are convex and non- decresca., and the weight matrices W have non-negative entries. We parameterize this by storing them elementwise applications of softplus operations precursor matrices of the same size, or, alternatively, by regularizing their negative part. Finally, much like ICNNs, all matrices at the K 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 (v, where y is a label vector, typically much smaller than that of vector z. Their PICNN is convex w.r.t. y. in order to easily recover, given datapoint I (c.g., image) the best label y that corresponds to using gradient descent subroutine, i.c. (x) = arg min, PICNN, (z, y). PICNN were therefore originally proposed to learn parameterized, implicit classification 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 In addition that swap, we do not use the convexity of the PICNN to define an implicit layer (or to carry out gradient descent). Indeed, it does not make sense in our setting to minimize (x, e) as function of x, since an observation. Instead, our work rests on the property that V Vi (1, c) describes a parameterized family of OT maps. We note that PICNNs were considered within the context of OT in (Fan al., 2021, Appendix B). In that work, PICNN provide an elegant reformulation for neural Wasserstein barycenters. Fan et al. (2021) considered a context vector e that was restricted to be small vector of probabilities.
3.3 Conditional Monge Map Architecture
Using PICNNs as base module, the CONDOT architecture integrates operations the contexts C. As seen in Figure 1, context values e may take various forms:
1. A scalar denoting strength or temporal effect. For instance, McCann's 's interpolation and its time parameterization, a₁ = ((1 - t)Id + +T)) (McCann, 1997) can be interpreted as a trivial conditional OT model that creates, from OT map the, set of maps parameterized by t, To(t) := x + V. ((1- t) |> + tz))
2. A covariate vector influencing the nature of the effect that led μ4 to vs, (capturing, e.g., patient
feature vectors).
3. One or multiple actions, possibly discrete, representing decisions or perturbations applied onto pa-
To provide a flexible architecture capable of modeling different types of conditions as well as conditions appearing in combinations, the more general CONDOT architecture consists of the

---

hypernetwork To that is fed a context vector through embedding and combinator modules. This generic architecture provides one-size fits all approach to integrate all types of contexts
Embedding Module To give greater flexibility when setting the context variable c, CONDOT contains an embedding module that translates arbitrary contexts into real-valued vectors. Besides simple scalars (Fig. la) for which embedding is required, discrete contexts can be handled with an embedding module Eg. When the set C is small, this can be done effectively using one-hot embeddings Eohe- For more complicated actions such as treatments, there is no simple way to vectorize context c. Similarly to action embeddings in reinforcement learning (Chandak et al., 2019; Tennenholtz and Mannor, 2019), we learn embeddings for discrete actions into learned continuous representation. This often requires domain-knn -nwwledgg on the context values. For molecular drugs, for example, we can learn molecular representations &nal such as chemical, motif-based (Rogers and Hahn, 2010) neural fingerprints (Rong et al., 2020; Schwaller et al., 2022). However, often this domain knowledge available. In this work, we thus construct so-called mode-of-action embeddings, by computing an embedding Emaa that encourages actions a with similar effect on target population v to have similar representation. In 5, we analyze several embedding types for different use-cases.
While we often have access to contexts e in isolation, it is crucial 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 efficacy (Kummar et al., 2010). these settings, the mode of operation between individual contexts e is often not known, and can thus not be directly modeled via simple arithmetic operations such min, max, sum, mean. While we test as a baseline the case, applicable to me-hot-emmeeddings where simple additions used to model these combinations, we propose to augment the CONDOT architecture with a parameterized combinator module Cs. If the order in which the actions are applied irrelevant or unknown, the corresponding network Co needs to be pernutation-invarian which can be achieved by using deep set architecture (Zaheer et al., 2017). Receiving flexible number of inputs from the embedding module Eg. CONDOT allows for joint training of the PICNN parameters 0, embedding parameters @, and combinator parameters & in a single, end-to-end differentiable architecture.
Combinator Module
To
inputs
C+
Exp
combinator
embedding
map
Figure 2: CONDOT Architecture and Modules. The embedding module & embeds arbitrary conditions c, which are then combined via module Cg. Using the processed contexts the map (c() acts on μ to predict the target measure v.
Training Procedure Given a dataset D = fca. (H4, 1/ of N pairs of populations before and after transport v; connected to context c,, we detail in Algorithm provided B, training loop that incorporates all of the architecture proposals described above. The training loss aims at making sure the map (c;) an OT map from 14 to Mi, where e; 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 defined in (Makkuva et al., 2020, Eq. 6) (16)-(16) in place of (2) in our overall loss (5). This involves training the CONDOT architecture using two PICNNs, i.e. PICNN,, and PICNNe, that share the same embedding/combiiatoo module, with a regularization (14) promot- ing that for any e, the PICNNa, (+, c) resembles the Legendre transform of the other, PICNN; (..c).
Initialization Strategies for Neural Convex Architectures
We address the problem of initializing the parameters of (P)ICNNs ensure their gradient evaluated at every point is (initially) meaningful in the context of OT, namely that it is able to map the first and second moments of measure μ into those of target measure v. 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 Ni, NV2 with means respectively my, my and covariance matrices E1, E₂ (where Σ₁ is assumed to be full rank), the
5

---

a.
b.
c.
1
Ag
C1
h Dr
ug softmax
c
V2
1/3
bs,.
h
20
h A2 b2 (I)
o 140
Aa
Figure 3: a. From measure js to several target measures V1, 1/, M3 provided with labels c1, C2, C3 we can extract three Gaussian (quadratic) potentials in closed form, b. whose gradients transport on first approximation j to areas in space that cover the three targets. c. Given new label vector c, we compare it to known labels to modulate the magnitude of each of the three potentials.
Brenier potential solving the OT problem from the first the second Gaussian reads:
fiLN = 1 A Ar
b" x t(A, b) = |1
1 + b" x + t(A,b) b), where,
A = Σ Σ -3 1/2 Σ 1/2
1/2
b b: m A Am1,
define both quadratic and linear terms and t(A,b, b) can be any constant. Importantly, note that we write the quadratic term factorized form A enforce psd-ness, as done by Korotin et al. (2020), not as usually done with single psd matrix (Peyré and Cuturi, 2019, Remark 2.31).
Our quadratic potentials only injected in the first state of hidden vector Zo. populate it with collection of relevant full-rank quadratic convex functions, with the goal of recovering an affine OT map from the start, as illustrated in the experiments from C.I.
Quadratic Potentials Lower Bounded by 0 Naturally, for any choice of L(A, b) one recovers the property that Sij 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Ñr.NG is O. This can be obtained by setting
N2. When used in deep architectures, the level of that constant does,
6(A,b) 6T (AT A)
(8)
which results in the following choice, writing m (A A) mm2.
fis (x) L (r (A* b) 113 = 11 (x -w) w)
(9)
To mimic these potential functions, we introduce a quadratic layer parameterized by a weight matrix M and a "bias" vector m, defined as q (r)) = [x m)11 By design, (M)(x) is a convex quadratic, non-negative layer. Finally, one has the following relationships.
Vgr.o, = Id, NN = N22
(10)
ICNN Initialization We explore two possible ICNN (3) initializers for OT.
Identity Initialization The first approach ensures that upon initialization the ICNN's gradient mimics the identity map, i.e., (z) = for any x. We do so by injecting in the initial hidden state zo the norm of the input vector |12 cast trainable layer qM initialized with M / and m Og, see (10). The remaining parameters are chosen to propagate that norm throughout layers using averages. This amounts the following choices:
1. Set all o, to be activations such that (u( x= for u large enough, e.g., (leaky) ReLU or softplus.
2. Introduce initialization layer, 2n 9M. (æ:),, itself initialized with M = I and = O₂-
3. Initialize all matrices W to Rs 1d /d1, where d1, d2 are the dimensions of these matrices.
4. Initialize all matrices T to Rs O.
5. Initialize biases b; to s1, where is a large enough value s so that d((s) R
Gaussian Initialization The second approach can be used to initialize an ICNN so that its gradient mimics the affine transport between the Gaussian approximations of μ and v. To this end, we follow all of the steps outlined above, except for step 2 where the quadratic layer q. is initialized instead with M A and = m, (AT -'m using notations in (7), (8), (9), where mi, m2, Et, En

---

Table Evaluation of drug effect predictions from control cells to cells treated with drug Givinostat when conditioning various covariates influencing cellular responses such as drug dosage and cell
type. Results are reported based on MMD and the le distance between perturbation signatures of marker genes in the 1000 dimensional gene expression space.
Method
Conditioned on Drug Dosage
Conditioned on Cell Line In-Sample
In-Sample
Out-of-Sammle
MMD
Ex(PS)
MMD
Ex(PS)
MMD
(a(PS)
CPA (Letfollahi et al., 2021)
0.1502 0.0769
2.47 4: 2.88
0.1568 1: 0.0729
2.65 2.75
0.2551 1: 0.006
2.71 1.51
ICNN OT (Makkuva al. 2020)
0.0365 0.0473
2.37 4: 2.15
0.0466 1: 0.0479
2.24 2.39
0.0206
0.0109
1.16 0.75
CONDOT (ddentity initialization)
0.0111
0.0055
0.63 + 0.09
0.0374 + 0.0052
2.02 + 0.10
0.0188
0.0078
0.39 0.06
CONDOT (Gaussia inidialization)
0.0128 0.0081
0.0062
1.84 0.14
0.0074 0.0146
0.41 0.07
ConKIT OT
CPA
(Lotfollshi al., 2021)
Wasserstin Distance
20
INSC20000165022.1
A549
MCF7
KS62
18-
0.04
CondOT
Target
16-
Source
34-
Demity
0.4
CPA of etal, 2021)
0.3-
3
2-
ICNN
0.2-
10-
0.1-
0.0 0.5 Expression Level
In-Sample
Out-of Sample
Setting
Expression Level
Expression Level
Figure 4: a. Predictive performance of CONDOT and baselines w.I.t the eetropy-regularized Wasser- stein distance on drug dosages in-sample, i.e., seen during training, and sut-le- i.e., unseen during training. b. Marginal distributions of observed source and target distributions, as well pre- dictions on perturbed distributions by CONDOT and baselines of an exemplary gene across different cell lines. Predicted marginals of each method should match the marginal of the target population.
replaced by the empirical mean and covariances of y and Throughout the experiments, the Gaussian and identity initialization. Further comparisons between the vanilla initialization and those introduced in this work can be found in C.1 (Fig. 8).
PICNN Initialization Recall for convenience that a K-layeer PICNN architecture reads:
Uk+1 = Th (Viu + 24)
1 = ak (Wr (zs [W un + b + W (x (W ux + bz) + W + (%)
(2, c) = =バ- In their original form (Amos et al., 2017, Eq. 3), PICNNs are initialized by setting uo = c and 20 0 to zero vector of suitable size. Intuitively, the hidden states uy act as context-dependent modulators, whereas vectors propagate, layer after layer, a collection of convex functions in x that are iteratively refined, while retaining the property that they are each convex in x. A reasonable initialization for PICNN that is provided a context vector e is that if ~ Cj (where j is in the training set), one has that V V (- c) maps approximately μ4 to Vj, which can be obtained by having ve, (,, c) mimic the closed-for form Brenier potential between the Gaussian approximations of μj, V/j. Alternatively, one may also default to an identity initialization as discusses above. To obtain either behavior, we make the following modifications, and refer to the illustration in Fig. 3:
1. The modulator (c)(c) softmax(cT M), where C is initialized as M = [cyر,- j. and V. = 82 0. 2. 20 = Iq,,.m, (x)l;. where weight matrices and bias (M,, m2) are either initialized to (1, 0) or as (A,, w3) recovered by solving the Gaussian affine map from μ4y to V using (9).
3. Modulator un is passed directly to hidden state upon first iteration W I, bí 0.
4. All subsequent matrices W are initialized to ze dd /d.. where dj d2 are their dimensions,
O, the biases by 1, by Rs O. 5. WE and W
Evaluation
Biological cells undergo changes their molecular profiles upon chemical, genetic, or mechanical perturbations. These changes can be measured using recent technological advancements in high-
7

---

MoA Embedding
b.
One-Hot Embedding
c.
One-Hot Emhedding
80
Wasserstin Loss ICNN OT
1000
70
Waseentein Loees of CIA
50-
40
50-
60
30
Wasserstein
40-
20
Knoown Nemarbattene
Krows Perterhatons
10-
Keiowaa Pottalasions
Uillasan Parzalations
Uiliooer Conbiaaaions
Colboowe Corniinatins
20
20 20
20
30
60
00
40
90
10
Wasserstein Loss of CondOT
Wasterstein
CondOT
Wassenstein I of CondOT
Figure 5: Comparison between a. CONDOT and ICNN OT (Makkuva et al., 2020) based on embedding Emoa b. as well as Eohe- and CONDOT and CPA (Lotfollahi et al., 2021) based on embedding Eohe on known and wnknown perturbations combinations. Results above the diagonal suggest higher predictive performance of CONDOT.
resolution multivariate single-cell biology. Measuring single cells in their unperturbed or perturbed state requires. however, to destroy them, resulting in populations and y that are unpaired. The relevance of OT to that comes from its ability resolve such ambiguities through OT maps, holding promises of a better understanding of health and disease. We consider various nigh-dimensional prob- lems arising from this scenario evaluate the performance of CONDOT ($ 3) versus other baselines.
5.1 Population Dynamics Conditioned Scalars
Upon application of molecular drug, the state of each cell 2; of the unperturbed population altered, and observed in population Molecular drugs are often applied different dosage levels 1, and the magnitude of changes in the gene expression profiles of single cells highly correlates with that dosage. We seek to learn global, parameterized transport map To 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-inflammatory, anti-angiogenic, and antincoplastic activities (Srivatsan et al., 2020), applied at different dosage levels, i.e., E {10 100 nM, 1, 000 nM, 10, 000 nM.. The dataset contains 3, 541 cells described with the gene expression levels of 1, 000 highly-variale how well CONDOT captures the drug effects at different dosage levels via distributional distances such as MMD (Gretton et al., 2012) and the f--norm between the corresponding perturbation signatures (PS), as well as the the metrics 50 marker genes, i.e., genes mostly affected upon perturbation. For more details on evaluation metrics, see $ E.2. To put CONDOT'S performance into perspective, we compare to current state-of-the-ard baselines (Lotfollahi et al., 2021) well parameterized Monge maps without context variables (Bunne et al., 2021; Makkuva et al., 2020, ICNN OT), see $ E. As visible in Table and Fig. 4a, CONDOT achieves consistently more accurate predictions 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 subset of dosages and we test CONDOT on out-oe-sample dosages. Table and Fig. 4a demonstrate that CONDOT 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 CONDOT, 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 CONDOT is able to reach (see Table 2).
measure genes. In first experiment,
Wasserstein distance (Cuturi, 2013). We compute
5.2 Population Dynamics Conditioned Covariates
Molecular processes are often highly dependent on additional covariates that steer experimental conditions, and which are not present the features measures in population μ or v. This can be, for instance, factors such as different cell types clustered within the populations. When the model can only

---

a.
Known Perturbations
b.
Unknown Combinations
c.
Unknown Combinations
27.5
Maximum Main Dicetepaccy
0.4
25.0
Wastenstiin Distancs
(Lorfollahi et al., 2021)
22.5
CondOT
20.0
(Combinator C+
17.5
0.2
CondOT
15.0
(Combinator Co
011
12.5
10.0
Difficulty Train / Test Split
Difficulty of Train Test Split
Difficulty Train 'est Split
Figure 6: Predictive performance for known perturbations, b. unknown perturbations in combina- tion w.r.t. regularized Wasserstein distance and e. MMD over different train / test splits of increasing difficulty for baseline CPA as well as CONDOT with different combinators Cahe and Ca For more details on the dataset splits, see $D.2.
be conditioned w.r.t. small and fixed set of metadata information, such cell types, it is sufficient to encode these contexts using a one-ho hot embedding module E To illustrate this problem, we consider cell populations comprising three different cell lines (A549, MCF7, and K562). As visible in Table 1, CONDOT 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. CONDOT captures the distribution shift from control to target populations consistently across different cell lines.
5.3 Population Dynamics Conditioned Actions
To recommend personalized medical procedures for patients, or to improve our understanding of genetic circuits, it is key be able predict the outcomes of novel perturbations, arising from combinations of drugs or of genetic perturbations. Rather than learning individual maps T predicting the effect of individual treatments, we aim learning global map which, given as input the unperturbed population μ as well as the action of interest, predicts the cell state perturbed by a. Thanks its modularity, CONDOT can not only learn a map To for all actions known during training, but also to generalize 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 CONDOT's ability to accurately predict phenotypes of genetic perturba- tions based on single-cell RNA-sequeneing pooled CRISPR screens (Norman et al., 2019; Dixit et al., 2016), comprising 98, 419 single-cell gene expression profiles with 92 different genetic perturbations, each cell measured via a 1, 500 highly-variable gene expression vector, As, in first step, we do not aim at generalizing beyond perturbations encountered during training, we utilize again a one-hot em- bedding Eohe condition To on each perturbation a. We compare our method other baselines capa- ble of modeling effects of large set of perturbations such CPA (Lotfollahi et al., 2021). Often, the effect genetic perturbations are subtle the high-dimensional gene expression profile of single cells. Using SANNNparameterized OT maps without context information, we can thus assess the gain in accu- maatext-wwwreness over simply pre-
racy of predicting the perturbed target population by incorporating dicting average perturbation effect. Figure 5a and b demonstrate that compared to OT ablation stud- ies, Fig. 5c and Fig. 6a for the current state-of-the-ar method CPA (Lotfollahi et al., 2021). Compared to both, CONDOT 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 per- turbations. This requires us to use an embedding ε, which can provide us with a representation an unknown action a'. As genetic perturbations further have no meaningful embeddings as, for example,

---

CondOT
ICNN OT Identity initialization
CPA (Lotfollahi al,, 2021)
Identity initialization
Porturbation KI.FI-MAPIKC
Data
Predictions
UMAP1
UMAPP
UMAP
Figure 7: UMAP embeddings of cells perturbed by the combination KLFI+MAP226 (gray) and predictions CONDOT (ours), ICNN OT (Makkuva et al., 2020), and CPA (blue). While CONDOT aligns well with observed perturbed cells, the baselines fail capture subpopulations.
molecular fingerprints for drugs, we resort to mode-of-action embeddings introduced in $ 3.3. Assum ing marginal sample access to all individual perturbations, compute multidimensional scaling (MDS)-based embedding from pairwise Wasserstein distances between individual target populations, such that perturbations with similar effects 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 ICNN OT only. As displayed in Fig. 5a, CONDOT accurately captures the response of unknown actions (BAKI, FOXFI, MAP2K6, MAP4K3), which were not seen during training, at similar Wasserstein loss perturbation effects seen during training. For more details, see E.
5.3.3 Actions Combination
While experimental studies often measure perturbation effects in biological systems in isolation, the combinatorial space of perturbations composition is too large to capture experimentally. Often, however, combination therapies comerstones of cancer therapy (Mokhtari al., 2017). In the fol- lowing, we test different combinator architectures to predict genetic perturbations in combination from single targets. Similarly to Lotfollahi al. (2021),
can emhed combinations by adding individual one-hot encodings of single perturbations (i.e., ) addition, we parameterize a combinator via a commutation-invarian deep set, introduced in $ 3.3, based mode-of-action embeddings of individ- ual perturbations (i.e., C We split the dataset into train / test splits of increasing difficulty (details on the dataset splits in $D.2). Initially containing all individual perturbations as well some combina- tions, 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 pertur- bations (Fig. 6a) and decreases with increasing difficulty of the train test split, CONDOT outperforms all baselines. When embedding these high-dimensional populations in low-dimensional UMAP space (MeInnes et al., 2018), one can see that CONDOT captures the entire perturbed population, while ICNN OT and CPA fail in capturing certain subpopulations in the perturbed state (see Fig. 7 and 11).
6 Conclusion
We have developed the CONDOT framework that able infer OT maps from not only one pair of measures, but many pairs that come labeled with context value. To ensure that CONDOT encodes optimal transports, we parameterize it as PICNN, an input-convex NN that modulates the values of its weights matrices according to sequence of feature representations of that context vector. We showcased the generalization abilities of CONDOT 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 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 manuscript.
10

---

References
D. Alvarez-Melis, Y. Schiff, and Y. Mroueh. Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks. arXiv Preprint arXiv:2106.00774 2021.
B. Amos, L., Xu, and Z. Kolter. Input Convex Neural Networks. In International Conference on Machine Learning (ICML), volume 34, 2017.
M. Arjovsky, S. Chintala, and I.. Bottou. Wasserstein Generative Adversarial Networks. In Interna- tional Conference on Machine Learning (ICML). PMLR, 2017.
Y. Brenier. Décomposition polaire et réarrangement monotone des champs de vecteurs. CR Acad. Sci. Paris Sér. Math., 305, 1987.
C. Bunne, S. G. Stark, G. Gut, J. S. del Castillo, K.-V. Lehmann, L. Pelkmans, A. Krause, and G. Ratsch. Learning Single-Cel Perturbation Responses using Neural Optimal Transport. bioRtiv, 2021.
C. Bunne, Y.-P. Hsieh, M. Cuturi, and A. Krause. Recovering Stochastic Dynamies via Gaussian Schrödinger Bridges, arXiv Preprint arXiv:2202.05722, 2022a.
C. Bunne, L. Mang-Papas A. Krause, and M. Cuturi. Proximal Optimal Transport Modeling of Population Dynamics. In International Conference on Artificial Intelligence and Statistics (AISTATS), volume 25, 2022b.
R. Caruana. Multitask Learning. Machine Learning, 28(1), 1997.
Y. Chandak, G. Theocharous, J. Kostas, S. Jordan, and P. Thomas. Learning Action Representations for Reinforcement Learning. In International Conference on Machine Learning (ICML), 2019.
Y. Chen, Y. Shi, and B. Zhang. Optimal Control Via Neural Networks: A Convex Approach. In International Conference Learning Representations (ICLR), 2019.
S. Chopra, R. Hadsell, and Y. LeCun. Learning similarity metric discriminatively, with application to face verification. In JEEE Conference on Computer Vision and Pattern Recognition (CVPR), volume 1. IEEE, 2005.
M. Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. In Advances in Neural Information Processing Systems (NeurIPS), volume 26, 2013.
M. Cuturi and G. Peyré. Semidual Regularized Optimal Transport. SIAM Review, 60(4):941-965. 2018.
M. Cuturi, L. Tools (OTT); JAX Toolbox for all things Wasserstein. arXiv Preprint arXiv:2201. 12324, 2022.
Y. Tian, C. Bunne, G. Davis, and O. Teboul. Optimal Transport
J. De Leeuw and P. Mair. Multidimensional Scaling Using Majorization: SMACOF in R. Journal of Statistical Software, 31, 2009.
A. Dixit, O. Parnas, B. Li, J. Chen, C. P. Fulco, L. Jerby-Amon, N. D. Marjanovic, D. Dionne, T. Burks, R. Raychowdhury, et al. Perturb-Seq: Dissecting Molecular Circuits with Scalable Single-Cell RNA Profiling of Pooled Genetic Screens. Cell, 167(7), 2016.
J. Fan, A. Taghvaei, and Y. Chen. Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks, In International Conference on Machine Learning (ICML), 2021.
M. Gelbrich. On a Formula for the I Wasserstein Metric between Measures on Euclidean and Hilbert Spaces. Mathematische Nachrichten, 147(1), 1990.
A. Genevay, L.. Chizat, Bach, M. Cuturi, and G. Peyré. Sample Complexity of Sinkhorn Divergences. In International Conference on Artificial Intelligence and Statistics (AISTATS), volume 22, 2019.
A. Gretton, K. M. Borgwardt, M. Rasch, B. Schölkopf, and A. Smola. Kernel Two-Sample Test. Journal of Machine Learning Research, 13, 2012.

---

D. Ha, A. Dai, and Q. Le. Hypemetworks. arXiv preprint arXiv: 09000, 2016.
T. Hashimoto, D. Gifford, and T. Jaakkola. Learning Population-Level Diffusions with Generative Recurrent Networks. In International Conference on Machine Learning (ICML), volume 33, 2016.
C.-W. Huang, R. T. Q. Chen, C. Tsirigotis, and A. Courville. Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization. In International Conference on Learning Representations (ICLR), 2021.
J.-C. Hütter and P. Rigollet. Minimax estimation of smooth optimal transport maps. The Annals of Statistics, 49(2), 2021.
L. Jacob, J. She, A. Almahairi, Rajeswar, and A. Courville. W2GAN: Recovering an Optimal Transport Map with GAN. In arXiv Preprint, 2018.
H. Janati, T. Bazeille, B. Thirion, M. Cuturi, and A. Gramfort. Multi-subject MEG/EEG source imaging with sparse multi-task regression. Neurolmage, 220, 2020.
L. Kantorovich. On the transfer of masses (in Russian). Doklady Akademii Nauk, volume 37, 1942.
D. Kingma and Ba. Adam: A Method for Stochastic Optimization. In International Conference on Learning Representations (ICLR), 2014.
A. Korotin, V. Egiazarian, A. Asadulaev, A. Safin, and E. Burnaev. Wasserstein-2 Generative Networks. In International Conference on Learning Representations (ICLR), 2020.
A. Korotin, L. Li, A. Genevay, J. M. Solomon, A. Filippov, and E. Burnaev. Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark. Advances in Neural Information Processing Systems (NeurIPS), 34, 2021.
S. Kummar, H. X. Chen, J. Wright, S. Holbeck, M. D. Millin, Tomaszewski, J. Zweibel, J. Collins, and J. H. Doroshow. Utilizing targeted cancer therapeutic agents in combination: novel approaches and urgent requirements. Nature Reviews Drug discovery, 9(11), 2010.
M. Lotfollahi, F. A. Wolf, and F. J. Theis. scGen predicts single-cell perturbation responses. Nature Methods, 16(8), 2019.
M. Lotfollahi, M. Naghipourfar, F. Theis, and F. A. Wolf. Conditional out-of-distribution generation for unpaired data using transfer VAE. Bioinformatics, 36, 2020.
M. Lotfollahi, A. K. Susmelj. C. De Donno, Y. Ji, I. L. Ibarra, F. A. Wolf, N. Yakubova, F. J. Theis, and D. Lopez-Paz. Compositional perturbation autoencoder for single-cell response modeling. bioRxiv, 2021.
R. K. Mahabadi, S. Ruder, M. Dehghani, and J. Henderson. Parameter-efficient Multi-task Fine- tuning for Transformers via Shared Hypernetwors.. In Proceedings of the Annual Meeting of the Association for Computational Linguistics (ACL), 2021.
A. Makkuva, A. Taghvaci, S. Oh, and J. Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning (ICML), volume 37, 2020.
R. J. McCann. A Convexity Principle for Interacting Gases. Advances in Mathematics, 128(1), 1997.
L. MeInnes, J. Healy, and J. Melville. UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction. arXiv Preprint, 2018
A. Mead. Review of the Development of Multidimenssonal Scaling Methods. Journal of the Royal Statistical Society: Series (The Statistician), 41(1), 1992.
T. Mikolov, K. Chen, G. Corrado, and J. Dean. Efficient Estimation of Word Representations in Vector Space. International Conference on Learning Representations (ICLR), Workshop Track, 2013a.
12

---

T. Mikolov, I. Sutskever, K. Chen, G. S. Corrado, and J. Dean. Distributed Representations of Words and Phrases and their Compositionality. Advances Neural Information Processing Systems (NeurIPS), 26, 2013b.
Mikolov, W.-t. Yih, and G. Zweig. Linguistic Regularities in Continuous Space Word Representa- tions. In Annual Conference of the North American Chapter of the Association for Computational Linguistics (NAACL-HLT)) 2013c.
R. B. Mokhtari, T. S. Homayouni, N. Baluch, E. Morgatskaya, S. Kumar, B. Das, and H. Yeger. Combination therapy in combating cancer. Oncotarget, 8(23):38022, 2017.
P. Mokrov, A. Korotin, L. Li, A. Genevay, J. Solomon, and E. Burnaev. Large-Scale Wasserstein Gradient Flows. Advances in Neural Information Processing Systems (NeurIPS), 2021.
G. Monge. Mémoire sur la théorie des déblais et des remblais. Histoire de l'Académie Royale des Sciences, pages 666-704, 1781.
T. M. Norman, M. A. Horlbeck, M. Replogle, A. Y. Ge, A. Xu, M. Jost, L. A. Gilbert, and J. S. Weissman. Exploring genetic interaction manifolds constructed from rich single-cell phenotypes. Science, 365(6455), 2019.
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pretten- hofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine Learning in Python. Journal of Machine Learning Research, 12, 2011.
G. Peyré and M. Cuturi. Computational Optimal Transport. Foundations and Trends in Machine Learning, 1 11(5-6), 2019, ISSN 1935-8245.
A.-A. Pooladian and J. Niles-Weed. Entropic estimation of optimal transport maps. arXiv preprint arXiv:2109. 12004, 2021.
N. Prasad, K. Yang, and C. Uhler. Optimal Transport using GANs for Lineage Tracing, arXiv preprint arXiv:2007. /2098, 2020.
A. Rambaldi, C. M. Dellacasa, G. Finazzi, A. Carobbio, M. L. Ferrari, P. Guglielmelli, E. Gattoni, S. Salmoiraghi, M. C. Finazzi, S. Di Tollo, et al. pilot study of the Histone- Deacetylae inhibitor Givinostat in patients with JAK2V617F positive chronic myeloproliferative neoplasms. British journal of haematology, 150(4):446 455, 2010.
D. Rezende and S. Mohamed. Variational Inference with Normalizing Flows. In International Conference on Machine Learning (ICML), 2015.
J. Richter-Powell J. Lorraine, and B. Amos. Input Convex Gradient Networks. arXiv preprint arXiv:2111.12187, 2021.
Rigollet and A. J. Stromme. On the sample complexity of entropic optimal transport, arXiv preprint arXiv:2206. 13472, 2022.
D. Rogers and M. Hahn. Extended-Coonectio Fingerprints. Journal of Chemical Information and Modeling, 50(5), 2010.
Y. Rong, Y. Bian, T. Xu, W. Xie, Y. Wei, W. Huang, and Huang. Self-Supervised Graph Transformer on Large-Scale Molecular Data. Advances in Neural Information Processing Systems (NeurIPS), 2020.
F. Santambrogio. Optimal Transport for Applied Mathematicians. Birkhäuser. NY, 55(58-63):94. 2015.
G. Schiebinger, J. Shu, M. Tabaka, B. Cleary, V. Subramanian, A. Solomon, J. Gould, S. Liu, S. Lin, Berube, et Transpo-t Analysis of Single-Cell Gene Expression Identifies Developmental Trajectories in Reprogramming. Cell, 176(4). 2019.
13

---

M. A. Schmitz. M. Heitz, N. Bonneel, F. Ngole, D. Coeurjolly, M. Cuturi, G. Peyré, and J.-L. Starek. Wasserstein dictionary learning: Optimal transport- -based unsupervised nonlinear dictionary learning. SIAM Journal on Imaging Sciences, 11(1):643-60 -788 2018.
P. Schwaller, A. C. Vaucher, R. Laplaza, C. Bunne, A. Krause, C. Corminboeuf, and T. Laino. Machine intelligence for chemical reaction space. Wiley Interdisciplinary Reviews: Computational Molecular Science, page e1604, 2022.
S. R. Srivatsan, J. L. McFaline-Figuero.i V. Ramani, L. Saunders, Cao, J. Packer, H. A. Pliner, D. L. Jackson, R. M. Daza, L. Christiansen, et al. Massively multiplex chemical transcriptomics at single-cell resolution. Science, 367(6473), 2020.
V. Stathias, M. Jermakowiez, M. E. Maloof, M. Forlin, W. Walters, R. K. Suter, M. A. Durante,
S. L. Williams,, J. W. Harbour, C.-H. Volmar, et al. Drug and disease signature integration identifies synergistic combinations in glioblastoma. Nature Communications, 91), 2018.
G. Tennenholtz and S. Mannor. The Natural Language of Actions. In International Conference on Machine Learning (ICML), 2019.
C. Villani. Topics in Optimal Transportation, volume 58. American Mathematical Soc., 2003.
A. Wolf, P. Angerer, and F. J. Theis. SCANPY: large-scale single-cell gene expression data analysis. Genome biology, 19(1), 2018.
K. D. Yang and C. Uhler. Scalable Unbalanced Optimal Transport using Generative Adversarial Networks. International Conference on Learning Representations (ICLR), 2019.
K. D. Yang, K. Damodaran, S. Venkatachalapathy, A. C. Soylemezoglu, G. Shivashankar, and C. Uhler. Predicting cell lineages using autoencoders and optimal transport. PLoS Computational Biology, 16(4), 2020.
M. Zaheer, Kottur, S. Ravanbakhsh, B. Poczos, R. R. Salakhutdinov, and A. J. Smola. Deep Sets.
In Advances in Neural Information Processing Systems (NeurIPS), volume 30, 2017.
14


Reviews:
[{"review_id": "raUpwwZIUJW", "reviewer": "Reviewer_P7bD", "summary": "This work concerns the problem of learning the optimal transport (Monge) map that push-forwards between a given pair of measures coupled with a given context. They propose a framework that composes of an input convex neural network and a context embedding, as well as two ways to effectively initialize the network. Empirically, they show that the proposed model can predict and generalize well in the application of population dynamics prediction coupled across various types of contexts.", "questions": "- In the experiments, the comparison with ICNN OT seems to be unfair as the whole context information is discarded. What happens if you feed the embedded context directly to the ICNN (not PICNN) or maybe try some trivial ways to incorporate the context (e.g., building a network that outputs $h(f(x), g(z))$ where $f$ is an ICNN and $g$ is a neural net and $h$ is some linear combination) - will it still perform poorly?\n\n- Should there be some out-sample results in the setting of covariate conditioning in the experiments?\n- Why are the UMAP embeddings of data in Figure 7 different in three plots? Should they be the same given the same data?\n- Figure 7 appears in the main text without being referenced.", "limitations": "The discussion on limitations is inadequate to some extent (see comments above). It seems to be no potential negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "- The paper is well-written and well-presented. \n- The given answer to the question \"how to make use of additional contexts in the traditional framework of ICNNs for learning OT Monge maps\" is straightforward: feeding embedded contexts as an additional input to the neural network. The challenge appears to be the design of a network that is only convex in the main input (support of the source measure) and not in the context input (supposedly for more expressiveness - but this is not discussed in the paper). The proposed design (called PICCN) of such a network is novel.\n- Another difficulty is said to lie in the training of such an architecture, for which the authors propose two straightforward approaches for good initialization, but there is no experiment showing the effectiveness of both methods compared to just initializing randomly. Moreover, there is also no discussion/recommendation on which initialization to choose in practice. From the experiments, it seems that the Gaussian initialization is slightly better than the identity one, but in my opinion, it is due to the former carrying more data information through the Gaussian approximation). In addition, I believe that there is likely not much difference between these two when the source distribution is complex).\n- The generalization claim of the proposed method is not strongly supported in some experiments. For example, in Table 1, the in-sample errors and out-sample errors are consistent for all considered methods, so it cannot be said that CondOT generalizes better.\n- There is a wide range of applications for the proposed method in modelling population dynamics given contexts, which is highly practical. Moreover, the results seem significant (as there is a substantial improvement over the SOTA).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "r3IjlojYh_", "reviewer": "Reviewer_sa6w", "summary": "The paper considers the problem of learning a map $\\mathcal T : \\mathcal C \\times \\mathbb R^d \\to \\mathbb R^d$ such that given observations $(c_i, (\\mu_i, \\nu_i)) \\in \\mathcal C \\times \\mathcal P(\\mathbb R^d)^2$ it satisfies ${\\mathcal{T}(c_i)}$\\#$\\mu_i \\approx \\nu_i$ for each $ i $. The primary motivation is the following: suppose we observe unpaired biological data before and after treatment that can be represented as point clouds $\\mu_i,\\nu_i$ respectively. The 'context' or metadata of the treatment (e.g. dosage, type of drug etc) is encapsulated by $c_i$. The key aspects of the authors' approach to the problem is as follows. \n1) Use optimal transport as a judicious way of selecting the push forward map.\n2) Parametrize $\\mathcal T=\\mathcal T_\\theta$ by a partially input convex neural network (nonconvex in the 'context' variable), motivated by the fact that the optimal transport map is the gradient of the unique convex function that trasports the source to the target.\n3) Transform the context variables to vectors in Euclidean space by a generic 'embedding module' which are combined by the generic 'combinator' module (that can take an arbitrary number of contexts as input which allows for multiple treatments, the most challenging/interesting application). \n4) Train the entire model by end-to-end gradient descent on a dual formulation of the OT problem over the parameters of the input convex neural networks, the embedding module and the combinator module. \n\nThe main contributions of the paper are: \n1) Describing a state of the art framework that allows for unseen/composite contexts and utilizes the 'continuity' of the context space. \n2) Describing an initialization procedure for the partially input convex neural network (with the aim of making training easier) such that the initial map is either (i) the identity (ii) the optimal transport between the Gaussian approximations.", "questions": "### Questions/Comments:\n1) Personally I would prefer that the main text include the algorithm (Algorithm 1) and the training objective (display (12-13)) for easier readibility. A short derivation of the training objective could also be added. Although under the space-constraints this may not be possible. \n2) In Section 2) clarify the relevant regularity requirements for the OT map to exist. \n3) What is _+ in display (4) and bottom of page 6?\n4) Line 129: The message here is not completely clear. The McCann interpolation parametrizes a geodesic, but in practice we have no reason to expect the curve $(t, (\\mu_t, \\nu_t))_{t\\in[0,1]}$ to form a geodesic (if, say, $t$ is the dosage)? \n5) Why not always take train_freq_f=2? (as in why is it beneficial to train $g$ more than $f$)\n6) What motivates perturbation signature metric? It is much cruder than Wasserstein distance or MMD - is it commonly used in the literature? The explanation in the paragraph starting on line 672 could be simplified to simply the $\\ell_2$ distance between means of prediction and truth. \n\n\n\n\n### Typos/minor suggestions:\n- Line 3: repeats the word optimization\n- Line 18: from of \n- Display (2) why introduce f instead of just use \\psi (also display (5))\n- Figure 6b perturbation vs combination unclear\n- Line 73: unique convex potential\n- Line 91: …we encode this by …\n- Paragraph under line 106: …leveraging the context values as much as possible …\n- same paragraph: …propose to learn a unique …. Unique in what sense?\n- Line 108: …not in the context value…\n- Line 119: replace ‘condition’ by ‘context’?\n- Line 123: …and the weight matrices…\n- Line 155: allows vs encourages\n- Line 320: …that is able to…\n- Line 538: ...in that …\n- Line 548: …learn our model…\n- Line 552: definition of \\tilde\\Psi seems wrong, shouldn’t be a product\n- Line 553: $\\int |f| d\\mu < \\infty$\n- Line 572: parentheses\n- Line 747: multiple typos", "limitations": "The limitations and potential negative societal impact have been adequately addressed. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1) The problem is well-motivated, interesting and important.\n2) The paper is clearly written and easy to follow. \n2) The empirical results suggest that the outlined framework is superior to prior work on common tasks.\n3) The outlined framework allows to generalize to unseens contexts and combine them, which wasn't possible with prior work. \n\nWeaknesses:\n1) A lot of the components of the framework have appeared in related literature in some form (partially input convex neural networks specifically for training OT maps with the same dual objective, permutation invariant network for combinator, encoding of contexts etc.). The main methodological contribution seems to be that (i) they don't train separately for each context, instead parametrize wrt the context continuously and thus leverage the intuition that OT maps for similar contexts should be similar and (ii) the initialization of the PICNNs. \n\n2) One of the main contributions (highlighted in the paragraph on line 49) is the initialization of the PICNN. However, from the text it is unclear whether the initialization improves test performance or training speed. In Table 1 the difference in test performance between identity and Gaussian initialization doesn't seem statistically significant, and later in Figure 9 only the Gaussian model is shown. Intuitively one might expect the Gaussian initialization to perform better than identity, which in turn should be better than a randomly initialized networks, commenting on this would make this aspect of the paper more convincing. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "lcvbAStM97m", "reviewer": "Reviewer_CtVg", "summary": "The paper proposes a generalization of Monge maps to the case where multiple pairs of labeled measures $\\\\{(c_i,(\\mu_i,\\nu_i))\\\\}$ are to be matched by a single, global map. The desired map has the context label $c_i$ as an input, and for each $c_i$ it serves as the OT map from $\\mu_i$ to $\\nu_i$. The paper starts by setting preliminaries and describing the objective, then provides a detailed account of implementation using partially input convex neural nets, and concludes with several experiments on treatment effects of drug combinations. ", "questions": "Major questions/concerns appear under `Weaknesses' above. Some smaller issues and specific questions are listed next:\n\n1. **Eq. (2):** Should both $f^\\star_{\\mu,\\nu}$ and $\\psi^\\star_{\\mu,\\nu}$ appear there? The authors seem to interchange between $f$ and $\\psi$ for OT potentials throughout the paper. I suggest picking one symbol and sticking with it. \n\n2. **Section 3.2, Item 1:** In the last sentence, $\\psi_\\theta$ is referred to as an 'OT map' but it should be 'OT potential'.\n\n3. **Section 4, 1st paragraph:** The authors say 'it is able to approximately map a measure $\\mu$...'. In what sense is the approximation here? Can this be quantified?\n\n4. **Eq. (5):** Should write $f^\\star_{\\mathcal{N}_1,\\mathcal{N}_2}(x)$ on the left-hand side, i.e., add the function argument to the notation.\n\n5. **Fig. 6:** Why is there no comparison to ICNN OT [Makkuva et al., 2020] for that experiment? \n\n6. **General suggestion:** Having the figures appear on the same page where they are discussed (or at most one page before) would improve presentation. Currently, figures are spread out widely across the text which makes reading harder. For instance, Fig. 4 appears on page 6 but is discussed the first time on p. 8. Jumping back and forth is quite inconvenient. I suggest clustering figures together, if necessary, and bringing them closer to the relevant text. ", "limitations": "Coupling the results of the paper with some theory, as suggested above, would provide a principled account of limitations. Currently, the method seems to work well but there are no assurances on how far the framework can be pushed or when it might fail. The authors are clear about what is being proposed and are honest in their presentation of the results, but the limitations of the framework remain an item for further exploration.\n\nI have no concerns about potential negative societal impact.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\n* The paper is very well-written: clear, streamlined, and concise. The discussions and descriptions of ideas are detailed enough for the reader to get all the information they need, but at the same time they are not bolstered or too long. The diagrams are helpful as well and make for nice visual aids. The amount of background provided was also just right, making the paper self-contained and with a good flow. The technical writing and notation, although not a dominant part of this paper, are also to a high level. Overall, I enjoyed reading the paper.\n\n* Another strength is that the problem is well-motivated. The authors successfully conveyed the importance of the proposed setting to the `treatment effect of combinations of drugs' problem and I am confident that their approach could come handy for additional applications. \n\n* The approach is novel and the detailed implementation description serve as a nice contribution. The empirical results are also very encouraging and demonstrate how the conditional Monge map outperforms competing methods on several tasks. Altogether this makes for good practical/implementational contribution in the space of applied OT for inference and learning. \n\n**Weaknesses:**\n\n* I think that the paper can be improved by providing more theoretical context and justification for the conditional Monge map. Specifically, is there a natural (conditional) OT problem that one can set up for which the conditional Monge map emerges as a solution? Providing such a formulation would make the proposed framework more principled and would enable further theoretical exploration thereof. \n\n* Proceeding from the above point, given such a formulation, the authors would be able to pursue various theoretical results that would strengthen and further justify the proposed approach. A few immediate ones that come to mind include:\n1) Statistical approximation results and empirical convergence rates: how well do empirical conditional Monge maps (i.e., maps learned from data) approximate the population map? At what rate? what is the dependence of this rate on the ambient dimension (I would expect a curse of dimensionality effect to occur)?\n2) Approximation by neural network results: For the population measures, wow well can a NN-based architecture approximate the actual conditional map?\n3) Formal optimization guarantees: Can gradient/subgradient based methods be shown to converge to a local optimum or a stationary point for the conditional OT problem, and under what assumptions? Can we obtain a Sinkhorn-like algorithm to an entropically regularized version of this problem?\n\n* Like mentioned above, I think that the considered problem is interesting and the writing is excellent, but the level of novelty of the paper is a concern. The authors do well to combine existing ideas together (PICNN, Brenier's theorem, Gaussian approximations, etc.) to obtain an end-product that is greater than the sum of its parts. That said, the number of original ideas in the paper is rather limited. I believe that accounting for some of the above points would significantly strengthen the paper. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "lMu0NXohLga", "reviewer": "Reviewer_zoPg", "summary": "This paper proposes to learn a single Monge map $T$ between different pairs of probability measures $\\{(\\mu_i,\\nu_i)\\}$ such that the $i$th particular transport depends on a context function $c_i$, which corresponds to $T(c_i)$#$\\mu_i \\approx\\nu_i$. The Monge map is parameterized by a convex neural network with partial input.\n", "questions": "- A bottleneck in your method seems to be the initialization. Could you elaborate more on why the proposed initializations (identity or Gaussian) work? Are there any theoretical results on this?\n\n- You could specify that the measures you are considering are point clouds.", "limitations": "The underlying assumptions needed on the pair of probability measures for the method to work (in particular the initialization) are missing in my opinion.\n", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I find the idea of the paper, which consists of generalizing the learning of a transport map to several pairs of measures, which are first transformed by a \"context\" function, well justified in terms of application.\n\nHowever, as far as I understand the method, several issues and limitations seem to arise, which I develop in the following points.\n- An important part of the paper is to propose an efficient initialization. It is not clear to me why the initialization with Gaussian potential would work well in practice, for general dataset of measures that cannot be approximated by Gaussian distributions. It seems that there is a hidden assumption about the \"Gaussian\" geometry of the probability measures.\n- Since by the properties of the pushforward operator one has for a deterministic function $c_i:\\mathbb{R}^d\\to\\mathbb{R}^d$ that $T_\\theta(c_i)$#$\\mu_i$ is equivalent to $T_\\theta$#$(c_i$#$\\mu_i)$, a framework of your proposal could correspond to finding a unique application (or Monge map) that maps (approximately) each $c_i$#$\\mu_i$ to $\\nu_i$. Therefore, for this to be possible, I would say that the pairs $(c_i$#$\\mu_i,\\nu_i)$ must share a strong structure or relationship for all $i$. This is the case in the examples shown in Figure 1.\n\nAs a final comment, the proposed paper is very applied and the contribution is limited in my opinion, since the use of a convex neural network to parametrize a transport map has been widely developed in the literature.\n\nMinor comments:\n- l.99 : problem in the syntax of the sentence.\n- l.128 : \"Figure 3\" should be \"Figure 1\", same l.144.\n- Fig. 7 is not referenced in the text.\n- Would a natural continuation of this work would be to consider normalising convex flow for the invertible Monge map in order to have similar uncertainty for both $\\mu_i$ and $\\nu_i$? This could be done based on the paper [Coeurdoux et al., 2022], which uses normalising flows to approximate a Monge map between a pair of empirical distributions, and then generalizes it to other pairs of point clouds", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "d1mRnSqdbjv", "reviewer": "Reviewer_kDhV", "summary": "This paper proposes to learn a global map that is conditioned on the contexts. They achieve this by minimizing the Wasserstein-2 distance between the pushforward distribution and target distribution. They assume the measures from source and target distributions are paired and parameterize their map as PICNN. Their contexts are quite flexible, including scalars, covariates, and discrete actions. To integrate the contexts, they design an embedding module and combinator module to help the map learn the conditional dependence. Depending on the complexity of the contexts, the embedding module can take one-hot embedding or multi-dimensional scaling embedding. The combinator module can take multi-hot combinator or deep set combinator. The paper mainly applies their method to the biology data.", "questions": "Why do you choose Makkuva's dual formula instead of the primal formula? Do you want to get bi-directional maps?\n\nFigure 1 is a bit confusing, especially 1b, according to the covariate condition experiment, when $c_i$ changes, $\\mu_i$ also changes, but in fig1b, all $c_i$ corresponds to the same $\\mu_i=\\mu$.\n\nSome questions about experimental results: \n- In Figure 5, 6, you use one-hot embedding for unknown combinations, but I thought one-hot doesn't generalize to unknown action. Does \"unknown combinations\" mean combining existing perturbations instead of adding unseen perturbations? If that's the case, why not add an example for combinations of unknown perturbations? \n\n- In Figure 6, a kind of conflicting result is that one-hot embedding is working better than moa embedding. Does this mean moa embedding is not necessary?\n\n- Is the combinator module only used when you have the action of combination?\n\n- The source distribution $\\mu_i$ are the same for $\\forall~ i$, which is just the control cells, except for the covariates condition?\n\nMisc:\n\nFigure 6 is out of margin border.\n\nTypo: Line 217: where \"M\" is initialized ..\n\nLine 131: Please write ||x||^2 in front of the backslash to avoid confusion.\n\nThe title of sections D.1, D.2 has a format issue: three dots. \n\nRow 286: The figure pointers are not correct\n\nIn equation below line 555, sup shouldn't involve f*.\n\nCaption of Fig 5, \"one\" instead of \"ohe\"\n\nThere exists one concurrent paper about the conditional transport map. It's up to the authors to add it to the literature part or not. \"Asadulaev A, Korotin A, Egiazarian V, Burnaev E. Neural Optimal Transport with General Cost Functionals. arXiv preprint arXiv:2205.15403. 2022 May 30.\"", "limitations": "It asks the contexted measure to be paired. I also conjecture it's not easy to extend to imbalanced contexts, e.g. the two distributions have different number of labels.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nThe paper parameterizes the dependence over the contexts in a flexible way. This makes their framework applicable to not only scalar labeled data but more complex ones, like multi-label (actions) or soft labels (covariates). Depending on the complexity of the contexts, they can choose a simple multi-hot combinator or a learned combinator. If using a learned combinator, they can generalize to unseen data. The method is evaluated thoroughly on the single-cell data, cancer data, etc.\n\nWeak points:\n\nLiterature: 1) The technique of initializing ICNN to be an identity map has been introduced in [1], especially the quadratic layer technique in DenseICNN. 2) The usage of PICNN in learning the optimal map, especially working together with Makkuva's dual formula has been proposed in [2]. \n\nI'm worried about the writing. The fundamental loss function and algorithm are postponed to the supplementary material. To me, this is not proper. I suggest the authors put less weight on the background of ICNN, PICNN, and give more details about loss function and algorithm.\n\nIf I understand correctly, when you run ICNN OT, you just put data from all $\\nu_i $ (or $\\mu_i$) into a combined distribution $\\nu$ (or $\\mu$), and learn the map between $\\mu$ and $\\nu$. If this is the case, when all $\\mu_i$ are equal, the pushforward distribution by ICNN OT map is fixed, which doesn't depend on the context. To me, this is an unreasonable baseline when all $\\mu_i$ are equal. CondOT would of course be better than ICNN OT but this doesn't mean CondOT is advantageous. I think a good baseline is to calculate ICNN map for each $(\\mu_i , \\nu_i)$ pair. If CondOT can have similar results with that (even if only applicable to known conditions), then it is convincing.\n\n[1] Korotin A, Egiazarian V, Asadulaev A, Safin A, Burnaev E. Wasserstein-2 generative networks. arXiv preprint arXiv:1909.13082. 2019 Sep 28.\n\n[2] Fan J, Taghvaei A, Chen Y. Scalable computations of Wasserstein barycenter via input convex neural networks. arXiv preprint arXiv:2007.04462. 2020 Jul 8.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept