How to align data scaling curves under different initialization scales

Phenomenon

In our previous post (Can We Derive Scaling Law From First Principles), we discussed how scaling laws actually arise. We added an experimental section to its preprint, containing a series of figures that show data scaling curves under different $\alpha$; one of them, for example, shows the relationship between loss and data size in the data-limited regime.

To ensure feature learning, in this figure we set the initialization std=0.01. But what happens if we increase the std to 0.05? The results are as follows.

We can see that with std=0.05, the data scaling curve shifts. What if we increase it further, say to 0.1? The results are as follows.

As the initialization std increases, the empirical line gradually deviates from the theoretical prediction. What if we plot the empirical slope as a function of std? The results are as follows.

This result is quite interesting. There could be many possible reasons, such as the std being too large and pushing the model into the lazy learning regime, or the increased std making optimization unstable (this second possibility is open to question, because I tested different numbers of epochs and the curve in the figure above was still reproduced reliably). But we can set the cause aside for now and first consider how to align the data scaling curves under different initialization scales.

Preliminaries

We consider a two-layer ReLU network whose forward computation is written as $g(W_1,W_2;x)=W_2\phi(W_1 x)$, where $\phi=\mathrm{ReLU}$ is positively homogeneous: for any $c>0$,

The input $x\in\mathbb R^K$ is set to be one-hot or a normalized vector satisfying $\lVert x\rVert^2=1$. The hidden layer first forms the pre-activation $h=W_1 x$, then passes it through ReLU to get the activation $a=\phi(h)=\mathrm{ReLU}(h)$, and the final output is $y=W_2 a\in\mathbb R^K$; the hidden layer width is denoted by $N$. Both weight matrices $W_1,W_2$ are initialized i.i.d. from $\mathcal N(0,\sigma^2)$. For any matrix $A\in\mathbb R^{d_{\mathrm{out}}\times d_{\mathrm{in}}}$ whose entries have mean $0$ and variance $\sigma^2$, the expected squared Frobenius norm satisfies $\mathbb E[\lVert A\rVert_F^2]=d_{\mathrm{out}}\cdot d_{\mathrm{in}}\cdot\sigma^2$, so in order of magnitude $\lVert A\rVert_F \propto \sigma$.

Attempt 1 (Failed)

With the learning rate lr fixed, increasing the initialization std leaves the absolute step size unchanged but shrinks the step size relative to $\lVert W\rVert_F$, and weight updates that are too small drive the model into the lazy learning regime. From this starting point, we need to compute the relative update ratio R:

This is the denominator part, which is straightforward. For the second layer weight $W_2 \in \mathbb{R}^{K \times N}$:

Taking the root mean square as the magnitude estimate:

Next, we estimate $\lVert\nabla W_2\rVert_F$. Each element $h_j$ of the pre-activation $h$ follows $\mathcal{N}(0,\sigma^2)$, so $\text{Var}(h_j) = \sigma^2$. For the activation output a: $a_j = \text{ReLU}(h_j)$. Since ReLU zeroes out the negative half-axis, the second moment is halved:

Then the expected squared norm of the activation vector, $\lVert a\rVert^2 = \sum_{j=1}^N a_j^2$, is:

For the second layer output y: $y_k = \sum_{j=1}^N W_{2, kj} a_j$. Assuming $W_2$ and $a$ are independent, and $\mathbb{E}[W_2]=0$, then

Substituting the known terms:

Then $\lVert y\rVert \approx \sqrt{\frac{N}{2}} \sigma^2 \propto \sqrt{N} \sigma^2$

The loss function is $L = \frac{1}{2} \lVert y - t\rVert^2$. The gradient formula is:

When $\sigma$ is large, the initial output $\lVert y\rVert$ is much larger than the target $\lVert t\rVert$ (which in our setting is one-hot, with magnitude 1). Therefore, the error term $\epsilon = y - t \approx y$. Now compute the norm of the gradient matrix $\nabla_{W_2} L$:

Substituting the result from step 2, $\lVert y - t\rVert \approx \lVert y\rVert \propto \sqrt{N} \sigma^2$, and $\lVert a\rVert \propto \sqrt{N} \sigma$. Multiplying gives the gradient magnitude:

Now we substitute the above results into the relative update ratio formula:

Substituting the magnitude relations:

Simplifying (ignoring the constant $K$):

To keep the relative update ratio constant, we need $\mathcal{R} = \text{Constant}$, i.e.:

Since in our experiments $N$ is fixed and only the initialization std is varied, we only need $\eta \propto \frac{1}{\sigma^2}$. So I immediately implemented this adjustment in the code. The results are as follows

Let’s put them in the same plot for comparison:

We can see that the adaptive learning rate $\eta \propto \frac{1}{\sigma^2}$ does mitigate the effect of initialization std on the empirical slope of the scaling curve to some extent. But when the initialization std is small, the effect is still hard to mitigate. I think the reason might be that the approximation \(\lVert\nabla W_2\rVert_F \approx N \sigma^3\) holds only when $\sigma$ is large (see the derivation above), so in the small initialization std regime, $\eta \propto \frac{1}{\sigma^2}$ may not be correct.

Why Attempt 1 Failed (probably)

Beyond the range in which the approximation holds, we can also analyze it as follows. Under the setting in our paper, writing the parameters as $W=\sigma U$, the risk corresponding to the model is

Note that the target “1” is fixed and does not scale with $\sigma$ (in most cases, the true label also does not scale with $\sigma$), and this step already breaks scale invariance.

Performing gradient descent on $W$

Switching to $U^t=W^t/\sigma$, we can derive the exact update rule

There is a factor that cannot be eliminated:

It depends on the current $U^t$ and the sample $k$, and is not a global constant. Overall, no scalar $\eta(\sigma)$ can achieve alignment. $\eta\propto 1/\sigma^2$ does slightly better in some regimes (because it aligns one dominant scale, such as the leading NTK term). But it cannot turn training at every $\sigma$ into the same problem; systematic drift and instability remain.

Attempt 2 (Successful)

What exactly needs to be done to align the data scaling curves for all $\sigma$?

To cancel out the $\sigma^2$ scaling that automatically appears in the output during the forward pass, we set

and set the learning rate $\eta_\sigma\propto\sigma^2$.

If we write $W=\sigma U$, then we have

The right-hand side contains no $\sigma$ at all.

Define the objective function

Let $U=W/\sigma$. Then it is easy to obtain:

1. The objective function itself does not contain $\sigma$:

1. Gradient scaling (chain rule):

1. Gradient descent trajectory:

Define $U^t:=W^t/\sigma$, then

Substituting $\nabla_W\mathcal J_\sigma(\sigma U^t)=\frac{1}{\sigma}\nabla_U\mathcal J_1(U^t)$, we get

This is the key point: we need to choose $\eta_\sigma=\sigma^2\eta_0$ to obtain

The generalization to momentum SGD is also strictly equivalent.

In summary, we need to set

Initialization:

Learning rate selection:

Only in this way can the loss-$D$ curves for different initialization stds be aligned. The experimental results are shown below.

This result can be easily generalized to multi-layer ReLU networks. Suppose the network satisfies:

• $h_0=x$
• $h_\ell=\phi(W_\ell h_{\ell-1})$, $\ell=1,\dots,K-1$
• $f(x;W)=W_K h_{K-1}$

All layers are initialized as $W_\ell\sim \mathcal N(0,\sigma^2)$. Define the normalized model output:

As long as we choose:

we can align the forward and backward passes at the same time.

Citation

If you need to cite this article, please refer to:

@article{zou2026scaling,
  title={如何对齐不同初始化大小下的 Data scaling 曲线},
  author={Zou, Jiaxuan},
  journal={Jiaxuan's Blog},
  year={2026},
  url={https://jiaxuanzou0714.github.io/blog/2026/data-scaling-and-std/}
}