Tree-Sliced Entropy Partial Transport

Viet-Hoang Tran, Thanh Tran, Thanh Chu, Tam Le, Tan Minh Nguyen

Keywords: tree-sliced wasserstein distance, optimal transport, partial transport, unbalanced optimal transport

TL;DR: We propose a tree-sliced framework to the Partial Transport setting by observing that the Partial Transport problem on tree-metric spaces can be reformulated as a standard Optimal Transport problem.

Abstract: Optimal Transport (OT) has emerged as a fundamental tool in machine learning for comparing probability distributions in a geometrically meaningful manner. However, a key limitation of classical OT is its requirement that the source and target distributions have equal total mass, limiting its use in real-world settings involving imbalanced data, noise, outliers, or structural inconsistencies. Partial Transport (PT) addresses this limitation by allowing only a fraction of the mass to be transported, offering greater flexibility and robustness. Nonetheless, similar to OT, PT remains computationally expensive, as it typically involves solving large-scale linear programs—especially in high-dimensional spaces. To alleviate this computational burden, several emerging works have introduced the Tree-Sliced Wasserstein (TSW) distance, which projects distributions onto tree-metric spaces where OT problems admit closed-form solutions. Building on this line of research, we propose a novel framework that extends the tree-sliced approach to the PT setting, introducing the Partial Tree-Sliced Wasserstein (PartialTSW) distance. Our method is based on the key observation that, within tree-metric space, the PT problem can be equivalently reformulated as a standard balanced OT problem between suitably modified measures. This reformulation enables efficient computation while preserving the adaptability and robustness of partial transport. Our method proves effective across challenging tasks such as outlier removal and addressing class imbalance in image-to-image translation. Our code is publicly available at [this https URL](https://github.com/thanhqt2002/PartialTSW).

Primary Area: General machine learning (supervised, unsupervised, online, active, etc.)

Submission Number: 10563