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

TL;DR

The 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; the estimator is precisely the entropic analogues of these converging maps.

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.

Figures (4)

• Figure 1: We observe that $(\mathsf{T2})$ asymptotically converges with a rate of $\mathcal{O}(t^2)$ convergence rate with randomly generated covariance matrices of block-type.
• Figure 2: MSE of the estimated map for the NN and EOT estimators for a multivariate Gaussian problem with increasing sample size $n$.
• Figure 3: Left: Joint samples of $\mu(x_1,x_2)$ with slices at the conditioning values of interest $x_1 \in \{-0.5,3\}$. Right: Generated samples from the EOT, NN, and MLP maps with the true density $\mu_{2|1}(\cdot|x_1)$ in black.
• Figure 4: Comparison of samples from the posterior distribution using the entropic estimator (left) and an adaptive MCMC sampler (right).

Theorems & Definitions (16)

• Theorem 1.1: Convergence to conditional Brenier maps
• Theorem 2.1: Brenier's theorem for rescaled quadratic costs
• Proposition 4.1: Closed-form expressions
• Theorem 4.2
• Remark 4.3
• Proposition 4.4
• Proposition 4.5
• Proposition 4.6
• Theorem 4.7
• proof : Proof of \ref{['prop:closedform']}