[Paper Review] On Robust Optimal Transport: Computational Complexity and Barycenter Computation

Khang Le, Huy Hoang Nguyen+4 more|arXiv (Cornell University)|Feb 13, 2021

Markov Chains and Monte Carlo Methods37 references4 citations

TL;DR

This paper proposes robust optimal transport (ROT) via KL-divergence relaxation of marginal constraints, enabling faster computation through Sinkhorn-based algorithms with complexity Õ(n²/ε). It introduces RobustIBP for robust barycenter computation, achieving improved complexity Õ(mn²/ε) for m=2, outperforming prior methods and enabling efficient, robust generative modeling and domain adaptation under data contamination.

ABSTRACT

We consider robust variants of the standard optimal transport, named robust optimal transport, where marginal constraints are relaxed via Kullback-Leibler divergence. We show that Sinkhorn-based algorithms can approximate the optimal cost of robust optimal transport in $\widetilde{\mathcal{O}}(\frac{n^2}{\varepsilon})$ time, in which $n$ is the number of supports of the probability distributions and $\varepsilon$ is the desired error. Furthermore, we investigate a fixed-support robust barycenter problem between $m$ discrete probability distributions with at most $n$ number of supports and develop an approximating algorithm based on iterative Bregman projections (IBP). For the specific case $m = 2$, we show that this algorithm can approximate the optimal barycenter value in $\widetilde{\mathcal{O}}(\frac{mn^2}{\varepsilon})$ time, thus being better than the previous complexity $\widetilde{\mathcal{O}}(\frac{mn^2}{\varepsilon^2})$ of the IBP algorithm for approximating the Wasserstein barycenter.

Motivation & Objective

• To address the computational inefficiency and lack of robustness in standard optimal transport under data outliers.
• To formalize robust optimal transport by relaxing marginal constraints using KL divergence, improving stability against noisy or corrupted data.
• To develop a scalable algorithm for computing robust barycenters with improved computational complexity over existing methods.
• To demonstrate the practical advantages of robust optimal transport in generative modeling and domain adaptation under data contamination.
• To establish theoretical complexity bounds for both robust OT and robust barycenter computation, showing improvements over prior Sinkhorn and IBP-based approaches.

Proposed method

• Proposes two variants of robust optimal transport: one with partial marginal relaxation and one fully relaxed via KL divergence, enabling robustness to outliers.
• Adapts the Sinkhorn algorithm to solve robust OT formulations, achieving Õ(n²/ε) time complexity, which is faster than standard OT and matches UOT complexity.
• Introduces the RobustIBP algorithm for computing fixed-support robust barycenters, inspired by iterative Bregman projections but adapted for robust formulations.
• Leverages low-rank approximation (Nyström method) to accelerate matrix-vector operations, reducing complexity to Õ(nr² + nr/ε) for rank-r approximation.
• Employs entropic regularization and iterative projection techniques to ensure convergence and numerical stability in barycenter computation.
• Uses Sinkhorn updates with k iterations to approximate robust divergences (e.g., RSD) in generative modeling, enabling end-to-end training with robustness to data corruption.

Experimental results

Research questions

• RQ1 Can robust optimal transport with KL-divergence relaxation be computed more efficiently than standard optimal transport?
• RQ2 What is the computational complexity of Sinkhorn-based algorithms for robust optimal transport, and can it be improved using low-rank approximation?
• RQ3 Can the iterative Bregman projection (IBP) framework be adapted to compute robust barycenters with better complexity than classical IBP?
• RQ4 How does robust optimal transport improve performance in generative modeling under data contamination?
• RQ5 Does the proposed RobustIBP algorithm achieve near-optimal complexity for the robust barycenter problem, especially for m=2?

Key findings

• The Sinkhorn-based algorithm for robust optimal transport achieves Õ(n²/ε) time complexity, which is faster than the standard OT complexity of Õ(n²/ε²) and matches UOT complexity.
• By applying low-rank approximation via the Nyström method, the computational complexity is further reduced to Õ(nr² + nr/ε), where r is the approximated rank.
• For the robust barycenter problem with m=2, the RobustIBP algorithm achieves Õ(mn²/ε) complexity, which is an improvement over the classical IBP complexity of Õ(mn²/ε²).
• Empirical results show that generative models trained with robust optimal transport (e.g., RSD) successfully ignore 10% corrupted data (e.g., uniform noise), while standard OT-based models generate noisy samples.
• On contaminated MNIST data, the robust formulation produces clean digit images, whereas standard Sinkhorn-based training generates mixed noise and digit patterns.
• The proposed RobustIBP is the first practical algorithm to achieve near-optimal complexity Õ(mn²/ε) for robust barycenter computation, even in the minimal case m=2.