SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zixuan Jiang, Jiaqi Gu, Hanqing Zhu, David Z. Pan
Chandra Department of Electrical and Computer Engineering *
Austin, Texas, 78712
{zixuan, jqgu, hqzhu}@utexas.edu, dpan@ece.utexas.edu
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\%$ .
Transformers have become a successful architecture for a wide range of machine learning applications, including natural language [39], computer vision [13], and reinforcement learning [7]. It is one of the foundation models [2], and pretrained Transformers [29] demonstrated impressive generalization results. Among the components of Transformers, normalization plays a critical role in accelerating and stabilizing the training process [21]. Layer Normalization (LayerNorm, LN) [1] and Root Mean Square Normalization (RMSNorm) [44] are two common normalization layers in Transformers. LayerNorm is in the original Transformer architecture [39], recentering and rescaling the input vector in  $\mathbb{R}^d$  to obtain a zero-mean and unit-variance output. RMSNorm only rescales the input vector with its RMS value, offering greater computational efficiency than LayerNorm.
*This work was done when all the authors were at UT Austin. Currently, Zixuan Jiang is at Google, and Jiaqi Gu is at Arizona State University.
The machine learning community does not reach a consensus regarding the preferred normalization technique for Transformers. LayerNorm demonstrates remarkable success in the milestone Transformers, such as GPT [30, 5] and ViT [13]. It is still the default normalization layer when building a new Transformer. On the contrary, RMSNorm is reported to accelerate the training and inference with similar performance as LayerNorm in Transformers. It has gained popularity in recent large language models, such as T5 [32], Gopher [31], Chinchilla [16], and LLaMA [38]. However, concerns persist regarding the potential negative impact of RMSNorm on the representation ability of Transformers. It remains an open question to determine the preferred normalization type for Transformers, requiring further theoretical and empirical investigation.
In this work, we aim to mitigate the discrepancy between LayerNorm and RMSNorm in Transformers. When delving into the prevalent Transformer architectures, Pre-LN and Pre-RMSNorm Transformers, we identify an opportunity to unify them by removing the inherent redundancy in Pre-LN models. In particular, the main branch vectors in the Pre-LN Transformers are always normalized before they are used, implying that the mean information is redundant. We can recenter the main branch without impact on the functionality of the models, which allows us to reduce LayerNorm to RMSNorm.
We further propose Compressed RMSNorm (CRMSNorm), which takes a vector in  $\mathbb{R}^{d - 1}$  as input, decompresses it to a zero-mean vector in  $\mathbb{R}^d$ , and applies the RMSNorm on the decompressed vector. Building upon this new normalization, we propose Pre-CRMSNorm Transformer that employs lossless compression on the zero-mean vectors. We apply such compression to zero-mean activations and parameters in Pre-RMSNorm Transformers, enhancing the efficiency of Pre-RMSNorm Transformers while preserving the equivalent arithmetic functionality.
We formally claim that Pre-LN, Pre-RMSNorm, and Pre-CRMSNorm Transformers are equivalent for both training and inference. Figure 1 visualizes the overview of their equivalence. Such equivalence directly enables more efficient training and deployment of Pre-LN Transformers. We can translate a Pre-LN model into an equivalent
Figure 1: Overview of the three equivalent Transformer variants.
Pre-(C)RMSNorm model, which can be readily deployed or adapted to downstream tasks. The conversion process incurs minimal costs. We can also train a Pre-(C)RMSNorm model directly as if we train an equivalent Pre-LN Transformer counterpart.
Although our relative improvement in the training and inference efficiency seems not large (up to  $10\%$  time reduction), we have the following arguments to support its significance. (1) Our proposed method can guarantee arithmetic equivalence and is a free lunch for Pre-LN Transformers. The efficiency improvement originates from removing the inherent redundancy in Pre-LN Transformers without introducing any fine-tuning or calibration. Our work can strictly push the performance-efficiency Pareto frontier of Pre-LN Transformers. (2) The modest relative improvement can be translated into significant absolute improvement given that Pre-LN Transformers are foundation models for the current and future data-centric [42] and generative [6] artificial intelligence. For instance, reducing the ChatGPT inference cost by  $1\%$  may save  $\$7,000$  per day [25]. (3) Our method is orthogonal and complementary to most work improving efficiency, such as efficient Transformer variants [35], quantization [22, 3], and distillation to smaller models [34].
We highlight our contributions as follows. Our code is available at https://github.com/ZixuanJiang/pre-rmsnorm-transformer.
- We achieve the first-ever unification of LayerNorm and RMSNorm in pre-normalization Transformers with proven arithmetic equivalence.
- We propose two variants: Pre-RMSNorm and Pre-CRMSNorm Transformers. The original Pre-LN Transformer and our proposed two variants are equivalent and can seamlessly interchange without affecting functionality.
- Our proposed architectures are  $1\% - 10\%$  more efficient than the original Pre-LN Transformer for both training and inference. Such efficiency gains are effortlessly obtained without the need for fine-tuning or calibration.
We introduce LayerNorm, RMSNorm, and their usage in Transformers. We provide an abstraction for Pre-LN Transformers.
2.1 LayerNorm and RMSNorm
Layer Normalization (LayerNorm, LN) [1] is a technique to normalize the activations of intermediate layers of neural networks. Given a vector  $\pmb{x} \in \mathbb{R}^d$ , LayerNorm normalizes it to obtain a zero-mean unit-variance vector,
$$
\operatorname {L a y e r N o r m} (\boldsymbol {x}) = \frac {\boldsymbol {x} - \mu (\boldsymbol {x}) \mathbf {1}}{\sqrt {\left\| \boldsymbol {x} \right\| _ {2} ^ {2} / d - \mu^ {2} (\boldsymbol {x}) + \epsilon}}, \text {w h e r e} \mu (\boldsymbol {x}) = \frac {\mathbf {1} ^ {T} \boldsymbol {x}}{d}, \epsilon > 0. \tag {1}
$$
LayerNorm recenters and rescales the activations and gradients in the forward and backward computations [41], which enables fast and robust training of neural networks.
Root Mean Square Normalization (RMSNorm) [44] is another technique used for normalizing the activations. It is similar to LayerNorm in that it aims to accelerate and stabilize the training but uses a different normalization approach. Instead of normalizing the inputs based on their mean and variance, RMSNorm normalizes them based on their root mean square (RMS) value. It is defined in the following equation,
$$
\operatorname {R M S N o r m} (\boldsymbol {x}) = \frac {\boldsymbol {x}}{\sqrt {\left\| \boldsymbol {x} \right\| _ {2} ^ {2} / d + \epsilon}}, \text {w h e r e} \epsilon > 0. \tag {2}
$$
RMSNorm only rescales the input vector and the corresponding gradients, discarding the recentering process. As shown in their definitions, RMSNorm is computationally simpler and more efficient than LayerNorm. It is reported that replacing LayerNorm with RMSNorm can achieve comparable performance and save training and inference time by  $7\% - 64\%$  [44].
Given a zero-mean vector  $\pmb{x}$ , these two kinds of normalization are equivalent. Formally, if  $\mu(\pmb{x}) = 0$  then  $\mathrm{LayerNorm}(\pmb{x}) = \mathrm{RMSNorm}(\pmb{x})$ . We may optionally introduce learnable parameters and apply an element-wise affine transformation on the output of LayerNorm and RMSNorm.
We focus on the computation of LayerNorm and RMSNorm instead of their optimization and expressivity in this paper.
2.2 Normalization in Transformers
Normalization plays a crucial role and has many variants in Transformers [39]. LayerNorm is widely used in Transformer architectures to address this issue. The position of LN within the architecture is essential for the final performance. While the initial Transformer uses Post-LN, most Transformers employ Pre-LN to achieve more stable training, even though this can result in decreased performance [40]. Pre-LN is the mainstream normalization in Transformers, especially the large models, such as ViT [13, 9], PaLM [8], and GPT-series models [30, 5].
RMSNorm is proposed as an alternative normalization technique in Transformers. Several large language models, such as Chinchilla [16] and LLaMA [38], use Pre-RMSNorm in their blocks [45]. RMSNorm can help accelerate the training and inference with similar performance in these large models. Specifically, the experiments in [26] show that RMSNorm improves the pre-training speed by  $5\%$  compared with the LayerNorm baseline.
2The following discussion on expressivity is out of the scope of this paper and is only for reference. Normalization helps accelerate and stabilize the training but hurts the model's expressivity. For example, we usually add learnable parameters  $\gamma ,\beta$  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, so it has less information loss.
Figure 2: Left. The original Pre-LN Transformer architecture. Middle and Right. Our proposed Pre-RMSNorm and Pre-CRMSNorm Transformer architectures. These three architectures are equivalent. The differences are highlighted in bold and green blocks.
It is challenging to convert Transformers with one normalization to the other type. It is not clear which version of normalization is more suitable for Transformers.
2.3 Pre-LN Transformer
Figure 2.a illustrates the Pre-LN Transformer architecture, which consists of three parts, preprocessing, a stack of  $L$  blocks, and postprocessing.
Preprocessing. We preprocess the raw inputs, which range from paragraphs in natural language [39], images [13], to state-action-reward trajectories in reinforcement learning problems[7]. We import special tokens (such as the classification token) and embeddings (such as positional embeddings), ultimately obtaining a sequence of token embeddings  $\boldsymbol{x}_0$ .
Transformer blocks. The main body of Pre-LN Transformer [40] is a stack of residual blocks
$$
\boldsymbol {x} _ {l + 1} = \boldsymbol {x} _ {l} + \mathcal {F} _ {l} (\boldsymbol {x} _ {l}), l = 0, 1, \dots , L - 1, \tag {3}
$$
where  $\pmb{x}_l$  is the input of the  $l$ -th block,  $\mathcal{F}_l$  is a sequence of operators  $LayerNorm \to Linear \to g_l \to Linear$ . We name  $\pmb{x}_l$  as the vectors on the main branch and  $\mathcal{F}_l$  as the residual branch [14]. The block  $\mathcal{F}_l$  is usually an attention or a multi-layer perceptron (MLP) module. If  $g_l$  is an activation function, such as GELU [15], then the block  $\mathcal{F}_l$  is a two-layer MLP. If  $g_l$  is a (masked) multi-head scaled dot product attention, then the block is the (casual) attention module [39]. These two linear layers are usually explicitly defined. Taking the attention module as an example, the input linear projection generates the query, key, and value vectors, while the output projection is applied to the concatenated results of all heads. If they are not explicitly defined, we can add an identity mapping, a special linear layer. Most of the learnable parameters of Transformers are in these two linear layers.
LayerNorm and postprocessing. We finally process the result  $\mathrm{LN}(\pmb{x}_L)$  to obtain the task-related results, such as classification probabilities. We apply the layer normalization on  $\pmb{x}_L$  since it usually has a high variance because it is the accumulation of all the Transformer blocks.
We propose Pre-RMSNorm and Pre-CRMSNorm Transformer variants, shown in Figure 2, and claim that Pre-LN, Pre-RMSNorm, and Pre-CRMSNorm Transformers are arithmetically equivalent.
$$
\text {P r e - L N T r a n s f o r m e r} = \text {P r e - R M S N o r m T r a n s f o r m e r} = \text {P r e - C R M S N o r m T r a n s f o r m e r .} \tag {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.
3The elementwise affine transformation of LayerNorm or RMSNorm is also a linear transformation. 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-RMSNorm Transformer and the pre-existing Pre-RMSNorm models (e.g., LLaMA [38]) are different. Ours has a recentering and special output linear layer, as shown in Figure 2b. However, they can share the same implementation for inference. We adopt the term "Pre-RMSNorm Transformer" to represent our proposed one in the discussions below.
3.1 Pre-LN Transformer  $=$  Pre-RMSNorm Transformer
LayerNorm is invariant to the shifting  $\mathrm{LN}(\pmb{x} + k\mathbf{1}) = \mathrm{LN}(\pmb{x}), \forall k \in \mathbb{R}$ . We observe that LayerNorm is applied to the main branch vectors before they are used in either residual branches or the final postprocessing in Pre-LN Transformers. Therefore, we can replace the main branch vectors  $\pmb{x}_l$  with  $\pmb{x}_l + k_l\mathbf{1}, \forall k_l \in \mathbb{R}$  without impact on the functionality of the Pre-LN Transformer. If  $k_l = -\mu(\pmb{x}_l)$ , we replace the main branch vectors with its recentered version  $\pmb{x}_l - \mu(\pmb{x}_l)\mathbf{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  $x_0$  before the first Transformer block, where  $\operatorname{Recenter}(\pmb{x}) = \pmb{x} - \mu(\pmb{x})\mathbf{1}$ .
2. For the output projection in residual branches, replace the weight  $\mathbf{A}_o$  and bias  $\mathbf{b}_o$  with  $\hat{\mathbf{A}}_o = \mathbf{A}_o - \frac{1}{d}\mathbf{1}\mathbf{1}^T\mathbf{A}_o$ ,  $\hat{\mathbf{b}}_o = \mathbf{b}_o - \mu (\mathbf{b}_o)\mathbf{1}$ , where  $d$  is the dimension of  $\mathbf{x}_0$ .
3. Replace LayerNorm with RMSNorm at the beginning of residual blocks and before postprocessing.
Since  $\mu(\boldsymbol{x}_{l+1}) = \mu(\boldsymbol{x}_l) + \mu(\mathcal{F}_l(\boldsymbol{x}_l))$ , we can keep zero-mean on the main branch if and only if the input of the first block  $\boldsymbol{x}_0$  and the output of each residual branch  $\mathcal{F}_l$  are re-centered with zero-mean. The first modification is to recenter  $\boldsymbol{x}_0$ , while the second modification is to recenter the output of residual branches. For the residual branch  $\mathcal{F}_l$ , the ending linear transformation enables us to recenter its output without extra computation, implied by Lemma 3.1. We can recenter the weight and bias of a linear layer to recenter its output.
Lemma 3.1 Given a linear transformation  $\pmb{y} = \pmb{A}\pmb{x} + \pmb{b},\pmb{x}\in \mathbb{R}^n,\pmb{A}\in \mathbb{R}^{m\times n},\pmb{b},\pmb{y}\in \mathbb{R}^m$  , we can decompose the output with two parts  $\pmb {y} = \hat{\pmb{A}}\pmb {x} + \hat{\pmb{b}} +\mu (\pmb {y})\mathbf{1}$  . The first part  $\hat{\pmb{A}}\pmb {x} + \hat{\pmb{b}} = \pmb {y} - \mu (\pmb {y})\mathbf{1}$  with zero mean, is another linear transformation with  $\hat{\pmb{A}} = \pmb {A} - \frac{1}{m}\pmb{1}\pmb{1}^T\pmb {A},\hat{\pmb{b}} = \pmb {b} - \mu (\pmb {b})\mathbf{1}$
The first two modifications ensure that we maintain zero mean on the main branch vectors. Given a zero-mean input, LayerNorm is equivalent to RMSNorm, which implies that the third modification has no impact on the functionality. With these modifications, we demonstrate that Pre-LN and Pre-RMSNorm Transformers are equivalent.
3.2 Pre-RMSNorm Transformer = Pre-CRMSNorm Transformer
For a zero-mean vector  $\pmb{x} \in \mathbb{R}^d$ , we can compress it losslessly by discarding its last element. In the decompression, we recover the discarded element with  $x_d = -\sum_{i=0}^{d-1} x_i$ . The decompression has an extra cost, while the compression does not induce extra computation. The space-saving ratio of this compression method is  $1/d$ .
We define Compressed Root Mean Square Normalization (CRMSNorm), which takes a vector  $\pmb{x} \in \mathbb{R}^{d-1}$  as input. CRMSNorm first decompresses the vector  $\pmb{x}$  to obtain a zero-mean vector in  $\mathbb{R}^d$ , then applies RMSNorm on the zero-mean vector. It can generate the normalized zero-mean results in either  $\mathbb{R}^{d-1}$  or  $\mathbb{R}^d$ . Its formal definition is in Equation 5.
$$
\operatorname {C R M S N o r m} (\boldsymbol {x}) = \frac {\boldsymbol {x}}{\sqrt {\left(\sum_ {i = 1} ^ {d - 1} x _ {i} ^ {2} + \left(\sum_ {i = 1} ^ {d - 1} x _ {i}\right) ^ {2}\right) / d + \epsilon}}, \text {w h e r e} \boldsymbol {x} \in \mathbb {R} ^ {d - 1}. \tag {5}
$$
We simplify the Pre-RMSNorm Transformers to obtain the Pre-CRMSNorm Transformers with the following modifications.
1. Compress the zero-mean main-branch vectors from  $\mathbb{R}^d$  to  $\mathbb{R}^{d - 1}$ . Preprocessing and postprocessing handle compressed vectors in  $\mathbb{R}^{d - 1}$ .
2. Replace RMSNorm with CRMSNorm.
3. Simplify the weight  $\mathbf{A}_i$  in the input projection layer. Let  $\mathbf{a}_d$  be the last column vector of  $\mathbf{A}_i$ .  $\hat{\mathbf{A}}_i = \mathbf{A}_i - \mathbf{a}_d\mathbf{1}^T$  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  $\hat{A}_o$  and the last element of the bias  $\hat{b}_o$  since they are used to generate the redundant last element.
We compress the zero-mean activations in our proposed Pre-RMSNorm Transformers and correspondingly simplify the two linear layers in residual branches.
We can fuse preprocessing, recentering, and compression. The fused preprocessing generates compressed vectors in  $\mathbb{R}^{d - 1}$  to represent zero-mean vectors in  $\mathbb{R}^d$ . Taking language models as an example, we can recenter and compress (discard the last element) the word and position embeddings.
For the input linear projection, its input is the output of RMSNorm, whose mean is zero. We compress the output of the RMSNorm and the weight of the linear layer. Specifically, if the linear layer takes a zero-mean vector as input, then
$$
\operatorname {L i n e a r} (\boldsymbol {x}) = \boldsymbol {A} _ {i} \boldsymbol {x} + \boldsymbol {b} _ {i} = \left(\boldsymbol {A} _ {i} - \boldsymbol {a} _ {d} \mathbf {1} ^ {T}\right) \boldsymbol {x} + \boldsymbol {b}. \tag {6}
$$
$\hat{\pmb{A}}_i = \pmb {A}_i - \pmb {a}_d\pmb{1}^T$  is the compressed weight matrix since its last column is a zero vector. The simplified linear layer only needs the compressed zero-mean vectors in  $\mathbb{R}^{d - 1}$  as input.
For the output linear projection, we only need to calculate the first  $(d - 1)$  elements for the vectors in  $\mathbb{R}^d$ . Thus, we can compress the weight  $\hat{A}_o$  and bias  $\hat{b_o}$ . As shown in Figure 2, the shape of  $\hat{A}_o$  is  $(d,d_o)$ , we can compress it directly to  $(d - 1,d_o)$  by discarding its last row. Similarly, we can compress  $\hat{b_o} \in \mathbb{R}^d$  by discarding its last element. This weight compression also reflects their redundancy. As shown in Section 3.1,  $\hat{b_o}$  and all column vectors of  $\hat{A}_o$  are zero-mean.
Figure 2 demonstrates the difference in vector dimension. We demonstrate the equivalence between Pre-RMSNorm and Pre-CRMSNorm Transformers.
3.3 Pre-LN Transformer = Pre-CRMSNorm Transformer
To translate a pre-trained Pre-LN Transformer to a Pre-CRMSNorm Transformer, we can convert it into a Pre-RMSNorm model and then finish the conversion. For the model implementation, we need two steps to convert a Pre-LN Transformer to Pre-CRMSNorm Transformer.
1. Reduce the hidden dimension from  $d$  to  $d - 1$ .
2. Replace LayerNorm with CRMSNorm.
Namely, Pre-LN Transformers with the main branch vectors in  $\mathbb{R}^d$  are equivalent to Pre-CRMSNorm Transformers in  $\mathbb{R}^{d - 1}$ , which further echoes the redundancy in the Pre-LN Transformers.
3.4 Training and Inference Efficiency
We show how to make conversions between the three variants. The conversions only consist of one-time parameter adjustments without expensive fine-tuning or calibration, similar to operator fusion [27]. We qualitatively analyze the efficiency improvement of our proposed Pre-(C)RMSNorm Transformers compared with equivalent Pre-LN models.
3.4.1 Pre-RMSNorm Transformer
We discuss the impact of three modifications on training and inference, as shown in Table 1.
The linear layer with zero-mean output. The modified linear layer will not induce extra computation for inference since we can replace the weight and bias in advance. During inference, we can treat it as a standard linear layer with equivalently transformed parameters.
<sup>4</sup>We reduce the hidden dimension in the preprocessing, main branch, two linear layers in residual branches, and postprocessing. We keep the  $d_{i}, d_{o}$  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	a little increase	a little increase	decrease
Inference	same	same
Table 1: The computation workload of our Pre-RMSNorm model compared with the original Pre-LN Transformer.