N ORM F ORMER : I MPROVED T RANSFORMER P RETRAINING WITH E XTRA N ORMALIZATION
A BSTRACT
During pretraining, the Pre-LayerNorm transformer suffers from a gradient magnitude mismatch: gradients at early layers are much larger than at later layers. These issues can be alleviated by our proposed NormFormer architecture, which adds three normalization operations to each layer: a Layer Norm after self attention, head-wise scaling of self-attention outputs, and a Layer Norm after the ﬁrst fully connected layer.
The extra operations incur negligible compute cost (+0.4% parameter increase), but improve pretraining perplexity and downstream task performance for both causal and masked language models ranging from 125 Million to 2.7 Billion parameters.
For causal language modeling, adding NormFormer on top of {+our strongest 1.3B parameter baseline achieves equal performance with 79% as much compute, and matches GPT-
 Large performance+} in 62% of the cost.
In the same compute budget, NormFormer comfortably beats both baselines in both pre-training perplexity and zero shot performance.
For masked language modeling, NormFormer improves ﬁne-tuned GLUE performance by 1.9% on average. Code to train NormFormer models is available in REDACTED.
I NTRODUCTION
The original transformer architecture (Vaswani et al., 2017) applies Layer Normalization (Ba et al., 2016) after each sublayer’s residual connection (“Post-LN”) in order to reduce the variance of the inputs to the following sublayer, i.e.:
[Equation]
with
[Equation]
where γ and β are trainable parameters, and (cid:15) is a small constant.
Recent work has shown empirically and theoretically that Post-LN transformers tend to have larger magnitude gradients in later layers compared to earlier layers (Xiong et al., 2020) and has advocated moving the LayerNorm operation to the beginning of each sublayer (“Pre-LN”; see Figure 1, left), i.e.:
[Equation]
Sublayer(LayerNorm( x )) . In practice Pre-LN transformers can be trained with larger learning rates, shorter learning rate warmup and often yield improved performance compared to Post-LN transformers (Xiong et al., 2020), so most recent, large pretrained language models tend to use Pre-LN transformers (Baevski & Auli, 2019; Radford et al., 2019; Raffel et al., 2020; Brown et al., 2020; Lieber et al., 2021).
In this work we show that, while Pre-LN improves stability over Post-LN, it has the opposite side effect: gradients at earlier layers tend to be larger than gradients at later layers, thereby limiting the learning rate.
1 We propose NormFormer , which alleviates the gradient magnitude mismatch by addingnormalization operations to each layer (see Figure 1, middle).
These operations reduce gradients to early layers and increase gradients to later layers, bringing their magnitudes closer together.
Compared to compute-matched, well-tuned Pre-LN baselines, NormFormer models reach target pretraining perplexities faster and achieve better pretraining perplexities and downstream task performance.
The rest of this paper is organized as follows: Section 2 describes the proposed modiﬁcations. Section 3 describes related work. Section 5 shows pretraining and downstream task performance for fully trained NormFormer models against well-tuned, compute-matched baselines.
Sectionshows the gradient mismatch introduced by Pre-LN and how NormFormer alleviates it.
Section 6. analyzes residual scaling, a related technique proposed to stabilize Post-LN architectures (Xiong et al., 2020; Zhu et al., 2021).
Section 7 shows that removing any of the added operations degrades performance and that NormFormer improves over the baseline at a wide range of hyperparameter conﬁgurations. Section 9.1 compares NormFormer to Related Work from other domains.
A PPROACH 2. N ORM F ORMER
NormFormer includes three modiﬁcations to the Pre-LN transformer:
First, we apply head-wise scaling inside the attention module and add two additional LayerNorm operations: one after the attention module and a second after the ﬁrst fully connected layer.
The modiﬁcations introduce a small number of additional learnable parameters, which provide a cost-effective way for each layer to change the magnitude of its features, and therefore the magnitude of the gradients to subsequent components.
The changes are visualized in Figure 1 and described below.
Scaling Attention Heads
The standard multi-head attention operation is deﬁned as:
MultiHeadAttention( Q, K, V ) =
Concat(h 1 , . .
, h n ) W O h [Equation] Attention(
[Equation] W
[Equation]
where n is the number of heads, i is the attention head index, d k is the dimensionality of the keys and W O , W Qi , W Ki , W Vi are learned projection matrices for the output, query, key and value, respectively.
We propose scaling the output of each attention head via learned scalar coefﬁcients γ i : HeadScaleMHA
[Equation]
where γ are learnable parameters initialized to 1.
Additional Layer Normalization and Putting it All Together
In the Pre-LN transformer each layer l modiﬁes an input x l as follows:
[Equation]
where
[Equation]
In this work σ is the GELU non-linear activation introduced in Hendrycks & Gimpel (2016).
Our overall method, NormFormer , instead modiﬁes each input x l as:
[Equation]
where
NormScaledMHA( x ) = x
[Equation]
LN
[Equation]
LN
[Equation]
x )
[Equation]
LN
[Equation]
x )
[Equation]
x )
[Equation]
where bolded operations are newly introduced.
R ELATED W ORK
Architectural Modiﬁcations GradInit (Zhu et al., 2021) introduces a set of scalars and biases for initialization based on a variance heuristic, and Admin (Liu et al., 2020) applies a similar heuristic in proﬁling and initialization stages.
These works also use variants of our ResScale operation, which we ﬁnd helpful at small scale and harmful at large scale.
Our approach, in contrast, only has new learnable parameters without variance heuristics, and has no extra stages or changes in initialization.
Shazeer (2020) proposes FFN-GeGLU, which includes scaling but no normalization, in the same position as our FFN LN.
Ding et al. (2021) propose related stabilization strategies for text to image generation tasks with larger models including a down-scaled embedding gradient, a slightly different LN formulation, LN after the ﬁnal fully connected layer, and the same post-attention LN. Section 9.1 compares NormFormer to these proposals, as well as the T5 LayerNorm Variant (Raffel et al., 2020), which removes the bias and the mean subtraction from the normalization.
Our HeadScale operation is related to that used in Chen et al.
(2021), but
used differently.
Whereas that work prunes attention heads with low γ parameters, we use the γ parameters to improve pretraining performance.
Press et al.
(2020a) proposes an architecture where instead of interleaving attention and feed forward sublayers, the attention all happens ﬁrst. This increases the number of late FFN parameters, rather than increasing their importance and gradient norm, as our FFN LN does, and does not impact stability. Residual Scaling Standard Post-LN transformers simply sum the previous output (residual) with the new output. Recent work attempts to stabilize transformers by weighting the residual connection for each layer (Zhu et al., 2021; Liu et al., 2020; Touvron et al., 2021).
We thus experiment with scaling the residual in each embedding dimension via learned scalar coefﬁcients ( λ resid ) i :
[Equation]
where ◦ is elementwise multiplication, and λ resid are learned parameters initialized to 1. While this can be applied at any normalization layer, we ﬁnd it it most effective for normalizing the feedforward network (FFN) submodule for the smaller sized language models. In this setting, [Equation]
[Table] For 1.3B parameter models and larger, scaling residuals hurts performance (see discussion in Section 6.1), so ResScale is not used in our 1.3B and 2.7B CLM results. Additionally, we experiment with initializing λ resid = 1 e − 5 , following (Touvron et al., 2021), as well replacing addition in residual connections with concatenation (Davis et al., 2021) in Section 9.1.  E XPERIMENTS
Causal Language Models We pretrain causal LMs (CLM) that roughly match the “Small” (125M parameter), “Medium” (355M), “Large” (1.3B) and “XL” (2.7B) sizes from Brown et al.
Our model architecture differs from Brown et al.
(2020) in two ways: (1) we use only dense attention, while they alternate between dense and locally banded sparse attention; (2) we train our models with sinusoidal positional embeddings, following Shortformer (Press et al., 2020b), since early experiments found this to produce comparable results with fewer learned parameters.
We train the baseline models for 300 billion tokens.
We train NormFormer models for an equivalent number of GPU hours, which typically results in 2-6% fewer steps and tokens due to the additional overhead of the normalization operations.
On our dataset, we ﬁnd that the learning rates proposed in GPT-3 are suboptimally low.
For both baseline and NormFormer at each size besides 2.7B, we tune the learning rate by training models for 50,000 steps and selecting the best performing learning rate among: { 1e − 4 , 6e − 4 , 3e − 4 , 6e − 4 , 1e − 3 , 3e − 3 } .
The learning rates we obtained from this process, shown in Table 1, are 3-5 times larger than those used in the GPT-3 paper.
Additionally, we have veriﬁed that the baseline and NormFormer both perform worse at the full training budget with the GPTlearning rates than with the higher learning rates.
Other hyperparameters do not differ from [-GPT-
-]GPT-3.
Large scale experiments We also train three large-scale models with 2.7B parameters.
Our ﬁrst baseline is a replicated version of GPT-3-2.7B with GELU activations, the published learning rate (1.6e-4) and the same number of training steps and tokens (286K steps; 300B tokens). This model slightly exceeds the reference zero shot performance (Brown et al., 2020). Next, we train two variants of GPT3-2.7B with Relu 2 activations (So et al., 2021), but use slightly fewer training steps (20% less) for compute efﬁciency. The ﬁrst of these uses the baseline learning rate (1.6e-4) and the second uses NormFormer-2.7B with a higher learning rate of 6e-4.
We note that training baseline 2.7B CLMs (i.e., without NormFormer modiﬁcations) with a higher 6e-4 learning rate diverged and failed to train.
However, as opposed to the smaller architectures, we did not exhaustively tune the learning rate, so it is possible that an intermediate value would perform better.
Zero Shot Evaluation In addition to validation perplexity, we evaluate CLMs on a subset of the tasks that GPT3 evaluated on in a zero-shot setting (Brown et al., 2020), with the same prompts.
We select WinoGrande (Sakaguchi et al., 2020),
StoryCloze (Mostafazadeh et al., 2016), OpenBookQA (Mihaylov et al., 2018), HellaSwag (Zellers et al., 2019) and PIQA (Bisk et al., 2020) because GPT3 showed strong performance on these tasks at small scale, as well as consistently improving performance with scale. See Table 2.1 in Brown et al.
Masked Language Models (MLM) We adopt the RoBERTa-base, Pre-LN architecture and hyperparameters used in Liu et al.
For the baseline, we pretrain for 2 million batches ofmillion tokens, about 14 of the training budget of the original roberta-base .
NormFormer runs through 1.92 million batches in the same amount of time.
Fine-Tuning We ﬁne-tune both the baseline MLM and NormFormer with learning rates 1e − 5 , 1e − 4 , 3e − 4 , 1e − 3 , 3e − 3 , 6e − 3 and report the best performance on the validation set for each GLUE task (Wang et al., 2019), following Liu et al.
Other ﬁne-tuning hyperparameters match those used for roberta-base in Liu et al.
Pretraining data We pretrain all models on a collection of English language text including the
English portion of the CC100 corpus (Conneau et al., 2020) as well as the data from Liu et al. , consisting of BookCorpus (Zhu et al., 2019), English Wikipedia and ﬁltered subsets of Common Crawl.
We encode our data with the byte-level Byte Pair Encoding (BPE) vocabulary from Liu et al.
(2019), originally introduced in Radford et al.
The combined dataset contains around 450GB of uncompressed text and 110B BPE tokens.
We hold out 40M BPE tokens from this data as a validation set on which we report pretraining perplexities.
Implementation details We train our causal and masked language models in fairseq (Ott et al., 2019; Paszke et al., 2019).
Although NormFormer introduces fewer than 0.07% additional parameters, it slows individual training updates and increases memory usage between 2% (2.7B model) to 6% (125M model) due to the FFN LNs.
Accordingly, we compare NormFormer to baseline models trained for an equal amount of GPU time, i.e., controlling for compute rather than the number of training updates.
Finally, we note that the HeadScale operation can be moved outside the self attention module to allow the use of the very efﬁcient pytorch F.multihead attention .
This change reduces overhead without noticeable performance degradation.
R ESULTS
We report pretraining perplexities for CLMs and MLMs as a function of training wall-time (GPU days) in Figure 2.
We observe that NormFormer trains signiﬁcantly faster and achieves better validation perplexities for a given training compute budget.
The blue stars mark the ﬁrst validation step where NormFormer matches the baseline’s lowest perplexity and shows that NormFormer matches
Pre-LN models while needing only 60% and 57% as much compute for CLM and MLM models, respectively.
This is particularly impressive since NormFormer models take 2-6% longer for each training step and thus see less data than Pre-LN models in this comparison. The left side blue line in Figure 2 shows the failed attempt to add ResScale to NormFormer-1.3B . [Table] [Table]
We observe a similar trend on downstream tasks.
In Table 2 we report zero shot accuracy for causal
LMs using the tasks and prompts from Brown et al.
NormFormer outperforms GPT-3 at all sizes.
The gains from Normformer extra parameters operations outpace the gains from normal scaling laws.
Changing the hidden dimension of a 125M parameter model from 768 to 780, for example, results in a 127 million parameter model that is only 0.08 perplexity better than the baseline whereas NormFormer-125M adds only 100,000 parameters and is 0.83 perplexity better than the baseline.
For MLM models, we report ﬁne-tuned accuracy on GLUE in Table 3.
We again ﬁnd that NormFormer MLM models outperform their Pre-LN counterparts on every task (rows 1 vs 2).
Adding ResScale improves improves pre-training performance marginally (3.29 valid PPL vs 3.31), but the gains to do not translate to ﬁnetuned performance.  A NALYSIS
Analysis of gradient norms by layer We begin by examining the magnitude of the gradients at different layers for Post-LN, Pre-LN and NormFormer models, since large magnitude differences in gradients across layers can destabilize training, particularly when training in mixed precision (Micikevicius et al., 2018).
Figure 3 shows the average L1 norm of the gradients to the second fully connected weight in various layers for a 12 layer, 125M parameter CLM model at the beginning of training.
As reported in past work (Xiong et al., 2020), we observe that the gradients to later layers in Post-LN models are much larger than for earlier layers, and that the gradients to early layers quickly vanish in the early stages of training.
Pre-LN models have the opposite behavior, with early layers instead receiving signiﬁcantly larger gradients than later layers.
NormFormer brings the average gradient norms closer together for different layers in the network.
In Figure 4 we present the distribution of scaling parameters learned by NormFormer models.
For the FFN LN, the γ parameters are smaller for earlier layers, reducing the magnitude of the inputs to early fully connected parameters, thereby decreasing the magnitude of their gradients.
The post attention LN, in the middle of Figure 4, all layers have γ coefﬁcients below 1, indicating downscaling.
4 The HeadScale γ parameters, shown in the rightmost plot in Figure 4 vary more than the others, and have no relationship with depth in the network.
We interpret this as evidence that the HeadScale parameters dynamically increase the importance of well initialized attention heads, as suggested in Chen et al.
Reducing gradient mismatch allows training stably with larger learning rates.
To measure the stability of an architecture, we train it on a learning rate schedule with a very large peak learning rate, so that the learning rate increases a little each step until the loss explodes.
Figure 5 shows that NormFormer models can survive for more updates in this environment than the baseline.
For the baseline 125M model (the left most blue dot), the loss eventually explodes, with the activations from multiplying the query and key features at layer 0 overﬂowing the FP16 range.
The down scaling of the attention outputs allows NormFormer to avoid this issue and remain stable with larger learning rates.
Figure 5 also shows that λ resid reduces the stability improvement at all sizes. 6. R ESIDUAL S CALING By comparing adjacent NormFormer-125M and NormFormer-355M rows in Table 2 we can see that adding ResScale to NormFormer improves perplexity and zero shot performance for small scale CLMs. For 125M parameter MLM, ResScale improves pre-training perplexity marginally, but hurts ﬁne-tuned performance.
At 1.3 billion parameter scale, however, adding ResScale to NormFormer does not improve performance (Figure 2). Although it’s not included in our tables, we ﬁnd that ResScale without NormFormer is stronger than the baseline at small scale, but not large scale. This suggests that the negative result is caused by scale, rather than interaction with
NormFormer . Figure 6 in the appendix shows the average λ resid weights at each layer of different sized CLMs. We can see that at 125M and 355M parameters, the weights in the later layers are lower, indicating down weighting of the residual connection, whereas at the largest scale, 1.3B, the weights are larger deeper into the network.
Adding the λ resid parameters to the other (earlier) residual connection in each layer, or using a scalar instead of a vector for each λ resid , does not ﬁx the large scale issue, but hurts small scale performance marginally. Additionally, Table 8 shows that initializing the network to place no weight on the layer outputs, and all the weight on residual connections, as proposed in (Touvron et al., 2021), does not improve performance.
A BLATIONS
This section provides evidence that removing any of our additions to the transformer block degrades performance on language modeling tasks, and that our additions improve language modeling performance across a wide range of hyperparameter settings.
Experiments use 125M parameter CLMs, and are run with the default hyperparameters given in Table 7 in the appendix for 470 V100 Hours (100,000 updates for the baseline) unless otherwise mentioned.
Removing any of the added operations hurts performance Table 4 shows that none of the four introduced operations can be removed without degrading performance.
Rows 2-5 remove each operation one at a time.
In all cases perplexity increases, with the removal of HeadScale being [Table] the most damaging and the removal of the Post-Attn LN being the least damaging.
In Row 6 ( + More LN )
we try to introduce more normalization inside self attention, applying LN to the query, key and value features in addition to our 3 other operations, for a total of 6 new operations.
In this setting, every other parameterized operation inside the transformer layer is an LN.
We ﬁnd that this does not change perplexities at a ﬁxed number of updates, but reduces training speed by another 5%.
This result suggests that there is not much upside to adding even more normalization on top of
NormFormer .
Other Experiments Table 8 in the appendix compares NormFormer CLMs to related architectural modiﬁcations, both in terms of the stability and pre-training perplexity at 1.3B parameter scale.
Table 5 in the appendix shows language modeling perplexities for 7 different hyperparameter conﬁgurations at 125M parameter scale, separated by horizontal lines.
NormFormer outperforms the baselines in all settings.
C ONCLUSION
We identify a mismatch in the gradients of Pre-LN transformer weights: earlier layers receive much larger gradients than later layers, while the optimal scaling of residuals is larger at earlier layers than at later layers.
We propose NormFormer , which alleviates these issues by adding 3 extra operations to each transformer layer.
These modiﬁcations help the gradient mismatch for fully connected parameters and improve validation perplexity and downstream task performance for both causal and masked language models.
None can be removed without degrading performance back towards the baseline, and adding more normalization – at least of the types we have tried – does not improve performance.
Since NormFormer primarily addresses the gradient mismatch by increasing the gradients to the last FFN layers while decreasing the gradient magnitudes in other parts of the network, future work could examine whether all 3 operations need to be added to every layer.
Additionally, the small computational overhead associated with NormFormer could be alleviated by fusing the
FFN LN with the preceding fully connected layer, with or without the mean centering and bias, which do not appear to improve pretraining perplexity.
In general, we have shown that adding small numbers of learnable parameters in the right places in our architectures can alleviate certain issues in current state of the art networks.
Future work should ascertain if there are additional similarly efﬁcient modiﬁcations that can bring gains, while helping us understand current deﬁciencies further.
[Figure]
[-Wikitext
-][Table]
Wikitext103 Table 6 shows that NormFormer can also provide gains on top of a well tuned language model in settings with much less data.
We simply add our three operations to the architecture and hyperparameters of Baevski & Auli (2019).
Convergence perplexity improves, and we reach the baseline perplexity in 70% as many steps.
In this setting, NormFormer does not improve in the last 30% of training, which suggests that with more tuning the perplexity gap could be widened. Steps to Final PPL PPL
[Table] [Figure]
9. C OMPARISON TO R ELATED W ORK
{+Understanding Table 
+}
The two right most columns of Table 8 contains describes the result of two experiments from the same 1.3B parameter architecture:
Stability indicates how many steps the experiment survived in the ”LR Stability Test”, where we increase LR from 0 to 0.1 linearly over 1,000 steps and Perf , where available, describes the validation perplexity of the model trained for 64 A100 days with sequence length 512.
The table is sorted by Stability, with more stable conﬁgurations lower in the table.
The left columns of the table indicate conﬁguration.
Row 15, which has completely empty conﬁguration columns, is the baseline.
The ﬁnal row is NormFormer .
The Architecture and LN Variant columns show whether we changed the model architecture completely and/or used a non-standard LayerNorm algorithm.
CogView[1]
(Ding et al., 2021): implement all of the changes proposed in Section 2.4: (1) reduce the gradients to the embeddings by a factor of 10, (2) add an LN after attention, (3) add an LN after the second fully connected layer (not the ﬁrst, like NormFormer) (4) change the layer norm formula to LN (
X Max ( X ) ) .
fairseq already uses the attention score stabilization trick by default.
In rows 19 and 25 where Arch is blank, we use the proposed LayerNorm formula but none of the other changes.
CatFormer[2] (Davis et al., 2021):
{+Since no code was released, we re-implement CatFormer and set (cid:15) = 
+}
and =.
448 to match 1.3B parameters.
We guessed that the number of attention heads should be ﬁxed in each layer and that LN positioning should not move, but these details are not clear from the paper.
ID
Arch.
Scale FC
Scale Attn λ Resid
Scale Heads
LN Variant
Stability
PerfPowerNorm[6] -
CatFormer[2] - - - - -
DeepInit[7] - -
DeepInit[7] - - -
FFNGeglu[4] 17.
LayerScale[3] - 17.
CogView[1]
CogView[1] - - - - -
T5[5] - -
FFNGeglu[4] - - - 17.
No Bias or γ - -
No Bias - -
No γ - - - 16.
CogView[1]
CogView[1] 17.
CogView[1] -
CogView[1]
CogView[1] - - - - -
PowerNorm[6] -
LayerScale[3] - - -
No γ - -
CogView[1] - -
No Bias or γ 16.
No Bias - - T5[5] - - - 16. 9.1 for details. LayerScale[3] (Touvron et al., 2021):
We use the layer scale formulation from Section 2, which is similar to the λ resid discussed earlier, but with a weight on each layer’s contribution to the main branch initialized at 1e-5, instead of a weight on the residual’s input to the main branch initialized at 1. 4.
FFNGeglu[4]
(Shazeer, 2020): replace the FFN LayerNorm with ‘FFNGeglu‘.
Although this is proposed as an activation function, it is also equivalent to just using the γ of LayerNorm, with no normalization or bias.
T5[5]
(Raffel et al., 2020): Switch LNs to the T5 variant, which removes the mean centering and bias.
PowerNorm[6] (Shen et al., 2020):
Switch LN to PowerNorm, which is a variant of BatchNorm that shows promising results for smaller NLP models.
DeepInit[7]
(Zhang et al., 2019):
Multiply the initialization of each weight parameter by a factor of 1 √ l , where l is the layer number starting from 1. 8.
No γ : Freeze γ = 1 in LN 9.
No Bias or γ : Freeze γ = 1 , bias=0 in LN 10.
No Bias: Freeze bias=0 in LN.
Results The results suggest that other proposals to mitigate instability in other domains, like vision, text to image generation and reinforcement learning do not improve stability or pre-training performance in the CLM setting.