Understanding Scaling Laws in Deep Neural Networks via Feature Learning Dynamics (2512.21075v1)

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.

• The paper constructs the first mathematical framework, Neural Feature Dynamics (NFD), to analyze feature and gradient evolution in joint infinite-width and infinite-depth regimes.
• It reveals new phenomena, including restored gradient independence and a depth-induced vanishing mechanism, which clarify the limits of scaling and hyperparameter transfer.
• The study proposes a depth-aware learning-rate correction for multi-layer residual blocks, empirically improving train/test loss and mitigating capacity saturation.

Feature Learning Dynamics and the Rigorous Foundations of Scaling Laws in Deep Neural Networks

Introduction and Motivation

The empirical observation that performance in deep learning increases predictably with neural network scale—model parameters, dataset size, and compute—has led to the prominent concept of scaling laws. These laws have shaped the development of large-scale architectures such as LLMs and ViTs, but they fail to pinpoint when and why the scaling-related gains break down due to training instabilities or diminishing returns. The core unresolved challenge is a rigorous, analytical understanding of feature learning—how representations evolve during training—in extremely deep architectures. The paper "Understanding Scaling Laws in Deep Neural Networks via Feature Learning Dynamics" (2512.21075) addresses this by constructing the first mathematically tractable framework, Neural Feature Dynamics (NFD), for analyzing feature and gradient evolution in the joint infinite-width and infinite-depth regime.

Theoretical Landscape: From Kernel Regimes to Feature Learning

Initial approaches to scaling laws often relied on simplified models predicting test loss via kernel eigenspectra, generally holding features fixed during training. The NTK theory formalized a regime where networks behave akin to kernel methods, training features remain essentially static ("lazy regime"), and scaling effects can be accurately captured via Gaussian process correspondence. However, NTK offers no framework for the nontrivial feature learning (FL) phenomena that characterize breakthroughs such as in-context learning, multi-modal reasoning, and chain-of-thought prompting.

Mean-field theory, maximal update parameterization (pP), and the Tensor Program (TP) framework advanced the state-of-the-art by mathematically characterizing feature evolution in infinite-width networks with nontrivial dynamics. While pP enables hyperparameter (HP) transfer and active feature learning across width, its extension to depth ("depth-pP") yields stable training only for single-layer residual blocks. In deeper blocks, depth-pP breaks down, motivating various corrective scalings (e.g., $1/L$) but lacking rigorous FL analysis at infinite depth.

Neural Feature Dynamics (NFD): Mechanistic and Mathematical Foundations

The paper introduces NFD as a novel coupled forward-backward stochastic differential equation (SDE) system that captures feature and gradient dynamics during SGD in ResNets with joint infinite-depth and infinite-width scaling. Key insights include:

• Stability of Feature Norms and Architectural Choice: Pre-activation ResNets, as opposed to post-activation designs, maintain stable feature norms under $1/\sqrt{L}$ residual scaling with common nonlinearities (e.g., ReLU), admitting a well-posed infinite-depth limit.
• Induced SDEs for Feature and Gradient Propagation: In the limit, each coordinate's trajectory converges to McKean-Vlasov SDEs for forward and backward propagation, providing sharp convergence rates $O(1/L + 1/n)$ and commutativity between depth and width limits for any Lipschitz nonlinearity.
• Capacity Ceiling and the Role of the Time Horizon: Model capacity is ultimately limited by the NFD SDE. Increasing width or depth only reduces the approximation error to this limit and cannot expanding the representational class; increasing the time horizon parameter $T$ can enlarge capacity but at the cost of reduced training stability.
• Gaussian Process Correspondence and RKHS Universality: In the kernel regime ($a = 1$), infinite-depth-and-width networks converge to a Gaussian process with a strictly positive definite, universal kernel. Enlarging $T$ yields nested RKHSs, contractively contained within each other, mapping out the possible gains and inherent saturation in scaling.

Depth-Induced Decoupling: Gradient Independence Assumption and Vanishing Mechanisms

A core technical achievement is the restoration of the Gradient Independence Assumption (GIA) during training at infinite depth. While GIA is rigorously valid at initialization in the infinite-width regime, it provably fails during training at finite depth due to forward-backward coupling. NFD reveals a new depth-induced vanishing mechanism: under $1/\sqrt{L}$ residual scaling, the forward and backward SDEs become driven by independent Brownian motions, dynamically suppressing correlations and making GIA valid again in the infinite-depth limit.

Empirically, this leads to convergence of standard and decoupled dynamics (where backward weights are independently sampled), improved performance, and restored analytical tractability for FL dynamics well beyond what NTK or mean-field theories alone provide.

Internal Learning Collapse in Multi-layer Residual Blocks and Corrective Scaling

For architectures with multi-layer residual blocks, including Transformers, a structural failure emerges: feature learning in the first internal layer vanishes at large depth, while the second layer (governing the residual-stream dynamics) remains active. This mechanistic insight explains long-standing empirical failures in depth-wise HP transfer for blocks with more than one internal layer.

To address this, the authors propose a depth-aware learning-rate correction: scaling the first internal layer's learning rate as $\eta_1 = \eta_0 \sqrt{L}$ while maintaining the standard scaling for the second layer. This correction restores nontrivial updates, effective feature learning, and robust HP transfer, as confirmed by consistent improvements in train/test loss and accuracy on CIFAR-10 benchmarks.

Implications for Neural Scaling, HP Transfer, and Capacity Saturation

The NFD formalism yields several strong results and implications:

• Empirical Restoration of GIA: Depth-wise pP and NFD enable HP transfer across widths and depths, aligning both standard and decoupled training trajectories and yielding consistent empirical improvements for deep ResNets.
• Capacity Saturation and Diminishing Returns: Performance gains derived from width and depth scaling saturate at the NFD limit, explaining observed diminishing returns in practice. The only avenue for further capacity expansion is through increasing $T$, subject to stability constraints.
• Rigorous Foundation for Scaling Analysis: The coupled SDE system enables the analysis of nontrivial feature learning across training trajectories, opening the door to mathematically rigorous feature-learning theory in deep architectures, including multiblock structures, broader parameterizations, and alternative optimizers.

Conclusion

This work establishes Neural Feature Dynamics (NFD) as a principled, mathematically rigorous mechanism for analyzing training dynamics and scaling laws in deep neural networks. NFD not only encompasses and generalizes prior kernel and mean-field analyses but also reveals new phenomena, such as the restoration of gradient independence and internal learning collapse in multiblock architectures. The depth-aware learning-rate correction provides a practical prescription for enabling HP transfer and stable scaling in architectures where existing regimes fail.

Theoretically, NFD underscores the intrinsic ceiling to scaling derived from the underlying SDE limit and the necessity for refined scaling strategies and parameterizations. Practically, these insights inform the design of training protocols and hyperparameter schedules for very deep networks, with implications for future developments in neural architecture scaling, optimization, and generative modeling.

Future research directions should include extending NFD to non-ResNet architectures (e.g., Transformers with attention), complex multi-modal or multi-objective training protocols, and the investigation of SGD variants or momentum-based optimizers in the infinite-scale regime. The rigorous characterization afforded by NFD sets the stage for more systematic understanding and predictability in the scaling and training of next-generation deep models.