Conditional Distribution Regression with Kernel Mean Embeddings

Why Regress to a Distributional Feature?

A conditional distribution contains more information than a conditional mean. For example, two predictor values can have the same response mean but very different spread or shape. One way to predict the whole conditional law is to map each response into a feature space, then regress that feature on the predictor. The resulting prediction is a conditional kernel mean embedding.

This approach turns distributional prediction into regression with values in a Hilbert space. Its attraction is that the response feature can encode nonlinear features of the entire law, while the regression step has an explicit regularized solution. The embedding itself is an element of a function space; it should not automatically be interpreted as a density or as a valid probability distribution.

The expectation is a Bochner conditional expectation: it is a Hilbert-space-valued expectation, not a pointwise density. Bounded \(l(y,y)\) is a convenient sufficient condition for integrability, since \(\|\psi(y)\|_{\mathcal H_l}^2=l(y,y)\). If \(l\) is characteristic, the mean embedding of a probability measure identifies that measure uniquely. Characteristicness is useful for interpretation, but it is not needed to derive the regression estimator.

The choice of output kernel determines which distributional features are represented. With the linear kernel \(l(y,y')=yy'\) for real responses, the embedding is just the conditional mean. A richer kernel, such as a Gaussian kernel, can encode more than the mean. For a finite categorical response and the delta kernel \(l(y,y')=\mathbf 1\{y=y'\}\), the feature vector records category probabilities exactly.

Conditional Embeddings as Regression Targets

The target \(\mu_{Y\mid X}\) is a random element of \(\mathcal H_l\). Under square-integrability, conditional expectation is the squared-error-optimal prediction of this random element from \(X\). This fact explains why ordinary squared-loss regression can estimate a conditional distribution embedding.

Proof. Write \(\psi(Y)-g(X)=[\psi(Y)-\mu_{Y\mid X}]+[\mu_{Y\mid X}-g(X)]\), expand the squared norm, and take expectations. The cross term is \(2E\langle \psi(Y)-\mu_{Y\mid X},\mu_{Y\mid X}-g(X)\rangle\). Conditionally on \(X\), the second factor is fixed and the conditional expectation of the first factor is zero. The cross term is therefore zero. The remaining first term does not depend on \(g\), and the second is nonnegative, with value zero precisely when \(g(X)=\mu_{Y\mid X}\) almost surely. \(\square\)

This result concerns prediction in embedding space. It does not say that every function \(g(x)\) is the embedding of a probability measure. To obtain a data-driven predictor, impose smoothness across predictor values by placing \(g\) in a vector-valued reproducing-kernel Hilbert space.

Regularized Kernel Regression

Use the operator-valued kernel \(\Gamma(x,x')=k(x,x')I_{\mathcal H_l}\), where \(I_{\mathcal H_l}\) is the identity operator on the output space. Its associated Hilbert space consists of \(\mathcal H_l\)-valued functions. Given independent observations \((x_i,y_i)\), \(i=1,\ldots,n\), define \(\widehat g\) to minimize empirical squared embedding loss plus a norm penalty.

Let \(K\) be the \(n\times n\) Gram matrix with \(K_{ij}=k(x_i,x_j)\), and let \(k_x\) be the vector with \(i\)th entry \(k(x_i,x)\). The estimate has the explicit form \[ \widehat\mu_{Y\mid x}=\widehat g(x)=\sum_{i=1}^n\alpha_i(x)\psi(y_i), \qquad \alpha(x)=(K+n\lambda I_n)^{-1}k_x. \] The inverse exists because \(K\) is positive semidefinite and \(n\lambda>0\). The coefficients depend on the predictor data and query point; the observed responses enter through their feature vectors.

For completeness, the coefficient equations do yield the minimizer, not merely a stationary candidate. Let \(g_*\) be the function defined by their solution, and let \(h\) be any function in the vector-valued Hilbert space. The equations imply \[ \frac{1}{n}\sum_i\langle g_*(x_i)-\psi(y_i),h(x_i)\rangle_{\mathcal H_l} +\lambda\langle g_*,h\rangle_\Gamma=0. \] Indeed, \(g_*(x_i)-\psi(y_i)=-n\lambda c_i\), while the reproducing property gives \(\langle g_*,h\rangle_\Gamma=\sum_i\langle c_i,h(x_i)\rangle_{\mathcal H_l}\). Expanding the objective at \(g_*+h\) now shows that its difference from the objective at \(g_*\) is \[ \frac{1}{n}\sum_i\|h(x_i)\|_{\mathcal H_l}^2+\lambda\|h\|_\Gamma^2\geq 0. \] Because \(\lambda>0\), equality requires \(h=0\) in the Hilbert space, proving uniqueness. \(\square\)

Covariance-Operator View

The same estimator can be written in terms of feature covariance operators. For vectors \(u,v\) in Hilbert spaces, write \(u\otimes v\) for the operator mapping \(a\) to \(u\langle v,a\rangle\). Define \[ \widehat C_{XX}=\frac{1}{n}\sum_i\phi(x_i)\otimes\phi(x_i), \qquad \widehat C_{YX}=\frac{1}{n}\sum_i\psi(y_i)\otimes\phi(x_i). \] Then the regularized operator prediction is \[ \widehat C_{YX}(\widehat C_{XX}+\lambda I_{\mathcal H_k})^{-1}\phi(x) =\sum_i\alpha_i(x)\psi(y_i). \]

This identity follows from the kernel Gram-matrix identity (or by applying the regularized normal equations in the span of the observed feature vectors). The factor \(n\lambda\) in \(K+n\lambda I_n\) corresponds to using an empirical covariance with factor \(1/n\). It is easy to lose this factor when switching between operator and matrix formulas.

At the population level, one sometimes writes \(C_{YX}C_{XX}^{-1}\phi(x)\), with \(C_{XX}=E[\phi(X)\otimes\phi(X)]\) and \(C_{YX}=E[\psi(Y)\otimes\phi(X)]\). This expression requires care: \(C_{XX}\) need not have a bounded inverse, and the conditional embedding need not be representable by a single bounded linear operator applied to \(\phi(x)\). Regularization replaces the unstable inverse by \( (C_{XX}+\lambda I)^{-1}\); it does not by itself establish an exact population operator identity.

Interpretation and a Common Pitfall

Conditional embedding regression matters when the target is an entire conditional law but a density model is inconvenient or restrictive. It supplies a direct regression estimator of distributional features and makes predictions through a kernel matrix. Kernel choice, regularization, and the input distribution all affect the result. An embedding prediction can be compared or used through its Hilbert-space representation, but its validity as a probability measure requires separate scrutiny.

• Pitfall: treating the regression weights as probabilities. One might see \(\widehat\mu_{Y\mid x}=\sum_i\alpha_i(x)\psi(y_i)\) and conclude that the coefficients form a probability distribution over observed responses. The formula imposes no nonnegativity or sum-to-one constraint. As the examples show, ridge shrinkage can change the sum and kernel interpolation can give a negative weight.
• Correction: The formula estimates a Hilbert-space conditional expectation under squared loss. If a probability-valued output is required, one must use or add a method that enforces the relevant constraints, and assess the resulting estimator on that basis.
• Pitfall: interpreting an output kernel as a density model. The feature map represents a response through inner products in \(\mathcal H_l\); it need not be a density or a pointwise estimate of one. The kernel determines which features of the law the embedding retains.