Hierarchical Optimal Transport for Multimodal Distribution Alignment

Abstract

In many machine learning applications, it is necessary to meaningfully aggregate, through alignment, different but related datasets. Optimal transport (OT)-based approaches pose alignment as a divergence minimization problem: the aim is to transform a source dataset to match a target dataset using the Wasserstein distance as a divergence measure. We introduce a hierarchical formulation of OT which leverages clustered structure in data to improve alignment in noisy, ambiguous, or multimodal settings. To solve this numerically, we propose a distributed ADMM algorithm that also exploits the Sinkhorn distance, thus it has an efficient computational complexity that scales quadratically with the size of the largest cluster. When the transformation between two datasets is unitary, we provide performance guarantees that describe when and how well aligned cluster correspondences can be recovered with our formulation, as well as provide worst-case dataset geometry for such a strategy. We apply this method to synthetic datasets that model data as mixtures of low-rank Gaussians and study the impact that different geometric properties of the data have on alignment. Next, we applied our approach to a neural decoding application where the goal is to predict movement directions and instantaneous velocities from populations of neurons in the macaque primary motor cortex. Our results demonstrate that when clustered structure exists in datasets, and is consistent across trials or time points, a hierarchical alignment strategy that leverages such structure can provide significant improvements in cross-domain alignment.

Cited by in corpus (11)

• Neural Latents Benchmark '21: Evaluating latent variable models of neural population activity
• Minibatch optimal transport distances; analysis and applications
• Learning Autoencoders with Relational Regularization
• Statistical Optimal Transport posed as Learning Kernel Embedding
• Making transport more robust and interpretable by moving data through a small number of anchor points
• G2DA: Geometry-Guided Dual-Alignment Learning for RGB-Infrared Person Re-Identification
• On Transportation of Mini-batches: A Hierarchical Approach
• Hierarchical Optimal Transport for Robust Multi-View Learning
• Variational Wasserstein Barycenters for Geometric Clustering
• CO-Optimal Transport
• An Information-Geometric Distance on the Space of Tasks