Neural Conditional Transport Maps

TL;DR

This paper introduces a neural framework for learning conditional optimal transport maps that adaptively process various conditioning variables, enhancing performance in complex probabilistic modeling tasks.

Contribution

The paper proposes a hypernetwork-based neural method for conditional OT maps, improving adaptability and performance over existing simpler conditioning approaches.

Findings

Outperforms baseline methods in ablation studies

Enables high-performance OT-based sensitivity analysis

Advances conditional optimal transport in high-dimensional domains

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.