PaperlayerView exact revision on arXiv

Source-linked AI summary

Root Mean Square Layer Normalization

Biao Zhang, Rico Sennrich

arXiv:1910.07467v1cs.LGcs.CLstat.ML

TL;DR

LayerNorm stabilizes training but incurs substantial computational overhead, creating a need for more efficient normalization. The paper proposes RMSNorm, which retains re-scaling invariance without re-centering, and reports comparable performance with 7%∼64% speed-ups across models.

  • Problem

    LayerNorm improves training stability and convergence, but its computational overhead significantly slows underlying networks, particularly RNNs.

  • Method

    RMSNorm normalizes summed neuron inputs using RMS alone, retaining re-scaling invariance; pRMSNorm estimates RMS from a subset of inputs.

  • Results

    Across diverse tasks and models, RMSNorm achieves performance comparable to LayerNorm while improving running speed by 7%∼64%.

  • Takeaways & Limitations

    RMSNorm is an efficient drop-in replacement for LayerNorm across different model architectures.

  • Takeaways & Limitations

    pRMSNorm can exhibit exploding gradients with small m, although satisfactory convergence is observed with a 6.25% partial ratio.

Abstract

from arXiv · show

Layer normalization (LayerNorm) has been successfully applied to various deep neural networks to help stabilize training and boost model convergence because of its capability in handling re-centering and re-scaling of both inputs and weight matrix. However, the computational overhead introduced by LayerNorm makes these improvements expensive and significantly slows the underlying network, e.g. RNN in particular. In this paper, we hypothesize that re-centering invariance in LayerNorm is dispensable and propose root mean square layer normalization, or RMSNorm. RMSNorm regularizes the summed inputs to a neuron in one layer according to root mean square (RMS), giving the model re-scaling invariance property and implicit learning rate adaptation ability. RMSNorm is computationally simpler and thus more efficient than LayerNorm. We also present partial RMSNorm, or pRMSNorm where the RMS is estimated from p% of the summed inputs without breaking the above properties. Extensive experiments on several tasks using diverse network architectures show that RMSNorm achieves comparable performance against LayerNorm but reduces the running time by 7%~64% on different models. Source code is available at https://github.com/bzhangGo/rmsnorm.

1 Introduction

LayerNorm stabilizes training but adds computational overhead, motivating RMSNorm, which keeps re-scaling invariance while removing re-centering. Across diverse tasks, RMSNorm matches LayerNorm comparably while improving running speed.

  • LayerNorm stabilizes deep-network training through mean and variance statistics and supports diverse architectures and variable-length RNN sequences.
  • LayerNorm’s computational overhead can offset its faster convergence, especially as networks become larger and deeper.
  • RMSNorm regularizes neuron inputs using the root mean square statistic alone.
  • RMSNorm preserves re-scaling invariance while estimating RMS from only a subset of summed inputs for partial RMSNorm.
  • 7%∼64%: RMSNorm’s speed-up across different models, with performance comparable to LayerNorm across machine translation, image classification, image-caption retrieval, and question answering.

2 Related Work

Prior normalization methods improve convergence or control parameterization, but often trade lower training-step counts for higher per-step cost. RMSNorm instead removes re-centering while retaining re-scaling invariance to improve efficiency.

  • BatchNorm stabilizes activations using mini-batch statistics, whereas LayerNorm estimates statistics within the same layer and handles variable-length sequences better.
  • Normalization methods can shorten convergence while consuming more time per running step, motivating data-independent and batch-renormalization alternatives.
  • Internal covariate shift is presented as a motivation for normalization, while later work also attributes normalization’s success to controlling activation growth.
  • RMSNorm differs from related methods by removing re-centering, retaining re-scaling invariance, and reducing computational overhead.

3 Background

The background formulates a feed-forward network around weight-summed inputs and describes LayerNorm as normalizing those inputs using mean and variance statistics before activation scaling.

  • A feed-forward network maps input vector x to output vector y through a linear transformation followed by an element-wise nonlinearity.
  • The weight-summed input a contains the neuron-wise quantities targeted for normalization, with w_i as the weight vector and b_i as the bias.
  • LayerNorm normalizes summed inputs to fix their mean and variance, addressing input-distribution changes that can delay convergence.
  • LayerNorm uses gain g to rescale standardized inputs, while μ and σ^2 are estimated from raw summed inputs.
  • LayerNorm decouples the norm of neurons from the inputs and weight matrix.

4 RMSNorm

RMSNorm removes LayerNorm’s re-centering operation and normalizes summed inputs using RMS alone, retaining re-scaling invariance while simplifying computation. The analysis links these invariance properties to activation and gradient stability.

  • RMSNorm hypothesizes that re-scaling invariance, rather than re-centering invariance, explains LayerNorm’s success.
  • RMSNorm regularizes summed inputs using the root mean square statistic alone.
  • RMSNorm removes mean normalization, sacrificing invariance to variable shifts; when summed-input means are zero, it equals LayerNorm.
  • 4.1 Invariance Analysis: RMSNorm is invariant to re-scaling of both the weight matrix and inputs, while not being invariant to all re-centering operations.
  • 4.2 Gradient Analysis: RMSNorm gradients for g and b are invariant to input and weight scaling, while normalized inputs stabilize g’s gradient magnitude.
  • 4.2 Gradient Analysis: RMSNorm reduces sensitivity of ∂L/∂W to input scaling and uses negative correlation with weight scaling as an implicit learning-rate adaptor.

5 pRMSNorm

pRMSNorm estimates the RMS statistic from a subset of summed inputs, using the same invariance principle as RMSNorm. Although small subsets can cause gradient instability, a 6.25% partial ratio can still achieve satisfactory convergence.

  • pRMSNorm estimates RMS from a subset of neurons, motivated by their assumed independent and identically distributed structure.
  • pRMSNorm retains RMSNorm’s invariance properties because RMS’s linearity property still holds for the partial estimate.
  • 6.25% partial RMS estimation achieves satisfactory convergence in practice despite theoretical approximation to RMSNorm.
  • Small m can produce gradient instability, with gradients tending to explode, because partial RMS is often an inaccurate biased estimate.

6 Experiments

Across diverse architectures and tasks, RMSNorm generally matches LayerNorm’s performance while reducing computational cost and training time. Experiments also examine partial RMS estimation, robustness to initialization, and activation statistics.

  • Experimental setup: RMSNorm was evaluated across RNN, convolutional, self-attentional, and other models implemented in TensorFlow, PyTorch, and Theano.Comparisons included unnormalized baselines, LayerNorm, and in some experiments BatchNorm, WeightNorm, or L2-Norm.
  • 6.1 Machine Translation: About 50% fewer training steps were needed for RNNSearch with LayerNorm or RMSNorm than with the baseline, while RMSNorm’s test accuracy was comparable to LayerNorm.In TensorFlow RNNSearch, LayerNorm was about 67% slower than the baseline, whereas RMSNorm improved speed over LayerNorm by approximately 25%.
  • 6.1 Machine Translation: 11%∼34% speedups over LayerNorm were obtained for RMSNorm in Theano and PyTorch RNNSearch while translation quality remained comparable.pRMSNorm with p = 6.25% was sometimes slower than RMSNorm because of non-optimal tensor-slicing implementations.
  • 6.1 Machine Translation: Changing the partial RMS ratio had little influence on final RNNSearch translation quality, so subsequent experiments used p = 6.25%.The partial ratio controls RMS-estimation accuracy and can affect training stability.
  • 6.1 Machine Translation: RMSNorm achieved BLEU scores comparable to LayerNorm and 7%∼9% speedups in Transformer experiments.The unnormalized Transformer baseline failed to train, and the relative normalization cost was lower because Transformer has fewer sequential normalization operations.
  • Activation Statistics: Both RMSNorm and LayerNorm stabilized hidden-state standard deviation, while RMSNorm’s unnormalized mean was more stable in practice than the baseline’s mean.This supports stabilization of recurrent activations without explicitly normalizing the mean.
  • 6.2–6.4 Other Tasks: RMSNorm remained comparable to or better than LayerNorm across retrieval, reading-comprehension, and image-classification results, with reported speedups of 40%∼64%, about 15%, and about 20.5%, respectively.For image classification, pRMSNorm added 2.6% speedup but sacrificed 1.54% test accuracy; under abnormal initialization, RMSNorm was more robust than LayerNorm.

7 Conclusion and Future Work

The paper concludes that RMSNorm preserves LayerNorm’s re-scaling invariance while removing re-centering, reducing computational overhead and maintaining comparable quality. Future work will investigate why RMSNorm succeeds and whether other norms can simplify normalization techniques.

  • RMSNorm normalizes summed inputs by their root mean square while preserving re-scaling invariance and removing re-centering invariance.
  • RMSNorm is a drop-in LayerNorm replacement that delivers comparable quality with lower computational overhead and observed speedups of 7%∼64%.Actual gains depend on framework, hardware, architecture, and the relative cost of other components.
  • Partial RMSNorm estimates RMS from a subset of summed inputs, but empirical speed improvements were not consistently observed.
  • Future work will analyze RMSNorm’s success, explore alternative norms, and simplify other normalization techniques such as BatchNorm.

A.1 Machine Translation

The machine-translation experiments evaluate RNNSearch and Transformer models across implementations, including comparisons with WeightNorm. WeightNorm converges more slowly and achieves lower test quality than LayerNorm and (p)RMSNorm.

  • The experiments cover WMT14 English-German translation using GRU-based RNNSearch and self-attention-based Transformer models.The training corpus contains 4.5M aligned sentence pairs, with newstest2013 for development and newstest2014/newstest2017 for testing.
  • WeightNorm converges more slowly, requires more training steps, and underperforms LayerNorm and (p)RMSNorm on test translation quality.Its test scores are 21.7 on Test14 and 23.5 on Test17.

A.2 CNN/Daily Mail Reading Comprehension

The CNN/Daily Mail reading-comprehension experiment trains LSTM models on the CNN corpus and evaluates RMSNorm within the LSTM architecture. The supplied passages identify the training setup and a SacreBLEU curve figure but do not report the figure’s outcome.

  • Figure 7 plots SacreBLEU over training steps on newstest2013 for RNNSearch models trained with Nematus in Theano.
  • Models are trained on the CNN corpus using an LSTM with hidden size 240 and RMSNorm applied within the LSTM for comparison with LayerNorm.Optimization uses Adam with batch size 64 and learning rate 8e−5.

A.3 Image-Caption Retrieval

The image-caption retrieval experiment uses a GRU-based sentence encoder and pretrained VGGNet image features in a pairwise-ranking model. Retrieval performance is reported with Recall@K across five test sets, while implementation differences can affect baseline comparisons.

  • The OE model encodes sentences with a GRU-based RNN and represents images using pretrained VGGNet outputs.It trains image-caption pairs with a pairwise ranking loss.
  • Image-caption retrieval is evaluated using Recall@K averaged across five test sets, each containing 1000 images and 5000 captions.Word embeddings have size 300 and GRU hidden states have size 1024.
  • Baseline comparisons may differ across frameworks because of implementation details and setup differences.
Loading 1910.07467v1…