Computing Optimal Transport Maps and Wasserstein Barycenters Using Conditional Normalizing Flows

2025-05-28

AI Summary

This work addresses scalability and stability challenges in computing optimal transport maps and Wasserstein barycenters in high-dimensional spaces. We propose an end-to-end framework based on Conditional Normalizing Flows (CNFs), which jointly maps multiple source distributions into a shared latent space via invertible implicit transport—directly modeling the transport maps and enabling closed-form barycenter computation, thereby circumventing conventional dual optimization and adversarial training. To our knowledge, this is the first method to employ CNFs for joint learning of multi-source transport maps, supporting efficient barycenter estimation across up to hundreds of input distributions. Experiments demonstrate that our approach significantly outperforms existing state-of-the-art methods on high-dimensional tasks, achieving substantial improvements in accuracy, computational efficiency, and training stability.

Technology Category

• Machine Learning: Learning with Manifolds
• Computer Vision: Learning & Optimization for CV
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

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

Abstract

We present a novel method for efficiently computing optimal transport maps and Wasserstein barycenters in high-dimensional spaces. Our approach uses conditional normalizing flows to approximate the input distributions as invertible pushforward transformations from a common latent space. This makes it possible to directly solve the primal problem using gradient-based minimization of the transport cost, unlike previous methods that rely on dual formulations and complex adversarial optimization. We show how this approach can be extended to compute Wasserstein barycenters by solving a conditional variance minimization problem. A key advantage of our conditional architecture is that it enables the computation of barycenters for hundreds of input distributions, which was computationally infeasible with previous methods. Our numerical experiments illustrate that our approach yields accurate results across various high-dimensional tasks and compares favorably with previous state-of-the-art methods.

Problem

• Efficiently compute optimal transport maps in high dimensions
• Calculate Wasserstein barycenters for hundreds of distributions
• Avoid complex adversarial optimization via gradient-based minimization

Innovation

• Uses conditional normalizing flows
• Solves primal problem via gradient minimization
• Computes barycenters for hundreds of distributions