TL;DR. We prove the Functional Scaling Law for SGD in infinite-dimensional linear regression by tracking coordinate errors. Each update reduces old errors and adds new ones. Under power-law assumptions and sufficiently small learning rates, this gives a loss formula for the whole learning-rate schedule. The key trade-off is simple: larger steps learn faster but add more error, and later steps determine how much of that error remains.
Problem setup. We consider an infinite-dimensional linear regression model, with coordinates indexed by $j\in\{1,2,\ldots\}$. Let $\boldsymbol x=(x_j){j\ge1}$ be the input, $\boldsymbol w*$ the target coefficients, and $\boldsymbol w$ the trainable coefficients. Here, the input coordinates are independent, with $x_j\sim N(0,\lambda_j)$ and $\lambda_j$ is the variance of coordinate $j$. Predictions and labels are
$$ f_{\boldsymbol w}(\boldsymbol x)=\sum_j w_jx_j, \qquad y=\sum_j w_{*,j}x_j+\xi. $$
The label noise $\xi\sim N(0,\sigma^2)$ is independent of the input. We allow $\sigma=0$. Assume $\sum_j\lambda_j<\infty$ and $\sum_j\lambda_jw_{*,j}^2<\infty$, so these labels $y$ are well-defined.
Start from $\boldsymbol w_0=\boldsymbol0$. At step $k$, draw a fresh pair $(\boldsymbol x_k,y_k)$. Assume squared loss, the update rule of SGD gives
$$ \boldsymbol w_{k+1} =\boldsymbol w_k-\eta_k \bigl(f_{\boldsymbol w_k}(\boldsymbol x_k)-y_k\bigr)\boldsymbol x_k. $$
In this section, we define the coordinate error and show that the test loss decomposes into the sum of coordinate errors and an irreducible label noise term.
<aside>
Definition 1 (Coordinate error). Let $e_{k,j}$ be the parameter error in coordinate $j$. Let $v_{k,j}$ be its expected square:
$$ e_{k,j}:=w_{k,j}-w_{*,j}, \qquad v_{k,j}:=\mathbb E[e_{k,j}^2]. $$
The expectation averages over the training samples used so far.
</aside>
Clearly, the coordinate error comes from the mismatch between our learned parameters and the target parameter $\boldsymbol{w}^\star$, and thus naturally contributes to the test loss.
<aside>
Lemma 1 (Expected excess test loss). Let $(\boldsymbol x,y)$ be a fresh test sample. The expected excess test loss is given by
$$ \begin{equation} \begin{aligned} L_k &:=\frac12\mathbb E\!\left[ \bigl(f_{\boldsymbol w_k}(\boldsymbol x)-y\bigr)^2 \right]-\frac{\sigma^2}{2}.\\ &=\frac12\sum_j\lambda_jv_{k,j}. \end{aligned} \end{equation} $$
The total expected test loss is $L_k+\sigma^2/2$.
</aside>
In this section, we characterize how the coordinate error evolves during SGD updates.
What one-step SGD changes. We first show that after one SGD update, each coordinate keeps part of its old error and receives a new error from the sampled gradient.
<aside>
Lemma 2 (Dynamics of coordinate error).
$$ \begin{equation} v_{k+1,j} =\underbrace{(1-2\eta_k\lambda_j+2\eta_k^2\lambda_j^2)v_{k,j}}{\text{old error that remains}} +\underbrace{\eta_k^2\lambda_j(2L_k+\sigma^2)}{\text{new error from this step}}. \end{equation} $$
</aside>
What n-step SGD gives. We then extend Lemma 2 to $n$ updates. For simplicity, we define
$$ a_{k,j}:=1-2\eta_k\lambda_j+2\eta_k^2\lambda_j^2, \qquad b_{k,j}:=\eta_k^2\lambda_j(2L_k+\sigma^2). $$