Understanding Scaling Laws in Deep Neural Networks via Feature Learning Dynamics

TL;DR

The paper advances the theoretical understanding of scaling in deep neural networks by deriving Neural Feature Dynamics (NFD), a forward–backward stochastic differential equation describing feature and gradient evolution in the joint infinite-width and infinite-depth limit under depth-$\mu$P. It reveals a depth-driven vanishing mechanism under $1/\sqrt{L}$ residual scaling that makes the gradient-independence assumption provably valid, enabling tractable end-to-end feature learning analyses and explaining diminishing returns. It further shows a learning-collapse in two-layer residual blocks, and proposes a depth-aware learning-rate correction that restores depth-wise hyperparameter transfer and improves performance in deep ResNets. The results provide a principled framework for depth-aware scaling, connect kernel and feature-learning regimes via NFD, and offer practical guidance for training very deep architectures with stable dynamics. Overall, the work bridges mean-field, TP, and stochastic-dynamics perspectives to explain when and why scaling laws succeed or fail, with implications for design and optimization of deep networks.

Abstract

The empirical success of deep learning is often attributed to scaling laws that predict consistent gains as model, data, and compute grow; however, large models can exhibit training instability and diminishing returns, suggesting that scaling laws describe what success looks like but not when and why scaling succeeds or fails. A central obstacle is the lack of a rigorous understanding of feature learning at large depth. While muP characterizes feature-learning dynamics in the infinite-width limit and enables hyperparameter transfer across width, its depth extension (depth-muP) breaks down for residual blocks with more than one internal layer. We derive Neural Feature Dynamics (NFD) for ResNets with single-layer residual blocks, characterizing feature learning via a coupled forward-backward stochastic system in the joint infinite-width and infinite-depth limit. In this regime, NFD identifies when scaling-law trends persist and explains diminishing returns. It also reveals a vanishing mechanism induced by the 1/sqrt(depth) residual scaling under which the gradient-independence assumption (GIA), known to fail during training at finite depth, becomes provably valid again at infinite depth, yielding an analytically tractable regime for end-to-end feature learning. Motivated by this insight, we study two-layer residual blocks and show that the same mechanism causes feature-learning collapse in the first internal layer at large depth, providing a structural explanation for the empirical failure of depth-muP. Based on this diagnosis, we propose a depth-aware learning-rate correction that counteracts the collapse and empirically restores depth-wise hyperparameter transfer, yielding stronger performance in deeper ResNets.

Figures (18)

• Figure 1: Pre- and post-act ResNets debate under depth-$\mu$P. In (a), the pre-act variant maintains stable feature across depth, whereas the post-act exhibits rapid growth. In (b)-(c), we train depth-64 width-128 ResNets with ReLU on CIFAR-10 under SGD (LR 0.01, batch size 128). The stability from pre-act design yields faster convergence and lower test loss with reduced variance across runs.
• Figure 2: Train loss with $L=5$
• Figure 3: Test loss with $L=5$
• Figure 4: Train loss (vary $L,T$)
• Figure 5: Test loss (vary $L,T$)

Theorems & Definitions (49)

• Proposition 1
• Proposition 2
• Proposition 3
• Proposition 4
• Proposition 5
• Proposition 6
• Proposition 7
• Proposition 8
• Definition 1: Neural Feature Dynamics (NFD)
• Theorem 1