Feature learning networks are/become balanced

Feature learning made simple!

Oct 18, 2025

I’ve been going through the literature on feature-learning linear networks, with the hope that it be useful to non-linear networks. Along the way, I came across “From Lazy to Rich: Exact Learning Dynamics in Deep Linear Networks”(Dominé et al., 2025), and while reading through, it seemed like they had identified a different type of feature learning than that of the maximum-update parameterization ($\mu$P) feature learning that I knew.

For reference $\mu$P (Yang & Hu, 2021) tells you how you need to scale all of your network’s (scale/learning-rate) hyperparameters in order for networks to stabily learn non-trivial functions even at infinite width, with this new perspective based on the ideas of network layers being ‘balanced’ (drawing back all the way to (Du et al., 2018)), with (Kunin et al., 2024) describing how different amounts of balancedness affect feature learning. Note the balancedness picture’s lacks of mentioning stability for large width networks: as long as you make two subsequent layers have ‘a similar scale,’ that’s all you need!

In this post, I want to discuss how the ideas of $\mu$P align with those of deep linear network balancedness in a unified way, and use this perspective to make predictions for how other ‘ML trickery’ that is often used in practice promotes feature learning.

Setup

Unless otherwise stated, all networks will be 2 layer (1 hidden layer) networks of the form $\mathbf{f}(\mathbf{x}) = \sum_{i} \mathbf{a}_i \sigma(\sum_jW_{ij}x_j)$. The input dimension of $d$ ($x \in \mathbb{R}^{d}$) and output dimension of $d_{out}$ ($\mathbf{f}(\mathbf{x})\,, \mathbf{y} \in \mathbb{R}^{d_{out}}$). The nonlinearity is to be taken as the identity function $\sigma: \mathbf{x}\mapsto \mathbf{x}$. The hidden dimension $w$, $W \in \mathbb{R}^{w \times d}$ is a parameter we can scale, where we wish for the network to be trainable for any value of $w$. The network is to be updated iteratively using gradient flow (without addons such as weight decay/momentum/etc.) $\Delta \mathbf{a} = -\eta_a \frac{\partial \mathcal{L}}{\partial \mathbf{a}}$, $\Delta W = -\eta_W \frac{\partial \mathcal{L}}{\partial W}$, ($\eta_a, \eta_W \ll 1$) under a square loss $\mathcal{L} = (\mathbf{f}(\mathbf{x})-\mathbf{y})^2$.¹

Occasionally, I will denote $X \in \mathcal{R}^{N\times d}$ and $Y \in \mathcal{R}^{N\times d_{out}}$ later during discussion. These are matrices with rows consisting of individual samples: $X_i/Y_i = \mathbf{x}^{(i)}/\mathbf{y}^{(i)}$.

How to achieve feature learning

$\mu$P

Based on (Yang et al., 2023), we should set up our network with the following conditions:

where $\mathbf{a}_i \sim \mathcal{N}(0, \sigma_a^2)$, $W_{i,j} \sim \mathcal {N}(0, \sigma_W^2)$. The broad intuition for this setup is such that we have (1) non-trivial updates to our weights (more specifically, the first weight matrix W), else we end up with a NNGP, and (2) not have our outputs blow up during training. Up to moving around the factors (à la mean-field parameterization), this has been identified as the only ($d$, $w$, $d_{out}$)-unique scaling of our network parameters.

Balancedness

If that is true, then why does (Dominé et al., 2025) have a seemingly different rule for achieving feature learning that is removed from the $\mu$P scaling? Their paper finds feature learning for a (deep) linear network can be achieved through the following condition,² modified slightly to accomidate layerwise learning rates:

with data being whitened: $X^\top X = I_d$.

A proof for this quantity remaining invariant can be found in both (Kunin et al., 2024) and below. With Gaussian initialization, this conserved equation is trivially achieved when the width w is large, making

where $\sim$ is used due to sub-leading order corrections not being included; from here on, I will be using $=$, with the understanding that this is to leading order. Thus far, we haven’t needed to worry about making sure update sizes won’t cause trivial or blown-up updates at all due to the trajectory of our network lying on an invariant manifold defined by Eq. 1, so what gives? The trick is in the value that $\lambda$ takes: the closer to 0, the more our solutions appear to be in the feature-learning regime.

While our solution space for feature learning is quite large, with 3 different independent knobs to turn, if we rewrite our feature learning condition in the following way,

for some small $\epsilon$ that’ll represent how close to the “true” feature learning network we get, we can see a ‘natural choice’ for how we should scale our hyperparameters; taking both portions of the equation to be invariant to our network parameters of $(d, d_{out}, w)$ forces $\sigma_a^2 \sim \eta_a/d_{out}$, $\sigma_W^2 \sim \eta_W/d$, representing a perfect match with how (Yang et al., 2023) told us to scale, up to actually choosing the layerwise learning rates and reintroducing a factor of $w$ (which I interpret as a byproduct of moving from a $\lambda$-rescaled diagonal $w\times w$ identity matrix to a single scalar $\lambda$).

For completeness, let’s plug in the $\mu$P scaling laws to see how small $\lambda$ is:

where $f_2$ and $f_1$ represent scalars that are $\Theta(1)$ with respect to the width. This quantity tends towards 0 as $w\rightarrow \infty$ and tells us that any scaling done to $d$ XOR $d_{out}$ should be reflected in an equivalent change to $w$: $w \sim d(_{out})$.

Lazy & ultra-rich scaling

On first glance, we have more freedom in the balancedness picture than the $\mu$P picture, so what gives? Lazy networks have thus far been more approachable from a theoretical perspective, so surely balancedness should be our preferred lens with which to view feature learning. Why are people so much more drawn to $\mu$P, outside of it being the first to introduce many of the feature learning concepts we know today?

Taking from the MFP literature, if we scale our network outputs as $f \mapsto f/\gamma$ and the overall learning rate as $\eta \mapsto \eta\cdot\gamma^2$ for $\gamma \leq 1$ and $\eta \mapsto \eta\cdot\gamma$ otherwise, a scaling law discovered in (Atanasov et al., 2025), we can transition a network into the lazy regime with $\gamma \ll 1$ and into a new ‘hyper-rich’ regime of $\gamma \gg 1$. In the balancedness picture, both the network output prefactor and learning rate rescaling can be absorbed globally; it’s unclear where to put the network output prefactor, but the learning rate can be easily included:

plugging in our $\mu$P relations, this gets

Some insight that comes from this setup is that lazy networks should scale their richness paramater $\gamma$ such that $\gamma \gg \sqrt{w}$, or else we run the risk of not leaving the rich regime. No such insight appears to be able to be gleamed from the ultra-rich case, where both the width and the richness parameter contribute to lowering $\lambda$. Finally, our equation recovers $\lambda \rightarrow 0 (\pm\infty)$ as $\gamma \rightarrow \infty(0)$, as expected.

This suggests that $\mu$P is not the only way to see feature-learning networks; we an also look at how balanced a network is throughout training in order to gain a deeper understanding of its feature learning properties.

Big picture: why (insert ML trick) lets you learn features

Up until now, we’ve considered an extremely simple network with no special bells or whistles. If we want to understand how different ML levers affect feature learning, we need to move beyond this, and now that we have the balancedness-$\mu$P picture, we have a greater capacity to add on extra parts.

Weight decay

Take (L2) weight decay, for example: the purpose is to decrease the weights of a network. Previous papers [links] have investigated …

There doesn’t exist an invariant when training is done with weight decay (except under the explicit case that the layerwise learning rates are identical and additionally the weight decay constant is global), but we can see how our earlier constant $\lambda$ varies with the weight penalty hyperparameter $\zeta$:

where we note that weight decay lowers the overall scale of weights by forcing the weights between layers to become balanced, forcing the network parameters onto a ‘more feature-learning manifold’.

Learning rate warmup

To analyze learning rate increases (and subsequent decays), we need to move away from gradient flow. We’ll still let $\lambda = \lvert\lvert a\rvert\rvert_F^2/\eta_a-\lvert\lvert W\rvert\rvert_F^2/\eta_W$ and now have $\Delta\lambda = \eta_a\lvert\lvert R(WX)^\top\rvert\rvert_F^2-\eta_W\lvert\lvert \mathbf{a}^\top R X^\top\rvert\rvert_F^2$ for residual $R \equiv Y-aWX$. If we define $M = XR^\top$, we can rewrite $\Delta \lambda$ as,

with the general interpretation being for any initial values $\lambda$ far from 0, learning up the learning rate allows the weights to take much faster progress towards $\lambda=0$; this needs to be followed up by a decrease in the learning rate, or else $\lambda$ will end up “bouncing” between the positive and negative of some final value $\pm \lambda_f$.

Deeper network balancedness

If we want to consider deeper networks, the story is largely the same. Take $Y=W_LW_{L-1}\cdots W_2W_1X$. Between two subsequent layers, we have

that is, $W_{\ell+1}^\top W_{\ell+1}$ and $W_{\ell} W_{\ell}^\top$ are simultaneously diagonalizable, giving

while we can’t generally compare $\frac{1}{\eta_i} W_{i}^\top W_{i} - \frac{1}{\eta_j} W_{j} W_{j}^\top$ ($j \neq i-1$) due to possible shape inconsistency, we can say

with the implication being that there is a general balancedness between all layers that feature learning networks satisfy: the $i$-th singular values of all layers in a feature learning network are all roughly equivalent.

Takeaway & Future Efforts

This blogpost was written as a dive into something I found that didn’t make sense: how can balanced nets fit into the $mu$P picture. While discussing my thoughts with Dan Kunin (whose work is the core of this post), we noted how I’m using the known $\mu$P to get balacedness. Right now, I’m working on extending this in the opposite direction: can we get a different perspective on feature learning through balacedness alone, no $\mu$P involved? This is (part of) my current research, so if that sounds interesting, please feel free to reach out!

Appendix

References