Embedding Empirical Distributions for Computing Optimal Transport Maps

TL;DR

A novel approach to learning transport maps for new empirical distributions is introduced, employing the transformer architecture to produce embeddings from distributional data of varying length that are fed into a hypernetwork to generate neural OT maps.

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.

Figures (7)

• Figure 1: The network architectures for direct computation of OT maps (left), and for our proposed method of generating OT maps through hypernetworks (right).
• Figure 2: Visualization of the transported data mapped via various methods. The ground truth data are shown in the top-left corner. For the other figures, the left (resp. right) subplots are generated by applying the forward maps $\hat{T}$ (resp. inverse maps $\hat{T}^{-1}$) to the corresponding ground truth data.
• Figure 3: One-to-one color transfer using HOTET.
• Figure 4: Multiple-to-one color transfer using HOTET. Showing both the training and in-context learning stages.
• Figure 5: Samples generated by the forward and inverse maps in $d=4$, compared with ground truth input distributions.