Wasserstein and Fréchet Regression of Distributions

Regression With Distribution-Valued Responses

In ordinary regression, a response is a number or vector. In some applications, however, each observation is itself a probability distribution: a patient may have an estimated distribution of future recovery time, or a site may have a distribution of measured environmental concentrations. The goal is then to predict a distribution, not just a mean or a few parameters.

A natural strategy is to treat distributions as points in a metric space and predict a distribution that is close, on average, to the observed responses at a given covariate value. The choice of distance matters: it determines what “close” means, and therefore what kind of average distribution the regression produces. Wasserstein distance is especially useful when the locations of probability mass matter, while a general Fréchet objective permits other distances and response spaces.

Definition: Let $(\mathcal M,d)$ be a metric space, and let $Y$ be a random response taking values in $\mathcal M$. For a covariate $X$, a conditional Fréchet mean at $x$ is any minimizer $$ m(x)\in\operatorname*{arg\,min}_{z\in\mathcal M} E[d^2(Y,z)\mid X=x]. $$ When $\mathcal M$ is a space of probability distributions equipped with Wasserstein distance, this is a Wasserstein regression target. A minimizer need not exist or be unique without additional conditions.

The target is a distribution $m(x)$, rather than the conditional distribution of the random distribution $Y$ given $X=x$. It summarizes that conditional law by minimizing expected squared distance. The distinction is important: a conditional Fréchet mean is one fitted response, whereas the conditional law describes variation among possible responses at the same covariate value.

Wasserstein Distance and Barycenters

For probability measures on $\mathbb R^d$ with finite second moments, write $\mathcal P_2(\mathbb R^d)$ for their space. If $\mu,\nu\in\mathcal P_2(\mathbb R^d)$, their quadratic Wasserstein distance is

W_2^2(\mu,\nu) = \inf_{\pi\in\Pi(\mu,\nu)} \int_{\mathbb R^d\times\mathbb R^d}\|u-v\|^2\,d\pi(u,v),

where $\Pi(\mu,\nu)$ is the set of joint probability measures with marginals $\mu$ and $\nu$. The coupling $\pi$ describes how mass in one distribution is paired with mass in the other; the distance is the smallest possible mean squared pairing cost.

Given distributions $\mu_1,\ldots,\mu_k$ and nonnegative weights $a_i$ summing to one, a Wasserstein barycenter is a minimizer of $\sum_{i=1}^k a_iW_2^2(\mu_i,\nu)$ over $\nu\in\mathcal P_2(\mathbb R^d)$. In regression, the weights can depend on the covariates and on the prediction point $x$. More generally, a Fréchet barycenter replaces $W_2$ by a chosen metric $d$.

Worked Example: A Local Wasserstein Regression Fit on the Line

Suppose the fitted response at a query covariate $x$ uses two observed distributions with weights $a_1=0.25$ and $a_2=0.75$. Let the first be uniform on $[0,2]$, and the second uniform on $[2,4]$. For $0<t<1$, their quantile functions are $Q_1(t)=2t$ and $Q_2(t)=2+2t$.

For one-dimensional distributions with finite second moments, the squared Wasserstein distance equals the squared $L^2(0,1)$ distance between quantile functions. Thus the weighted squared-distance objective is minimized by the weighted average quantile. Here,

$$ Q_x(t)=0.25Q_1(t)+0.75Q_2(t) =0.25(2t)+0.75(2+2t) =1.5+2t. $$

This is the quantile function of the uniform distribution on $[1.5,3.5]$, so that is the fitted distribution. It is not the mixture $0.25\mu_1+0.75\mu_2$: the mixture puts mass on both intervals, whereas the barycenter transports and averages locations of mass. The two answers summarize the observations in different ways.

Why Quantiles Give an Explicit One-Dimensional Fit

For a probability distribution $\mu$ on the real line, use the generalized inverse $Q_\mu(t)=\inf\{y:F_\mu(y)\ge t\}$, for $0<t<1$. The one-dimensional Wasserstein identity is

W_2^2(\mu,\nu)=\int_0^1\{Q_\mu(t)-Q_\nu(t)\}^2\,dt.

This identity expresses Wasserstein geometry on the line as squared distance between quantile functions. One way to see why matching quantiles is optimal is the rearrangement principle: among couplings of two real random variables, pairing their ordered values together maximizes the expected product, and hence minimizes expected squared difference. For bounded nonnegative variables $U,V$, the layer-cake identity gives $E[UV]=\int_0^\infty\int_0^\infty P(U>s,V>t)\,ds\,dt$. Each joint tail probability is at most the smaller of its two marginal tail probabilities. The quantile coupling attains this bound for every $s,t$, because its two threshold events are nested intervals in the same uniform variable. Shifting bounded variables does not change which coupling maximizes the product. Truncation and convergence in $L^2$ extend the conclusion to variables with finite second moments. Since $E[(U-V)^2]=E[U^2]+E[V^2]-2E[UV]$, maximizing the product minimizes squared difference.

Theorem: Let $Y$ be a random element of $\mathcal P_2(\mathbb R)$ with $E\int_0^1 Q_Y(t)^2\,dt<\infty$. A Wasserstein Fréchet mean of $Y$ is the distribution whose quantile function is $Q_*(t)=E[Q_Y(t)]$, with equality understood for almost every $t$. The mean is unique as a probability distribution.

The formula applies conditionally as well: at a covariate value $x$, if the conditional second-moment condition holds, the fitted quantile is $Q_x(t)=E[Q_Y(t)\mid X=x]$, for almost every $t$. To verify the theorem, take any candidate distribution $\nu$ with quantile $q(t)$. The quantile identity and Tonelli’s theorem for the nonnegative squared loss give

E[W_2^2(Y,\nu)] = \int_0^1 E[(Q_Y(t)-q(t))^2]\,dt.

For almost every $t$, the conditional second-moment assumption makes $Q_Y(t)$ square-integrable. Expanding around its expectation yields

$$ E[(Q_Y(t)-q(t))^2] = E[(Q_Y(t)-Q_*(t))^2]+(Q_*(t)-q(t))^2, $$

because $E[Q_Y(t)-Q_*(t)]=0$, so the cross term vanishes. Integrating shows that the objective at $\nu$ equals a term independent of $\nu$ plus $\int_0^1(Q_*(t)-q(t))^2\,dt$. It is minimized when $q=Q_*$ almost everywhere. The expectation of nondecreasing quantile functions is nondecreasing wherever it is finite; choosing its usual monotone representative gives a quantile function of a probability distribution. The moment condition ensures that this distribution has finite second moment. Equality of quantiles almost everywhere identifies the same distribution, proving uniqueness. The conditional assertion follows by the same calculation under the conditional law. $\square$

Worked Example: Gaussian Responses and the Meaning of the Fit

Suppose two distribution-valued responses are $N(0,1)$ and $N(4,9)$, with weights $0.25$ and $0.75$. Their standard deviations are $1$ and $3$. The quantile function of $N(m,\sigma^2)$ is $m+\sigma\Phi^{-1}(t)$, where $\Phi$ is the standard normal distribution function. Averaging the quantiles gives

0.25\{0+\Phi^{-1}(t)\} +0.75\{4+3\Phi^{-1}(t)\} =3+2\Phi^{-1}(t).

The fitted distribution is therefore $N(3,4)$: its mean is $3$, and its standard deviation is $2$, so its variance is $2^2=4$. In this family, the Wasserstein barycenter averages means and standard deviations, not variances. The result reflects the geometry of $W_2$, rather than a universal rule for averaging distribution parameters.

Existence for General Fréchet Objectives

A useful sufficient condition for existence is compactness of bounded sets in the response space. The condition is not necessary, but it gives a direct argument that applies to general metric-valued responses, not only distributions.

Theorem: Let $(\mathcal M,d)$ be a complete metric space in which every closed bounded set is compact. Let $Y$ be a measurable $\mathcal M$-valued random element and suppose $E[d^2(Y,z_0)]<\infty$ for some $z_0\in\mathcal M$. Then $F(z)=E[d^2(Y,z)]$ attains its minimum on $\mathcal M$.

Proof. The objective is finite at $z_0$, so its infimum is finite. Choose a sequence $z_j$ whose objective values approach the infimum; after discarding finitely many terms, the values $F(z_j)$ are bounded by a constant $C$. The triangle inequality and the $L^2$ triangle inequality imply

d(z_j,z_0) \le \{E[d^2(z_j,Y)]\}^{1/2} +\{E[d^2(Y,z_0)]\}^{1/2} \le \sqrt C+\{E[d^2(Y,z_0)]\}^{1/2}.

Thus the sequence lies in a bounded set. Its closure is compact by assumption, so a subsequence converges to some $z_*$. For each realization of $Y$, continuity of the metric gives $d^2(Y,z_*)=\lim d^2(Y,z_{j_\ell})$ along that subsequence. Fatou’s lemma then gives $F(z_*)\le\liminf_\ell F(z_{j_\ell})$, which equals the infimum. Hence $z_*$ is a minimizer. $\square$

This theorem supplies a general sufficient condition, not a claim that every Wasserstein response space satisfies its compactness hypothesis. On the real line, the quantile result above directly supplies existence under its stated moment condition. For a particular multivariate Wasserstein regression, existence should be justified by conditions appropriate to its response space and objective; finite second moments alone should not be treated as a general compactness argument for arbitrary metric spaces.

Nonuniqueness and What a Fitted Distribution Means

Existence does not imply uniqueness in a general metric space. Even for Wasserstein distance, multivariate barycenters can be nonunique when there are different optimal ways to match mass. The following example makes the issue explicit.

Worked Example: Two Distinct Wasserstein Barycenters

In $\mathbb R^2$, let $\mu_1$ put equal mass at $(-1,0)$ and $(1,0)$, and let $\mu_2$ put equal mass at $(0,-1)$ and $(0,1)$. Every point of the support of $\mu_1$ is squared distance $2$ from every point of the support of $\mu_2$, so $W_2^2(\mu_1,\mu_2)=2$.

Pairing $(-1,0)$ with $(0,-1)$ and $(1,0)$ with $(0,1)$, then taking pairwise midpoints, gives $\nu_A$, with equal mass at $(-0.5,-0.5)$ and $(0.5,0.5)$. The crossed pairing gives $\nu_B$, with equal mass at $(-0.5,0.5)$ and $(0.5,-0.5)$. Each measure is at squared Wasserstein distance $0.5$ from each of $\mu_1$ and $\mu_2$: the midpoint moves each paired point a distance $\sqrt{0.5}$.

For any candidate $\nu$, the triangle inequality and $(a+b)^2\le 2a^2+2b^2$ imply

W_2^2(\mu_1,\mu_2) \le \{W_2(\mu_1,\nu)+W_2(\nu,\mu_2)\}^2 \le 2W_2^2(\mu_1,\nu)+2W_2^2(\nu,\mu_2).

Consequently, the equal-weight objective $\frac12W_2^2(\mu_1,\nu)+\frac12W_2^2(\mu_2,\nu)$ is at least $W_2^2(\mu_1,\mu_2)/4=0.5$. Both $\nu_A$ and $\nu_B$ attain this bound, so both are barycenters. They are distinct distributions. Thus a fitted distribution can be set-valued unless the geometry or additional conditions ensure uniqueness.

A fitted barycenter is a squared-distance summary, not a promise that the fitted distribution is itself an observed response or that it preserves every feature of the conditional law. In one dimension its quantile function averages the observed quantile functions. In other settings, the minimizing distribution depends on the geometry, the weights, and potentially the choice among multiple minimizers.

Regression Weights and a Common Pitfall

Given training responses $Y_1,\ldots,Y_n$ in a metric space and nonnegative weights $a_i(x)$ summing to one, a weighted Fréchet regression fit is any minimizer of

\hat m(x)\in\operatorname*{arg\,min}_{z\in\mathcal M} \sum_{i=1}^n a_i(x)d^2(Y_i,z).

For one-dimensional Wasserstein responses, its quantile is $\hat Q_x(t)=\sum_i a_i(x)Q_{Y_i}(t)$. Nonnegative weights ensure this is a nondecreasing weighted average of quantile functions and define a genuine barycenter. The choice of weights determines which responses influence the prediction; the metric determines how they are averaged.

Pitfall: treating signed weights as barycentric weights. A method may produce weights that sum to one but include negative values. That fact alone does not make them valid barycenter weights. A negative linear combination of quantile functions need not be nondecreasing, and the resulting expression may not be the quantile function of any probability distribution.
Correction. Check both nonnegativity and normalization before interpreting a weighted objective as a barycenter. If weights are signed, the displayed weighted sum does not automatically define a probability measure or minimize a convex combination of squared distances; the estimator needs a separate justification and may require a different constrained construction.
Pitfall: calling the barycenter a mixture. A mixture combines probability mass from the component distributions. A Wasserstein barycenter instead minimizes transport distance and can move mass to intermediate locations, as the first worked example shows.

Key takeaway: Distributional regression is a metric-space regression problem. Specify the response space, the distance, the conditional or weighted squared-distance objective, and the conditions supporting existence and uniqueness. In one dimension, $W_2$ turns the fit into an average of quantile functions; in general spaces, the geometry can produce different answers or multiple minimizers.

Check Your Understanding

Use the definitions and results in this tutorial to answer the following.

How does a conditional Fréchet mean differ from the conditional law of a distribution-valued response?
For one-dimensional Wasserstein regression with nonnegative weights summing to one, what is averaged to obtain the fitted distribution?
In the Gaussian example, why is the fitted variance $4$, rather than the weighted average of the two variances?
Which compactness and moment conditions support the general Fréchet-mean existence theorem stated here?
Why can two different multivariate Wasserstein barycenters arise from the same pair of response distributions?
Why are weights that sum to one but include negative values not automatically valid barycentric weights?