Title: OTCOP: Learning optimal transport maps via constraint optimizations
Authors: Xiaokai Huo
Hailiang Liu
Venue: ICLR 2023
Keywords: optimal transport
constraint optimization
Monge problem

Abstract: The approximation power of neural networks makes it an ideal tool to learn optimal transport maps. However, existing methods are mostly based on the Kantorovich duality and require regularization and/or special network structures such as Input Convex Neural Networks (ICNN). In this paper, we propose a direct constraint optimization algorithm for the computation of optimal transport maps based on the  Monge formulation. We solve this constraint optimization problem by using three different methods: the penalty method, the augmented Lagrangian method, and the alternating direction method of multipliers method (AMDD). We demonstrate a significant improvement in the accuracy of learned optimal transport maps on benchmarks. Moreover, we show that our methods reduce the regularization effects and accurately learn the target distributions at lower transport cost.


TL;DR:  We integrates constraint optimization algorithm and neural networks for the computation of optimal transport maps based on the  Monge formulation.

Review discussion:
The strengths of the current work
1. Propose novel algorithm to solve optimal transport map method directly using Monge's formulation with constraint optimization.
2. Three different optimization algorithms are formulated, evaluated and compared.
3. Theoretic proofs for the existence of the solution (Theorem 1 and Theorem 2), and the convergence of ADMM (Theorem 3)
The mathematical formulation and proofs are clean, rigorous and thorough.
The weakness of the current work is that the experiments are conducted on simple datasets (MNIST), it is desirable to test the theoretic results on real world datasets.
Another weakness is that by using Yau and Gu's geometric variational algorithm, the problem of finding Monge Optimal transport map is reduced to a convex optimization problem, so the theoretic setup is easier and the optimization can be carried out using Newton's method. It is helpful to compare the proposed method with this convex optimization approach.
Another weakness is that the proposed method can not find the singularity of the OT map precisely, since it use soft constraints (mollifiers) which will obscure the boundary of the supports.
The idea of using constraint optimization to solve optimal transport problem has been existed for a well. By adding the penalty terms to the objective function with Lagrange multiplier is conventional. The ADMM algorithm is relatively novel than the other two methods. The theoretic proofs are convincing, especially theorem 3 guarantees the convergence of the proposed algorithm, this rigorous result has high practical value as well.
But the scope of theorem 2 is limited, only for normal distributions. The paper is well written, all the mathematical formulations, key concepts, proofs are very clean and succinct. The algorithms are given in details and the experimental results are easy to reproduce.
NO.
This research is fundamental and mainly focuses on theoretical aspects.
Many thanks for reviewing our paper and providing the valuable comments. We thank you for appreciating the presentation of our content. We have uploaded a revised version of our paper. Below is a detailed response to your concern.
1. As the reviewer suggested, we are working on applying our methods to more complicated problems. Considering the theoretical results and successful application in simple datasets, we expect our methods to be applicable to more complicated problems.
2. Thanks for referring Yau and Gu's geometric variational algorithm to us. Since they transfer the problem into a convex one, the constraint optimization technique can also be applied and we expect better theoretical results. Different from our work, they use the Kantorovich duality in getting the optimization problem. We will work on this in the future.
3. Since we assume the distributions to be absolutely continuous singularity of the transport map might not be an issue theoretically.
However, this may become an issue when applying our methods to more complicated problem. A preprocessing of the dataset may be needed before applying our methods, for example dequantization can make the discrete dataset continuous.
4. We have added a nonlinear version of applying the constraint techniques (use test function as multiplier) in the modified version, which better utilizes the nonlinearity nature of neural networks.
Strength:
(1) The article has sufficient theoretical derivations and proofs for the three algorithms.
Weaknesses:
(1)  Although it is illustrated that traditional numerical methods are not applicable to high-dimensional problems, relevant comparative experiments and graphs are missing.
(2) Lack of comparative experiments with existing learning-based methods.
In this paper, the authors combine traditional numerical methods with deep learning methods so as to learn the optimal transmission on a high-dimensional benchmark. The theoretical derivation is relatively sound, but the relevant experiments are missing.
NO.
Many thanks for reviewing our paper and providing the valuable comments.  We thank you for appreciating the presentation of our content. We have uploaded a revised version of our paper.
Below is a detailed response to your concern.
We have added some numerical experiments to demonstrate our algorithm. There are already some other work doing the comparison of different optimal transport solvers, for example "Korotin, Alexander, et al. "Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark." Advances in Neural Information Processing Systems 34 (2021): 14593-14605." As for the traditional methods that use PDE to solve optimal transport problems, they work probably only under dimension 3, since in higher dimensions, one needs a grid for computation and the computation cost can be very high.
As you suggested, we also plan to do more numerical experiments to see how our method work in more complex problems. Considering the theoretical results and our benchmark experiments, we expect the method to be also applicable to larger scale problems.
Strength:
- the paper is well-written
- the discussion on the Gaussian example is very helpful in understanding the proposed approach, its advantages and disadvantages
Weakness:
- the paper is motivated by shortcomings of the existing approaches in terms of special neural network structure or regularization, but the proposed approach seems to suffer from the same issues.
- the most challenging step of the approach is how to compute the penalty and perform the ADMM penalty step. It is not clear how to compute the KL divergence or any other type of distance (W1, IPM,..) in general. Minimizing the penalty term is the subject of all GAN papers and the way to do it exactly suffers from the issues that the paper is trying to avoid.
The paper is clearly written and the results seem reproducible.
NO.
Many thanks for your reviewing our paper and that you find value in our work. We appreciate your valuable comments. We have modified the paper based on your suggestions. Below is a detail response to your concern.
1. We donot intend to solve the shortcomings of GANs. Our focus is on how to solve the optimal transport problem with a novel strategy.
2. We have added section 3.4 to describe how the KL divergence and Wasserstein distance are computed. KL divergence is computed using the network of normalizing flow.
Moreover, inspired by your suggestion, we introduce a test function as Lagrangian multiplier and obtain a new way of regularization of GAN.
Instead of use regularization on the discriminator network, we show in figure 6 that the transport cost regularization on the generator network also work and lead to a lower transport cost for the learned generative map.
3. Our methods can apply to different type of neural networks.
Indeed, we use normalizing flow in our implementation with the KL divergence and use GAN for Wasserstein distance. One could use other distance function for probabilities and apply constraint optimization techniques on it.
4. We add Appendix 6 for the implementation details and provide the anonymous github link for the code.
Strength:
- It is interesting to solve the KL regularized OT problem. Combining them together should give us more understanding of the both problems.
Weakness:
- The method can only solve continuous OT problem supported on the whole Euclidean space, or it will encounter severe convergence problem since the KL divergence can easily be $\infty$.
- If the target distribution is discrete like MNIST, how to compute the KL divergence?
- For the Alg 1,2,3, it is unclear how to conduct the optimization, by sampling?
- From the current experimental results, either Fig. 4 or Fig. 5, it is hard to say that the proposed methods work well.
- No implement details.
Clarify: The paper is generally not very clear and hard to follow. It makes a lot of hidden assumptions that need to be proven to be true. There is also no implement details, which make it difficult to evaluate its real performance and hard to reproduce.
NO.
Many thanks for carefully reading our manuscript and providing valuable comments. We have uploaded the modified version of our paper. Below is the detailed response to your concerns.
1. The KL divergence is not necessarily needed in our work, we can use any distance/divergence functional. We add section 3.4, illustrating how to apply the Wasserstein distance in our methods.
2. For discrete dataset such as MNIST, one can apply dequantization technique to transfer the problem to continuous space.
For example, see "https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/tutorial11/NF_image_modeling.htmlhttps://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/tutorial11/NF_image_modeling.html". To compute the KL divergence, we use a neural network of normalizing flow and the KL divergence can be computed, see section 3.4 of the updated manuscript for details.
3. For Alg. 1, 2, 3, the optimization is performed by sampling from the initial and target distributions. The optimization is done using gradient descent. We have add an anonymous github page for the code. We have added appendix A.6 for the implementation details and the github link.
4. We have modified the algorithm about GAN, see section 3.4 and figure 7 for details.
Also the transport cost reduces significantly by using our methods while the target distribution is still learned. Please pay attention to the value of the axis since the range of plot varies from figure to figure.
5. The implement details are added to Appendix A. 6 with the anonymous  github link for the code.
For the theoretical side of the paper, we do make some assumptions.
For example, we assume all measures are absolutely continuous.
While these assumptions may not be always realistic, they allow us to prove the convergence of the algorithms in certain limited situations.
I thank the authors for the response. After carefully reading the rebuttals and other reviewers' comments, I decide to keep my ratings.
Thanks for taking the time in reading our work and providing the comments that help improve our work. We hope our rebuttal version address your concerns.

Meta-review:
role: PC
data: metareview:_summary,_strengths_and_weaknesses: The paper considers  a  constrained optimization method to solve the original Monge's formulation of OT.
Three different algorithms are considered, proposed and partially evaluated. 

Most reviewers agree that the paper needs to be improved in several points (see reviewer's comment) and especially in 
its experimental evaluation parts and its comparison with other approaches for learning Monge map. 

justification_for_why_not_higher_score: * better description of the algorithm and the way to compute and handle the penalty.
* better empirical evaluations (other datasets and comparisons with other Monge map estimation approaches)
justification_for_why_not_lower_score: N/A

Decision: Reject