Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The results in this paper apply to the computation of optimal transport maps for the Wasserstein-2 distance in Euclidean space. Under weak assumptions, these maps are given by gradients of convex functions. Moreover, the convex function and its Fenchel-Légendre conjungation are a solution to the dual problem. The paper presents different frameworks for performing the conjugation via a differentiable optimization problem. Moreover, the author also provides an efficient and paralelizable "Armijo-type" algorithm for performing the conjugation. Experiments suggest that this approach leads to significant improvements in the computation of optimal transport maps. 

 
strength_and_weaknesses: *Strenghts*

- The experimental results look promising. 
- The paralellized Armijo is a cool idea. 

*Weaknesses*

- I found the writing a bit difficult for a non-expert (eg. terms like "amortization" are never defined, and is only really explained in the code in the appendix).  
- There's no theory to back up the methods.  
- Some parts of notation are not well defined (eg. the letter for the dimension changes from n to D in the text).
- It is not clear to me how the method should be tuned in the abscence of a ground truth. 

clarity,_quality,_novelty_and_reproducibility: The paper is mostly *clear*, although, as mentioned above, some parts of the text are unclear. (An additional point: eq. (2) is missing a square in the RHS.)

*Reproducibility* is partial/laborious: the appendix contains written code for the examples, but, to my knowledge, there's no way to download and run code directly. 

As far as I can tell, the present paper is *novel*: the first to try and compute conjugate maps for optimal transport via "fast methods".

*Quality*: the main evidence that the method is good is the experiments presented, which are adequate. However, I find them a bit incomplete, in that I'd like to see a "real task" where OT is typically used and the present method leads to improvements.  Also, there's the tuning issue mentioned above. 
 
summary_of_the_review: The paper presents novel methods for the conjugation step in dual Wasserstein-2 optimal transport. Experiments suggest that the methods are good, but the paper is a bit hard to read, and leaves some questions regarding tuning unresolved.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: Monge's primal Wasserstein-$2$ transport map optimization problem admits a (Kantorovich) dual with appealing properties which enable the use of several learning/optimization algorithms to (approximately) solve the problem. However, the dual objective function includes a term involving the convex conjugate of the potential function (which is the dual optimization variable), which is hard to compute in many settings, and even impossible in some others. The author proposes taking an "amortized optimization" view for the conjugate computation problem within the iterations of estimating the potential function (e.g., while approximating it with NNs). The amortized optimization view is shown to be useful since the potential function estimation stage will involve solving the conjugate sequentially/iteratively.

**Note:** My score increased from "6" to "8" after the rebuttal.
strength_and_weaknesses: **Strength:** Firstly, I would like to thank the author for such a clear motivation, literature review, numerical experiments, implementation, and discussions. The paper is very "fun" to read, and I had a good time reading it. The focus on computing the conjugate is very essential to the relevant community. 

**Weaknesses:** 
*I list my major and minor concerns that hold me from being *strictly* positive about this paper. However, I still think the paper has some interesting ideas, therefore recommending an acceptance. I am also going to stay active during the discussion period in case the author has updates or questions regarding my review.*

***Major Weaknesses/Questions***
- Most of the remarks and optimization-related discussions are very redundant and they do not add much to the paper. For example, in Remark 8 it is discussed that the first-order conditions would not give a global solution for nonconvex functions. Do we even need this? I do not think any reader who wouldn't already know this will understand even page 1 of the paper.
- The tone of the author may be taken as subjective/informal or even offensive for some people. There are several **bold highlighted** sentences that are subjective and these may make the paper look like some course notes or informal discussion. Remark 7 is a good example. Also, Remark 6 "in my experiments, updating .. hurts ..." is not formal.
- Some results are not very new, but the way they are presented may sound like a breakthrough from the tone. I would like to kindly ask the author whether Section 4.1 adds anything new. Isn't this just a parallel line search implementation?
- I would be happier to see further optimization formality. Validity of $\min$ instead of $\inf$ is typically not argued. The existence of feasible solutions for, for example, duality arguments, is not discussed.
- Some problems are claimed to be "easy", for example, "**computing the exact conjugate is easy**". Is it really an easy problem? Or is that a numerical observation?
- Sometimes there are overly positive claims for simple/standard evidence. For example, in Section 3.1, the selection of a gradient for the amortization of a function is said to be well-motivated because the $\arg\min$ of the conjugate is also a derivative.
- Finally, a general question to the author: The OT setting looks relevant because we need to know the conjugate and because we would apply gradient-based learning to solve approximations of (2), but in general, this solution technique is a way of estimating the conjugate of a potential function. I am wondering, is there anything else that makes the underlying OT setting unique? Moreover, amortized optimization is already being used a lot in similar settings, may I please ask the author what makes this setting special? Apologies if I am not seeing something obvious, and again, many thanks for the work.
 
***Minor Weaknesses/Questions***
- Abstract has "an amortized approximations". 
- Whenever $:= \arg\min$ notation is used, could the author please discuss that the $\arg\min$ set is a singleton? Otherwise, $\in \arg\min$ would perhaps be more suitable.
- "solving" the conjugation sounds a little confusing in Section 1, especially since $f$ is an optimization function. Maybe the context behind the parametrization of $f$ should be discussed first.
- Page 2: "the the exact conjugate" (typo)
- Equation (4): $f_\theta^\star(x)$ has a typo I believe. 
- Equation (4): $J_{f_\theta} (x(y))$ -> this shorthand is not explained yet (shorthand for $J_{f_\theta} (x(y); y)$
- When Algorithm is first mentioned in Section 1, the "initial dual potential $f_\theta$" is not explained yet (which appears in the algorithm). Similarly, the same algorithm uses $\tilde{x}_{\varphi}(y_j)$ which is not yet discussed. Also, the same algorithm, maybe $N$ can be shown as an input, too?
- Remark 1: "$W_2$ coupling" is not abbreviated yet. Section 3.1: "MLP" not abbreviated.
- Remark 2: Perhaps the author should discuss the feasibility $x \in \mathcal{X}$, too (in order to make a weak-duality sort of argument)?
- Last paragraph of Section 2: "by predicting the solution with a model $\tilde{x}_{\varphi}(y_j)$". Solution of "what" is not clear, and $\tilde{x}_{\varphi}(y_j)$ is not introduced.
- Section 3.1, element 1: "directly maps to the solution by" -> please also state "solution of what"
- Section 3.2.1: "as optimal as possible" -> this is not a very usual terminology
- I am curious: how does the approach mentioned right after (6) relate to the standard SAA-like techniques in Stochastic Optimization?
- Page 4: ".. et al ... proposes" -> 'propose'
- Remark 7: "... state that they are not a maximin optimization procedure" -> this is not clear to me
- Question: Could the author please discuss how someone would still be interested in 3.2.3 even when the true solution to the conjugate problem is known? (for example, when we use NN based approximations for the potential?)
- Page 12, Nhan Dam et al. (2019) citation -> please capitalize "GAN".
- Conclusions: The sentence that starts with "In non-Euclidean ...." is hard to grasp.
clarity,_quality,_novelty_and_reproducibility: For the details of these points, please see my review of "Strength and Weakness". However, to give a high-level answer:

**Clarity:** The overall idea is very clear. Optimization terminology is sometimes vague.

**Quality:** The paper has high quality for sure. The experiments, visualizations, and results are thorough.

**Novelty:** The perspective of looking at the conjugate estimation problem as an amortized optimization problem is novel. However, I am not sure how useful it is, as there is not much theoretically appealing property rather than experimental numerical results. This is just **a** way of solving this problem.

**Reproducibility:** The paper provides all the source codes, and the main algorithm is very clear. Hence I cannot foresee any issues with reproducibility.
summary_of_the_review: The paper is in general written well with a very accurate literature review, the topic studied is very interesting with high relevance to both optimal transport and optimization communities. Some of the results that are discussed in detail are already well-known (if not they are easy to derive as there is no new theory, but rather a combination of known methods), the paper is written in a subjective/informal way, and in general, the paper looks like a review paper rather than a novel one. I recommend "marginally above acceptance threshold", but I do not have strong positive feelings either. I decided in favor of acceptance as the paper provides a complete set of reusable algorithms, which the community can benefit from.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper studies efficient method for computing convex conjugate arising in wasserstein OT problem. Although exactly and approximately computing so is believed to be hard in prior work, this work proposes a new method based on amortized approximation scheme could be used to computing the exact conjugate. This amortization with fine-tuning also has favorable practical performance, improving over all Wasserstein-2 benchmark tasks, and produce stable results on synthetic dataset.
strength_and_weaknesses: This paper studies an important problem in OT and provides interesting idea and compelling practical implementation of the proposed method. The experiment part looks solid and convincing. The theoretical part feels less convincing for the following reasons:
- If not mistaken, it feels the key idea of amortizing largely follows the amortizing optimization framework in Amos [2022] and other prior work, and the main part of Section 3 is mostly restating this framework in the particular setting of OT. Most new designs tailored to the setting are discussed in Section 4, but most of it is also still using well known optimizer like Adam and L-BFGS to directly solve the Conjugate function. In that case, the contribution of the theoretical part of the paper reads unclear to me.
- Though I am unsure possible or not, it may be helpful to have at least some convergence guarantee, stability analysis, or even hard instance which explicitly shows why *all methods* need to be sensitive to model's hyper-parameters.
- There are many remarks in the paper, some of them seem a bit vague and in particular, require a very good knowledge on previous work to be able to understand, e.g. Remark 6, Remark 11. It may be helpful to discuss their formulations first and then distinction to this work in detail in the main text.
clarity,_quality,_novelty_and_reproducibility: The paper is structured nicely. The presentation and writing of the paper could be improved a little bit to facilitate readers not very familiar with the area, and save sometime from memorizing all the equations and settings in different prior works while reading the paper. The paper has done a lot of comparison with prior works all over the paper and maybe doing so more systematically will help the presentation too.

In terms of novelty, the idea is interesting yet the theoretical contribution over prior works is not very clear (see comments above).
summary_of_the_review: I am unfamiliar with this area and especially for the computational perspective for 2-wasserstein OT. The paper reads like trying to address an important problem, and the practical performance is quite compelling. My slight concerns are regarding the novelty of the algorithmic idea, lack of support in theory, and a potential for a better presentation. I'd love to hear from the author's perspective before giving a final assessment from my side.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper suggests fine-tuning a learnable conjugate map when solving the Wasserstein-2 distance dual formulation (where the potential map is parameterized using a neural network) can lead to significant improvement of the resulting transport map as demonstrated on the Wasserstein-2 benchmark by Korotin et al. (2021a). To facilitate the optimization during fine-tuning, a new parallel line search for L-BFGS is implemented in Jax which is much faster than Jax's default line search method.

strength_and_weaknesses: ## Strengths:
* The problem of computing the conjugate map effectively is a central one in neural OT.
* The proposed fine-tuning step is simple in concept/implementation yet very effective.
* Careful software engineering is done to obtain good experimental results.
* Lots of good empirical advice from the remarks in the paper for OT practitioners.


## Weaknesses:
* The novelty of the proposed method is quite limited. There is no new methodological formulation or new theoretical results. To me, it seems mostly like a nice engineering hack (which works quite well).
* The novelty of the new line search method is also limited. The proposed line search improvement is a straightforward way to convert a classical sequential algorithm involving conditional branches to a batched version.
* In a couple of the remarks there are some new perspectives (e.g. Remark 5, 7) of viewing existing formulation in different ways, but I do not find these new perspectives reveal particularly interesting new insights.
* The paper does not reveal a new understanding of why certain combinations of parts of the algorithms work better than others. Some heuristics are provided but could be better explained (e.g. Remark 8 and Figure 1).

## Detailed comments:
* When citing Brenier's theorem below (3), it should be mentioned that it only holds when one distribution has a density (which I believe is the assumption for this paper and most of the related works).
* The message of Remark 8 and Figure 1 is not clear to me. Does it just mean solving (9) is easier than non-convex (6) and (8)?
* In Table 2, ICNN performs very poorly compared to MLP, even with fine-tuning. What could be an explanation for this, other than ICNN is just hard to train as commonly understood? I figure fine-tuning using regression loss should make ICNN easier to train.


clarity,_quality,_novelty_and_reproducibility: The clarity of the paper is very good and sufficient code snippets and algorithm boxes were provided for reproducibility. I personally like the style of using bold text. For novelty, see the comments above.

summary_of_the_review: Overall I think this is a good paper, despite not having a lot of novelty as mentioned above. The experiments are carefully done. OT practitioners can all learn something useful from reading this paper.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: This work proposes a novel method for computation of convex conjugates in optimal transport problems with Euclidean costs.

The reviewers and I agree that this is a very good submission that presents new, interesting ideas, in a well-exposed manner.
justification_for_why_not_higher_score: see meta-review - this is an interesting submission, but with a somewhat narrow scope - still a very 'cool' idea, as written by one of the reviewers.
justification_for_why_not_lower_score: see meta-review
Decision: Accept: poster