The optimal transportation problem defines a geometry of probability measures which leads to a definition for weighted averages (barycenters) of measures, finding application in the machine learning and computer vision communities as a signal processing tool. Here, we implement a barycentric coding model for measures which are supported on a graph, a context in which the classical optimal transport geometry becomes degenerate, by leveraging a Riemannian structure on the simplex induced by a dynamic formulation of the optimal transport problem. We approximate the exponential mapping associated to the Riemannian structure, as well as its inverse, by utilizing past approaches which compute action minimizing curves in order to numerically approximate transport distances for measures supported on discrete spaces. Intrinsic gradient descent is then used to synthesize barycenters, wherein gradients of a variance functional are computed by approximating geodesic curves between the current iterate and the reference measures; iterates are then pushed forward via a discretization of the continuity equation. Analysis of measures with respect to given dictionary of references is performed by solving a quadratic program formed by computing geodesics between target and reference measures. We compare our novel approach to one based on entropic regularization of the static formulation of the optimal transport problem where the graph structure is encoded via graph distance functions, we present numerical experiments validating our approach, and we conclude that intrinsic gradient descent on the probability simplex provides a coherent framework for the synthesis and analysis of measures supported on graphs.
Figures & tables
Figure 1: A schematic of the barycentric coding model for measures supported on a graph. At the extremal points of the simplex, which correspond to the vertices of the triangle, are the reference measures; away from the extremal points are barycentric interpolations of the reference measures for certain weights λ . The coordinates recovered by solving the analysis problem are displayed and labeled with λ^ .
Figure 2: Weighted center of mass on a non-Euclidean space; νλ is the synthesis of the points ν1,ν2,ν3 with weights λ=(0.2,0.5,0.3)
Figure 3: To solve the analysis problem, we compute the tangent vectors associated to the unit time geodesic curves connecting the target measure to each reference measure, then form the associated Gram matrix.
Figure 4: Comparison between a classic optimal transport plan and an entropically regularized one. Left: two samples of 100 points each in the plane (normally distributed with different means and variances). Center: optimal assignment between points to minimize transport cost, obtained via computation of a linear program. Right: entropic transport plan for minimizing entropic transport cost. Notice the difference in sparsity between the two plans: entropic transport results in a diffusive effect.
Figure 5: The Chambolle-Pock routine is used to approximate energy minimizing curves in P(X) , ultimately producing a rectified curve encoded as a sequence of points and a collection of momentum vector fields. We leverage these vector fields to approximate the gradient of ( 8 ).
Routine
δg
N
δb
ε
Chambolle-Pock
✓
✓
Synthesis
✓
✓
✓
✓
Analysis
✓
✓
Table 1: Table of hyperparameters associated to each algorithm
Figure 6: Barycenter of 4 measures on the graph formed by the edges and vertices of a hypercube. Bottom left: edges of the hypercube are uniformly weighted with ωi,j≡1 ; bottom right: edges are randomly weighted (weights are symmetric to preserve the reversibility condition of the Markov chain). Parameters used to compute barycenters: ε=0.1,N=10,δg=10−10,δb=10−10 .
Figure 7: Barycenter of 4 measures on the graph formed by the spatial geography node constructed from the Massachusetts state house districts. Left: 4 randomly generated geographically concentrated reference measures; right: barycenters with weights λ=[0.4,0.3,0.2,0.1] and recovered coordinates; top right, barycenter computed with δg=10−10 ; bottom right: barycenter computed with δg=10−6 . We observe that more rigorous convergence tolerance in the Chambolle-Pock routine leads to superior coordinate recovery in the analysis problem. Common parameters used to compute barycenters: ε=0.1,N=2,δb=10−8 . Data retrieved from https://www.mass.gov/info-details/massgis-data-massachusetts-house-legislative-districts-2021 .
Figure 8: Comparison of barycenters on a 2 node graph computed by solving an ODE vs computed by intrinsic gradient descent.
Figure 9: Comparison between barycenters of two measures computed as geodesics via Chambolle-Pock and via intrinsic gradient descent
Figure 10: Distribution of 1000 computed relative norm differences between (a) the barycenter νb of two randomly generated measures ν0i,ν1i on a triangle graph and (b) the corresponding point on the geodesic between ν0i and ν1i , νg . When computing the relative error we assume that the point on the geodesic is a genuine representative of the barycenter, and therefore compute ∥νg−νb∥π/∥νg∥π
Figure 11: Barycenter of three locally concentrated measures on a grid graph on 49 vertices. Bottom left: computed via intrinsic gradient descent and the dynamic formulation of discrete transport; bottom center: computed via entropic regularization with cost given by the squared shortest path metric; bottom right: computed via entropically regularized transport, with the cost matrix given by the squared diffusion distance for t=12 . Parameters used to compute barycenters: ε=0.1,N=3,δg=10−12,δb=10−10 .
Figure 12: Barycenter of three geographically concentrated measures on a graph representing the 48 contiguous states + DC. Bottom left: computed via intrinsic gradient descent and the dynamic formulation of discrete transport; bottom center: computed via entropic regularization with cost given by the squared shortest path metric; bottom right: computed via entropically regularized transport, with the cost matrix given by the squared diffusion distance for t=11 . Parameters used to compute barycenters: ε=0.1,N=3,δg=10−12,δb=10−10 .
Figure 13: Barycenter of several reference measures computed with differing initializations for the descent scheme. The maximum observed difference in ∥⋅∥π between the computed barycenters was 1×10−9 , the parameters used to compute the barycenters were N=16,δg=10−6,ε=0.25,δb=10−10.
Figure 14: Barycenter of several reference measures computed with differing initializations for the descent scheme. The maximum observed difference in ∥⋅∥π between the computed barycenters was 8×10−10 , the parameters used to compute the barycenters were N=5,δg=10−6,ε=0.25,δb=10−10.
Figure 15: Distribution of 1000 observed relative errors in recovery of coordinates from a barycenter synthesized via intrinsic gradient descent for a grid graph on 16 nodes, 3 random measures, a random set of coordinates λ∈Δ2 , with hyperparameters δg=10−9,N=10,ε=0.5,δb=10−10 .
Figure 16: Recovery of coordinates for synthesized barycenters.
Figure 17: Log-linear scale of the absolute ℓ2 norm difference between consecutive iterates and the number of iterations, or equivalently, the norm of the W gradient at each iteration multiplied by the stepsize ε . Hyperparameters for this experiment: ε=0.25 , δb=10−10 , N=10 , δg=10−9 .
We introduce, to our knowledge, the first deep generative modeling framework for probability distributions continuously supported on compact metric graphs. Given source and target measures on a metric graph, our method embeds the graph into a smooth ambient space, solves an entropic Kantorovich problem via a neural semidual parameterization, and projects generated samples back onto the original graph. We study two embedded geometries: an extrinsic Euclidean realization and the intrinsic tropical Abel--Jacobi embedding into the Jacobian torus. In both cases, the resulting generator is graph-supported by construction. We prove that, in the joint limit of increasing neural expressivity, the learned generator converges weakly to a valid transport coupling between the original graph measures. Empirically, across a range of geometrically distinct graphs, our method matches or improves upon heuristic transport baselines based on discrete graph OT, while scaling more favorably. Finally, we demonstrate scalability on real-world urban mobility data by training our model on one million Uber pickup locations in Manhattan, New York City.
Alessandro Micheli, Yueqi Cao, Anthea Monod +1
Imperial College London London, UK · KTH Royal Institute of Technology Stockholm, Sweden · Statens Serum Institut Copenhagen, Denmark +1
Barycenter models, based on stochastic and deterministic optimal transport bridges between probability measures, have proven to be useful for supervised and unsupervised learning for (probability) measure-valued data. In this paper, we utilize the directional asymmetry and dual drift-diffusion control arising in the recently proposed Schrödinger-Bass Bridge (SBB) between two measures, and propose SBB-encoder, SBB-decoder, SBB-symmetric, and SBB-Markovian barycenter models. For each of these novel models, our first theoretical result establishes the existence of the SBB-barycenter under mild assumptions. Our second theoretical result characterizes the decoder SBB-barycenter of Gaussians where we show that for Gaussian input measures, the decoder formulation admits a Gaussian minimizer, reducing barycenter computation to a finite-dimensional optimization and yielding a compact algorithm. For general input measures, we derive a hybrid Wasserstein gradient particle descent scheme for all four formulations using the SBB dual representation. Finally, we utilize these algorithms for solving distributional-inverse/interpolation problems arising in the context of spatial images and single-cell RNA sequencing data, where we show that the SBB-decoder and Markovian barycenter allow us to interpolate between successive time points. Results show that our method outperforms existing barycenter models for measure interpolation and inference.
Ngoc-Hai Nguyen, Thuan Nguyen, Prakash Ishwar +1
Department of Electrical and Computer Engineering Tufts University · Department of Engineering, Engineering Technology East Tennessee State University · Department of Electrical and Computer Engineering Boston University +1
Defining a weighted mean over probability measures under probability metrics is a central tool in probabilistic machine learning. Under the Wasserstein metric, these are called Wasserstein barycenters. While most approaches compute barycenters for a fixed weight vector, approximating the whole family of barycenters over the simplex, which we call the Wasserstein simplex, remains underexplored. We refer to this problem as Universal Barycenter Approximation, and propose BaryFM, a flow matching model transporting the marginal measures into any barycenter in the Wasserstein simplex. Once trained, the network can draw samples from measures in the Wasserstein simplex through an ordinary differential equation. We validate our method on 4 downstream tasks: domain adaptation, generalization, Bayesian posterior aggregation and algorithmic fairness. BaryFM achieves the best average rank among 15 competing methods across 10 domain adaptation benchmarks, matching or surpassing non-universal solvers.