Efficient Neural Network Approaches for Conditional Optimal Transport with Applications in Bayesian Inference

2023-10-25

arXiv.org

AI Summary

This paper addresses Bayesian conditional inference under likelihood-free settings by proposing two neural network-based approaches for solving Conditional Optimal Transport (COT): (1) a static mapping leveraging partially input-convex neural networks, introducing a novel gradient approximation strategy for efficient training and accurate conditional sampling; and (2) a dynamic flow mapping based on regularized neural ordinary differential equations (ODEs), achieving significantly accelerated sampling while preserving modeling flexibility. Both methods operate within a unified framework capable of conditional density estimation and high-dimensional conditional distribution modeling. Evaluations on benchmark datasets and Bayesian inverse problems demonstrate that the static method attains faster training convergence, whereas the dynamic method delivers superior sampling efficiency and fidelity. Collectively, the proposed approaches substantially outperform existing COT and conditional generative methods in terms of accuracy, scalability, and computational efficiency.

Abstract

We present two neural network approaches that approximate the solutions of static and dynamic $unicode{x1D450}unicode{x1D45C}unicode{x1D45B}unicode{x1D451}unicode{x1D456}unicode{x1D461}unicode{x1D456}unicode{x1D45C}unicode{x1D45B}unicode{x1D44E}unicode{x1D459}unicode{x0020}unicode{x1D45C}unicode{x1D45D}unicode{x1D461}unicode{x1D456}unicode{x1D45A}unicode{x1D44E}unicode{x1D459}unicode{x0020}unicode{x1D461}unicode{x1D45F}unicode{x1D44E}unicode{x1D45B}unicode{x1D460}unicode{x1D45D}unicode{x1D45C}unicode{x1D45F}unicode{x1D461}$ (COT) problems. Both approaches enable conditional sampling and conditional density estimation, which are core tasks in Bayesian inference$unicode{x2013}$particularly in the simulation-based ($unicode{x201C}$likelihood-free$unicode{x201D}$) setting. Our methods represent the target conditional distribution as a transformation of a tractable reference distribution. Obtaining such a transformation, chosen here to be an approximation of the COT map, is computationally challenging even in moderate dimensions. To improve scalability, our numerical algorithms use neural networks to parameterize candidate maps and further exploit the structure of the COT problem. Our static approach approximates the map as the gradient of a partially input-convex neural network. It uses a novel numerical implementation to increase computational efficiency compared to state-of-the-art alternatives. Our dynamic approach approximates the conditional optimal transport via the flow map of a regularized neural ODE; compared to the static approach, it is slower to train but offers more modeling choices and can lead to faster sampling. We demonstrate both algorithms numerically, comparing them with competing state-of-the-art approaches, using benchmark datasets and simulation-based Bayesian inverse problems.

Problem

• Approximating conditional optimal transport solutions
• Enabling conditional sampling and density estimation
• Addressing computational challenges in Bayesian inference

Innovation

• Neural networks approximate conditional optimal transport
• Static method uses gradient of convex neural network
• Dynamic approach employs neural ODE flow map