Conditional expectation

Supervised learning is a prediction problem: given features $X$, guess the target $Y$. Before any model is trained, probability theory already says what the best possible guess is, if the joint distribution were known, and how much error would remain even then. The answer is the conditional expectation $E[Y \mid X]$. The multivariate Gaussian note found it for one family: a linear function. This note defines it in general, proves that it minimizes squared error, splits the variance of $Y$ into the part $X$ explains and the part it cannot, and shows which other losses lead to which other targets.

Notation. $m(x) = E[Y \mid X = x]$ is a number for each $x$, the regression function. $E[Y \mid X] = m(X)$ is the same function applied to the random $X$, so it is a random variable. $\operatorname{Var}(Y \mid X = x)$ is the variance of the conditional distribution.

Conditional expectation as a function#

The conditional expectation is the expectation of the conditional distribution:

$$E[Y \mid X = x] = \sum_y y\, p_{Y \mid X}(y \mid x) \quad \text{or} \quad \int y\, f_{Y \mid X}(y \mid x)\, dy .$$

With two dice, $X$ the first and $Y$ the larger of the two, the conditional distribution given $X = 3$ puts $3/6$ on 3 and $1/6$ on each of 4, 5, 6, so $E[Y \mid X = 3] = (9 + 4 + 5 + 6)/6 = 4$. Over all six values,

$$m(1), \dots, m(6) = \tfrac72, \tfrac{11}3, 4, \tfrac92, \tfrac{31}6, 6 .$$

For the Gaussian vector of the previous note, $m(x_1) = 0.6\, x_1$. In general $m$ can be any function; it is linear for jointly Gaussian variables and in general need not be.

Reading $E[Y \mid X]$ as a random variable, $m(X)$, is what makes the concept work. Before $X$ is observed, the best prediction of $Y$ is itself uncertain, and it has a distribution, a mean and a variance of its own.

The tower property#

Averaging the conditional means over the distribution of $X$ gives back the unconditional mean:

$$E\big[E[Y \mid X]\big] = E[Y] .$$

For the dice, $\frac16 (\frac72 + \frac{11}3 + 4 + \frac92 + \frac{31}6 + 6) = \frac{161}{36}$, the mean of the larger value found in the joint distributions note. The tower property, also called the law of total expectation, is the law of total probability for expectations: to compute a mean, condition on whatever makes the problem easy, then average.

A random number of terms. A shop has $N \sim \mathrm{Poisson}(20)$ customers in an hour, and each spends an amount with mean 30, independent of the other customers and of their number. Given $N = n$ the expected revenue is $30n$, so the expected revenue is $E[30N] = 30 \cdot 20 = 600$. The same step gives the expected cost of a random number of retries, or the expected number of tokens generated before a stop token.

The best predictor#

For any function $g$, add and subtract $m(X)$:

$$E\big[(Y - g(X))^2\big] = E\big[(Y - m(X))^2\big] + E\big[(m(X) - g(X))^2\big] .$$

The cross term $2E[(Y - m(X))(m(X) - g(X))]$ vanishes: conditioning on $X$ turns $m(X) - g(X)$ into a constant and $E[Y - m(X) \mid X] = 0$, and the tower property averages the zero. The first term does not depend on $g$; the second is zero only for $g = m$. So $m$ is the best predictor, and every other predictor pays the extra squared distance from $m$.

Figure

The conditional mean is the best predictor

The decomposition separates the error of any predictor into two parts. The first term, the noise in $Y$ that $X$ does not determine, is irreducible: $0.090$ in the figure, whatever the model. The second is the squared distance of $g$ from $m$. For the best predictor $g^*$ inside a model class it is the approximation error of that class: the best straight line cannot follow a sine and pays $0.255 - 0.090 = 0.165$ on top. A model fitted to a finite training set lands at some other $g$ in the class, further from $m$ than $g^*$; the extra distance is the estimation error, plus an optimization error if training stops short of the best fit. The estimators note splits the estimation part into bias and variance.

The law of total variance#

The variance of $Y$ splits in the same way:

$$\operatorname{Var}(Y) = \underbrace{E\big[\operatorname{Var}(Y \mid X)\big]}_{\text{unexplained}} + \underbrace{\operatorname{Var}\big(E[Y \mid X]\big)}_{\text{explained}} .$$

The first term is the average spread of $Y$ around the prediction, the irreducible error of the previous section; the second is how much the prediction itself varies with $X$. For the dice:

Figure

The law of total variance on two dice

The ratio $\operatorname{Var}(E[Y \mid X]) / \operatorname{Var}(Y)$ is the largest $R^2$ any model can reach using $X$, here $0.39$. For the Gaussian of the previous note it checks out with the formulas found there: $\operatorname{Var}(X_2) = 1 = 0.28 + 0.6^2 \cdot 2$, the conditional variance plus the variance of the linear conditional mean.

The law also says that replacing a random quantity by its conditional expectation never increases the variance, since the dropped term $E[\operatorname{Var}(\cdot \mid X)]$ is non-negative. Applied to an estimator, by conditioning it on a sufficient statistic, a summary of the data given which their distribution no longer depends on the unknown parameter (for coin flips, the number of heads), this is the Rao–Blackwell theorem: averaging out irrelevant randomness never makes an estimate worse and usually makes it better.

Other losses, other targets#

The conditional mean is the answer for squared loss. Other losses ask for other summaries of the same conditional distribution:

Loss	Best prediction given $X = x$	In practice
squared error $(y - g)^2$	conditional mean $E[Y \mid X = x]$	regression with MSE
absolute error $\lvert y - g \rvert$	conditional median	robust regression, MAE
pinball loss at level $q$, $(y - g)(q - \mathbb 1[y < g])$	conditional $q$-quantile	prediction intervals, quantile regression
0–1 loss for classes	most probable class $\arg\max_k P(Y = k \mid X = x)$	the Bayes classifier
log loss (cross-entropy)	the whole conditional distribution $P(Y \mid X = x)$	classifiers, language models

The last row is why classifiers and language models output distributions rather than single answers. When the conditional distribution has several modes, the mean is a poor summary: a model trained with squared error to predict the next video frame averages the possible futures into a blur, which is one of the reasons generative models learn the conditional distribution and sample from it.

What changed#

	Without $X$	With $X$
Best prediction of $Y$	the constant $E[Y]$	the function $m(X) = E[Y \mid X]$
Its mean squared error	$\operatorname{Var}(Y)$	$E[\operatorname{Var}(Y \mid X)]$
Computing a mean	directly	tower property: condition, then average
Variance of $Y$	one number	explained + unexplained
Error of a trained model	one number	irreducible + squared distance from $m$ (approximation, estimation)

Limitations#

• $m$ is defined by the unknown joint distribution. Supervised learning is, in essence, the problem of estimating it from finitely many samples, and in many dimensions that needs either many samples or strong assumptions about the shape of $m$.
• A mean is not a decision. When the conditional distribution is skewed or multimodal, the conditional mean can be a value that rarely or never occurs; the loss, not habit, should choose the target.
• Conditioning is not intervention. $E[Y \mid X = x]$ describes what $Y$ looks like among cases where $X$ happens to be $x$, not what $Y$ would become if $X$ were set to $x$.

What you should remember#

• $E[Y \mid X = x]$ is the mean of the conditional distribution; $E[Y \mid X] = m(X)$ is a random variable.
• Tower property: $E[E[Y \mid X]] = E[Y]$.
• $m(X)$ minimizes $E[(Y - g(X))^2]$; $E[(Y - g(X))^2] = E[(Y - m(X))^2] + E[(m(X) - g(X))^2]$.
• $\operatorname{Var}(Y) = E[\operatorname{Var}(Y \mid X)] + \operatorname{Var}(E[Y \mid X])$: unexplained plus explained; the irreducible error is the first term.
• Squared loss targets the mean, absolute loss the median, pinball loss a quantile, 0–1 loss the most probable class, log loss the whole conditional distribution.

Code#

The dice in exact fractions, the best predictor and a random sum by simulation:

import numpy as np
from fractions import Fraction

# two dice, X the first and Y the larger: E[Y | X = x], the tower property, the law of total variance
m = {x: Fraction(x * x + sum(range(x + 1, 7)), 6) for x in range(1, 7)}
v = {x: Fraction(x**3 + sum(y * y for y in range(x + 1, 7)), 6) - m[x] ** 2 for x in range(1, 7)}
EY = sum(m.values()) / 6
print([str(m[x]) for x in m], EY)                     # ['7/2', '11/3', '4', '9/2', '31/6', '6'] 161/36
unexplained, explained = sum(v.values()) / 6, sum((m[x] - EY) ** 2 for x in m) / 6
print(unexplained, explained, unexplained + explained, round(float(explained / (unexplained + explained)), 3))
# 259/216 1001/1296 2555/1296 0.392: 1.20 + 0.77 = 1.97, and X explains 39%

# the best predictor sin(1.5x) against the best straight line, on simulated data
rng = np.random.default_rng(0)
x = rng.uniform(-2, 2, 1_000_000)
y = np.sin(1.5 * x) + 0.3 * rng.normal(size=x.size)
b = np.cov(x, y)[0, 1] / np.var(x, ddof=1)
print(round(float(np.mean((y - np.sin(1.5 * x)) ** 2)), 3), round(float(b), 3), round(float(np.mean((y - b * x) ** 2)), 3))
# 0.09 0.518 0.254

# the tower property for a random sum: N ~ Poisson(20) customers, each spending 30 on average
N = rng.poisson(20, 20_000)
print(round(float(np.mean([rng.exponential(30, n).sum() for n in N])), 1))   # 599.9, against 20 * 30 = 600

Next#

The section on several random variables ends here: joint distributions, the Gaussian, and the best prediction. Everything so far has assumed the distributions are known. In practice they are seen only through samples, and the question becomes how close an average of samples is to the expectation it estimates. The next note gives the first guarantees, inequalities that bound the probability of a large deviation using only the mean, the variance or the range.

References#

• J. K. Blitzstein, J. Hwang. Introduction to Probability, 2nd ed. CRC Press, 2019. Chapter 9 (conditional expectation, Adam’s law, Eve’s law, prediction).
• T. Hastie, R. Tibshirani, J. Friedman. The Elements of Statistical Learning, 2nd ed. Springer, 2009. Section 2.4 (statistical decision theory: the regression function and the Bayes classifier).
• R. Koenker, G. Bassett. Regression Quantiles. Econometrica, 46(1), 1978.