[Paper Review] Approximating 1-Wasserstein Distance with Trees

Makoto Yamada, Yuki Takezawa+4 more|arXiv (Cornell University)|Jun 24, 2022

Anomaly Detection Techniques and Applications4 citations

TL;DR

This paper proposes a Lasso-based method to learn optimal edge weights in tree-structured embeddings for approximating 1-Wasserstein distance (WD) with linear-time complexity. By formulating tree-based shortest-path distances as a linear model and solving a non-negative Lasso problem, the method achieves highly accurate approximation of exact 1-WD—outperforming baseline trees like QuadTree and ClusterTree in both MAE and correlation, while enabling sparsity via L1 regularization.

ABSTRACT

Wasserstein distance, which measures the discrepancy between distributions, shows efficacy in various types of natural language processing (NLP) and computer vision (CV) applications. One of the challenges in estimating Wasserstein distance is that it is computationally expensive and does not scale well for many distribution comparison tasks. In this paper, we aim to approximate the 1-Wasserstein distance by the tree-Wasserstein distance (TWD), where TWD is a 1-Wasserstein distance with tree-based embedding and can be computed in linear time with respect to the number of nodes on a tree. More specifically, we propose a simple yet efficient L1-regularized approach to learning the weights of the edges in a tree. To this end, we first show that the 1-Wasserstein approximation problem can be formulated as a distance approximation problem using the shortest path distance on a tree. We then show that the shortest path distance can be represented by a linear model and can be formulated as a Lasso-based regression problem. Owing to the convex formulation, we can obtain a globally optimal solution efficiently. Moreover, we propose a tree-sliced variant of these methods. Through experiments, we demonstrated that the weighted TWD can accurately approximate the original 1-Wasserstein distance.

Motivation & Objective

• To address the high computational cost of exact 1-Wasserstein distance estimation in NLP and CV applications.
• To improve the accuracy of tree-based Wasserstein approximations, which often suffer from high approximation error due to suboptimal tree structures.
• To develop a differentiable, convex optimization framework for learning optimal edge weights in trees to better approximate 1-WD.
• To extend the method to a tree-sliced variant for enhanced performance and robustness.
• To demonstrate that L1-regularized weight learning enables both high accuracy and model compression through sparsity.

Proposed method

• Formulate the 1-Wasserstein distance approximation problem as a shortest-path distance approximation on a fixed tree structure.
• Represent the shortest-path distance as a linear model of edge weights, enabling regression-based optimization.
• Formulate edge weight learning as a non-negative Lasso-based optimization problem to ensure convexity and global optimality.
• Use off-the-shelf Lasso solvers (e.g., SPAMS) to efficiently compute globally optimal edge weights.
• Propose a tree-sliced variant by averaging multiple random tree projections, improving correlation without increasing computational cost.
• Apply L1 regularization to prune unimportant edges, enabling compact, sparse tree representations with minimal performance loss.

Experimental results

Research questions

• RQ1 Can a convex optimization framework based on Lasso learning improve the accuracy of tree-based 1-Wasserstein distance approximations?
• RQ2 How does the proposed method compare to standard tree structures like QuadTree and ClusterTree in approximating exact 1-WD?
• RQ3 To what extent does the tree-sliced variant enhance performance, especially in terms of correlation with the ground-truth 1-WD?
• RQ4 Can L1 regularization effectively compress the tree while preserving approximation accuracy?
• RQ5 Does the performance of the method scale favorably with increasing training sample size?

Key findings

• The proposed qTWD and cTWD methods achieved significantly lower mean absolute error (MAE) than QuadTree and ClusterTree across all datasets, with MAE values an order of magnitude smaller than QuadTree.
• For the Amazon dataset, Sliced-cTWD with λ=10⁻¹ achieved a PCC of 0.870 using only 29,540 nodes, outperforming cTWD (PCC=0.785) with 32,642.3 nodes.
• Even with λ=10⁻¹ (10% of nodes), qTWD and cTWD maintained high PCC values comparable to those of full QuadTree and ClusterTree.
• The number of nonzero tree weights was drastically reduced via L1 regularization, demonstrating effective sparsity without performance degradation.
• Increasing the number of training samples improved PCC, with convergence observed at 100,000 samples, while MAE remained low even with smaller training sets.
• The proposed method outperformed both non-sliced and sliced baselines in both MAE and PCC, showing consistent superiority in approximating the exact 1-Wasserstein distance.