ResNets of All Shapes and Sizes: Deciphering the Large-Scale Training Limit

ResNets of All Shapes and Sizes: Convergence of Training Dynamics in the Large-scale Limit

Summary
Problem
Method
Results
Takeaways
Abstract

This paper establishes the first rigorous quantitative convergence proof for the training dynamics of Residual Neural Networks (ResNets) in the joint infinite depth (L), width (M), and embedding dimension (D) limit. The authors identify the Maximal Local feature Update (MLU) regime and derive an error rate of O(1/L + √D/√ML + 1/√D) between finite models and their large-scale limit.

Executive Summary

TL;DR: This paper provides a rigorous mathematical proof that ResNet training dynamics converge to a stable "infinite-size" limit as depth (), hidden width (), and embedding dimension () all grow to infinity. Moving beyond heuristic arguments, it establishes an empirically tight error bound of , proving that modern massive models are disciplined approximations of a well-defined underlying dynamical system.

This work sits at the intersection of Mean-Field Theory and Neural ODEs, effectively bridging the gap between width-scaling and depth-scaling research. It moves us from "what happens when we scale" to a quantitative "how fast do we converge to the ideal limit."

The Scaling Dilemma: Why Is This Hard?

In the quest for larger models, researchers have observed that "shape matters." Simply increasing parameters isn't enough; the ratio between width and depth influences everything from training stability to hyperparameter sensitivity.

Prior theoretical frameworks like the Neural Tangent Kernel (NTK) simplified neural networks into linear models, but they failed to capture feature learning—the process where neurons adapt their representations to the data. While newer "Maximal Update" (P) regimes allowed for feature learning, they were primarily focused on width () or depth () in isolation. The synergy of scaling , and simultaneously introduces a nightmare of correlations: every weight update depends on the features, which in turn depend on the history of all previous weight updates.

Methodology: The Core Innovations

The authors tackle this using two sophisticated mathematical structures: Skeleton Maps and the Cavity Method.

1. Skeleton Maps: Tracking the Path

To handle the training dynamics over time, the authors define "Skeleton Maps." These are functions that express weight updates () as dependencies on past hidden states. By viewing the entire training process as a path-dependent system, they can formulate the limit as a fixed-point of these maps.

2. The Cavity Method: Capturing Correlations

The hardest part of the proof is the Central Limit Theorem (CLT) scaling. In the MLU regime, hidden states tend toward Gaussian processes. Standard CLT fails because the weights and features are not independent. The authors use the Cavity Method (borrowed from spin-glass physics) to isolate the influence of a single coordinate (the "cavity") and perform a Taylor expansion to show that the correlations effectively "cancel out" or converge at a rate of .

Model Architecture and Training Dynamics The fundamental ResNet structure investigated, using depth-two blocks which are standard in modern architectures.

Experimental Results & Tightness

The theory isn't just abstract—it's precise. The authors tested their predicted rates on ResNets of various sizes.

  • Embedding Space: The error rate perfectly matches the predicted .
  • Output Space: Interestingly, the authors found their bound is loose for the final output, conjecturing a faster convergence of .
  • Optimal Prescription: For a fixed parameter budget , the math suggests that the widths () should scale like the square of the depth (). This mirrors the design of models like Llama 3, where width is prioritized over extreme depth.

RMS Error vs Dimension D Figure 1: Empirical validation showing the tightness of the convergence rate in embedding space.

Hidden State Trajectories Figure 3: Hidden representations in a finite ResNet (histograms) versus the Gaussian limit (red line) over training iterations .

Critical Analysis & Conclusion

While the paper is a masterclass in rigour, it has its limits:

  • Time and Sample Scaling: The constants grow super-exponentially with the number of training steps () and samples (). Improving these dependencies is a clear path for future work.
  • Activation Functions: It requires smooth activation functions (up to 5 derivatives for the full proof), which technically excludes the raw ReLU, though smooth approximations (like GeLU or SwiGLU) fit perfectly.

Takeaway: This work signifies a turning point where we can finally treat the "alphabet soup" of model shapes () as a unified field. It provides the mathematical "North Star" for scaling up the next generation of architectures, ensuring that as models grow, their training remains predictable and efficient.

Find Similar Papers

Try Our Examples

  • Find recent papers addressing the "Maximal Local feature Update" (MLU) regime or μP scaling laws in Transformer architectures.
  • What are the original papers proposing the "Cavity Method" in the context of Dynamical Mean Field Theory (DMFT) for neural networks?
  • Search for studies that investigate the training instabilities of very deep ResNets and how they relate to the large-scale limits described here.
Contents
ResNets of All Shapes and Sizes: Deciphering the Large-Scale Training Limit
1. Executive Summary
2. The Scaling Dilemma: Why Is This Hard?
3. Methodology: The Core Innovations
3.1. 1. Skeleton Maps: Tracking the Path
3.2. 2. The Cavity Method: Capturing Correlations
4. Experimental Results & Tightness
5. Critical Analysis & Conclusion