Embedding Empirical Distributions for Computing Optimal Transport Maps

2025-04-24

AI Summary

Existing neural optimal transport (OT) methods are limited to modeling single distribution pairs, exhibiting poor generalization across multiple empirical distributions. Method: We propose the first generalized neural OT framework capable of learning OT mappings for arbitrary unseen empirical distributions. Our approach introduces a Transformer-based encoder to embed variable-length discrete distributions and designs a neural hypernetwork that dynamically generates distribution-pair-specific OT mappings—thereby departing from conventional pairwise modeling paradigms. Contribution/Results: By end-to-end optimizing the Wasserstein distance, our method achieves high-accuracy, low-latency OT mapping generation across diverse synthetic and real-world benchmarks, significantly outperforming state-of-the-art baselines. The implementation is publicly available.

Technology Category

• Search and Optimization: Learning to Search
• Computer Vision: Learning & Optimization for CV
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Experiences and lessons learnt from Web-based algorithms and system deployments
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Distributional data have become increasingly prominent in modern signal processing, highlighting the necessity of computing optimal transport (OT) maps across multiple probability distributions. Nevertheless, recent studies on neural OT methods predominantly focused on the efficient computation of a single map between two distributions. To address this challenge, we introduce a novel approach to learning transport maps for new empirical distributions. Specifically, we employ the transformer architecture to produce embeddings from distributional data of varying length; these embeddings are then fed into a hypernetwork to generate neural OT maps. Various numerical experiments were conducted to validate the embeddings and the generated OT maps. The model implementation and the code are provided on https://github.com/jiangmingchen/HOTET.

Problem

• Computing optimal transport maps for multiple distributions
• Learning transport maps for new empirical distributions
• Generating neural OT maps using transformer embeddings

Innovation

• Transformer embeds varying-length distributional data
• Hypernetwork generates neural optimal transport maps
• Validated through diverse numerical experiments