Understanding and Minimising Outlier Features in Transformer Training

We thank the reviewer for their time and effort in carefully reviewing our work. We are pleased that the review demonstrates a clear understanding of our contributions, and in particular how our work builds on existing work to relate OFEs to areas such as signal Propagation and entropy collapse. We are also grateful that the reviewer has spent time digging deep into our paper’s appendix, beyond the main 9 pages.

Overall the review is extremely positive, and we are happy that our work was well received by Reviewer cUUo. Indeed, we share the belief that “this work will be appreciated by multiple communities for its wealth of insights” and that our insights will “have effects on transformer training procedures going forward.”

Weaknesses We will make sure to use an extra page to include more detail of our experimental setup, and thank the reviewer for spotting the typos, which we will amend. Regarding scale, we agree with the reviewer that scaling beyond 1.2B is “unnecessary but would further increase the significance of this work”. However, to address scale (which is also mentioned as a weakness by R7V85) in Figures 1* and 2* of the additional pdf page we present loss and kurtosis curves of LLaMa-style transformers trained at 7B parameters for 6B tokens. We did not have enough compute to perform hyperparameter tuning, but we still see that the OP block closely matches the performance of Pre-Norm, and is vastly better in terms of kurtosis. This mirrors our findings at smaller scale in the submission. We provide more context and details of the experimental setup in our response to R7V85.

Given the thoughtful nature of the review, we would like to use the remainder of our rebuttal to engage in further discussion with the reviewer on some topics/questions/experiments that may be of interest:

Quantisation experiments In Table 1* of the additional pdf page, we go back to our original motivation in studying OFs and assess how our various architecture and optimisation suggestions affect quantisation. We take the OPT-125m setting of Bondarenko et al 2023 [14] and vary architecture and optimiser. For architecture, we compare OP to Pre-LN and also the Gated Attention baseline of [14], and for optimiser choices we take those identified in Section 5 (increased LR, and increased Adam epsilon), as well as removing dropout regularisation present in [14]. After training on BookCorpus+Wikipedia in mixed FP16/32 we quantise the models to W8A8 (int8 weights and activations) and assess the quantisation error (in perplexity lost). We also present the average kurtosis across layers in Table 1*.

In Table 1* we see that our kurtosis metric for OFs is indeed highly correlated with quantisation error across different architectures and optimiser choices. For example, Pre-LN has consistently high kurtosis, and consistently catastrophic quantisation error. Moreover, the OP block has consistently low kurtosis across different optimisation hyperparameters, and also low quantisation error. Finally, the quantised Gated Attention baseline is better than Pre-LN but struggles with aggressive hyperparameters that improve FP16/32 performance but at the expense of OFs (like large LRs as identified in Section 5, and removed dropout reg). However, both kurtosis and quantisation are improved in Gated Attention when increasing Adam epsilon, as identified in section 5. We believe this experiment ties together our findings and motivation throughout the paper. We present a more detailed discussion of our quantisation results in the global response.

Underselling in title/abstract We thank the reviewer for their kind and positive words regarding underselling. Throughout the paper, we were very careful to avoid overclaiming in our wording. This is reflected also in our choice of title and abstract. The reviewer is likely correct that we may have undersold in the process.

Stated limitations We are glad that the reviewer appreciates our stated limitations, and sees the strengths of our submission based on their own merit. We chose to include a thorough limitations section in the spirit of honest research and progressing the field in terms of future directions. While we stand by this decision, we also accept that including our limitations may have come to our own detriment in the review process.

Disparity across reviews We note that there is an unusually high variance in the reviews for our work. In particular, many of the criticisms of Reviewer Lb1w are in direct opposition of Reviewer cUUo. We find such criticisms to be unfounded, and we have used several quotes from RcUUo’s review in our rebuttal to Reviewer Lb1w.

We hope reviewer cUUo will continue to support our work throughout the discussion period, and thank them once again for their careful review. We also welcome any additional feedback.

AC, I think this paper should be strongly considered for an oral presentation, and it clearly deserves at least a spotlight award.

All, I have read the other reviewers' comments and would like to maintain my high score. Based on the rebuttal, I increased my confidence from 3 to 4. I would also like to clarify that, to the best of my recollection, this is the first time in a couple of years that I have given a score higher than 8 to an ICML, ICLR, or NeurIPS submission.

I had a similarly high opinion of this submission before I saw the rebuttal. The rebuttal then showed that the original submission's findings hold at larger (7B) model sizes, and that the original submission's insights can (a) explain failures of quantization in a new set of experiments, and (b) provide hyperparameter/architecture tuning guidance that mitigates these failures.

The reviewer is likely correct that we may have undersold in the process.

Maybe you didn't "undersell", but you actually delivered on what you described in the title and abstract much better than expected.

In Table 1* of the additional pdf page, we go back to our original motivation in studying OFs... We believe this experiment ties together our findings and motivation throughout the paper.

I agree completely. It will make an excellent addition to the final version of the paper, tying everything together while corroborating key findings.

We thank the reviewer for their constructive review. We are pleased the reviewer feels we “made a sufficiently large contribution to the research of OFs”, appreciates our experimental design, and finds our paper “well written and relatively easy to understand considering its topic”. We address the raised concerns in order.

Background knowledge Thanks for the suggestion. We assume familiarity with the terms “post-norm” and “pre-norm” as they are ubiquitous to transformer design, with relevant citations on lines 113-114. However, following the suggestion we will add more discussion to formalise these terms. For entropy regularisation we included a thorough discussion in appendix A, but will formalise the notation. For Signal Prop we believe there is sufficient relevant background and formalism in appendix A, and would be curious to know if the reviewer agrees.

Scale Scale is discussed by both Reviewers 7V85 and cUUo. We accept the shared point of Reviewers 7V85 and cUUo that larger scales would “strengthen the message” or “further increase the significance” of our work. However, we disagree with R7V85 that scale is a ‘significant concern’, and agree with RcUUo that scaling further is “unnecessary” for the paper. We are at an academic lab, and our results at 1.2B are prohibitive for most in academia, so we disagree that the scales in our paper are “relatively small”.

Moreover, the OP block itself is not the main contribution of this work. As on lines 164-171 and 346-348, as well as our title, our main focus is to understand and minimise Outlier Features, and we introduce the OP block to test the impact of normalisation layers on OFE. We believe we present an extensive study on this question, and that the scales we test at are sufficient to validate our other findings on OFs e.g. the links between OFs and signal propagation or optimisation choices.

Having said that, we managed, with much difficulty, to access compute to train 7B parameter models during the author response period. Due to short notice, we spent the weekend preparing/running this experiment in a setup based on HF nanotron, with tensor+data parallelism. As such, we did not tune hyperparameters, and our hyperparameters are from the default Pre-Norm mode (see below). Moreover, we use fixed residual gains of beta=0.05 and did not implement trainable betas in the OP, which as discussed in our response to Reviewer Yyp1 would give a small improvement in perplexity.

Despite this, in Figure 1* the OP block closely matches the Pre-Norm performance in loss over steps with a minimal gap at 7B scale, which we believe would be closed with hyperparameter tuning. The plot for kurtosis, averaged across layers, is Figure 2*. We see at 7B scale the OP block also has significantly lower kurtosis than Pre-LN, mirroring smaller scales. We implemented the OP block before the kurtosis logging, and when rerunning for logging our compute allowance unfortunately ran out. As such, Figure 2* stops after around 2B (of 6B tokens). Despite this, we clearly see in this case the OP block’s kurtosis has stabilised (around 8) whereas the Pre-Norm block’s kurtosis is still 450+ and increasing.

7B hyperparameters: These autoregressive language models have depth 32, and hidden width 4096, like the LLaMa-7B models. Standard LLaMa design choices (AdamW betas, RoPE, RMSNorm, cosine decay etc) are used apart from a standard GeLU MLP (instead of gated MLP), and we use max LR=3e-4 like in LLaMa-7B. We trained on the FineWeb-Edu dataset for 20K steps with batch size 72 and context length 4096, giving ~6B tokens.

We hope this experiment alleviates scaling concerns. As we wrote in lines 348-349, we had no reason to suspect the OP block would struggle at bigger scales because the OP block is designed with scaling principles like signal propagation in mind. We maintain this position. We will endeavour to empirically scale to further parameter and token counts, and also investigate other settings like ViTs beyond the language settings that are typically studied in the OF literature, in future work.

Section 5 Thanks for this point. Section 5 is very important because OFs only arise during training. Thus, to understand OFs one must consider optimisation choices. This leads to several interesting findings as the Reviewer states, particularly the effect of adaptivity.

The brevity of Section 5 is due to the page limit. We point the Reviewer to Appendix F which discusses why the different optimisation suggestions we propose (smaller LR and bigger adam epsilon) result in reduced OFs, by breaking down the updates to the kurtosis into different moment statistics. Reviewer cUUo even suggested moving the discussion in Appendix F to the main paper.

Besides the results in the submission, we present in Table 1* a new experiment that studies the quantisation effect of our proposed optimisation and arch choices, as well as their effect on kurtosis. We take the setting of Bondarenko et al. [14] on OPT-125m, which trains on BookCorpus+Wikipedia and quantises post-training from mixed FP16/32 to W8A8 (weight-and-activation int8). We see that our optimisation interventions (bigger LR/Adam epsilon) have the same effect on kurtosis as found in Section 5, across different architectures, which further reinforces Section 5.

Table 1* also shows kurtosis is highly correlated to quantisation error e.g. our OP block has lowest kurtosis and also lowest quantisation error, across optimiser choices. A complete discussion of Table 1* is in the global response. We believe this experiment brings to full circle our paper's narrative for studying the roles of different architecture/optimisation choices on OF, in which Section 5 plays a crucial role.

We believe we have addressed the reviewer's concerns thoroughly in our response and would appreciate reconsideration of the score if so. We thank the reviewer again for their constructive feedback and would welcome any additional feedback.

I would like to thank the authors for the rebuttal.

Background knowledge Not too surprising given its title, Appendix A is written like a related work section rather than a background section. I think the paper would be more accessible if the most important concepts were properly defined in a unified mathematical notation and with the most important equations not being inlined such that they are easy to find. This does not have to include all concepts cited in Appendix A, but I think methods that are used in the paper, such as QK-Norm should be defined in formal mathematical notation. After reading Section 3 again, I also noticed that the OP Block is only described as a diagram, and via the 3 changes to the Pre-Norm block. I think the paper/appendix should include a proper mathematical definition of the forward pass of the OP block since I find it difficult to think about this only in terms of changes to a pre-norm block that is not even defined anywhere in the current version of the paper. The diagram helps but I believe mathematical notation is the best way to express this and could help people that try to reimplement the OP block on their own.

Scale I disagree with the authors that scale is not a concern when it comes to LARGE language models. That being said, I understand the issues with running large-scale experiments in academia and I understand that the analysis of OFE is the main contribution of this paper. I also want to thank the authors for running these additional experiments and I think that demonstrating that OP norm works well for a 7B model and significantly reduces kurtosis is a strong contribution since this is the standard size for small-scale LLMs that are being used in practice. That being said, I would encourage the authors to train at least one ViT for the camera-ready version since a ViT-Bor a ViT-L are much cheaper to train than LLMs and to my knowledge, the phenomenon of OFEs has been mostly studied in the language literature.

Section 5 I agree with the other reviewer that extending Section 5 in the main paper with some of the results from the Appendix and the rebuttal would be a good idea. I also liked the quantisation experiments here in particular.

In conclusion, I would recommend acceptance given the author's detailed rebuttal.

We thank the reviewer for their time and effort in reviewing our work. We are pleased that the reviewer writes that our “findings are supported by empirical validation across various datasets, demonstrating the effectiveness of the proposed metrics and strategies”. The reviewer’s concerns largely centre around the beta residual scaling gain parameter, in the context of our OP block introduced in Section 3. We address the raised concerns in order:

Beta before residual connection not after The reason for beta before the residual is that as per Signal Propagation theory we care about reducing the relative weight of the residual branch relative to the skip branch in order to reproduce the initialisation benefits of pre-normalisation. We outline relevant citations for this fact in lines 127-129, and perhaps the single most relevant citation for this is De and Smith 2020 (which is citation [45] in the submission). Once the skip and residual branches have been summed, it is not possible to change their relative weightings.

In the inputs to the first block, e.g. output of the embedding layers in Transformers, we do not tend to observe the most extreme OFE, as can be seen across different settings in Figures 13, 18, 21, 23, 25, 27, 30, 32 (note the lack of log scale in Figures 21,23). Intuitively, this makes sense because preceeding the input to the first block there is only a single linear matrix of parameters (the embedding matrix), compared to other layers where there are the trainable weights in the non-linear attention and MLP sub-blocks.

Trainable beta or not We are not the first to propose the residual scaling idea to enable scaling to large depths, so we based our experiments off existing practice. Many have proposed trainable beta initialised to a small value, like SkipInit (De and Smith 2020, [34]), NF-Nets (Brock et al 2021, [47]), Wortsman et al 2023 [10]. From Signal Propagation theory, Hayou et al 2021 [22] and Noci et al 2022 [25] show that $\beta=O(1/\sqrt{\text{depth}})$ is needed for a well-behaved infinite depth limit at initialisation, so the question of trainable beta is outside the current theoretical scope of the literature.

Empirically, we find that trainable betas lead to slight improvements in perplexity/loss, but stress that they are not essential for the OP block to be performant and scalable, as seen in our 7B scale experiment in Figure 1* of the additional pdf page. We provide a plot of the evolution of betas in our 1.2B Languini experiments in Figure 3* of the additional pdf page, where we see the trainable betas evolve during training but stay roughly in the range of their initialisation 0.1 for both OP and Pre-LN blocks. The slight improvement with trainable betas intuitively makes sense, because beta can be seen as scaling the learning rate for parameters on the residual branch (e.g. https://arxiv.org/abs/2205.15242 or He and Hofmann [28]), and so trainable beta amounts to block-wise adaptive LRs. We think analysing the training dynamics of beta (or proposing modifications to ensure it remains $O(1/\sqrt{\text{depth}})$) is an interesting direction for future work.

Increasing kurtosis with tokens seen We address the point on kurtosis increasing (to comparatively low values) with OP block in Section 4 (particularly on lines 236-245). There we show that kurtosis can be seen to closely track with signal propagation dynamics, which we attribute to the model creating structure in its hidden layer representation space to fit the task at hand. Identifying this relationship between signal propagation and OFs is one of our most significant contributions, as noted by Reviewers 7V85 and cUUo, in addition to the OP block.

For increased “tokens seen”, in our new quantisation experiment, Table 1*, we take the OPT-125m setting of Bondarenko et al. [14], which at 12B is more than double the number of tokens seen in Figure 4. Even with this increased training length, we again see that the OP-block has significantly lower kurtosis, across various training hyperparameters, than not only the Pre-LN block but also the Gated Attention method of [14], which was designed to reduce OFs. Moreover, we note that this reduced kurtosis directly translates to reduced quantisation errors with Post-training quantisation from mixed FP16/FP32 to int8 weights-and-activations, which motivated our study of OFs in the first place.

We believe we have addressed the reviewer's concerns in our response and would appreciate reconsideration of the score if so. We thank the reviewer again for their review and would welcome any additional feedback.

We thank Reviewer Lb1W for their time. However, the reviewer has made several assertive yet unsubstantiated criticisms of our work, which oppose all the other reviewers. We are concerned by these criticisms as they lack constructive feedback and could be perceived as overly critical from our view. All other reviewers gave scores leaning towards accept, with RcUUo giving the maximum score 10, in contrast to RLb1W’s score 3. We hope our rebuttal will address these points effectively and convince RLb1W of our paper’s merit, hopefully leading to an improved score.

Writing Quality: RLb1W claims the paper is “poorly written” and “overall lacks fluidity and coherence across sections”. This is in stark contrast to R7V85 and RcUUo who praise the paper for being “well written”, and “extremely well written”. RLb1W provides no evidence for this purported poor writing, besides line 59 where the verb form of the noun “matrix multiplication” is used. As such, this criticism is unsubstantiated. With regards to fluidity/coherence, we gave significant thought to the paper’s structure/flow, as seen via the key takeaway boxes between sections. The effect of this, quoting R7V85, is: “the paper is relatively easy to understand considering its topic”.

Insufficient Evidence for OP Block: We contrast RLb1W’s concern with RYyp1 and RcUUo, who write: “findings are supported by empirical validation across various datasets, demonstrating the effectiveness of the proposed metrics and strategies”, and “this is a thorough and careful study, making claims that are well supported by the experiments”.

More concretely, Fig 4 plots all different layers (the legend samples layers to save space) which addresses weakness 2a and we discuss later layers and their connection to signal propagation in lines 236 to 245 of the submission. Besides Fig 4, we show the effectiveness of the OP block for OF reduction in five other Figs (2, 13, 14, 32, and 38) which constitute extensive ablations across datasets, entropy regulation mechanism, centring in kurtosis metric, and the effect of norm in MLP. This answers weakness 2b regarding training schemes, and so we view weaknesses 2a+b to be unfounded. We respond to weakness 2c-d (and provide additional evidence for 2b) in the next point.

Also, in Figures 1* and 2* of the additional pdf page we show the effectiveness of the OP block at 7B scale.

Quantisation Experiments In response to weakness 2c-2d, we present new quantisation results in Table 1* of the additional pdf page, which compares OP to Pre-LN and the best baseline (Gated Attention) from Bondarenko et al [14]. To our knowledge [14] is the only existing paper that has studied architecture and OFs. We test these models across training schemes, and find the OP block has significantly lower kurtosis across different optimiser choices. This lower kurtosis directly translates to improved quantisation, as seen in the significantly smaller drops in performance when going from mixed FP16/32 to W8A8 (weight and activation in int8) with OP, compared to Pre-LN and Gated Attention. We provide a complete discussion of the results of Table 1* in the global response.

Lack of novelty of relation between OF and Signal Propagation: To us, this is again unfounded: according to RcUUo, we "make strides towards understanding OFE by combining ideas from Signal Prop and entropy collapse (e.g.), in ways that are insightful, novel, and excellently motivated.” We take it as a strength not weakness that RLb1W sees the link between OF and Signal Prop unsurprising, as it implies the significant effort taken to make the paper accessible was successful in Section 4 (going back to the point on writing quality). This criticism also contradicts R7V85: who “liked the connection between signal prop and OFs in Section 4. The section was formal enough while still being easy to understand." We add that a result that is unsurprising in hindsight is not necessarily unsurprising.

Lack of depth in training design analysis: We disagree: due to the page limit, we could not fit all our optimiser choices analysis into Section 5, but this does not mean our analysis lacks depth. As also discussed in line 339, we refer the reviewer to Appendix F, where we dive deeper into the reasons why our proposed optimisation changes lead to reduced OFs, by breaking down the kurtosis updates into different moment statistics. Reviewer cUUo read Appendix F and suggested we include it in the main paper.

Moreover, our claims regarding the importance of large adaptive learning rates for OFs are not “rough findings”, but instead substantiated across multiple architectures and datasets. Namely, Figures 7, 24 and 25 show the effect of large learning rates, and our claims regarding adaptivity with Adam epsilon are supported by Figures 8, 27, 30. As such, we again fail to see the basis of this criticism.

In addition to the existing depth of analysis of optimiser choices in the submission, we refer to our new quantisation experiment in Table 1*, which provides further evidence of the effect of optimiser choices we identify in Section 5 in terms of OFE, on a new dataset and across multiple architectures.

More like a technical report than a conference paper ready for publication: We are completely surprised by this assertion. The quote from lines 353-354 is misquoted without context, and thus misleading. The missing context clearly refers to theoretical understanding; we give many empirical insights into signal propagation training dynamics in Section 4. We leave RLb1W with quotes from R7V85 and RcUUo, who write: “I feel like the authors made a sufficiently large contribution to the research of OFs” and “this work will be appreciated by multiple communities for its wealth of insights”.

We believe we have addressed the reviewer's concerns thoroughly in our response and would appreciate reconsideration of the score if so. We would also welcome any additional feedback.