AutoNorm: Understanding Adaptive Normalization in Transformers through Differentiable Gating
2026-07-12 • Machine Learning
Machine Learning
AI summaryⓘ
The authors studied how different methods of normalizing data inside Transformers affect training. They found that using a gating method with high randomness works better on changing language tasks but struggles on stable vision tasks because it makes learning unstable. To fix this, they introduced AutoNorm-S, which gradually stops the gating randomness, helping the model train more reliably. Their method performs well on several natural language and vision benchmarks, showing it can adapt normalization without causing training problems.
TransformerNormalizationLayer NormalizationGumbel-SoftmaxAdaptive NormalizationGradient VarianceRouting MechanismLanguage ModelingVision TasksGate Freezing
Authors
Piyush Kaushik Bhattacharyya, Divyanshu Rai, Swastik Singh, Kumar Aakash, Ayush Ranjan, Krutika Verma
Abstract
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.