[Paper Review] Tree-Wasserstein Barycenter for Large-Scale Multilevel Clustering and Scalable Bayes

Tam Le, Viet Huynh+3 more|arXiv (Cornell University)|Oct 10, 2019

Anomaly Detection Techniques and Applications42 references4 citations

TL;DR

This paper proposes the tree-Wasserstein barycenter, a novel approach that leverages tree-structured ground metrics to efficiently compute Wasserstein barycenters for large-scale multilevel clustering and scalable Bayesian inference. By exploiting the tree metric's structure, the method enables fast, memory-efficient computation via $ε$-regularized optimization and tree-based projections, achieving up to 50% faster runtime than Sinkhorn-based baselines while maintaining high accuracy on high-dimensional data with thousands of supports.

ABSTRACT

We study in this paper a variant of Wasserstein barycenter problem, which we refer to as tree-Wasserstein barycenter, by leveraging a specific class of ground metrics, namely tree metrics, for Wasserstein distance. Drawing on the tree structure, we propose an efficient algorithmic approach to solve the tree-Wasserstein barycenter and its variants. The proposed approach is not only fast for computation but also efficient for memory usage. Exploiting the tree-Wasserstein barycenter and its variants, we scale up multi-level clustering and scalable Bayes, especially for large-scale applications where the number of supports in probability measures is large. Empirically, we test our proposed approach against other baselines on large-scale synthetic and real datasets.

Motivation & Objective

• To address the scalability limitations of traditional Wasserstein barycenter computation in high-dimensional, large-support settings.
• To develop an efficient algorithmic framework for computing Wasserstein barycenters using tree metrics as ground distances.
• To scale up multilevel clustering and scalable Bayesian inference for large-scale datasets with many probability measure supports.
• To demonstrate superior computational efficiency and accuracy compared to existing baselines like Sinkhorn and sliced Wasserstein barycenters.

Proposed method

• Uses tree metrics as ground distances in optimal transport, enabling closed-form solutions and fast computation via tree-structured projections.
• Applies entropic regularization to the tree-Wasserstein barycenter problem, allowing efficient optimization using Sinkhorn-like algorithms on tree-mapped probability measures.
• Employs clustering-based tree metric construction with farthest-point sampling to build hierarchical tree structures that preserve data geometry.
• Leverages the center-of-mass property on trees to relax constraints on the number of supports in the barycenter, enabling bounded-support variants.
• Uses weighted samples from the barycenter to approximate posterior distributions in scalable Bayesian inference.
• Employs no-U-turn sampling (NUTS) for posterior sampling and compares barycenter-based approximations across multiple tree configurations.

Experimental results

Research questions

• RQ1 Can tree-structured ground metrics enable faster and more memory-efficient computation of Wasserstein barycenters than standard OT or sliced OT?
• RQ2 How does the tree-Wasserstein barycenter perform in large-scale multilevel clustering with high-dimensional data and many supports?
• RQ3 To what extent can the tree-Wasserstein barycenter improve scalability and accuracy in Bayesian inference with large posterior samples?
• RQ4 How sensitive is the approximation quality to the number of trees and tree depth in the metric construction?

Key findings

• The tree-Wasserstein barycenter achieves less than half the running time of the Sinkhorn-based baseline when estimating posteriors with only 30 samples.
• With 100 trees, the method reduces absolute error in estimating the random effects covariance matrix to 0.67425, outperforming the Sinkhorn baseline (0.751236) and sliced Wasserstein (0.678854 at k=30).
• Increasing the number of trees improves estimation confidence, reducing error means and standard deviations in covariance matrix estimation.
• The method maintains high accuracy even with large numbers of supports (up to 3,000), where Sinkhorn and sliced baselines show performance degradation or stagnation.
• The proposed algorithm is robust to high-dimensional data and preserves structural properties of the original support distributions better than one-dimensional projection methods like sliced Wasserstein.
• The tree-Wasserstein barycenter enables scalable multilevel clustering and Bayesian inference without requiring explicit specification of the number of supports in the barycenter.