In-Sample Moments "Generalize" under Overfitting

As a hitchhiker in learning theory, I constantly bump into a weird kind of problems, where an overfitted hypothesis still predicts certain population quantities accurately. One specific instance is about the 2nd moment of prediction which will be today’s topic.¹ After being haunted by the issue for years,² I finally figured out a way to nicely characterize the phenomenon.

---

Problem Setup

Consider a standard least-square regression problem where we have i.i.d. sample $X_{1:n}, Y_{1:n}$. Let $\hat{f}$ be ERM from function class $\mathcal{F}$, i.e.,

$$\hat{f} \in \arg\min_{f\in\mathcal{F}} L_n(f) := \mathbb{E}_n[(f(X)-Y)^2],$$

where $\mathbb{E}_n$ denotes empirical expectation. Let’s assume realizability, i.e., $f^\star \in \mathcal{F}$ where $f^\star(X) = \mathbb{E}[Y\mid X]$ is the Bayes-optimal predictor. All $f \in \mathcal{F}$ and $Y$ are bounded in $[0, 1]$. The quantity of interest is

$$\mathbb{E}_n[\hat{f}^2].$$

The goal is to show that $\mathbb{E}_n[\hat{f}^2]$ is bounded by something not large. (As we will see, this something may not be fixed!)

---

Standard Approach

The standard approach is to show that $\mathbb{E}_n[\hat{f}^2] \approx \mathbb{E}[(f^\star)^2]$ since $\hat f \to f^\star$ with enough data. To ensure this, we need to pay $\textup{comp}(\mathcal{F})$, the complexity of the function class (e.g., $\log \vert \mathcal{F} \vert$ for finite classes, or dimensionality $d$ for linear): with high probability,

$$\mathbb{E}[(\hat{f} - f^\star)^2] \lesssim \frac{\textup{comp}(\mathcal{F})}{n}.$$

Define $\|\cdot\|_n^2 = \mathbb{E}_n[(\cdot)^2]$, then $\mathbb{E}_n[\hat{f}^2] = \|\hat{f}\|_n^2$, and

$$\|\hat{f}\|_n \le \|f^\star\|_n + \|\hat{f} - f^\star\|_n.$$

The second term is controlled by $\sqrt{\textup{comp}(\mathcal{F})/n}$. For the first term, we can use standard concentration to relate $\mathbb{E}[(f^\star)^2]$ to $\mathbb{E}_n[(f^\star)^2]$ and this term does not involve $\textup{comp}(\mathcal{F})$ since it’s w.r.t. a fixed function $f^\star$. As $n$ increases, the upper bound converges to $\mathbb{E}[(f^\star)^2]$.

---

The Overfitting Regime

The problem with the above bound is that when $\mathcal{F}$ is extremely expressive, the bound can grow indefinitely with $\textup{comp}(\mathcal{F})$ all the way to infinity. However, if we completely overfit, then we’d have $\hat{f}(X_i) = Y_i$, and

$$\mathbb{E}_n[\hat{f}^2] = \mathbb{E}_n[Y^2],$$

which clearly concentrates towards $\mathbb{E}[Y^2]$ in a $\textup{comp}(\mathcal{F})$-free manner! Believe or not, I have run into a number of scenarios like this, where an upper bound close to $\mathbb{E}[Y^2]$ is totally acceptable.³ (NB: this is a different target compared to the previous analysis where the upper bound is close to $\mathbb{E}[(f^\star)^2]$.) What I’d really love to have is a single bound that covers both cases: in the good-fit regime ($\textup{comp}(\mathcal{F})$ small), the bound should be close to $\mathbb{E}[(f^\star)^2]$); in the over-fit regime ($\textup{comp}(\mathcal{F})$ large), the bound is close to $\mathbb{E}[Y^2]$). Better still, perhaps the bound can even provide characterization for the continuous spectrum in between.

---

Unified Analysis through Empirical Pythagoreans

I gave this weird problem some thought every time I ran into it, but always gave up due to its messiness and lack of urgency. Thanks to LLMs, I was able to dig into it deeply recently and figured out a solution. The key is the empirical Pythagoreans that⁴

$$\mathbb{E}_n[Y^2] = \mathbb{E}_n[\hat{f}^2] + \mathbb{E}_n[(Y - \hat{f})^2],$$

which holds as long as $\mathcal{F}$ is closed under scaling, i.e., $a f \in \mathcal{F}, \forall f\in\mathcal{F}$ for constant scalar $a$. To see why this is true, note that the identity will generally be violated by a cross-term

$$2\mathbb{E}_n[(Y - \hat{f})\hat{f}].$$

On the other hand, given $\hat{f}$ is ERM, we can consider the derivative of the empirical loss $L_n(\cdot)$ along the direction of $a \hat{f}$, which should be $0$ at $a=1$:

$$0 = \frac{\partial \mathbb{E}_n[(Y - a\hat{f})^2]}{\partial a} \Big\vert_{a=1} = 2 \mathbb{E}_n[(Y - a \hat{f}) \hat{f}] \big\vert_{a=1} = 2\mathbb{E}_n[(Y - \hat{f})\hat{f}],$$

which establishes the empirical Pythagoreans. Then

$$\mathbb{E}_n[\hat{f}^2] = \mathbb{E}_n[Y^2] - \mathbb{E}_n[(Y - \hat{f})^2]. \qquad (1)$$

For the first term, $\mathbb{E}_n[Y^2]$ concentrates to $\mathbb{E}[Y^2]$ and this clearly does not involve $\mathcal{F}$; in fact, we need to pay this anyway even in the perfect overfitting case ($\hat{f}(X_i) = Y_i$). The second term is the empirical loss of $\hat{f}$, $L_n(\hat{f})$; let $L(\cdot)$ be the population loss. Therefore, to obtain an upper bound on $\mathbb{E}_n[\hat{f}^2]$, we need a lower bound on $L_n(\hat{f})$. One trivial (but later useful) lower bound is $0$. Another comes from bounding

$$\lvert L_n(\hat{f}) - L(f^\star)\rvert \le \lvert L_n(\hat{f}) - L(\hat{f})\rvert + \lvert L(\hat{f}) - L(f^\star) \rvert \lesssim \sqrt{\textup{comp}(\mathcal{F})/n},$$

as the first term is bounded by concentration + union bound, and the second can be similarly bounded by introducing $- L_n(\hat{f}) + L_n(f^\star)$. Therefore, the second lower bound of $L_n(\hat{f})$ is

$$L_n(\hat{f}) \ge L(f^\star) - O(\sqrt{\textup{comp}(\mathcal{F})/n}).$$

Put together into Eq.(1), we have

$$\mathbb{E}_n[\hat{f}^2] \le \mathbb{E}[Y^2] - \max\left\{0, L(f^\star) - O(\sqrt{\textup{comp}(\mathcal{F})/n})\right\} + O(1/\sqrt{n}).$$

This perfectly captures both the good-fit and over-fit regimes: when $\textup{comp}(\mathcal{F})$ is very large, the $0$ is active inside $\max$ and we have an upper bound close to $\mathbb{E}[Y^2]$; when $\textup{comp}(\mathcal{F})$ is small, $L(f^\star)$ becomes active; when combined with $\mathbb{E}[Y^2]$, this recovers

$$\mathbb{E}[Y^2] - L(f^\star) = \mathbb{E}[(f^\star)^2].$$

As conclusion, the “$\max\{0, \ldots\}$” is not as elegant as I hoped, but given the numerous failures I am pretty satisfied with the current form. Also it does provide characterization for the intermediate regime where $\sqrt{\textup{comp}(\mathcal{F})/n}$ gradually decreases below $L(f^\star)$. Finally, I believe the analysis also holds in the misspecified case where we replace $f^\star$ with the population best fit (which may not be unique), as long as we have closure under rescaling and $\mathbb{E}[f^\star(Y - f^\star)]=0$. This enables the population version of Pythagoreans for $f^\star$ which is the only property we need.

Footnote Remarks