Paper: Neural Optimal Transport with General Cost Functionals (ICLR-2024)

Reviewer comment (Novelty):
The paper heavily draws upon the prior work of (Korotin et al. 2023a) for its theoretical foundations. Approximately 5 out of 9 pages are dedicated to presenting these theoretical results, which may not constitute the primary novelty of the paper. Based on my personal reading, it seems that the authors might be overselling their theoretical contributions, potentially leading to a less reader-friendly introduction.

Author response:
We appreciate the reviewer's observation and concerns, but we believe that the theorems, corollaries, and results presented in Sections 3 are novel and important. These results are form the theoretical backbone of our work and crucial in establishing a theoretically justified algorithm that surpasses prior OT methods in practice (Section 5) as it allows using additional label information. Moreover, the unique error analysis via duality gaps for the general cost in Section 4 represents a novel aspect of our approach, distinguishing it from existing methods.

Reviewer comment (Presentation):
The paper exceeds the strict upper limit of 9 pages for the main text of the submission by including a section on reproducibility (Section 7) on Page 10. As a reviewer, this doesn't pose an issue for me, but I would advise adhering to the prescribed page limits as a matter of following the rules.

Author response:
Please note that according to ICLR 2024 rules, it is allowed to include the reproducibility statement as an additional 10th page, see the reproducibility section in https://iclr.cc/Conferences/2024/AuthorGuide.

Reviewer comment (Novelty):
A direct comparison of the theoretical findings with those presented in (Korotin et al. 2023a) is essential. Readers are likely seeking a consolidated presentation rather than having to review two separate papers with overlapping theoretical content.

Author response:
To be honest, we do not clearly understand your concern. Korotin et. al. (2023a) proposes kernel regularizers which we use as an example of a strongly convex regularizer in **Appendix D** of our paper. Korotin et. al. (2023a) uses only the *strict* convexity of this functional, while we derive its **strong** convexity for our purposes. Please see the detailed discussion before Proposition 3 (Page 28).

Reviewer comment (Novelty):
From my perspective, the paper's novelty appears to lie more in its practical applications, particularly in utilizing a cost functional to map data distributions while maintaining class-wise structures. In light of this, it would be beneficial for the paper to allocate more of its main text to discussing practical implementation aspects.

Author response:
Thank you for acknowledging our practical contribution. However, it is important to note, as mentioned in our response to Q1, that our theoretical contribution is equally significant. General costs constitute a fruitful field in discrete OT, and our work extends them to the continuous case, effectively bridging a gap between OT and deep learning. We have established theoretical properties that specifically arise in the continuous case. We believe that our analytical insights can prove valuable for future studies in neural optimal transport.

The technical details are presented in **Appendix C.1** and the source code is in the supplementary material. *Could you please recommend, which practical implementation aspects to move to the main text from Appendix?*

Reviewer comment (Novelty):
The novelty of the paper, in comparison to Ref[1], is weak. The theoretical novelty, in comparison to the existing theoretical result in Ref[1], is not explained well (what is the new approach or new tool that is being used here). What are the challenges of considering a general cost functional that the previous approach could not handle.

Author response:
**Q1.1:  The novelty of the paper, in comparison to Ref[1] (...) what is the new approach or new theoretical tool that is being used here.**

- When deriving our max-min duality formula for general OT (section 3.1/3.2), we do not introduce new principal tools which have not been considered before. We just derive a max-min reformulation of general OT; it generalizes previously known classical and weak costs Ref[1], see the discussion at the end of Section 3.2. This form enables us to establish an algorithm for general OT which **can be applied to problems where the predecessors are not directly applicable** (*see the answers to Q1.2, Q2 below*).


- The **newly introduced theoretical tool** is the duality gap analysis of Section 3.3. Please see the discussion around the Theorem 3. Previously known error analysis works exclusively with the classical OT and operate only under certain restrictive assumptions such as the convexity of the dual potential. In contrast, our error analysis is free from assumptions on the dual variable and *is applicable* not only to general OT but also *to weak OT Ref[1], for which there is currently no existing error analysis*. Please reconsider the end of **Section 3.3** (relation to prior works) and **Appendix D**.


**Q1.2. What are the challenges of considering a general cost functional that the previous approach could not handle?**

Our general cost functional-based algorithm can use **both labeled** and **unlabeled** target samples for training, which can be useful for the data transfer tasks. Existing continuous OT approaches do not handle a such type of training. Indeed, suppose we have additional information (labels) in the dataset and try to solve the class-guided mapping using the Ref[1] (as you mentioned in your Q4 below). In this scenario, we can train Ref[1] using only the labeled samples (10 separate maps in case of MNIST). In this case, the **unlabeled data immediately becomes useless**. Indeed, using unlabeled data for a class during training for that class implies that we know the labels for that class, which is a contradiction. Thus, training Ref[1] using both labeled and unlabeled target samples, as in our approach, is not possible.

Reviewer comment (Novelty):
The computational algorithm is also very similar to that in Ref[1]. The NOT algorithm block in Ref[1] can be simply extended to a general cost functional. The proposed algorithm block is very similar, with the only difference that it is written for a class-guided functional.

Author response:
**Q2. The computational algorithm is also very similar. The NOT algorithm block in Ref[1] can be simply extended to general cost functional. (...) the only difference that it is written for a class-guided functional.**

We agree that our algorithms may look similar, but this is because Ref[1] is a special case of our general approach. Our algorithm provides a better flexibility for using continuous OT. As we write in **Section 4**, the **differences** lie is the estimation of the cost functional $\mathcal{F}$ to perform the gradient updates. This difference is crucial, because such estimation, for example, allows to exploit both labeled and unlabeled data (see the answer to your previous question).

Reviewer comment (Experiments):
It would be interesting to see and discuss how the algorithm performs when the data from classes overlap.

Author response:
**Q5: It is interesting to see and discuss how the algorithm performs when the data from classes overlap.**

To specifically address the performance in scenarios where classes overlap, we conducted a class-guided mapping task on a Gaussian dataset with overlapping classes. For more detailed insights, please refer to **Appendix C.2**. Our method effectively handles this scenario, showcasing the robustness of our model in dealing with overlapping classes.

Reviewer comment (Theory):
Regarding the discussion at the beginning of Sec 3.2., the reviewer questions why there is a "measurable" map for each coupling. They understand existence for each x, but are not sure how to argue existence of a measurable map as a function of x and z.

Author response:
**Q6: Regarding the discussion at the beginning of Sec 3.2., why is there a "measurable" map for each coupling? I understand existence for each x, but not sure how to argue existence of a measurable map as a function of x and z.**

This fact follows from [2]. In particular, Theorem 5.3 yields that the mapping $x \rightarrow \pi(\cdot \vert x)$ is a so-called probability kernel (see definition in [2]). After that, Lemma 2.22 establishes that there exist some measurable function $T : \mathcal{X}\times [0, 1] \rightarrow \mathcal{Y}$ (we utilize our notation here) such that $T(x, \cdot)\sharp U(0, 1)$ has distribution $\pi(\cdot \vert x)$. One can easily substitute the uniform distribution $U(0, 1)$ with our considered latent distribution $\mathbb{S}$.

Reviewer comment (Experiments):
The setup is very toyish: the datasets are toy datasets, and the labels are entirely unrelated, e.g., matching the digit '1' to 'trouser'. This is fine as a first toy experiment but there must be some more interesting ML applications where the labels from P and from Q have a relationship and are not paired randomly. Having only this artificial experiment is underwhelming and does not convince the reader that it is actually an interesting problem for machine learning.

Author response:
**Q1: the datasets are toy datasets, and the labels are entirely unrelated: why match the digit '1' to 'trouser'?  (...) there must be some more interesting ML applications where the labels from P and from Q have a relationship and are not paired randomly.**

These experiments are for illustrative purposes only. Our goal here is to show that even when classes are *unrelated* ("1" to "trouser", etc.), our method still can learn an class-preserving map. For experiments on the *related* domains, when the labels from $\mathbb{P}$ and from $\mathbb{Q}$ have a relationship, please see **Appendix C.3**. 

In our paper, we also considered the application of our method to the well-known problem in biology called the *batch effect*, see **Appendix C.12**. This problem often occurs in genomic studies such as gene expression microarrays, RNA-seq, and proteomics analyses.

For **new additional experiments** involving high-resolution images, please see the **general answer** provided to all reviewers and **Appendix E**.

Reviewer comment (Evaluation):
The computation of metrics for stochastic outputs could be clearer. It is unclear whether they are averaged over z.

Author response:
**Q3: Same question for the metrics: how it is computed for stochastic outputs could be clearer. is it averaged over z?**

Initially, we tried using multiple $z$ per $x$, generating $T(x,z)$, and averaging the metric over them. Later, we found that the result is the same if we sample only one $z$ per $x$. So we finally decided to sample a single $z$ per $x$ for simplicity. We have clarified this in the revised version. (Appendix C.1).

Reviewer comment (Evaluation):
It is unclear whether FID is computed per class or on the whole dataset; this should be clarified.

Author response:
**Q4: Is FID computed per class or on the whole dataset? this should be clarified? All fonts are too small in the figures**

The FID was calculated on the entire dataset, not separately for each class. Following you suggestion, we clarified this in the **Metrics** paragraph of the paper (Section 5).

Reviewer comment (Theory):
The cost function proposed in Proposition 1 contains a quadruple sum over samples; a discussion about its variance would be welcome.

Author response:
**Q5: Another point is that the cost function proposed in Prop.1 contains a quadruple sum over samples: a discussion about its variance would be welcome. Also, the proposition mentions that it is an estimator: in which sense?**

The reported quantity in Eq. (15) is the estimator of energy distance $\mathcal{E}$ in the sense that its expectation w.r.t. batch samples $X_n, Z_n, Y_n$ yields $\Delta \mathcal{E}^2(T_{\pi}\sharp (\mathbb{P}_n \times \mathbb{S}), \mathbb{Q}_n)$. 

It equals to $\mathcal{E}^2$ up to $T_\pi$-independent constant $C = - \frac{1}{2} \mathbb{E} \Vert Y_2 - Y_2' \Vert_2$, where $Y_2, Y_2'$ are independent copies from the target distribution $\mathbb{Q}_n$. 

Regarding the quadruple sum, we refer the reviewer to [1] for the analysis of similar quantities. Note that in practice we did not experience any problems with estimation of Eq. (15). 

[1] Sutherland et. al. Unbiased estimators for the variance of MMD estimators.

Reviewer comment (Theory):
Proposition 1 mentions that the cost function is an estimator; it should be clarified in which sense.

Author response:
**Q5: Another point is that the cost function proposed in Prop.1 contains a quadruple sum over samples: a discussion about its variance would be welcome. Also, the proposition mentions that it is an estimator: in which sense?**

The reported quantity in Eq. (15) is the estimator of energy distance $\mathcal{E}$ in the sense that its expectation w.r.t. batch samples $X_n, Z_n, Y_n$ yields $\Delta \mathcal{E}^2(T_{\pi}\sharp (\mathbb{P}_n \times \mathbb{S}), \mathbb{Q}_n)$. 

It equals to $\mathcal{E}^2$ up to $T_\pi$-independent constant $C = - \frac{1}{2} \mathbb{E} \Vert Y_2 - Y_2' \Vert_2$, where $Y_2, Y_2'$ are independent copies from the target distribution $\mathbb{Q}_n$. 

Regarding the quadruple sum, we refer the reviewer to [1] for the analysis of similar quantities. Note that in practice we did not experience any problems with estimation of Eq. (15). 

[1] Sutherland et. al. Unbiased estimators for the variance of MMD estimators.

Reviewer comment (Writing):
The statement "The performance of such methods in high dimensions is questionable" would benefit from a reference to elaborate on this.

Author response:
**Q6.2 "The performance of such methods in high dimensions is questionable": a reference to elaborate on this would be welcome.**

The general problem with these methods (wavelet/kernel-based, barycentric)  is actually the curse of dimensionality. Application of these methods bears similarity to application of kernel density estimators for approximating a probability distribution. It is not applicable to the distribution of images. This is a general problem with non-parametric setup. Regarding more theoretical justification, see [1]. They showed that as the dimensionality $d$ of data grows, the quality of the recovered map in the worst case behaves like $\approx N^{-c/d}$ for a constant $c$, where $N$ is the number of empirical samples. 

[1] Hutter et. al., Minimax rates of estimation for smooth optimal transport maps.

Reviewer comment (Theory):
The statement "it does not always have a solution with each P_n exactly mapped to Q_n" would benefit from a few more words about why this is the case.

Author response:
**Q6.3: "it does not always have a solution with each $\mathbb{P}_n$ exactly mapped to $\mathbb{Q}_n$" a few more words about why this is the case would be nice.**


The problem may occur with the existence of the solution since we try to map all class-specified distributions $\mathbb{P}_n$ to their counterparts $\mathbb{Q}_n$ with a **singe** stochastic map $T$. If, say, the source distributions are the same while the target ones are different, there is obviously no suitable map $T$.

Reviewer comment (Experiments):
The baseline accuracy of each ResNet used in the experiments is not provided.

Author response:
**Q6.4 what is the baseline accuracy of each resnet used in the experiments?**

The accuracy of the ResNet18 classifiers is 99.17 for the MNIST and 95.56 for the USPS datasets. For the KMNIST and MNITST, the accuracy's are 99.39 and 97.19, respectively. We have included this information in Appendix C.1.