Conditional Wasserstein Distances with Applications in Bayesian OT Flow Matching

2024-03-27

arXiv.org

AI Summary

In Bayesian inverse problems, inaccurate quantification of posterior distributions hinders reliable uncertainty estimation. Standard joint measure optimization fails to control the Wasserstein distance between conditional posteriors. Method: We introduce the **conditional Wasserstein distance (cWd)**—defined as the Wasserstein distance minimized over a constrained coupling set—and rigorously prove its equivalence to the expected Wasserstein distance under the posterior. We derive its dual formulation, geodesic structure, and velocity field representation, and establish its theoretical connection to the conditional WGAN loss. Leveraging cWd, we propose a conditional generative framework based on optimal transport flow matching, integrating ODE-based modeling with flow matching. Results: Experiments demonstrate significant improvements in both inverse problem reconstruction accuracy and class-conditional image generation quality. The method combines theoretical rigor—rooted in optimal transport—with practical superiority, offering a principled approach to conditional posterior approximation.

Technology Category

• Constraint Satisfaction and Optimization: Distributed CSP/Optimization
• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Experiences and lessons learnt from Web-based algorithms and system deployments

Abstract

In inverse problems, many conditional generative models approximate the posterior measure by minimizing a distance between the joint measure and its learned approximation. While this approach also controls the distance between the posterior measures in the case of the Kullback--Leibler divergence, this is in general not hold true for the Wasserstein distance. In this paper, we introduce a conditional Wasserstein distance via a set of restricted couplings that equals the expected Wasserstein distance of the posteriors. Interestingly, the dual formulation of the conditional Wasserstein-1 flow resembles losses in the conditional Wasserstein GAN literature in a quite natural way. We derive theoretical properties of the conditional Wasserstein distance, characterize the corresponding geodesics and velocity fields as well as the flow ODEs. Subsequently, we propose to approximate the velocity fields by relaxing the conditional Wasserstein distance. Based on this, we propose an extension of OT Flow Matching for solving Bayesian inverse problems and demonstrate its numerical advantages on an inverse problem and class-conditional image generation.

Problem

• Addresses limitations of Wasserstein distance in conditional generative models
• Introduces conditional Wasserstein distance for posterior approximation
• Proposes OT Flow Matching extension for Bayesian inverse problems

Innovation

• Conditional Wasserstein distance via restricted couplings
• Dual formulation resembling conditional Wasserstein GAN losses
• Extended OT Flow Matching for Bayesian inverse problems