Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: 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. 
strengths_and_weaknesses: **Strengths:**

* 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.

* 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. 

* 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. 

**Weaknesses:**

* 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. 

* 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:
1) 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)?
2) Approximation by neural network results: For the population measures, wow well can a NN-based architecture approximate the actual conditional map?
3) 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?

* 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. 
questions: Major questions/concerns appear under `Weaknesses' above. Some smaller issues and specific questions are listed next:

1. **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. 

2. **Section 3.2, Item 1:** In the last sentence, $\psi_\theta$ is referred to as an 'OT map' but it should be 'OT potential'.

3. **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?

4. **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.

5. **Fig. 6:** Why is there no comparison to ICNN OT [Makkuva et al., 2020] for that experiment? 

6. **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.

I have no concerns about potential negative societal impact.
review_title: null

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: 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. 
1) Use optimal transport as a judicious way of selecting the push forward map.
2) 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.
3) 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). 
4) 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. 

The main contributions of the paper are: 
1) Describing a state of the art framework that allows for unseen/composite contexts and utilizes the 'continuity' of the context space. 
2) 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.
strengths_and_weaknesses: Strengths:
1) The problem is well-motivated, interesting and important.
2) The paper is clearly written and easy to follow. 
2) The empirical results suggest that the outlined framework is superior to prior work on common tasks.
3) The outlined framework allows to generalize to unseens contexts and combine them, which wasn't possible with prior work. 

Weaknesses:
1) 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. 

2) 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. 
questions: ### Questions/Comments:
1) 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. 
2) In Section 2) clarify the relevant regularity requirements for the OT map to exist. 
3) What is _+ in display (4) and bottom of page 6?
4) 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)? 
5) Why not always take train_freq_f=2? (as in why is it beneficial to train $g$ more than $f$)
6) 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. 




### Typos/minor suggestions:
- Line 3: repeats the word optimization
- Line 18: from of 
- Display (2) why introduce f instead of just use \psi (also display (5))
- Figure 6b perturbation vs combination unclear
- Line 73: unique convex potential
- Line 91: …we encode this by …
- Paragraph under line 106: …leveraging the context values as much as possible …
- same paragraph: …propose to learn a unique …. Unique in what sense?
- Line 108: …not in the context value…
- Line 119: replace ‘condition’ by ‘context’?
- Line 123: …and the weight matrices…
- Line 155: allows vs encourages
- Line 320: …that is able to…
- Line 538: ...in that …
- Line 548: …learn our model…
- Line 552: definition of \tilde\Psi seems wrong, shouldn’t be a product
- Line 553: $\int |f| d\mu < \infty$
- Line 572: parentheses
- Line 747: multiple typos
limitations: The limitations and potential negative societal impact have been adequately addressed.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: 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.
strengths_and_weaknesses: - The paper is well-written and well-presented. 
- 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.
- 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).
- 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.
- 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).
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?

- Should there be some out-sample results in the setting of covariate conditioning in the experiments?
- Why are the UMAP embeddings of data in Figure 7 different in three plots? Should they be the same given the same data?
- 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.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: 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.

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.

However, as far as I understand the method, several issues and limitations seem to arise, which I develop in the following points.
- 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.
- 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.

As 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.

Minor comments:
- l.99 : problem in the syntax of the sentence.
- l.128 : "Figure 3" should be "Figure 1", same l.144.
- Fig. 7 is not referenced in the text.
- 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
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?

- 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.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: 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.
strengths_and_weaknesses: Strengths:

The 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.

Weak points:

Literature: 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]. 

I'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.

If 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.

[1] Korotin A, Egiazarian V, Asadulaev A, Safin A, Burnaev E. Wasserstein-2 generative networks. arXiv preprint arXiv:1909.13082. 2019 Sep 28.

[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.
questions: Why do you choose Makkuva's dual formula instead of the primal formula? Do you want to get bi-directional maps?

Figure 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$.

Some questions about experimental results: 
- 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? 

- 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?

- Is the combinator module only used when you have the action of combination?

- The source distribution $\mu_i$ are the same for $\forall~ i$, which is just the control cells, except for the covariates condition?

Misc:

Figure 6 is out of margin border.

Typo: Line 217: where "M" is initialized ..

Line 131: Please write ||x||^2 in front of the backslash to avoid confusion.

The title of sections D.1, D.2 has a format issue: three dots. 

Row 286: The figure pointers are not correct

In equation below line 555, sup shouldn't involve f*.

Caption of Fig 5, "one" instead of "ohe"

There 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.
review_title: null

Meta-review: metareview: The paper considers the problem of learning a single Monge map between different pairs of probability measures such that the particular transport depends on a given context. The reviewers have found the paper well written and motivated. In addition the approach is novel and interesting. I therefore recommend acceptance. 
confidence: Certain
Decision: Accept