Unbalanced Low-rank Optimal Transport Solvers

Summary

The authors focus on the problem of efficiently compute discrete Optimal Transport (OT) problems, which are known to have a cubic complexity w.r.t. the number of input samples. They propose to approximate it by assuming that the transport plan is a low-rank matrix and propagate this property in the computations, such that the complexity of matrix products is reduced. This idea has already been proposed to approximate Balanced OT and Gromov-Wasserstein (GW) problems. The authors extend this to variants called Unbalanced OT, as well as Unbalanced GW and Fused-Unbalanced GW. The authors introduce these unbalanced OT variants, they derive in each setting the associated low-rank optimization algorithms using mirror descent. Then they perform experiments on brain and cell biology data, where unbalanced OT was extremely powerful, to assess the performance of these algorithms in an applied setting.

Strengths

- The contributions might look incremental at first (combining unbalanced OT with low-rank OT), but are actually very thorough. They do not restrict to this one combination, but consider all recent variants of unbalanced OT (Unbalanced OT, GW and Fused GW) which have been developped in the literature in the last years. They also leverage one recent optimization idea of 'translation-invariance' which improves the computational efficiency of computing Unbalanced OT problems. All in all, the authors aggregate several works altogether to propose a set of interesting algorithms and tools for practitioners, especially for the biology field.

Weaknesses

- The low-rank approximation is not sufficiently discussed in this paper. I mentioned that it is considered to accelerate the computations. However, one might use this variant because it makes sense to have a prior of low-rank transport plan in their applications, because the data has some structural property which could be leveraged. To my mind, using a low-rank plan seems contradictory with the property that the optimal plan is a (full rank) permutation in some setting. Could the authors discuss this in detail ? Why does it make sense to have a low-rank plan ? Which kind of data is relevant with such assumption ? This is unclear to me. - Some experimental illustrations could be provided on the methodological side of their algorithms, for completeness. The authors propose some optimization problems and algorithms to compute them. To my mind, the authors do provide complexity per iteration of their algorithms, but the experimental aspect of measuring the time performance of their algorithms is omitted. In particular, have you checked experimentally that this 'translation-invariance' variant in Section (3.2) does accelerate the computations in your low-rank setting ? Could you provide a plot of the time to compute OT problems (standard or Sinkhorn vs. low-rank) as a function of the number of samples ? Also, the use of unbalanced OT involves two extra parameters $(\tau_1, \tau_2)$ which requires cross-validation. Could the authors give plots of the performance of their biology task as a function of $(\tau_1, \tau_2)$. This would provide interesting insights on the interpretation of these parameters, and their impact on learning tasks.

Questions

See my question above.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

They adressed the societal impact of their work.

Summary

In this paper, the authors combine two variants of optimal transport (OT) known as unbalanced OT and low-rank OT. While the former relaxes the marginal constraints to ease the modelling rigidities and discard possible outliers from input measures, the latter helps reduce the computational cost of $\mathcal{O}(n^3)$, and therefore, makes OT scalable. In addition to standard OT which quantifies the discrepancy between two measures in the same dimensional space, they also apply this combination to the scenarios when two input measures belong to distinct dimensional spaces, which are Gromov-Wasserstein and Fused-Gromov-Wasserstein. Finally, the authors choose the spatial transcriptomic matching problems to justify the practical usage of their proposed methods.

Strengths

- Originality: This work is a novel combination of two versatile and scalable variants of optimal transport (OT), which are unbalanced OT and low-rank OT. - Quality: The derivations of algorithm proposed in this work are associated with theoretical proof. The empirical performance of the proposed method is demonstrated via the spatial transcriptomic matching problems. - Significance: the results in this paper are important to some extent as it allows practitioners to find a low-rank solvers for the problem of quantifying the discrepancy between two arbitrary (not necessarily probability) measures.

Weaknesses

- Originality: Although the combination of unbalanced OT and low-rank OT is novel, key tools and techniques (e.g. reparametrization of low-rank couplings, Dykstra algorithm) used in this paper have been already introduced in [1] and [2]. Thus, I think this work is incremental to some extent. To address this concern, I suggest the authors should highlight the main challenges of solving unbalanced low-rank OT compared to its balanced counterpart more clearly. - Clarity: The presentation of this paper is not good as there are a lot of notations which are either not carefully defined or define inaccurately (see Requested Changes). Additionally, there are some important parts, namely the initialization of Algorithm 1, which should be presented in this paper rather than merely refer to another paper. **References** [1] Meyer Scetbon, Marco Cuturi, and Gabriel Peyré. Low-rank sinkhorn factorization. In International Conference on Machine Learning, pages 9344–9354. PMLR, 2021. [2] Meyer Scetbon, Gabriel Peyré, and Marco Cuturi. Linear-time Gromov-Wasserstein distances using low rank couplings and costs. ICML, 2022.

Questions

1. In line 95, what is the definition of nonnegative rank of $P$? Please add this definition to the revision of this paper. 2. In equation (2), the term $Q\text{diag}(g)R$ seems to be incorrect. Should it be $Q\text{diag}(1/g)R^{\top}$? 3. In equation (7), are there any differences between $\mathrm{KL}(\cdot,\cdot)$ and $\mathrm{KL}(\cdot | \cdot)$? If any, they should be defined explicitly in the paper. 4. Are there any theoretical guarantees that the tuples $(Q_{k+1},R_{k+1},g_{k+1})$ approximate the solution of the optimization problem in equation (6)? 5. In the proposed algorithms for solving ULOT, ULGW and ULFGW, which values of the rank hyperparameter $r$ should we choose? Would it be as low as possible? **Requested Changes**: 1. References: In lines 38 and 39, when mentioning the usage of Sinkhorn algorithm in solving OT, the authors should cite more relevant papers, namely [1]. Similarly, in line 78, regarding solving unbalanced OT using entropic regularization, the authors should cite the paper [2]. 2. In the definition of cost matrix $C$ in Section 2, the index notation $1\leq i,j\leq n,m$ is inacurrate. It should be changed to $1\leq i\leq n$ and $1\leq j\leq m$. 3. Notations: When introducing the unbalanced OT in equation (3), the authors should explain the notations $\mathrm{KL}$, $\tau_1$ and $\tau_2$ rather than assume that readers implicitly understand. Analogously, the notation $A^{\odot2}$ in equation (4) should be defined explicitly. 4. Typo: in line 96, repamatrization --> reparametrization. **References** [1] Khang Le, Huy Nguyen, Quang M Nguyen, Tung Pham, Hung Bui, and Nhat Ho. On robust optimal transport: Computational complexity and barycenter computation. Advances in Neural Information Processing Systems, 2021. [2] K. Pham, K. Le, N. Ho, T. Pham, and H. Bui. On unbalanced optimal transport: An analysis of sinkhorn algorithm. In ICML, 2020.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The limitations are not discussed in this paper.

Summary

The authors proposed unbalanced low-rank OT (ULOT) solver, which is an extension of the balanced counterpart. They show how to adapt this solver to the other unbalanced low-rank settings, namely translation-invariant and GW. The experiments compare the performances of multiple low-rank solvers, and of the Unbalanced Fused GW.

Strengths

The extensions from the balanced to unbalanced LOT, as well as from ULOT to ULFGW are not trivial and require some calculation effort. I also find that the writing is very instructive and all technical details are clearly presented and on point.

Weaknesses

It seems to me that - The content in the main paper and appendix is not adequately partitioned, namely the section 3 is somewhat too long that there is few space left for the experiment section, and some experiment details in the Appendix should be put in the main. - The main contribution of the paper is quite incremental. More precisely, + While adding translation-invariant improves the convergence of the usual unbalanced OT, it is unclear about the real improvement in the experiments. IMHO, it is more or less an add-on of the ULOT solver and should be moved to Appendix because it would dilute the main message of the paper (which is about ULOT). It is enough that Algo 4 and 5 use the ULR-Dykstra solver (Algo 6), instead of Algo 3. + The Algo 4 is merely a special case of Algo 5, where alpha = 0, thus should be removed. While it is more instructive to start with the GW setting before moving to the fused GW one, all the details can be moved to the Appendix and section 3.3 and 3.4 should be restructured or/and merged. + It seems that the experiment section is not sufficiently diverse because the experiments and competing methods are restricted to just a family of the low-rank based methods (except FUGW). IMHO, since the contribution on the methodogoly (i.e. section 3) is not very significant, I would love to see more comparison with other methods, like (unbalanced) mini-batch OT, or quantized GW, and maybe on more experiments (e.g. on graphs). + I find it weird that the section 4.3 is unreasonably brief and not sufficiently discussed, while most of its details are moved to the Appendix. This makes the section 4.3 look incomplete in the main paper.

Questions

In Algorithm 2: typo in 1 / gamma / tau1

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The authors do not discuss the limitations of their work.

Summary

In its original exact formulation, optimal transport (OT) suffers from an $\mathcal{O}(n^3)$ cost when applied to clouds of $n$ points. In addition, the exact constraint on the marginals makes it sensitive to outliers. One solution for each of these problems exist: - unbalanced OT [Schiebinger et al., 2019] is less sensitive to outliers than regular OT, as it relaxes the marginal constraints into a KL-penalized form - entropic OT allows for the use of the celebrated Sinkhorn algorithm but still requires large matrix vector multiplications. Low rank OT [Forrow et al., 2018] is a promising direction to further improve upon entropic OT, by using a low rank cost matrix and low rank constraints of the plan. Building on a these two seminal papers and a line of work by Scetbon, Cuturi and coauthors, the paper combines low rank OT with unbalanced OT (Formulation in Eq 5). A dedicated algorithm is proposed, interpreted as proximal gradient descent in the KL geometry. The latter is improved based on the works of Séjourné et al [2022]. The approach is extended to the Gromov-Wasserstein and Fused Gromov-Wasserstein OT problems.

Strengths

OT has found many applications in practice in the last decade; alleviating its cost extends its range of applications. The proposed formulation is sound, the treatment is relatively easy to follow and the derivations are made clear.

Weaknesses

- Although the paper heavily sells low rank optimal transport, isn't the formulation in eqs 5/6 non convex? If so, doesn't it make the exact solution intractable? How to ensure convergence towards a good critical point? - Convergence of the algorithm - Known results about Dykstra's algorithm are used to compute the solution of Equation 7 via its dual. Can the authors provide a reference for the convergence of the whole algorithm, namely the iterations defined by 7? I may have missed it, but I could not find reference of convergence towards a critical point; - convergence seems to be measured in terms of difference between two successive iterates going to zero. It is well-known that this does not guarantee convergence of the iterates (e.g. take the harmonic series). To clarify, do the author have a classical convergence result?

Questions

- In the last cost of eq(2) I believe it should be diag(1/g) not diag(g), and $R^T$ not $R$ Typos: - reparamaterized, repamatrization - mass conversation for mass conservation - misuse of \cite or \citep vs \citet , e.g. "of [Frogner et al., 2015, Chizat et al., 2018]", and several other places e.g. "notations from [Scetbon et al., 2021]" - towards a stationary points - such that with probability at least 0.99 that: extra "that" - Algo 1: (check me) - paper uses \star and * for $\lambda^\star$

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

not applicable

Response to the authors

I thank the authors for their response. It helps clarify most of my concern. I do agree that the two suggested competing methods (minibatch OT and quantized GW) don't really fit in the context of the paper. However, comparing to the prior work on LOT and low-rank GW, I would still consider the contribution of ULOT (and its extensions) to be incremental. I see it as a mere stack of many prior works (low rank, unbalanced, GW) and there has not yet been any theoritical analysis of the proposed method in the paper. Even if it is possible that some results in the prior works can be applied, I'm not convinced that it is enough for a real contribution, except something really stands out. For this reason, I expect to see more contribution on the application side to compensate for the lack of theoritical understanding. In conclusion, I decide to increase the score to 4, while still being more prone to the reject.

Response

Dear Authors, I thank you for your rebuttal, and for your additional experiment comparing the gain of translation-invariant Sinkhorn into your algorithm. While I understand the concerns of other reviewers, I tend to disagree with their focus on the 'incrementality' of your work. To my mind, this is compensated by the thoroughness of the UOT variants you compute, which are as many tools available for future works of the ML community and beyond. You could have tackled each formulation in a different paper to trade quality of work for quantity of publications, which you did not. I could slightly increase my score to support your work, but I'm not sure it would change the global consensus among all reviewers. Side remark: In your answer to fjNH, you say that unbalanced OT solvers are non-convex. It is a convex problem (like balanced OT and unlinke GW), thus all unbalanced OT solvers are also convex. The non-convexity should only arise from your low-rank asumption.