Enjoy Your Layer Normalization with the Computational Efficiency of RMSNorm

TL;DR

This paper introduces a framework to replace layer normalization with RMSNorm in deep neural networks, enabling faster inference without changing model outputs by identifying foldable LNs.

Contribution

It proposes a graph-based detection algorithm to identify foldable LNs, allowing exact or approximate conversion to RMSNorm for efficiency gains.

Findings

Exact conversion yields 2-12% inference speedup.

Many LNs in popular architectures are foldable.

The method remains competitive even when exact equivalence is not maintained.

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.