# Critical Initialization of Wide and Deep Neural Networks through Partial Jacobians: General Theory and Applications

arXiv: https://arxiv.org/abs/2111.12143 · HTML: https://mlsanity.com/paper/2111.12143/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 55 of its 71 references that are in the map.
- First posted on arXiv 23 Nov 2021.

## Builds on (most cited first, 30 of 55)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [RandAugment: Practical automated data augmentation with a reduced search space](https://mlsanity.com/paper/1909.13719/) (2019)
- [Identity Mappings in Deep Residual Networks](https://mlsanity.com/paper/1603.05027/) (2016)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [MLP-Mixer: An all-MLP Architecture for Vision](https://mlsanity.com/paper/2105.01601/) (2021)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent](https://mlsanity.com/paper/1902.06720/) (2019)
- [Deep Neural Networks as Gaussian Processes](https://mlsanity.com/paper/1711.00165/) (2017)
- [On Exact Computation with an Infinitely Wide Neural Net](https://mlsanity.com/paper/1904.11955/) (2019)
- [MetaFormer Is Actually What You Need for Vision](https://mlsanity.com/paper/2111.11418/) (2021)
- [Gaussian Process Behaviour in Wide Deep Neural Networks](https://mlsanity.com/paper/1804.11271/) (2018)
- [Stronger generalization bounds for deep nets via a compression approach](https://mlsanity.com/paper/1802.05296/) (2018)
- [ReZero is All You Need: Fast Convergence at Large Depth](https://mlsanity.com/paper/2003.04887/) (2020)
- [Exponential expressivity in deep neural networks through transient chaos](https://mlsanity.com/paper/1606.05340/) (2016)
- [Understanding and Improving Layer Normalization](https://mlsanity.com/paper/1911.07013/) (2019)
- [Deep Information Propagation](https://mlsanity.com/paper/1611.01232/) (2016)
- [Sensitivity and Generalization in Neural Networks: an Empirical Study](https://mlsanity.com/paper/1802.08760/) (2018)
- [Dynamical Isometry and a Mean Field Theory of CNNs: How to Train 10,000-Layer Vanilla Convolutional Neural Networks](https://mlsanity.com/paper/1806.05393/) (2018)
- [The Principles of Deep Learning Theory](https://mlsanity.com/paper/2106.10165/) (2021)
- [Bayesian Deep Convolutional Networks with Many Channels are Gaussian Processes](https://mlsanity.com/paper/1810.05148/) (2018)

## Built on by (most cited first, 8 of 8)

- [Geometric Dynamics of Signal Propagation Predict Trainability of Transformers](https://mlsanity.com/paper/2403.02579/) (2024)
- [Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos](https://mlsanity.com/paper/2311.02076/) (2023)
- [Effective Theory of Transformers at Initialization](https://mlsanity.com/paper/2304.02034/) (2023)
- [Two failure modes of deep transformers and how to avoid them: a unified theory of signal propagation at initialisation](https://mlsanity.com/paper/2505.24333/) (2025)
- [Field theory for optimal signal propagation in ResNets](https://mlsanity.com/paper/2305.07715/) (2023)
- [Gradient Flossing: Improving Gradient Descent through Dynamic Control of Jacobians](https://mlsanity.com/paper/2312.17306/) (2023)
- [AutoInit: Automatic Initialization via Jacobian Tuning](https://mlsanity.com/paper/2206.13568/) (2022)
- [Towards Distributed Neural Architectures](https://mlsanity.com/paper/2506.22389/) (2025)

## Abstract (arXiv)

Deep neural networks are notorious for defying theoretical treatment. However, when the number of parameters in each layer tends to infinity, the network function is a Gaussian process (GP) and quantitatively predictive description is possible. Gaussian approximation allows one to formulate criteria for selecting hyperparameters, such as variances of weights and biases, as well as the learning rate. These criteria rely on the notion of criticality defined for deep neural networks. In this work we describe a new practical way to diagnose criticality. We introduce partial Jacobians of a network, defined as derivatives of preactivations in layer l with respect to preactivations in layer l0≤ l. We derive recurrence relations for the norms of partial Jacobians and utilize these relations to analyze criticality of deep fully connected neural networks with LayerNorm and/or residual connections. We derive and implement a simple and cheap numerical test that allows one to select optimal initialization for a broad class of deep neural networks; containing fully connected, convolutional and normalization layers. Using these tools we show quantitatively that proper stacking of the LayerNorm (applied to preactivations) and residual connections leads to an architecture that is critical for any initialization. Finally, we apply our methods to analyze ResNet and MLP-Mixer architectures; demonstrating the everywhere-critical regime.