Low-Rank Optimal Transport through Factor Relaxation with Latent Coupling

Summary

As the paper clearly outlines in the introduction, Optimal Transportation (OT) is widely used in various fields of machine learning, however, the cost for computing OT is computationaly expensive (quadratic scaling), even after Sinkhorn algorithm employing entropy regularization significantly alleviated the cost. The purpose of Low-rank Optimal Transport (LOT) is to reduce the complexity by constraining the transportation plan to possess a certain low-rank structure. The key contribution of the paper is (1) new parametrization of LOT using a Latent Coupling (LC) factorization, (2) The eponymous "Balance FRLC" algorithm for resolving the optimization procedure for computing LOT with respect to various objectives associated with (LOT), (3) Theoretical guarantees for the smoothness of the regularized objectives, as well as upper bounds for the convergence criteria. The algorithm is applied to illustrative toy examples and then to a more realistic example (spatial transcriptomics alignment) and the results indicate improvements over previous LOT.

Strengths

The contributions of the paper is significant. The writing is excellent. The LC factorization, adapted to the LOT framework from a previous work, is well-motivated. The Balanced FRLC algorithm appears a sound and non-trivial extension of previous approaches for LOT (Sinkhorn, semi-relaxed OT). The algorithm is particularly simple and shown to be effective. The experimental results are clear and illustrative, improvements are convincing.

Weaknesses

What this reviewer found missing is discussion of the case the optimal transport map is itself not low-rank. It is well-known that the optimal transportation plan is not necessarily low-rank (as indicated e.g. by their connection to differential equations with non-regular solutions), so LOT may converge very slowly to the full-rank optimal transportation plan. In section 4.2, the authors provide a simple illustrative example, but this is when the transportation plan is known to be low-rank - in this case the transportation plan seems to stabilize after some threshold rank - however, what happens when the full-rank plan does not have an inherently low-rank structure?

Questions

- line 103: Would it be good to say the what complexity "massively improved" to? - Can the authors comment on how to pick the rank for the LOT? (This of course ties in to the point above the rank of the transportation plan.) - Perhaps a more discussion of of how different the LC factorization is to, say the previous factored coupling with (the diagonal version $P = Q \text{diag}(1/g) R^T$). This relation seems important when considering "more" LC factorizations: If this LC factorization improves the result while simultaneously allowing a simple optimization algorithm, would it then be natural to discuss further factorizations that are nested this way, with additional LC factorizations for $T$ and so on?

Rating

9

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

This work builds on the LOT framework - so perhaps the authors comment on some scenarios where the LC factorization can potentially yield no better results, when compared to previous LOT?

Summary

This work introduces a new low rank formulation of optimal transport based on the latent coupling decomposition introduced in https://arxiv.org/pdf/2012.11589 . Compared to previous formulations of low rank OT, this formulation allows easier extensions to unbalanced and Gromov Wasserstein settings.

Strengths

I found that the presented FRLC method is a quite elegant approach to low rank OT. It decouples optimization into entropy regularized semi relaxed OT problems for $Q$ and $R$, and $T$ is updated solving entropy regularized OT of size $r$. This decoupling is enabled by the form of the inner coupling which is no longer assumed to be diagonal. This construction alleviates the need to use the Dykstra machinery as in (Scetbon 2021). With the proposed approach, $Q$ and $R$ are not required to have the same inner marginal. From what I understand, the optimization over inner marginals is handled by the semi relaxed steps and does not require a separate step as in previous approaches. Another strength of the approach is that it generalizes smoothly to the unbalanced setting. Other strengths include: * I carefully looked at the presented method and the appendix and the algorithmic approach made sense. * non asymptoptic convergence bounds are derived for the method. * the appendix includes a complete review of the necessary technical background presented in a clear way, easy to follow. * limitations are clearly stated (section N in appendix). * a new initialization scheme for LC decomposition is proposed. * there is a gain in interpretability compared to (Scetbon 2021) as the inner coupling captures the coupling between clusters in the data.

Weaknesses

* from section 3.1 to 3.2, the new aspects of the approach compared to https://arxiv.org/pdf/2012.11589 should be stated more clearly. The contributions of the paper should be separated more clearly from previous works. I had to read the other paper to better understand each contribution. * In algorithm 1, it would be interesting to have the computational complexity of each step (presented in the algorithm directly).

Questions

* How to choose $r$ in practical scenarios ? * Would it make sense to remove the entropic penalties (at the expense of smoothness) to recover results about hard clustering of the points via the decomposition of the OT coupling ? Put differently, what is the effect of entropic smoothing on the interpretability of the coupling in terms of clustering ?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

/

Summary

The paper introduces a novel framework called Factor Relaxation with Latent Coupling (FRLC) which is based on coordinate mirror descent to compute the low-rank LC factorization.The algorithm decouples the optimization into three sub-problems, offering greater flexibility, interpretability, and linear space complexity. FRLC is applicable to various OT objectives (Wasserstein, Gromov-Wasserstein, Fused Gromov-Wasserstein) and marginal constraints (balanced, unbalanced, and semi-relaxed). Theoretical results support its effectiveness, and empirical tests demonstrate superior performance in applications like graph clustering and spatial transcriptomics.

Strengths

The introduction of the LC factorization and its integration into the FRLC algorithm represents an innovation in reducing the complexity of optimal transport problems. The FRLC framework is versatile, handling multiple OT objectives and marginal constraints, making it applicable to a wide range of practical problems. The latent coupling provides a interpretable description of the transport plan, which is beneficial for understanding and visualizing the results.

Weaknesses

While the empirical results are good, the paper could benefit from more extensive comparisons with additional baseline methods to solidify the claims of superiority. The theoretical results are limited in that, even though $\Delta_k(x_k, x_{k+1})$ is small, it cannot be guaranteed that the iteration points converge to the optimal solution or even a stationary point.

Questions

I suggest that Propositions E.3, E.5, E.6 should be stated in the main text to make the theoretical part easier to follow. Table 1 should be put in the main text to show the efficiency of FRLC in terms of runtime. It is hard to say that the method shows significant improvement because the runtime of LOT is already very short (less than one second). I suggest the authors conduct experiments on more challenging and diverse examples.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

N.A.

Summary

The paper presents an approach for low-rank optimal transport (OT) leveraging a latent coupling (LC) factorization and solving it with mirror descent. This approach offers a new parameterization of the low-rank OT problem, providing advantages such as decoupling the problem into three OT problems and enhancing flexibility and interpretability. The authors introduce the Factor Relaxation with Latent Coupling (FRLC) algorithm, which utilizes coordinate mirror descent to compute the LC factorization. FRLC accommodates various OT objectives and marginal constraints. The paper includes theoretical results and demonstrates the performance on selected tasks like graph clustering and spatial transcriptomics.

Strengths

- Generality: handles different OT objective costs and relaxations of the marginal constraints. - Significance: the proposed method, and improving OT using factorizations, has the potential to address the scalability issue of OT with large datasets, making it highly relevant for applications in machine learning and data science.

Weaknesses

**Presentation and logical organization of the ideas** I found it difficult to discern the paper's particular contribution, especially concerning previous art like Latent Optimal Transport (Lin et al., 2021), Forrow et al. (2019), and Scetbon et al. (2021). For example, the main contribution is said to be "compute a minimum cost low-rank transport plan using the LC factorization," where LC factorization is precisely the factorizations proposed by Lin et al. (2021), with the further contribution being the computation of this factorization using mirror-descent. I wondered how much of the description also applies to Lin et al. (2021) and which part is particular to the novel method. Furthermore, according to Lin et al. (2021), there is an algorithm for computing low-rank factorized plans based on factorized costs using a projection method and an algorithm to factorize the costs and select anchor points without the need to have a similar number of anchor points at source and target latent space. How does your approach compare with the approach? **Comparison and empirical validation** The first empirical comparison analyzes LOT and FRLC, capturing how both methods behave as the factorization rank changes. This is a good way of understanding the effect of the rank on each case. Nevertheless, in the second setup, the comparison with LOT is not present anymore, which begs the question of how LOT performs in the second empirical evaluation. Furthermore, an extensive empirical study with the structured approach of Lin et al. (2021) would help clarify the particular contributions of the type of factorization and factoring algorithm used. I found some other issues lacking in a more in-depth discussion of the computational complexity, possible implementation difficulties, and overall limitations in the proposal.

Questions

- Is the main contribution here actually the mirror descent formulation of the OT problem with the factorization with a different set of latent source and target points proposed by Lin et al. (2021)? - Furthermore, is the main difference between the fact that Lin et al. (2021) use a two-stage process and the novel mirror-descent method solves the OT problem in a more straightforward unified process? - What are the theoretical complexity of the algorithms, and what practical considerations are relevant when implementing the method? One presentation improvement that would change my evaluation is a summary table locating the differences in the type of factorization, whether the cost or transport is being factorized, the algorithm used for computing the factorization, etc.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

I found it necessary to discuss in more depth the limitations and the tradeoffs present in the proposal compared with previous art.