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

Thank you so much for your detailed and constructive review, and your kind words regarding the appeal and relevance of the paper’s core idea.

On the clarity of the paper

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.

Unfortunately, despite the known use of GNS as detailed in the Related Work section, there is limited research published on the topic since McCandlish et al. (2018). However, GNS estimation is of practical value in our work and we believe it will be for the community.

On Einstein notation

Drop the Einstein notation, make things explicit instead of implicit.

We agree that Einstein notation may be unclear. The most straightforward way to address this problem would be to write out the sums explicitly. However, we find the utility of the Einstein notation to be the ability to express that the order of the summation is flexible. Another better solution would be to write example matrix operations. For the equation on line 75, we could have written

as one possible contraction. We will add these examples to the paper to make the notation more clear.

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?

Thanks for this suggestion. Based on this we have created a diagram to shown in our attached pdf in Figure 1. We think this may help prime the reader for the concepts in the paper.

On Figure messaging

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.

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.

Thank you for this feedback. We agree that the figure captions should state the message of the figures. For example, in Figure 1, we will say, "For the same number of samples processed, a smaller $B_{small}$ always has a lower standard error (dashed), while the size of the large batch $B_{large}$ does not affect the standard error." The solid lines in this figure are the estimate of the GNS at each number of samples processed, supposed to illustrate the uncertainty caused by the finite number of samples. To simplify the figure it will be clearer to remove the solid lines.

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?

We agree that it would be beneficial to provide more context for Figures 2 and 3. We will add the following messages to the captions:

• Figure 2: "...The FLOP cost of Simultaneous per-example gradient norms is strictly dominant to alternative methods (left) and the ratio of this additional cost to the FLOP cost of processing the entire model does not depend on context length (right).
• Figure 3: "...The I/O cost of Simultaneous per-example norms is greater for models of 10B parameters or more but approximately equivalent for models of 1B parameters or less, depending on context length."

In addition we have tested normalized I/O costs in the Figure and found that it makes that Figure more clear. We will use this Figure in the updated paper.

Figure 5: Again, maybe tell us what is expected? What is the takeaway?

To Figure 5 we will add the message, "This Figure replicates an experiment from McCandlish et al. (2018) showing how the ratio of $\epsilon / B$ causes changes in the measured GNS but only due to changes in the learning rate. The batch size does not have the predicted effect."

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

We thank the reviewer for noticing that the description of the EMA on lines 177-178 does not mention how the alpha value controls smoothing. We will add an Appendix making the role of the alpha value and add EMA pseudo code, because there are many ways to implement it.

For Figure 6 we will add the message, "The total GNS (black) on the left is predicted well by individual layer types as indicated by the correlation coefficients (right), however the type with slope closest to 1 is LayerNorm (center), only overestimating the GNS by less than 40% across EMA alpha values."

On batch size scheduling

Batch Size Schedule: ... 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?

This is a great question! The batch size schedule we used linearly increased during training because we had observed this trend in the GNS and this is illustrated in Figure 3 of our attached pdf. We agree that an automatic online batch size schedule using GNS would be ideal but we didn't want to overcomplicate the paper. The schedule is supposed to mimic what a practitioner watching the GNS might do; in practice an operator adjusts the batch size manually for expensive runs.

Figure 8: This should probably be done using learning rate schedules of different lengths for statements like 2x to be meaningful.

We agree that varying hyperparameters would make this result more robust. At this scale, we observed that the batch size scheduled run was outperforming a tuned baseline, so we did not pursue this further

Thank you for your helpful review. Your suggestions will definitely help improve the paper. It is also gratifying to know that the work is well motivated and the application of per-example gradient norms to Transformers training is useful and clear.

Regarding revising the related work

perhaps a better way to represent this section would have been to segregate the related work with smaller sections with headers

We thank the reviewer for this suggestion. We will add \paragraph headers to the related work section to make it easier to navigate.

Regarding clarifying the contributions

Overall, I feel the main contributions of the paper is not well presented in this draft

We plan to address this suggestion through a new figure in Section 1 that highlights the unique contributions of our method. For example, using something like Rebuttal PDF Figure 1 and the figure caption, we will contrast our (layerwise) per-example norms with other methods that must aggregate gradients at a coarser level based on their specific data-parallel configuration. We will also do a better job of highlighting the specific contibutions in the introduction, previewing the experimental investiation that unfolds over the rest of the paper.

We will also do a better job of motivating the use of per-example gradient norm applications beyond GNS estimation (such as for differential privacy). We should have also made it more clear that GNS itself is rarely used because it traditionally requires large DDP setups to be useful. With our method, it can be applied both on the smallest MNIST experiment or extremely large scale language model training.

Regarding per-example norms in related work

only a few papers like Li et al. [29] and Goodfellow [22] are discussed … it would be nice to have … more related work to have a thorough pros and cons discussion.

This is a good suggestion, thanks. While related work in this area is limited, following your suggestion, we did discover "Efficient Per-Example Gradient Computations in Convolutional Neural Networks"(). This paper should have been mentioned in the related work section and we will add it in the final version. As backpack makes clear, this also reduces to a 3D tensor regime, so will be equivalent to our method. Li et al. [29] is the only method we know of that focuses directly on per-example norms in 3D tensor regimes, which motivated our comparison in Section 3 (lines 111-148). We agree that a discussion on pros and cons would be valuable, for example Goodfellow [22] is optimal in 2D tensor regimes, and we will add this in the final work. Thanks again!

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

We agree that a thorough discussion of this comparison would be valuable. At present the paper only provides a short comparison in Appendix A (line 389). Appendix A currently only compares the methods in the context of GNS estimation. To improve it, we will list pros and cons of all known per-example gradient norm estimation methods [22, 29, rochette2019efficient, gray2023efficient] directly in the final version and discuss their usage in GNS estimation.

Regarding notation and citations

Notation writing could be improved … 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 ?

We agree that the interpretation of $B$, $I$, and $K$ should be included. We will add notes that they correspond to the batch size, input dimension, and output dimension, respectively. Additionally, the $u$ index shares the same index space $u \in (1, ..., T)$ as $t$, which may cause confusion. We will clarify this in the final version. Good catch!

Section 2.1: it may be shown [32] -> It may be shown McCandlish et al. [32]. And you may remove McClandish elsewhere.

We thanks the reviewer for pointing this out, we mechanically applied \citet without considering how many times "McCandlish et al." would get printed. We will address this in the final version.

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.

We thank the reviewer for pointing out this gap. On line 50 $B_{big}$ and $B_{small}$ have an incomplete definition. In the final version we will explain that $B_{big}$ is typically the full batch size and $B_{small}$ is a fraction of that, typically the batch size executed on DDP nodes or during gradient accumulation.

Regarding Algorithm 1

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

We will expand the discussion in the paper to make the contribution of line 4 clear. We agree that the correction on line 4 of Algorithm 1 is only a consequence of backpropagation of losses that have been reduced by a mean; it is not strictly part of the per-example gradient norm estimation. Lines 123-128 describe why this correction exists and we found that it was often a source of error in new implementations.

The form of the correction in Algorithm 1 is because it is safer to apply the correction to the norm after taking the mean of the squared norm $s_w$ (it would be equivalent to sum the loss and then scale the update gradients, but this is more resource intensive and could introduce loss scale issues).

Thank you very much for your thoughtful feedback, and for your support of the paper's core idea and experimental approach.

Regarding larger-scale models

The largest model included in the case study only has 111M parameters... this work could benefit from experiments on larger architectures

Yes, we agree that demonstrating our findings on larger models would be beneficial. Following your suggestion, we have conducted experiments on a 1.3B parameter model. The results are in our attached pdf document in the response to all reviewers.

Regarding architectures beyond transformers

This study is limited to transformers. However, many other architectures do not use normalization sub-layers. This point is also acknowledged by the authors.

We agree that this is a limitation. However, we should have also mentioned in the paper that our per-example gradient estimation methods for other layers still apply (Algorithms 1 and 3). Although perhaps not as performant as gathering statistics from the LayerNorms -- depending on the kernel used -- the increase in runtime may be acceptable. In experiments we have run it was at worst 30% slower to gather per-example gradient norms for all layers. Gathering only LayerNorm per-example gradients did not slow down training at all. We will revise the paper to make this observation. Thanks!

Thank you for your thoughtful comments and suggestions, and for positively noting the efficiency of our GNS estimator.

Regarding the scope of the evaluations

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.

We agree that sufficient experimental evidence is an important consideration. Of course, in this paper we do not establish the relationship between GNS and critical batch size, as McCandlish et al [32] provided exhaustive experiments on this topic on many data and model types. The focus of our work is validating our novel unbiased estimator of the GNS statistics. To this end, following your suggestion, we will include additional studies for other model types, sizes, and datasets in the revised paper. We propose to expand on Figures 4-6 as follows:

• Figure 4 illustrates the separation of GNS measurements by layer type and Index; this demonstrates the granularity of the estimators in practice. We have performed the same experiment on many different model types and we will include those results in the camera version of the paper. We will also add additional plots for other model sizes to the Appendices.
• Figure 5 is a replication of prior work that opens questions about an experiment from McCandlish's Appendices. We will perform a replication of the original experiment on SVHN in order to check for methodological differences.
• Figure 6 relates the measurements of GNS across layer types. The results are presented for one model size (111M parameters). We observe this result for language models specifically because that's one large-scale setting where it may be useful. Additional studies for other model types would be valuable here and we will include other model types and datasets in the revised paper.

Additionally, the results were not compared to other batch-size scheduling methods.

Note that Figure 8 is presented purely as an example use case for GNS in large-scale language model training. The actual batch size schedule is not novel, and we do not claim to establish a SOTA batch size scheduler.

Regarding the generalization of the findings

The finding that the normalization layer effectively predicts total GNS behaviour in transformers has not been thoroughly tested on other transformers … the authors also do not provide any intuition as to why [the predictability] might be the case … could the author add more experiments on different models and datasets to validate the generalizability?

We agree that it would be interesting to know if this result generalizes to other model types, such as the original post-LayerNorm Transformer, or even to image models, and we will include such studies in the final paper.

Regarding why LayerNorm gradients predict total GNS behaviour, we are currently exploring whether the per-example layerwise gradients are correlated for the reason that they ultimately reflect variation in the per-example loss. While a full theoretical understanding of this phenomenon may prove elusive, we agree that some additional experiments to gather greater intuition would be valuable. Thank you for this suggestion.

Regarding the paper clarity

Some concepts could be better explained or motivated. For example, it's not clear how per-example gradient norms help estimate GNS... It's also challenging to extract key information from the figure captions.

We regret that Section 2 failed to explain the link between per-example gradient norms and GNS measurements. We suspect the issue may be line 80, where we need to clarify that $B_{big}$ is the full batch and $B_{small}$ are fractions of that same batch. We will revise section 2 to note that gradient norms are measured for both batch sizes, but typically this can be done cheaply when training DDP because the small batches gradients are on each node and the large batch gradient can be obtained after synchronisation before the update. Per-example gradient norms are intended to obtain a gradient norm for $B_{small}=1$ without requiring any specific training setup (ie DDP). Figure 1 shows that this would give the most accurate measurement of GNS.

We also agree that key information is difficult to extract from Figures. As described in our rebuttal to Reviewer 32m2 we will revise the Figure captions to include the important message of each Figure.

Regarding the role of gradient noise in optimizer performance

In section 2.3, it is claimed that the variance of gradient norms determines whether SGD outperforms adaptive optimizers. However, subsequent studies … have shown that the performance gap persists even in full-batch scenarios without stochasticity

Thank you for referring us to this paper! Indeed, Kunstner et al. (2023) does strongly suggest that gradient variance is likely not the determining factor for why Adam prevails over SGD in Transformer models. Our overall point was that our work can support such studies by providing tools to efficiently collect gradient statistics. However, as you mention, this particular area is perhaps now less well-motivated, and we will revise this part of the related work accordingly. Thanks again!

Thank you for addressing my concerns and conducting additional experiments; I will raise my score.