Signature
linkedinemailgithubgoogleScholar

Conditional Expectation under Squared-Error Loss

On this page

Conditional Expectation: Optimality, Proof, and Applications

Which function of the features minimizes mean squared error (MSE)? If the outcome YY has a finite second moment, the answer is E[Y∣X]\mathbb{E}[Y \mid X], its conditional mean. We prove this by separating prediction error into the outcome's variation around that mean and the predictor's deviation from it. We then use the same reasoning to explain how estimating a predictor from data adds bias and variance.

1. Theorem: Optimality of the Conditional Expectation

Statement:
Let (X,Y)∼P(X,Y)\sim P on Rd×R\mathbb R^d\times\mathbb R, with EP[Y2]<∞\mathbb E_P[Y^2]<\infty. Write PXP_X for the marginal distribution of XX and PY∣X=xP_{Y\mid X=x} for a conditional distribution of YY.

Throughout the proof, E\mathbb E means EP\mathbb E_P. Conditioning on X=xX=x means keeping the input fixed and averaging over Y∼PY∣X=xY\sim P_{Y\mid X=x}. An outer expectation also averages over X∼PXX\sim P_X. Thus E[Y∣X]=m(X)\mathbb E[Y\mid X]=m(X) is a random variable, while m(x)m(x) is its value at a fixed input.

We compare measurable, fixed predictors f:Rd→Rf:\mathbb R^d\to\mathbb R satisfying EX∼PX[f(X)2]<∞\mathbb E_{X\sim P_X}[f(X)^2]<\infty. Their MSE loss is

L(f)=E(X,Y)∼P[(Y−f(X))2].\mathcal{L}(f) = \mathbb{E}_{(X,Y)\sim P}\Bigl[(Y - f(X))^2\Bigr].

Then a minimizing predictor is

f∗(X)=E[Y∣X].f^*(X) = \mathbb{E}[Y \mid X].

The minimizer is unique up to equality almost surely, not necessarily at every possible input. More precisely,

f minimizes L⟺f(X)=E[Y∣X] almost surely.f\text{ minimizes }\mathcal{L} \quad\Longleftrightarrow\quad f(X)=\mathbb{E}[Y\mid X]\text{ almost surely}.

2. Proof of the Theorem

Error Decomposition

To compare a predictor ff with the conditional mean, add and subtract that mean. For any admissible function ff in the class above,

Y−f(X)=(Y−E[Y∣X])⏟Error around the mean+(E[Y∣X]−f(X))⏟Deviation of f(X) from the conditional mean.Y - f(X) = \underbrace{\bigl(Y - \mathbb{E}[Y \mid X]\bigr)}_{\text{Error around the mean}} + \underbrace{\bigl(\mathbb{E}[Y \mid X] - f(X)\bigr)}_{\text{Deviation of } f(X) \text{ from the conditional mean}}.

Squaring both sides gives:

(Y−f(X))2=(Y−E[Y∣X])2+(E[Y∣X]−f(X))2+2(Y−E[Y∣X])(E[Y∣X]−f(X)).\begin{aligned} (Y-f(X))^2 &= \bigl(Y-\mathbb{E}[Y\mid X]\bigr)^2\\ &\quad+\bigl(\mathbb{E}[Y\mid X]-f(X)\bigr)^2\\ &\quad+2\bigl(Y-\mathbb{E}[Y\mid X]\bigr) \bigl(\mathbb{E}[Y\mid X]-f(X)\bigr). \end{aligned}

Taking Expectations and Eliminating the Cross Term

Squaring introduces a cross term. Taking expectations gives:

E[(Y−f(X))2]=E[(Y−E[Y∣X])2]+E[(E[Y∣X]−f(X))2]+2E[(Y−E[Y∣X])(E[Y∣X]−f(X))].\begin{aligned} \mathbb{E}[(Y-f(X))^2] &=\mathbb{E}\bigl[(Y-\mathbb{E}[Y\mid X])^2\bigr]\\ &\quad+\mathbb{E}\bigl[(\mathbb{E}[Y\mid X]-f(X))^2\bigr]\\ &\quad+2\mathbb{E}\bigl[(Y-\mathbb{E}[Y\mid X]) (\mathbb{E}[Y\mid X]-f(X))\bigr]. \end{aligned}

By the linearity of conditional expectation,

E[Y−E[Y∣X] ∣ X]=E[Y∣X]−E[Y∣X]=0,\mathbb{E}\Bigl[Y - \mathbb{E}[Y \mid X] \,\Big|\, X\Bigr] = \mathbb{E}[Y \mid X] - \mathbb{E}[Y \mid X] = 0,

The other factor is a function of XX. Conditioning on XX and applying the tower property therefore removes the cross term:

E[(Y−E[Y∣X])(E[Y∣X]−f(X))]=E[(E[Y∣X]−f(X))E[Y−E[Y∣X]∣X]]=0.\begin{aligned} &\mathbb{E}\bigl[(Y-\mathbb{E}[Y\mid X])(\mathbb{E}[Y\mid X]-f(X))\bigr]\\ &\quad=\mathbb{E}\bigl[(\mathbb{E}[Y\mid X]-f(X)) \mathbb{E}[Y-\mathbb{E}[Y\mid X]\mid X]\bigr]=0. \end{aligned}

The cross term vanishes because the residual has conditional mean zero. The expected squared error therefore separates into:

E[(Y−f(X))2]=E[(Y−E[Y∣X])2]+E[(E[Y∣X]−f(X))2].\begin{aligned} \mathbb{E}[(Y-f(X))^2] &=\mathbb{E}\bigl[(Y-\mathbb{E}[Y\mid X])^2\bigr]\\ &\quad+\mathbb{E}\bigl[(\mathbb{E}[Y\mid X]-f(X))^2\bigr]. \end{aligned}

Minimization

The term E[(Y−E[Y∣X])2]\mathbb{E}\bigl[(Y - \mathbb{E}[Y \mid X])^2\bigr] is independent of ff and represents the irreducible error given XX. The second term is nonnegative and equals zero if and only if f(X)=E[Y∣X]f(X) = \mathbb{E}[Y \mid X] almost surely. Therefore, the unique minimizer in L2L^2 is:

f∗(X)=E[Y∣X].f^*(X) = \mathbb{E}[Y \mid X].

3. Connection to the Bias–Variance Tradeoff

The proof identifies the best prediction target. We next ask how estimating that target from data contributes to prediction error. Let D∼PDD\sim P_D be random training data independent of a new test pair (X,Y)∼P(X,Y)\sim P, and let f^D\hat f_D be the fitted predictor. For nn independent and identically distributed training pairs from PP, PD=P⊗nP_D=P^{\otimes n}.

Write m(x)=EY∼PY∣X=x[Y]m(x)=\mathbb E_{Y\sim P_{Y\mid X=x}}[Y] and σ2(x)=Var⁡Y∼PY∣X=x(Y)\sigma^2(x)=\operatorname{Var}_{Y\sim P_{Y\mid X=x}}(Y) for the conditional mean and variance. Subscripts DD and XX below mean that we average over D∼PDD\sim P_D and X∼PXX\sim P_X, respectively. At a fixed input xx, the joint average is over PD⊗PY∣X=xP_D\otimes P_{Y\mid X=x}: the training data and outcome vary, but xx stays fixed. Assuming the displayed second moments are finite,

ED,Y∣X=x[(f^D(x)−Y)2]=(ED[f^D(x)]−m(x))2⏟Squared bias+Var⁡D(f^D(x))⏟Estimation variance+σ2(x)⏟Irreducible error.\begin{aligned} \mathbb{E}_{D,Y\mid X=x}[(\hat f_D(x)-Y)^2] &=\underbrace{\bigl(\mathbb{E}_D[\hat f_D(x)]-m(x)\bigr)^2}_{\text{Squared bias}}\\ &\quad+\underbrace{\operatorname{Var}_D(\hat f_D(x))}_{\text{Estimation variance}} +\underbrace{\sigma^2(x)}_{\text{Irreducible error}}. \end{aligned}

Total Expected Error Decomposition

To obtain overall test error, average the pointwise decomposition over the test-input distribution. With b(x)=ED[f^D(x)]−m(x)b(x)=\mathbb{E}_D[\hat f_D(x)]-m(x) and v(x)=Var⁡D(f^D(x))v(x)=\operatorname{Var}_D(\hat f_D(x)), this gives

ED,X,Y[(f^D(X)−Y)2]=EX[b(X)2]+EX[v(X)]+EX[σ2(X)].\mathbb{E}_{D,X,Y}[(\hat f_D(X)-Y)^2] =\mathbb{E}_X[b(X)^2]+\mathbb{E}_X[v(X)]+\mathbb{E}_X[\sigma^2(X)].

Optimality of the Conditional Expectation

If the conditional mean f∗(x)=E[Y∣X=x]f^*(x)=\mathbb{E}[Y \mid X=x] were known, no fitting would be needed:

  • Bias:
    Bias(f∗(x))=ED∼PD[f∗(x)]−m(x)=f∗(x)−m(x)=0.\text{Bias}\Bigl(f^*(x)\Bigr) = \mathbb{E}_{D\sim P_D}[f^*(x)] - m(x) = f^*(x)-m(x)=0.
  • Variance:
    At a fixed input, f∗(x)f^*(x) does not depend on the training data, so its estimation variance is zero. This does not mean that the outcome itself has zero variance.
  • Total Error:
    E[(f∗(x)−Y)2∣X=x]=Var⁡(Y∣X=x).\mathbb{E}\Bigl[(f^*(x)-Y)^2\mid X=x\Bigr] = \operatorname{Var}\Bigl(Y \mid X=x\Bigr). This is the minimum possible error achievable, showing that f∗(x)=E[Y∣X=x]f^*(x)=\mathbb{E}[Y \mid X=x] is optimal in the L2L^2 sense.

4. Use Cases and Practical Implications

A. Linear Regression

In linear regression, we restrict the predictor to be linear:

fβ(X)=β0+β1X1+⋯+βpXp.f_\beta(X) = \beta_0 + \beta_1 X_1 + \cdots + \beta_p X_p.

Ordinary least squares (OLS) chooses the coefficient vector β\beta that minimizes squared error in the sample. At the population level, the best linear predictor approximates E[Y∣X]\mathbb{E}[Y\mid X] within the chosen linear class. It need not equal the true conditional mean if that mean is not linear.

The Gauss–Markov theorem concerns coefficient estimation rather than prediction over all functions. Under a correctly specified linear mean, a full-rank design, and errors with zero conditional mean and covariance proportional to the identity, OLS has the smallest covariance among linear unbiased estimators of the coefficient vector. This is the best linear unbiased estimator (BLUE) guarantee; it does not make OLS the best predictor among all functions.

B. Deep Learning and Transformer Models

For a continuous response LL, a neural network trained with squared error uses the sample objective

min⁡θ ∑i(Li−fθ(Xi))2,\min_\theta\, \sum_{i}\Bigl(L_i - f_\theta(X_i)\Bigr)^2,

Under squared loss, the unrestricted population target is E[L∣X]\mathbb{E}[L\mid X]. A fitted network fθ(X)f_\theta(X) need not equal this target: the model class, data, regularization, and optimization determine how well it approximates the conditional mean.

  • Uncertainty Modeling:
    A conditional mean does not specify a full predictive distribution. One possible simulation model adds a sampled residual: Lt(sample)=L^t+εt,L_t^{(\text{sample})} = \widehat{L}_t + \varepsilon_t, where L^t≈E[Lt∣history]\widehat{L}_t \approx \mathbb{E}[L_t \mid \text{history}] is the predicted conditional mean. Conditional on the fitted model and a pool of residuals e1,…,eme_1,\ldots,e_m, drawing εt∼P^ε=m−1∑j=1mδej\varepsilon_t\sim\widehat P_\varepsilon=m^{-1}\sum_{j=1}^m\delta_{e_j} samples uniformly from that empirical distribution; δe\delta_e denotes a point mass at ee. Pooling residuals across histories assumes that the pooled distribution is appropriate at each simulated history. It can miss history-dependent variance or other changes in the error distribution.

C. Counterfactual Outcome Prediction

For a concrete example, G-Transformer uses MSE to fit continuous covariate predictions and samples residuals from a holdout set when simulating trajectories. The conditional-mean target E[Lt∣history]\mathbb{E}[L_t\mid\text{history}] explains the role of squared error in that procedure. Residual sampling adds a distributional modeling step; it is not a consequence of MSE optimality. Moreover, interpreting the simulated trajectories causally requires the paper's identification assumptions, including consistency, sequential exchangeability, and positivity.

Conclusion

For square-integrable outcomes, E[Y∣X]\mathbb{E}[Y\mid X] is the unique MSE-minimizing predictor up to almost-sure equality. The proof separates the variation that remains around this target from the error of a candidate predictor. Once the target must be estimated from data, the bias–variance decomposition separates squared bias and estimation variance from irreducible error. Keeping the target and its fitted approximation distinct is essential when interpreting either a regression model or a neural network.

References and Further Reading