Fr\'echet Regression on the Bures-Wasserstein Manifold

TL;DR

This paper advances Fréchet regression on the Bures-Wasserstein manifold by establishing existence conditions, analyzing the optimization landscape, and proposing scalable algorithms validated on biological and imaging data.

Contribution

It provides the first comprehensive analysis of conditional barycenters in the Bures-Wasserstein space, including existence, landscape characterization, and scalable computation methods.

Findings

Existence condition for conditional barycenters in Bures-Wasserstein space

Objective function is free of local maxima under certain conditions

Proposed algorithms perform well on biological and imaging datasets

Abstract

Fr\'echet regression, or conditional Barycenters, is a flexible framework for modeling relationships between covariates (usually Euclidean) and response variables on general metric spaces, e.g., probability distributions or positive definite matrices. However, in contrast to classical barycenter problems, computing conditional counterparts in many non-Euclidean spaces remains an open challenge, as they yield non-convex optimization problems with an affine structure. In this work, we study the existence and computation of conditional barycenters, specifically in the space of positive-definite matrices with the Bures-Wasserstein metric. We provide a sufficient condition for the existence of a minimizer of the conditional barycenter problem that characterizes the regression range of extrapolation. Moreover, we further characterize the optimization landscape, proving that under this condition, the objective is free of local maxima. Additionally, we develop a projection-free and provably correct algorithm for the approximate computation of first-order stationary points. Finally, we provide a stochastic reformulation that enables the use of off-the-shelf stochastic Riemannian optimization methods for large-scale setups. Numerical experiments validate the performance of the proposed methods on regression problems of real-world biological networks and on large-scale synthetic Diffusion Tensor Imaging problems.