Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper extends CellOT to the case of unbalanced optimal transport, which is essential to modelling most cell systems. This extension is accomplished through the formalism of semi-couplings, learning reweighted distributions which are then transported with balanced transport, in this case using the ICNN framework. A novel algorithm is presented to solve this problem efficiently through gradient descent. This algorithm is evaluated on a toy dataset and a single-cell drug perturbation dataset captured with 4i imaging. NubOT fits the data better in that it is able to predict cell states for unseen cells (but in seen conditions) better than comparable methods.
strength_and_weaknesses: Strengths:

- Important extension to the CellOT formulation which accounts for relative size changes in cell populations, which is extremely important in heterogenous cell systems.
- Well grounded theoretical formulation of unbalanced OT with semi-couplings.
- Good dataset choice for the evaluation of the unbalanced transport with a ground truth measurement of the number of cells at a population level before and after treatment.

Weaknesses:

- The connection between Algorithm 1 and the presented theorems is not entirely clear. From Algorithm 1 I would have expected the transformation for new samples to be $(x, u) \to (\nabla g(x), \eta(x) \cdot u)$. Could the authors clear up what is correct, algorithm 1 or figure 1?
- To me “forecasting heterogenous responses of multiple cancer cell lines to various drugs” would imply extrapolation to unseen drugs or cell line or at least combinations, instead of the forecasting the response of an unobserved cell in a seen drug and cancer cell line as is done here. Especially considering a separate model was trained for each treatment. Would it be possible to tone down the language on applications? In particular the "This is an unprecedented achievement in the field of single-cell biology" line? Modelling in distribution responses with proliferation has been tackled in prior work for time series, applying it to a seen perturbation response is in my view a relatively straight forward extension.
- The application is extremely specific, it is unclear how interesting it is to a wider ICLR audience. It would be useful to discuss other applications for unbalanced OT such as domain alignment, transfer learning etc.

Comments / Questions:

- In the “Updating Mappings” paragraph: “Since these are guaranteed to be balanced due to the argument above”. Could the authors clarify this please? Why is this guaranteed? How is guaranteed that $\int \tilde{\mu} = \int \tilde{\nu}$?
- There is something that is not lining up for me between the optimization in algorithm 1 and how samples are transformed. To me it seems like you would want to weight $\hat{y}$ by
- With a batchsize less than the entire data this is similar to a mini-batch OT problem. The unbalanced minibatch OT problem is in some senses creates a similar (full) map to the (full) balanced OT problem (Fatras et al. ICML 2021). Further, it doesn’t have the same weights as the regular non-batched problem. Does batch size have any affect on your results? I think this should at least be mentioned as not solving the same problem as UBOT.
- I’m curious how the learned weights compares to a discrete UBOT algorithm. I understand that there are advantages of learning a continuous map, however, I would still like to see a comparison to understand how accurately the proposed NubOT algorithm solves UBOT.
- I don’t believe that there are any constraints on the total sum of the transported measure $Z = \int_{(x,u)} \eta(x) \cdot u \cdot \zeta( \nabla(g(x))^{-1})$. The weighted MMD is invariant to Z, thus there is no evaluation whether the total measure is accurate. Is this total mass at all important to your application?
- Is there (albeit small) data leakage from data normalization in the single cell data to the 75% level?
- How was the relaxation penalty chosen, and do the authors expect it to be tuned on new datasets? What was the batch size chosen?

Minor Comments:

- Line 7 of Algorithm 1 should have an $x$ instead of $y$. Line 1 might want to be batches instead of epochs?
- I would note that $L^1$ norm unbalanced OT is not exactly the same as partial OT (As stated in A.1). Partial OT is slightly different. Instead this $L^1$ norm is sometimes “Robust optimal transport” (Mukherjee et al. ICML 2020).
- Could the true weights in Table 1 be depicted in Figure 2? Its difficult to understand that these are the correct weights without scrolling to Table 1.
- What is the weighted kernel MMD? What is the kernel? Gaussian with what bandwidth?
clarity,_quality,_novelty_and_reproducibility: The experimental sections are well motived and clear. However, the presentation of the algorithm, and how it relates to the UBOT problem could be made more clear. 

The novelty is modest, combining work on ICNN based OT and unbalanced transport via rescaling as in Yang and Uhler 2019 and Tong et al. 2020. 

The experiments are largely reproducible.
summary_of_the_review: Overall a useful improvement concept with a clear and specific application. There are a few clarifications on the algorithm and the experiments. If these things are cleared up I have no problem raising my score.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: First, I am an emergent reviewer and not an expert on OT. Therefore, I'm not sure whether the literature is reviewed properly and whether the evaluation setup is fine.

This work presents a formulation of the unbalanced optimal transport problem. To solve the unbalance problem, it relies on the formalism of semi-couplings to account for the creation and destruction of mass. 

It further presents a cycle-consistent training procedure upon neural OT to estimate such semi-couplings and generalize out-of-sample, named NUBOT. 

The method is evaluated on a synthetic dataset and the single-cell perturbation task. It is compared to CellOT (a balanced neural OT method) and UBOT GAN (an imbalanced OT approach).  

According to the reweighted MMD (and other metrics), NUBOT outperforms the baselines on both datasets. Moreover, on the challenging single-cell perturbation task, NUBOT is able to successfully predict perturbed cell states, while explicitly modeling death and proliferation. 
strength_and_weaknesses: **Strength**

1. The paper is well-organized and clearly written. 

2. The empirical results (both qualitative and quantitative ones) seem pretty good.

**Question**

I have several questions here.

1. Why use MMD instead of the OT distance as the evaluation metric in Fig. 4? It is more natural to evaluate the objective on the validation dataset for machine learning algorithms.

2. The method presented in Section 3 seems complex to me. It would be better to discuss the necessity for each part in depth and present the empirical contribution of each part in the experiments.
clarity,_quality,_novelty_and_reproducibility: Clarity: 

The paper is well-organized and clearly written. 

Quality&Novelty:

The paper is of good quality for me and the technical contributions are clearly stated. However, I am an emergent reviewer and not an expert on OT. I'm not sure whether the literature is reviewed properly and whether the evaluation setup is fine. Therefore, I cannot justify its quality and novelty confidently. 

Reproducibility:

It presents the experimental details including the data preprocessing, network architectures, and hyperparameter tuning process. It seems reproducible.
summary_of_the_review: This is a solid paper on addressing the unbalanced OT problem and I tend to accept it now.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: Optimal transport distances (a.k.a. Wasserstein distance) have recently drawn ample attention in statistics and machine learning communities as powerful discrepancy measures for probability distributions. The standard formulation of OT assumes equality of mass to be transported from the source to the target, which is referred as to the conversation mass constraint. In other applications, this constraint is not filled; hence application of standard OT is not done directly. Existing lines of research investigate this problem by handling the non-equality of mass transportation. They consist of unbalanced Optimal transport (UBOT). These approaches are given by the sum of an OT cost and two constraints involving the non-equality of mass transportation. The second term is often written using divergence, for example, a Kullback-Leibler divergence, etc.

This paper proposes NUBOT, a neural unbalanced OT that relies on the formalism of semi-couplings to account for creation and destruction of mass. NUBOT takes advantage of the neural optimal transport, where it was shown that the Monge plan can be learned via the Kantorovich potentials of the dual problem, which can be parametrized via input convex neural networks (ICNNs). Towards this end, the authors suggest learning a proxy measures $\tilde{\mu} = \eta \cdot \mu$ and  $\tilde{\nu} = \zeta \cdot \nu$ to rescale the mass. The functions $\eta, \zeta$ are also parametrized via NN. An algorithm of NUBOT is given with an extensive numerical study with an accent on cell-biology applications.
strength_and_weaknesses: ### Strength ###
- Formulation UBOT with a neural approach based on semi-couplings proxy measures. The Kantorovich potentials are learned with ICNNs. 
- Extensive numerical study on synthetic and real data with an accent on cell-biology applications: prediction of cell proliferation and death, and the perturbation responses on the level of single cells.

### Weaknesses ###
- Steps 5 and 7 in Algorithm 1 deal with calculating an unbalanced Sinkhorn: it is a little be confused since the ultimate goal of this paper is to tackle UBOT. 
- The paper lacks comparisons with the SOTA approaches of UBOT: NUBOT is only compared to ubOT GAN? One can compare it to a vanilla generalized Sinkhorn algorithm (Chizat et al. 2018). 
- What is the computational complexity of NUBOT. I think this point should be clarified or at least the authors can give some insight bout it. Steps 5 and 7 in Algorithm 1 might be the most expensive ones.
- In Algorithm 1, it is not clear to me the update of $\eta$ and $\zeta$.

clarity,_quality,_novelty_and_reproducibility: - The paper is well-written and easy to follow. 
- Code is not available.

summary_of_the_review: 

### Typos ###
- Page 3: Eq. (3): there is no dependency on $\varepsilon$.
- Page 3: Eq. (5): OT$_$n. I don't understand this dependence on n.
- Page 4: The notation of the set of semi-coupling between measures $Gamma(\mu, \nu)$ is the same as the set of couplings in the nonnegative Radom measures (page 2). 
- Page 4: 
$\Gamma \gets \text{UBOT}(\boldsymbol{u}, T(\boldsymbol{x}_i), \boldsymbol{v}, \boldsymbol{y}_j)$ --> $\Gamma \gets \arg\min\text{UBOT}(\boldsymbol{e}\odot\boldsymbol{u}, T(\boldsymbol{x}_i), \boldsymbol{v}, \boldsymbol{y}_j)$
- Page 4: 
$\Gamma \gets \text{UBOT}(\boldsymbol{v}, S(\boldsymbol{y}_j), \boldsymbol{u}, \boldsymbol{x}_i)$ --> $\Gamma \gets \text{UBOT}(\boldsymbol{z}\odot\boldsymbol{v}, S(\boldsymbol{y}_j), \boldsymbol{u}, \boldsymbol{x}_i)$
- Page 6: "intervention respnses" --> "intervention responses"
review_title: null

final_score: rating: 3: reject, not good enough
review_content: summary_of_the_paper: The paper proposes a method for finding optimal transport maps between unbalanced distributions, i.e., distributions with different total masses. The method is composed of two steps repeated iteratively. First, it seeks for scaling factors for source and target measures, such that the re-scaled measures are balanced. This problem is considered as a discrete unbalanced optimal transport problem and solved using the Sinkhorn algorithm. These factors are used to train neural networks to predict scaling factors for arbitrary points. At the second step,  the method finds OT maps between re-scaled measures using standard balanced OT formulation (Makkuva et al., 2020). 

The method is tested on the synthetic problem of finding a map between unbalanced mixtures of Gaussians, and the real-world problem of testing cells' response to different drugs which may induce proliferation (quantitative growth) or death of cells, i.e., change in distribution mass. 
strength_and_weaknesses: **Strengths.** The paper proposes an optimal transport approach for an important unbalanced setup, which is not yet fully explored.

**Weaknesses.**
1. The proposed algorithm is based on the solutions of the discrete unbalanced OT problem. Thus, it is unclear to what extent the method approximates the actual underlying solution of the continuous problem. This might lead to bias in method solutions.
2. The proposed approach lacks comprehensive quantitative evaluation. It is not shown that it learns actual scaling factors and OT map between distributions.

clarity,_quality,_novelty_and_reproducibility: **Quality.** The proposed method uses the solutions of the discrete unbalanced OT problem to train networks to predict scaling factors. However, the solutions of discrete OT methods applied to batches are biased regarding the solutions of continuous / full batch problems. It is unclear to what extent their solutions and the learned networks match the actual solutions. Specifically, it is not evident how well the networks approximate the actual scaling factors, and, thus, how well the method approximates an actual OT map. In the proposed framework, I  think that it is necessary to show that such approximations are reasonable. The authors do not conduct a convincing experiment showing that the method calculates what is intended, i.e., finds an optimal transport solution. At the same time, the method is not supported by theoretical analysis.

*(a)* In the setup of the synthetic experiment, the authors consider the task of mapping unbalanced mixtures of Gaussians, for which, as the authors write, “a ground truth matching is available”. However, they did not include any comparison of the proposed method results with this “ground truth matching”. And how do you derive the ground truth OT maps between the mixtures? 

*(b)* I recommend the authors to quantitatively evaluate the method using unbalanced Gaussians for which both the precise scaling factors and precise OT maps seem to be analytically known [1]. Thus, it is possible to rigorously quantitatively evaluate whether the proposed method actually recovers the precise OT map and scaling factors. 

[1] Janati, H., Muzellec, B., Peyré, G., & Cuturi, M. (2020). Entropic optimal transport between unbalanced Gaussian measures has a closed form. Advances in neural information processing systems, 33, 10468-10479.

The main experiment section 4.3 of the paper focuses on the task of predicting single-cell responses to different treatments, which raises the problem of mapping unbalanced distributions of cells before and after the treatment. The authors provide many biological details regarding this problem which explain the choice of the task-specific metrics for comparison with other methods. Despite the fact that the method improves these metrics, I am concerned about this experiment because it is not transparent what happens inside the method (see my above comments). I am not familiar with the problem considered, so I wonder whether there are any non-OT based methods for it to compare with. The discussion of this aspect is limited.

Comparison with other unbalanced GAN OT method (Yang & Uhler, 2019) also might seem to be incomplete in some sense. Can the proposed method be applied to the other unbalanced OT tasks originally considered in (Yang & Uhler, 2019), e.g., in CelebA-Young/CelebA-Aged case? A comment about this is appreciated.

**Clarity.** In general, I think that some of the paper results are not well explained. From Section 4.1, subsection ‘Updating rescaling functions’, I do not completely understand why the measures re-scaled using obtained scaling factors are balanced? Section 4.3 is very hard to follow due to the large amount of biological terminology. In my view, this is a serious drawback of the paper because the experimental part may be hard to parse for ML practitioners without significant background in biology.

**Reproducibility.** The authors do not provide the code of their method. Thus, the reproducibility of their method is questionable. (However, some technical details are given in Appendix.).

**Minor remarks.**

(1) In the Equation (5) - squares of norms are missing, subscript n is undefined. 

(2) In the Equation (7) - two subscripts m.

(3) In Section 3, subsection ‘Updating rescaling functions’ - typo in ‘our goal is to find w’.

(4) In Section 3, subsection ‘Transforming new samples’ - ‘mapped backed’ should be changed to ‘mapped back’.

(5) In Figure 3 - undefined parameters R and P.

(6) In section 4.1 typo in ‘(see Fig. 10, 12, as well as Fig. 10, 12).
summary_of_the_review: The authors propose a method for finding an optimal transport map between unbalanced distributions. However, from the construction of the method it is unclear how well it approximates actual scaling factors and OT map between the distributions. The experimental section also does not clarify this issue. Due to this, I think that not enough evidence is given to convince the reader that the method computes what it is intended to do. Moreover, the entire focus of the method is a particular biological task, which is hard to follow for an unprepared reader. Because of all these issues, I am currently more on the negative side about the paper. If the authors solve these evaluation/clarity issues during the rebuttal phase, I will be happy to reconsider my assessment.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: In this submission, the authors proposed a new member of the neural network-assisted OT algorithms for aligning unpaired samples in an unbalanced setting. Experiments on single-cell data are designed to demonstrate the usefulness of the proposed method.

Strengths:
(1) The idea is clear, and the derivation is easy to follow. Parametrizing OT plans via neural networks is an attractive topic in the community.
(2) The application considered by this submission, i.e., single-cell data analysis, is important.

Weaknesses:
(1) The novelty of the proposed method is limited, as the reviewers mentioned.
(2) The numerical stability and the scalability of the proposed method are not verified, especially when the sample size is large and the feature dimension is high.
justification_for_why_not_higher_score: Essentially, this work is a new member of the neural network-assisted OT algorithms for aligning unpaired samples in an unbalanced setting. In my opinion, it is a combination of existing techniques (e.g., ICNN, semi-coupling, and unbalanced Sinkhorn), and the main contribution is combining existing techniques to improve the performance of a significant application. Besides the lack of novelty, my main concerns include: (1) as some reviewers mentioned, the complexity and the scalability of the method may be questionable because of the usage of the unbalanced Sinkhorn --- when the batch size is large, some numerical issues will happen according to my experience; (2) in the comparison on MMD, the performance of the proposed method is just slightly better than CellOT and GaussApprox. To my knowledge, GaussApprox is much more computationally-efficient than the proposed method when the dimension of the feature is not very large. After reading the authors' reply, my concerns above are not solved.

The reviewers also proposed some comments about the novelty and the complexity of the proposed method. Although the authors tried to resolve these concerns, some reviewers are not very satisfied with the replies. Two reviewers raised their scores from 5 to 6, but one had low confidence.
Decision: Reject