Technical Summary: Enjoy Your Layer Normalization with the Computational Efficiency of RMSNorm

Problem Statement

Layer Normalization (LN) is a foundational component in modern deep learning architectures, particularly Transformers, where it stabilizes training and improves optimization. However, during inference, LN requires per-sample centering (subtracting the mean) and scaling, which introduces non-negligible computational overhead compared to normalization-free alternatives.

RMSNorm was proposed as a more efficient alternative by removing the centering operation, retaining only scaling. While RMSNorm offers significant speedups (reported reductions of 7%–64% in running time in prior work), it discards the centering operation. This removal may lead to less stable output ranges and larger activation variations, potentially discarding benefits associated with centering established in prior literature (LeCun et al., 1990; Schraudolph, 1998; Ba et al., 2016). Previous attempts to equate LN and RMSNorm (e.g., Jiang et al., 2023) were restricted to specific pre-norm Transformer structures, leaving open the question of whether a general framework exists for arbitrary Deep Neural Networks (DNNs) to replace LN with RMSNorm without altering model function or training dynamics.

Methodology

The paper proposes a general framework to determine if an LN layer in an arbitrary DNN can be replaced by RMSNorm while maintaining mathematical equivalence in both inference and training. The core strategy involves "folding" the centering operation of LN into upstream linear layers.

1. Column-Centered Constraint (CCC) and Column-Based Weight Centering (CBWC)

The authors first analyze the case where LN immediately follows a linear layer. They observe that if the linear layer's output has a zero mean across neurons, the centering step in LN becomes redundant.

• Column-Centered Constraint (CCC): A linear layer satisfies CCC if the elements in each column of its weight matrix sum to zero. Under this constraint, the output of the linear layer is guaranteed to have a zero mean for any input.
• Column-Based Weight Centering (CBWC): To enforce CCC during training without restricting the model's capacity, the authors introduce a reparameterization technique. A proxy parameter matrix $W$ is used, which is transformed into an effective weight matrix $V$ via a projection: $V = (I - \frac{1}{m}\mathbf{1}_m\mathbf{1}_m^\top)W$. This transformation ensures the effective weights satisfy the zero-sum column constraint while preserving the optimization dynamics of the original model.

Proposition: A model using a linear layer with CBWC followed by RMSNorm is mathematically equivalent to a model using the same linear layer (without reparameterization) followed by standard LN.

2. Generalization to Arbitrary DNNs: Foldable LNs

The framework is extended to arbitrary network structures using a graph-based analysis.

• Foldable LN: An LN is defined as "foldable" if its input can be guaranteed to have a zero mean through constraints on upstream layers.
• Zero-Mean Graph: The authors construct a directed graph representing the data flow to an LN. By backtracking from the LN, they identify "leaf nodes" (upstream layers).
  • If all leaf nodes correspond to general linear layers (including convolutions and RNNs with shared weights) or layers that inherently guarantee zero-mean outputs, the LN is foldable.
  • Scalar operations (e.g., dropout in inference mode, temperature scaling) preserve the zero-mean property and do not break foldability.
  • Residual connections (addition) preserve foldability if each branch is individually foldable.
  • Non-linear operations (e.g., ReLU, Softmax) or concatenations generally break the zero-mean property, preventing folding unless auxiliary centering is inserted.

3. Detection Algorithm

The paper presents an algorithm (Algorithm 1) that performs a forward traversal of the computation graph to automatically detect foldable LNs and identify the specific upstream general linear layers that require CBWC. This process is computationally lightweight and requires only a single forward pass.

Key Contributions

1. General Framework: A theoretical framework to determine the replaceability of LN with RMSNorm in arbitrary DNNs, extending beyond the limitations of previous work restricted to pre-norm Transformers.
2. Mathematical Equivalence: The introduction of CBWC, which allows for the exact folding of the LN centering operation into upstream weights, ensuring that the CBWC+RMSNorm scheme produces identical outputs and gradients to the original LN-based model during both training and inference.
3. Foldable LN Detection: The definition of "foldable LNs" and the development of a graph-based detection algorithm that identifies which LNs in a given architecture can be converted.
4. Practical Extension: A discussion on handling cases where exact folding conditions are not met (e.g., due to dropout during training), suggesting that CBWC+RMSNorm remains competitive with vanilla LN even when exact equivalence is partially broken.

Experimental Results

The authors validated their method across multiple architectures (GPT-2, BERT, OPT, Phi, ViT, Swin-Tiny) and tasks (text classification, translation, image classification, fine-tuning).

• Inference Acceleration:

  • The method achieves end-to-end inference speedups of 2% to 12% on representative models (e.g., GPT-2, BERT, OPT) without changing model predictions.
  • The acceleration is more pronounced with longer sequence lengths (up to 24.95% speedup for BERT at 16,384 tokens).
  • Throughput increases significantly (e.g., +1,119 to +3,385 tokens/s on A100 GPUs).
  • The overhead of the folding algorithm is minimal (amortized after a few inference steps) and does not increase memory footprint.

• Training and Fine-Tuning Performance:

  • In scenarios where exact folding conditions hold (e.g., inference), the method is identical to the baseline.
  • In practical training settings where conditions like dropout break exact equivalence, CBWC+RMSNorm consistently performs between standard LN and vanilla RMSNorm.
  • On image classification (ImageNet-100 with Swin-Tiny), the method achieved the best test accuracy and loss among the three variants (LN, RMSNorm, CBWC+RMS).
  • In text generation and fine-tuning tasks (GPT-3 on Alpaca), the method achieved the best test loss and perplexity while offering efficiency gains.

Significance and Claims

The paper claims to provide a principled way to reduce the computational cost of Layer Normalization while preserving its theoretical and practical benefits as much as possible.

• Exact Equivalence: For "foldable" LNs, the method offers a mathematically exact replacement, enabling deployment of RMSNorm in models originally trained with LN without retraining or performance degradation.
• Efficiency: It delivers measurable end-to-end inference speedups (2%–12%) on widely used architectures, which is significant for large-scale deployments of Large Language Models (LLMs) and Vision-Language Models (VLMs).
• Robustness: Even when exact folding conditions are not fully met (e.g., during training with dropout), the method remains competitive with standard LN, offering a favorable trade-off between fidelity and efficiency.
• Generality: Unlike previous approaches, this framework applies to arbitrary DNNs, not just specific Transformer variants, and includes a detection algorithm to automate the process.

The authors acknowledge that their analysis covers representative architectures but does not exhaust all modern large-scale models. They note that exact equivalence may fail in architectures with modules like dropout during training, suggesting future work could refine treatments for such modules.