Title: Pre-RMSNorm and Pre-CRMSNorm Transformers: Equivalent and Efficient Pre-LN Transformers
Venue: NeurIPS 2023
Accepted: yes
Keywords: Transformer

Normalization

Layer Normalization

RMSNorm

Efficient Machine Learning
Abstract: Transformers have achieved great success in machine learning applications.
Normalization techniques, such as Layer Normalization (LayerNorm, LN) and Root Mean Square Normalization (RMSNorm), play a critical role in accelerating and stabilizing the training of Transformers.
While LayerNorm recenters and rescales input vectors, RMSNorm only rescales the vectors by their RMS value.
Despite being more computationally efficient, RMSNorm may compromise the representation ability of Transformers.
There is currently no consensus regarding the preferred normalization technique, as some models employ LayerNorm while others utilize RMSNorm, especially in recent large language models.
It is challenging to convert Transformers with one normalization to the other type.
While there is an ongoing disagreement between the two normalization types,
we propose a solution to unify two mainstream Transformer architectures, Pre-LN and Pre-RMSNorm Transformers.
By removing the inherent redundant mean information in the main branch of Pre-LN Transformers, we can reduce LayerNorm to RMSNorm, achieving higher efficiency.
We further propose the Compressed RMSNorm (CRMSNorm) and Pre-CRMSNorm Transformer based on a lossless compression of the zero-mean vectors.
We formally establish the equivalence of Pre-LN, Pre-RMSNorm, and Pre-CRMSNorm Transformer variants in both training and inference.
It implies that Pre-LN Transformers can be substituted with Pre-(C)RMSNorm counterparts at almost no cost, offering the same arithmetic functionality along with free efficiency improvement.
Experiments demonstrate that we can reduce the training and inference time of Pre-LN Transformers by 1% - 10%.

Pre-RMSNorm and Pre-CRMSNorm Transformers: Equivalent and Efficient Pre-LN Transformers

Abstract

Transformers have achieved great success in machine learning applications. Normalization techniques, such as Layer Normalization (LayerNorm, LN) and Root Mean Square Normalization (RMSNorm), play a critical role in accelerating and stabilizing the training of Transformers. While LayerNorm recenters and rescales input vectors, RMSNorm only rescales the vectors by their RMS value. Despite being more computationally efficient, RMSNorm may compromise the representation ability of Transformers. There is currently no consensus regarding the preferred normalization technique, as some models employ LayerNorm while others utilize RMSNorm, especially in recent large language models. It is challenging to convert Transformers with one normalization to the other type. While there is an ongoing disagreement between the two normalization types, we propose a solution to unify two mainstream Transformer architectures, Pre-LN and Pre-RMSNorm Transformers. By removing the inherent redundant mean information in the main branch of Pre-LN Transformers, we can reduce LayerNorm to RMSNorm, achieving higher efficiency. We further propose the Compressed RMSNorm (CRMSNorm) and Pre-CRMSNorm Transformer based on a lossless compression of the zero-mean vectors. We formally establish the equivalence of Pre-LN, Pre-RMSNorm, and Pre-CRMSNorm Transformer variants in both training and inference. It implies that Pre-LN Transformers can be substituted with Pre-(C)RMSNorm counterparts at almost no cost, offering the same arithmetic functionality along with free efficiency improvement . Experiments demonstrate that we can reduce the training and inference time of Pre-LN Transformers by 1% -10% .

1 Introduction

Transformers have become a successful architecture for a wide range of machine learning applications, including natural language [39], computer vision [13], and reinforcement learning [7]. It is one of the foundation models [2], and pretrained Transformers [29] demonstrated impressive generalization results. Among the components of Transformers, normalization plays a critical role in accelerating and stabilizing the training process [21]. Layer Normalization (LayerNorm, LN) [1] and Root Mean Square Normalization (RMSNorm) [44] are two common normalization layers in Transformers. LayerNorm is in the original Transformer architecture [39], recentering and rescaling the input vector in R d to obtain a zero-mean and unit-variance output. RMSNorm only rescales the input vector with its RMS value, offering greater computational efficiency than LayerNorm.

The machine learning community does not reach a consensus regarding the preferred normalization technique for Transformers. LayerNorm demonstrates remarkable success in the milestone Transformers, such as GPT [30, 5] and ViT [13]. It is still the default normalization layer when building a new Transformer. On the contrary, RMSNorm is reported to accelerate the training and inference with similar performance as LayerNorm in Transformers. It has gained popularity in recent large language models, such as T5 [32], Gopher [31], Chinchilla [16], and LLaMA [38]. However, concerns persist regarding the potential negative impact of RMSNorm on the representation ability of Transformers. It remains an open question to determine the preferred normalization type for Transformers, requiring further theoretical and empirical investigation.

In this work, we aim to mitigate the discrepancy between LayerNorm and RMSNorm in Transformers. When delving into the prevalent Transformer architectures, Pre-LN and Pre-RMSNorm Transformers, we identify an opportunity to unify them by removing the inherent redundancy in Pre-LN models. In particular, the main branch vectors in the Pre-LN Transformers are always normalized before they are used, implying that the mean information is redundant. We can recenter the main branch without impact on the functionality of the models, which allows us to reduce LayerNorm to RMSNorm.

We further propose Compressed RMSNorm (CRMSNorm), which takes a vector in R d -1 as input, decompresses it to a zero-mean vector in R d , and applies the RMSNorm on the decompressed vector. Building upon this new normalization, we propose Pre-CRMSNorm Transformer that employs lossless compression on the zero-mean vectors. We apply such compression to zero-mean activations and parameters in PreRMSNormTransformers, enhancing the efficiency of Pre-RMSNorm Transformers while preserving the equivalent arithmetic functionality.

We formally claim that Pre-LN, Pre-RMSNorm, and Pre-CRMSNorm Transformers are equivalent for both training and inference. Figure 1 visualizes the overview of their equivalence. Such equivalence directly enables more efficient training and deployment of Pre-LN Transformers. We can translate a Pre-LN model into an equivalent

Pre-LN

Pre-RMS

Pre-CRMS

Sec. 3.2

Equivalence

remove redundant

mean information

Our Proposed Transformer Variants

Sec. 3.4

Higher Efficiency

Pre-(C)RMSNorm model, which can be readily deployed or adapted to downstream tasks. The conversion process incurs minimal costs. We can also train a Pre-(C)RMSNorm model directly as if we train an equivalent Pre-LN Transformer counterpart.

Although our relative improvement in the training and inference efficiency seems not large (up to 10% time reduction), we have the following arguments to support its significance. (1) Our proposed method can guarantee arithmetic equivalence and is a free lunch for Pre-LN Transformers. The efficiency improvement originates from removing the inherent redundancy in Pre-LN Transformers without introducing any fine-tuning or calibration. Our work can strictly push the performance-efficiency Pareto frontier of Pre-LN Transformers. (2) The modest relative improvement can be translated into significant absolute improvement given that Pre-LN Transformers are foundation models for the current and future data-centric [42] and generative [6] artificial intelligence. For instance, reducing the ChatGPT inference cost by 1% may save $7 , 000 per day [25]. (3) Our method is orthogonal and complementary to most work improving efficiency, such as efficient Transformer variants [35], quantization [22, 3], and distillation to smaller models [34].

We highlight our contributions as follows. Our code is available at https://github.com/ ZixuanJiang/pre-rmsnorm-transformer .

We achieve the first-ever unification of LayerNorm and RMSNorm in pre-normalization Transformers with proven arithmetic equivalence.

We propose two variants: Pre-RMSNorm and Pre-CRMSNorm Transformers. The original Pre-LN Transformer and our proposed two variants are equivalent and can seamlessly interchange without affecting functionality.

Our proposed architectures are 1% -10% more efficient than the original Pre-LN Transformer for both training and inference. Such efficiency gains are effortlessly obtained without the need for fine-tuning or calibration.

2 Background

We introduce LayerNorm, RMSNorm, and their usage in Transformers. We provide an abstraction for Pre-LN Transformers.

2.1 LayerNorm and RMSNorm

Layer Normalization (LayerNorm, LN) [1] is a technique to normalize the activations of intermediate layers of neural networks. Given a vector x ∈ R d , LayerNorm normalizes it to obtain a zero-mean unit-variance vector,

LayerNorm recenters and rescales the activations and gradients in the forward and backward computations [41], which enables fast and robust training of neural networks.

Root Mean Square Normalization (RMSNorm) [44] is another technique used for normalizing the activations. It is similar to LayerNorm in that it aims to accelerate and stabilize the training but uses a different normalization approach. Instead of normalizing the inputs based on their mean and variance, RMSNorm normalizes them based on their root mean square (RMS) value. It is defined in the following equation,

RMSNorm only rescales the input vector and the corresponding gradients, discarding the recentering process. As shown in their definitions, RMSNorm is computationally simpler and more efficient than LayerNorm. It is reported that replacing LayerNorm with RMSNorm can achieve comparable performance and save training and inference time by 7% -64% [44].

Given a zero-mean vector x , these two kinds of normalization are equivalent. Formally, if µ ( x ) = 0 , then LayerNorm ( x ) = RMSNorm ( x ) . We may optionally introduce learnable parameters and apply an element-wise affine transformation on the output of LayerNorm and RMSNorm.

We focus on the computation of LayerNorm and RMSNorm instead of their optimization and expressivity in this paper. 2

2.2 Normalization in Transformers

Normalization plays a crucial role and has many variants in Transformers [39]. LayerNorm is widely used in Transformer architectures to address this issue. The position of LN within the architecture is essential for the final performance. While the initial Transformer uses Post-LN, most Transformers employ Pre-LN to achieve more stable training, even though this can result in decreased performance [40]. Pre-LN is the mainstream normalization in Transformers, especially the large models, such as ViT [13, 9], PaLM [8], and GPT-series models [30, 5].

RMSNorm is proposed as an alternative normalization technique in Transformers. Several large language models, such as Chinchilla [16] and LLaMA [38], use Pre-RMSNorm in their blocks [45]. RMSNorm can help accelerate the training and inference with similar performance in these large models. Specifically, the experiments in [26] show that RMSNorm improves the pre-training speed by 5% compared with the LayerNorm baseline.

𝒙𝒙

𝐿𝐿

∈ ℝ 𝑑𝑑

LayerNorm

Linear in

ℊ𝑙𝑙

Linear out

𝒙𝒙

𝑙𝑙

𝒙𝒙

𝑙𝑙+1

Preprocess

Block 1

Block

L

Postprocess

LayerNorm

𝒙𝒙0

∈ ℝ 𝑑𝑑

Residual

ℱ𝑙𝑙

(a) Pre-LN Transformer

𝜇𝜇 ≠

0

ℝ 𝑑𝑑

ℝ 𝑑𝑑

ℝ 𝑑𝑑𝑖𝑖

ℝ 𝑑𝑑𝑜𝑜

ℝ 𝑑𝑑

𝒙𝒙

𝑙𝑙

∈ ℝ 𝑑𝑑

𝜇𝜇 ≠

0

𝒙𝒙

𝐿𝐿

∈ ℝ 𝑑𝑑

RMSNorm

Linear in

ℊ𝑙𝑙

Linear out

𝒙𝒙

𝑙𝑙

𝒙𝒙

𝑙𝑙+1

Preprocess

Block 1

Block

L

Postprocess

RMSNorm

𝒙𝒙0

∈ ℝ 𝑑𝑑

Residual

ℱ𝑙𝑙

(b)

Our Pre-RMSNorm Transformer

𝝁𝝁

=

𝟎𝟎

ℝ 𝑑𝑑

ℝ 𝑑𝑑

ℝ 𝑑𝑑𝑖𝑖

ℝ 𝑑𝑑𝑜𝑜

ℝ 𝑑𝑑

𝒙𝒙

𝑙𝑙

∈ ℝ 𝑑𝑑

𝝁𝝁

=

𝟎𝟎

Recenter

CRMSNorm

Linear in

ℊ𝑙𝑙

Linear out

𝒙𝒙

𝑙𝑙

𝒙𝒙

𝑙𝑙+1

Preprocess

Block 1

Block

L

CRMSNorm

𝒙𝒙0

∈ ℝ 𝑑𝑑-1

Residual

ℱ𝑙𝑙

(c)

Our Pre-CRMSNorm Transformer

ℝ 𝑑𝑑-1

ℝ 𝑑𝑑-1

ℝ 𝑑𝑑𝑖𝑖

ℝ 𝑑𝑑𝑜𝑜

ℝ 𝑑𝑑-1

𝒙𝒙

𝑙𝑙

∈ ℝ 𝑑𝑑-1

𝒙𝒙

𝐿𝐿

∈ ℝ 𝑑𝑑-1

(a) = (b) = (c) Equivalence; (a) < (b) < (c) Efficiency

Postprocess

It is challenging to convert Transformers with one normalization to the other type. It is not clear which version of normalization is more suitable for Transformers.

2.3 Pre-LN Transformer

Figure 2.a illustrates the Pre-LN Transformer architecture, which consists of three parts, preprocessing, a stack of L blocks, and postprocessing.

Preprocessing. We preprocess the raw inputs, which range from paragraphs in natural language [39], images [13], to state-action-reward trajectories in reinforcement learning problems[7]. We import special tokens (such as the classification token) and embeddings (such as positional embeddings), ultimately obtaining a sequence of token embeddings x 0 .

Transformer blocks. The main body of Pre-LN Transformer [40] is a stack of residual blocks

where x l is the input of the l -th block, F l is a sequence of operators LayerNorm → Linear → g l → Linear . We name x l as the vectors on the main branch and F l as the residual branch [14]. The block F l is usually an attention or a multi-layer perceptron (MLP) module. If g l is an activation function, such as GELU [15], then the block F l is a two-layer MLP. If g l is a (masked) multi-head scaled dot product attention, then the block is the (casual) attention module [39]. These two linear layers are usually explicitly defined. Taking the attention module as an example, the input linear projection generates the query, key, and value vectors, while the output projection is applied to the concatenated results of all heads. If they are not explicitly defined, we can add an identity mapping, a special linear layer. Most of the learnable parameters of Transformers are in these two linear layers. 3

LayerNorm and postprocessing. We finally process the result LN ( x L ) to obtain the task-related results, such as classification probabilities. We apply the layer normalization on x L since it usually has a high variance because it is the accumulation of all the Transformer blocks.

3 Method

We propose Pre-RMSNorm and Pre-CRMSNorm Transformer variants, shown in Figure 2, and claim that Pre-LN, Pre-RMSNorm, and Pre-CRMSNorm Transformers are arithmetically equivalent.

Pre-LN Transformer = Pre-RMSNorm Transformer = Pre-CRMSNorm Transformer . (4)

We will show the equivalence of these three architectures and then analyze the computational efficiency improvement by our proposed model variants. We discuss the Post-LN in Appendix B.

To clarify, our proposed Pre-RMSNorm Transformer and the pre-existing Pre-RMSNorm models (e.g., LLaMA [38]) are different. Ours has a recentering and special output linear layer, as shown in Figure 2b. However, they can share the same implementation for inference. We adopt the term "Pre-RMSNorm Transformer" to represent our proposed one in the discussions below.

3.1 Pre-LN Transformer = Pre-RMSNorm Transformer

LayerNorm is invariant to the shifting LN ( x + k 1 ) = LN ( x ) , ∀ k ∈ R . We observe that LayerNorm is applied to the main branch vectors before they are used in either residual branches or the final postprocessing in Pre-LN Transformers. Therefore, we can replace the main branch vectors x l with x l + k l 1 , ∀ k l ∈ R without impact on the functionality of the Pre-LN Transformer. If k l = -µ ( x l ) , we replace the main branch vectors with its recentered version x l -µ ( x l ) 1 . We can explicitly maintain zero-mean main branches with the same arithmetic functionality.

We propose three modifications to the original Pre-LN Transformer to obtain an equivalent PreRMSNorm Transformer.

Recenter the x 0 before the first Transformer block, where Recenter( x ) = x -µ ( x ) 1 .

For the output projection in residual branches, replace the weight A o and bias b o with ˆ A o = A o -1 d 11 T A o , ˆ b o = b o -µ ( b o ) 1 , where d is the dimension of x 0 .

Replace LayerNorm with RMSNorm at the beginning of residual blocks and before postprocessing.

Since µ ( x l +1 ) = µ ( x l ) + µ ( F l ( x l )) , we can keep zero-mean on the main branch if and only if the input of the first block x 0 and the output of each residual branch F l are re-centered with zero-mean. The first modification is to recenter x 0 , while the second modification is to recenter the output of residual branches. For the residual branch F l , the ending linear transformation enables us to recenter its output without extra computation, implied by Lemma 3.1. We can recenter the weight and bias of a linear layer to recenter its output.

Lemma 3.1 Given a linear transformation y = Ax + b , x ∈ R n , A ∈ R m × n , b , y ∈ R m , we can decompose the output with two parts y = ˆ Ax + ˆ b + µ ( y ) 1 . The first part ˆ Ax + ˆ b = y -µ ( y ) 1 , with zero mean, is another linear transformation with ˆ A = A -1 m 11 T A , ˆ b = b -µ ( b ) 1 .

The first two modifications ensure that we maintain zero mean on the main branch vectors. Given a zero-mean input, LayerNorm is equivalent to RMSNorm, which implies that the third modification has no impact on the functionality. With these modifications, we demonstrate that Pre-LN and Pre-RMSNorm Transformers are equivalent.

3.2 Pre-RMSNorm Transformer = Pre-CRMSNorm Transformer

For a zero-mean vector x ∈ R d , we can compress it losslessly by discarding its last element. In the decompression, we recover the discarded element with x d = -∑ d -1 i =0 x i . The decompression has an extra cost, while the compression does not induce extra computation. The space-saving ratio of this compression method is 1 /d .

We define Compressed Root Mean Square Normalization (CRMSNorm), which takes a vector x ∈ R d -1 as input. CRMSNorm first decompresses the vector x to obtain a zero-mean vector in R d , then applies RMSNorm on the zero-mean vector. It can generate the normalized zero-mean results in either R d -1 or R d . Its formal definition is in Equation 5.

We simplify the Pre-RMSNorm Transformers to obtain the Pre-CRMSNorm Transformers with the following modifications.

Compress the zero-mean main-branch vectors from R d to R d -1 . Preprocessing and postprocessing handle compressed vectors in R d -1 .

Replace RMSNorm with CRMSNorm.

Simplify the weight A i in the input projection layer. Let a d be the last column vector of A i . ˆ A i = A i -a d 1 T is the compressed weight matrix.

Simplify the weight and bias in the output projection layer. We discard the last row of the weight matrix ˆ A o and the last element of the bias ˆ b o since they are used to generate the redundant last element.

We compress the zero-mean activations in our proposed Pre-RMSNorm Transformers and correspondingly simplify the two linear layers in residual branches.

We can fuse preprocessing, recentering, and compression. The fused preprocessing generates compressed vectors in R d -1 to represent zero-mean vectors in R d . Taking language models as an example, we can recenter and compress (discard the last element) the word and position embeddings.

For the input linear projection, its input is the output of RMSNorm, whose mean is zero. We compress the output of the RMSNorm and the weight of the linear layer. Specifically, if the linear layer takes a zero-mean vector as input, then

ˆ A i = A i -a d 1 T is the compressed weight matrix since its last column is a zero vector. The simplified linear layer only needs the compressed zero-mean vectors in R d -1 as input.

For the output linear projection, we only need to calculate the first ( d -1) elements for the vectors in R d . Thus, we can compress the weight ˆ A o and bias ˆ b o . As shown in Figure 2, the shape of ˆ A o is ( d, d o ) , we can compress it directly to ( d -1 , d o ) by discarding its last row. Similarly, we can compress ˆ b o ∈ R d by discarding its last element. This weight compression also reflects their redundancy. As shown in Section 3.1, ˆ b o and all column vectors of ˆ A o are zero-mean.

Figure 2 demonstrates the difference in vector dimension. We demonstrate the equivalence between Pre-RMSNorm and Pre-CRMSNorm Transformers.

3.3 Pre-LN Transformer = Pre-CRMSNorm Transformer

To translate a pre-trained Pre-LN Transformer to a Pre-CRMSNorm Transformer, we can convert it into a Pre-RMSNorm model and then finish the conversion. For the model implementation, we need two steps to convert a Pre-LN Transformer to Pre-CRMSNorm Transformer.

Reduce the hidden dimension from d to d -1 . 4

Replace LayerNorm with CRMSNorm.

Namely, Pre-LN Transformers with the main branch vectors in R d are equivalent to Pre-CRMSNorm Transformers in R d -1 , which further echoes the redundancy in the Pre-LN Transformers.

3.4 Training and Inference Efficiency

We show how to make conversions between the three variants. The conversions only consist of one-time parameter adjustments without expensive fine-tuning or calibration, similar to operator fusion [27]. We qualitatively analyze the efficiency improvement of our proposed Pre-(C)RMSNorm Transformers compared with equivalent Pre-LN models.

3.4.1 Pre-RMSNorm Transformer

We discuss the impact of three modifications on training and inference, as shown in Table 1.

The linear layer with zero-mean output. The modified linear layer will not induce extra computation for inference since we can replace the weight and bias in advance. During inference, we can treat it as a standard linear layer with equivalently transformed parameters.

We have to pay the extra cost for training for the parameter change. The parameter optimizer manages A o , b o , but we use their recentered version ˆ A o , ˆ b o during training. The induced cost is small for several reasons. (1) The size of parameters A o , b o is relatively small since they do not depend on the batch size and sequence length. For reference, the computation cost of LayerNorm in the original Pre-LN Transformer is proportional to the batch size and the sequence length. (2) Obtaining ˆ A o , ˆ b o can be done ahead of time since it does not depend on the input. We may leverage the idle time of accelerators to compute them. (3) In data-parallel distributed training, the parameter server [19] manages the parameters. Each worker will receive ˆ A o , ˆ b o from the server and then pass the gradients of loss to ˆ A o , ˆ b o to the server. Only the server needs to maintain and update the original A o , b o . In short, it is much easier to process the model parameters than the intermediate activations.

Recentering. It is possible to fuse the recentering with the preprocessing, which usually handles the sum of several kinds of embeddings. For example, the input embeddings of the BERT model [11] are the accumulation of the token embeddings, the segmentation embeddings, and the position embeddings. Since Recenter ( x + y ) = Recenter ( x )+ Recenter ( y ) , recentering the input is equivalent to recentering each embedding before the addition. We can recenter the accumulated embeddings or each embedding separately before the accumulation. Suppose an embedding is from a linear layer. In that case, we can modify the linear layer such that it generates the zero-mean output, similar to how we edit the output linear projection.

For inference, we can recenter the related embeddings or linear layers in advance such that no extra computation is induced. For training, recentering induces extra cost since it is on the fly.

Replacing LayerNorm with RMSNorm. Section 2.1 introduces that RMSNorm can achieve speedup compared with LayerNorm, as demonstrated by the previous models. This replacement can help us accelerate the training and inference of the Pre-LN Transformer.

3.4.2 Pre-CRMSNorm Transformer

CRMSNorm, an extension of RMSNorm, saves execution time as it is more computationally efficient than LayerNorm. Additionally, CRMSNorm further reduces the hidden dimension from d to d -1 , which can reduce the model size, computation, communication, and memory consumption by 1 /d in theory. However, most accelerators can not efficiently handle the vectors in R d -1 when d is a large even number, especially a power of two (e.g., 1024, 4096). This limitation arises because these accelerators are typically optimized for arithmetic with even dimensions. In some cases, handling R d -1 vectors may take much more time than R d vectors. Thus, we need to examine if the accelerators can support R d -1 vectors efficiently. If not, we have the following alternatives. For inference, we may either (1) add zero embeddings, or (2) decompress the vectors and translate the model into the Pre-RMSNorm variant. For training, we suggest keeping the hidden dimension d , which is equivalent to a Pre-LN model with the hidden dimension d +1 . In this way, the CRMSNorm can help us increase the model representability and computation efficiency at the same time.

3.5 Training and Inference Equivalence

Taking Pre-LN and Pre-RMSNorm Transformers as examples, we further explain the arithmetic equivalence.

Inference equivalence. Let f be a Pre-LN Transformer with parameter θ , and g be a Pre-RMSNorm Transformer with parameter ϕ . We demonstrate that for any input x ∈ R d , we can always have f ( x , θ ) = g ( x , ϕ = h ( θ )) and we show how we conduct the conversion ϕ = h ( θ ) . Thus, we claim that f and g are equivalent in terms of arithmetic functionality. Namely, when calculating f ( x , θ ) during inference, we can always compute the equivalent counterpart g ( x , ϕ = h ( θ )) . Our

proposed method is a re-parameterization technique [12], which builds equivalent models with different parameters.

Training equivalence. Despite f ( x , θ ) = g ( x , ϕ = h ( θ )) , there may be a large difference in the convergence speed and stability when training f and g with SGD and its variants. We avoid this issue by applying the gradient updates on θ instead of ϕ . Specifically, during the training process, we maintain a Pre-LN Transformer model f and its parameter θ . In the forward and backward computation, we convert the model and its parameter into g and ϕ to save training time. We do not update ϕ directly. Instead, we calculate the gradients ∇ θ and call the SGD optimizer θ = θ -η ∇ θ , where η is the learning rate. From the perspective of optimization, the training is still on the Pre-LN Transformer model f and its parameter θ . Consequently, the training performance and stability are preserved and equivalent to the original Pre-LN Transformers.

4 Experiments

We claim that our major contributions are the unification and equivalence of the three Transformer variants. We do not report the task-related performance, such as the classification accuracy and perplexity, since the training and inference equivalence are guaranteed. We discuss how we verify the task-related performance in Appendix D.4.

The efficiency is a free lunch accompanying the equivalence. Now that the efficiency of RMSNorm over LayerNorm has been shown in the previous work [44, 26], we focus on analyzing the efficiency of each component of our method.

We conduct experiments on ViT [13, 36] and GPT-3 [5] since they represent two mainstream architectures of Transformers, encoder-only and casual decoder. Other Transformer variants can be treated as an extension of these two architectures, such as encoder-decoder [39], and non-causal decoder [43]. Also, ViT and GPT cover the areas of computer vision and natural language, where Transformers have been popular and achieved great success.

We abstain from utilizing pre-trained weights. We replicate the ViT and GPT-3 architectures as described in their respective papers, employing dummy data to assess the training and inference speed. We use PyTorch 2.0 [28] to build the training and inference pipeline. We run iterations at least 100 times and report the 25th, 50th (median), and 75th percentile since these quartiles are more robust than the mean. We use the automatic mixed precision [24] for both training and inference.

4.1 Experiments on ViT

The ViT takes images with 3 channels and a resolution of 224 × 224 and generates a classification over 1000 classes, which is the standard setting for ImageNet training and inference.

In the training recipe of ViT [13], dropout [33] is added at the end of each residual branch, which breaks the zero-mean property of the residual output. Hence, in Pre-RMSNorm Transformers, we need to recenter the output vectors explicitly at the end of the residual branch. 5 The Pre-CRMSNorm Transformers are compatible with dropout since it uses vectors in R d -1 to represent the compressed zero-mean vectors in R d . On the contrary, DeiT [36] disables the dropout and can guarantee the zero-mean output, in spite that the stochastic depth [17] and LayerScale [37] are applied. We follow the settings in DeiT [36] in this paper.

Inference. Figure 3 illustrates the inference time with different batch sizes and model sizes on a single A100 GPU. Taking the Pre-LN Transformer as the baseline, our proposed Pre-RMSNorm Transformer can reduce the inference time by 1% -9% . This reduction takes effect for ViTs with different models and various mini-batch sizes. The left subfigure demonstrates the inference latency, which is the inference time when the batch size is 1.

The proportion of LayerNorm in the total computation is 12% -18% in these experiments. As discussed in Sections 2.1 and 3.4.1, replacing LayerNorm with RMSNorm can help us accelerate the inference. RMSNorm can reduce the inference time of LayerNorm by 20% -60% , thus helping us achieves a stale faster inference for Pre-RMSNorm models.

85.5%

87.1%

87.7%

89.0%

87.1%

87.7%

97.4%

97.2%

97.7%

97.3%

97.7%

96.4%

0.8

0.85

0.9

0.95

1

Ti/16

S/16

B/16

H/14

G/14

Normalized Inference Time

ViT Variants

Inference time with batch size 1

85.5%

83.0%

84.3%

86.5%

81.8%

82.2%

96.7%

96.1%

95.8%

96.2%

93.9%

93.4%

0.8

0.85

0.9

0.95

1

1

4

16

64

256

1024

Normalized Inference Time

Inference time of ViT-Ti/16

w/o norm     Pre-LN     Pre-RMS     Pre-CRMS

87.4%

87.4%

87.4%

87.4%

12.6%

9.3%

10.9%

0.1%

0.9%

0.8

0.85

0.9

0.95

1

w/o norm

Pre-LN

Pre-RMS

Pre-CRMS

Training time breakdown

others

normalization

recenter

linear

(a) Training time breakdown on ViT-Ti/16.

97.7%

97.7%

98.3%

99.1%

100.0%

100.6%

98.3%

99.6%

98.4%

99.4%

98.6%

98.9%

0.95

1

Ti/16

S/16

B/16

H/14

G/14

Normalized Training Time

ViT Variants

Training time of ViT Variants

Pre-LN

Pre-RMS

Pre-CRMS

(b) Normalized training time

For Pre-CRMSNorm, the GPU cannot efficiently handle vectors in R d -1 in some cases, where we add zero padding to obtain vectors in R d . With zero padding, Pre-CRMSNorm models are less efficient than Pre-RMSNorm ones due to the extra decompression. However, for the cases where R d -1 vectors can be accelerated efficiently, we can achieve an even 10% time reduction.

We also conduct experiments (1) with other precisions, (2) on CPUs, (3) with JAX [4] and observe similar performance. For these ViT variants, we have achieved an average of 3 . 0% inference time reduction. Please see Appendix D for more details.

Training. We train ViTs on a single workstation with 4 A100s with data parallel training [20], following the DeiT training recipe. Each training iteration consists of forward and backward computation on all the workers, gradient all-reduce, parameters update with an optimizer, and broadcasting the new parameters. Figure 4a visualizes the breakdown of the related components. In the Pre-LN Transformer, the layer normalization accounts for 12 . 6% of total training time. For the Pre-RMSNorm variant, we have to modify the output projection and recenter the input, which induces the 0 . 89% and 0 . 09% extra computation time. Then LayerNorm is reduced to RMSNorm, whose computation cost is 9 . 27% . Overall, the Pre-RMSNorm reduces the training time by 2 . 36% .

We train the Pre-CRMSNorm Transformers with d as the hidden dimension since we do not obtain a speedup from the compressed dimension since the GPUs cannot handle R d -1 vectors efficiently. The CRMSNorm is more computationally expensive than the RMSNorm given the same input but takes less time than LayerNorm. The Pre-CRMSNorm Transformer does not need the recentering and special linear layers to generate zero-mean results at the end of residual branches. Above all, training the Pre-CRMSNorm variant is 1 . 74% faster than the Pre-LN model.

Figure 4b illustrates the training time of different ViTs. We have achieved a speedup of 1% -2 . 5% , which is smaller than the inference speedup. Considering only the forward and backward computation, the speedup is similar between training and inference. Nevertheless, the training needs extra time on gradient all-reduce, optimizer update, and parameters broadcast, which shrinks the percentage of

normalizations in the whole computation. For reference, the percentage is 10% -15% for training these ViTs.

4.2 Experiments on GPT

We measure the inference time on the GPT-3 [5] variants with batch size 1 and sequence length 512. The results are shown in Figure 5. In small GPT-3 models, layer normalization takes considerable time. Hence, replacing the LayerNorm with (C)RMSNorm can reduce the inference time by up to 10% . However, as the model grows, the attention and MLP take charge of the main part of the computation [18]. The normalization takes < 1% of the inference time for GPT-3 XL and larger models, which is the upper bound of the speedup with our methods.

Applying quantization may mitigate this issue to some extent. By applying int8 matrix multiplication [10], the percentage of the layer nor-

0.85

0.9

0.95

1

S

M

L

XL

2.7B

6.7B

Normalized  Inference Time

GPT3 variants

w/o norm

Pre-LN

Pre-RMS

Pre-CRMS

malization increases to 10% for GPT-3 XL and 2.7B. Our Pre-RMSNorm can reduce the inference time by 4% . Training performance is similar to the ViT one. We have achieved 1 . 5% and 1 . 8% time reduction for GPT-3 Small and Medium, respectively.

5 Conclusion

In this paper, we propose two equivalent and efficient variants for the widely used Pre-LN Transformers. We point out the inherent redundancy in the Pre-LN Transformer. By maintaining zero-mean on the main branch vectors and thus simplifying LayerNorm, we obtain the Pre-RMSNorm architecture. We further apply a lossless compression on the zero-mean vectors to obtain the Pre-CRMSNorm model. For the first time, We unify these normalization variants within the Transformer model.

We enable the more efficient utilization of Pre-LN Transformers, allowing for seamless transitions between normalization techniques with minimal overhead. We can replace a Pre-LN Transformer with an equivalent Pre-(C)RMSNorm Transformer with better training and inference efficiency, which is a free lunch. We strictly push the performance-efficiency Pareto frontier of foundational Pre-LN Transformers. As a result, pre-trained Pre-LN Transformers (such as ViT and GPT) can be deployed more efficiently, and new equivalent or superior models can be trained with higher throughput.

Extensions. We believe that our proposed CRMSNorm technique has the potential for application in other neural architectures. Further exploration of its usage in different contexts would be beneficial. Additionally, while our focus has been on pre-normalization Transformers, it would be valuable to investigate the application of our methods to other related architectures, especially the foundation models. Finally, integrating our proposed method into machine learning compilers could directly enable the generation of equivalent and simplified computation graphs.

Limitations. Detailed implementations and specific optimizations will play a crucial role in achieving practical performance gains. It is necessary to focus on developing highly optimized implementations of (C)RMSNorm to bridge the gap between theoretical and practical efficiency improvements. By addressing these challenges, we can fully leverage the potential benefits of (C)RMSNorm and enable more efficient utilization of Pre-LN Transformers.

Acknowledgments and Disclosure of Funding

We acknowledge NVIDIA for donating its A100 GPU workstations and the support from TILOS, an NSF-funded National Artificial Intelligence Research Institute.

References

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258 , 2021.

Yelysei Bondarenko, Markus Nagel, and Tijmen Blankevoort. Understanding and overcoming the challenges of efficient transformer quantization. arXiv preprint arXiv:2109.12948 , 2021.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Yihan Cao, Siyu Li, Yixin Liu, Zhiling Yan, Yutong Dai, Philip S Yu, and Lichao Sun. A comprehensive survey of ai-generated content (aigc): A history of generative ai from gan to chatgpt. arXiv preprint arXiv:2303.04226 , 2023.

Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems , 34:15084-15097, 2021.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311 , 2022.

Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. arXiv preprint arXiv:2302.05442 , 2023.

Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 (): 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.07339 , 2022.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

Xiaohan Ding, Xiangyu Zhang, Ningning Ma, Jungong Han, Guiguang Ding, and Jian Sun. Repvgg: Making vgg-style convnets great again. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 13733-13742, 2021.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 , 2020.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770-778, 2016.

Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415 , 2016.

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556 , 2022.

Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part IV 14 , pages 646-661. Springer, 2016.

Cheng Li. Llm-analysis: Latency and memory analysis of transformer models for training and inference. https://github.com/cli99/llm-analysis , 2023.

Mu Li, David G. Andersen, Jun Woo Park, Alexander J. Smola, Amr Ahmed, Vanja Josifovski, James Long, Eugene J. Shekita, and Bor-Yiing Su. Scaling distributed machine learning with the parameter server. In Proceedings of the 11th USENIX Conference on Operating Systems Design and Implementation , OSDI'14, page 583-598, USA, 2014. USENIX Association.

Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, et al. Pytorch distributed: Experiences on accelerating data parallel training. arXiv preprint arXiv:2006.15704 , 2020.

Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. arXiv preprint arXiv:2004.08249 , 2020.

Zhenhua Liu, Yunhe Wang, Kai Han, Wei Zhang, Siwei Ma, and Wen Gao. Post-training quantization for vision transformer. Advances in Neural Information Processing Systems , 34:28092-28103, 2021.

Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.

Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al. Mixed precision training. arXiv preprint arXiv:1710.03740 , 2017.

Aaron Mok. Chatgpt could cost over $700,000 per day to operate. microsoft is reportedly trying to make it cheaper. https://www.businessinsider.com/ how-much-chatgpt-costs-openai-to-run-estimate-report-2023-4 . Accessed: 2023-05-01.

Sharan Narang, Hyung Won Chung, Yi Tay, William Fedus, Thibault Fevry, Michael Matena, Karishma Malkan, Noah Fiedel, Noam Shazeer, Zhenzhong Lan, et al. Do transformer modifications transfer across implementations and applications? arXiv preprint arXiv:2102.11972 , 2021.

Wei Niu, Jiexiong Guan, Yanzhi Wang, Gagan Agrawal, and Bin Ren. Dnnfusion: accelerating deep neural networks execution with advanced operator fusion. In Proceedings of the 42nd ACM SIGPLAN International Conference on Programming Language Design and Implementation , pages 883-898, 2021.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, High-Performance Deep Learning Library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32 , pages 8024-8035. Curran Associates, Inc., 2019.

Xipeng Qiu, Tianxiang Sun, Yige Xu, Yunfan Shao, Ning Dai, and Xuanjing Huang. Pre-trained models for natural language processing: A survey. Science China Technological Sciences , 63(10):1872-1897, 2020.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11446 , 2021.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research , 21(1):5485-5551, 2020.

Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. Journal of Machine Learning Research , 15(56):1929-1958, 2014.

Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-following llama model. https://github.com/tatsu-lab/stanford_alpaca , 2023.

Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys , 55(6):1-28, 2022.

Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning , pages 10347-10357. PMLR, 2021.

Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 32-42, 2021.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning , pages 10524-10533. PMLR, 2020.

Jingjing Xu, Xu Sun, Zhiyuan Zhang, Guangxiang Zhao, and Junyang Lin. Understanding and improving layer normalization. Advances in Neural Information Processing Systems , 32, 2019.

Daochen Zha, Zaid Pervaiz Bhat, Kwei-Herng Lai, Fan Yang, Zhimeng Jiang, Shaochen Zhong, and Xia Hu. Data-centric artificial intelligence: A survey. arXiv preprint arXiv:2303.10158 , 2023.

Biao Zhang, Behrooz Ghorbani, Ankur Bapna, Yong Cheng, Xavier Garcia, Jonathan Shen, and Orhan Firat. Examining scaling and transfer of language model architectures for machine translation. In International Conference on Machine Learning , pages 26176-26192. PMLR, 2022.

Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems , 32, 2019.

Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, et al. A survey of large language models. arXiv preprint arXiv:2303.18223 , 2023.

Appendices

A Proof for Lemma 1.

Given a linear transformation y = Ax + b , x ∈ R n , A ∈ R m × n , b , y ∈ R m , we have

where ˆ A = A -k 11 T A , ˆ b = b -µ ( b ) 1 , f ( x , k ) = k 1 T Ax + µ ( b ) .

If k = 1 /m , then we obtain

The ˆ A is the recentered matrix of A , and all its column vectors have zero-mean. We decompose the output into two parts.

The first part ˆ Ax + ˆ b = y -µ ( y ) 1 , with zero mean, is another linear transformation with ˆ A = A -1 m 11 T A , ˆ b = b -µ ( b ) 1 .

The second part corresponds to the mean information µ ( y ) 1 = ( 1 m 1 T Ax + µ ( b )) 1 .

B Post-LN Transformers

Different from the Pre-LN Transformers, the Post-LN Transformers have the following blocks.

Layer normalization is on the main branch instead of the beginning of residual branches. We can keep a zero-mean branch on the main branch without impacting the functionality.

For the residual branch F l , we can apply the same method in the Pre-LN Transformer. We can modify the output linear projection to obtain ˆ F l , which generates the zero-mean part of the original result.

The recentering operation ˆ x l = x l -µ ( x l ) 1 requires extra computation. If elementwise affine transformation is disabled in LayerNorm, x l is the output of a normalization such that µ ( x l ) = 0 and ˆ x l = x l . If the transformation is enabled, x l is not guaranteed zero-mean such that explicit recentering is necessary.

C Revertible Conversions

The conversions of Pre-LN → Pre-RMSNorm and Pre-RMSNorm → Pre-CRMSNorm are listed in Sections 3.1 and 3.2. These steps are fully reversible. We write the inverse steps explicitly below.

Remove the recentering.

A o = ˆ A o + 1 c T , where c can be a random vector, b o = ˆ b o + d 1 where d can be a random scalar.

Replace RMSNorm with LayerNorm.

Coverting Pre-CRMSNorm into Pre-RMSNorm Transformers.

Decompression vectors in R d -1 into zero-mean vectors in R d . The compression is lossless, so the decompression is its invertible operation.

Replace CRMSNorm with RMSNorm.

A i = ˆ A i + a d 1 T , where a d is a random vector.

Recover the last row of ˆ A o and the last element of ˆ b o such that each column of ˆ A o and ˆ b o is zero-mean.

D Experiments

D.1 Implementation of Normalization

We have provided our implementation with JAX and PyTorch in the supplementary material. The reported results are based on the following implementations.

For JAX, we use the APIs of LayerNorm and RMSNorm in the flax library. For PyTorch, we use the implementations of LayerNorm and RMSNorm from NVIDIA's apex extension 6 . For CRMSNorm, we use our own customized implementations. We also provide our customized implementation of LayerNorm and RMSNorm.

We notice that there are lots of APIs for the standard LayerNorm and RMSNorm. For example, PyTorch has provided the official LayerNorm API but lacks RMSNorm implementation. These different implementations are mixed. We do not find one implementation dominant over others for all the cases. For instance, torch.nn.LayerNorm is usually faster than apex.normalization.FusedLayerNorm when the input vectors are small in inference, while it is slower than the apex when the input vectors are large. PyTorch's official implementation is also slower than the apex for training.

D.2 Extended Experiments in ViT

Table 2 lists the architecture parameters of Vision Transformer. We first measure the inference time . We sweep these 6 ViTs with 6 batch sizes (1, 4, 16, 64, 256, 1024) and collect the medians of these

55

60

65

70

75

80

85

100

200

300

400

500

600

700

800

Test accuracy (%)

epochs

Train ViT-S/16 on ImageNet-1k from scratch

Pre-LN

Pre-RMSNorm

Pre-CRMSNorm

2.5

3

3.5

4

0

5

10

15

Train loss

steps (10 3 )

Train loss of GPT-2 Small

Pre-LN

Pre-RMSNorm

Pre-CRMSNorm

15

20

25

30

35

40

0

5

10

15

Evaluation perplexity

steps (10 3 )

Evaluation perplexity of GPT-2 Small

Pre-LN

Pre-RMSNorm

Pre-CRMSNorm

(a) Training loss

2.5

3

3.5

4

0

5

10

15

Train loss

steps (10 3 )

Train loss of GPT-2 Small

Pre-LN

Pre-RMSNorm

Pre-CRMSNorm

15

20

25

30

35

40

0

5

10

15

Evaluation perplexity

steps (10 3 )

Evaluation perplexity of GPT-2 Small

Pre-LN

Pre-RMSNorm

Pre-CRMSNorm

(b) Evaluation perplexity

36 data points. We report the average of these 36 experiments in Table 3. We conduct inference on a single A100 with automatic mixed precision (amp) in PyTorch. We further change the precision (disabling the amp), computation platforms (16 threads in AMD EPYC 7742 CPUs), and machine learning frameworks (JAX).

D.3 Numerical Issue

The theoretical arithmetic equivalence cannot be fully translated into equality in practical numerical computation if we use floating numbers. An intuitive example is that µ ( x + y ) = µ ( x )+ µ ( y ) always holds for any vectors x , y . However, if these two vectors are represented as (low precision) floating numbers, this equality is not guaranteed in real-world numerical computation. It is possible that these small discrepancies may be accumulated and enlarged in the large models, further degrading the numerical stability. In our proposed method, we cannot ensure the exact zero-mean in the main branch numerically.

The numerical issue is a common problem in machine learning. A typical example is operator reordering and layer fusion. PyTorch provides a related API officially, named torch.ao.quantization.fuse_modules . We can fuse the convolution layer and its following batch normalization layer to simplify the computation. These two layers are separate in training and can be fused to accelerate the inference. The fusion does not break the arithmetic equivalence but changes the numerical results. In spite of the numerical difference, the fusion usually has a neutral impact on task-related performance, such as classification accuracy, even in large models. Fine-tuning or calibration may be helpful in case there is severe performance degradation.

Our proposed methods encounter a similar issue as layer fusion since we modify partial parameters. In our experiments, we can convert the pre-trained Pre-LN ViT-H/14 into Pre-(C)RMS variants without any accuracy change on the ImageNet validation dataset. We observe that replacing PyTorch's official LayerNorm implementation with the apex may have a larger impact on the model performance.

D.4 Verification of Training Equivalence

For validation and reference, we train ViT-S/16 as Pre-LN, Pre-RMSNorm, and Pre-CRMSNorm models on ImageNet-1k following the DeiT-3 setting and achieve similar accuracy. The test accuracy over 800 epochs is shown in Figure 6. Remarkably, these Transformer variants yield similar training performance. We also train three variants of GPT-2 Small (12 layers, 768 hidden dimensions) on the wikitext-103-raw-v1 dataset [23] from scratch. The training and evaluation losses are demonstrated in Figure 7. These performance curves exhibit remarkable similarity, providing empirical evidence for their training equivalence.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper analyses the two common types of normalization layers in Transformers: LayerNorm and RMSNorm (root-mean-square-norm). LayerNorm scales all vectors to be of the same norm, while changing the vectors' \"directions\". RMSNorm, in contrast, keeps the same direction, but just rescales the vectors to the same constant norm.\nThe authors \"unify\" both approaches, by showing how to force LayerNorm to act like RMSNorm. Then, since RMSNorm is faster to compute, the authors argue that this conversion (LayerNorm -> RMSNorm) can get the expressivity benefits of LayerNorm with the speed of RMSNorm.\nFinally, the authors propose C-RMSNorm, which saves 1 dimension. For example, if the model hidden size is 1024, C-RMSNorm allows reducing it to 1023 while calculating the norm. \nOverall, the authors argue that these changes can reduce the training and inference time of LayerNorm transformers by 10% (although without discussing whether there's a reduction in performance/accuracy).\n\nSTRENGTHS:\n## Strengths\n+ The authors provide a thorough explanation of LayerNorm vs. RMSNorm.\n+ The authors show the relation between LayerNorm vs. RMSNorm, which I haven't seen in any other paper.\n+ The proposed approaches can sometimes reduce inference and training time, which is an important problem that solving can save millions and save carbon emmission.\n\nWEAKNESSES:\n## Weaknesses\n- I'm not sure about the soundness of some of the claims. The authors argue that LayerNorm is **equivalent** to RMSNorm and even write explicitly `PreLN Trasnformer = Pre-RMSNorm Transformer` (Eq 4). However, they only show one direction of this equality: the authors show how to take a Pre-LayerNorm transformer, and by imposing specific values to weights and enforcing constraints it can be equivalent to RMSNorm. That is, by constraining LayerNorm you can achieve RMSNorm. To my understanding, this means that Pre-LayerNorm is **at least** as expressive as Pre-RMSNorm. But what about the other direction? Maybe Pre-LayerNorm is more expressive if we do not impose these constraints? Can we take a Pre-RMSNorm and make it a Pre-LayerNorm?\n- Similarly, I am not sure that the authors prove both directions when claiming in Section 3.2 that `Pre-RMSNorm = Pre-C-RMSNorm`.\n- Novelty - The authors claim to \"propose Pre-RMSNorm\". Is this proposal novel? aren't models such as LLaMA already using Pre-RMSNorm?\n- Another issue with soundness: the authors write multiple times that these conversions are \"free efficiency improvement\" and \"free lunch\". However, in Line 228 they write: \"We have to pay the extra cost for training for the parameter change\". So, what exactly do the authors mean by \"free lunch\" if there is an extra training cost, and the accuracy is not guaranteed to be the same?\n- Many evaluation details are missing. \n    - On which datasets do the authors perform the experiments? \n    - How can the authors perform training and inference with **GPT-3**? As far as I know, GPT-3 is not open-sourced. Further, the authors mention they used \"**GPT3 XL and 2.7B**\". What exactly do they mean here? What is \"**GPT3 XL**\"? and what is \"**GPT3 2.7B**\", Is there such a model? The sizes of GPT-3 that I know about are much bigger than 2.7B. Where exactly this model was taken from and what exactly do the authors mean?\n\nQUESTIONS:\n## Questions\n1. Section 3.1 says that \"LayerNorm is applied to the main branch vectors before they are used in either residual branches or the final postprocessing. Therefore, we can replace the main branch vectors $x_l$ with $x_l + k_l\\mathbb{1}$\" - I am not sure I agree. I agree that since LayerNorm is invariant to the shifting $LN(x) = LN(x_l + k_l\\mathbb{1})$, the replacement will not change the output of the **residual branch**. However, if we replace the main branch vectors as suggested, the summation of the main branch and the residual branch that happens at the end of each transformer layer - is no longer identical to its original version.\n2. In most of the cases in Figures 3+4, C-RMSNorm is slower than RMSNorm. So what is the benefit of C-RMSNorm?\n3. Evaluation - the evaluation compares only the speed of training/inference of different normalization layers, without discussing the result of the training. Maybe Pre-RMSNorm is faster than Pre-LayerNorm, but needs to be trained for 10% **more steps** to reach the same training loss? Maybe even after training longer, it does not achieve the same test accuracy? I will not be surprised if that's the case: since the authors showed that RMSNorm is a special case of LayerNorm, I would expect LayerNorm to be more expressive and possibly reach better results, or reach these results within fewer steps.\n4. On which datasets were the experiments conducted? What was the downstream accuracy of each model?\n5. What exactly are the GPT-3 models mentioned in the evaluation?\n6. Are the claims and findings in this paper related to the findings in the following papers regarding LayerNorm? \n\n    * [\"On the Expressivity Role of LayerNorm in Transformers' Attention\", Brody et al., EMNLP'2023](https://lessw.medium.com/what-layernorm-really-does-for-attention-in-transformers-4901ea6d890e)\n    * [\"Understanding and Improving Layer Normalization\", Xu et al., NeurIPS 2019](https://arxiv.org/pdf/1911.07013.pdf)\n    * [\"On Layer Normalization in the Transformer Architecture\", Xiong et al., ICML 2020](https://arxiv.org/pdf/2002.04745.pdf)\nI'm not asking this question to make the authors just cite these papers, I am actually trying to understand whether this paper is consistent with, is orthogonal to, or contradicts these papers, since these are some of the most significant papers that were written about the analysis of LayerNorm that I know of, and none of them is discussed.\n\n## Summary\nOverall, the paper provides a thorough analysis and comparison of LayerNorm and RMSNorm, and provides insights that lead to 10% reduced inference and training time. If these improvements are valid and real, a 10% reduction is significant and important. \nHowever, there are many questions regarding the validity of the evaluation and the soundness of the claims. Therefore, I'm currently voting for a borderline reject, until the authors clarify all the issues above. I will raise my score if all issues are explained and resolved.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2023-06-16"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis study explores the relationship between Pre-RMSNorm and Pre-LN Transformers, demonstrating that these two variants can be theoretically reparameterized into one another. Additionally, the authors introduce a novel Transformer variant called Pre-CRMSNorm, which reduces one hidden dimension while maintaining the same representation power as Pre-LN and Pre-RMSNorm. Experimental evaluations on ViT and GPT models reveal that the proposed approach significantly improves computation efficiency during both model training and inference stages.\n\nSTRENGTHS:\nThe findings are interesting and reasonable. The proposed approach is novel and demonstrates a new direction to improve model efficiency via parameterization. The authors conduct experiments on both image data (i.e., ViT) and natural language (i.e., GPT).\n\nWEAKNESSES:\n1. The study would benefit from additional evaluations of the proposed method beyond computation efficiency. While the comparisons conducted on computation efficiency are useful, it is crucial to include discussions on the effectiveness of the resulting model. The absence of such results makes it challenging to assess the proposed method and its claims, particularly regarding training speedup. For example, it remains unclear whether the proposed reparameterization yields comparable training performance and stability when compared to Pre-LN Transformers.\n\n2. The efficiency gain diminishes as the model size scales larger, which affects the empirical contribution of the proposed method. Additionally, without sufficient discussions on model performance, it is unclear whether the proposed method is compatible with low-precision data for model inference and training, which plays an important role in making deep learning efficient. Incorporating experiments with fp16 and bf16, as well as exploring the use of quantized models, would significantly enhance the paper and provide valuable insights.\n\n3. The main finding of the study fails to provide a very deep insight beyond the reparameterization connection. Also, I feel the equivalent claim is a bit ambiguous, and would recommend the author phrase the claim as reparameterization. This is crucial because different reparameterizations for the same network can lead to substantial variations in training stability, training performance, and even inference performance. Emphasizing the reparameterization aspect would address this ambiguity.\n\nQUESTIONS:\nWhat would be the computation efficiency difference between the LayerNorm and RMSNorm from NVIDIA’s apex package and the customized implementation?\n\nLIMITATIONS:\nThe authors mentioned the potential compatibility issue with low-precision data types but did not provide adequate discussions on this issue.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes modifications to the popular transformer architecture's normalization mechanism in order to improve efficiency without sacrificing performance. It starts out with the baseline architecture Pre-LN (LayerNorm), and derives two new architectures: Pre-RMSNorm and Pre-CRMSNorm, which are inspired by the use of RMSNorm instead of LayerNorm mechanism in certain transformers in the literature to improve efficiency without noticeably affecting performance.\n\nThe paper formally proves the arithmetic equivalence of all three architectures. It also evaluates the two proposed architectures over the baseline by conducting experiments on ViT and GPT3.  It reports modest improvements in training time (around 2%) and slightly bigger improvements in inference time (1-9%) for a variety of different hyperparameter settings.\n\nSTRENGTHS:\nThe paper sheds light on normalization mechanisms in transformer architectures by relating two of the most popular such mechanisms: LayerNorm and RMSNorm.\n\nBy starting out with the empirical observation in the literature that RMSNorm is more efficient than LayerNorm, it derives two new architectures based on RMSNorm and proves them to be arithmetically equivalent to the LayerNorm architecture.\n\nThe empirical evaluation is quite thorough, using both ViT and GPT settings, and demonstrates modest improvements in both training and inference times across a wide range of hyperparameter settings.\n\nFinally, the presentation is crystal clear: the paper is extremely well written even for an informed outsider to follow.\n\nWEAKNESSES:\nThe single major weakness is that the savings in training and inference time is not too significant (around 1-3% in training time and about the same or slightly more in inference time, depending on the setting).\n\nThe authors acknowledge this by noting that normalization is not the elephant in the room, and that attention and MLP dominate the computation especially as the model size grows.  They even show an empirical upper bound on the speedup (which is the time taken by the normalization mechanism as a fraction of the overall time -- around 12-18%).\n\nA minor nit: The last sentence of the abstract \"Experiments demonstrate that we can reduce the training and inference time of Pre-LN Transformers by up to 10%\" is a bit misleading.  I would appreciate it if you paraphrase this sentence, by 1) separating out the training and inference time numbers, and 2) providing both lower and upper bounds instead of just the 10% upper bound.\n\nQUESTIONS:\nNone\n\nLIMITATIONS:\nYes, the authors have adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes two novel modifications to the Pre-Layer Normalization (Pre-LN) Transformers, introducing the Pre-Root Mean Square Normalization (Pre-RMSNorm) and Pre-Compressed Root Mean Square Normalization (Pre-CRMSNorm) Transformers. The authors aim to improve computational efficiency by simplifying LayerNorm to RMSNorm and introducing a lossless compression technique for zero-mean vectors. They claim these changes maintain the arithmetic functionality of the original models, reducing training and inference time by up to 10%. The paper includes extensive mathematical proofs and technical diagrams to support their claims. Experiments using Vision Transformer (ViT) and GPT models demonstrate the improvements in speed.\n\nSTRENGTHS:\nThe paper presents novel ideas in Transformer optimization, leading to improved computational efficiency. The use of compression in the Pre-CRMSNorm Transformer is particularly inventive. The detailed mathematical proofs and the experimental results provide a strong basis for the claims made, enhancing the paper's quality and clarity. Also, the potential for the method's applicability in various domains enhances the paper's significance.\n\nWEAKNESSES:\nAlthough the paper presents comprehensive mathematical proofs, the practical effectiveness of the proposed model is not extensively demonstrated. More empirical results across different domains and data types would further substantiate the authors' claims. The authors acknowledge that the proposed modifications might slightly increase the training workload, which could be a limitation for some applications. Lastly, the handling of vectors in $\\mathbb{R}^{d-1}$ by accelerators is assumed, which might not hold true in all scenarios.\n\nQUESTIONS:\n1. Could the authors provide more empirical evidence of the model's performance on real-world tasks, including various data types and domains?\n2. How do the proposed models handle scenarios where the input and output of certain blocks are not zero-mean?\n3. Are there any specific scenarios where the proposed models might underperform compared to the original Pre-LN Transformer?\n\nLIMITATIONS:\nThe paper acknowledges that the input and output of certain blocks are assumed to be zero-mean, which might not be the case in all scenarios. There might also be potential issues in the decompression stage of the CRMSNorm method that are not discussed. Further empirical validation is necessary to fully understand the impact of these modifications on various use cases and data types. Also, the proposed models might not be as efficient when the handling of vectors in $\\mathbb{R}^{d-1}$ by accelerators is not optimal.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-12"}