Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

On Barycenter Computation: Semi-Unbalanced Optimal Transport-based Method on Gaussians

Ngoc-Hai Nguyen, Dung Le, Hoang-Phi Nguyen, Tung Pham, Nhat Ho

Abstract

We explore a robust version of the barycenter problem among $n$ centered Gaussian probability measures, termed Semi-Unbalanced Optimal Transport (SUOT)-based Barycenter, wherein the barycenter remains fixed while the others are relaxed using Kullback-Leibler divergence. We develop optimization algorithms on Bures-Wasserstein manifold, named the Exact Geodesic Gradient Descent and Hybrid Gradient Descent algorithms. While the Exact Geodesic Gradient Descent method is based on computing the exact closed form of the first-order derivative of the objective function of the barycenter along a geodesic on the Bures manifold, the Hybrid Gradient Descent method utilizes optimizer components when solving the SUOT problem to replace outlier measures before applying the Riemannian Gradient Descent. We establish the theoretical convergence guarantees for both methods and demonstrate that the Exact Geodesic Gradient Descent algorithm attains a dimension-free convergence rate. Finally, we conduct experiments to compare the normal Wasserstein Barycenter with ours and perform an ablation study.

On Barycenter Computation: Semi-Unbalanced Optimal Transport-based Method on Gaussians

Abstract

We explore a robust version of the barycenter problem among centered Gaussian probability measures, termed Semi-Unbalanced Optimal Transport (SUOT)-based Barycenter, wherein the barycenter remains fixed while the others are relaxed using Kullback-Leibler divergence. We develop optimization algorithms on Bures-Wasserstein manifold, named the Exact Geodesic Gradient Descent and Hybrid Gradient Descent algorithms. While the Exact Geodesic Gradient Descent method is based on computing the exact closed form of the first-order derivative of the objective function of the barycenter along a geodesic on the Bures manifold, the Hybrid Gradient Descent method utilizes optimizer components when solving the SUOT problem to replace outlier measures before applying the Riemannian Gradient Descent. We establish the theoretical convergence guarantees for both methods and demonstrate that the Exact Geodesic Gradient Descent algorithm attains a dimension-free convergence rate. Finally, we conduct experiments to compare the normal Wasserstein Barycenter with ours and perform an ablation study.

Paper Structure

This paper contains 23 sections, 18 theorems, 204 equations, 5 figures, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Wasserstein distance
  4. Unbalanced Optimal Transport
  5. Riemannian Optimization with Gradient Descent
  6. Semi-Unbalanced Optimal Transport-based Barycenter
  7. Studies On Centered Gaussians
  8. Semi-Unbalanced Optimal Transport has a Closed-Form Expression for Gaussians
  9. Exact Geodesic Gradient Descent for SUOT-based Barycenter
  10. Hybrid Algorithm for Barycenter UOT
  11. Numerical Experiments
  12. Related work
  13. Conclusion
  14. Proof for Theorem \ref{['theo: bary is gaussian']}
  15. Proofs for Theorem \ref{['UOT Plan']} and Theorem \ref{['UOT Entropic']}
  16. ...and 8 more sections

Key Result

Theorem 1

Let $(\alpha_i)_{i=1}^n$ be zero-mean Gaussian distributions in $\mathbb{R}^d$ having covariance matrices $(\Sigma_i)_{i=1}^n$. Let $\Sigma_\beta \in \mathbb{S}_{++}^{d}$ is a SPD matrix. Consider the problem SUOT-based Barycenter in uot framework where $\mathcal{P}(\Sigma_\beta)$ be the set of zero-mean probability distribution in $\mathbb{R}^d$ having covariance matrix $\Sigma_\beta$. Then, $\be

Figures (5)

  • Figure 1: Overview of hybrid updated iteration
  • Figure 2: 2D Contour Plot Gaussian Mixture Distribution. From left to right: two Gaussians with their barycenter (blue); noise is added to one Gaussian on the left; normal Wasserstein Barycenter (green); SUOT-based Barycenter (red).
  • Figure 3: Loss on $L(\Sigma_\beta)$ through iterations with different step sizes. (Top): GD. red: Hybrid Gradient Descent, blue: Exact Geodesic Gradient Descent, orange: Auto Geodesic Gradient Descent, green: Auto Geodesic Gradient Descent with Momentum. (Bottom): SGD. red: Hybrid Stochastic Gradient Descent, blue: Exact Geodesic Stochastic Gradient Descent, orange: Auto Stochastic Geodesic Gradient Descent, green: Auto Stochastic Geodesic Gradient Descent with Momentum.
  • Figure 4: Ablation study of the dependence between the SUOT-based Barycenter and parameter $\tau$. From top to bottom, left to right, we calculate the barycenter with values of $\tau$ as $0.005, 0.01, 0.02, 0.05, 0.1, 0.5, 1, 5, 10, 50, 100$. The bottom right corner subfigure is normal Wasserstein Barycenter in Figure 2 of the main manuscript
  • Figure 5: Bures-Wasserstein Distance from SUOT-based Barycenters to the barycenter learnt by chewi2020gradient

Theorems & Definitions (36)

  • Theorem 1
  • Theorem 2
  • Theorem 3
  • Theorem 4
  • Theorem 5
  • Theorem 6
  • Proposition 1
  • proof
  • proof
  • Lemma 1
  • ...and 26 more