Research map / Learning theory and algorithms / Optimal transport
Estimating Barycenters of Distributions with Neural Optimal Transport
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.
Proposes a bi-level max-min neural OT solver for Wasserstein barycenters with general (weak) cost functions, with duality-gap error bounds, and demonstrates it on 2D, shape-color, and Ave, celeba! benchmarks.
Problem
Computing Wasserstein barycenters in the continuous setting typically requires tri-level adversarial objectives or is restricted to quadratic cost and entropic/quadratic regularization; a scalable method handling general cost functions is needed.
Approach
Derives a max-min dual formulation of the weak OT barycenter problem with a congruence constraint on potentials (sum of weighted potentials equals zero), which is then optimized with the bi-level Neural OT procedure using stochastic maps (or deterministic maps) and congruent NN potentials, giving quality bounds on recovered plans.
Key results
- On the Ave, celeba! data-space benchmark, the method reaches FID 39.0/38.6/39.8 for k=1,2,3, beating WIN (49.3/46.9/61.5) and SCWB (56.7/53.2/58.8).
- In the manifold-constrained setting, the method (FID ~30.7-31.7) is outperformed by EgBary (FID 8.4-10.2), which the authors attribute to EgBary's Langevin sampling.
- On 2D Twister, the unregularized solver recovers the ground-truth zero-centered Gaussian barycenter.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use OURS instead of WIN or SCWB | Neural OT barycenter solver in data space · Ave, celeba! data space k=1, k=2, and k=3 | FID: 56.7 -> 39.0; 49.3 -> 39.0; 53.2 -> 38.6; 46.9 -> 38.6; 58.8 -> 39.8; 61.5 -> 39.8 | Helped |
| use OURS instead of EgBary | Neural OT barycenter solver in manifold space · Ave, celeba! manifold k=1, k=2, and k=3 | FID: 8.4 vs 30.7; 8.7 vs 31.0; 10.2 vs 31.7 | Hurt |
| use ε-KL cost functions instead of classical cost functions (ablation) | OURS on the StyleGAN manifold · Ave, celeba! manifold k=1, k=2, and k=3 | FID: 30.7 -> 34.5; 31.0 -> 34.9; 31.7 -> 35.7 | Hurt |
| use γ-Energy cost functions instead of classical cost functions (ablation) | OURS on the StyleGAN manifold · Ave, celeba! manifold k=1, k=2, and k=3 | FID: 30.7 -> 38.3; 31.0 -> 37.8; 31.7 -> 37.6 | Hurt |
Surprising findings
- In the StyleGAN manifold setting the proposed method is clearly worse than the EgBary baseline (FID about 31 vs 8-10).
Limitations
- Strong convexity assumption needed for quality bounds with classical cost is not guaranteed in practice.
- Gaussian plan parameterization is an approximation that may bias solutions.
- Adversarial optimization may be unstable and require hyperparameter tuning.
Experimental setup
Datasets: CelebA MNIST Ave, celeba! benchmark
Benchmarks: Ave, celeba! barycenter benchmark 2D Twister
Metrics: FID
Baselines: WIN SCWB EgBary
Builds on: Neural OT (Korotin et al.) weak optimal transport Sion's minimax theorem Fan et al. barycenter dual formulation
Closest papers by meaning
- Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks (ICML 2021 · similarity 0.91)
- Optimal Transport Barycenter via Nonconvex-Concave Minimax Optimization (ICML 2025 · similarity 0.89)
- Robust Barycenter Estimation using Semi-Unbalanced Neural Optimal Transport (ICLR 2025 · similarity 0.89)
- Computing Optimal Transport Maps and Wasserstein Barycenters Using Conditional Normalizing Flows (ICML 2025 · similarity 0.89)
- Energy-Guided Continuous Entropic Barycenter Estimation for General Costs (NeurIPS 2024 · similarity 0.89)
- A Unified Approach for Computing Wasserstein Barycenters of Discrete and Continuous Measures (NeurIPS 2026 · similarity 0.86)
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.