NormFormer: Improved Transformer Pretraining with Extra Normalization

1 INTRODUCTION . 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 . : PostLN ( x ) = LayerNorm ( x+ Sublayer ( x ) ) , with LayerNorm ( x ) = x− E [ x ] √ V ar [ x ] + · γ + β , where γ and β are trainable parameters , and 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 . : PreLN ( x ) = x+ 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 adding 3 normalization 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 . 1Intuitively , training stably requires that the largest weight update not be too large , while training efficiently requires large weight updates . 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 modifications . Section 3 describes related work . Section 5 shows pretraining and downstream task performance for fully trained NormFormer models against well-tuned , compute-matched baselines . Section 6 shows the gradient mismatch introduced by Pre-LN and how NormFormer alleviates it . Section 6.1 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 configurations . Section 9.1 compares NormFormer to Related Work from other domains . 2 APPROACH . 2.1 NORMFORMER . NormFormer includes three modifications 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 first fully connected layer . The modifications 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 defined as : MultiHeadAttention ( Q , K , V ) = Concat ( h1 , . . . , hn ) W O hi = Attention ( QW Q i , KW K i , V W V i ) Attention ( Q , K , V ) = softmax ( QKT√ dk ) V , where n is the number of heads , i is the attention head index , dk is the dimensionality of the keys and WO , WQi , W K i , W V i are learned projection matrices for the output , query , key and value , respectively . We propose scaling the output of each attention head via learned scalar coefficients γi : HeadScaleMHA ( Q , K , V ) = Concat ( γ1h1 , . . . , γnhn ) W O where γ are learnable parameters initialized to 1 . Additional Layer Normalization and Putting it All Together In the Pre-LN transformer each layer l modifies an input xl as follows : xPreLNl+1 = FFN ( MHA ( xl ) ) where MHA ( x ) = x+MultiHeadAttention ( LN ( x ) , LN ( x ) , LN ( x ) ) FFN ( x ) = x+ σ ( LN ( x ) W1 + b1 ) W2 + b2 LN ( x ) = LayerNorm ( x ) In this work σ is the GELU non-linear activation introduced in Hendrycks & Gimpel ( 2016 ) . Our overall method , NormFormer , instead modifies each input xl as : xNormFormerl+1 = NormFFN ( NormScaledMHA ( xl ) ) where NormScaledMHA ( x ) = x+ LN ( HeadScaleMHA ( LN ( x ) , LN ( x ) , LN ( x ) ) ) NormFFN ( x ) = x+ LN ( σ ( LN ( x ) W1 + b1 ) ) W2 + b2 where bolded operations are newly introduced . 3 RELATED WORK . Architectural Modifications 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 profiling and initialization stages . These works also use variants of our ResScale operation , which we find 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 final 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 first . 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 coefficients ( λresid ) i : ResScale ( x ) = λresid ◦ x+ Sublayer ( LayerNorm ( x ) ) where ◦ is elementwise multiplication , and λresid are learned parameters initialized to 1 . While this can be applied at any normalization layer , we find it it most effective for normalizing the feedforward network ( FFN ) submodule for the smaller sized language models . In this setting , NormFFN ( x ) = λresid ◦ x+ LN ( σ ( LN ( x ) W1 + b1 ) ) W2 + b2 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 = 1e − 5 , following ( Touvron et al. , 2021 ) , as well replacing addition in residual connections with concatenation ( Davis et al. , 2021 ) in Section 9.1 . 4 EXPERIMENTS . 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 . ( 2020 ) . 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 find that the learning rates proposed in GPT-3 are suboptimally low.2 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 verified 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-3.3 Large scale experiments We also train three large-scale models with 2.7B parameters . Our first 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 Relu2 activations ( So et al. , 2021 ) , but use slightly fewer training steps ( 20 % less ) for compute efficiency . The first of these uses the baseline learning rate ( 1.6e-4 ) and the second uses NormFormer-2.7Bwith a higher learning rate of 6e-4 . We note that training baseline 2.7B CLMs ( i.e. , without NormFormer modifications ) 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 . 2The difference in optimal learning rates may be due partly to architectural differences between our baseline and GPT-3 ( e.g. , not using locally banded sparse attention ) . 3See Table 2.1 in Brown et al . ( 2020 ) . Masked Language Models ( MLM ) We adopt the RoBERTa-base , Pre-LN architecture and hyperparameters used in Liu et al . ( 2019 ) . For the baseline , we pretrain for 2 million batches of 1 million 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 fine-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 . ( 2019 ) . Other fine-tuning hyperparameters match those used for roberta-base in Liu et al . ( 2019 ) . 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 . ( 2019 ) , consisting of BookCorpus ( Zhu et al. , 2019 ) , English Wikipedia and filtered 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 . ( 2019 ) . 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 efficient pytorch F.multihead attention . This change reduces overhead without noticeable performance degradation .

Summary: This paper aims to improve pretraining Pre-LayerNorm transformers by alleviating two issues: early layers have much larger gradients than later ones, and naive residual learning can't provide optimal weighting. To this end, it proposes to add two LayerNorms after the multi-head attention and the GELU non-linear activation in FFN, respectively. It also adds learnable scaling coefficients for the FFN residual and the attention head outputs. The four modifications are applied to both casual and masked language modeling with improvements observed in downstream tasks.