NormFormer: Improved Transformer Pretraining with Extra Normalization
Authors: Sam Shleifer
Myle Ott
Venue: ICLR 2022 Submitted (2022)
Keywords: Language Modeling
NLP
Transformer
Zero Shot Learning

Abstract: 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 significantly stronger zero shot performance.
For masked language modeling, NormFormer improves fine-tuned GLUE performance by 1.9\% on average.

Under review as a conference paper at ICLR 2022
NORMFORMER: IMPROVED TRANSFORMER PRETRAINING WITH EXTRA NORMALIZATION
ABSTRACT
During pretraining, the Pre-LayerNorm transformer suffers from gradient mag- nitude mismatch: gradients early layers are much larger than at later layers. These issues be alleviated by proposed NormFormer architecture, which adds three normalization operations to each layer: a Layer Norm after self atten- tion, head-wise scaling of self-attten celf-attention outputs, and Layer Norm after the first fully connected layer. The extra operations incur negligible compute cost (+0.4% parameter increase), but improve pretraining perplexity and downstream task per- formance for both causal and masked language models ranging from 125 Million 2.7 Billion parameters. For causal language modeling, adding NormFormer on top of our strongest 1.3B parameter baseline achieves equal performance with 79% much compute, and matches GPT-3 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 fine-tuned GLUE performance by 1. 1.9% on av- erage. Code to train NormFormer models is available in REDACTED.
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(r Sublayer(x)).
with
E
LayerNorm(z) [arlaz)+ + 7 + 3,
where "y and B are trainable parameters, and e a small constant. Recent work has shown empir- ically and theoretically that Post-LN transformers tend to have larger magnitude gradients in later layers compared to carlier layers (Xiong et al., 2020) and has advocated moving the LayerNorm operation the beginning of each sublayer ("Pre-LN":; see Figure 1, left), i.e.:
PreLN(x)
Sublayer(LayerNoo(ee)
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,! We propose No NorFFo o rmer, r, which alleviates the gradient magnitude mismatch by adding normalization operations to each layer (see Figure 1, middle). These operations reduce gradients to carly layers and increase gradients to later layers, bringing their magnitudes closer together.
Intuitively, training stably requires that the largest weight update not be too large, while training efficiently requires large weight updates.

---

Under review as a conference paper at ICLR 2022
Residual Connection"
Residual Connection
FG2
HeadScale (:3g)
FC2
FFN LayerNorm
MatMul
FC1
LayerNorm
LayerNorm
SoftMax
Residual Connection
Connection
MatMull
Post Attn Layer/oorm
Qn:
Ka,
Vhu
Multi-Heed
HeadScale
Linear
Linear
Linear
Attention
MHA
Scalled Attention Head
LayerNorm
LayerNorm
Pre-LN Layer
NormFormer
Figure I: Left: a baseline Pre-LayerNorm transformer layer. Center: NormE ormer, with the three proposed additions in bold. Right: single attention head with our proposed HeadScal e operation applied prior to the output projection with trainable parameters "/:- When applied, residual scaling impacts the second residual connection in each layer.
Compared to compute-matee matphed, well-tuned Pre-LN baselines, NormForme models reach target pretraining perplexities faster and achieve better pretraining perplexities and downstream task per- formance.
The rest of this paper is organized as follows: Section describes the proposed modifications. Section 3 describes related work. Section shows pretraining and downstream task performance for fully trained NormF ormer models against well-tuned, compute-matched baselines. Section shows the gradient mismatch introduced by Pre-LN and how NormF omer alleviates it. Section 6. analyzes residual scaling, a related technique proposed to stabilize Post-LN architectures (Xiong al., 2020; Zhu et al., 2021). Section shows that removing any the added operations degrades performance and that No rmer improves over the baseline at a wide range of hyperparameter configurations. Section 9.1 compares NormForme or Related Work from other domains.
APPROACH
2.1 NORMFORMER
FormFormer 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 small number of additional learnable parameters, which provide 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 Figure and described below.
Scaling Attention Heads The standard multi-head attention operation is defined as:
MultiHeadAttennionn( Q, K, V) Concat (hi...., W
h, Attention(QW, KW.K, vW."
QKT
Attention K, softmax
V.
Vdgg
where n is the number of heads, is the attention head index, de is the dimensionality the keys and Wo, W W WV are learned projection matrices for the output, query, key and value, re- spectively.
2

---

Under review as a conference paper at ICLR 2022
We propose scaling the output of each attention head via learned scalar coefficients % HeadScaleMHAAQ, K, V) (o t Tn W
where y are learaalle parameters initialized to 1.
Additional Layer Normalization and Putting All Together the Pre-LN transformer each layer / modifies an input 27 as follows:
Pre18 = FFN(MHA((1))
1+1
where
MHA(2) (x) AutttlaalkAttrrooooCAAAAA,LL,LLL
FFN ()) = x + (LN(()))1 + b + b2
LN(x) TayerNorm(z))
In this work is the GELU non-linear activation introduced in Hendrycks & Gimpel (2016).
Our overall method, No rmF o ormer, instead modifies each input x; as:
NormFormer
NormFFFN(NommSaaceelllAtttt)))))
F1+1
where
NormScaledMHA((2 = x LDN
NormFFN(z) LN(r(LN((>WW b,)) 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 variance heuristic, and Admin (Liu et al., 2020) applies similar heuristic in profiling and initialization stages. These works also use variants of our ResScale operation, which find helpful 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, slightly different LN formulation, LN after the final fully connected layer, and the same post-attention LN. Section 9.1 compares roF ormer r to these proposals, as well as the T5 LayerNorm Variant (Raffel al., 2020), which removes the bias and the mean subtraction from the normalization.
Our HeadScale operation is related to that used Chen al. (2021), but used differently. Whereas that work prunes attention heads with low parameters, we use the parameters to im- prove 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 late FFN parameters, rather than increasing their importance and gradient norm, 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 al., 2021). We thus experiment with scaling the residual in each embedding dimension via learned scalar coefficients (Aresid)::
ResScale(x) = Aresid o + Sublayer (LayerNorm(z))
where elementwise multiplication, and Aresid are learned parameters initialized to
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,
(F))FFN(z) Aresid + LN( + b,)) + by

---

Under review as a conference paper at ICLR 2022
Model Size
GPT-3 Paper
Baseline
NormFormer
125M
6e-4
3e-3
3e-3
355M
3e-4
le-3
1e-3
1.3B
2e-4
6e-4
6e-4
Table I: Searching for learning rates on our dataset results in higher values than reported in Brown et al. (2020), providing stronger baselines to compare to our NormFormer architecture.
For 1.3B parameter models and larger, scaling residuals hurts performance (see discussion in Sec- tion 6.1), so ResScale not used in our 1.3B and 2.7B CLM results. Additionally, we experiment with initializing Aresid = le - 5, following (Touvron et al., 2021), as well replacing addition in residual connections with concatenation (Davis al., 2021) Section 9.1.
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).
only dense at- tention, while they alternate between dense and locally banded sparse attention; (2) we train our models with sinusoidal positional embeddings, following Shortformer (Press al., 2020b), since early experiments found this to produce comparable results with fewer learned parameters.
Our model architecture differs from Brown al. (2020) in two ways: (1)
We train the baseline models for 300 billion tokens. alent number of GPU hours, which typically results in 2-6% fewer steps and tokens due to the additional overhead of the normalization operations.
We train NormF oreer models for equiv-
On our dataset, we find that the learning rates proposed in GPT-3 are suboptimally low. For both baseline and NormFormer 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: (1c-4, 6c-4, 3e-4, 6c-4, 1.-3, 3,-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 the full training budget with the GPT-3 learning rates than with the higher learning rates. Other hyperparameters do not differ from GPT-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 Relu² 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. 7B with a higher learning rate of 6e-4. We note that training baseline 2.7B CLMs (i.e., without NormF r 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, is possible that intermediate value would perform better.
Zero Shot Evaluation In addition validation perplexity, we evaluate CLMs on subset of the tasks that GPT3 evaluated on in zero-shot setting (Brown et al., 2020), with the same prompts. We select WinoGrande (Sakaguchi et al., 2020), StoryCloze (Mostafazadeh et al., 2016), Open- BookQA (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.
The difference in optimal learning rates may due partly to architectural differences between our baseline and GPT-3 (c.g., not using locally banded sparse attention). "ee Table 2.1 in Brown al. (2020).
4

---

Under review as a conference paper at ICLR 2022
Causal LM (1.3B)
Masked LM (125M)
Nammhorme
Validation Perplexity
Normforme + real
Validation Perplexity
40-
10-
44-
320
40
200
120
140
100
A100 Days
A100 Days
Figure 2:
Pretraining perplexity on held-out validation data for Causal and Masked Language Mod- els as function of training compute (GPU days). The blue stars show the point where model matches the baseline's lowest perplexity.
Masked Language Models (MLM) We adopt the RoBERTa-base, Pre-LN architecture and hy- perparameters used in Liu et al. (2019). For the baseline, we pretrain for 2 million batches of million tokens, about & 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 1c-5, 1, 3e-4, 1e-3, 3, 3. 6 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 hyperparam- eters match those used for robert -baee in Liu et al. (2019).
Pretraining data We pretrain all models 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 al. (2019), originally introduced in Radford et (2019). The combined dataset contains around 450GB of uncompressed text and 110B BPE tokens. We hold out 40M BPE tokens from this data as validation set on which we report pretraining perplexities.
collection of English language text including the
Implementation details We train our causal and masked language models in fai rseq (Ott et al., 2019; Paszke et al., 2019). Although NormFormer introduces fewer than 0.07% additional parame- ters, 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 equal amount of GPU time, i.c.. controlling for compute rather than the number of training updates. Finally, we note that the HeadScal le operation can be moved outside the self attention module allow the use of the very efficient pytorch F. mul t hhw.haaaateetete ion. This change reduces overhead without noticeable performance degradation.
5 RESULTS
We report pretraining perplexities for CLMs and MLMS as function of training wall-time (GPU days) in Figure 2. We observe that NormFormer trains significantly faster and achieves better vali- dation perplexities for given training compute budget. The blue stars mark the first validation step where NormFormer matches the baseline's '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 add ResSca l NormF rmer-1..3 3R.
5

---

Under review as a conference paper at ICLR 2022
[0)
LR
Relu
Aveaid
Steps
PPL
HS
Pl
WG
SC
OB
Avg
Random Baseline
25.0
50.0
50.0
50.0
25.0
40.0
GPT3-125M (paper)
124.4
6c-4
572K
33.7
64.6
52.0
63.3
35.6
49.8
GPT3-125M (replicated)
124.4
6c-4
572K
21.11
33.7
66.5
52.2
66.1
35.4
50.8
GPT3-125M (High LR)
124.4
3o-3
572K
21.09
35.3
67.5
50.5
66.3
35.0
50.9
NormFormer- 125M
124.5
3e-3
540K
20.34
34.9
67.1
52.3
66.3
38.0
51.7
NormFormer-12255
124.5
3e-3
539K
20.11
34.9
65.9
53,4
67.5
40.0
52.3
GPT3-355M (paper)
354.7
3e-4
572K
43.6
70.2
52,1
68.5
43.2
55.5
GPT3-355M (replicated)
354.7
3e.4
572K
15.41
46.1
70.8
54.6
71.1
41.2
56.8
GPT3-355M (High LR)
354.7
lo-3
572K
14.85
71.7
53.8
73.3
43.4
58.1
NomFormer-35555
355.0
1e-3
552K
14.54
49.7
71.8
56.0
73.8
43.6
.9.0
NornFormer
355.0
1e-3
550K
14.52
49.7
72.0
56.7
73.2
43.8
59.1
GPT3-1.38 (paper)
1313.5
2e-4
286K
54.7
75.1
58.0
73.4
46.8
61.6
GPT3-1.3B (replicated)
1313.5
2c-4
286K
12.56
58.5
74.6
58.1
76.8
49.4
63.5
GPT3-1.3B (High LR)
1313.5
62-4
286K
12.21
57.5
74.3
59.3
76.3
50.8
63.6
NormFormet-11.3
1314.0
6c-4
275K
11.94
60.5
74.5
60.1
77.5
50.8
64.7
GPT3-2.7B (paper)
2648.7
1.6c-4
286K
62.8
75.6
62.3
77.2
53.0
66.2
GPT3-2.7B (replicated)
2648.7
1.6c44
286K
10.92
65.9
76.6
61.4
78.2
49.6
66.3
NormFormer-2.755
2649.5
6c-4
277K
10.55
68.1
78.1
64.4
79.4
53.4
68.7
GPT3-2.7B- Relu
2648.7
1.6e-4
230K
10.99
65.9
76.1
63.2
79.3
49.4
66.8
GPT3-2.7B-RGC
2648.7
6e-4
diverged
NormfFormer-2.775
2649.5
6c-4
222K
10.73
67.4
77.2
64.4
78.9
52.6
68.1
Table 2: Zero-Shot Accuracy for Causal LMs for the following tasks: HS: HellaSwag, PI: PIQA, WG: WinoGrande, SC: StoryCloze, OB: OpenBookQA. PPL. is validation perplexity during pretrain- ing. GPT-3 (paper) results taken from Brown et al. (2020). Horizontal lines group compute-matched runs. High LR corresponds to using a larger learning rate than reported in Brown et al. (2020). Arestd indicates whether residual scaling was used. Xresid did not help at 1.3B scale, as shown 2, but that run is not compute matched so it is not included here. Model size (10) is reported in millions of parameters.
Model Size
Areaid
PPL
CoLA
MNLI
MRPC
QNLI
QQP
RTE
SST-2
Avg
Bascline 125.42
3.42
74.3
85.9
84.6
91.6
90.7
66.4
92.9
83.77
NormFormer 125.50
3.31
82.6
86.3
86.0
91.9
91.3
67.9
93.8
85.69
NormFormer 125.51
3.29
80.9
86.2
85.3
91.5
91.2
62.8
94.2
84.59
Table 3: Masked LM: Pretraining validation perplexity (PPL) and fine-tuned performance on GLUE tasks for Pre-LN and NormFormer models. Note that models are trained for an equal amount of compute, which is less than the publicly-released roberta-base models.
We observe a similar trend on downstream tasks. Table we report zero shot accuracy for causal LMs using the tasks and prompts from Brown et al. (2020). NormForner 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 780, for example, results in a 127 million parameter model that is only 0.08 perplexity better than the baseline whereas No rmF reee r- 25MM adds only 100,000 parameters and is 0.83 perplexity better than the baseline.
For MLM models, we report fine-tuned accuracy on GLUE Table 3. We again find that Norm- Former 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 3.31), but the gains to do not translate to finetuned performance.
6 ANALYSIS
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 (Mi- cikevicius al., 2018). Figure 3 shows the average L.1 norm of the gradients to the second fully

---

Under review as a conference paper at ICLR 2022
Post-LN
Pre-LN (Baseline)
NormFormer
600
600
600
Layer
Layer
Layer
Grad Norm
400
10
400
10
400
10
11
200
200
200
230
230
400
430
300
350
600
230
500
330
400
430
500
000
230
330
400
430
500
330
600
Step
Step
Step
Figure 3: Average LI norm of gradients to the second fully connected weight for layers 0,1,6,10 and 11, early in training.
FFN LayerNorm
Post Attn LayerNorm
HeadScale
2.00
2.00
1.73
1.75
1.75
1.50
1.50
1.50
Weight
1.25
1.25
1.25
1.00
1.00
1.00
0.75
0.75
0.75
0.50
0.50
0.5
0.25
10
10 11
Layer
Layer
Layer
Figure 4: Distribution of learned scaling parameters in three of the added operations. For FFN LN, earlier layers receive downscaled inputs, keeping their gradients in the same range as the gradients of later layers. This plot discussed in detail in Section
connected weight in various layers for 12 layer, 125M parameter CLM model at the beginning of training. As reported in past work (Xiong et al., 2020), 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 significantly larger gradients than later layers. FormFormer 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 7 coefficients below indicating downscaling." 4 The HeadScale 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 Chen al. (2021).
parameters, shown in the rightmost plot in Figure 4 vary more
Reducing gradient mismatch allows training stably with larger learning rates. To measure the sta- bility of an architecture, we train 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 overflowing the FP16 range. The down scaling of
'The downscaling is also apparent in Figure 7 in the Appendix, which plots the change in grad norm for each operation at each layer. It shows that adding extra normalization reduces the gradient norm for all attention parameters at every layer. Only FFN parameters at later layers, have increased gradient norms.

---

Under review as a conference paper at ICLR 2022
Stability Test Results
Updates before Explosion
Architecture
300
Baseline
NormFormer
NormFormer r.resid
250
200
150
100
50
125.0M
355.0M
1.38
2.7B
6.7B
13B
Model Size
Figure 5: LR Stability Test: learning rate starts from 0 and linearly increases by 5e-5 at each train- ing step until training destabilizes. NormFormer reaches higher learning rate before destabilizing. Each data point the median of runs with a different random seed.
the attention outputs allows NormFormer to avoid this issue and remain stable with larger learning rates. Figure also shows that Aresid reduces the stability improvement at all sizes.
6.1 RESIDUAL SCALING
By comparing adjacent NormFormer- I 125M and NormFFrrmr-335M rows in Table 2 can that adding ResScale to NormForme improves perplexity and zero shot performance for small scale CLMs. For 125M parameter MLM, ResScale improves pre-training perplexity marginally, but hurts fine-tuned performance. At 1.3 billion parameter scale, however, adding ResSca l to No corFormee does not improve performance (Figure 2). Although it's included in our tables, we find 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 NormFornee
Figure the appendix shows the average Aresid weights each layer 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 Xresid parameters to the other (earlier) residual connection in each layer, or using a scalar instead of a vector for each Aresid, does not fix 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 al., 2021), does not improve performance.
ABLATIONS
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 per- formance across wide range of hyperparameter settings. Experiments and are run with the default hyperparameters given in Table the appendix for 470 V100 Hours (100,000 updates for the baseline) unless otherwise mentioned.
125M parameter CLMs,
Removing any of the added operations hurts performance Table shows that none of the four introduced operations can be removed without degrading performance. Rows 2-5 remove each operation time. In all cases perplexity increases, with the removal of HeadScall e being
8

---

Under review as a conference paper at ICLR 2022
Architecture
Valid PPL.
NormFormer-ResScal
15.88
Post-Attn LN
15.92
-FFN LN
16.14
Head Scale
16.22
Res Scale
16.20
+ More LN
15.88
Baseline
16.37
Table 4: 125M parameter Language Modeling Validation perplexities after 470 V100 Hours of pretraining. Removing any our proposed additions degrades performance (Rows 2-5). Adding more normalization inside the Multi Headed Attention (Row 6) does not impact perplexity at fixed number of updates, but reduces throughput such that the model can only complete 87,500 updates vs. 92,500 for Rows 1-5 and 100,000 for Row 7. Note that these PPL scores not directly comparable to other tables they use a different validation set.
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 our other operations, for total of 6 new operations. In this setting, every other parameterized operation inside the transformer layer is an LN. We find that this does change perplexities at a fixed 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 NormE ormer.
Other Experiments Table 8 in the appendix compares NormFormes CLMs to related architec- tural modifications, 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- figurations at 125M parameter scale, separated by horizontal lines. Norm r outperforms baselines in all settings.
CONCLUSION
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 larger earlier layers than at later layers. We propose NormF ormer, which alleviates these issues by adding 3 extra opera- tions each transformer layer. These modifications 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 gra- dients to the last FFN layers while decreasing the gradient magnitudes in other parts of the network, future work could examine whether all operations need to be added to every layer. Additionally, the small computational overhead associated with NormForner 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 improve pretraining perplexity. In general, 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 networks. Future work should ascertain if there are additional similarly efficient modifications that can bring gains, while helping us understand current deficiencies further.
9

---

Under review as a conference paper at ICLR 2022
9
APPENDIX
1.0
0.9
A Aesid (Avg.)
0.8
0.7-
0.5
1.38/44
125M/12
0.4
355M/24
0.6
0.8
0.9
1.0
Layer (Depth)
Figure 6: Average Aresid weights each layer of different sized CLMs in the NormFormes Avesid setting. Depth is layer number total layers.
Learning Rate
Setting Changes
Valid PPL
Baseline
0.001
16.80
NormFormer
0.001
16.33
Baseline
0.003
16.37
NormFormer
0.003
15.88
Baseline
0.006
16.58
NormFormer
0.006
16.22
Baseline
0.003
Longer Warmup
16.50
NormFormer
0.003
Longer Warmup
16.06
Baseline
0.003
GPT3
16.29
NormFormer 0,003
GPT3
15.88
Baseline 0.003
Clip Grad Norms 0.1
16.46
NormFormer 0.003
Clip Grad Norms 0.
16.14
Table 5: Longer Warmup: increase LR Warmup 6,000 steps (from 500). GPT3: increase sequence length to 2048, increase dropout to 0.1, increase training budget to 1,000 V100 hours. Grad Clip: clip gradient norms 0.1. NormFormer outperforms the baseline in all settings.
Wikitext103 Table shows that NormFormer can also provide gains on top of well tuned lan- guage 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, Normb ormer does not improve in the last 30% of training, which suggests that with more tuning the perplexity gap could be widened.
9.1 COMPARISON RELATED WORK
Understanding Table 8 The two right most columns of Table 8 contains describes the result of two experiments from the same 1.3B parameter architecture: Stabilit 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 the model trained for 64 A100 days with sequence length 512. The table is sorted by Stability, with more stable configurations lower in the table. The left columns of the table indicate configuration. Row 14, which has completely empty configuration columns, is the baseline. The final row No rmFormer..
10

---

Under review as a conference paper at ICLR 2022
Steps to Target PPL.
Final PPL.
A100 Hours
Baseline NormFormer
279.893
18.70
288
223.904
18.65
237
Table 6: Wikitext 103 results following Baevski & Auli (2019). Steps to Target PPI: at what percentage of the 280K steps did the model reach 18.70 perplexity. Final PPL: Best Per- plexity
A100 Hours Cost of reaching Target PPL.
Change in Grad Norm Scale
Operation
krpr
[_proj V_proj
out pro) f:1
R22
Layer
Figure 7: Change in grad norm with each operation of NormFormer compared the baseline. Norms are the average between step 950 and 1000, normalized to control for different losses. 2.0 on the Y axis means the gradient to parameter is twice large the baseline, on average. The NormFormer increases the norm to fully connected parameters in later layers, while reducing the gradient norm attention parameters at all layers. The results are discussed in detail in Section
The Architecture and LN Variant columns show whether we changed the model architecture com- pletely and/or used non-standard LayerNorm algorithm.
1. Cog View [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 first, like NormFormer) (4) change the layer norm formula to LN X )) fairseq already uses the attention score stabilization trick by default. rows 19 and 25 where Arch is blank, we use proposed LayerNorm formula but none of the other changes.
2. CatFormer[2] (Davis et al., 2021): Since no code was released, we re-implement CatFormer and set 2 and =448 match 1.3B parameters. We guessed that the number of attention heads should be fixed in each layer and that LN positioning should not move, but these details are not clear from the paper.
3. LayerScale[3] (Touvron et al.. 2021): We the layer scale formulation from Section 2, which is similar to the Aressd discussed earlier, but with a weight on each layer's 's contribu- tion to the main branch initialized 1c-5, instead of weight 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 Layer- Norm, with no normalization or bias.
5. T5[5] (Raffel et al., 2020): Switch LNs to the T5 variant, which removes the mean centering and bias.
6. PowerNorm[6] (Shen al., 2020): Switch LN PowerNorm, which variant of Batch- Norm that shows promising results for smaller NLP models.
DeepInit[7] (Zhang et al., 2019): Multiply initialization of each weight parameter by / where is the layer number starting from 1.
factor of 1

---

Under review as a conference paper at ICLR 2022
Learning Rate
0,003
Batch Size
524K Tokens 124M+
Parameters
Layers
12
Layer Dimension
768
Dropout
Warmup Updates
LR Scheduler
Linear Decay
Sequence Length
1024
Train Budget
470 V100 Hours
Table 7: Hyperparameters for ablations in Tables and This train budget allows the baseline model run for 100,000 updates.
ID
Arch.
Scale FC
Scale Attn
XReaid
Scale Heads
LN Variant
Stability
Perf
0
PowerNorm[6]
15
1
CatFormer[2]
2
52
3
DeepInit[7]
52
4
DeepInit[]]
58
5
FFNGeglu(4)
67
17.1
6
7
LayerScale(3]
76
17.5
8
76
17.1
9
17.1
10
[ogView[1]
CogView[1]
86
11
91
12
TS[5]
93
13
FFNGeglu[4]
94
14
17.1
15
No Bias or y
95
16
No Bias
96
17
No
18
112
16.8
19
CogView[1]
CogView[11
118
17.1
20
Cog [ie
122
21
CogView[1]
Cog ViewW]
124
22
128
23
PowerNorm[6]
135
24 LayerScale[3]
148
25
No
157
26
CogView[I]
165
27
No Bias or y
178
16.9
28
No Bias T5[5]
184
29
189
30
192
16.7
31
200
16.6
Table 8: Stability and Performance for different architectures for 1.3B parameter CLMss see Sec-
tion 9.1 for details.
12

---

Under review as a conference paper at ICLR 2022
8. No y: Freeze y 1 in LN
9. No Bias or y Freeze = 1, bias=O in LN
10. No Bias: Freeze bias=0 in LN.
Results The results suggest that other proposals mitigate instability in other domains, like vi- sion, text to image generation and reinforcement learning do not improve stability or pre-training performance in the CLM setting.
13

---

Under review as a conference paper at ICLR 2022
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. In International Conference on Learning Representations, 2019. URL https: /ttp///pppeeeeeeee//eeeeee net futt/fozee
Yonatan Bisk, Rowan Zellers, Ronan Le bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. Proceedings of the AAAI Conference on Artificial Intelligence, 34(05)7433--4499 Apr. 2020. doi: 10.1609/aaai, wt41://aaaiii430...... 6,299, URL https://
ojs. aaai. org/ it////ndexxxx/// php/ AAAI/ wwt.cceevieeeee//eeeeeeee 6239.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler. Ma- teusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Ad- vances in Neural Information Processing Systems, volume 33, pp. 1877-1901, Curran Asso- ciates, Inc., 2020. URL 1457c0d6bfcbb 91116661118111811144224221 pdf.
Xiaohan Chen, Yu Cheng, Shuohang Wang, 7he Gan, Zhangyang Wang, and Jingjing Liu. Earlybert: Efficient bert training via early-bird lottery tickets, 2021.
Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume Wenzek, Francisco Guzmán, Edouard Grave, Myle Ott, Luke Zettlemoyer, and Veselin Stoyanov. Un- supervised cross-lingual representation learning at scale. 2020.
Jared Q Davis, Albert Gu, Krzysztof Choromanski, Tri Dao, Christopher Re, Chelsea Finn, and Percy Liang. Catformer: Designing stable transformers via sensitivity analysis. Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Leurning Research, pp. 2489-2499, PMLR, 18-24 Jul 2021. URL https:tt//.......... httpt:t:://prooeee http///oeeednieee ngs mwl...mlr...... prw.s /v1 http//l19///aaa davi a21a.html.
Ming Ding, Zhuoyi Yang, Wenyi Hong, Wendi Zheng, Chang Zhou, Da Yin, Junyang Lin, Xu Zou, Zhou Shao, Hongxia Yang, al. Cogview: Mastering text-to-image generation transformers. arXiv preprint arXiv:2105. 10022000.... 2021.
Dan Hendrycks and Kevin Gimpel.
Gaussian error linear units (gelus).
arXiv preprint
arXiv: 160::666.000 wew.16606811.. 2016.
Opher Lieber, Or Sharir, Barak Lenz, and Yoav Shoham. Jurassic-1: Technical details and evalua- tion. Technical report, A121 Labs, August 2021.
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the diffi- culty of training transformers. arXiv preprint 2XXx:2004400244 2020.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907, 1692, 2019.
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed precision training. In International Conference on Learning Representations, 2018
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor con- duct electricity? a new dataset for open book question answering. In Proceedings of the 2018 Conference Empirical Methods in Natural Language Processing. pp. 2381-2391, Brussels, Belgium, 1/188 URL https::/p///. : wwwww//www../////////////w// acl web.orgg htt.://nooollllyyoaeoaaa
2018. Association for Computational Linguistics. doi:
14

---

Under review as a conference paper at ICLR 2022
Nasrin Mostafazadeh, Nathanael Chambers, Xiaodong He, Devi Parikh, Dhruv Batra, Lucy Van- derwende, Pushmeet Kohli, and James Allen. A corpus and cloze evaluation for deeper under- standing of commonsense stories. In Proceedings of the 2016 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, pp- 839- 849, San Diego, California, June 2016. Association for Computational Linguistics. doi: 10.18653/v 1/N1//////1111/ 1098. URL https: www aclweb. org/ ha//atthoologgg
wwt.wee....//atthoooog..... /N16 1098.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. FAIRSEQ: A fast, extensible toolkit for sequence modeling. 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury. Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. pp. 8024-8035, 2019.
Ofir Press, Noah A. Smith, and Omer Levy. Improving transformer models by reordering their sublayers, 2020a.
Ofir Press, Noah A Smith, and Mike Lewis. Shortformer: Better language modeling using shorter inputs. arXiv preprint arXiv:20/2.15833 2020b.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and lly Sutskever. Language models are unsupervised multitask learners. Technical report, OpenAI, 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. 21:1-67, 2020.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An ad- versarial winograd schema challenge at scale. Proceedings of the AAAI Conference on Artifi- cial Intelligence, 34(00):8732-877 Apr. 2020. doi: 10. www00//aai..3440....o.. URL. https: //ojs. aaai org/ ndex php/ AAAT htw.A///aaaeeeeeee////ee
Noam Shazeer. Glu variants improve transformer, 2020.
Sheng Shen, Zhewei Yao, Amir Gholami, Michael W. Mahoney, and Kurt Keutzer. Powernorm:
Rethinking batch normalization transformers, 2020.
David R. So, Wojeiech Mańke, Hanxiao Liu, Zihang Dai, Noam Shazeer, and Quoc V. Le. Primer: Searching for efficient transformers for language modeling, 2021.
Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going
deeper with image transformers, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, 2017.
Alex Wang, Amanpreet Singh. Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. 2019.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng. Shuxin Zheng, Chen Xing. Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie-Ya Yan Liu. On layer normalization in the transformer architec- ture, 2020.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Associ- ation for Computational Linguistics, pp. 4791-4800, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v 115/33/11/1111177777 1472. URL https:::// www. aclweb. org/ anthology /P1 / 1912-
Biao Zhang, Ivan Titov, and Rico Sennrich. Improving deep transformer with depth-scaled initial- ization and merged attention, 2019.
15

---

Under review as a conference paper at ICLR 2022
Chen Zhu, Renkun Ni, Zheng Xu, Kezhi Kong, W Ronny Huang. and Tom Goldstein. Gra- dinit: Learning to initialize neural networks for stable and efficient training. arXiv preprint arXiv:2102.08008, 2021.
Yukun Zhu, Ryan Kiros, Richard Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba. and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. arXiv preprint 4rX14:4506.06724. 2019.
16


Reviews:
[{"review_id": "zMxxguTKqCd", "reviewer": "Reviewer_uRwh", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n1. The paper is well-written and easy to follow.\n2. The four operations are easy to implement and applicable to many tasks using transformer architectures.\nSome visualizations, e.g., Figure 3, provide straightforward comparisons.\n\nWeaknesses\n1. Although the four operations are shown effective in accelerating pretraining convergence or boosting downstream performance, they look more likely engineering tricks. These tricks have been widely used in prior works, and I don't see much insight in being used here.\n2. There is a lack of analysis of why the gradient mismatch issue exists in Pre-LayerNorm transformers and why the proposed four operations can alleviate it. The paper mainly focuses on what it does rather than justifying why doing that.\n3. There are two scaling operations added in Figure 1, but three are plotted in Figure 4. It is unclear how the left two subfigures in Figure 4 correspond to the scaled residual connection in Figure 1.", "summary_of_the_paper": "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.", "summary_of_the_review": "In summary, this paper proposes to add four operations, two LayerNorms and two scaling parameters, in the Pre-LN transformer layers. The operations are simple and can help stabilize pretraining and improve downstream performance in several NLP tasks. However, I think the method novelty is not enough. Moreover, the paper doesn't provide convincing motivation why the four operations help handle the issues.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "zIDNSxfTdlj", "reviewer": "Reviewer_ust2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper ​proposes some modifications including two normalizations and two new scaling operations to mitigate the issue of gradient magnitude mismatch inside the transformer. The resulting Normformer improves pretraining perplexity and downstream task performance with a negligible increase of computing cost.\n\n(1) The paper is concisely written and easy to understand. However, the contributions are not clearly presented in the Introduction. Besides, since the section of related work is put at the end of the paper, the readers may get confused about how the proposed method differentiates from existing works.\n\n(2) The issue of gradient magnitude mismatch is not well explained. For example, how the gradient magnitude mismatch affects the training. Moreover, how the optimal weighting of residuals is connected with this issue is unclear. As they are the main motivations of the proposed method, it is suggested to provide a detailed explanation about how gradient magnitude mismatch, optimal weighting connects with the proposed method.\n\n(3) As shown in Fig.3 of  [1], the PreLN with warmup can effectively tackle the problem of gradient magnitude mismatch inside the FFN of the transformer. Can you explain why the issue still emerges in transformers with a common learning schedule? For example, can you provide some intuitions on why gradients at early layers are much larger than at later layers in PreLN transformers? As shown in Fig.3, gradnorm in only several layers are visualized.\n\n(4) I appreciate that the ablation study is conducted to show the performance gain of each component in  NormFormer. However, since four modifications are proposed to solve a single issue, it is expected to see how these four modifications affect the gradient magnitude and how these modifications interact with each other.\n\n(5) The effect of HeadScale parameters is overclaimed. The readers cannot see any evidence showing that HeadScale parameters can adjust the importance of attention heads.\t\n\n(6) Are there other studies investigating training instability? Please discuss more in related work. It is good to show the effectiveness of NormFormer through extensive experiments. \n\n(7) The results seem convincing and competitive.\n", "summary_of_the_paper": "Competitive empirical results but four modifications (two norms and two scalings) are not well motivated and explained.", "summary_of_the_review": "\nOverall, this paper proposes NormFormer to improve the PreLN transformer by introducing some modifications. The proposed techniques are empirically successful but not well motivated. I hope that the authors can firstly present the issue inside the PreLN transformer clearly and then provide a detailed discussion about why the proposed modifications can mitigate this issue. \n\t\t", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "lHgESZhwyNT", "reviewer": "Reviewer_o2vf", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "In general, this is a good paper with sufficient and exhaustive experiment results. The 60% time-saving is impressive. The ablation studies show that both rescaling and more layer normalizations are contributing to better quality in MaskLM accuracies and zero-shot tasks. The 3 GPT-3 models of different size scales prove the generalizability of the techniques.\n\nThere are some small points that I think the authors can improve to make their statement stronger:\n* The current framework, though claimed to be Transformer-related modification, is conducted in GPT-3 only. It would be helpful to see whether this Transformer modification can be used in other frameworks such as BERT and T5 or a single Transformer model.\n* The Peak LR experiment does show that the current modification will allow the NormFormer to have up to 0.0275 learning rate stability, which compared to the original 0.02 baseline does not have a magnitude difference. It would not be that a strong statement to argue for the added robustness if higher LR would not always lead to better quality in general. I hope the authors can show more evidence if this is the case and the reason why they care about this aspect.\n", "summary_of_the_paper": "This paper proposes a framework that adds extra layer normalizations and scaled residual layer to the normal GPT-3 model. Empirical results demonstrate faster convergence speed and its robustness towards the learning rate.\n", "summary_of_the_review": "The paper shows a simple yet effective way of modifying the existing GPT-3 framework by increasing only a small number of parameters and computation per epoch, which can be a potential addon to many existing Transformer architectures.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "SJzg51N48ua", "reviewer": "Reviewer_jDo1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I think this is mostly solid work. The writing and methods are easy to understand and the contributions are clear and convincing. However I think the baselines could be more challenging and the analyses more convincing.\n\n# Feedback\n\nI appreciate the authors' clear communication and presentation of their work. I also applaud their use of GPU-hour-matched baselines.\n\nMy main concern with the present work is that it is evaluated against baselines with no stability/trainability interventions. Such interventions exist; the authors mention a number of them in the *Related Work* section. But without including any of them, it's difficult to directly assess whether the present work is an improvement over existing methods.\n\nI’m a bit confused. The caption of Figure 2 says “The green star shows the point where NormFormer outperforms the baseline’s lowest perplexity, reflecting 22% and 43% compute savings for CLM and MLM models, respectively.”, while the first paragraph of results says “NormFormer matches Pre-LN models while needing only 60% and 57% as much compute for CLM and MLM models, respectively”. Either I’m misunderstanding your metrics, or one of these statements is incorrect. Regardless, I would encourage the authors to pick one efficiency metric and use it consistently and exclusively. Personally, I think “x% as much compute” or “x% the amount of compute” are more intuitive, as they clearly mean “(normformer/baseline)*100 amount of compute”.\n\nI don’t find the gradient norm analysis (Section 4.1) particularly convincing. It’s hard to know whether the effect of NormFormer on gradient norms is meaningful, and even whether it's causally related to the performance improvements. It would be great to see some experiments demonstrating the causal effect of gradient norm changes on trainability. At the very least I think readers would appreciate some theoretical justification or intuition. Why is it desirable to have similar gradient norms across layers? Also in Figure 3, most of the Layer 0 gradient norms are out of the frame in both Pre-LN plots, making comparison difficult. I suggest changing all three plots to a logarithmic y-axis.\n\nI think it would be very informative to examine the scaling parameters over the course of training—especially early in training. The early phase of transformer training seems to be the most fraught. This is when the loss tends to explode, and this is when you (and others) show the largest changes in gradient magnitudes. Analyzing only the final learned parameter values ignores the most important part of the story. \n\nThere seems to be an implicit assumption that the converged scaling parameters (e.g. Figure 4) are “good for training”. Does initializing the NormFormer parameters to be of similar magnitude to what is observed at convergence yield improvements over the current initialization scheme?\n\nLayerNorm consists of a scale and bias, but only the scale parameter is analyzed. It’s possible that the bias could counteract the effect of the scaling. Although I think this scenario is unlikely, it should nevertheless be accounted for. \n\nThe ResScale results (end of Section 4) should be expanded and plotted. And why did the authors choose only to look at the minimum value of lambda_resid? The layerwise analysis of lambda_resid should be repeated with different metrics.\n\nPlease include quantitative results from the **Other Experiments** section.\n\nThere are a few relevant papers that might be worth discussing (or at least referencing). Zhang et al. conducted earlier research on the importance of initialization (*Improving Deep Transformer with Depth-Scaled Initialization and Merged Attention*, 2019). GradInit (Zhu et al., 2021) is recent, and also quite relevant. PowerNorm (Shen et al., 2020) and Catformer (Davis et al., 2021) both also take similar approaches to solving transformers’ trainability issues. LayerScale (from Touvron et al.’s *Going Deeper with Image Transformers*, 2021) examines the trainability problem in vision transformers and find a solution that appears similar to one the present innovations. Finally, Brock et al. (*Characterizing signal propagation to close the performance gap in unnormalized ResNets*, 2021) examine CNNs, but their approach and findings are quite relevant.\n", "summary_of_the_paper": "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.", "summary_of_the_review": "Transformers are notoriously unstable to train. Progress on this problem is of clear value to the field. The baselines could be more challenging and the analyses more convincing, but if these issues are addressed—which I am optimistic about—I think the paper could be a meaningful contribution to the field.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This submission proposes a few small changes to the (PreLN) Transformer architecture that enable training with higher learning rates (and therefore can result in faster convergence). The changes include the addition of two layer norm operations as well as a learnable head scaling operation in multi-headed attention. The proposed operations add only a small computational overhead and should be simple to implement. Experiments are conducted on language modeling and masked language modeling, with improved results demonstrated at various scales and according to various evaluation procedures. The paper also includes a good amount of ablation study as well as some analysis. Reviews on the paper were mixed, and a great deal of changes were made to the paper during the rebuttal period. To summarize the concerns and recommendations, reviewers requested
- better connection between the proposed changes and the purported issue (gradient scale mismatch between early/late layers)
- better analysis of why gradient scale mismatch is a major issue and investigation of where it comes from
- better comparison to existing techniques that allow for higher learning rate training of Transformers
- additional experiments on different model types and ideally different codebases/implementations

I think overall this is a solid submission, since it proposes a simple change that is reasonably likely to be helpful (or at least not harmful). However, I think that there are enough concerns with the current draft and there were enough changes made during rebuttal that this paper should be resubmitted to a future conference. I would suggest the authors take the final updated form from this round, add additional motivation/analysis/experiments, and resubmit, and I suspect a positive outcome.
Decision: Reject