Embedding Empirical Distributions for Computing Optimal Transport Maps

Optimal Transport (OT) has evolved as a fundamental tool for measuring distributional differences, with applications spanning image processing, generative models, and domain adaptation. Classical OT algorithms like Monge and Kantorovich formulations laid the groundwork, while recent deep learning approaches, including ICNN-based neural OT, have improved flexibility. However, these methods typically handle only single source-target pairs, and their computational complexity scales poorly with multiple distributions. The rise of deep neural architectures has enabled approximate solutions, but the challenge remains to efficiently learn mappings across many distributions simultaneously, especially in high-dimensional settings. Existing solutions like CONDOT and MetaOT attempt to address this but are limited by pairwise training and scalability issues.

The core challenge is to develop a scalable, generalizable framework capable of learning OT maps from multiple empirical distributions to a common target efficiently. Traditional methods require training separate neural networks for each pair, leading to high computational costs and poor scalability. Moreover, existing models struggle to embed variable-length, unordered samples into a fixed representation that supports downstream learning. Addressing these issues is critical for applications requiring rapid adaptation to new distributions, such as real-time image editing, transfer learning, and large-scale data alignment. The difficulty lies in designing an embedding that captures the essence of empirical samples while maintaining permutation invariance and compatibility with neural potentials.

This work introduces three key innovations: 1) embedding empirical distributions via Transformer models, leveraging their permutation invariance and ability to handle variable-length data; 2) parameterizing multiple ICNN potentials with hypernetworks conditioned on these embeddings, enabling multi-source mapping with shared parameters; 3) employing a dual optimization framework that incorporates target distribution information during training, improving generalization to unseen distributions. These innovations collectively allow for a unified, scalable approach to multi-distribution OT, reducing the need for pairwise training and enabling rapid prediction of new mappings.

• �� Input: multiple empirical samples from source distributions; • Transformer module processes each sample set, removing positional encoding, to produce a fixed-dimensional embedding vector; • Hypernetworks F and G generate parameters for ICNN potentials based on these embeddings; • ICNNs approximate the convex potentials for OT, enabling map computation; • During training, the model uses the target distribution to compute loss via a base OT solver, updating hypernetworks and embedding modules; • For new distributions, the trained hypernetworks predict OT maps directly from sample embeddings, supporting fast inference.

The authors evaluate HOTET on Gaussian mixture datasets, comparing against MetaOT and MM-B baselines using W2B metrics. They train models with embedding dimension 64, ICNN layers 3, and transformer heads 4. Experiments include predicting OT maps for unseen distributions, color transfer tasks, and ablation studies on embedding modules. Metrics such as L2-UVP and cosine similarity quantify accuracy. Results show HOTET achieves lower errors and higher similarity scores, demonstrating superior generalization and efficiency. Visualizations of color transfer confirm practical viability. Hyperparameters are tuned via grid search, and robustness is verified through multiple runs.

HOTET’s L2-UVP scores are consistently below 3.5% on Gaussian mixtures, outperforming baseline methods. The cosine similarity between predicted and true OT maps exceeds 0.91, indicating high accuracy. In high-dimensional settings, the model maintains errors under 3%, significantly better than existing approaches. Color transfer experiments show that HOTET can adapt to new images with minimal fine-tuning, producing visually coherent results. These findings validate the framework’s scalability, accuracy, and practical relevance across diverse tasks.

HOTET can be directly applied to multi-image style transfer, domain adaptation in computer vision, and generative modeling where distribution alignment is crucial. Its ability to quickly predict OT maps from samples makes it suitable for real-time applications, such as video editing, augmented reality, and large-scale data integration. The framework’s flexibility supports extension to other modalities, including text and audio, facilitating cross-domain transfer tasks. Its efficiency enables deployment in resource-constrained environments, broadening the impact of deep OT in industry.

Despite promising results, HOTET faces challenges in many-to-many scenarios with numerous distributions, where parameter sharing may limit expressiveness. The reliance on the base OT solver’s accuracy means that poor solver performance can degrade overall results. Handling extremely high-dimensional data or highly imbalanced samples remains difficult, requiring further robustness enhancements. Additionally, the current approach assumes the availability of sufficient sample data for embedding, which may not hold in sparse data regimes. Future work should address these limitations to improve scalability and robustness.