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: NormFormer improves on Pre-LN transformers by making the following modifications: learnable scaling parameters for each dimension of the output of each attention head prior to concatenation across heads (*Scaled Attention*); layer norm on the attention output (*Post Attn LN*); layer norm on the FFN nonlinearity output (*FFN LN*); and learnable scaling parameters for each dimension of the skip connection around the FFN (*Scaled Residuals*). They apply NormFormer to GPT3- and RoBERTa-style model configurations, and find that NormFormer models reach baseline iso-accuracy in 22%-43% less time, and achieve notably lower (higher) iso-time perplexity (accuracy). They then conduct a number of analyses to attempt to understand why NormFormer works.