Pre-RMSNorm and Pre-CRMSNorm Transformers: Equivalent and Efficient Pre-LN Transformers
Authors: Zixuan Jiang
Jiaqi Gu
Hanqing Zhu
David Z. Pan
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Transformer
Normalization
Layer Normalization
RMSNorm
Efficient Machine Learning
TL;DR: We unify LayerNorm and RMSNorm in pre-normalization Transformers and propose equivalent and efficient Transformer variants.

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
Zixuan Jiang, Jiaqi Gu, Hanqing Zhu, David Z. Pan Chandra Department of Electrical and Computer Engineering The University of Texas Austin Austin, Texas, 78712
{zixuan, jqgu, hqzhu)@utexas. edu, dpan@ece btan@ece.utexa.... wtw.us.eeee.....
Abstract
Transformers have achieved great success in machine learning applications. Nor- malization 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 be- ing 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 convert Transformers with one normalization to the other type. While there an ongoing disagreement between the two normalization types, propose solution to unify two mainstream Transformer architectures, Pre-LN and Pre-RMSNorm Transform- 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 PrMS(O) counterparts at almost no cost, offering the same arithmetic functionality along with free effi- ciency improvement. Experiments demonstrate that we can reduce the training and inference time of Pre-LN Transformers by 1% 10%.
Introduction
Transformers have become successful architecture for a wide range of machine learning applications. including natural language [39], computer vision [13], and reinforcement learning [7]. 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. LayenNorm is in the original Transformer architecture [39], recentering and rescaling the input vector i2d to obtain a zero-mean and unit-variancs output. RMSNorm only rescales the input vector with its RMS value, offering greater computational efficiency than LayerNorm.
* 'This work done when all the authors UT Austin. Currently, Zixuan Jiang is Google, and Jiaqi Gu Arizona State University.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

The machine learning community does not reach a consensus regarding the preferred normalization technique for Transformers. LayerNorm demonstrates remarkable success in the milestone Trans- formers, such as GPT [30, 5] and VIT [13]. It is still the default normalization layer when building new Transformer. On the contrary, RMSNorm 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- Fre-RMSNorm Transformers, we identify an opportunity to unify them by removing the inherent redundancy 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 the functionality of the models, which allows to reduce LayerNorm to R.MNNrmm.
We further propose Compressed RMSNorm (CRMSNorm), which takes a vector in Rd-1 as input, decompresses tor in Rd and applies the RMSNorm on the de- compressed vector. Building upon this new nor- malization, we propose Pre-CRMSNorm Trans- former that employs lossless compression on the zero-mean vectors. We apply such compression to zero-mean activations and parameters Pre- RMSNNrm Transformers, enhancing the efficiency of Pre-RMSNorm Transformers while preserving the equivalent arithmetic functionality.
Pre-LN
Equivalence
to a zero-mean vec-
Sec. 3.1
Sec. 3.3
remove redundant
mean information
Pre-RMS Our Proposed Transformer Variants
Sec. 3.2
Pre-CRMS
We formally claim that Pre-LN, Pre-RM Tre-RMSNorm, and Pre- Fre-CRMMSNrm Transformers are equivalent for both training and inference. Figure visu- alizes the overview of their equivalence. Such equivalence directly enables more efficient train- ing and deployment of Pre-LN Transformers. We can translate a Pre-LN model into an equivalent Pre-(C)RMSSorm model, which can be readily deployed or adapted to downstream tasks. The conversion process incurs minimal costs. We can also train Pre-(C)RMSNorm model directly as train an equivalent Pre-LN Transformer counterpart.
Sec. 3.4 Higher Efficiency
Figure 1: Overview of the three equivalent Trans- former variants.
Although our relative improvement 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 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 cefformannn--fffciienee 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 S7, 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://gitbe com/ Zi aaaaanninn//reeee. reteermsnomm--r wwa.t--rraanrreeeeeeoee
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 Trans- former 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 technique normalize the activations of intermediate layers of neural networks. Given a vector z E IR, LayerNorm normalizes to obtain a zero-mean vector,
p(z)
1"x
LayerNorm(z)=
where u(x) = d 00
(I)
| (2(x))
LayerNorm recenters and rescales the activations and gradients in the forward and backward compu- tations [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 different normalization approach. Instead of normalizing the inputs based on their mean and variance, RMMNNrm normalizes them based on their root mean square (RMS) value. It is defined in the following equation,
x
RMSNorm(z)
where >
| /
R.SNNrm 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 reported that replacing LayerNorm with RMSNorm can achieve comparable performance and save training and inference time by 7% 64% [44].
Given mean vector these two kinds of normalization equivalent. Formally, if μхx²) 0, then LayerNorm(a x) RMSNorm(z). We may optionally introduce learnable parameters and apply an element-wise affine transformation 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 Normalization in Transformers
Normalization plays 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 Transformers, especially the large models, such as ViT [13, 9], PaLM [8], and GPT-series models [30, 5].
RMSNorm is proposed
alternative normalization technique in Transformers. Several large language models, such as Chinchilla [16] and LLaMA [38], use Pre- Tre-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.
"The following discussion expressivity is out of the scope of this paper and is only for reference. Normalization helps accelerate and stabilize the training hurts model's expressivity. For example, we usually add learnable parameters 7, 8 after normalization to recover the model expressivity. LayerNorm removes the mean and variance information of the input vector, such that all the output must have zero mean and unit variance. RMSNorm only removes the scale information, it has less information loss.

---

(a) Pre-LN Transformer
(b) Our Pre-RMSNorm Transformer
(c) Our Pre-CRMSNorm Transformer
Preprocess
Residual Fi
Preprocess
X]
Residual F;
Preprocess
X
Residual Fj
LayerNorm
Recenter
RMSNorm
CRMSNorm
Xy g
B
Xxg g
E 1g
B 2-
Block
Linear
Block
Linear
Block 1
Linear
x
Ni
E R
Block
Block
Block L
6%
R
xx,
Xg
d-1
Linear
Linear out
Linear
LayerNorm
RMSNorm
CRMSNorm
Postprocess
Postprocess
Postprocess
X141
(a) = (b)
(o) Equivalence:
(b) (c) Efficiency
Figure 2: Left. The original Pre-LN Transformer architecture. Middle and Right. Our proposed Pre-
R.M.Nrrm and Pre-CRMSS -remCRmSNorm Transformer architectures. These three architectures equivalent. The differences are highlighted in bold and green blocks.
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 architecturr, 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[7J. We import special tokens (such as the classification token) and embeddings (such as positional embeddings), ultimately obtaining a sequence of token embeddings zo.
Transformer blocks. The main body of Pre-LN Transformer [40] is stack of residual blocks
E1+1 xt + Fix), = 0, l...., 1,
(3)
where æi is the input of the 1-th block, Fi a sequence of operators LayerNorm + Linear >> 9t + Linear. We name a the vectors the main branch and Fi the residual branch [14]. The block Fi is usually an attention or a multi-layer perceptron (MLP) module. If gi is an activation function, such GELU [15], then the block F, is two-layer MLP. If gi is (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 explicitly defined, we can add an identity mapping, special linear layer. Most of the learnable parameters of Transformers are in these two linear layers.
LayerNorm and postprocessing. We finally process the result LN(æ1) to obtain the task-related results, such as classification probabilities. We apply the layer normalization on σ since usually has 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, Prr-RM. and Fre-CRMSNorm Transformers are arithmetically equivalent.
Pre-LN Transformer Pre-RMSNorY Transformer Pre- Crre-RMSNorm 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.
'The elementwise affine transformation of LayerNorm or RMSNorm is also linear transtormation. Thus, it can be fused with the input linear projection. We disable this transformation in the normalization layers to simplify the analysis.

---

To clarify, our proposed Pre-RMSNome Transformer and the pre-existing Pre-RMSSonn 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 the discussions below.
3.1 Pre-LN Transformer Pre-RMSNorm Transformer
LayerNorm is invariant to the shifting LN + k1) = LN(x), Ak E R. We observe that LayerNorm applied the main branch vectors before they are used either residual branches or the final postprocessing in Pre-LN Transformers. Therefore, we can replace the main branch vectors æ; with an 1, vky E IR without impact on the functionality of the Pre-LN Transformer. If k = - (u(r)) replace the main branch vectors with its recentered version æ; =(xx))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 Pre- RMSNorm Transformer.
1. Recenter the zo before the first Transformer block, where Recenter(a:) x p(x)1.
2. For the output projection in residual branches, replace the weight A. and bias b, with
Å. A. 1 A., b, b, ((b,)) 1, where d is the dimension of æt-
3. Replace LayerNorm with RMSNorm at the beginning of residual blocks and before postpro- cessing.
Since ((æ1++1)= μ(x;) + (x:F( we can keep zero-mean on the main branch if and only the input of the first block 20 and the output of each residual branch Fi are re-centered with zero-mean. The first modification is to recenter æo, while the second modification is to recenter the output of residual branches. For the residual branch Fi, 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 linear transformation y Aa b, R A Rmxn b, y E m
decompose the output with two parts y Ax + b p(y) The first part Ax + y p(y)1, with zero mean, another linear transformation with Á A 1 A, b (bb)
The first two modifications ensure that maintain zero mean the main branch vectors. Given zero-mean input, LayerNorm is equivalent to RMSNorm, which implies that the third modification has impact on the functionality. With these modifications, demonstrate that Pre-LN and re-RRMSNorm Transformers are equivalent.
3.2 Pre-RMSNorm Transformer = Pre- CRMSNorM Transformer
For zero-mean vector IRd, compress it losslessly by discarding its last element. In the decompression, we recover the discarded element with va =- -/ 24. 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 d input. CRMSNorm first decompresses the vector to obtain zero-mean vector in Rd then applies RMSNorm on the zero-mean vector. can generate the normalized zero-mean results in either dR or Rdd, Its formal definition is in Equation 5.
CRMSNorm(z)
where Rd-1
(5)
2d
i=
We simplify the Pre- Fre-RMmNorm Transformers obtain the Pre-CRMSNorm Transformers with the following modifications.
Compress the zero-mean main-branch vectors from I&d Rd-1, Preprocessing and postpro- cessing handle compressed vectors in d
5

---

2. Replace RMSNorm with CRMSNorm.
3. Simplify the weight A, in the input projection layer. Let ag be the last column vector of A,-
A; = A; - a;1" "1 is the compressed weight matrix.
4. Simplify the weight and bias in the output projection layer. We discard the last row of the weight matrix A, and the last element of the bias 6, since they are used to generate the redundant last element.
We compress the zero-mean activations in our proposed Fre-RMSNorm Transformers and corre- spondingly simplify the two linear layers in residual branches.
We fuse preprocessing. recentering, and compression. The fused preprocessing generates compressed vectors in Rd- to represent zero-mean vectors in R2 Taking language models as an recenter and compress (discard the last element) the word and position embeddings.
example,
For the input linear projection, its input the output of RMSNorm, whose mean zero. We compress the output of the RMSNorm and the weight of the linear layer. Specifically, if the linear layer takes zero-mean vector as input, then
Linear(x) A,æ + b. (A, aa1")az ) +
(6)
Á, A, a is the compressed weight matrix since its last column is a zero vector. The simplified linear layer only needs the compressed zero-mean vectors Rd-1 as input.
For the output linear projection, we only need to calculate the first (d - 1) elements for the vectors in Id. Thus, we can compress the weight A, and bias b.. As shown in Figure 2, the shape of A. is (d, d.). we can compress it directly to (d - 1, d,) by discarding its last row. Similarly, we compress b. Rd by discarding its last element. This weight compression also reflects their redundancy. As shown in Section 3.1, 6, and all column vectors of A, are zero-mean.
Figure demonstrates the difference vector dimension. We demonstrate the equivalence between Fre-RMSNorm and Fre-CRMSNorm Transformers.
3.3 Pre-LN Transformer = Pre-CRMSNorm Transformer
To translate a pre-trained Pre-LN Transformer to Pre-CRMSNorm Transformer, can convert into a Pre-RMSNone model and then finish the conversion. For the model implementation, we need two steps convert Pre-LN Transformer Pre-C CRMSNorm Transformer.
1. Reduce the hidden dimension from d to
2. Replace LayerNorm with CRMSNorm.
Namely, Pre-LN Transformers with the main branch vectors Rdd are equivalent to Pre-CRMSNorm Transformers in Rd-1 which further echoes the redundancy in the Pre-LN Transformers.
3.4 Training and Inference Efficiency
We show how 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 (CMMSNomo 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 replace the weight and bias in advance. During inference, can treat it
for inference since as a standard linear layer with equivalently transformed parameters.
"We reduce the hidden dimension in the preprocessing. main branch, two linear layers in residual branches. and postprocessing. We keep the d,, d, in residual branches. For instance, we maintain the MLP dimension and head dimension (dimension of query, key, and value vectors).

---

Recenter
Zero-Mean Linear
RMSNorm
Training
little increase
a little increase
decrease
Inference
same
same
Table 1: The computation workload of our Pre- Rr-MRMMNorm model compared with the original Pre-LN Transformer.
We have to pay the extra cost for training for the parameter change. The parameter optimizer manages Ao: bo but we use their recentered version A., b, during training. The induced cost is small for several reasons. (1) The size of parameters A,, b, 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.. b. can be done ahead of time since it does not depend on the input. We may leverage the idle time of accelerators compute them. (3) In data-parallel distributed training, the parameter server [19] manages the parameters, Each worker will receive A., b. from the server and then pass the gradients of loss to Aor b, to the server. Only the server needs maintain and update the original Aor bo In short, it is much easier to process the model parameters than the intermediate activations.
Recentering. It is possible 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 x++y)) = Recenter +eeeeterr(() 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, 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 induced. For training, recentering induces extra cost since on the fly.
Replacing LayerNorm with RMSNorm. Section 2.1 introduces that RMSNorm can achieve speedup compared with LayerNorm, demonstrated by the previous models. This replacement can help us accelerate the training and inference of the Pre-LN Transformer.
3.4.2 Pre-CRMMNorm 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 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 Rd-1 when d is 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 Rd-11 vectors may take much more time than Rdd vectors. Thus, we need to examine if the accelerators can support d 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-RMNNomt variant. For training, suggest keeping the hidden dimension d, which is equivalent to a Pre-LN model with the hidden dimension d+1 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 6, and be Pre-RMSNorm Transformer with parameter Φb We demonstrate that for any input æ E Rd we can always have f(x 0) g(ær, & h())) and we show how we conduct the conversion ø h(().. Thus, we claim that and g are equivalent terms of arithmetic functionality. Namely, when calculating f 0) during inference, we can always compute the cquivalent counterpart g(x, á h(0)). Our
7

---

proposed method is ee-parameterization technique [12], which builds equivalent models with different parameters.
Training equivalence, Despite 0) g h())), there may be a large difference in the convergence speed and stability when training and g with SGD and its variants. We avoid this issue by applying the gradient updates on 0 instead of 6. Specifically, during the training process, we maintain Pre-LN Transformer model f and its parameter 0. In the forward and backward computation, we convert the model and its parameter into 9 and ø to save training time. We do not update @ directly. Instead, we calculate the gradients V0 and call the SGD optimizer where 1) is the learning rate. From the perspective of optimization, the training is still on the Pre-LN Transformer model f and its parameter A. Consequently, the training performance and stability are preserved and equivalent to the original Pre-LN Transformers.
- v/.
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 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-decode [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 resolution of 224 x 224 and generates 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. The Pre-CRMSNon Transformers are compatible with dropout since uses vectors in dd- to represent the compressed zero-mean vectors in Rd 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 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-RMSNom Transformer can reduce the inference time by 1% - 9%. This reduction takes effect for ViTs with different models and various mini-baich sizes. The left subfigure demonstrates the inference latency, which is the inference time when the batch size is
The proportion of LayerNorm in the total computation is 12% 18% in these experiments. As discussed in Sections 2.1 and L. 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 Tre-RMMNorm models.
Droooout only impacts training and has impact on inference. Thus, the fe-RRMNorm variant can also be used for inference.
8

---

Inference time with batch size
Inference time of VIT-Ti/66
Normalized Inference Time
Normalized Inference
0.95
4.%
97 22%
97
77%
97
3%
97
7%
0.95
96
96.7%
96.1%
96.2%
93. 93.9%
93.4%
0.9
89.0%
0.85
87. 1%
87.7%
87 1%
87.7%
0.85
86.5%
85 5%
85.5%
84
3%
83.0%
81.8%
82 82.2%
0.8
0.8
Ti/16
S/16
B/16
L16
H/14
G/14
16
84
256
1024
VIT Variants
w/o norm
Pre-LN
Prr--MS
Pre-CRMS
Batch Size
Figure 3: Normalized inference time on ViT with different model sizes and batch sizes.
Training time breakdown
Training time of ViT Variants
Normalized Training Time
Pre-CRMS
87.4%
10.9%
Pre-RMS
87.4%
9.3%
0.1%
0.9%
98.3%
99.6% 8%
T
9%
99.4%
98.6%
100.69
96, 98.9%
Pre-LN
87.4%
12.6%
98.3%
99.1%
97 97.7%
97.7%
w/o norm
87.4%
0.95
TV/6
S/16
B/16
L/16
H/14
G/14
0.8
0.85
0.9
0.95
Variants
others
normalization
recenter linear
Pre-LN
Pre-RMS
Pre-CRMS
(a) Training time breakdown on ViT-TV16.
(b) Normalized training time
Figure 4: Training time comparison and breakdown ViT variants
For Pre-CRMSNorm, the GPU cannot efficiently handle vectors in dR in some cases, where add zero padding to obtain vectors in Rd. With zero padding, Pre-CRMSNorm models are less efficient than Tre-RMSNorm ones due the extra decompression. However, for the cases where Rd--1 vectors can be accelerated efficiently, we can achieve an even 10% time reduction.
We also conduct experiments (1) with other precisions, (2) 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 Appendix D for more details.
Training. We train ViTs on single workstation with 4 A100s with data parallel training [20], follow- ing 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 broadeasting the new parameters. Figure 4a visualizes the breakdown of the related components. In the Pre-LN Trans- former, the layer normalization accounts for 12.6% of total training time. For the Pre- Ree-RMMNorm variant, 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 9.27%. Overall, the Pre-RMSSorr reduces the training time by 2.36%.
We train the Pre-CRMSNorm Transformers with as the hidden dimension since we do not obtain speedup from the compressed dimension since the GPUs cannot handle dR vectors efficiently. The CRMSNorm is more computationally expensive than the RMSNorm given the same input but takes less time than LayerNorm. The Fre-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-CRM Pre-CRMSNorm variant is 1.74% faster than the Pre-LN model.
Figure 4b illustrates the training time 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 and sequence length 512. The results are shown in Figure 5. In small GPT-3 models, layer normalization takes considerable time. Hence, replacing the Layer- Norm with RM)SMoor can reduce the infer- ence time by up 10%. However, the model grows, the attention and MLP take charge of the main part of the computation [18]. The nor- malization takes < 1% of the inference time for GPT-3 XL and larger models, which is the upper bound of the speedup with our methods.
w/o norm ---.Pre- Pre-Pre-LN
Pre-RMS
Pre-CRMS
Normalized Inference Time
0.95
0.85
XL
2.7B
6.7B
GPT3 variants
Applying quantization may mitigate this issue to some extent. By applying int8 matrix multi- plication [10], the percentage of the layer nor- 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.
Figure 5; GPT-3 inference performance
5 Conclusion
In this paper, propose two equivalent and efficient variants for the widely used Pre-LN Transform- ers. We point out the inherent redundancy in the Pre-LN Transformer, By maintaining zero- zero-mean mean on the main branch vectors and thus simplifying LayerNorm, we obtain the Pre-RMMNorm 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 (OMMMSom Transformer with better training and inference efficiency, which is a free lunch. We strictly push the performance-efficieng Pareto frontier of foundational Pre-LN Transformers. As result, pre-trained Pre-LN Transformers (such 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 different contexts would be beneficial. Additionally, while focus has been on pre-normalization Transformers, it would be valuable 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.
Linitations. Detailed implementations and specific optimizations will play crucial role in achieving practical performance gains. It is necessary to focus on developing highly optimized implementations of (CSS. to bridge the gap between theoretical and practical efficiency improvements. By addressing these challenges, we can fully leverage the potential benefits of (C)RMSNorn and enable more efficient utilization of Pre-LN Transformers.
Acknowledgments and Disclosure of Funding
We acknowledge NVIDIA for donating its 1100 GPU workstations and the support from TILOS, an NSF-funded National Artificial Intelligence Research Institute.
10

---

References
[1] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton, Layer normalization, arXiv preprint arXiv:1607.06450, 2016.
[2] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.0720 2021.
[3] Yelysei Bondarenko, Markus Nagel, and Tijmen Blankevoort. Understanding and overcoming the challenges of efficient transformer quantization. arXiv preprini arXiv:2 109. /2948, 2021.
[4] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Mihne and Qiao Zhang. JAX: composable transformations of Python+NummPy programs, 2018.
[5] 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 33:1877-19011 2020.
[6] 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.044. 240.0026,. 2023.
[7] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learnin via sequence modeling, Advances in neural information processing systems, 34:15084-15097,
2021.
[8] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmano, et al. Palm: Scaling language modeling with pathways. arXiv preprint urXiv:2204.02311 2022.
[9] 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.
[10] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 (): 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.073338 2022.
[11] Jacob Devlin, Ming-Wos Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 1rXiv:1810.04800 2018.
[12] Xiaohan Ding, Xiangyu Zhang, Ningning Ma, Jungong Han, Guiguang Ding, and Jian Sun. Repygg: Making vgg-style convnets great again. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 13733-13742, 2021.
[13] 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.1 19999, 2020.
[14] 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.
[15] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint 16X1v:1606.08415, 2016.
[16] 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. 2022.
[17] Gao Huang, Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger: Deep networks with stochastic depth. In Computer Vision 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part IV 14, pages 646-661. Springer, 2016.
[18] Cheng Li. LIm-analysis: Latency and memory analysis of transformer models for training and inference. https: / thub. cow/clii///////c/ooooocccocl 10119/11mmm 11mm wllmmmaaasss,a...m.mm 2023.

---

[19] Mu Li. David G. Andersen, Jun Woo Park, Alexander J. Smola, Amr Ahmed, Vanja Josifovski, James Long, Eugene J. Shekita, and Bor- Ying 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.
[20] 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. 10066.157000 15704, 2020.
[21] Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. arXiv preprint arXiv:2004. 08249,, 2020.
[22] Zhenhua Liu, Yunhe Wang, Kai Han, Wei Zhang, Siwei Ma, and Wen Gao. Post-training quantization for vision transformer. Advances in Newral Information Processing Systems, 34:28092-281033 103, 2021.
[23] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
[24] Paulius Micikevicius, Sharan Narang. Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, al. Mixed precision training. arXiv preprint 17 2017.
[25] Aaron Mok. Chatgpt could cost over $700,000 per day to operate. microsoft is reportedly trying to make it cheaper. weppmmtteeerrrt-2233-3-- 2023-4. Accessed: how-much-chatere whw.hhhttcssts--peeattee to--eaaiittoouutttiiimee www---rueesttttteeeepeeee 2023-05-01.
https://wwrr businessiniidddr.. com/
[26] 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 modifi- cations transfer across implementations and applications? arXiv preprint 2rX11:2102.11977, 2021.
[27] 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.
[28] 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 Detito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, Pigh-Peeffrraacc Deep Learning Library. In H. Wallach, H. Larochelle, A. Beygelzimer, d' Alché Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 8024 8035. Curran Associates, Inc., 2019.
[29] 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.
[30] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[31] 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 1xi1:2112..14446 2021.
[32] 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-5511 2020.
[33] 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): 2014.
[34] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An fnolloition-folloiinng llama model. https: / /i thub. com ttsr//atsuulaas
ord, alpaca, 2023.
12

---

[35] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey, ACM Computing Surveys, 55(6):1-28, 2022.
[36] 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 machine learning, pages 10347-10357. PMLR, 2021.
[37] Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Go- ing deeper with image transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 32 32-42, 2021.
(38) Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et Llama: Open and efficient foundation language models. arXiv preprint arXiv:2300.. 13971, 2023.
[39] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[40] 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.
[41] Jingjing Xu, Xu Sun, Zhiyuan Zhang, Guangxiang Zhao, and Junyang Lin. Understanding and improving layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
[42] 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. 2023.
[43] 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.
[44] Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
[45] 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 2ivv 18223, 2023.
13

---

Appendices
A Proof for Lemma 1.
Given linear transformation y Aa b, R" A E Rmxn b. R we have
y = Ar + b
(7)
(A k1 T A)æ: + k11" Ax + (b ((b) + p(b)
(8)
(A k11 A) a) + (b (b)1) + 1(1² Ax)1 + (bb)
(9)
(A h11 A)æ + (b-14) (0)1 + (k1" Ax ar++((b))
(10)
Aæ + b+ (x, k)1
(11)
where Á A- k11 8 = ((b)1, fzz k) k1 Ax + p(b).
1 = 1/m, then we obtain
1 1 (A
1 -117 A) - (1" - 1 1"AAAAAAA A)æ =0
1
u ((xx)
(12)
m7
m
m7
(1y = (m(Ax) - -m(b) + - 1/m)11
(13)
=0 + 00 f(x == 1
(14)
1 1" Ax u(b)
(15)
m
The Å 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 + y with zero mean, another linear transformation with
Å A 1 11 A, b b- (b)11
The second part corresponds to the mean information ((y)1 = (L117 Ar + ((b))
B Post-LN Transformers
Different from the Pre-LN Transformers, the Post-LN Transformers have the following blocks.
3t+1 = LN(#; + Fix))) 0, 1,.,, L 1,
(16)
Layer normalization is 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.
F1+1 LN(æ: + Fi(z:))
(17)
LN(( (u(x:) + (Fi(m)) - ((F((m))
(18)
LN(E; + Fi(x))
(19)
RMSNorm(2:+ f(x)))
(20)
For the residual branch Fú, apply the same method in the Pre-LN Transformer. We modify the output linear projection obtain Fi, which generates the zero-mean part of the original result.
The recentering operation ④; I; μ(x;)1 requires extra computation. elementwise affine transformation is disabled in LayerNorm, is the output of normalization such that μ(œ;) 0 and क æt. If the transformation is enabled, ay 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 listed in Sections 3.1 and 3.2. These steps are fully reversible. We write the inverse steps explicitly below.
Coverting Pre-R Ro--RMMNorm into Pre-LN Transformers.
14

---

1. Remove the recentering.
2. A. A. 1c", where can be random vector, b, 6 + d1 where d can be random scalar.
3. Replace RMSNorm with LayerNorm.
Coverting Pre-CRMSNorm into Pre-RMSNorm Transformers.
1. Decompression vectors in Rd-1 into zero-mean vectors in Rd. The compression is lossless, so the decompression its invertible operation.
2. Replace CRMSNorm with RMSNorm.
3. A; A; a.1", where ag is random vector.
4. Recover the last row of A, and the last element of 6, such that each column of A, and b. 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 the following implementations.
For JAX, we use APIs of LayerNorm and RMSNorm in the flax library. For PyTorch, we use the implementations of LayerNorm and RMSNorm from NVIDIA's apex extension 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 ex- ample. PyTorch has provided the official LayerNorm API but lacks RMSNorm implementa- tion. These different implementations are mixed. We do not find one implementation domi- nant over others for all the cases. For instance, torch nn LayerNorm is usually faster than apex. normal izattiin....... 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
Name
Dimension
Depth
Heads
MLP Dimension
Tiny-16
192
12
3
192 x
Small-16
384
12
384 x
Base-16
768
12
12
768 X 4
Large-16 16
1024
16
1024
Huge-14 14
1280
32
16
1280 4
Giant- 14
1664
8192
Table 2: ViTs with different sizes. The number in the model name the patch size.
no norm
Pre-LN
Pre-RMS
Pre-CRMS
Porcc,, single A100, amp amp - float32
0.8567
1.000
0.9699
0.9783
0.9353
1,000
0.9850
0.9951
single A100 + 16-thread CPU
0.8697
1.000
0,9012
0.8857
Poroorc > JAX
0.9610
1.000
0.9873
1.0005
Table 3: Normalized inference time of VIT.
Table 2 lists the architecture parameters of Vision Transformer. We first measure the inference time. We sweep these ViTs with batch sizes (1, 4, 16, 64, 256, 1024) and collect the medians of these
"haps//githubbcom///iiDaaa
15

---

Train VIT-S/16 on ImageNet- 1k from scratch
85
Test (%
80
75
70
65
60
65
100
200
300
400
500
700
800
epochs
Pre-LN
Pro-RMSNorm
Pre-CRMSNom
Figure 6: The test accuracy of training ViT-S/16 on ImageNet- 1k from scratch.
Train loss of GPT-2 Small
Evaluation perplexity of GPT-2 Small
perplexity
35
Train loss
25
2.5
15
15
steps (10?)
steps (10")
.... ....Prre-NN
Pre-RMSNorm
Pre-CRMSNorm
Pre-LN
Pre-RMSNorm
Pre-CRMSNorm
(a) Training loss
(b) Evaluation perplexity
Figure 7: We train GPT-2 Small the wikitext-103rraww// dataset from scratch.
36 data points. We report the average of these 36 experiments in Table 3. We conduct inference on single A100 with automatic mixed precision (amp) 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 p(x+ + = p(x) )) + ((y) always holds for any vectors y, However, if these two vectors represented as (low precision) floating numbers, this equality is not guaranteed in real-world numerical computation. 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-meat sero-mean in the main branch numerically.
is possible that
The numerical issue is a common problem in machine learning. A typical example is operator reordering and layer fusion. PyTorch provides related API officially, named torch. ao quantization on. fuse modules. We can fuse the convolution layer and its following batch normalization layer to simplify the computation. These two layers 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 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 similar issue as layer fusion since we modify partial parameters. In our experiments, 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'ss 's official LayerNorm implementation with the apex may have larger impact on the model performance.
16

---

D.4 Verification of Training Equivalence
For validation and reference, we train ViT-S/I6 as Pre-LN, Tre-RMMNomm, and Pre- CRMSNorm models on ImageNet-Ie 1k following the DeiT-3 setting and achieve similar accuracy. The test accuracy over 800 epochs is shown in Figure 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- 100rrx---00--ww-v 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.
17


Reviews:
[{"review_id": "CElTaNIjWg", "reviewer": "Reviewer_AcCQ", "summary": "The 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.", "questions": "1. 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?", "limitations": "The 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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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.", "weaknesses": "Although 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for your insightful comments. We appreciate the opportunity to address your comments in this response.\n\n### **Question 1, more empirical results on various domains and data types**\n\n**Domains.**\nWe present experimental results on ViT and GPT3, covering the field of computer vision and natural language. These two models also represent two mainstream architectures of Transformers, encoder-only and casual decoder. Other Transformers can be treated as an extension of these two architectures, e.g., encoder-decoder and non-causal decoder.\n\nWe also conduct our experiments on Decision Transformers [1] in reinforcement learning. It adapts the GPT2 model. On the OpenAI Gym benchmark, the model is small with only 3 layers, 1 attention head, and an embedding dimension of 128. Consequently, our method significantly reduces offline training time (2-6%) and testing time (5-13%).\n\n**Data Types.**\nWe provide three levels of precisions in our paper.\n1. By default, we use automatic mixed precision (AMP) for training and inference (Line 276).\n2. Lines 338 - 342. By applying int8 matrix multiplication, the percentage of LayerNorm increases to 10% for GPT3 XL and 2.7B. Our Pre-RMSNorm can reduce the inference time by 4%. As a reference, the rate and the reduction are <1% in AMP.\n3. Appendix C.2. We provide results of disabling AMP (i.e., using float32 precision). The normalization takes less percentage, leading to a smaller speedup with Pre-RMSNorm.\n\nWe believe that our paper encompasses representative model architectures, diverse domains, and various data types. We are always open to adding more experiments if you can suggest specific settings.\n\n[1] Chen, L., Lu, K., et al. Decision transformer: Reinforcement learning via sequence modeling. NeurIPS 2021.\n\n### **Question 2, zero-mean assumption**\nWe discuss the general Pre-LN Transformer architecture without zero-mean assumption. In Figure 2a, we mark $\\mu \\neq 0$ in the main branch and the residual block.\n\nOur analysis of the Pre-LN Transformer uncovers that the mean information on the main branch is inherently redundant. We (1) add a Recenter after Preprocess, and (2) enforce the output linear layer to generate zero-mean output in the Pre-LN model. These two modifications can guarantee a zero-mean main branch and have no impact on arithmetic functionality. To achieve our Pre-RMSNorm Transformer, we further replace LayerNorm with RMSNorm since they are equivalent on zero-mean input. We mark $\\mu = 0$ in Figure 2b.\n\nIn summary, zero-mean is not an assumption in Pre-LN Transformers. It is an intentional architectural modification to achieve our desired Pre-RMSNorm models without impacting arithmetic functionality. We provide an example to show this conversion in our submitted code.\n\n### **Question 3, when proposed models are worse than Pre-LN Transformers**\nWe discuss the limitations at the end of the paper. Less-optimized implementations of (C)RMSNorm may eclipse our proposed methods. Highly optimized implementations of (C)RMSNorm are critical and will further improve our models’ efficiency.\n\n### **Vectors in $\\mathbb{R}^{d-1}$ on accelerators**\nWe discuss this issue in Pre-CRMSNorm Transformer in Section 3.4.2 and experimental analysis. In our experiments, the Pre-CRMSNorm in $\\mathbb{R}^{d-1}$ is usually slower than Pre-RMSNorm in $\\mathbb{R}^d$, since the accelerators cannot handle $\\mathbb{R}^{d-1}$ vectors efficiently when $d$ is a large even number (e.g., 1024).\n\nFor inference, we may either (1) add zero paddings, 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\nto a Pre-LN model with the hidden dimension $d+1$. In this way, the CRMSNorm can increase model representability and computation efficiency at the same time.\n\nWe have trained two groups of models in our early exploration.\n\n1. Pre-LN in $d = 1024$, Pre-CRMSNorm in $d - 1 = 1023$.\n2. Pre-LN in $d + 1 = 1025$, Pre-CRMSNorm in $d = 1024$.\n\nThe models in each group are equivalent. The accelerators can only handle $d = 1024$ efficiently, with poor performance on $d - 1$ and $d + 1$. To avoid unfair comparison, we only report the results of the first group throughout the paper. For reference, Pre-CRMSNorm in $d$ is 0.9% faster than Pre-LN in $d$.\n\n### **Slight increase in training workload**\nWe provide qualitative analysis in Table 1 and quantitative experimental results in Figure 4a.\n* Table 1. Although recentering and the special linear layer induce extra workload for training, RMSNorm can reduce the training workload significantly such that the total workload is reduced.\n* Figure 4a. The recentering and the special linear induce 1% overhead, while RMSNorm save 3.3% of training time. Overall, the Pre-RMSNorm reduces the training time by 2.3%.\n\nThere is no computation overhead for inference since the extra workload can be fused or done before inference.\n\n### **Decompression in the CRMSNorm**\nOur compression on the zero-mean vectors is lossless, so our decompression can recover the original vectors without loss. We appreciate your attention to this aspect and would be grateful for further details if you can pinpoint any specific concerns.\n\nThank you so much for reviewing our paper and this response. We are looking forward to further discussions with you."}, {"review_id": "Xv4p1dkhS1", "reviewer": "Reviewer_d6Jr", "summary": "This 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", "questions": "None", "limitations": "Yes, the authors have adequately addressed the limitations.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The 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", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for your insightful reviews. We are truly grateful for your acknowledgment of our paper's excellence in terms of soundness, presentation, and contribution.\n\nWe totally agree that our relative efficiency improvement is modest. As you correctly pointed out, our quantitative analysis shows that normalization takes a small portion of the computation. Moreover, we elaborate on three arguments to justify the significance of this improvement in Section 1 (Lines 65 - 78).\n> 1. The improvement comes from equivalence instead of tradeoff. We can strictly push the performance-efficiency Pareto frontier.\n2. The modest relative saving can be translated into significant absolute improvement, especially when applied to widely-used Transformers.\n3. Our method is orthogonal and complementary to most work improving efficiency. Our method can effectively benefit from other methods. For example, we have larger efficiency improvement with low precision since low precision increases the percentage of normalization (Section 4.2 and Appendix C.2).\n\nWe genuinely value your suggestions and will diligently polish the paper, particularly enhancing the abstract to provide a clearer summary.\n\nOnce again, we express our gratitude for your thoughtful review of our paper and response. We are looking forward to further discussions with you.\n"}, {"review_id": "iqrOj66VmE", "reviewer": "Reviewer_j4qQ", "summary": "This 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.", "questions": "What would be the computation efficiency difference between the LayerNorm and RMSNorm from NVIDIA’s apex package and the customized implementation?", "limitations": "The authors mentioned the potential compatibility issue with low-precision data types but did not provide adequate discussions on this issue. ", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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). ", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you sincerely for your insightful feedback. We appreciate the opportunity to address your questions in this response.\n\n### **Weakness 3, reparameterization and equivalence**\nWe would like to initially clarify these two terms. Let $f$ be a Pre-LN Transformer with parameter $\\theta$, and $g$ be a Pre-RMSNorm Transformer with parameter $\\phi$. We demonstrate that for any input $x \\in \\mathbb{R}^d$, we can always have $f(x,\\theta) = g(x,\\phi=h(\\theta))$ and we show how we conduct the conversion $\\phi=h(\\theta)$. Thus, we claim that $f$ and $g$ are equivalent in terms of arithmetic functionality. Namely, when calculating $f(x,\\theta)$, we can always compute the equivalent counterpart $g(x,\\phi=h(\\theta))$.\n\nWe find that reparameterization in machine learning usually refers to a trick to enable the differentiability of discrete variables [1]. To avoid confusion and misunderstanding, we do not use that term.\n\nBeyond the equivalence and reparameterization, we would like to highlight one of our contributions: identifying the inherent redundancy in Pre-LN Transformers, which we believe could offer valuable insights to the community.\n\n### **Weakness 1, impact on training performance and stability**\nYour concern is that although $f(x,\\theta) = g(x,\\phi=h(\\theta))$, the convergence speed and stability may have a large difference with SGD and its variants (e.g., Adam). We avoid this issue by applying the gradient updates on $\\theta$ instead of $\\phi$.\n\nSpecifically, during the training process, we maintain a Pre-LN Transformer model $f$ and its parameter $\\theta$. In the forward and backward computation, we convert the model and its parameter into $g$ and $\\phi$ to save training time. We do not update $\\phi$ directly. Instead, we calculate the gradients $\\nabla \\theta$ and call the optimizer $\\theta = \\theta - \\eta \\nabla \\theta$. From the perspective of optimization, the training is still on the Pre-LN Transformer model $f$ and its parameter $\\theta$. Consequently, the training performance and stability are preserved and equivalent to the original Pre-LN Transformers. Therefore, our paper did not extensively discuss training performance, as the training results remain the same.\n\nIn the supplementary material, we provide both training and inference implementations for your reference.\n\n### **Weakness 2 and Limitation 1, low-precision**\nWe appreciate your attention to our use of automatic mixed precision (AMP) for both training and inference (Line 276). We agree that low-precision calculation plays a pivotal role in our method's performance.\n\nOur method can benefit from low precision. With high precision in Transformers, matrix multiplications in the self-attention and MLP dominate the computation, while the normalization only takes a small portion. With a lower precision, matrix multiplications can be significantly accelerated, implying that normalization will take a larger portion. The larger percentage normalizations take, the larger reduction we can achieve.\n\nTo illustrate this further, we have presented two examples in the paper.\n1. Lines 338 - 342. By applying int8 matrix multiplication, the percentage of the layer normalization increases to 10% for GPT3 XL and 2.7B. Our Pre-RMSNorm can reduce the inference time by 4%. As a reference, the percentage and the reduction are < 1% in AMP.\n2. Appendix C.2. We provide results of disabling AMP (i.e., using float32 precision). The normalization takes less percentage, leading to a smaller speedup with Pre-RMSNorm.\n\n### **Question on the computation efficiency difference**\nWe measure the inference time of the normalization layer on an input tensor with shape [1024, 1024]. The first 1024 is the batch size, and the second one is the dimension where normalization applies. RMSNorm is faster than LayerNorm since it does not need recentering.\n\n|  time (us) | LayerNorm | RMSNorm |\n|:----------:|:---------:|:-------:|\n|    apex    |   55.44   |  49.39  |\n| customized |   58.25   |  54.58  |\n\nPlease refer to Appendix C.1 regarding implementations of normalization. Our submitted code also involves all kinds of normalization implementations in both PyTorch and JAX.\n\nThank you sincerely for reviewing our paper and response. We are looking forward to further discussions with you.\n\n[1] Jang, E., Gu, S., & Poole, B. (2016). Categorical reparameterization with gumbel-softmax. ICLR 2017\n"}, {"review_id": "d9PoQcO5N0", "reviewer": "Reviewer_NSno", "summary": "The 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).", "questions": "## 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", "limitations": "NA", "rating": 7, "confidence": 3, "soundness": 1, "presentation": 3, "contribution": 3, "strengths": "## 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.", "weaknesses": "## 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks a lot for your insightful comments. Due to the word limit, we have to redirect you to our response to other reviewers. We sincerely appreciate your understanding.\n\n### **Question 1**\n\nThe original Transformer block is \n* $x_{l+1}=x_l+F_l(x_l)$\n* $x_{l+2}=x_l+F_l(x_l)+F_{l+1}(x_l+F_l(x_l))$\n\n$F_l(x_l) = F_l(x_l + k_l \\mathbb{1})$ since $F_l$ begins with a LayerNorm. The proposed block with shifting is\n* $y_{l+1} = y_l+k_l \\mathbb{1}+F_l(y_l+k_l\\mathbb{1})=y_l+F_l(y_l)+k_l\\mathbb{1}$\n* $y_{l+2} = y_l+F_l(y_l) + F_{l+1}(y_l+F_l(y_l))+(k_l+k_{l+1})\\mathbb{1}$\n\nThe main branch vectors are different $y_l=x_l+\\sum_{i=0}^l k_l\\mathbb{1}$. However, the difference has no impact on arithmetic functionality. $y_l$ is used in residual branches and final postprocessing. These two blocks begin with a LayerNorm, so they are invariant to this difference.\n\nWe also give an example to show the arithmetic equivalence in our submitted code.\n\n### **Question 2**\n\nPlease check our response to **Reviewer AcCQ on Vectors in $\\mathbb{R}^{d-1}$ on accelerators**.\n\n### **Question 3**\n\nPlease check our response to **Reviewer j4qQ on Weaknesses 3 and 1**. In short, we maintain and update Pre-LN model parameters to guarantee training equivalence.\n\n### **Question 4, Weakness 5.1**\nSince we guarantee training equivalence, we do not focus on datasets and accuracy.\n\nFor validation and reference, we train ViT-S/16 as a Pre-RMSNorm and a Pre-CRMSNorm model on ImageNet-1k following the DeiT-3 setting and achieve similar accuracy (81.5%, 81.4% for our models, and 81.4% for baseline).\n\n### **Question 5, Weakness 5.2**\nGPT-2 is publicly available at https://github.com/openai/gpt-2. In Section 2.1 of the GPT-3 paper [1], the authors describe that\n> We use the same model and architecture as GPT-2, including the modified initialization, pre-normalization, and reversible tokenization described therein, ……, we train 8 different sizes of model, ……\n\nThe hyperparameters of these 8 models (Small, Medium, Large, XL, 2.7B, 6.7B, 13B, 175B) are defined in Table 2.1 of the GPT-3 paper. We follow the naming and description in their paper to build GPT-3 models.\n\n[1] Brown, Tom, et al. \"Language models are few-shot learners.\" NeurIPS 2020.\n\n### **Question 6**\n\nWe focus on the computation of LayerNorm and RMSNorm without touching their optimization and expressivity. We discuss how we calculate them rather than their role in optimization. Hence, our work is orthogonal and complementary to your mentioned work, which focuses on expressivity [2], the effect on optimization [3], and the location [4] of LN in Transformers. We have also discussed two of them.\n* Lines 94-95. LN takes effect in both forward and backward computation [3].\n* Lines 109-114. Pre-LN Transformers are easier to optimize than Post-LN variants [4], which is the major reason that Pre-LN is much more popular than Post-LN. Hence, we focus on Pre-LN architectures and only discuss Post-LN in Appendix B.\n\n[2] On the Expressivity Role of LayerNorm in Transformers' Attention. EMNLP 2023.\n\n[3] Understanding and Improving Layer Normalization. NeurIPS 2019.\n\n[4] On Layer Normalization in the Transformer Architecture. ICML 2020.\n\n### **Weaknesses 1, 2, and 3**\n\nWe first define two modifications. M1: add a Recenter after Preprocess. M2: adjust parameters of the output linear layer in residual branches to generate zero-mean output. We define four models.\n* LN: The general Pre-LN model in Figure 2a\n* LN-M: Pre-LN with two modifications M1 & M2\n* RMS: The existing Pre-RMSNorm model, e.g., LLaMA\n* RMS-M: Our proposed Pre-RMSNorm model with M1 & M2 in Figure 2b\n\nWe prove that `LN = LN-M = RMS-M`, where `=` means arithmetic equivalence. `LN = LN-M` because the mean information is redundant on the main branch. The modifications do not hurt the expressivity of Pre-LN models. `LN-M = RMS-M` because `LN(x) = RMSNorm(x)` if x is zero-mean.\n\n`RMS` and `RMS-M` are not equivalent. Model definitions of them for training are different since M1 and M2 need to be handled explicitly. Conversely, they can share the same implementation for inference since M1 can be fused with preprocessing, and M2 can be done before inference (Section 3.4.1).\n\n**Model expressivity.** Although the expressivity is beyond the scope of our paper, we provide a new perspective on that.\n* If the two modifications (M1 & M2) harm the model expressivity, the expressivity ranking is `LN = LN-M = RMS-M < RMS`.\n* If M1 & M2 are helpful, the ranking is `LN = LN-M = RMS-M > RMS`.\n\n**Model conversions in the inverse direction.**\nPre-LN -> Pre-RMS needs 3 steps (Lines 156-160), and Pre-RMS -> Pre-CRMS needs 4 steps (Lines 185-192). These steps are fully reversible. We write the inverse steps explicitly below.\n\nPre-RMS -> Pre-LN\n1. Remove the recentering \n2. $A_o=\\hat{A_o} + \\mathbb{1}c^T$, where c can be a random vector, $b_o=\\hat{b_o}+d\\mathbb{1}$ where d can be a random scalar\n3. RMSNorm -> LayerNorm\n\nPre-CRMS -> Pre-RMS\n1. Decompression vectors in $\\mathbb{R}^{d-1}$ into zero-mean vectors in $\\mathbb{R}^{d}$. The compression is lossless, so the decompression is its invertible operation.\n2. CRMSNorm -> RMSNorm\n3. $A_i=\\hat{A_i}+a_d\\mathbb{1}^T$, where $a_d$ is a random vector\n4. Recover the last row of $\\hat{A_o}$ and the last element of $\\hat{b_o}$ such that each column of $\\hat{A_o}$ and $\\hat{b_o}$ is zero-mean.\n\n`{LN, LN-M, RMS-M, CRMS}` can be converted in any direction. `RMS` cannot be converted to and from them.\n\n**Normalization.** For zero-mean vector, the equation `LayerNorm(x) = RMSNorm(x)` holds. For any vector x, we only have `LayerNorm(x) = RMSNorm(Recenter(x))`.\n\n### **Weakness 4**\nPlease check our response to **Reviewer AcCQ on Slight increase in training workload**. In short, the extra training cost is a local overhead, while the efficiency improvement refers to a global reduction.\n\nThanks again for reviewing our paper. We are anticipating further discussions with you.\n"}]

Author rebuttal:
We sincerely appreciate the constructive feedback provided by all the reviewers. We are truly grateful for their acknowledgment of our paper's contribution and presentation. Engaging with the reviewers is a privilege we highly cherish. Should any queries arise, please do not hesitate to reach out to us. Thank you!

Meta-review: This paper proposes a solution to the ongoing debate between Layer Normalization (LayerNorm) and Root Mean Square Normalization (RMSNorm) techniques in Transformers. The authors introduce Pre-LN and Pre-RMSNorm Transformers and demonstrate that by removing redundant mean information, LayerNorm can be reduced to RMSNorm, resulting in higher efficiency without compromising representation ability. Experimental results showcase up to 10% reduction in training and inference time for Pre-LN Transformers. Overall, this paper provides a valuable solution for efficiently normalizing Transformer architectures, offering practical benefits for large language models and other Transformer-based applications. This paper is highly recommended for acceptance as a spotlight paper.
Decision: Accept (spotlight)