Normalization Layer Per-Example Gradients are Sufficient to Predict Gradient Noise Scale in Transformers

Summary

This paper explores efficient ways of computing the gradient noise scale (GNS) when training transformers. The main practical relevance of the GNS is that it can be used to estimate the critical batch size, where the larger batch sizes become computationally inefficient. The authors discuss different ways of doing this at a fine granularity both time and layerwise. They show that the backwards pass through a layernorm can be modified to estimate their GNS for essentially free and that this predicts the GNS of other layers. The authors then vary the batch size throughout training based on this, showing this can potentially save compute for a given target model performance.

Strengths

- The core idea of estimating the GNS in a cheap way based on normalization layers is interesting and relevant to the community. - The proposed layernorm based estimation can be performed very cheaply.

Weaknesses

- Some more work is required to make the proposed method practical due to numerical issues. Although the idea is interesting it would be much more impactful if the kernel worked as a drop-in replacement. - The 18% speed improvement claimed in the abstract may be an overclaim due to the numerical issues and presumably only being applicable to certain types of training setups (maybe one GPU doing gradient accumulation rather than distributed setups). - The paper is quite hard to follow, despite the relatively simple ideas presented. - The use of Einstein notation contributes to this, it is fine for brevity but not clarity. - The paper relies too heavily on McCandlish et al 2018. Despite spending significant effort on trying to summarize the relevant portions of this paper, many things are still unclear in the later sections. - The figures are hard to interpret, both due to very short captions that do not summarize the high level idea and takeaway, and insufficient labels on the figures themselves.

Questions

Overall I think the core ideas are interesting but this paper would strongly benefit from being resubmitted with improved writing and overall presentation. Suggestions for improvement: - Drop the Einstein notation, make things explicit instead of implicit. - Consider rewriting the introduction to gently introduce the GNS, maybe with a diagram and summarize the high level importance of this before diving in. What are the high level ideas needed to understand your contribution? Do you really need to go into all these details upfront? - Improve the description of your figures. Take Figure 1 for example, it is very hard to parse. The caption does not summarize the message, it should also emphasize how the dashed and plotted lines are related. - For Figures 2 and 3, maybe emphasize more what Li et al is, what the differences are and what you expect to see? Why not normalize the IO costs as well in the comparison (maybe you need some assumptions on the other kernels. Without some sort of reference number it is hard to see how much these IO overheads matter in practice. - Figure 5: Again, maybe tell us what is expected? What is the takeaway? - Figure 6: It is very unclear what is going on in the two plots on the right. Almost no description is given in the text or the caption of the “EMA alpha values”. - Batch Size Schedule: It is not clear how “a batch size schedule that increases linearly with the number of tokens processed to the original batch size” relates to your GNS estimations. Do you set the length of the schedule based on the GNS? Is this done based on prior estimates of the GNS or in an online fashion? - Figure 8: This should probably be done using learning rate schedules of different lengths for statements like 2x to be meaningful.

Rating

4

Confidence

4

Soundness

2

Presentation

1

Contribution

3

Limitations

Some limitations discussed, negative societal impact not applicable

Summary

The paper proposes a method for efficient computation of per-example gradient norm for the broader usecase of computing gradient noise scale. Further the authors showcase the usecase of gradient noise scale (GNS) in Transformers and showcases that GNS of only normalization layers in transformer models suffices, rather than the total GNS. Using this proposal, the paper proposes a batch size scheduling resulting in training time reduction.

Strengths

The paper for the most part is well written and has shown useful applications of gradient noise scale in Transformer architectures which can be further extrapolated to other architectures. e.g. state space models. - Related work section discussion is well motivated to discuss the utility of gradient noise scaling, perhaps a more better way to represent this section would have been to segregate the related work with smaller sections with headers - one discussing practical utility and other definitions used in literature for gradient noise scaling and any recent efforts towards efficient computation for per-gradient norms.

Weaknesses

Overall, I feel the main contributions of the paper is not well presented in this draft and in many portions the writing is hard to follow and seems disconnected from other sections in the paper. On one aspect the authors propose existing methods of efficient per-example gradient norms in Section 2.2 but only a few papers like Li et al. [29] and Goodfellow [22] are mostly discussed. It would be nice to have a much more better discussion of more related work to have a thorough pros and cons discussion. Notation writing could be improved to a better extent. For instance in Sec 2.2, line 74 which space does B, I, K indicate or is just a placeholder for 3D tensor dimensions. Similarily line 75 $x_{bti}y^{'}_{btk}x_{bui}y^{'}_{buk}$ , how does the index indicate u ? Related work discussion on other methods on per-gradient norm efficient computation is missing. For instance https://openreview.net/forum?id=xINTMAvPQA is cited but this one or similar other approaches how they have approached this problem and how the authors proposal differs in that respect , that thorough discussion would be appreciated. **Minor fixes in Writing** - Section 2.1: it may be shown [32] -> It may be shown McCandlish et al. [32]. And you may remove McClandish elsewhere. - Section 2.1 : I do not see any specific discussion on $B_{big}$ and $B_{small}$. How do they differ? Any references to later section in the paper where they define the size differences.

Questions

Step 4. of Algorithm 1. (what is the exact correction) and is very limited in scope in terms of contribution.

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

1

Limitations

Yes, I appreciate that the authors have adequately addressed the limitations of this work in terms of their empiricial evaluations only on Transformer architectures and not on other similar architectures like RNNs or state-space models, which is possibly left as a future work. Apart from this there are no limitations or any potential negative social impact of their work.

Summary

This paper proposes a more efficient method for computing per-example gradient norms without significant computational overhead in terms of FLOPs and I/O cost. This method accurately estimates the gradient noise scale (GNS), useful for neural network batch size selection and scheduling. Additionally, it observes that the normalization layer effectively predicts GNS in transformer models. These findings, along with the proposed algorithm, are applied in a case study on a large language model, where effective batch size scheduling reduces training time.

Strengths

- The proposed method efficiently computes per-example gradient norms and rapidly estimates the transformer gradient noise scale. - This work also investigates a downstream application of batch size scheduling and demonstrates its time-saving benefits during training.

Weaknesses

- If the proposed method is only supported by heuristic and empirical observations, and if the primary practical use of GNS is to estimate the critical batch size, then the experiment should be much more comprehensive. Only one set of experiments was conducted on a fixed model/dataset, making it unclear if the results are robust without ablation studies. Additionally, the results were not compared to other batch-size scheduling methods. - The finding that the normalization layer effectively predicts total GNS behaviour in transformers has not been thoroughly tested on other transformers, and it's unclear if this generalizes beyond transformers. The authors also do not provide any intuition as to why this might be the case. - The paper needs to be more clearly written and easier to read. Some concepts could be better explained or motivated. For example, it's not clear how per-example gradient norms help estimate GNS, what other ingredients are needed, and what other sources of noise there are that should be taken into consideration. It's also challenging to extract key information from the figure captions.

Questions

- Could the author add more experiments on different models and datasets to validate the generalizability of the method? To ensure robustness, it is also important to include ablation studies. Additionally, the results should be compared with other established batch-size scheduling methods to highlight the advantages or disadvantages of the proposed approach. - In section 2.3, it is claimed that the variance of gradient norms determines whether SGD outperforms adaptive optimizers. However, subsequent studies [Kunstner et al, Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be, ICLR 2023] have shown that the performance gap between Adam and SGD persists even in full-batch scenarios without stochasticity. This suggests that there might not be applications of GNS in this context. It would be beneficial for the author to address this discrepancy.

Rating

5

Confidence

2

Soundness

2

Presentation

2

Contribution

2

Limitations

Limitations are well-addressed

Summary

This work proposes a method to compute per example gradient norms as a means to compute GNS. It shows that not all layers are necessary to estimate the GNS and that the per-example gradient norms can be computed for normalized layers without any overhead.

Strengths

- This work provides an efficient technique for computing the GNS. The technique is supported by various experiments, and elucidates that the GNS of the model is highly correlated between layer types. Furthermore, for LayerNorm layers, this paper develops a custom kernel to compute the backward pass and the per-example gradient, and find that the the throughput overhead of gathering the per-example gradient is 0 (which outperforms PyTorch’s Layernorm). - The paper also replicates prior GNS observations, which helps support the method. - I found the experimental results interesting, especially the case study with dynamic batch sizes.

Weaknesses

- As the authors mention, estimating gradient noise scale is useful in training large scale models. The largest model included in the case study only has 111M parameters. I think that this work could benefit from experiments on larger architectures, given the size of language models used in practice today (e.g., at least 7B). - This study is limited to transformers. However, many other architectures do not use normalization sub-layers. This point is also acknowledged by the authors.

Questions

Perhaps the work can benefit from experiments with larger model sizes, given that even small open source models have 7B parameters.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

As acknowledged by the authors, this paper is limited to transformers, which inherently use normalization sublayer (other architectures do not conventionally use such layers).

Response from Reviewer

I thank the authors for responding to some of the questions I had. Based on their responses regarding clarity and other existing queries I had, I am happy to increase my rating