Understanding Scaling Laws in Deep Neural Networks via Feature Learning Dynamics

2025-12-24

AI Summary

This work addresses training instability and diminishing returns in deep residual networks, attributing these issues to gradient coupling degradation arising from mismatched feature learning dynamics. We propose the Neural Feature Dynamics (NFD) theoretical framework, which—under the joint infinite-width-and-depth limit—unifies forward and backward propagation as coupled stochastic evolutionary processes. Our analysis reveals that 1/√depth residual scaling restores gradient statistical independence, explaining the structural failure of depth-muP in multi-layer residual blocks; based on this insight, we design a depth-aware learning rate correction scheme. We rigorously prove that the Gradient Independence Assumption (GIA) holds again in the infinite-depth limit under NFD, enabling principled hyperparameter transfer across depths. Empirically, our method significantly improves both generalization and hyperparameter transferability of ResNets.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Natural Language Processing: Learning & Optimization for NLP
• Computer Vision: Learning & Optimization for CV

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Responsible Web: Machine-in-the-loop, human agency and autonomy

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.

Problem

• Analyzes scaling laws in deep neural networks via feature learning dynamics.
• Explains diminishing returns and instability in deep ResNets during scaling.
• Proposes a correction to restore hyperparameter transfer in deep ResNets.

Innovation

• Derives Neural Feature Dynamics for ResNets in infinite limits
• Identifies vanishing mechanism enabling gradient-independence validity
• Proposes depth-aware learning-rate correction to restore hyperparameter transfer