Bayesian Multiple Multivariate Density-Density Regression
Khai Nguyen, Yang Ni, Peter Mueller
TL;DR
Bayesian MDDR addresses regressing a distribution-valued response on multiple distribution-valued predictors by modeling the fitted distribution as the sliced Wasserstein barycenter ($\mathrm{SWB}$) of push-forwards and evaluating fit with the sliced Wasserstein distance ($\mathrm{SW}$). Inference is performed within a generalized Bayes framework, leveraging a differentiable $\mathrm{SWB}$ and gradient-based MCMC to estimate regression maps $f_k$, barycenter weights $\pi$, and uncertainty. Theoretical results establish the stability of the SWB, a parametric $\mathcal{O}(n^{-1/2})$ rate for the generalized likelihood, and posterior-consistency under standard identifiability and regularity assumptions, while practical computation relies on empirical SWB approximations and MALA-style sampling. Empirically, the approach yields accurate fits and predictive performance on simulations and a population-scale single-cell dataset, enabling interpretable, sparse cell–cell communication networks via posterior barycenter weights. This framework extends distribution-to-distribution regression to the multivariate setting without reliance on Riemannian geometry and provides principled uncertainty quantification for complex density-based regression tasks.
Abstract
We propose the first approach for multiple multivariate density-density regression (MDDR), making it possible to consider the regression of a multivariate density-valued response on multiple multivariate density-valued predictors. The core idea is to define a fitted distribution using a sliced Wasserstein barycenter (SWB) of push-forwards of the predictors and to quantify deviations from the observed response using the sliced Wasserstein (SW) distance. Regression functions, which map predictors' supports to the response support, and barycenter weights are inferred within a generalized Bayes framework, enabling principled uncertainty quantification without requiring a fully specified likelihood. The inference process can be seen as an instance of an inverse SWB problem. We establish theoretical guarantees, including the stability of the SWB under perturbations of marginals and barycenter weights, sample complexity of the generalized likelihood, and posterior consistency. For practical inference, we introduce a differentiable approximation of the SWB and a smooth reparameterization to handle the simplex constraint on barycenter weights, allowing efficient gradient-based MCMC sampling. We demonstrate MDDR in an application to inference for population-scale single-cell data. Posterior analysis under the MDDR model in this example includes inference on communication between multiple source/sender cell types and a target/receiver cell type. The proposed approach provides accurate fits, reliable predictions, and interpretable posterior estimates of barycenter weights, which can be used to construct sparse cell-cell communication networks.
