[Paper Review] Conditional Wasserstein Barycenters and Interpolation/Extrapolation of Distributions

Jianing Fan, Hans‐Georg Müller|arXiv (Cornell University)|Jul 20, 2021

Geometric Analysis and Curvature Flows40 references4 citations

TL;DR

This paper introduces conditional Wasserstein barycenters for multivariate distributions using the 2-Wasserstein metric, enabling interpolation and extrapolation of distributional responses based on scalar or vector predictors. Under regularity conditions, the method ensures that geodesics in predictor space map to geodesics in Wasserstein space, with convergence guarantees via Sinkhorn-regularized computation and asymptotic consistency for both global and local models.

ABSTRACT

Increasingly complex data analysis tasks motivate the study of the dependency of distributions of multivariate continuous random variables on scalar or vector predictors. Statistical regression models for distributional responses so far have primarily been investigated for the case of one-dimensional response distributions. We investigate here the case of multivariate response distributions while adopting the 2-Wasserstein metric in the distribution space. The challenge is that unlike the situation in the univariate case, the optimal transports that correspond to geodesics in the space of distributions with the 2-Wasserstein metric do not have an explicit representation for multivariate distributions. We show that under some regularity assumptions the conditional Wasserstein barycenters constructed for a geodesic in the Euclidean predictor space form a corresponding geodesic in the Wasserstein distribution space and demonstrate how the notion of conditional barycenters can be harnessed to interpolate as well as extrapolate multivariate distributions. The utility of distributional inter- and extrapolation is explored in simulations and examples. We study both global parametric-like and local smoothing-like models to implement conditional Wasserstein barycenters and establish asymptotic convergence properties for the corresponding estimates. For algorithmic implementation we make use of a Sinkhorn entropy-penalized algorithm. Conditional Wasserstein barycenters and distribution extrapolation are illustrated with applications in climate science and studies of aging.

Motivation & Objective

• To develop a statistical framework for regression of multivariate distributions on scalar or vector predictors using the 2-Wasserstein metric.
• To extend the concept of Wasserstein barycenters to conditional settings where the response is a distribution and the predictor is a covariate.
• To establish conditions under which geodesics in the predictor space induce geodesics in the Wasserstein distribution space, enabling meaningful interpolation and extrapolation.
• To provide asymptotic convergence results for both global parametric-like and local smoothing-like models in the conditional barycenter framework.
• To enable practical computation via Sinkhorn entropy-penalized algorithms with convergence to true Wasserstein estimates as regularization diminishes.

Proposed method

• Uses the 2-Wasserstein metric as a Riemannian structure on the space of probability measures to define conditional barycenters as Fréchet means of distributional responses.
• Defines conditional barycenters via minimization of weighted sum of squared 2-Wasserstein distances between estimated and observed distributions.
• Applies Sinkhorn algorithm with entropy regularization to approximate the computationally intensive optimal transport plans, reducing complexity and enabling scalable implementation.
• Establishes theoretical convergence by showing that Sinkhorn-regularized estimates converge to true Wasserstein barycenters as the regularization parameter tends to infinity.
• Develops two modeling approaches: a global parametric-like model and a local smoothing-like model, both grounded in conditional barycenter estimation.
• Proves that under regularity conditions, the conditional barycenter path induced by predictor geodesics forms a geodesic in the Wasserstein space, enabling valid interpolation and extrapolation.

Experimental results

Research questions

• RQ1 Can conditional Wasserstein barycenters be used to interpolate and extrapolate multivariate distributions based on predictors in a statistically sound way?
• RQ2 Under what conditions does a geodesic in the predictor space induce a geodesic in the Wasserstein distribution space?
• RQ3 How can the computational complexity of optimal transport be reduced while preserving statistical consistency in distributional regression?
• RQ4 What are the asymptotic convergence rates of global and local models for conditional Wasserstein barycenters?
• RQ5 How do Sinkhorn-regularized estimates converge to the true Wasserstein barycenter estimates as regularization diminishes?

Key findings

• Conditional Wasserstein barycenters form a geodesic in the Wasserstein space when the predictor space follows a geodesic, enabling natural interpolation and extrapolation of multivariate distributions.
• The Sinkhorn-regularized barycenter estimates converge almost surely to the true Wasserstein barycenter as the regularization parameter tends to infinity.
• For both global and local models, the proposed estimators achieve asymptotic convergence with well-defined rates under regularity assumptions on the underlying measures.
• The method successfully interpolates and extrapolates distributions in real-world applications, including climate data and longitudinal aging studies.
• Theoretical results confirm that the conditional barycenter path preserves geodesic structure, validating the use of Wasserstein-based regression for distributional responses.
• Simulations and data examples demonstrate the method's robustness and practical utility in modeling complex, high-dimensional distributional responses.