# GradInit: Learning to Initialize Neural Networks for Stable and Efficient Training

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

## Key facts

- Cited by 44 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 42 references that are in the map.
- First posted on arXiv 16 Feb 2021.

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

- [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)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Densely Connected Convolutional Networks](https://mlsanity.com/paper/1608.06993/) (2016)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Wide Residual Networks](https://mlsanity.com/paper/1605.07146/) (2016)
- [Improved Regularization of Convolutional Neural Networks with Cutout](https://mlsanity.com/paper/1708.04552/) (2017)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [fairseq: A Fast, Extensible Toolkit for Sequence Modeling](https://mlsanity.com/paper/1904.01038/) (2019)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [On the Variance of the Adaptive Learning Rate and Beyond](https://mlsanity.com/paper/1908.03265/) (2019)
- [Exact solutions to the nonlinear dynamics of learning in deep linear neural networks](https://mlsanity.com/paper/1312.6120/) (2013)
- [signSGD: Compressed Optimisation for Non-Convex Problems](https://mlsanity.com/paper/1802.04434/) (2018)
- [High-Performance Large-Scale Image Recognition Without Normalization](https://mlsanity.com/paper/2102.06171/) (2021)
- [ReZero is All You Need: Fast Convergence at Large Depth](https://mlsanity.com/paper/2003.04887/) (2020)
- [Understanding the Difficulty of Training Transformers](https://mlsanity.com/paper/2004.08249/) (2020)
- [Fixup Initialization: Residual Learning Without Normalization](https://mlsanity.com/paper/1901.09321/) (2019)
- [Dissecting Adam: The Sign, Magnitude and Variance of Stochastic Gradients](https://mlsanity.com/paper/1705.07774/) (2017)
- [All you need is a good init](https://mlsanity.com/paper/1511.06422/) (2015)
- [Traditional and Heavy-Tailed Self Regularization in Neural Network Models](https://mlsanity.com/paper/1901.08276/) (2019)
- [Characterizing signal propagation to close the performance gap in unnormalized ResNets](https://mlsanity.com/paper/2101.08692/) (2021)
- [Theoretical Analysis of Auto Rate-Tuning by Batch Normalization](https://mlsanity.com/paper/1812.03981/) (2018)
- [Spherical Motion Dynamics: Learning Dynamics of Neural Network with Normalization, Weight Decay, and SGD](https://mlsanity.com/paper/2006.08419/) (2020)

## Built on by (most cited first, 30 of 44)

- [H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models](https://mlsanity.com/paper/2306.14048/) (2023)
- [MetaFormer Baselines for Vision](https://mlsanity.com/paper/2210.13452/) (2022)
- [The Devil is in the Detail: Simple Tricks Improve Systematic Generalization of Transformers](https://mlsanity.com/paper/2108.12284/) (2021)
- [Parameter Prediction for Unseen Deep Architectures](https://mlsanity.com/paper/2110.13100/) (2021)
- [NormFormer: Improved Transformer Pretraining with Extra Normalization](https://mlsanity.com/paper/2110.09456/) (2021)
- [Learning to Grow Pretrained Models for Efficient Transformer Training](https://mlsanity.com/paper/2303.00980/) (2023)
- [Cramming: Training a Language Model on a Single GPU in One Day](https://mlsanity.com/paper/2212.14034/) (2022)
- [Hyper-Representations as Generative Models: Sampling Unseen Neural Network Weights](https://mlsanity.com/paper/2209.14733/) (2022)
- [Sparse MoE as the New Dropout: Scaling Dense and Self-Slimmable Transformers](https://mlsanity.com/paper/2303.01610/) (2023)
- [Fast Certified Robust Training with Short Warmup](https://mlsanity.com/paper/2103.17268/) (2021)
- [No Train No Gain: Revisiting Efficient Training Algorithms For Transformer-based Language Models](https://mlsanity.com/paper/2307.06440/) (2023)
- [A Loss Curvature Perspective on Training Instability in Deep Learning](https://mlsanity.com/paper/2110.04369/) (2021)
- [How to Train Vision Transformer on Small-scale Datasets?](https://mlsanity.com/paper/2210.07240/) (2022)
- [Can We Scale Transformers to Predict Parameters of Diverse ImageNet Models?](https://mlsanity.com/paper/2303.04143/) (2023)
- [Each Test Image Deserves A Specific Prompt: Continual Test-Time Adaptation for 2D Medical Image Segmentation](https://mlsanity.com/paper/2311.18363/) (2023)
- [Training Your Sparse Neural Network Better with Any Mask](https://mlsanity.com/paper/2206.12755/) (2022)
- [On Efficient Training of Large-Scale Deep Learning Models: A Literature Review](https://mlsanity.com/paper/2304.03589/) (2023)
- [Learngene: Inheriting Condensed Knowledge from the Ancestry Model to Descendant Models](https://mlsanity.com/paper/2305.02279/) (2023)
- [Complexity Control Facilitates Reasoning-Based Compositional Generalization in Transformers](https://mlsanity.com/paper/2501.08537/) (2025)
- [Initialization is Critical to Whether Transformers Fit Composite Functions by Reasoning or Memorizing](https://mlsanity.com/paper/2405.05409/) (2024)
- [NAR-Former: Neural Architecture Representation Learning towards Holistic Attributes Prediction](https://mlsanity.com/paper/2211.08024/) (2022)
- [Towards Theoretically Inspired Neural Initialization Optimization](https://mlsanity.com/paper/2210.05956/) (2022)
- [PVG: Progressive Vision Graph for Vision Recognition](https://mlsanity.com/paper/2308.00574/) (2023)
- [Transferring Core Knowledge via Learngenes](https://mlsanity.com/paper/2401.08139/) (2024)
- [From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics](https://mlsanity.com/paper/2510.06954/) (2025)
- [From Deterioration to Acceleration: A Calibration Approach to Rehabilitating Step Asynchronism in Federated Optimization](https://mlsanity.com/paper/2112.09355/) (2021)
- [BranchNorm: Robustly Scaling Extremely Deep Transformers](https://mlsanity.com/paper/2305.02790/) (2023)
- [AutoInit: Analytic Signal-Preserving Weight Initialization for Neural Networks](https://mlsanity.com/paper/2109.08958/) (2021)
- [FINE: Factorizing Knowledge for Initialization of Variable-sized Diffusion Models](https://mlsanity.com/paper/2409.19289/) (2024)
- [Principled Architecture-aware Scaling of Hyperparameters](https://mlsanity.com/paper/2402.17440/) (2024)

## Abstract (arXiv)

Innovations in neural architectures have fostered significant breakthroughs in language modeling and computer vision. Unfortunately, novel architectures often result in challenging hyper-parameter choices and training instability if the network parameters are not properly initialized. A number of architecture-specific initialization schemes have been proposed, but these schemes are not always portable to new architectures. This paper presents GradInit, an automated and architecture agnostic method for initializing neural networks. GradInit is based on a simple heuristic; the norm of each network layer is adjusted so that a single step of SGD or Adam with prescribed hyperparameters results in the smallest possible loss value. This adjustment is done by introducing a scalar multiplier variable in front of each parameter block, and then optimizing these variables using a simple numerical scheme. GradInit accelerates the convergence and test performance of many convolutional architectures, both with or without skip connections, and even without normalization layers. It also improves the stability of the original Transformer architecture for machine translation, enabling training it without learning rate warmup using either Adam or SGD under a wide range of learning rates and momentum coefficients. Code is available at https://github.com/zhuchen03/gradinit.