Table of Contents
Fetching ...

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.

Bayesian Multiple Multivariate Density-Density Regression

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 () of push-forwards and evaluating fit with the sliced Wasserstein distance (). Inference is performed within a generalized Bayes framework, leveraging a differentiable and gradient-based MCMC to estimate regression maps , barycenter weights , and uncertainty. Theoretical results establish the stability of the SWB, a parametric 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.
Paper Structure (17 sections, 7 theorems, 93 equations, 4 figures, 3 tables)

This paper contains 17 sections, 7 theorems, 93 equations, 4 figures, 3 tables.

Key Result

Lemma 1

For any $F_1,\ldots,F_K,F_1',\ldots,F_K'\in \mathcal{P}_p(\mathbb{R}^d)$ and $\pi=(\pi_1,\ldots,\pi_K) \in \Delta^K, \pi'=(\pi_1',\ldots,\pi_K') \in \Delta^K$, (a) Let $\bar{F}=\text{SWB}_p(F_1,\ldots, F_K,\pi)$ and $\bar{F}'=\text{SWB}_p(F_1',\ldots, F_K',\pi)$. Then (b) Let $\bar{F}=\text{SWB}_p(F_1,\ldots, F_K,\pi)$ and $\bar{F}'=\text{SWB}_p(F_1',\ldots, F_K',\pi')$. Then where $M$ is a cons

Figures (4)

  • Figure 1: The first two columns display two randomly selected in-sample regression examples ($i=1,2$), and the last two columns show two randomly selected out-of-sample examples ($i=N+1,N+2$). The first row presents the observed responses $G_i$, along with the posterior mean fitted distributions $\mathbb{E}(\widetilde{G}_i \mid data)$ for DDR and MDDR. The second row shows the observed responses, the fitted distributions $\widetilde{G}_i$ for DDR and MDDR under the final posterior sample of the chain, and the push-forward $f_{\phi_k} \sharp \widehat{F}_\ell$ of the predictors based on that final posterior sample.
  • Figure 2: The first two columns show two random in-sample donors and the last two columns show two random out-sample donors. The first row presents results for T Cells including the observed responses, and the posterior means of the KDEs of fitted distributions of DDR and MDDR. The second row presents similar results for B Cells. We use PCA for visualization.
  • Figure 3: (a) Weighted graph using posterior means of the barycenter weights (weights of edges) under the MMDR model. (b) Sparse graph obtained from the weighted graph by retaining only the edges whose weights exceed 1/3.
  • Figure 4: The first two columns show two random in-sample donors and the last two columns show two random out-sample donors. The first row presents results for NK Cells including the observed responses, and the posterior means of the KDEs of fitted distributions of DDR and MDDR. The second row presents similar results for Monocytes.

Theorems & Definitions (11)

  • Lemma 1: Stability of sliced Wasserstein barycenter
  • Theorem 1
  • Theorem 2
  • Lemma 2: Projection of sliced Wasserstein barycenter
  • proof
  • Lemma 3
  • proof
  • Lemma 4: Uniform Law of Large Numbers
  • proof
  • Proposition 1: Convexity of Sliced Wasserstein Barycenter
  • ...and 1 more