Conditional simulation via entropic optimal transport: Toward non-parametric estimation of conditional Brenier maps

TL;DR

This paper introduces a scalable, non-parametric estimator for conditional Brenier maps using entropic optimal transport, providing theoretical guarantees and demonstrating superior performance on benchmarks and Bayesian tasks.

Contribution

It develops the first entropic optimal transport-based estimator for conditional Brenier maps with statistical and algorithmic guarantees.

Findings

Estimator converges to true conditional Brenier maps in Gaussian setting.

Outperforms existing machine learning and non-parametric methods on benchmarks.

Provides heuristic for choosing the entropic regularization parameter.

Abstract

Conditional simulation is a fundamental task in statistical modeling: Generate samples from the conditionals given finitely many data points from a joint distribution. One promising approach is to construct conditional Brenier maps, where the components of the map pushforward a reference distribution to conditionals of the target. While many estimators exist, few, if any, come with statistical or algorithmic guarantees. To this end, we propose a non-parametric estimator for conditional Brenier maps based on the computational scalability of \emph{entropic} optimal transport. Our estimator leverages a result of Carlier et al. (2010), which shows that optimal transport maps under a rescaled quadratic cost asymptotically converge to conditional Brenier maps; our estimator is precisely the entropic analogues of these converging maps. We provide heuristic justifications for choosing the scaling parameter in the cost as a function of the number of samples by fully characterizing the Gaussian setting. We conclude by comparing the performance of the estimator to other machine learning and non-parametric approaches on benchmark datasets and Bayesian inference problems.