Research map / Probabilistic and scientific modeling / Variational and generative sampling
Computing Optimal Transport Maps and Wasserstein Barycenters Using Conditional Normalizing Flows
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Computes optimal transport maps and Wasserstein-2 barycenters by training conditional normalizing flows (conditional Real NVP/Glow) that map a shared Gaussian latent to each input distribution, solving the primal problem with gradient descent instead of adversarial dual training.
Problem
Computing optimal transport maps and Wasserstein-2 barycenters for continuous, high-dimensional distributions, where existing methods rely on dual formulations with complex adversarial (bi-level or tri-level) optimization and scale poorly with the number of input distributions.
Approach
Parametrize each input distribution as the pushforward of a common Gaussian latent via a conditional normalizing flow, enforce pushforward constraints by likelihood maximization, and minimize the L2 transport cost in latent space with decreasing weights; barycenters become an expected conditional variance minimization solved by the same conditional flow.
Key results
- On location-scatter Gaussian data, BW2^2-UVP of 0.004-0.094% across d=2-128 versus 0.01-0.38% for WIN and 0.07-1.28% upper-bound LUVP style comparisons, with better or equal training time than SCW2B and WIN.
- Fair regression on Communities and Crime reduces Pearson correlation between the predicted variable and the sensitive feature from 0.43 to 0.0003, using a barycenter of 100 input distributions.
- Barycenter computation over 128 input distributions in d=64 keeps performance roughly constant (LUVP 0.073-0.119%).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use conditional normalizing flows instead of SCW2B | Location-scatter Gaussian barycenters, d=128 · Location-scatter Gaussian barycenters, d=128 | LUVP: 1.280 -> 1.55 ± 0.002 | Hurt |
| Use conditional normalizing flows instead of WIN | Location-scatter Gaussian barycenters, d=128 · Location-scatter Gaussian barycenters, d=128 | Bures-Wasserstein UVP: 0.380 -> 0.094 ± 0.001 | Helped |
| Use conditional normalizing flows instead of SCW2B | Location-scatter uniform barycenters, d=128 · Location-scatter uniform barycenters, d=128 | LUVP: 1.380 -> 5.506 ± 0.004 | Hurt |
| Use conditional normalizing flows instead of WIN | Location-scatter uniform barycenters, d=128 · Location-scatter uniform barycenters, d=128 | Bures-Wasserstein UVP: 0.460 -> 0.362 ± 0.003 | Helped |
| Add a residual connection for the conditioning input (ablation) | Conditional Real NVP · Numerical experiments | performance | Helped |
| Use a joint conditional model rather than two separate normalizing flows | OT map estimation · OT map experiments | performance | Helped |
| Increase the number of input distributions | Conditional normalizing flows · High-dimensional Gaussian barycenters | LUVP: 0.057, 0.119, 0.081, 0.073 | No clear effect |
Surprising findings
- Adding a residual connection for the conditioning input improved performance over prior conditional architectures.
Limitations
- Restricted to Wasserstein-2 barycenters; extension to p != 2 is open.
- Training may be computationally intensive for very high-dimensional data (d >> 1000).
Experimental setup
Datasets: Swiss roll high-dimensional Gaussian location-scatter high-dimensional uniform location-scatter MNIST Communities and Crime
Benchmarks: Swiss roll location-scatter Gaussian location-scatter uniform MNIST Communities and Crime
Metrics: LUVP upper bound BW2^2-UVP lower bound training time correlation with sensitive feature
Baselines: SCW2B WIN
Builds on: Real NVP Glow Wasserstein barycenter fixed-point algorithm SCW2B
Closest papers by meaning
- Optimal Transport Barycenter via Nonconvex-Concave Minimax Optimization (ICML 2025 · similarity 0.91)
- Estimating Barycenters of Distributions with Neural Optimal Transport (ICML 2024 · similarity 0.89)
- Energy-Guided Continuous Entropic Barycenter Estimation for General Costs (NeurIPS 2024 · similarity 0.88)
- A Provably Convergent and Practical Algorithm for Gromov–Wasserstein Optimal Transport (NeurIPS 2026 · similarity 0.88)
- Dynamic Conditional Optimal Transport through Simulation-Free Flows (NeurIPS 2024 · similarity 0.88)
- Differentiable Generalized Sliced Wasserstein Plans (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.