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, while the optimal weighting of residuals is larger at earlier than at later layers.
These issues can be alleviated by the addition of two normalization and two new scaling operations inside each layer.
The extra operations incur negligible compute cost (+0.5% parameter increase), but improve pretraining perplexity and downstream task performance for both causal and masked language models of multiple sizes.
Adding NormFormer on top of the GPT3-Medium architecture can reach the SOTA perplexity 22% faster, or converge 0.33 perplexity better in the same compute budget. This results in signiﬁcantly stronger zero shot performance.
For masked language modeling, NormFormer improves ﬁne-tuned GLUE performance by 1.9% on average.
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 observed 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]
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).
Indeed, most recent large pretrained language models 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 also has the opposite side effect: gradients at earlier layers tend to be larger than gradients at later layers.
We propose NormFormer , which alleviates the gradient magnitude mismatch by adding 2 extra normalization operations to each layer (see Figure 1, middle), and 2 extra sets of scaling parameters (Figure 1, right).
These operations reduce gradients to early layers and increase gradients to later layers, bringing their magnitudes closer together, and overall improve the scaling of intermediate hidden states.
Compared to well-tuned Pre-LN baselines, NormFormer models reach target pretraining perplexities faster and achieve better pretraining perplexities and downstream task performance when controlling for compute.
The rest of this paper is organized as follows: Section 2 describes the proposed modiﬁcations, Section 3 shows pretraining and downstream task performance for fully trained NormFormer models against well-tuned baselines.
Section 4 analyses our proposed modiﬁcations, and shows the gradient mismatch introduced by Pre-LN and how NormFormer alleviates it.
Section 5 shows that NormFormer improves over the baseline at a wide range of hyperparameter conﬁgurations, and that removing any of the added operations degrades performance.
A PPROACH 2. N ORM F ORMER
We propose four modiﬁcations to the Pre-LN transformer and call the resulting architecture
NormFormer .
First, we apply head-wise scaling inside the attention module (Sec. 2.1.1). Second, we add another set of scaling weights to learn an improved mixing of the inputs and residuals that are fed into the
LayerNorms (Sec. 2.1.2).
Finally, we add two additional LayerNorm operations: one after the attention module and a second after the ﬁrst fully connected layer (Sec. 2.1.3). The four 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.
2.1. S CALING A TTENTION H EADS
The standard multi-head attention operation is deﬁned as:
[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( Q, K, V ) = Concat( γ 1 h 1 , . . . , γ n h n ) W O where γ is initialized as a vector of 1s.
2.1. S CALING R ESIDUALS
Standard Post-LN transformers simply sum the previous output (residual) with the new output.
As we will see, the optimal mixture of these two inputs varies from layer to layer, with earlier layers typically requiring larger residual weightings.
Moreover, they also vary per dimension of the hidden representation.
We thus propose scaling the residual in each embedding dimension via learned scalar coefﬁcients ( λ resid ) i :
[Equation]
where ◦ is elementwise multiplication, and λ resid are initialized as a vector of 1s.
While this can be applied at any normalization layer, we ﬁnd it it particularly effective for normalizing the feedforward network (FFN) submodule. This can also be interpreted to a degree as alleviating the need for the FFN to produce a pertinent linear map which might be difﬁcult depending on the non-linear activation used, and concentrating instead on nonlinear operations with its two-layer network, where the linear map is learnt using λ resid instead. 2.1. A DDITIONAL L AYER N ORMALIZATION AND P UTTING IT A LL T OGETHER
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
[Equation]
where bolded operations are newly introduced.
This incorporates both additional scaling parameters introduced in the previous two subsections, and the two new normalization layers, which we subsequently refer to as Post Attn LN and FFN LN to differentiate them.
2. E XPERIMENTS
Causal Language Models We pretrain causal LMs (CLM) that roughly match the “Small” (125M parameter) and “Medium” (355M parameter) model sizes introduced in Brown et al.
We show that the conclusions hold for larger models in Section 5.
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 6% fewer steps and tokens due to the additional overhead of the normalization operations. [Table] We ﬁnd that, on our dataset, the learning rates proposed in GPT-3 are suboptimally low.
For both baseline and NormFormer at each size, we selected the best learning rate from 6 50,000 step runs at 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. This may also explain our baseline improving over GPT-3 zero shot results in Table 2.
Additionally, we have veriﬁed that the baseline and NormFormer {+both perform worse at the full training budget with the GPT-3 learning rates than with the higher learning rates. Other hyperparameters do not differ from GPT-
+}
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.
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 192 batches in the same amount of time.
Fine-Tuning We ﬁne-tune both the baseline MLM and NormFormer with learning rates 1e − 5 , 1e − 3e − 1e − 3 , 3e − 3 , 6e − 3 and report the best performance on the validation set for
See Table 2. Brown et al. [Table]
[Table] 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 introduced in
Liu et al. (2019), consisting of the BookCorpus (Zhu et al., 2019), Wikipedia and Common Crawl, as well as the English portion of the CC100 corpus (Conneau et al., 2020).
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). Although NormFormer introduces fewer than 0.07% additional parameters, it slows individual training updates and increases memory usage by 2-6%, primarily 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.
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 green 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 (see Section 2.2) and thus see less data than Pre-LN models in this comparison. 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 both the 125M parameter and 355M parameter 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 .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.  A NALYSIS 4. A NALYSIS 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.
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.
Finally, introducing an additional LN between fully connected layers brings the 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 pattern for the post attention LN, in the middle of Figure 4, is only visible at layer 0, but all layers have coefﬁcients below 1, indicating downscaling.
3 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.
One result of reducing the gradient mismatch, besides better perplexities and downstream task performance, is the ability to train stably with larger learning rates.
To measure the stability of an [Table] architecture, we train it on a learning rate schedule with a very large peak learning rate and very long warmup, so that the learning rate increases a little each step until the loss explodes.
Tableshows that NormFormer models are stable at a higher peak learning rate than the baseline.
For the baseline, loss explodes because the activations produced by multiplying the query and key features at layer 0 are too large to be represented in 16 bits.
The down scaling of the attention outputs allows NormFormer to avoid this issue and remain stable with larger learning rates.
After training with the ResScale approach, one can also view the vectors that are learnt. Taking the minimum value over λ resid for each layer we see the values are the smallest for the last layer (around 0.8) while being close to 1 for the ﬁrst layer. According to an experiment at 125M parameters, all the beneﬁt seems to result from the scaling attached to the ﬁnal residual connection of each layer for the FFN. Adding scaling parameters for the MHA residual actually increased PPL slightly by .03.
We thus only apply this approach to the FFN.
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 and model sizes.
Experiments use CLMs, and are run with the default hyperparameters given in Table 8 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 5 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 the most damaging.
In Row 6 ( + 3 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 non LN operation inside the transformer layer is followed by an LN.
We ﬁnd that this does not change perplexities at a ﬁxed number of updates, but
[Table]
reduces training speed by another 5%.
This result suggests that there is not much upside to adding even more normalization after our additions.
We see similar improvement from the addition of extra normalization as we scale model size. Other Hyperparameter Settings Table 6 shows language modeling perplexities for 7 different hyperparameter conﬁgurations, separated by horizontal lines.
NormFormer outperforms the baseline in all settings. Other Experiments Replacing the FFN LN with the FFNGeGlu proposed in Shazeer (2020), which includes scaling but no normalization, degraded performance in our setting.
We also ﬁnd that the LN variant proposed in Raffel et al.
(2020), which removes the bias and the mean substraction from the normalization
[Table]
performs equally well to our LN and has fewer trainable parameters
[Table]
but is about 2x slower than the FusedLayerNorm implementation we use.
We therefore do not adopt it.
[Table]
R ELATED W ORK
Layer normalization (Ba et al., 2016) is an important component of the transformer architecture.
Xiong et al.
(2020) shows that for Post-LN: gradients are too big for later layers and solves this problem with Pre-LN.
We build on the Pre-LN architecture to make it even more stable and efﬁcient.
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.
We note that there are other solutions from other domains for enforcing stability, such as in the ﬁeld of reinforcement learning (Parisotto et al., 2020).
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.
Our ResScale operation is related to Highway Networks (Srivastava et al., 2015), which advocate nonlinear mappings of the residual, while the original ResNet paper, with experiments on convolutional nets, also suggests a full linear map is possible (He et al., 2016). We found linear maps were unstable in our transformers, while elementwise scaling worked well.
These approaches are also related to techniques for initializing neural networks: 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. Similarly, some other approaches targeted initialization as well, in particular ReZero (Bachlechner et al., 2020),
FixUp (Huang et al., 2020) and LookLinear (Balduzzi et al., 2017). We note that DALL-E (Ramesh et al., 2021) also added a per residual scaling factor (only during backprop). Our approach, in contrast, only has new learnable parameters without variance heuristics, and has no extra stages or changes in initialization.
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 than at early layers.
We propose NormFormer , which alleviates these issues by adding 4 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 – 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 4 operations need to be added to every layer. Additionally, the small computational increase 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, whilst helping us understand current deﬁciencies further.
A PPENDIX
{+Wikitext
+}
Table 7 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 four 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.
[Table] [Table]