Infinite width and depth limits match during neural network training
Width and Depth Limits Commute During Training
Machine Learning
Summary
When building very large neural networks, it matters whether you make them wider (more neurons per layer) or deeper (more layers). This paper shows that, for certain types of networks and scaling rules, the final behavior is the same whether you make them infinitely wide first, infinitely deep first, or both at once. The authors prove this holds not only when the networks start training but also during training, as long as some precautions like gradient clipping are used. Without those precautions, the convergence is weaker but still happens in a probabilistic sense.
What this means in practice
- •For deep learning engineers: Design large residual neural networks with confidence that scaling width and depth leads to predictable outputs during early training.
- •For machine learning framework developers: Implement training algorithms supporting Depth-μP and μP parameterizations to ensure stable training for very wide and deep networks.
A theory result. No direct application yet.
Authors
Ruijia Zhang, Soufiane Hayou
Abstract
Recent work established that the infinite-width and infinite-depth limits of residual networks commute at initialization when residual branches are scaled by $1/\sqrt{depth}$. In particular, the network output converges to the same limit whether width and depth tend to infinity sequentially or jointly. We extend this commutativity result to training in linear residual networks. Under the same $1/\sqrt{depth}$ residual-branch scaling (Depth-$μ$P) and the $μ$P width parameterization, we prove that, for any fixed number of gradient-descent steps with gradient clipping, the network output converges in $L^2$ along every joint width--depth growth path as both width and depth tend to infinity. Consequently, the infinite-width and infinite-depth limits commute throughout any finite training horizon. Without gradient clipping, the corresponding $L^2$ convergence does not hold in general, but we establish the weaker convergence in probability.