Neural Conditional Transport Maps

TL;DR

This work tackles learning conditional optimal transport maps between distributions under auxiliary conditioning variables, enabling adaptive transport in high-dimensional problems. It extends Neural Optimal Transport (NOT) by introducing a hypernetwork-based conditioning mechanism that generates condition-specific transport parameters, and uses discrete embeddings plus continuous positional encodings within an encoder–decoder framework for the transport map T and the critic f. The approach is validated through extensive ablations and two real-world applications: climate-economy distribution emulation and OT-based global sensitivity analysis in Integrated Assessment Models (IAMs), demonstrating scalability and accuracy with a single conditioned model. The authors also provide open-source code and data, discuss limitations, and outline future directions including multi-modal conditioning and connections to diffusion-like models.</text>

Abstract

We present a neural framework for learning conditional optimal transport (OT) maps between probability distributions. Our approach introduces a conditioning mechanism capable of processing both categorical and continuous conditioning variables simultaneously. At the core of our method lies a hypernetwork that generates transport layer parameters based on these inputs, creating adaptive mappings that outperform simpler conditioning methods. Comprehensive ablation studies demonstrate the superior performance of our method over baseline configurations. Furthermore, we showcase an application to global sensitivity analysis, offering high performance in computing OT-based sensitivity indices. This work advances the state-of-the-art in conditional optimal transport, enabling broader application of optimal transport principles to complex, high-dimensional domains such as generative modeling and black-box model explainability.

Figures (8)

• Figure 1: A diagram of our conditional OT map framework. Our transport network $T$ takes as input samples $x\sim \mathbb{P}$ and transports them to $\mathbb{Q}$, conditioned by $c$. Optionally, $T$ can receive additional noise inputs $z \sim \mathbb{S}$, from a known probability distribution $\mathbb{S}$, introducing stochasticity.
• Figure 2: Architecture of the conditional networks (transport $T$ and critic $f$). The encoder processes input data ($[x,z]$ for $T$ or $x$ for $f$) through residual blocks to produce latents $h_L \in \mathbb{R}^d$. We combine discrete variable embeddings $\mathbf{e}_k$ with positional encoding of continuous values $PE(p)$ through an MLP to produce the unified conditioning vector $\mathbf{c} \in \mathbb{R}^{d_c}$. The conditioning module $\mathcal{C}$ transforms this latent into $h' \in \mathbb{R}^{d'}$, which the decoder processes into the final output. The noise $z$ is only used in $T$.
• Figure 3: Results of our ablation study comparing different model configurations on an unconditional transport setting. The table presents quantitative metrics, while the right panel shows the visualization.
• Figure 4: Results of our climate damages model, under different SSP scenarios, on a particular country the dataset. We show the ground truth distribution of damages and samples from our model.
• Figure 5: Comparison of simplex and our neural transport across three variables for the IAM dataset. The first three panels show costs across partition values for klogistic, gamma, and kw_2 variables (simplex in blue, neural in red). The rightmost panel shows average costs for each variable.