Abstract:Layer normalization (LN) is a fundamental component in modern deep learning, but its per-sample centering and scaling introduce non-negligible inference overhead. RMSNorm improves efficiency by removing the centering operation, yet this may discard benefits associated with centering. This paper propose a framework to determine whether an LN in an arbitrary DNN can be replaced by RMSNorm without changing the model function. The key idea is to fold LN's centering operation into upstream general linear layers by enforcing zero-mean outputs through the column-centered constraint (CCC) and column-based weight centering (CBWC). We extend the analysis to arbitrary DNNs, define such LNs as foldable LNs, and develop a graph-based detection algorithm. Our analysis shows that many LNs in widely used architectures are foldable, enabling exact inference-time conversion and end-to-end acceleration of 2% to 12% without changing model predictions. Experiments across multiple task families further show that, when exact equivalence is partially broken in practical training settings, our method remains competitive with vanilla LN while improving efficiency.
| Comments: | 33 pages, 21 figures |
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2605.14521 [cs.LG] |
| (or arXiv:2605.14521v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2605.14521 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Yunhao Ni [view email]
[v1]
Thu, 14 May 2026 08:05:39 UTC (3,780 KB)