Sum-of-Squares Programming for Ma-Trudinger-Wang Regularity of Optimal Transport Maps

Authors: Sachin Shivakumar, Georgiy Antonovich Bondar, Gabriel Khan, Abhishek Halder
Primary area: optimization
Keywords: Optimal transport, sum-of-squares programming, Ma-Trudinger-Wang tensor

Abstract:
For a given ground cost, approximating the Monge optimal transport map that pushes forward a given probability measure onto another has become a staple in several modern machine learning algorithms. The fourth-order Ma-Trudinger-Wang (MTW) tensor associated with this ground cost function provides a notion of curvature in optimal transport. The non-negativity of this tensor plays a crucial role for establishing continuity for the Monge optimal transport map. It is, however, generally difficult to analytically verify this condition for any given ground cost. To expand the class of cost functions for which MTW non-negativity can be verified, we propose a provably correct computational approach which provides certificates of non-negativity for the MTW tensor using Sum-of-Squares (SOS) programming. We further show that our SOS technique can also be used to compute an inner approximation of the region where MTW non-negativity holds. We apply our proposed SOS programming method to several practical ground cost functions to approximate the regions of regularity of their corresponding optimal transport maps.

Decision: reject
Reviewer scores: 6, 6, 6, 5, 5

Reviews and discussion:
<thread>
creation date:
1729125722049

thread_type:
official_review

<thread_object>
creator: reviewer_o9bv

<content>
summary:
This paper presents a computational method for assessing the regularity structure of optimal transport plans. Specifically, the source and target distribution are continuous distributions on a manifold, the cost is some smooth function, and the regularity in question is the regularity of the pushforward map from source to target distributions. In addition, the computational tool computes the region where the MTW holds.

strengths:
1. The paper presents an under-studied aspect of regularity of Monte map. Moreover, the paper contains numerical algorithm and is practical.
2. The inverse problem is well thought through. 
3. The numerics seems promising for a relatively difficult problem in OT.

weaknesses:
1. While it is not an issue for the author per se, it is unfortunate that the SOS condition only applies to semialgebraic sets for the manifold. 

2. The writing doesn't seem to include sufficient focus on the gap between the non-negativity condition and the SOS counterpart. It doesn't seem to be clear whether SOS is too strong for this case. 

3. The presentation is not clear for this paper. For example, the indicting convention for something like $c_{ij, p}$ is quite confusing. The only mention of $c_{ij, kl}$ in the earlier part is too far away, and the readers can not be expected to find where the notation is and also generalize from $c_{ij, kl}$ to $c_{ij, p}$. This is too confusing for this conference.

4. Theorem 5 seems currently wrong: the function $F$ in (5) is matrix-valued. It doesn't seem correct to somehow assess if this matrix-valued rational function belongs to \sum_{SOS}[x, y]. *Unless this issue is either resolved or explained, this reviewer cannot increase the score above the acceptance threshold.*

5. The author doesn't seem to use certain terms in differentiable geometry correctly. For x, y on different points of $\mathcal{M}$, one cannot directly apply an "inner product"/contraction between the tangent plane on $x$ and the contingent plane of $y$. This issue is resolved, however, if $\mathcal{M}$ is a subset of $\R^{n'}$ and the differentiable structure comes from the Euclidean space. The author is advised to change the writing on this and make sure no further major mistake such as this is made.

Minor comments:
- The logic at line 293 is wrong: \eta(\xi) = 0 should belong to after \forall.

questions:
1. The manuscript doesn't contain (A) a discussion on the computational complexity for checking the NNCC condition and the two types of MTW conditions. Also, the discussion on (B) the complexity of the inverse problem is missing. For (A), the author is advised to provide an analysis. For (B), the author is advised to also provide a runtime in the work, i.e. the time it takes to plot the figures.

2. What is the relationship between Monge OT map (which is Borel) and a Brenier map (which is point to point)? Is regularity of Brenier map, supposing it exists, also something that this formulation can answer?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730306365172

thread_type:
official_review

<thread_object>
creator: reviewer_fqtp

<content>
summary:
In the context of OT, the fourth-order Ma-TrudingerWang (MTW) tensor associated with this ground cost function provides a notion
of curvature. The non-negativity of this tensor plays a crucial role for establishing continuity for the Monge optimal transport map. In general, it is difficult to analytically verify this condition for any given ground cost. This paper proposes a provably correct computational approach which provides certificates of non-negativity for the MTW tensor using Sum-of-Squares (SOS) programming. The authors further show that their SOS technique can also be used to compute an inner approximation of the region where MTW non-negativity holds. They apply this proposed SOS programming method to several practical ground cost functions to approximate the regions of regularity of the corresponding OT maps. They also evaluate the proposed SOS computational framework for both the forward and the inverse problems.

strengths:
This is a mathematically solid paper and resolve an interesting theoretical problem. It proposes a provably correct computational framework that can certify or falsify the non-negativity of the MTW tensor associated with a given ground cost under the assumptions that the ground cost is a rational and semialgebraic function. The proposed approach is based on sum-of-squares (SOS) programming and can be of independent interests. The authors also demonstrate that the proposed computational framework can be applied to non-rational ground cost function given that the elements of the MTW tensor are rational and can be used to solve the inverse problem.

weaknesses:
The main concern is that this paper seems irrelevant to the ICLR community. In practice, the common ground cost function would be Euclidean and I am not really sure if it is practically important to conduct the computational verification of OT regularity for a general class of non-Euclidean ground cost functions. I encourage the authors to elaborate on potential ML applications or benefits of their work on non-Euclidean optimal transport. For example, would you like to discuss how your method could enhance practical ML systems that use OT, or to provide concrete examples of where non-Euclidean costs arise in ML problems? 

Another concern is the poor quality of writing. In particular, there are many advanced mathematical notations, such as semialgebraic functions and Archimedean sets, which are not accessible to the ICLR audience. Both Section 2 and Section 3 are written in a technical way without sufficient intuitive explanations or examples alongside the formal mathematical notation. In my humble opinion, the major contribution of the paper would be the SOS formulations for computing the MTW tensors, which is certainly nontrivial, but this paper would much better fit the applied mathematics oriented journal.

questions:
1. I encourage the authors to elaborate on potential ML applications or benefits of their work on non-Euclidean optimal transport. For example, would you like to discuss how your method could enhance practical ML systems that use OT, or to provide concrete examples of where non-Euclidean costs arise in ML problems? 

2. I encourage the authors to improve accessibility of technical parts (e.g., the description of forward problems and inverse problems), such as adding more intuitive explanations or examples alongside the formal mathematical notation.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730484690340

thread_type:
official_review

<thread_object>
creator: reviewer_xnx2

<content>
summary:
In this paper, the authors propose the first computational approach to certify the regularity of the Monge map for optimal transport problems with specific conditions on the transport cost and the state spaces. To be more precise, they evaluate the non-negativity of the fourth-order Ma-Trudinger-Wang (MTW) tensor associated to the transport cost, which has been proved to be a sufficient condition to establish the continuity of the Monge transport map under proper conditions on the marginals of the transport plan. In this work, they consider three versions of this non-negativity condition, previously considered by related works. Their method consists in reformulating the MTW condition (for each of the three versions) into a sum-of-squares program defined on a semialgebraic setvia Putinar's Positivstellansatz, which can then be solved with efficient software. In particular, their approach assumes that the transport cost (or at least the corresponding MTW tensor) is a rational function defined over a two-state semialgebraic space (or at least, a two-state space that contains a semialgebraic space). They apply their framework to verify if the transport cost verifies the MTW condition or to find the largest semialgebraic set on which the transport cost verifies the MTW condition. They propose several convincing numeric experiments in small dimension for a large variety of non-trivial transport costs.

strengths:
- Although I am no expert of SOS programming, the paper is well written so that it can be read by a large audience. In particular, the notation is easy to understand and Section 2 provides the most essential theoretical elements from OT and SOS programming domains to introduce the method.
- Given the elements of Section 2, the idea of proving the MTW non negativity via SOS programming seems to be a very good (and natural) idea. This work seems to be the first to answer this question with relatively moderate theoretical and computational frameworks.
- The formulation of the inverse problem is very interesting and once again well introduced and explained.
- The diversity of numerical experiments (i.e. non trivial transport costs) definitely proves the theoretical statements.

weaknesses:
- As it seems crucial to apply SOS programming, the transformation of the non-negativity condition into a SOS representation in Equation (8) would deserve more explanation, in the appendix for example. For non expert readers, this relation is hard to understand.
- The dimension of the numerical experiments is relatively low, while OT aims at solving large-scale problems.
- Although the problem tackled in this paper is interesting from a theoretical perspective, I am quite concerned by the effective application of this work to OT problems.

questions:
- Could you please provide examples of real-world OT applications where the knowledge of the regularity of the Monge map is crucial ?
- Could you please bring more details on the equivalence between non-negativity on polynomial terms and SOS representation ?
- I think it would be of interest to provide the results on the regularity of the Monge transport map given the three types of non-negativity conditions given in the paper.
- Have you considered experiments with higher dimension ? Is there any computational burden ?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730699935502

thread_type:
official_review

<thread_object>
creator: reviewer_aecx

<content>
summary:
The paper presents a sum-of-squares programming based approach to verifying the Ma-Trudinger-Wang (MTW) condition. Precisely, both the forward problem of identifying if a given cost function and domains satisfy the MTW condition and the inverse problem of finding the largest semialgebraic domain on which the MTW condition holds are considered. The corresponding problems can be solved via standard SOS solvers on modest hardware. The paper concludes with a numerical study which validates the theoretical findings.

strengths:
To my knowledge, this is the first paper which explores the question of numerically verifying the MTW condition. The paper is overall written well, and the theoretical details look correct.

weaknesses:
The limitation of the paper regards the assumption that the cost function is rational or that the elements of the MTW tensor are rational. I believe it would be useful to provide general examples of when it holds/does not hold in the text to further clarify how strong/weak the assumption really is. 

I believe, however, that the implications of this work are not quite fully fleshed out.

questions:
My main question pertains to how the authors see this work fitting within the broader optimal transport literature. In effect, though some results require regularity of the optimal transport map to hold, these results typically pertain to questions of statistical estimation. In these settings, population measures are estimated based on samples and so (i) absolute continuity cannot be verified a priori, (ii) upper and lower bounds for the density cannot be verified, and (iii) the support of the distributions are unknown. It is thus unclear to me how the content of the current paper fits within the previous context.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730700105730

thread_type:
official_review

<thread_object>
creator: reviewer_p2ky

<content>
summary:
This paper presents a computational approach using Sum-of-Squares (SOS) programming to verify and approximate regions of regularity for optimal transport maps, specifically focusing on the Ma-Trudinger-Wang (MTW) tensor. Regularity of optimal transport maps is critical in many machine learning applications, and this regularity can be assured by the non-negativity of the MTW tensor. However, verifying this condition analytically for general cost functions is challenging. The authors propose using SOS programming to generate certificates of non-negativity for the MTW tensor across a broader range of cost functions, potentially providing computationally verified regions of regularity. Their method is applied to both verifying non-negativity conditions (the "forward problem") and to computing inner approximations of regularity regions (the "inverse problem") for several ground cost functions, demonstrating the effectiveness of SOS programming in this context. This computational framework contributes a systematic approach to certifying regularity in optimal transport, potentially facilitating its application in various machine learning tasks. The paper concludes by applying the proposed framework to several common examples in the literature.

strengths:
The paper’s strength lies in its innovative application of Sum-of-Squares (SOS) programming to address the longstanding challenge of verifying the regularity of optimal transport (OT) maps through the Ma-Trudinger-Wang (MTW) tensor. SOS programming is a well-established tool in optimization and control, but this work extends it to OT regularity, opening new possibilities for computational verification of the MTW conditions in general cases where analytic approaches are intractable. The paper also showcases the practical efficacy of the approach by applying it to various cost functions, demonstrating its flexibility and adaptability across different scenarios.

weaknesses:
A missing key aspect in the paper is the time complexity analysis for the proposed framework. What's the computational efficiency of SOS programming in verifying regularity of the different OT problems? While the authors showcase the method’s application to specific examples and shared the wall-clock time, a time complexity discussion could be a good addition to the paper.

questions:
1. What are some other regularity verification methods? how does the SOS programming compare to them in terms of accuracy and efficiency?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732215114174

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of improvements

comment:
We are grateful to the reviewers for their careful reading of our work and for the perceptive inputs provided.

Here is a summary of the major improvements we have made to address the reviewers' comments. For the reviewers' convenience, all revisions in the manuscript are marked in blue.

- **(Accessible writing with pedagogical examples)** Based on the reviewers' comments, we have included a *new Appendix A titled "Nonnegative Polynomials and Sum-of-Squares Programming"* (cited in Sec. 2.2) that explains the SOS polynomial and its decomposition (Sec. A.1), the matrix-valued SOS polynomial and its decomposition (Sec. A.2), and the SOS polynomials and Archimedean semialgebraic sets (Sec. A.3). In addition to definitions and explanations, this includes several examples to illustrate the connections between different ideas (e.g., gap between SOS and nonnegative polynomial). The Sec. A.3 is specifically purposed to explain the equivalence between (7) and (8) for Archimedean semialgebraic sets. In Sec. 2.2, an example is included right after the definition of semialgebraic set. We believe these additions will help the readers who may not be familiar with some of the background ideas.

- **(Computational complexity)** Following the suggestions from multiple reviewers, we performed a detailed computational complexity analysis for the NNCC/MTW forward problems in terms of the dimension $n$ and the number of semialgebraic constraints $\ell$. They scale polynomially in $n$ and sub-quadratic in $\ell$. The analysis is included in *new Appendix D titled "Computational Complexity"* (cited in the end of Sec. 3.1) where Sec. D.1 and D.2 detail the complexity analysis for the NNCC and the MTW forward problems. respectively. In the individual responses to the reviewers, we explain that these results are for off-the-shelf SOS/SDP solvers that do not exploit our problem structure, and thus further speed-ups should be possible by customized solvers. Per suggestion of Reviewer o9bV, the runtimes for the inverse problems are now reported at the end of the numerical examples 3 and 4. Here too, most of the computation overhead was found to be in problem parsing and SDP setup to deploy off-the-shelf solvers. These results are very encouraging for a first computational work on OT regularity.

- **(Improved explanations)** In Sec. 1 (Introduction), we slightly rephrased the 3rd paragraph, included a new paragraph after assumption **A1** to explain why this assumption is benign. The latter also mentions with citation that one of the early driving factors for OT regularity theory was the engineering problem of reflector antenna design which was reformulated as an OT problem with non-Euclidean ground cost $c(x,y)=-\log\|x-y\|$ over the sphere. In Sec. 2.1, paragraph after Definition 2, we clarified some differential geometric issues raised by Reviewer o9bV.

- **(Notational improvement and simplifications)** We fixed the matrix-SOS notation throughout to distinguish it from scalar SOS notation, clarified that as well as the index notation in the revised Table 1. Some additional notations are also included in that Table. Some formulas were simplified in Proposition 13 within Appendix B.

- **(Better positioning the work w.r.t. OT literature)** In addition to mentioning the antenna design problem, we added a new paragraph at the end of Sec. 2.1 explaining why certifying NNCC or MTW condition are of interest in designing algorithms for solving unregularized OT problems with general costs, i.e., such certifications are of interest beyond the regularity of OT map. We wrote a detailed response to Reviewer fQTp on the growing relevance of solving OT problems with non-Euclidean ground costs in ML applications. *We also included the Appendix F listing the OT with non-Euclidean ground cost examples with references, pointing out that most of these are amenable to the proposed SOS framework*.

**Scope and fit for ICLR:**

While SOS programming is well-known in ML and optimization, the novelty this work is its application to automatically certify/falsify OT regularity, and computationally discover regions where local regularity holds. This is an existing problem: its need is recognized in both ML and OT theory literature, but existing approaches rely on unwieldy analytical calculations hand-crafted for very few costs. As a result, these calculations and the related techniques in the existing literature do not generalize. Being the first work on computational verification of OT regularity, this work opens door for a new research direction in computational OT, especially for non-Euclidean OT that is increasingly finding applications in ML including significant ICLR footprints in the recent years. We believe this work--its topic, style, novelty and significance of results--are well within the scope of ICLR.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper presents a computational approach using Sum-of-Squares (SOS) programming to certify the Ma-Trudinger-Wang regularity condition for optimal transport (OT) maps. While technically sound and offering an interesting contribution, the reviewers exposed several weaknesses. A significant concern, particularly emphasized in the final discussions, is the lack of clear exemplification of why this work is relevant for machine learning.. The paper does not convincingly demonstrate how the proposed method provides new insights into ML problems where OT is used or enhances practical ML applications. Moreover, reviewers raised concerns about the high computational complexity of the approach, its limited scalability to higher dimensions, and its focus on semialgebraic cost functions, which may restrict practical applicability. Additional weaknesses include the technical presentation, which lacks accessibility for a broader ML audience, and unclear positioning within the OT and ML literature. Despite the efforts made in revisions, these issues were not sufficiently addressed, and the paper fails to bridge the gap between theoretical contributions and practical relevance.