# AutoNorm: Understanding Adaptive Normalization in Transformers through Differentiable Gating

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 30 of its 36 references that are in the map.
- First posted on arXiv 12 Jul 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [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)
- [Squeeze-and-Excitation Networks](https://mlsanity.com/paper/1709.01507/) (2017)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features](https://mlsanity.com/paper/1905.04899/) (2019)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [FiLM: Visual Reasoning with a General Conditioning Layer](https://mlsanity.com/paper/1709.07871/) (2017)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [DARTS: Differentiable Architecture Search](https://mlsanity.com/paper/1806.09055/) (2018)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Instance Normalization: The Missing Ingredient for Fast Stylization](https://mlsanity.com/paper/1607.08022/) (2016)
- [In Search of Lost Domain Generalization](https://mlsanity.com/paper/2007.01434/) (2020)
- [Deep Networks with Stochastic Depth](https://mlsanity.com/paper/1603.09382/) (2016)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [Going deeper with Image Transformers](https://mlsanity.com/paper/2103.17239/) (2021)
- [Dynamic Convolution: Attention over Convolution Kernels](https://mlsanity.com/paper/1912.03458/) (2019)
- [AdaBelief Optimizer: Adapting Stepsizes by the Belief in Observed Gradients](https://mlsanity.com/paper/2010.07468/) (2020)
- [ReZero is All You Need: Fast Convergence at Large Depth](https://mlsanity.com/paper/2003.04887/) (2020)
- [Escaping the Big Data Paradigm with Compact Transformers](https://mlsanity.com/paper/2104.05704/) (2021)
- [Fixup Initialization: Residual Learning Without Normalization](https://mlsanity.com/paper/1901.09321/) (2019)
- [DeepNet: Scaling Transformers to 1,000 Layers](https://mlsanity.com/paper/2203.00555/) (2022)
- [Transformers without Normalization](https://mlsanity.com/paper/2503.10622/) (2025)
- [Differentiable Learning-to-Normalize via Switchable Normalization](https://mlsanity.com/paper/1806.10779/) (2018)
- [Activate or Not: Learning Customized Activation](https://mlsanity.com/paper/2009.04759/) (2020)

## Abstract (arXiv)

Normalization is a critical component for stabilizing Transformer training, yet the choice between static strategies such as Layer Normalization (LN) and adaptive alternatives remains largely task-dependent. In this paper, we investigate a key optimization challenge in differentiable normalization gating. Our experiments show that, on relatively stationary vision tasks, the high gradient variance introduced by Gumbel-Softmax gating can hinder convergence of the routing mechanism, causing learned gates to underperform simple random selection. In contrast, on non-stationary language modeling and classification tasks, sustained gating diversity enables the model to learn more effective layer-wise normalization policies. Motivated by these observations, we propose AutoNorm-S (Stabilized), a training strategy that mitigates optimization instability through a gate-freezing schedule. AutoNorm-S achieves competitive or improved performance across multiple benchmarks, outperforming adaptive normalization baselines on NLP datasets, including PTB and SST-2, while remaining competitive on standard vision benchmarks. These results suggest that decoupling normalization selection from optimization noise provides a practical and principled approach for adaptive normalization in Transformer architectures.