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

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. The 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. Finally, 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. Overall, 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).

Strengths

## Strengths + The authors provide a thorough explanation of LayerNorm vs. RMSNorm. + The authors show the relation between LayerNorm vs. RMSNorm, which I haven't seen in any other paper. + 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 - 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? - Similarly, I am not sure that the authors prove both directions when claiming in Section 3.2 that `Pre-RMSNorm = Pre-C-RMSNorm`. - Novelty - The authors claim to "propose Pre-RMSNorm". Is this proposal novel? aren't models such as LLaMA already using Pre-RMSNorm? - 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? - Many evaluation details are missing. - On which datasets do the authors perform the experiments? - 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?

Questions

## Questions 1. 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. 2. In most of the cases in Figures 3+4, C-RMSNorm is slower than RMSNorm. So what is the benefit of C-RMSNorm? 3. 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. 4. On which datasets were the experiments conducted? What was the downstream accuracy of each model? 5. What exactly are the GPT-3 models mentioned in the evaluation? 6. Are the claims and findings in this paper related to the findings in the following papers regarding LayerNorm? * ["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) * ["Understanding and Improving Layer Normalization", Xu et al., NeurIPS 2019](https://arxiv.org/pdf/1911.07013.pdf) * ["On Layer Normalization in the Transformer Architecture", Xiong et al., ICML 2020](https://arxiv.org/pdf/2002.04745.pdf) I'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. ## Summary Overall, 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. However, 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

1 poor

Presentation

3 good

Contribution

3 good

Limitations

NA

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.

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. 2. 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. 3. 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.

Questions

What would be the computation efficiency difference between the LayerNorm and RMSNorm from NVIDIA’s apex package and the customized implementation?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors mentioned the potential compatibility issue with low-precision data types but did not provide adequate discussions on this issue.

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. The 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.

Strengths

The paper sheds light on normalization mechanisms in transformer architectures by relating two of the most popular such mechanisms: LayerNorm and RMSNorm. By 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. The 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. Finally, the presentation is crystal clear: the paper is extremely well written even for an informed outsider to follow.

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). The 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%). A 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.

Questions

None

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

Yes, the authors have adequately addressed the limitations.

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.

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.

Questions

1. Could the authors provide more empirical evidence of the model's performance on real-world tasks, including various data types and domains? 2. How do the proposed models handle scenarios where the input and output of certain blocks are not zero-mean? 3. Are there any specific scenarios where the proposed models might underperform compared to the original Pre-LN Transformer?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

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.

Response to authors

Thank you for your response. >Q3. our contributions: identifying the inherent redundancy in Pre-LN Transformers I am not sure I agree. Pre-LN Transformers change the "direction" of the vector, while Pre-RMS don't. I thus argue that Pre-RMS is inherently less expressive than Pre-LN. In some special cases, when the input is "centered" in a certain way, they are equivalent, but I was not convinced that this holds for the general case. > Q4. Since we guarantee training equivalence, we do not focus on datasets and accuracy. Even if there is equivalence in theory, I would like to be sure that there is no empirical degradation in practice. Why can't you show training and test curves? Are all training loss and test loss curves completely identical? > Q5. GPT-2 is publicly available at https://github.com/openai/gpt-2 ... So, do you mean that this is just the 2.7B architecture, which can be referred to as either GPT-2 or GPT-3? Did you use the pretrained weights, or just the bare architecture which you trained from scratch? ## Weaknesses I am not sure I am following, since the explanations in your response do not use the same terms as in the paper. I still don't understand how can you claim equivalence, and write explicitly that `PreLN Transformer = Pre-RMSNorm` , if one of the directions applies only when the input is recentered? Moreover, can you explicitly refer to my question in the original review: >Novelty - The authors claim to "propose Pre-RMSNorm". Is this proposal novel? aren't models such as LLaMA already using Pre-RMSNorm?

Rebuttal Response

Thanks for the clarification on weakness 1. I have updated my score. As to the re-parameterization, the main reason for suggesting this term is that I found the proposed method behaves in a similar way to the RepVGG line of work. RepVGG: Making VGG-style ConvNets Great Again