Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: 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.

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.

There are some small points that I think the authors can improve to make their statement stronger:
* 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.
* 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.

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.
review_title: null

final_score: rating: 8: accept, good paper
review_content: 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.
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.

# Feedback

I appreciate the authors' clear communication and presentation of their work. I also applaud their use of GPU-hour-matched baselines.

My 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.

I’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”.

I 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.

I 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. 

There 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?

LayerNorm 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. 

The 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.

Please include quantitative results from the **Other Experiments** section.

There 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.

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.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: Competitive empirical results but four modifications (two norms and two scalings) are not well motivated and explained.
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.

(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.

(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.

(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.

(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.

(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.	

(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. 

(7) The results seem convincing and competitive.

summary_of_the_review: 
Overall, 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.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: 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.
main_review: Strengths
1. The paper is well-written and easy to follow.
2. The four operations are easy to implement and applicable to many tasks using transformer architectures.
Some visualizations, e.g., Figure 3, provide straightforward comparisons.

Weaknesses
1. 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.
2. 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.
3. 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_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.
review_title: null

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