This repository contains the pytorch implementation of the paper "Wasserstein Gradient Flows for Scalable and Regularized Barycenter Computation", published at UAI 2026. This code implements a discrete (in time and space) Wasserstein Gradient Flow of the barycenter functional,
Wasserstein barycenters provide a principled approach for aggregating probability measures, while preserving the geometry of their ambient space. Existing discrete methods are not because as they assume access to the complete set of samples from the input measures. Meanwhile, neural network approaches do scale well, but rely on complex optimization problems and cannot easily incorporate label information. We address these limitations through gradient flows in the space of probability measures. Through time discretization, we achieve a scalable algorithm that i) relies on mini-batch optimal transport, ii) accepts modular regularization through task-aware functions, and iii) seamlessly integrates supervised information into the ground-cost. We empirically validate our approach on domain adaptation benchmarks that span computer vision, neuroscience, and chemical engineering. Our method establishes a new state-of-the-art Wasserstein barycenter solver, with labeled barycenters consistently outperforming unlabeled ones.
This project uses uv.
# Install the library and its dependencies into a local environment
uv sync
# Optional: include plotting helpers (bgf.utils)
uv sync --extra vizOr install into an existing environment with pip:
pip install -e .import torch
from bgf.measures import EmpiricalMeasure
from bgf.flows.euler import EmpiricalBarycenterFlow
from bgf.optimal_transport.losses import WassersteinDistance, BarycenterFunctional
# Create source measures
X1 = torch.randn(1000, 2)
X2 = torch.randn(1000, 2) + 3
measures = [EmpiricalMeasure(X1), EmpiricalMeasure(X2)]
# Set up barycenter computation
base_loss = WassersteinDistance()
barycenter_functional = BarycenterFunctional(measures, base_loss)
# Run gradient flow
flow = EmpiricalBarycenterFlow(
measures=measures,
n_samples=500,
n_iter=100,
time_step=0.01,
barycenter_functional=barycenter_functional,
)
flow.fit(optimizer='sgd', momentum=0.9)
# Access computed barycenter
barycenter_points = flow.barycenterfrom bgf.measures import LabeledEmpiricalMeasure
from bgf.flows.euler import LabeledEmpiricalBarycenterFlow
from bgf.optimal_transport.losses import JointWassersteinDistance, LabeledBarycenterFunctional
from bgf.optimal_transport.ground_costs import LabeledMinkowskiCost
# Create labeled measures
source_measure = LabeledEmpiricalMeasure(X_source, Y_source)
target_measure = LabeledEmpiricalMeasure(X_target, Y_target)
# Set up joint Wasserstein distance
ground_cost = LabeledMinkowskiCost(p=2, power=2, beta=1.0)
base_loss = JointWassersteinDistance(ground_cost=ground_cost)
barycenter_functional = LabeledBarycenterFunctional(
[source_measure, target_measure], base_loss
)
# Run labeled flow
flow = LabeledEmpiricalBarycenterFlow(
measures=[source_measure, target_measure],
n_samples=1000,
n_iter=200,
time_step=0.1,
barycenter_functional=barycenter_functional,
)
flow.fit(optimizer='sgd', momentum=0.9)
# Extract barycenter
X_barycenter, Y_barycenter = flow.barycenterbarycentric-gradient-flows/
├── bgf/
│ ├── flows/ # Gradient flow implementations
│ │ ├── euler.py # Empirical barycenter flows
│ │ └── generic.py # Base flow classes
│ ├── optimal_transport/ # Optimal transport components
│ │ ├── losses.py # Distance functionals
│ │ ├── ot_solver.py # OT problem solvers
│ │ ├── ground_costs.py # Distance metrics
│ │ ├── sinkhorn.py # Sinkhorn algorithms
│ │ └── barycenters.py # Barycenter algorithms
│ ├── measures.py # Data measure classes
│ ├── noise_schedule.py # Noise schedules
│ └── utils/ # Plotting / typesetting helpers
├── pyproject.toml
└── README.md
WassersteinDistance: standard Wasserstein-2 distance.JointWassersteinDistance: joint distance for labeled data.BarycenterFunctional: objective for barycenter computation.SinkhornBarycenterFunctional: entropic-regularized variant.
EmpiricalBarycenterFlow: direct optimization on point clouds.LabeledEmpiricalBarycenterFlow: flows for labeled data.
EmpiricalMeasure: point-cloud measures.LabeledEmpiricalMeasure: labeled point-cloud measures.- Synthetic measures:
SwissRoll,LabeledSwissRoll,ImbalancedBlobs,TwoMoons.
- Sinkhorn algorithm for entropic regularization.
- GPU-accelerated implementations.
from bgf.optimal_transport.losses import LabeledPotentialEnergy, LabeledInteractionEnergy
# Define a custom potential energy
def v(x, y):
"""Encourage margin separation between classes."""
# ... your code here
return value
potential_energy = LabeledPotentialEnergy(v)
# Define a custom interaction energy
def u(x, y):
"""Custom interaction between points."""
# ... your code here
return value
interaction_energy = LabeledInteractionEnergy(u)
# Use in a flow
flow = LabeledEmpiricalBarycenterFlow(
measures=measures,
potential_energy_functional=potential_energy,
interaction_energy_functional=interaction_energy,
# ... other parameters
)This project is licensed under the Apache License, Version 2.0. See the LICENSE file for details.
@misc{montesuma2026wassersteingradientflowsscalable,
title={Wasserstein Gradient Flows for Scalable and Regularized Barycenter Computation},
author={Eduardo Fernandes Montesuma and Yassir Bendou and Mike Gartrell},
year={2026},
eprint={2510.04602},
archivePrefix={arXiv},
primaryClass={stat.ML},
url={https://arxiv.org/abs/2510.04602},
}