# Robust Training of Neural Networks Using Scale Invariant Architectures

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

## Key facts

- Cited by 29 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 43 references that are in the map.
- First posted on arXiv 2 Feb 2022.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (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)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [Know What You Don't Know: Unanswerable Questions for SQuAD](https://mlsanity.com/paper/1806.03822/) (2018)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [Instance Normalization: The Missing Ingredient for Fast Stylization](https://mlsanity.com/paper/1607.08022/) (2016)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [On the Convergence of Adam and Beyond](https://mlsanity.com/paper/1904.09237/) (2019)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networks](https://mlsanity.com/paper/1602.07868/) (2016)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [Why gradient clipping accelerates training: A theoretical justification for adaptivity](https://mlsanity.com/paper/1905.11881/) (2019)
- [The large learning rate phase of deep learning: the catapult mechanism](https://mlsanity.com/paper/2003.02218/) (2020)
- [Understanding the Difficulty of Training Transformers](https://mlsanity.com/paper/2004.08249/) (2020)
- [Understanding Gradient Clipping in Private SGD: A Geometric Perspective](https://mlsanity.com/paper/2006.15429/) (2020)
- [L2 Regularization versus Batch and Weight Normalization](https://mlsanity.com/paper/1706.05350/) (2017)
- [An Exponential Learning Rate Schedule for Deep Learning](https://mlsanity.com/paper/1910.07454/) (2019)
- [Three Mechanisms of Weight Decay Regularization](https://mlsanity.com/paper/1810.12281/) (2018)
- [Beyond Convexity: Stochastic Quasi-Convex Optimization](https://mlsanity.com/paper/1507.02030/) (2015)
- [Theoretical Analysis of Auto Rate-Tuning by Batch Normalization](https://mlsanity.com/paper/1812.03981/) (2018)
- [Norm matters: efficient and accurate normalization schemes in deep networks](https://mlsanity.com/paper/1803.01814/) (2018)
- [Reconciling Modern Deep Learning with Traditional Optimization Analyses: The Intrinsic Learning Rate](https://mlsanity.com/paper/2010.02916/) (2020)
- [The Power of Normalization: Faster Evasion of Saddle Points](https://mlsanity.com/paper/1611.04831/) (2016)
- [On the Periodic Behavior of Neural Network Training with Batch Normalization and Weight Decay](https://mlsanity.com/paper/2106.15739/) (2021)

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

- [Fine-Tuning Language Models with Just Forward Passes](https://mlsanity.com/paper/2305.17333/) (2023)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [Understanding Gradient Descent on Edge of Stability in Deep Learning](https://mlsanity.com/paper/2205.09745/) (2022)
- [A Kernel-Based View of Language Model Fine-Tuning](https://mlsanity.com/paper/2210.05643/) (2022)
- [Understanding the Generalization Benefit of Normalization Layers: Sharpness Reduction](https://mlsanity.com/paper/2206.07085/) (2022)
- [Why Do We Need Weight Decay in Modern Deep Learning?](https://mlsanity.com/paper/2310.04415/) (2023)
- [Implicit Bias of AdamW: ℓ_∞ Norm Constrained Optimization](https://mlsanity.com/paper/2404.04454/) (2024)
- [Understanding Edge-of-Stability Training Dynamics with a Minimalist Example](https://mlsanity.com/paper/2210.03294/) (2022)
- [Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks](https://mlsanity.com/paper/2305.17212/) (2023)
- [Replacing softmax with ReLU in Vision Transformers](https://mlsanity.com/paper/2309.08586/) (2023)
- [How to set AdamW's weight decay as you scale model and dataset size](https://mlsanity.com/paper/2405.13698/) (2024)
- [CAME: Confidence-guided Adaptive Memory Efficient Optimization](https://mlsanity.com/paper/2307.02047/) (2023)
- [Symmetry in Neural Network Parameter Spaces](https://mlsanity.com/paper/2506.13018/) (2025)
- [Universality and Limitations of Prompt Tuning](https://mlsanity.com/paper/2305.18787/) (2023)
- [Weight Decay may matter more than muP for Learning Rate Transfer in Practice](https://mlsanity.com/paper/2510.19093/) (2025)
- [Training Scale-Invariant Neural Networks on the Sphere Can Happen in Three Regimes](https://mlsanity.com/paper/2209.03695/) (2022)
- [HSR-Enhanced Sparse Attention Acceleration](https://mlsanity.com/paper/2410.10165/) (2024)
- [Toward Equation of Motion for Deep Neural Networks: Continuous-time Gradient Descent and Discretization Error Analysis](https://mlsanity.com/paper/2210.15898/) (2022)
- [The Feature Speed Formula: a flexible approach to scale hyper-parameters of deep neural networks](https://mlsanity.com/paper/2311.18718/) (2023)
- [On Learning Verifiers and Implications to Chain-of-Thought Reasoning](https://mlsanity.com/paper/2505.22650/) (2025)
- [In Good GRACEs: Principled Teacher Selection for Knowledge Distillation](https://mlsanity.com/paper/2511.02833/) (2025)
- [A Minimalist Example of Edge-of-Stability and Progressive Sharpening](https://mlsanity.com/paper/2503.02809/) (2025)
- [Convergence of variational Monte Carlo simulation and scale-invariant pre-training](https://mlsanity.com/paper/2303.11602/) (2023)
- [Efficient Language Model Architectures for Differentially Private Federated Learning](https://mlsanity.com/paper/2403.08100/) (2024)
- [Optimized Speculative Sampling for GPU Hardware Accelerators](https://mlsanity.com/paper/2406.11016/) (2024)
- [Towards the Transferable Audio Adversarial Attack via Ensemble Methods](https://mlsanity.com/paper/2304.08811/) (2023)
- [Softplus Attention with Re-weighting Boosts Length Extrapolation in Large Language Models](https://mlsanity.com/paper/2501.13428/) (2025)
- [Twin: Tuning Learning Rate and Weight Decay of Deep Homogeneous Classifiers without Validation](https://mlsanity.com/paper/2403.05532/) (2024)
- [Predictability-Aware Compression and Decompression Framework for Multichannel Time Series Data with Latent Seasonality](https://mlsanity.com/paper/2506.00614/) (2025)

## Abstract (arXiv)

In contrast to SGD, adaptive gradient methods like Adam allow robust training of modern deep networks, especially large language models. However, the use of adaptivity not only comes at the cost of extra memory but also raises the fundamental question: can non-adaptive methods like SGD enjoy similar benefits? In this paper, we provide an affirmative answer to this question by proposing to achieve both robust and memory-efficient training via the following general recipe: (1) modify the architecture and make it scale invariant, i.e. the scale of parameter doesn't affect the output of the network, (2) train with SGD and weight decay, and optionally (3) clip the global gradient norm proportional to weight norm multiplied by √tfrac2λη, where η is learning rate and λ is weight decay. We show that this general approach is robust to rescaling of parameter and loss by proving that its convergence only depends logarithmically on the scale of initialization and loss, whereas the standard SGD might not even converge for many initializations. Following our recipe, we design a scale invariant version of BERT, called SIBERT, which when trained simply by vanilla SGD achieves performance comparable to BERT trained by adaptive methods like Adam on downstream tasks.