Understanding and Minimising Outlier Features in Transformer Training
Authors: Bobby He
Lorenzo Noci
Daniele Paliotta
Imanol Schlag
Thomas Hofmann
Venue: NeurIPS 2024 poster (2024)
Keywords: Outlier Features
Training Dynamics
Normalisation Layers
Signal Propagation
NN Optimisers
NN Architectures
Transformers
Understanding Deep Learning
Primary area: deep_learning_architectures
TL;DR: We study why Outlier Features occur during standard NN training dynamics, and which algorithmic choices one can make to minimise them.

Abstract: Outlier Features (OFs) are neurons whose activation magnitudes significantly exceed the average over a neural network's (NN) width. They are well known to emerge during standard transformer training and have the undesirable effect of hindering quantisation in afflicted models. Despite their practical importance, little is known behind why OFs emerge during training, nor how one can minimise them.

Our work focuses on the above questions, first identifying several quantitative metrics, such as the kurtosis over neuron activation norms, to measure OFs. With these metrics, we study how architectural and optimisation choices influence OFs, and provide practical insights to minimise OFs during training. As highlights, we introduce a novel unnormalised transformer block, the Outlier Protected block, and present a previously unknown benefit of non-diagonal preconditioning optimisers, finding both approaches to significantly reduce OFs and improve quantisation without compromising convergence speed, at scales of up to 7B parameters. Notably, our combination of OP block and non-diagonal preconditioner (SOAP) achieves 14.87 weight-and-activation int8 perplexity (from 14.71 in standard precision), compared to 63.4 int8 perplexity (from 16.00) with a default OF-prone combination of Pre-Norm model and Adam, when quantising OPT-125m models post-training.

Understanding and Minimising Outlier Features in Transformer Training
Bobby Hel+ Lorenzo Noci
Daniele Paliotta² Imanol Schlag
Thomas Hofmann'
Department of Computer Science, ETH Zürich 2 Machine Learning Group, University of Geneva
Abstract
Outlier Features (OFs) are neurons whose activation magnitudes significantly exceed the average over neural network's (NN) width. They well known emerge during standard transformer training and have the undesirable effect hindering quantisation in afflicted models. Despite their practical importance, little is known behind why OFs emerge during training, nor how one can minimise them.
Our work focuses on the above questions, first identifying several quantitative metrics, such as the kurtosis over neuron activation norms, measure OFs. With these metrics, we study how architectural and optimisation choices influence OFs, and provide practical insights to minimise OFs during training. As highlights, we introduce novel unnormalised transformer block, the Outler Protected block, and present a previously unknown benefit of non-diagonal preconditioning optimisers, finding both approaches significantly reduce OFs and improve quantisation with- out compromising convergence speed, at scales of up to 7B parameters. Notably, our combination of OP block and non-diagonal preconditioner (SOAP) achieves 14.87 biigh--and-actvatiion int8 perplexity (from 14.71 standard precision), com- pared to 63.4 int8 perplexity (from 16.00) with default OF OF-prone combination of Pre-Norm model and Adam, when quantising OPT-125m models post- post-training
Introduction
Despite their widespread use, our understanding deep neural networks (NNs) and their training dy- namics is very much incomplete. This, in part, reflects the complexity of traversing high-dimensional non-convex loss landscapes is also symptomatic of the myriad design choices, such as NN architecture and optimiser typerparameters, that practitioner must take before training. While standard choices architecture and optimiser exist, is often unclear how these choices affect model performance or the emergence of various empirically observed phenomena during NN training.
Outlier Features (OF) are one such training phenomenon. OFs are neurons whose activation magni- tudes are significantly larger than average in the same layer, i.e. across NN width [13]. They have been widely observed in the popular transformer NN architecture 4 6 we verify in Fig I and
are of practical interest because their existence hinders quantisation 3 7-12. In particular, OFs cause large dynamic ranges in activations across width, leading to high quantisation errors low precision matrix multiplications. As such, Outlier Feature Emergence (OFE) during training hinders low-precision training and inference, and minimising OFE could yield significant efficiency gains.
In this paper, we tackle OFE from two related angles: by (1) proposing interventions to minimise OFE without affecting model convergence or training stability, using insights motivated through (2) enhancing our understanding of why OFs appear during training. We argue that it is important to first understand why OFs appear during standard NN training dynamies in order identify which design
"Correspondence to bobby.he@inff TET AI Center
wwwinfffff..... ethz.ch
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Pythia 160M
Pythia 1.4B
Pythia 699
Activation Kurtosis
1.39
127
Layer
20k
4Ck
2024
Step
Step
Step
Figure 1: Outlier Features appear in open-source transformers [66 during training, as measured by our Kurtosis metric Eq (T. Our work investigates the design choices that influence their emergence.
choices influence OFE, and how. Though progress has been made 1 behind OFE remain largely unknown.
L3
14
T0
15
the mechanisms
Alongside the practical motivation of model quantisation, we believe understanding OFs, and their causes during training, to be an interesting research question for several reasons. The emergence of Outlier Features in standard transformer training regimes raises the question if OFs are simply an artifact of certain design choices or a fundamental property of transformer training, essential for best performance. Understanding (the causes of) OFE better helps to better understand NN training dynamics in general, and the roles played by different design choices. Moreover, we shed light the differences between models at initialisation compared during, after, training. While NN initialisation is more commonly studied owing to analytic tractability [17)30], understanding trained NNs is arguably more important they exhibit rich feature learning behaviour (31.34, like OFs, that arise during training. In the case of OFs, this has potentially wide-reaching implications, including for NN interpretability, which often focuses on the roles of individual neurons 35
Our contributions Overall, we show that OFE can mitigated relative to standard practices, and highlight key design choices to do so. We start by introducing OFs, and in particular quantitative metrics to measure OFs in Sec In SecB, we study the role of normalisation layers for OFE, and find that existing hypotheses do not fully capture the OF phenomenon. We proceed to show that removing normalisation through our Outlier Protected transformer block minimises OFs, without loss of convergence speed training stability compared to standard transformer blocks. In Sec 4 consolidate our findings by identifying signal propagation as an important object that can predict OFs during training, and that choices that improve signal propagation during training also minimise OFE. In Sec 5) we consider optimisation hyperparameters, and highlight the importance of large diagonal adaptive learning rates for OFE. Finally, in Sec we demonstrate the performance of our proposals to minimise OFs both at larger scales, up to 7B parameters, and in terms of improved quantisation performance. In the interests of space, in App A.3 we discuss additional related work.
2 Problem Setting
Consider an activation matrix X Raxd obtained from some neural network layer, where n the number of batch inputs/sequence positions, and d is the number of neurons across NN width. In typical NN layer, we matrix multiply X by weight matrix W dddd to give XW e Rexd with (a.j).th element: Ed= X., W k. This fundamental operation is central to NN computation and can be seen sum over d terms, one for each neuron.
Several works have established that if the magnitudes of the summands {Xa.k Wk. 1 have large variations, then it becomes difficult to compute their sum in low precision, thereby precluding potential efficiency gains from "vector- wise" quantised training or inference (though significant progress has These works have shown that (pre-) )trained transformer (37 models possess such a deficiency, which is attributed to the existence of Outlier Features (OFs) whose activations are much larger in magnitude compared to the other d - 1 neurons.
been made on the latter, 8 36

---

Measuring OFs Existing works have measured OFs in architecture-speciie ways or using n.j We argue that measuring OFs should be independent of architecture/actiiation scale: barring exploding/vanishing activation scales, the relative difference in summands is what causes issues for vector-wise quantisation. We use two metrics to measure OFs:
activation scales | del
Easn.jaa X
1. Kurtosis of neuron activation RMS: Let s R such that s) = =V - be the
vector of root mean mean-squared activations across inputs Then, let Kurt (X) be the ratio of the fourth moment m4 to the squared second moment m2 over the empirical distribution of
Kurt(X] (X)
M4
def
(I)
m2
( (3E)=1s)) s))
We that min(Kurt(s)) when all sj equal and outlier features exist, and max (Kurt(X)) = d. which is the limit when neurons have activation magnitudes dominated by single outlier feature.
2. Max-Median Ratio (across neurons): A metric for OFs more aligned with the original motivation of studying variation in summand magnitudes. Specifically, we compute:
max, Xa.il
MMR(X) def Aggregate,
median, X
in words, the neuron divided by the median absolute neuron, aggregated in some permutation invariant way across inputs. We typically use the mean to aggregate over inputs, but could also take e.g. median or max. MMR takes a minimum value 1 when all activations are identical in magnitude, and is unbounded when a dominant outlier feature exists.
Variants of Kurt(X) have previously been proposed [14 3. but our formulation in Eq ① aggregates activations over inputs first, which allows us to link OFs and signal propagation in Sec Though
we focus our analysis on Kurt(X), Figs 11 13 and 17 show that both our OF metrics are highly correlated. In this work, we measure OFs on the residual stream across different lavers/blocks. For example, for Pre-Norm Post-Norm models this is Xour in the notation of App A.1
Experimental Setup Throughout this work, we train transformers on the next-token language
modelling task, and study OFs, range of datasets, including: 1) CodeParrot 5
2) Languini Books F and 4) FineWeb-Edu 41] Unless stated otherwise
[39]. 3) BookCorpus 40 and English Wikipedia our experimental results are conducted on CodeParrot, but importantly our conclusions regarding
and
OFs are consistent throughout across language modelling datasets. In App E. we also explore OFs in image classification settings with other architectures like Vision Transformers
[2 and MLPs.
In terms of architecture and optimiser, our default choices are the Pre-Norm transformer A. and AdamW [33 respectively, which are known to be prone to OFs, e.g. Fig 1 Our default architecture scale has width d 768 and 6 layers, giving around 130M parameters, but we demonstrate our findings continue to hold at larger scales (up to 7B parameters) in Secs 3 and 6 In Secs 3 and we examine alternatives the Pre-Norm architecture and their effects on OFs, keeping Adam W as optimiser, while from See 3 onwards we additionally consider modifications to Adam W. Further experimental details and results beyond the main paper can be found in Apps D and respectively.
3 Normalisation Layers and Outlier Features
Several works have highlighted the architectural choice of Layer Normalisation (LN) 44 as a cause of OFE 15 LN belongs to a family of normalisation (Norm) layers commonly used in sequence models, which normalise a representation vector x E Rd across the width dimension independently
We do centre X in s, for simplicity. Fig shows centring makes no qualitative difference for Kurt(X). As MMR(X) is invariant normalisation layers like RMSNorm without trainable parameters 38 there is some redundancy here in where exactly we take OF measurements.
train
using the same setup as 15

---

for different sequence positions. In general, for a centring scalar E (0, a Norm layer maps to:
Norm (2) -c((z) ((x)
1 2 Z dx))=
1 d - (z, (z, cx(x))
B p(x) -
(3)
LN when c= 1, with a trainable scale y and bias /3 vectors initialised all 1s and Os respectively.
Previous works have attributed OFE to the 7. parameters of LN incurring outliers during training 717. It is therefore natural to ask if simpler Norms with different formulations of Eq remove OFF. In particular, Root Mean Square Normalisation (RMSNorm) 45 is a commonly used Norm known be as performant as LN in Transformer training 46 47 Compared LN, RMSNorm fixes the bias / = O and removes the centring by setting c = O, which highlights that centring is not crucial operation in modern sequence modelling practices. One step further would be remove trainable parameters entirely by fixing y = 1, thus simply projecting æ to the hypersphere of norm vd. This is dubbed Simple RMSNorm (SRMSNorm) by Qin et al. (38], who find that SRMSNorm has minimal performance degradation but more computationally efficient than LN and RMSNorm.
We compare these different Norms in Fig 2 where we see that independent of Norm choice, all Pre-Norn transform- ers incur OFE: the peak kurtosis during training across Norms is over 4 orders of magnitude larger than initialisa-
Activation Kurtosis
104
101
tion. We also show OFE not only in Pre-Norm 48 49 but
10>
Posi-Lh
also Post-Norm (37 blocks (more details on transformer blocks in App[A.I) highlighting OFE occurs independent of where Norms are placed. In this experiment, the Pre- SRMSNorm model has highest Kurtosis, despite its lack of trainable Norm weights.
10-3
Fre-SaM5
Pro. LH
1012
DP (ours)
20K
4.6K
60c
Training Step
Figure 2: Kurtosis becomes large (i.e. OFE) when training with different Norms 130M scale. We plot the residual stream entering the 2nd of 6 blocks. Other layers in Fig 14
Having established that removing trainable weights in Norms still results in OFE, the next question ask is: how does removing standard Norms entirely influence Outlier Feature emergence?
Recovering training benefits in unnormalised Transformers This is challenging question to answer, not least because comparing OFE in architectures that converge at different speeds may not be a fair comparison: Norms are well known to be an important component in most NN architectures, providing various benefits for initialisation, convergence speed, and training stability. Thus, answer the above question, we must first review different hypotheses for the benefits of Norms in transformer training dynamics in order to motivate novel transformer block that matches the Pre-Norm block in convergence speed, while eschewing standard Norm layers.
OP (ours)
28 30 have observed that the initialisation benefits of Pre-Norm architectures can be recovered in unnormalised residual models using downweighted residual branches, through a theory known as Signal Propagation (Signal Prop) [17 18 (6). Notably, Brock al. [53] achieve state of the art performance on the ImageNet benchmark using unnormalised convolutional architectures. However, it has been observed that fixing Signal Prop at initialisation is not sufficient to fully capture the benefits of Norms for training dynamics in unnormalised transformers 288 30), which implies that Norms have training benefits specific to the self-attention based transformer model.
Several works 550555 23
27
Non
4B
Proj
H
Attention
EniReg
At the same time, Zhai et al. [57] show Entropy Collapse, where the Stochastic attention matrix has rows with low entropy and each sequence position attends to only one position instead of many, be a key transformer training instability (see Eq (TO). Entropy collapse occurs because large attention logits saturate the softmax, and several Entropy Regulation (EntReg) mechanisms have been proposed to control the attention logits and thus prevent entropy collapse. Existing entropy regulating methods include QK-Norm (58 59 lanh thresholding (Grok-1), o oReparam (57) and clamping the QK logits (DBRX). In standard Pre/Post-Norm attention blocks, Norm layer appears before Query and Key weights and implicitly regulates attention entropy, to extent.
The Outlier Protected Transformer Block. We remove Pre- Norms and replace them with an Entropy Regulation mechanism prevent entropy collapse, as well as downscaling residuals with / < 1.
Figure 3:

---

Early layers
Middle layers
Final layers
Activation Kurtosis
1.2
137
134
16-3
Pre- 16
0.50
1.50 2.00 2.50
3.00
3.50
4.00
0.20
2.00
1.50 2.00
3,00 353 4.00
0.53
3.00
1.20 2:30
2.00
4.20
Tokens seen
Tokens seen
Tokens seen
Figure 4: Our OP block mitigates OFE. We plot activation kurtosis of the residual stream across layers. Experiments are 1.2B scale on Languini Books using a max AdamW learning rate of 0.001 with linear warmup for the first 1.5% steps and linear decay thereafter. Notice the shared log-scaled y-axis: activation kurtosis is consistently (up 4 orders of magnitude) lower in OP block, particularly in earlier layers. Also, peak kurtosis during training is always higher in Pre-LN. The OP model also removes the final LN before unembedding; the effect of the final LN on OFE is shown in Fig 10
Our key insight is to combine ideas from Signal Propagation and Entropy Collapse prevention to remove Normalisation layers while keeping their training benefits. This brings us to our Outlier Protected (OP) Block, F which replaces the Pre-Norm block by removing its normalisation layers in both Attention and MLP sub-blocks, and making three additional changes: 1) downweighting residual branches with some
/depth < 1 to recover Signal Prop benefits of Pre-Norms (51,23 27. 30]. 2) adding an Entropy Regulation mechanism to prevent Entropy Collapse; we mainly use QK-Norm as it relatively simple and performed well in all of our settings, but present experiments with tanh in App[E.2) and 3) (optionally) scaling the inputs before the MLP nonlinearity by a scalar a ensure the nonlinearity inputs are of order 1, as derived by Brock et al. (53) using straightforward Signal Prop arguments. App B presents a mathematical description of the OP block.
In T show that our Outlier Protected block matches the standard Pre-LN block in terms of convergence speed at scales up to 1.2B pa- rameters when trained with next token prediction on the Languini books dataset (39 for nearly 4.5B tokens In App E.2 we ablate our OP block and show that the lack of entropy regulation mechanism without normalisation layers causes training instabilities. This demonstrates that preventing entropy collapse is necessary to match training stability and convergence speed in unnormalised Transformers.
Table 1: OP matches Pre- LN performance at scales up 1.2B params, on Languini Books [39]
Params
Block
Eval PPL
100M
Pre-LN
19.1
OP
18.9
320M
Pre-LN
16.2
We note that independent of OFs, the OP block is interesting in its own right because it shows that the initialisation-time Signal Prop and Entropy Collapse benefits of Norms in Transformers can be disentangled, and also reveals what was missing in previous methods that used Signal Prop arguments correct initialisation defects in simplified unnormalised Transformers However, we now focus on the benefits of the Outlier Protected block in reducing outlier features.
OP
16.2
1.2B Pre-LN
13.9
OP
13.9
28 30
Removing Norms mitigates Outlier Features In Fig 2 we see that the Outlier Protected (OP) block greatly reduces OFE compared to standard blocks. Fig 4 presents the corresponding plots in our 1.2B parameter experiments using our kurtosis metric, for different layers. We draw several consistent conclusions: 1) peak kurtosis across the course of training is consistently higher in Pre-LN, sometimes by over 2 orders of magnitude, across different layers; 2) kurtosis across training is usually higher in Pre-LN (up 4 orders of magnitude here), especially at early training times and in earlier layers; 3) OFE (measured via our metrics) does not need to be monotonic in training time. Together, these findings suggest that the OP block will lead to more quantisable models compared standard Pre-Norm, as we will show in Sec 6
Tab 4 abbates the effect of Norm positioning on OFE.
Nevertheless, we observe in Fig 4 that kurtosis still slightly increases in our OP blocks (to relatively modest values, maximum around 20), usually monotonically throughout training. Moreover, the question of why normalisation layers cause outlier features is still unanswered despite the clear evidence that removing them mitigates OF prevalence. We investigate these questions next.
We train for 4.2B tokens at 1.2B scale as this took 24 hours on 4 A100 80GB GPUs; we were unable to train for longer due to compute constraints. Scales under 1B were trained on single A5000 RTX-208017 GPU, taking around 2 days for 3.3B tokens (or equivalently, 50K steps at batch size 128 and sequence length 512).
5

---

Sec 3 key takeaways: normalisation layers and OFE.
OFE still occurs for weight-less or uncentred Norms, & both Pre/Post-Nomn (Figs) and 17
The OP Block (Fig 3 matches Pre-LN training speed/stability (Tabs | and 3) without standard
Norms. It does so through an Entropy Regulation method to prevent attention entropy collapse.
The OP Block greatly reduces OFE compared to standard blocks (Figs
and 13
Signal Propagation and Outlier Features
To better understand why OF' still appear (albeit greatly reduced) in the OP block, and why normali- sation layers cause OFs, we examine Signal Propagation behaviour during training and its effect on OFE. This will also clarify why modifications that improve Signal Propagation reduce OFE [T]. Sig-
nal Propagation [77 18 56 23 25 27 28 studies the input-wise Gram matrix Σn XX E R=**. and how Σn evolves in deep NNs for different layer features E Raxd
On the other hand, as we will see below, our kurtosis metric is related to the feature-wise Gram
matrix Σp def X X Raxd Recall our kurtosis is normalised 4th moment of X E R, normalised by the square of the second moment m2 (X) = 1 Losn.jed X2 a.j = nd | Because kurtosis is scale-invariant we can consider the setting where m2 (X) = 1 and the average squared activation is 1 without loss of generality In this case, Tr(2i) = Tr Tr(Ep) nd by the cyclic trace property.
nd
Then, our kurtosis, Eq (I, is Kurt(X) = à ビ ( E.= x2.)² =
(HEP)) 1 which is
simply a second moment (or average squared value) of diagonal entries of the feature-wise Gram matrix Ep. At the same time, again by the cyclic property of the trace, we have:
Tr(2p2iF) Tr(X xx
X) Tr(XX XX Tr(EET))
(4)
ndd (u)( +
Σ L.Sdi
(I+)
Σ (Σ))
a.
(5)
In words, Eq (4) tells that the sum of squared elements of EF is equal the sum of squared elements of Σg On the left of Eq (5) we decompose Eq (4) into our feature-wise kurtosis (Eq (I). of interest for OFE), plus the
of squared off diagonal elements of Σp, equal to the sum of squared elements of ΣS on the right. Hence, it is clear that Signal Propagation is relevant for OFE. Contrary to most existing works in Signal Propagaton, Eq 5 true throughout training, not only initialisation.
In particular, we see that the right-hand side of Eq (5) captures both the (normalised) activation norms across inputs Eosn (Σi)² a.a from the diagonal terms, and inner products between inputs
Ea.8Smane (E.) a. the off-diagonals. If X is the output of Norm layer, then 1 _ becomes a cosine similarity matrix with diagonals equal to 1. Deep NNs, and Transformers in particular, are well known be susceptible to particular Signal Prop defect called rank collapse [60] 25 where this cosine similarity matrix Σn degenerates to the all ones matrix and all inputs look identical to a deep layer. Noci et al. [27] and He et al. (28) demonstrate that, at least initialisation, the off- diagonals of Σj are positive and increase monotonically with depth in deep Transformers towards rank collapse, even with Signal Prop inspired modifications that ensure non-degenerate deep limit exists.
Bad Signal Prop encourages OFF For OFE, the upshot of these observations is that poor Signal Propagation (in terms of large off-diagonal values of Σ, close to rank collapse) will make the right-hand side of Eq (5) large (the rank collapsed limit has RHS n compared to nd2 when the inputs are orthogonal and Σn is diagonal). In turn, this puts pressure the LHS, which contains the feature kurtosis, to be large, hence OFE. This argument is not fully rigorous because the off-diagonals
Enggasiy (E which captures correlations between different neuron features, could increase on the LHS to allow the kurtosis to remain low Theoretically predicting this behaviour deep into modern NN training is outside the scope of this work; we note that while it is possible to write down training trajectories in feature learning regimes 33 34 most works interpreting feature learning in
II all experiments concerning signal propagation (i.e. input-wise correlations equivalently, the elements of E.), we first scale X down by / /m2(X) to give m²(X)= = and make X scale invariant. Indeed, Fig 33 shows the link between OFs and signal propagation depends on the diagonality of optimiser.

---

Pre-IN
Pro-LN
downscaled res diddu
vation Katosls
corndation
Actvvation
corritation
84
17
12
200
538 128 155 3M0 100 2.54
80
A4S 180 1.1 2- L25 1.5
514 140
155 1:18
CAN 1.1 2:00 2.6 120 1.5 Toceer seen
Totern seen
Tokeen seen
Tobern ween
Figure 5: Adam-trained Pre-LN layers at 1.2B scale with extreme OFE (left) are those with bad Signal Prop close to rank collapse during training (centre left). (Right vs. left two plots) Downweighting residual branches improves signal propagation during training and results in smaller OFs, particularly in early layers. Respective plots for OP (with & without final LN before unembedding) in Fig TO
NNs focus only a single gradient step 61+64 Having said that, we formalise the intuition of bad Signal Prop leading to larger feature kurtosis in the context of Gaussian features in Prop G.I
In any case, we can empirically study the link between bad signal propagation and OFEs, which we do in Figs [0 for Pre-LN& & OP blocks trained with AdamW at 1.2B scale on Languini Books. We plot both the layerwise evolution of the kurtosis on the left and the average off-diagonal entry of (i.e. the average imput-wise correlation) on the right, normalised so that m²(X) 1.
-Xx
As suggested by Eq (5). we see a strong association between kurtosis and Signal Propagation: the layers with larger kurtosis tend to be the ones with larger input correlations, and vice versa. In particular, in Fig we see that the Pre-LN layer (2 in this case) with the most extreme OFE (kurtosis peaking over 1000) is precisely the one with the worst Signal Propagation closest to rank collapse (average input correlation peaking over 0.8) during training. Moreover, the trajectory of kurtosis closely tracks the trajectory of input correlations throughout training, with their peaks appearing at similar training steps, across layers. Fig 12 shows that the Adam -aau-teiined Pythia models [16] are very close to rank collapse, which partially explains their large OFs in [i
Given that Signal Propagation characteristics during training depict how a model creates structure (through increasing or decreasing the inner product for different inputs) in its layer representations to best learn the task at hand, our results suggest that OFs occur partly due to the inherent nature of the task that the model is trained on, particularly in architectures that are less prone to OFs, such as our OP block. In Transformers, this appears most apparent in the inputs to the final unembedding layer, which are linearly projected to the predictions: they tend to have similar kurtosis levels in both OP and Pre-Norm blocks, and the most extreme OFE rarely occurs in the final layers, (Figs [ [4] and 18). We hypothesise this is because extreme OFE in late layers would imply high kurtosis which could imply representations close to rank collapse by Eq (5). from which it may be hard to learn useful linear predictions with optimisers like Adam.
The correlation we identify between OFE and Signal Propagation also allows us to observe that interventions that worsen Signal Propagation during training cause increased OFE. Likewise, methods improving Signal Propagation throughout training help to mitigate OFE. This can be seen in Fig 5 for downscaled residuals, h(x) = x x+ x+) with some B < 1, which Wortsman et al. [10] show improve quantisation on vision-language models. We explore this link further in terms of normalisation layers and other architectural choices inspired by Signal Prop in App C
5 Optimisation Choices and Outlier Features
So far, we have focused on the impact of architecture for OFE, As a result, up until now all models have been trained with Adam 43) optimiser, which is an adaptive diagonal preconditioner, with default hyperparameters e.g. B1 = 0.9, 82 = 0.999, e = 10 8, As OFE is a training phenomenon, is important to also consider the role of optimsation choices, which we now explore.
Learning Rate Perhaps unsurprisingly, we find that using smaller L.s leads to reduced OFE during training (Figs 6 25 2nd across different architectures. In these cases, slightly reducing the maximum LR in our scheduler (e.g. 0.001 0.0003 in Fig 6) did not lead to a loss in convergence speed (Fig 27). highlighting that one should use a smaller I.R to avoid OF' if convergence is not affected.
10²
Activation Kurtosis
307
101
307
Liwl033
10t4
rours
10
23K
GOK
Training Step
Figure 6:
Smaller LRs lead to
7
smaller OFs across different blocks.

---

Adaptivity Having seen the importance of LR in OFE, we now assess the impact of adaptive LRs through the e hyperpa-
107
Train Loss
Adam
301
rameter in Adam, where the Adam update is -n/m// ( <),
? is global IR, and my and v, denote first and second moconds of each parameter's s gradient, respectively. c dampens adaptive preconditioning, with larger E reducing adaptivity for param- eters with smaller Ut. In Figs 7,29 and [nd Tab 2 we show that increasing& e also reduces OrE. Thus, one should increase to reduce OFE, if convergence is not impacted (like Fig 28
10-4
3e-04
10-1
23K
100
62K
BXC
Training Step
Figure 7: Larger Adam reduces OFs in 130M Pre-LN transformers.
Non-Diagonal Preconditioners To push the question of adap- tivity further consider the effect of diagonal adaptivity and OFE. First-order optimisers like Adam [43] or AdaFactor 65 are the de-facto optimisers in deep learning, acting as diagonal preconditioners where each parameter has its own
Activation Kurtosis
80
Addemil
Roteted Ads Tacter] in Shampos)
Rotated Adam8)--2244)
adaptive learning rate to scale its gradient. On the other hand,
popular second-order optimisers like K-FAC [66] or Shampoo 67 68) are non-diagonal preconditioners acting on the full gradient. Second-order optimisers are known to converge faster per-update than first-order methods, but first-order optimisers are much more widespread due to the additional overheads of non-diagonal preconditioning. We provide background on different NN optimisers in App A.2
2K 3K 46
Training Step
Figure 8: Diagonal optimisation on rotated parameters reduces OFs.
Recently, Vyas et al. [69] established a precise connection between first and second-order optimisers, namely: Shampoo 667 can be seen as running the diagonal AdaFactor [65] method after first rotating into the eigenbasis of Shampoo's 's preconditioner, before rotating back. This insight disentangles two effects of non-diagonal preconditioners: 1) transforming the parameter space in which one optimises via a rotation, and 2) using a diagonal optimisation method in the rotated space. Vyas et al. [69] use this insight to propose the SOAP optimiser, which applies AdamW in the rotated parameter space obtained from Shampoo's 's eigenbasis. SOAP is shown to converge slightly faster per step than Shampoo, which itself converges faster per step than Adam W/AdaFactor (verified in Fig 32).
In Fig we compare OFE in popular first-order optimisers, 1) AdamW and 2) AdaFactor, to rotated non-diagonally preconditioned alternatives, 3) SOAP and 4) AdaFactor in Shampoo's 's eigenbasis (akin to Shampoo as shown by Vyas et al. 69)), trained using 130M Pre-Norm transformers on CodeParrot. We clearly see that rotating the parameter space in which one optimises, as done in SOAP/Shampoo, mitigates OFEs. This effect is independent of the diagonal preconditioner used and occurs even spite of the Pre-Norm layers, which we know are susceptible OFs, in Sec 3 10 To further highlight the importance of diagonal adaptivity, in App E. we compare SGD to Adam on a setting where SGD can match Adam's is convergence speed: image classification with MLP. There, we again see that the non-adaptive SGD suffers less from OFs compared to the diagonal Adam.
Elhage al. [4] show a similar result as Fig g using random rotations, while QuaRot M applies random Hadamard rotations to remove OFs for post-training quantisation (PTQ) using computational invariance. SpinQuant [70] extends QuaRot using learnt rotations but again only considers inference time, after OFs have emerged during training. Both QuaRot and SpinQuant incur additional overheads to apply rotations in the forward pass inference time. In contrast, non-diagonal preconditioners optimise using "learnt" rotations that adapt during training, which leaves the forward pass intact post training and enables the improvement in convergence speed per step that Shampoo/SOAP enjoy relative to Adam W/AdaFactor (seen in Fig 32 and Tab (.. In sum, our results reveal an additional appeal of second-order optimisers: not only are they faster to converge per step, but they also lead to models that are not as prone to OFs and are thus easier to quantise, as we will see in Sec b
Breaking down kurtosis updates The findings in this section point to the importance of large diagonal adaptive LRs for OFE This motivates us to break down the updates to kurtosis into terms of different powers in the learning rate 7/. in App F. We find that sub-leading order updates (in terms of LR) are the key driver in increasing kurtosis, providing a consistent mechanism for OFE that encapsulates our different observations concerning the roles of optimiser and architecture
This,, coupled with our findings with the OP block in See 3 suggests that extreme OFs found in LLMs are due to an interaction of the architectural choice of Norm layers and the optimiser choice of diagonal adaptive preconditioners. Note Pre- Tre-SRMSNomm in Fig which doesn't 't have trainable Norm parameters.
8

---

Sec 5 key takeaways: Optimisation Choices and OFE.
Large diagonal adaptive LRs can lead to Outlier Features during training (Figs 6
25
26 and g
Reducing Adam W adaptivity via increased e hyperparameter reduces OFs (Figs
29 and 37
The non-diagonal preconditioning of second-order optimisers greatly minimises OFE (Fig and Tab 2). even with OF-prone architectures like Pre-Norm.
6 Additional Experiments
We conclude our study with additional experiments regarding scaling and quantisation properties of suggestions to minimise OFs. Further experimental details can be found in App
1.28 param less (90B takens)
7B param loss (68 tokens)
1.28 param kurtesis (90B takens)
7B param kurtocis (68 tokens)
Model
33
The Morm
Marm Adamli/ 20
331
155
54B
720
300
246
Tollcnd Figure 9: Loss (left two plots) and kurtosis (right two plots) curves for different models trained with AdamW. Our conclusions on OFs continue to hold when scaling both token count and model scale.
Scale Up until now, we have studied OFE and how we can minimise it scales up 1.2B parameters and 5B tokens. We now consider scaling both model size and training length in terms of loss performance and OFE. We compare scaling the OP and Pre-Norm blocks keeping Adam as optimiser; it would be interesting to considering scaling experiments using second -order optimisers
(in distributed settings e.g. [68 (7]) in future work. The dataset is FineWeb-Eda [41. We warmup the LR for 5% of training steps maximum value (0.001 and 0.0003 for 1.2B and 7B respectively),
before cosine decay. Due to computational cost, very little hyperparameter tuning was performed with these experiments and all default hyperparameters were optimised for the Pre-Norm baseline.
So far we have studied OFE training token counts that are relatively small compared to the "compute-optimal" Chinchilla recipe [72). In Fig
9 (first& & third subplots), we scale the number of training tokens at 1.2B parameter scale to around 90B, which beyond the Chinchilla prescription. We see that the OP block is indeed able to closely match Pre-Norm loss at longer token counts, but still benefits from significantly reduced OFs (at least an order of magnitude lower kurtosis throughout training) despite the aggressive Adam W LR and long training horizon.
In Fig 9 (second and fourth subplots), we additionally scale the model size 7B parameters, scale that has previously been hypothesised to be a sharp cutoff above which "systematic" OFs emerge [ Due to cost, we only train for 6B tokens, which is more than enough for extreme OFs to emerge (over 600 kurtosis averaged across residual layers) with Pre-Norm layers. On the other hand, the OP block has peak kurtosis under 10 and yet still matches Pre-Norm loss performance at 7B scale. Although increasing Adam c from 10 to 10-5 also reduces peak kurtosis to under 10 with the OF-prone Pre-Norm model, leads to a significant decrease in convergence speed in this setting.
Quantisation Returning to our original motivation, we investigate the effect of our suggested architectural and optimisation choices to minimise OFs in terms of quantisation. In Tab we take the OPT-125m 66 setting of Bondarenko et al. [1], training models using Adam W in standard mixed FP16/FP32 precision on BookCorpus+ Wikipedia for around 12B tokens. Post training, we quantise (PTQ) to int8 reight-and-acciivtion using the same quantisation recipe as [15). We report standard deviation over 3 seeds for PTQ, as random subsets of data are used to estimate the quantiser range.
Tab compares both the standard precision perplexity (FP16/32) and also 8-bit quantised perplexity (W8A8) across architecture and optimiser choices. We additionally present our kurtosis metric, Eq (L. calculated after training and averaged across layers. We compare 3 differrnt transformer blocks: a) standard Pre-LN, b) the Gated Attention block proposed by Bondarenko et al. [15] to reduce OFs. and c) our OP block, as well 5 different optimisation setups that are added one after another: 1) the default hyperparameters of [15], 2) removing dropout regularisation, 3) increasing the maximum LR from 4 x 10-4 + 10-3 3 4) increasing AdamW e from 10- + 10-5 and 5) changing Adam to SOAP optimiser (keeping = 10- Optimiser choices 2) and 3) were designed to improve standard precision performance, albeit potentially at the detriment of quantisation performance due to
increased OFs. Optimiser choices 4) and 5) were chosen to reduce OFs, from our findings in See S

---

Table 2: Average kurtosis across layers, plus standard precision (FP16/32) and quantised int8 (W8A8)
perplexity of various 125m OPT models [6] trained on Seassoorps+Wiiipeeee [T5]. Kurtosis strongly correlates with int8 error across settings, and the best int8 setup combines our architectural (OP) and optimiser (SOAP) suggestions, with only 1.2 kurtosis and 0.16 perplexity increase.
Architecture
Optimiser Hyperparameters
Kurtosis
FP16/32 (1)
W8A8 (4)
Pre-LN
Default from [5]
25.6
16,00
63.44-50.1
-Dropout Regularisation
46.7
15.53
105, 105.9425.0
+Big LR (4 x 10 +100-
61.4
15.40
59.0410.7
+ -Big Adam (10 (10-6+100 1000--1---
43.6
15.19
216,2:887.6
+ -on-diag Precond (SOAP)
5.1
14.80
16.43
Gated Attention 15
Default from 15
4.5
15.63
16.2+0.09
-Dropout Regularisation
28.8
15.08
18.7+0.09
+ +Bgg LR (4 x 10 +10
28.8
14.99
17.04-10.09
+ Big Adam (10- -100-
16.0
14.78
15.54
Non-diag Precond (SOAP)
16.7
14.65
15.64:0.01
OP (ours)
Default from [5]
3.6
15.64
16.01 +0000
Dropout Regularisation
7.1
15.15
15.78
+ Big LR (1 x 10 10-1~1003)
12.0
14.96
15.60
+ Big Adam (10 >=100=>
6.0
14.89
15.62
+ -Non-diag Precond (SOAP)
1.2
14.71
14.87 :00.01
As seen in Tab 2 our findings throughout the rest of our paper are validated. Firstly, our kurtosis metric to measure OF' is indeed highly correlated with quantisation error, which we define as the increase perplexity from FP16/FP32 to W8A8. For example, the Pre-LN model without SOAP has consistently high kurtosis (over 25), and also consistently poor performance at W8A8 (over 45 quantisation error across all AdamW optimiser settings). Secondly, our OP block has consistently low kurtosis (below 12) compared to the other models, and this directly translates to low quantisation error (below 0.73 across all optimiser settings). This low kurtosis/quantisttinn error with OP block holds true even for aggressive optimiser choices, like large diagonal adaptive LRs, that improve standard precision perplexity but also increase kurtosis. Moreover, the baseline Gated Attention model of [55] struggles with OFs when dropout is removed and a large learning rate is used, leading to increased quantisation error (2-4 perplexity increase), but increasing Adam W as suggested in Sec reduces kurtosis (29 to 16) and quantisation error 0.76, whilst also improving FP16/32 perplexity.
Finally, changing AdamW to SOAP optimiser either dramatically reduces kurtosis (in the case of Pre-LN and OP) matches the kurtosis reduction of increasing Adam (for Gated Attention), while also improving mixed precision performance for all models 11 This leads to the only non-catastrophic W8A8 perplexity (16.43) with Pre-LN, and the best overall W8A8 model when combining SOAP optimiser with our OP architecture (14.87 perplexity. with only 0.16 degradation from standard precision). This result highlights the combination of our architectural and optimiser suggestions for minimising OFs as a promising approach to training fast-converging and easily quantisable models.
7 Discussion
The goal of this work was to better understand the emergence of Outlier Features during standard NN training, and propose architectural and optimisation interventions that minimise their prevalence. On the architectural side, have shown that Normalisation layers can have unwanted effeets on OFs during training. Removing standard Norms through our Outlier Protected transformer block minimises OFs during training without loss of convergence speed training stability. On the optimisation side, we highlight that large diagonal adaptive learning rates are crucial for OFs, and non-diagonal preconditioners offer an appealing combination of reduced OFs and improved convergence speed. We demonstrate our methods to minimise OFs are effective at scales of up to 7B parameters, and also directly translate improved post-training quantisation performance. Overall, our results reveal the complex interactions between architecture and optimiser that lead to the OFs widely observed in LLMs. In future work, would be interesting to consider applying our methods to minimise OFs for post-training quantisation performance at larger scales, and also low-precision training.
We keep the batch size 196 & context length 512 same as 15 for fair comparison. is likely the FP16/FP32 gains of non-diagonal preconditioners e.g. SOAP would increase with higher token counts per step (73).
10

---

Acknowledgements
We would like thank Sam Smith for helpful suggestions at the beginning of this project, Tiago Pimentel for constructive feedback on an early version of this manuscript, and Saleh Ashkboos for various insightful discussions around rotations and quantisation. We are also grateful for the mostly constructive feedback we received from anonymous reviewers. Finally, we would also like to thank the Swiss National Supercomputing Centre for access GPU resources to perform some of the experiments in this work, via grant under project ID a06 on Alps as part of the Swiss Al Initiative. We thank Alex Hägele for help setting up our experiments on Fine FineWeb-FFF Web-Edu.
Reproducibility Statement
Our code for experiments on the CodeParrot dataset can be found at https:// com/ httphh//ssimmplffieeeeeerreaaaee
References
[1] Olga Kovaleva, Saurabh Kulshreshtha, Anna Rogers, and Anna Rumshisky. Bert busters: Outlier dimensions that disrupt transformers. Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021, 2021.
[2] William Timkey and Marten van Schijndel. All bark and no bite: Rogue dimensions in transformer language models obscure representational quality. In Marie-Francine Moens, Xuanjing Huang. Lucia Specia, and Scott Wen-tau Yih, editors, Proceedings of the 2021 Conference Empirical Methods in Natural Language Processing, pages 4527-4546, On- line and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10. 10.18653/v 1188536/1/2021.... http://12222..1mm//.// www.leemmnmmmm..... 2021. emnlp- mapo.main... m72
URL http:tttt///acaaatooooo org/
[3] Yelysei Bondarenko, Markus Nagel, and Tijmen Blankevoort. Understanding and overcoming the challenges of efficient transformer quantization. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 7947-7969, 2021.
[4] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 2018.
[5] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Lan- guage models are unsupervised multitask learners. 2019.
[6] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.0108 2022.
[7] Xiuying Wei, Yunchen Zhang, Xiangguo Zhang, Ruihao Gong. Shanghang Zhang, Qi Zhang, Fengwei Yu, and Xianglong Liu. Outlier suppression: Pushing the limit of low-bit bit transformer language models. Advances in Neural Information Processing Systems, 2022.
[8] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Gpt3. int8 O): 8-bit matrix multiplication for transformers at scale. Advances in Neural Information Processing Systems, 35:30318 30332, 2022.
[9] Aohan Zeng, Xiao Liu, Zhengxiao Du, Zihan Wang, Hanyu Lai, Ming Ding, Zhuoyi Yang, Yifan Xu, Wendi Zheng, Xiao Xia, Weng Lam Tam, Zixuan Ma, Yufei Xue, Jidong Zhai, Wenguang Chen, Zhiyuan Liu, Peng Zhang, Yuxiao Dong, and Jie Tang. GLM-130b: An open bilingual pre-trained model. In The Eleventh International Conference on Learning Representations, 2023. URL htttttt://opeereiiieeeeeee wttp:/rrreii...ettttormmiii--e-.
[10] Mitchell Wortsman, Tim Dettmers, Luke Zettlemoyer, Ari Morcos, Ali Farhadi, and Ludwig Schmidt. Stable and low preciicciooon training for large-scale vision- languag models, Advances 2023.
in Neural Information Processing Systems,
11

---

[11] Saleh Ashkboos, Amirkeivan Mohtashami, Maximilian L Croci, Bo Li, Martin Jaggi, Dan Alistarh, Torsten Hoefler, and James Hensman. Quarot: Outlier-free 4-bit inference rotated llms. arXiv preprint arXiv:2404. 00456, 2024.
[12] Aniruddha Nrusimha, Mayank Mishra, Naigang Wang, Dan Alistarh, Rameswar Panda, and Yoon Kim. Mitigating the impact of outlier channels for language model quantization with activation regularization. arXiv preprint arXiv:2404.003660 2024.
[13] Giovanni Puccetti, Anna Rogers, Aleksandr Drozd, and Felice Dell'Oriect Outliers di- mensions that disrupt transformers are driven by frequency. Findings of EMNLP 2022. Association for Computational Linguistics, 2022.
[14] Nelson Elhage, Robert Lasenby, and Christopher Olah. Privileged bases in the transformer residual stream, 2023. URL https://ransffrmer-cinnaaa pubb20233privlleeeebsisiiii html. Accessed, pages 08-07, 2023.
[15] Yelysei Bondarenko, Markus Nagel, and Tijmen Blankevoort. Quantizable transformers: Removing outliers by helping attention heads do nothing. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[16] Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, al. Pythia: suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pages 2397-2430. PMLR, 2023.
[17] Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Expo- nential expressivity in deep neural networks through transient chaos. D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016.
[18] Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. In International Conference on Learning Representations, 2017.
[19] Alexander G de G Matthews, Mark Rowland, Jiri Hron, Richard E Turner, and Zoubin Ghahramani. Gaussian Process Behaviour in Wide Deep Neural Networks. In International Conference on Learning Representations, volume 4, 2018.
[20] Jachoon Lee, Jascha Sohl-dickstein, Jeffrey Pennington, Roman Novak, Sam Schoenholz, and Yasaman Bahri. Deep Neural Networks Gaussian Processes. In International Conference on Learning Representations, 2018.
[21] Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[22] Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nngp and ntk for deep attention networks. In International Conference on Machine Learning, pages 4376-4386 PMLR, 2020.
[23] Soufiane Hayou, Eugenio Clerico, Bobby He, George Deligiannidis, Arnaud Doucet, and Judith Rousseau. Stable resnet. In International Conference on Artificial Intelligence and Statistics, pages 1332. PMLR, 2021.
[24] Lorenzo Noci, Gregor Bachmann, Kevin Roth, Sebastian Nowozin, and Thomas Hofmann. Precise characterization of the prior predictive distribution of deep relu networks. Advances in Neural Information Processing Systems, 14:20851-20862 2021.
[25] James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszez, Valentin Dalibard, Jascha Sohl-Dickstein, and Samuel S Schoenholz. Rapid training of deep neural networks without skip connections or normalization layers using deep kernel shaping. arXiv preprint arXiv:2110.0165,, 2021.
12

---

[26] Mufan Bill Li, Mihai Nica, and Daniel M Roy. The neural covariance sde: Shaped infinite depth-and- width networks at initialization. arXiv preprint arXiv:2206.02668, 2022.
[27] Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. urXiv preprint arXiv:2206.03 03 2022.
[28] Bobby He, James Martens, Guodong Zhang, Aleksandar Botev, Andrew Brock, Samuel Smith, and Yee Whye Teh. Deep transformers without shortcuts: Modifying self-attention for faithful signal propagation. In The Eleventh International Conference on Learning Represen- tations, 2023. URL httpths:://peeeeeiiieeeeee httppp/eeeiiii..eeeefotrrrrrr/// http:/oorrnndnnPP//s/
[29] Lorenzo Noci, Chuning Li, Mufan Bill Li, Bobby He, Thomas Hofmann, Chris Maddison, and Daniel M Roy. The shaped transformer: Attention models in the infinite depth-and- limit. arXiv preprint arXiv:2306.17755 2023.
[30] Bobby He and Thomas Hofmann. Simplifying transformer blocks. In The Twelfth International Conference on Learning Representations, 2024. URL httppttt////ppereeiiiee wtwp://wrwiiiieeeeeorrrrre. net/forum?
id=RtDok9eS3s
[31] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training differentiable program- ming. Advances in neural information processing systems, 32, 2019.
[32] Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representutions, 2020.
[33] Greg Yang and Edward Hu. Feature learning in infinite-wine winth neural networks, arXiv preprint arXiv:201 1.4522, 2020.
[34] Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evo- lution in wide neural networks. Advances in Neural Information Processing Systems, 35: 32240-32256, 2022.
[35] Chris Olah, Alexander Mordvintsev, and Ludwig Schubert. Feature visualization. Distill, 2 (11):e7, 2017.
[36] Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pages 38087-38099. PMLR, 2023.
[37] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[38] Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Baohong Lv, Fei Yuan, Xiao Luo, et al. Scaling transnormer 175 billion parameters. arXiv preprint arXiv:2307.1499 2023.
[39] Aleksandar Stanić, Dylan Ashley Oleg Serikov, Louis Kirsch, Francesco Faccio, Jürgen Schmidhuber, Thomas Hofmann, and Imanol Schlag. The languini kitchen: Enabling language modelling research at different scales of compute. arXiv preprini arXiv:2309. / 197, 2023.
[40] Yukun Zhu. Aligning books and movies: Towards sike visual explanations by watching movies and reading books. arXiv preprint 1rXiiv:5566. 067244 2015.
[41] Guilherme Penedo, Hynek Kydlíčck, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data at scale. arXiv preprint arXiv:2406. /7557, 2024.
[42] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner. Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2020.
13

---

[43] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.
[44] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:/607.06450, 2016.
[45] Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
[46] Jack W Rae, Sehastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young. et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:21 12. 1446, 2021.
[47] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 139971111 2023.
[48] Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. In International Conference on Learning Representations, 2018.
[49] Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
[50] Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: without normalization. In International Conference on Learning Representations, 2018.
Residual learning
[51] Soham De and Sam Smith. Batch normalization biases residual blocks towards the identity function deep networks. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 19964-19975. Curran Associates, Inc., 2020.
[52] Xiao Shi Huang, Felipe Perez. Jimmy Ba, and Maksims Volkovs. Improving transformer optimization through better initialization. In Hal Daumé III and Aarti Singh, editors, Pro- ceedings of the 37th International Conference on Machine Learning, volume 119 of Pro- ceedings of Machine Learning Research, pages 4475-4483 PMLR, 13-18 Jul 2020. URL
https://proceeding mlr. press / / www1///hhhgg0000 htmll
[53] Andy Brock, Soham De, Samuel L Smith, and Karen Simonyan. High-perffrmance large- scale image recognition without normalization. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 1059-1071. PMLR, 18-24 Jul 2021.
[54] Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian McAuley. Rezero all you need: Fast convergence large depth. In Uncertainty in Artificial Intelligence, pages 1352-1361. PMLR, 2021.
[55] Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 32-42, 2021.
[56] Soufiane Hayou, Amaud Doucet, and Judith Rousseau. On the impact of the activation function on deep neural networks training. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 2672-2680. PMLR, 09-15 Jun 2019.
[57] Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Joshua M. Susskind. Stabilizing transformer training by preventing attention entropy collapse. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Reseurch, pages 40770-40803 PMIR, 23-29 Jul 2023. URL. https:://// htaa:::s//pocceednge mlr, mwt.ggssmm.rrrrss//// press/
v202 zhail3a. html
14

---

[58] Alex Henry, Prudhvi Raj Dachapally, Shubham Shantaram Pawar, and Yuxuan Chen. Query- key normalization for transformers. In Findings of the Association for Computational Linguis- tics:
EMNLP 2020, pages 4246 4253, 2020.
[59] Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. In International Conference on Machine Learning, pages 7480-7512. PMLR, 2023.
[60] Yihe Dong, Jean -aatiste Cordonnier, and Andreas Loukas, Attention not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning, pages 1793-2802 PMLR, 2021.
[61] Greg Yang, Edward J Hu, Igor Babuschkin. Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. In A. Beygelzimer, Y. Dauphin, P. Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL https //openrevier net/ wwt//neetoornnee for/mmirddddd-oe (x BX6QKUBM2AD
[62] Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Zhichao Wang. Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https //openrevier net/forummidde akddwRG6EGi axch
[63] Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit. In The Twelfth International Conference on Learning Representations, 2024. /tttpps:s///peeereeeee
net/forum?1d-k2Jehhe
[64] Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs VI: Feature learning in infinite depth neural networks. In The Twelfth International Conference on Learning Representations, 2024. BtLLittts///perrriieeee net www.ee//ooooniddi//paaaa
[65] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596-4604. PMLR, 2018
[66] James Martens and Roger Grosse. Optimizing neural networks with factoee approximate curvature. In International conference on machine learning, pages 2408-2417, PMLR, 2015.
[67] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pages 1842-1850. PMLR, 2018.
[68] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09011, 2020.
[69] Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409.1 1131,, 2024.
[70] Zechun Liu, Changsheng Zhao, Igor Fedorov, Bilge Soran, Dhruv Choudhary, Raghuraman Krishnamoorthi, Vikas Chandra, Yuandong Tian, and Tijmen Blankevoort. Spinquant-IIm quantization with learned rotations. arXiv preprint arXiv:2405. 10....4050..66, 2024.
[71] Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Twasaki, Jose Hallego-PPoadaa Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, and Michael Rabbat. A distributed data- matalparallel pytorch implementation of the distributed shampoo optimizer for training neural networks at-scale. arXiv preprint arXiv:2309.064: 2023.
[72] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15555 2022.
15

---

[73] Guodong Zhang, Lala Li, Zachary Nado. James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger Grosse. Which algorithmic choices matter at which batch sizes? insights from noisy quadratic model. Advances in neural information processing systems, 32, 2019.
[74] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng. Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang. and Tieyan Liu. On layer normalization the transformer architecture. In International Conference on Machine Learning, pages 10524-10533. PMLR, 2020.
[75] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1412 2014.
[76] Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, and Alberto Bietti. Heavy- tailed class imbalance and why adam outperforms gradient descent on language models. arXiv preprint arXiv:2402. 19449, 2024.
[77] Jun Gao, Di He, Xu Tan, Tao Qin, Liwei Wang, and Tie- Yan Liu. Representation degeneration problem in training natural language generation models, 2019.
[78] Mingjie Sun, Xinlei Chen, J Zico Kolter, and Zhuang Liu. Massive activations in large language models. arXiv preprint 0rXiv:2402.17762, 2024.
[79] Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers. In The Twelfth International Conference on Learning Representations, 2024. URL htt.tttts///peeeeviiee net/ fodu/rmmddddd2200LL11
[80] Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. arXiv preprint 20XV::2012.09816. 2020.
[81] Bobby He and Mete Ozay. Feature kernel distillation. In International Conference on Learning Representations, 2022.
[82] Jerry Chee, Yaohui Cai, Volodymyr Kuleshov, and Christopher De Sa. QulP: 2-bit quantization of large language models with guarantees. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[83] Albert Tseng, Jerry Chee, Qingyao Sun, Volodymyr Kuleshov, and Christopher De Sa. Quipft: Even better Ilm quantization with hadamard incoherence and lattice codebooks. arXiv preprint XXXv:2400204439 2024.
[84] Jerry Yao-Chieh Hu, Pei-Hsuan Chang, Haozheng Luo, Hong- Yu Chen, Weijian Li, Wei- Po Wang, and Han Liu. Outlier-effiieen hopfield layers for large transformer-basee mod- els. In Forty-first International Conference on Machine Learning, 2024. URL https //openreview. wtwpiieww.eeeeeeeeeeeeee mett:/oouumddd-LLDMGG......
[85] Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. International conference on machine learning, pages 448-456 pmlr, 2015.
[86] Guodong Zhang, Aleksandar Botev, and James Martens. Deep learning without shortcuts: Shaping the kernel with tailored rectifiers. In International Conference on Learning Represen- tations, 2022.
[87] Samuel L Smith, Andrew Brock, Leonard Berrada, and Soham De. Convnets match vision transformers at scale. arXiv preprint arXiv:2310. 100164410.16664, 2023.
[88] Alex Henry, Prudhvi Raj Dachapally, Shubham Shantaram Pawar, and Yuxuan Chen. Query- key normalization for transformers. In Trevor Cohn, Yulan He, and Yang Liu, editors, Findings of the Association for Computational Linguistics: EMNLP 2020, pages 4246-4253, On- line, November 2020. Association for Computational Linguistics. doi: 10. 10.18853/v//002. URL httlltttss//aaaateettee http://aaaaaattyy....o org/ 2020. ffi.fiindgss weindings-emnl..... 379
16

---

[89] Yuandong Tian, Yiping Wang. Zhenyu Zhang, Beidi Chen, and Simon Shaolei Du JoMA: Demystifying multilayer transformers via joint dynamics of MLP and attention. In The Twelfth International Conference on Learning Representations, 2024. URL Utthttps::// httpss//opeereviee net/forummidd-bJJgRavYoY
[90] Greg Yang. Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, and Samuel S. Schoen- holz. A mean field theory of batch normalization. International Conference on Learning Representations, 2019. SRKLitttps://operrvvies http://pppereeeeeeeeeeeeerraaaaa wwt..tt//oormmmmoddsssttaaccc
[91] Lechao Xiao, Jeffrey Pennington, and Samuel Schoenholz. Disentangling trainability and generalization in deep neural networks. In International Conference on Machine Learning. pages 10472. PMLR, 2020.
[92] Kenta Oono and Taiji Suzuki. Graph neural networks exponentially lose expressive power for node classification. In International Conference on Learning Representations, 2019.
[93] Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pen- nington. Dynamical isometry and a mean field theory of cnns: How train 10.000-layer vanilla convolutional neural networks. In International Conference on Machine Learning, pages 5393-5402. PMLR, 2018.
[94] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[95] Yizhang Lou, Chris E Mingard, and Soufiane Hayou. Feature learning and signal propagation in deep neural networks. In International Conference on Machine Learning, pages 14282.. PMLR, 2022.
[96] Aristide Baratin, Thomas George, César Laurent, R Devon Hjelm, Guillaume Lajoie, Pas- cal Vincent, and Simon Lacoste-Julien. Implicit regularization via neural feature align- ment. Arindam Banerjee and Kenji Fukumizu, editors, Proceedings of The 24th In- ternational Conference on Artificial Intelligence and Statistics, volume 130 of Proceed- ings of Machine Learning Research, pages 2269-2277. PMLR, 13-15 Apr 2021. URL http://wwww........me.eemmmm
https://proceedinng mlr. www.ppreeevviie wttps//ve///aaaaaeeeeeeeeeeeeeee wttl:///aa/nnlaa..a.t.aa html
[97] Mariia Vladimirova, Jakob Verbeek, Pablo Mesejo, and Julyan Arbel. Understanding priors in bayesian neural networks the unit level. In International Conference on Machine Learning, pages 6458-6467. PMLR, 2019.
[98] Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice. Advances in neural information processing systems, 30, 2017.
[99] Amir Joudaki, Hadi Daneshmand, and Francis Bach. On the impact of activation and normal- ization in obtaining isometric embeddings at initialization. In seventh Conference Neural Information Processing Systems, 2023.
[100] Alexandru Meterez, Amir Joudaki, Francesco Orabona, Alexander Immer, Gunnar Rätsch, and Hadi Daneshmand. Towards training without depth limits: Batch normalization without gradient explosion. arXiv preprint arXiv:2310.02012 2023.
[101] Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based on gemini research and technology. arXiv preprint arXiv:2403.08295. 2024.
[102] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063 2024.
[103] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transform- ers. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pages 12104-12113, 2022.
17

---

[104] Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham Kakade. Decon- structing what makes good optimizer for language models. arXiv preprint arXiv:2407.07972. 2024.
[105] Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, et al. Minicpm: Unveiling the potential of small language models with scalable training strategies. arXiv preprint ar arXiv:2404.06.
2024.
[106] Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro Von Werra, and Martin Jaggi. Scaling Laws and Compute-Optinll Training Beyond Fixed Training Durations. Advances in Neural Information Processing Systems, 2024. URL http::tt/aaii..rrreeee..
2405. 18392
[107] Maxim Fishman, Brian Chmiel, Ron Banner, and Daniel Soudry. Scaling fp8 training to trillion-token llms. arXiv preprint arXiv:2409.12517 2024.
[108] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning, volume 139, pages 10347-10357, July 2021.
[109] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359 2022.
[110] Mihaela Claudia Rosca. On discretisation drift and smoothness regularisation in neural network training. arXiv preprint arXiv:2310.14036, 2023.
[111] Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5 r.S-ropoopp coursera: Neural networks for machine learning. University of Toronto, Technical Report, 6, 2012.
18

---

Appendix
Table of Contents
A Additional Background Knowledge
19
A.1 Mathematical Description of Transformer Blocks
19
A.2 Background on NN optimisers
20
A.3 Additional Related Work
21
B
Mathematical description of OP block
23
Modifications That Affect Signal Prop During Training Affeet OFE
24
Additional Experimental Details
25
Additional Experiments
27
F.1
Image Classification Experiments
38
Ablating the components of the OP block
43
F
Orders Activation Updates Kurtosis
47
Worse Signal Prop Means Higher Activation Kurtosis Gaussian Features
53
A
Additional Background Knowledge
A.1
Mathematical Description of Transformer Blocks
Transformer architecture (37 formed by sequentially composing Transformer blocks. The two popular Transformer blocks Pre-Norm and Post-Norm. The Pre-Norm Transformer block
48 49 is nowadays more widespread than the original Post-Norm block (37 due to advantageous
depth-scaling properties 77 51).
For an input sequence representation X Txxx with tokens and dimension d, the output Xoat of Pre-Norm Transformer block is:
X
X
MLP(Norm, X)), where X X,, MHA (Xin))
(6)
On the other hand, the Post-Norm Transformer block can be expressed as:
X
Norma X MLP(X)), where X Norm, (X,, + MHA(X
Here, "MHA" stands for Multi-Head Attention (detailed below), and "Norm" denotes a normalisation that the Pre-Norm transformer block consists of two sequential sub-blocks (one attention and one MIP),, with normalisation layers and residual connections for both sub-blocks. Crucially the Norm layers in Pre-Norm blocks are placed within the residual branch, whereas in Post-Norm blocks the Norm layers are placed after the residual connection.
layer like LayerNorm 44] RMSNorm 445 In words,
When we say "Pre-LN" we mean that the block is Pre-Norm and the Norm is LayerNorm likewise "Post-RMSNorm" would mean that the block is Post-Norm and the Norm is RMSNorm. and so on in Fig 2 There is no notion of "Pre" or "Post" with our OP block because there is no difference between Pre-Norm and Post-Norm if one removes the Norms.
and
In our work, the MI.P architecture single hidden-layer with hidden dimension that is 4d (as 37 and acts each token in the sequence independently.
19

---

The MHA sub-block uses softmax self-attention to share information between tokens. For given input sequence X, the softmax self-attention mechanism computes:
Attn X) = A(X)XWV where ( Softmax (v
+X) +M)
(8)
where s ( XW
are the pre-softmax attention scores/logits, and W W E Rddd, and W E Rdddd are trainable query, key and value parameters respectively.
The attention matrix A E RT*T allows different tokens to "attend" to each other, with mask M E RT'x determining which tokens any given token is allowed "attend" to. For causal auto- regressive transformers like GPT, Mij = 0ifi 2 and -oo else, which prevents a token from obtaining information from future tokens.
The Mullii-Haad Attention name arises due to the fact that in practice is typical to apply self-attention on H different "heads" with d = dę Hf before concatenating the heads, as follows:
(X) Concat (oncat(Attm) (Attni (X),... A (X)) WP,
where WP E Raxd denotes trainable matrix that combines different attention heads via a projection.
A.2 Background on NN optimisers
In this subsection we provide additional background on different deep learning optimisers to accom- pany for completeness.
Consider a weight matrix E RR(xr with scalar loss function L(W), and corresponding gradient G Vw L E (x We denote 17 as scalar learning rate and as a scalar "damping" hyperparameter that prevents numerical instabilities. We outline the three families of optimisers that we explore at various points our work, in terms their effect on outlier features: 1) SGD (with momentum), 2) diagonal preconditioners, and 3) non-diagonal preconditioners.
SGD with Momentum keeps an (exponential) moving average of gradients G, M E Rixr, and updates W as:
W
W n.M.
Diagonal preconditioners like Adam 775 are arguably the most popular deep learning optimiser family due to their combination of improved convergence modern architectures like transformers, compared to SGD, and computational efficiency, compared to non-diagonal preconditioners.
Adam works by maintaining not only exponential moving average of gradients M E RR' like SGD, also an element-wise second moment of gradients @ G. Rixr.
Then, the Adam update rule is:
M
W - W
VV+c
+
where all operations (e.g. division square root) element-wise This dlement-wise nature leads us to describe Adam as diagonal (as can easily be written out with matrix multiplication using diagonal matrices). Note that can view the elements of vV+c as per parameter adaptive learning rates, compared to SGD which just has a global shared LR 1).
AdamW W (43 is a popular variant of Adam that decouples weight decay. AdaFactor [65 is another variant of Adam, which replaces V by rank-1 approximation, V' Rixr, in the interests of memory. Despite the reduced rank of V", AdaFactor is similarly diagonal like Adam in the sense that there is unique element V for each element W. and all operations are element-wise.
Non-diagonal preconditioners obtain their update rule by instead applying a full non-diagonal (inverse) matrix preconditioner P P R to the (flattened) first order moment flat(M) R to give update- - (M))
Due to the exorbitant cost of storing, updating, inverting, and applying Ir x Ir matrices, a popular approach introduced by Martens and Grosse (66] is use kronecker factored preconditioner P == L & R, where L |gxx and R E xr to give update rule (with dampening):
20

---

W - W - - M MMR (R,)
In Shampoo 66
I is a moving average of GGT and moving average of GT G, although
different choices of exponents (e.g. 114 or - -1/2) are more common besides -1. For example, the
Shampoo update rule is W + W - (L e el,) for exponent 11/2.
As we discuss in Sec Vyas et al. [69] identify connection between Shampoo with exponent -1/2 and AdaFactor, and this insight gives rise to the SOAP optimiser. We refer the reader to Vyas et al. for additional background on different NN optimisers and their connections.
69
A.3
Additional Related Work
Understanding Outlier Features Kovaleva et al. M. Timkey and van Schijndel [2] first identified Outlier Features in trained Transformers and demonstrated that OFs are critical for representational quality and performance. Puccetti et al. [3 highlight the importance of token frequency [76] for OFs in transformers trained on language data, which is related to the representation degeneration phenomenon of Gao et al. [77]. and certain "vertical" structures appearing in attention matrices during training. Bondarenko et al. [15] term this vertical structure "no-op" behaviour, where uninformative tokens are given high attention weights, and show that modifying attention to encourage no-op behaviour can mitigate OFs. Dettmers al. [8] show that the effect of OFs is more pronounced at larger parameter scales, and Wortsman et al. [I0) suggest that OFs are related to increasing activation scales during training, motivating their use of downweighted residuals. Kovaleva et al. [I. Wei et al. 7 attribute OFs to the trainable parameters in Layer Normalisation. Nrusimha et al. [12] show that
OFs occur early training, and are stronger residual stream layers. Sun et al. (78 demonstrate the existence of "massive activations" and show they act as bias terms in transformers. Darcet et al.
779 show that outlier tokens with large activation norms lead to non-smoon attention maps in vision transformers, and propose additional "register" tokens in order to concentrate the outliers and yield smoother attention maps. Allen-Zht and Li [80], He and Ozay [8]] study a theoretical framework where sparse activations naturally appear and grow with gradient descent, owing to certain "lucky" neurons being correlated with useful features at initialisation, in order to study ensembling and knowledge distillation in two-layer convolutional NNs.
Outlier Features and Quantisation Wei et al. (7.. Bondarenko et al. 3 identified Outlier Features issue for quantised NNs. Most work in this area has focused on (weight) quantisation of already trained transformers S 36 11, for efficiency gains at inference time. Dettmers et al. [ keep outlier features in full precision to avoid their quantisation errors, while Xiao al. (36) propose migrate the quantisation difficulty of outlier features to their corresponding weights using some scaling factors. Chee et al. [2)] introduce the idea of "incoherence processing", where pre-trained weights are rotated with random orthogonal matrices to remove outliers, which is provably and empirically shown to make quantisation easier with their method QuIP. Tseng et [83] extend QuIP to use random Hadamard matrices (among other changes), which are more efficient and have better theoretical properties than random orthogonal matrices. Ashkboos et al. [T] combine incoherence processing with "computational invariance" to rotate the feature vectors in addition to pre-trained weights whilst preserving the forward pass, thereby removing OFs in the rotated features and achieving state of the art performance in beight-and-activation quantisation at inference time.
In terms of quantised training. Bondarenko al. [15] show that encouraging "no-op" behaviour can mitigate OF' and enable low-precision training, while Wortsman et al. 110] employ downweighted residuals (among other techniques) for quantised CLIP training. We discuss how our findings relate and extend these insights in the main text. Hu et al. (84) propose nutlier-efficien Hopfield Layers as an alternative to traditional attention mechanisms to improve post-training quantisation. Nrusimha et al. [12] propose to regularise the kurtosis of the outputs of a lincar layer for low- low-precision training. which the authors argue prevents migrating quantisation difficulty to the weights. We employ kurtosis to measure OFs, but focus on the kurtosis of the inputs to linear layer.
Normalisation Layers
Normalisation Layers have been near ever-present in NNs since their
introduction 85 44 owing to their training benefits. Many works since have considered removing Normalisation layers, by finding alternative mechanisms that keep their benefits. De and Smith [51) identify beneficial implicit effect of Normalisation layers in Pre-Norm Eq O architectures is to
21

---

downweight residual branches, and that explicit recreating this effect enables training deep NNs without Normalisation. Hayou et al. [23] show this theoretically using Signal Propagation theory, and propose downweighting residuals with scale factor de///))thh) to do so, which Noci et al. 27 corroborate in the transformer setting. Martens al. [25], Zhang et al. [86] demonstrate how to remove residual connections alongside normalisation layers in convolutional NNs using "transformed" activations, which He et al. [28] extend to the Transformer architecture by making attention more identity-like (see also "shaped" attention, Noci et al. (29). Brock et al. [53], Smith et al. [87 propose NFNett, and achieve state of the art performance on the ImageNet benchmark in an unnormalised residual convolution architecture, highlighting that Normalisation layers not necessary for best performance in convolutional models. NFNets employ downweighted residual branches to fix Signal Propagation at initialisation, among other techniques including adaptive gradient clipping. However, He et al. (28) He and Hofmann [30] find that removing Normalisation Layers, even with Signal Propagation modifications like downweighting residuals, leads to loss of performance in simplified Transformer blocks, implying that transformer training has different instabilities to convolutional models, and Normalisation layers have other training benefits in transformers.
Entropy Collapse and QK-Norm Zhai et al. (57 identify entropy collapse as a key training instability in transformers, where attention logits grow large during training. This causes the rows of the post-softmax attention matrix to become one-hot vectors and the attention weights are non-zero on only a single sequence position.
Mathematically, given a stochastic attention matrix A E RTXT (in the notation of Eq 8 can compute the entropy ((A), averaged over sequence locations, as:
we
H(A) Σ A.tt log(As.:),
(10)
and define Entropy Collapse to be the situation where H(A) tends to 0 during training. This occurs when the rows A., become one-hot vectors and token only attends to one other token, for all s. We treat 0 log(O) to be O.
To remedy entropy collapse, it is important to control the logits entering softmax from growing too large, and Zhai et al. [57] propose oReparam which regularises the spectrum of Query-Key weights in order do so.
As an alternative, Query-Key Normalisation 88 where the Queries and Keys are normalised using e.g. LayerNorm or RMSNorm after the Query/Key weight matrix has seen growing popularity, particularly in ViT-22B [59] where it was crucial for stable training.
Mathematically, instead of standard attention logits s(X) = XW (XWk) T (following the notation of Eq 8 ), QK-Norm first normalises the queries and keys across the de dimension:
XX) = Norm X Noom)NommX. JT
(11)
and is most commonly used as an addition on top of the Pre-Norm block Eq (6) (without removing other normalisation layers like in our OP block).
Other "entropy regulating" mechanisms include tanh thresholding (Grok-1) and clamping attention logits (DBRX). The training stability benefits of controlling attention entropy through QK-Nomm were shown at smaller scales in Wortsman et al. [0, who argue that the quadratic dependence in the attention logits (on the queries and keys), causes large attention logits to appear during training, hence entropy collapse. This is as opposed to convolutional/MIIP models which depend linearly on their inputs. Tian et al. [89] propose joint MLP/Attention dynamics to predict attention entropy during training. We note that the "vertical" or "no-op" attention structures discussed previous OF works [13 15] have collapsed attention entropy, and can be thus be seen as undesirable from the perspective of other existing works.
Signal Propagation Signal Propagation studies how different inputs evolve through a deep NN, and how their feature representation magnitudes and cosine similarities evolve with depth. Our work focuses on forward signal propagation (which studies the forward-pas. activations), as opposed to backward signal propagation (which studies the backward-pass activation derivatives), in line with the study of Outlier Features.
22

---

For an input activation matrix X E Raxd of n inputs and width d, mapped to an activation matrix xa layer signal propagation theory studies the evolution of the input- Gram matrix x E R=*= for increasing depths I. This is key object in an NN, as it tracks the "geometric information" that is conveyed in deep layer, through inner products between different inputs. The diagonal elements of Ej indicate the activation norms, and the off-diagonal elements indicates how similar a deep layer views two inputs to be.
X, x E
At initialisation, Yj can be tracked through its large d limits 20 several issues that will afflict badly designed NNs [8] 56 90
By studying Mj, one can see that affect either the diagonal elements, the off- -figonal elements or both at large depths. For example, the diagonal elements of Σn could blow up, which indicates exploding activation norms. For transformers, a particular degeneracy, known as rank collapse [60], can appear where the off- off-diagonals of Mj become positive and large, and 2j becomes proportional to the all ones matrix if activation norms are constant. Rank collapse is also possible in MLPs/CNNs Schoenholz et al. (18), Hayou et al. [56), Xiao et al. (91). Martens et al.
19
60
255, and is equivalent to the over-smoothing phenomenon in graph NNs [92] Martens et al. [25] argue that rank collapse will lead to vanishing gradients, which Noci et al. [27] show specifically for query and key parameters in transformers. As a result, when we refer to bad signal propagation, we mean that the off diagonals of Σn are large and positive, close rank collapse. This can be either
through the RMS of input correlations, n(R-T] Eazs ( a as we show in the appendix, or the mean, min-1) Edeg (Z a..9 as we show in Figs and 10
possible to design modifications to NN architectures and initialisations that correct potential degeneracies and/or yield simpler and/or more scalable architectures [93] [3
By applying Signal Propagation theory at initialisation,
[7] 28 (3).. But the vast majority of existing works in the literature do not theoretically study training beyond initialisation, and those that do are usually
25
86
restricted to the NTK (94] regime [2 25 which precludes feature learning, and OFs. Lou et al.
[95] suggest that the feature alignment 96 phenomenon during training is correlated to the rate at which signal propagation converges to its limit in a deep NN. Even at initialization, the distribution of the neurons becomes more heavy-tailed with depth 1971. thus making outliers more likely. Noci et al. [24] gives a precise description of the kurtosis for ReLU networks, showing that grows exponentially with depth. Together with the results presented in this work, there is empirical and theoretical evidence that depth has the double effect of increasing both the correlations and making large activations more likely, which we observe to be detrimental to outliers. However, the theoretical treatment of the emergence of outliers during training is still an open question.
B Mathematical description of OP block
13 we present a mathematical description of the OP block (Fig 3 Recall the OP block replaces the Pre-Norm block by removing its normalisation layers in both Attention and MLP sub-blocks, and making three additional changes: 1) downweighting residual branches with some 3MLP, .Amm = (de/vhepth)
For completeness, in Eqs 12 and
based on the notation in Apple A.
recover Signal Prop benefits of Pre-Norms 23 27 30), 2) adding an Entropy Regulation mechanism to prevent Entropy Collapse e.g. QK-Norm or tanh-softcapping. and 3) (optionally) scaling the inputs before the MLP nonlinearity by scalar OMLP to ensure the nonlinearity inputs are of order 1, following Brock et al. [53).
Mathematically, this can be expressed as:
X = X + (LLLMM)() where x= X + MAi,
(12)
With QK-Norm as entropy regulation mechanism, MHA is multihead attention using QK-Nomm in the attention heads:
(1
Attn X)= A(X)XWV
where A A(X) Softmax
Nomm(X99)Nom(XX))
+ M
Vd
(13)
23

---

final
Activation Surtods
input-vise aslen
vation
0.2
022
633
538 1.5 3.1 7.06 3.6. 1:0 142 239
048
644 3
1.8 2.5
048 1.0
140 2.20 2:8 1.5 11 1:4
0.0
C55 100 1. ADA
Tokers seer
Tobens seen
ToCert ween
Figure 10: OP layers at 1.2B scale with worse Signal Propagation (i.e. higher imput correlations) during training (centre left) have higher feature kurtosis (left). (Right vs. left two plots) Introducing a final LN before unembedding causes larger input correlations and feature kurtosis in later layers, even with the OP block. NB: y-axes values here are significantly smaller than Fig with Pre-LN.
C Modifications That Affect Signal Prop During Training Affect OFE
In this section we continue our discussion of the connection between Signal Propagation and OFE in Sec towards the practical implication that methods designed to improve Signal Propagation also improve OFE as a result.
To the best of our knowledge, in the Signal Propagation literature, most works have focused on characterising and improving Signal Propagation at initialisation due to analytic convenience. At initialisation, only the architecture plays role and not the optimiser. In particular, a practical focus of such works is to design architectural modifications that allow non-degenerate deep limits for models whose input cosine similarities can be well approximated by their large- width limits at initialisation 198 93 23 25 28, 30). Those considering training dynamics often reside in the kernel regime [94] and are thus not compatible with feature learning [31 33 which is necessary for OFE and Signal Prop dynamics during training. Our results connecting Signal Prop and OFE highlight the importance to the community of understanding Signal Prop dynamics during training in feature learning regimes, beyond initialisation. We note Tian et al. [89] predict attention entropy dynamics through joint MLP/Atteentio. In any case, empirically study the impact of initialisation iitialliattionnissiree Signal Prop architectural modifications in terms of OFE during training.
Signal Prop modifications, the most prevalent is downweighting residual branches h(x)- = x + (z( with some 3 < 1 [51] 23 27
Downweighted residuals Of initialisation-inspiii
12
In Fig 5
we see that downweighting residuals (with trainable scalar 3 initialised to 0.1) improves Signal Propagation in 24-block 1.2B Pre-LN model, not only at initialisation but also during training, thereby reducing OFF (peak kurtosis an order of magnitude lower). Having said that, Pre-LN with downscaled residuals still leads to higher kurtosis across layers than our OP block in Fig TO] Downscaling Pre-LN residuals leads to a small loss in performance of 0.2 perplexity. We show the corresponding results at 130M scale in Figs 19 to 21 Our results are consistent with previous work by Wortsman et al. 10 who observe that downweighted residuals help for low precision training in CLIP models, motivated as way to prevent OFs arising through increasing activations scales ||X|| during training. Given that standard models have Norm layers that are scale invariant (as are our OFE and Signal Prop metrics), we complement this argument by highlighting that the feature learning process of OFE is not only associated with increasing activation scales but also worsening Signal Propagation during training. Figs and 48 show that [X] r does not always correlate with OFs.
To
Normalisation layers On the other hand, for Norms, the difference between OP and standard blocks with Norms Figs 5] 10 and [8 respecttvv is already clear evidence that standard Norm placements can lead to worse Signal Propagation (and OFE) during training. To the best of our knowledge, this observation has not been made previously. To test this further, we reintroduce the final LN right after the final OP block (just before the unembedding layer) into an OP model, with no Pre-Norms, Fig TO We see that the final LN causes some layers to see increases both kurtosis and input correlations, and these layers correspond precisely to the final few blocks immediately preceding the LN. On the other hand, earlier layers further away from the final LN are largely unchanged in terms of both Signal Propagation and OFE during training. The model with a final LN performed slightly worse (0.1 perplexity difference).
"Typically, the theory indicates that / O( /dkph enables well-behaved infinite-depth limit.
24

---

Several works have discussed the effect of Norms on Signal Propagation theory at initialisation. The Deep Kernel Shaping (25 framework is compatible with LN (and also RMSNorm) layers. but makes other modifications (in weight initialisation and activation functions) that mean LN has no effect at initialisation in the wide limit. Other works show centred Norms in fact improve Signal Propagation at initialisation in MLPs by correcting imbalances in input activation norms to improve Isometry [99 [000] but consider non-standard architectures that are not residual and have Norm immediately following nonlinear activations, whereas standard Norms take the residual stream as input. Our work shows that initialisation and training can have very different Signal Prop behaviours.
Other Signal Prop modifications In Figs 22 and 24, we consider the effect of other initialisation- inspired Signal Propagation modifications in terms of OFE. In particular, we consider "transforming" activations to be more linear [25 86 26). and "shaping" attention to be more identity-like (28)30]. Although not predicted by initialisation theory, we find that these modifications mostly also reduce OFE and improve Signal Prop during training as well as initialisation. The latter finding is related to the work of Bondarenko et al. [15] who show that "no-op" heads that place large attention weights on shared uninformative tokens encourage OFs: large attention weights on shared tokens also worsen attention, which can be seen as a "no-op" that
signal propagation 33 compared to
instead encourages token to attend to itself.
App C key takeaways: Signal Propagation and OFE.
Signal Propagation is fundamentally connected to OFE: worse Signal Prop generally implies higher kurtosis and vice versa, throughout training (Eq 5 Prop G.I and Figs 5
10 and
(8).
The OP block's mild OFs can be traced increasing input correlations while training (Fig
Choices that improve Signal Prop during training (e.g. scaled residuals) also reduce OFs (Fig
.
Removing standard Norms can improve Signal Prop, & OFE, during training (Figs 5 and 10
D Additional Experimental Details
CodeParrot As discussed, all of our experiments at 130M scale (6 layers, width 768 transformers)
are on next token prediction with the CodeParrot dataset, with 50K vocabulary size. We use a similar
setup to He and Hofmann (30, and have updated their codebase to include the OP block 14 We train with Adam optimiser
43 and weight decay 0.1, (31, 32) (0.9, 0.999), and e le 8 unless otherwise stated, and clip the maximum gradient norm to 1. We do not tie embeddings, and remove the final layer norm before unembedding layer. When plot metrics (kurtosis, signal propagation, MMR ett) we plot the residual stream entering the attention sub-block (plots for the residual stream before the MLP sub-block are qualitatively the same). The only exception is the last layer, which is the input to the unembeddings. When we downweight residuals we set / = 0.3 in both attention and MLP
unless otherwise stated. We do not train residual scalings 3. Unless otherwise stated, we train with sequence length 128 and batch size 32 for 80K steps. with linear warmup to maximum learning rate 1e 3. for 5% of the steps, before linear decay. We keep the standard parameter initialisations to N(0, std = 0.02) but upweight the input embeddings by a factor of 50 in order make the average squared input at initialisation, similar to considerations made by the Gemma model [01]. We use ReLU activations and do not scale inputs with an a, c.f. Fig because ReLU is I-homogeneous.
Languini For Languini 33 our 100M, 320M, and 1.2B model sizes follow the "small" (depth 12, width 768), "medium" (depth 24, width 1024), and "XL"" (depth 24, width 2048) model sizes provided by the authors, respectively. Our setup follows the authors in terms of codehase and tokeniser. We train with sequence length 512 and batch size 128, again with a maximum learning rate of 1e unless otherwise stated. This learning rate was the largest stable and best performing choice on logarithmic grid. We train with AdamW, using weight decay of 0.1 and clip the maximum gradient norm to 0.5. We use linear warmup and linear decay after 1000 steps. We additionally use RoPE 102 with GeLU nonlinearities in the MLPs. We use the same method as Brock et al. [53 to
!'Consider the extreme case when attention weights are placed onto single token (say the first one): all attention outputs will be equal to the first token's value representation so all token-wise cosine similarities are Mhttps:/ / //t///gittt git......... http/bbbby--e//immmlitete trwpeiddddtrasssseeerere
25

---

calculate a to scale inputs to the GeLU. When we downweight residuals, we initialise 3 = 0.1 and allow them to be trainable. When we plot layer-wise metries like kurtosis, we plot the outputs of the Pre-Nonmalisation layer (if there is one), otherwise, we treat the Normalisation layer as the identity and plot the residual stream going into the attention sub-block. We use tied embeddings. We also keep the standard parameter initialisations to (0, std = 0.02) but upweight the input embeddings by a factor of 50 in order to make the average squared input I at initialisation.
For our MLP experiments on CIFAR-10, we train using batch size 2048 for 200 epochs. As described in Seeß] the model has 6 Pre-Norm layers with width 1024, giving 15M parameters. We zero initialise the last layer, and additionally downweight the output layer by / akin to μP [33], to encourage feature learning. We train with MSE loss and use LR 3 for SGD and 3e-3 for Adam. We use standard betas and epsilon for Adam and we do not use weight decay, We warm up the LR for 200 steps before cosine decay. We additionally found that it was important to whiten the inputs in order to observe OFE in the residual stream. We note that transformer embeddings are independently initialised, which can be thought of as implicitly whitening the embeddings for different tokens. Whitened inputs correspond to signal propagation with zero input correlations. This again suggests that signal propagation (and properties of the data) are important for OFs, but we leave further understanding of this to future work. We use PCA to whiten inputs.
CIFAR-10
Non-diagonal preconditioners For our non-diagonal preconditioning experiments, we take the implementation provided by SOAP (6) is We keep the default hyperparameters 31 0.95, 82 0.95 and preconditioning frequency to 10, as suggested by the authors. As in our other experiments, we use weight decay 0.1. It is likely that further hyperparameter tuning would further improve our results with diagonal preconditioners.
For our implementation of "AdaFactor in Shampoo's ' eigenbasis" (shown by [69 to be akin to Shampoo [67]). we do not rotate the input emhedding dimension, which has (vocabulary) size 50k in our CodeParrot setup, but continue store rank-1 moving averages for the squared gradient in that dimension, which are applied to the unprojected parameters in that dimension, as in AdaFactor. For both AdaFactor and its rotated counterpart, we do not use the parameter-norm dependent learning
rate schedule proposed by [65]. instead keeping the same update rule as [103 104 69] AdaFactor closest rank-1 approximation of Adam.
which keeps
For Fig we set a larger batch size of 4096 in order to demonstrate more clearly the convergence speed per step benefits of SOAP/Shampoo 773), as seen in Fig 32 We additionally use the Warmup- Stable-Decay (WSD) scheduler [105] 106) with LR=0.001 (we found this also to be the maximal trainable LRR for SOAP/Shampoo in our setting), 5% warmup and 20% decay steps. WSD maintains a higher LR for longer and increases kurtosis in OF-prone settings like Adam/ Aaaa/d as a result, as expected from e.g. Fig 6
Web-EEu Our longer 1.2B and 7B experiments (Fig 9 using Wed-Wud 41 were con- ducted on fork of nanotron [6] The model architectures are based off LLaMa (47 but we use single hidden layer GeLU MIPs instead of SwiGLU, as SwiGLU MLPs are themselves source of OFs [107]. Our 7B model has hidden width 4096 and 32 layers, and our 1.2B model has hidden width 2048 and 24 layers.
For our 6B token 7B model runs, we train for 20K steps with global batch size 72 (node batch size on 12 nodes of 4 GH200 GPUs, with tensor parallelism) and context length 4096. This gives 72 x 4096 x 20000 x 6B tokens. For our 90B token IB model runs, we train for 50K steps with global batch size 440 (device batch size on 22 nodes of 4x GH200 GPUs) and context length 4096 to give 440x x 4096 50000 90B tokens. We add LayerScale [55] to implement the downweighted residuals in OP block, with trainable residual gain vectors initialised to 0. and 0.05 for our 1. 1.2B and 8B experiments respectively. We use RMSNorm for the QK-Norm in the OP block. Like in our Languinn experiments, we upweight the input embeddings by factor of 50 in order to make the average squared input at initialisation.
Quantisation As discussed, our quantisation experimental setup closely follows Bondarenko et al. 15 for fair comparison, training OPT models [ on BookCorpus and English Wikipedia before post
15 ahatppr//pccabbbcc////iaiaaa///p
26

---

training quantisation. We also use their excellent codebase 17 are outlined in 2 Other hyperparameter are set defaults e.g. (3,, 52) (0.9, 0.95), 0.1 weight decay, or ReL.U activation function, as in Zhang et al. 66
The hyperparameters that we change
reads down the rows for a given architecture, the optimisation hyperparameters are changed one after another on top of each other. The only exception is SOAP, where we keep AdamW's s e 10- (recall SOAP is just Adam in the eigenbasis of Shampoo's preconditioner, so e is still used). The model architecture is OPT 6 and have 125m parameters with width 768 and 12 layers. We initialise the trainable scalar residual gains to 0.2 in the OP block and use LN for the QK-Norm Like [[]], we use symmetric quantisation for the weights and asymmetric quantisation for the activations int8, but also do not quantisation the final unembedding layer. We refer the reader to Bondarenko et al. [ for more details on the quantisation procedure.
In Tab 2
E Additional Experiments
In this section, we include all additional experiments not included in the main paper.
Pythia 160M
Pythia 1.48
Pythia 6.98
Avernge MMR
Layer
20
11
22
20k
40c
10c
230k
30k
402
60t
82c
130k
20k
42c
Edk
022
120k
Step
Step
Step
Figure 11: Max Median Ratio metric for Pythia, equivalent to Fig over inputs
We take the mean to aggregate
RMS of Inpattenit correlations
Pythia 160M
Pythia 1.4B
Pythia 6.98
0.
14
0.4
28
LIVe
17
LES: Layer
Step
Step
Step
Figure 12: Signal Prop for Pythia, equivalent Fig
Early layers
Middle layers
Final layers
200
Prcch
Layar
175
10
33
Ure
Peeth Pre- 23
Average MYR
350
hhuhh
OP
325
200
6.50
3,00
2.50
3.00
3.50
4.00
0.53
1.00
3.50 1.50 2500
3.00
3.52 1.00 1.53 2.00 3.50
3.00 3.53 4.00 Tokens seen
Takens seen
Tokens seen
Figure 13: Average MMR metric comparing Pre-LN and OP blocks at 1.2B scale on the Languini Books dataset 39 equivalent Fig
27

---

1st Block
2nd Block
3rd Block
4th Block
Activation Kurcosis
10a
101
109
10-1
10
20K
40K 60K
B0K
20K
40K
60K
B0K
20K
40K
60K
B0K
20K
40K
60K
80K
Training Step
Training Step
Training Step
Training Step
5th Block
6th Block
Unembedding
Activation Kurtosis
10a
101
Pre RMS
Post-LN
Pro SRMS
10-1
(ours)
10
20K
40K
60K
BOK
20K
40K
60K
BOK
20K
40K
60K
BOK
Training Stop
Training Step
Training Step
Figure 14: Kurtosis dynamics in different layers using different Norms and Norm locations CodeParrot at 130M scale. Equivalent of Fig 2 but for the remaining layers. Fig corresspons to the 2nd block.
1st Block
2nd Block
3rd Block
4th Block
Activatton Kurtocis
102
101
1000
10-2
20K
40K
60K
BOK
20K
40K
60K
80K
20K
40K
60K
80K
20K
40K
60K
00K
Training Stop
Training Step
Training
Training
5th Block
6th Block
Unembedding
Activatton Kurtoeis
102
10a
Pre-SIMS
100
Pro
10-2
Post
20K
40K
60K
80K
60K
00K
20K
40K
60K
BOK
Training Step
Training Step
Training
Figure 15: Equivalent of Fig 14 but with centred activations (centred along the width dimension).
Notice there is qualitative difference kurtosis dynamics when centring activations.
28

---

1st Block
2nd Block
3rd Block
4th Block
(78 |X1r
Activation scale
40K 60K
20K
40K
60K
Training Step
Training Step
Training Step
Training Step
5th Block
6th Block
Unembedding
Pre- RMS
Activation scale (18
Pro SRMS
(ours)
20K
60K
20K
40K
60K
Training Stop
Training Stop
Training Step
Figure 16: Equivalent of Fig 14 but for activation scale X trajectories through training. We see that activation scales do not correlate as well with OFs (Fig 14] as signal propagation (Fig[I8). For example, Post-LN has smaller activation scales than the OP block in all blocks besides the first one, much worse kurtosis in Fig 14
Block
2nd Block
3rd Block
Block
MMR (maz over batch Inpics
6.00
400
300
200
20K
40K
60K
B0K
20K
40K
60K
80K
20K
40K
60K
80K
20K
40K
60K
BOK
Training Step
Training Step
Training Stop
Training Step
MMR (max over batch Inputs)
Block
Unembedding
RMS
Post-LN
Pre SRMS
400
300
OP (ours)
200
20K
40K
60K
BOK
0
20K
40K
60K
BOK
20K
40K
60K
80K
Training Step
Training Step
Training Stop
Figure 17: Equivalent of Fig 14 but for the MMR metric (aggregated using maximum over the batch).
29

---

RMS of inpirt wise correlations
1st Block
2nd Block
3rd Block
4th Block
0.8
0.6
0.2
20K
40K 60K
80K
40K 60K
80K
60K
S0K
80K
Training Step
Training Step
Training Step
Training Step
RMS of Inpirt wise correaatons
5th Block
6th Block
Unembedding
0.8
Pre RMS
Post-LN
0.6
Pre LN
0.2
20K
40K
60K
60K
80K
Training Stop
Training Stop
Training Step
Figure 18: Equivalent of Fig 14 but for Signal Propagation (in terms of RMS of input correlations).
1st Block
2nd Block
3rd Block
4th Block
Activation Kirbosis
10a
101
1.0
10-1
10-2
20K Training Step
Training Stop
Training Stop 6th Block
Training Stop
5th Block
Unembedding
Activation kurbosis
102
20a
Resid scale 0.1
100
10-4
0.7
10
60K
60K Training Step
Training Step
Training Step
Figure 19: Downweighted residual scalings, h(x)= =z)==++f((2 with 3 < 1, reduce OFs 130M scale on CodeParrot. All models are Pre-LN, We downweight both MLP and Attention residuals.
30

---

RMS of input- wise correlations
1st Block
2nd Block
3rd Block
4th Block
0.4
20K
40K 60K
80K
40K 60K
80K
60K
80K
BOK
Training
Training Step
Training Step
Training Step
RMS of input-wise wite conrelations
5th Block
6th Block
Unembedding
scale
0.1
0.6
0.7
0.4
20K
40K
60K
80K
40K
60K
80K
60K
80K
Training Step
Training Step
Training Step
Figure 20: Residual scalings improve Signal Prop at 130M scale. Equivalent to Fig 19
1st Block
2nd Block
3rd Block
4th Black
Activation scale 2 1X
20K 40K
60K
20K
40K
60K
BOK
40K
60K
BOK
20K Training Step
80K
Training
Training
Training
5th Block
6th Block
Unembedding
Activatio scale 2 14
0.1
0.7
20K
40K Training Step
20K
40K Training
40K
BOK
Training
Figure 21: Residual scalings reduce activation scales at 130M scale. Equivalent to Fig 19
31

---

1st Block
2nd Block
3rd Block
4th Block
Activation Kurtoris
300
200
10K
15K
20K
10K Training Step
10K
15K
Training Step
Training Step
Training Step
5th Block
6th Block
Unembedding
Lewky/ReLU slope
Activaaton Kurtoris
300
0.0
200
0.4
0.8
100
10K 15K
20K
Training Step
Training Stop
ng Step
Figure 22: Increasing LeakyReL.U slope, so that the nonlinearity is more linear mostly improves kurtosis during training, as one might expect from Signal Prop initialisation theory [86] 26] Here our notation is ((akyReLU(z) = max(æ, sx) for slope < 1. The exception is when the slope is 0, i.c. ReLU, the kurtosis is actually better during training, but this is reflected in the signal propagation during training too (3). We hypothesise this is because zero neurons get no gradient with ReLU, and this behaves fundamentally differently to non-zero LeakyReLU slope. The plots show the average over 5 seeds, and we plot the first 20K steps (of 80K). The models are Pre-LN and we downweight the attention residual branch with factor 3 0.2 to reduce kurtosis contributions from the attention sub-block, but do not downweight the MLP residual. Note we do not use log-scaled y-axis to make the di differences between LeakyReLU slopes clearer. Experiment 130M scale on CodeParrot.
RMS input- W se correlations
1st Block
2nd Block
3rd Block
4th Block
0.8
0.6
0.4
0.2
5K
10K
15K
20K
5K
10K
15K
20K
5K
10K
15K
20K
10K
15K
20K
Training
Training Stop
Training Stop
Training Stop
RMS of Inpur-wie 50 correlations
5th Block
6th Block
Unembedding
0.8
LeakyReLL slope
0.0
0.6
0.6
0.4
0.2
5K
10K
15K
20K
5K
30K
15K
20K
5K
10K
15K
Training
ninningg Step
Training Step
Figure 23: Effect of different LeakyReLU slopes on signal propagation during training, equivalent to Fig 22 Surprisingly, ReLU (i.e. slope 0) has the best signal propagation (lowest input-wise correlations) during training in this setting, even though it has the worst signal prop at initialisation in later layers, compared all other LeakyReLU variants. This initialisation effect was predicted by Zhang et al. [86), but our findings regarding training were previously unknown and require further research.
32

---

Block
2nd Block
Block
Block
250
Activation Kurtosis
150
10K
15K
20K
5K
10K
15K
20K
5K
10K
15K
20K
10K
15K
20K
Training Step
Training Scep
Training Step
Training Step
5th Block
6th Block
Unembedding
250
Value Soip-Skippnit
Activatio Riroois
0.7
150
1.5
15K Training Step
10K
20K
5K
10K
15K
20K
5K
10K
15K
20K
Training Step
Train maaining Step
Figure 24: Reducing / in Value-SkipInit (8)) which replaces Attention matrix A < aI BA and makes attention more identity-like also reduces OFs. We do not train 3 in Value-Skiplnit and fix 1. The models are Pre-LN and we downweight the MLP residual branch with factor 0.2 reduce kurtosis contributions from the MLP sub-block, but do not downweight the attention residual. Each curve an average over 5 seeds and we plot only the first 20K steps (of 80K). Experiment is 130M scale on CodeParrot.
Early loyers
Middi layers
layers
Laplr
Small
Mo1h
Figure 25: Smaller LR (max value from 0.001 - 0.0006) reduces OFE in a Pre-LN model at 1.2B scale on Languini [39]. Models are slightly different from the Pre-LN model in Fig as we do not upweight the input embeddings described in A Still, we do also observe large increases in kurtosis during training. and that smaller LR reduces this. In this experiment, reducing the max LR to 0.0006 did not impact convergence speed.
33

---

1st Block
2nd Block
3rd Block
4th Block
Activation Kurtosis
10a
101
10-1
10
20K
40K 60K
00K
20K
60K
80X
20K
40K
60K
B0K
20K
40K
B0K
00K
Training Step
Training Step
Training Step
Training Step
5th Block
6th Block
Unembedding
IR=1e-3
Activation Kurtosis
102
UR=3e
Pre-RMS
101
Post-LN
Pre-IN
10
20K
40K
60K
B0K
20K
40X
60K
BOX
20K
40K
60K
BOK
Training Step
Training Step
Training Step
Figure 26: Smaller LRs means reduced OFs, for different Norms and Norm locations. Equivalent of Fig but with all layers. Experiment on CodeParrot at 130M scale.
3.00
LR=1e-3
Pre-LN
LR=3e-4
Post-LN
2.75
OP (ours)
Pre-SRMS
Eval loss
2.50
Pre Pre-RMS
2.25
2.00
1.75
1.50
20K
60K
Training Step
Figure 27: Convergence speed for the runs in Figs 6 and 26
comparing the effect of reduced LRs.
4.0
Adam eps
1e-10
1e-06
Train Loss
3.0
1c:05
3e-05
2.5
1e- 04
3e 04
2.0
20K
40K
60K
80K
Training Step
Figure 28: Train loss plot with different Adam epsilon, equivalent to Fig 29 There not a noticeable difference in convergence speed for 3e in this experiment.
34

---

1st Block
2nd Block
3rd Block
4th Block
Activation Kurtosis
10a
101
10-4
10-2
60K
20K
40K
60K
Training Stop
Training Step
Training Step
Training Stop
5th Block
6th Block
Unembedding
ops
Activation Kirzosis
10a
1e:10
101
---
1e 06
1o-04
10-4
10-2
20K
40K
80K
20K
40K
60K
80K
20K
40K
80K
Training Stop
Training Stop
Training Stop
Figure 29: Kurtosis plot with different Adam epsilons CodeParrot at 130M scale. Each curve is an average over seeds. We see that increasing from le- to 3e 4 monotonically decreases OFE. At values of smaller than le- there is less of difference in OFE between different values.
RMS of Imput-wwe wise corelations
1st Block
2nd Block
Block
Block
0.8
0.6
0.4
0.2
20K
40K
60K
BOK
20K
40K
60K
80K
20K
40K
60K
8OX
20K
40K
60K
BOK
Training Step
Training Step
Training Step
Training Step
RMS of Imput WN Se correiatons
5th Block
6th Block
Unembedding
0.8
Adam
10:10
1e:0 08
0.6
le 05
0.4
0.2
3e-04
70K
40K
60K
80K
20K
40K
60K
BOK
20K
40K
60K
BOK
Training Step
Training Stop
Training Step
Figure 30: Signal Prop plot with different Adam epsilon. Equivalent of Fig
35

---

1st Block
2nd Block
3rd Block
W
Activation KurtoR
60
20
Training Step
Training Step
Training Step
4th Block
5th Block
6th Block
AdaFactor
Activation kurtoeis
AdamW
AdaFactort Shampoo)
60
Rotated AdOMPN(-SAPP)
20
Training Step
Paining Step
Training Step
Figure 31: Kurtosis evolution in all 6 layers of a Tre-SRRMNorm transformer trained on CodeParrot with different optimisers. Equivalent to Fig
2.4
AdaFactor
2.2
AdamW
Rotated AdaFactor( Shampoo)
2.0
Train loss
Rotated AdamW(=SOAP)
1.8
1.6
1.4
1.2
1.0
0
1K
2K
3K
4K
Training Step
Figure 32: CodeParrot training loss curves for experiments comparing diagonal and non-diagonal preconditioners. Equivalent to the runs found in Fig g
36

---

RMS of input-uw wi -i correlatons
1st Block
2nd Block
Block
Adafactor
0.8
Rotated AdaFactori - Shampoo)
0.6
Rotated Adami SOAP)
0.4
0.2
1K
4K
3K
2K Training Step
Training Stop
Training Step
RMS of inpurt wise correlations
4th Block
5th Block
Block
0.8
0.6
1K
JK
4K
JK
2K
JK
Training Step
Training Step
Training Step
Figure 33: RMS of wnse correlation (i.e. signal propagation, equivalently the RHS of Eq (5) trajectories in all layers Pre-SRMSNorm transformer trained on CodeParrot with different optimisers. These optimisers the diagonal Adam and AdaFactor, and their non-diagonal rotated versions (SOAP and something close to Shampoo, c.f. [69]), like in Sec Notice that the connection between signal propagation and kurtosis identify in Sec is not apparent with non-diagonal preconditioners. In other words, the non-diagonal rotated versions of Adam/Ada Aaat/or have higher input-wise correlations than diagonal Adam/AdaFactor despite having better kurtosis properties (seen Fig 31 This observation is explained by Fig 34 which looks at the corresponding feature-wise correlations. These training runs are equivalent to the runs found in Fig
RMS Feature- wice correlations RMS Feature- wice correlations
1st Block
2nd Block
3rd Block
Adafactor
0.8
AdamW/
Rotated Adafactori Shampoo)
0.6
Rolated Adammi SOAP)
0.4
0.2
3K
2K
Training Stop
Training Step
Training Step
4th Block
5th Block
6th Block
0.8
0.6
0.4
0.2
2K Training Step
Training Step
Training Step
Figure 34: RMS of feature-wise correlation trajectories, equivalent to Fig Or more explic-
itly, Enjsdat (Er)?s in the notation of Eq 5 Notice that for non-diagonal preconditioners (SOAP/Shampoo), the curves in Fig 34 closely track the signal propagation trajectories in Fig 33 (and both relatively high), meaning that their difference (the kurtosis in Eq (5) is small. On the other hand, for diagonal preconditioners (Adam/AdaFactor), the trajectories Fig 33 are noticeably higher than in Fig 34 and this difference is precisely the increased kurtosis. As discussed in See theoretically predicting these feature learning behvaiours (dependent on the choices of architecture and optimiser) during training is beyond the scope of current tools in deep learning theory.
37

---

E.1 Image Classification Experiments
E.1.1 Adam SGD
In Sec[5] we saw that the diagonality of the preconditioning in Adam and AdaFactor was important for the emergence of outlier features. To test this further, we consider the effect of replacing Adam with SGD on OFE. SGD does not precondition gradients equivalently, SGD preconditions gradient with the identity matrix. In comparison to optimisers like SOAP or Shampoo that diagonally precondition in rotated parameter space, SGD also be thought optimise in rotated parameter space (for every possible rotation), precisely because preconditions with the identity matrix.
As transformers are difficult to train (fast) with SGD, we consider OFs in a much simpler architecture and task: an MLP on CIFAR- 10 image classification. Like with Shampoo and SOAP (Fig [). in Fig 35 we see that SGD is not susceptible to OFs as Adam, even with OF-prone architecture choices, like Normalisation layers. fact, in this experiment SGD kurtosis actually decreases during training with Pre-Norm. Fig shows that SGD matches Adam convergence speed in this setting. The model is a 6-layer Pre-Norm residual MLP of width 1024; we remove Pre-Norms for normless models. This also highlights that OFs are not specific to the Transformer model.
1st Layer
2nd Layer
3rd Layer
Adam, Pre-SRMS
Activation Kurtosis
Adam, Normless
SRMS
SGD, Normless
1K
2K
3K
Training Stop 4th Layer
Training Step
Training Step
5th Layer
6th Layer
Activation Kurtoiis
2K
3K
Training Step
Training Step
Training Step
Figure 35: OFs of SGD vs Adam in an MLP on CIFAR-10. Although normalisation layers lead to higher kurtosis for given optimiser, Adam always has higher OFs than SGD.
MLP on CIFAR-10
80
Training Accuracy
70
60
Adam, Pre-SRMS
SGD, Pre SRMS
SED,
Training Step
Figure 36: Train accuracy plot with SGD Adam of MLP CIFAR-10, 10, corresponding Fig
Adam e is the default value 1e-
38

---

We note that the levels of kurtosis reached by AdamW in CIFAR-I0 image classification MLP settings do not reach the same heights (peaking around 40 in individual layers in F as in our transformer language modelling settings, [88] Having said that, we also show in this subsection that many of our findings concerning OFs (in terms of kurtosis) carry through to the image classification setting, including: the effect of Adam (Figs 37 and 38) and the correlation between signal propagation and kurtosis for Adam and SGD optimisers (Fig 39
E.1.2 Vision Transformer Experiments
To assess whether the difference in architecture between our language modelling and image clas- sification experiments explains this observed difference in peak kurtosis, we consider the Vision Transformer (ViT) 42), which is largely the same Pre-LN transformer architecture in our lan- guage modelling experiments. Instead of processing sequences of tokens like its language modelling counterpart, ViTs turn an image into sequence of patches, and alternate self-attention and MLP blocklocks (as in App A.I to process patches.
We consider the ViT-B model, which has 12 blocks and width 768 giving 86M parameters, which around the same parameter scale as our 130M CodeParrot experiments and also Pythia- 160M in Fig 19 We train on ImageNet-1K 1 using the codebase 20 and default hyperparameters of DelT [108] but only train for 150 epochs (instead of 300) and do not use distillation in the interests of compute time and because in this experiment we are not interested in state of the art results but rather the OF properties of ViTs. We use Adam W optimiser with a max LR of 3.75e-4 (maximum stable LR on a logarithmically spaced grid) with batch size of 384 (data parallel across 6 RTX2080-Tis) and warm up the LR for epochs before cosine decay.
We compare 3 different ViT transformer blocks: Pre-LN, Pre-LN with LayerScale (55) (LayerScale is a variant of SkipInit [51] that downweights the residual branches with vector of trainable gains), and our OP block (using LayerScale to downweight residual branches). We use the default initialisation for residual branch downweighting factors LayerScale, of le 4. QK-Norm implemented with LN for the OP block.
In Fig 40 averaged across layers) compared to the OP block (around 4.5 averaged across layers), but the Pre-LN kurtosis decreases sharply as training progresses, unlike in language modelling settings, and actually ends up being lower than the OP block at the end of training. In addition, the kurtosis values on this ViT image classification task are again far lower than transformer language modelling counterparts at similar parameter scales (e.g. [ or Fig 14]. Adding LayerScale (55) reduces Pre-LN kurtosis (to even below that of the OP block in terms peak value), as expected from our findings in Sec In
see that Pre-LN ViT-Bs do still suffer from higher peak kurtosis values (around 20
Fig 41 we see that all three blocks have similar test accuracy performance, with the OP block having a slight advantage of 0.2%.
From these preliminary experiments with ViTs we conclude that, despite similarities, there are qualitative differences between OFs (measured via kurtosis) in image classification and language modelling tasks that are not explained through the choice of architecture (or indeed the choice of AdamW optimiser). We leave a more in-depth study to future work. remains be seen if OFs will be a barrier to quantising ViTs at scale, as it has been the case for LLMs, though our findings showing reduced kurtosis (both after and peak during training) suggest this may not be the case. We note Darcet et al. [79 identify artifacts in the attention maps of ViTs, which may potentially be related to OFs.
1his difference is necessarily explained by scale, our MLP residual stream width of 1024 is actually wider than our 130M transformer width of 768, which is the scale that produced the much higher kurtosis values in, for example, Figs 2 and [4]
This is actually deeper than default 6 layer transformer in 130M CodeParrot experiments, and shares the same width. The parameter difference with language that a significant fraction of parameters in language modelling are in the embedding layer.
https://githubbc "ttpps//gituub.co/ffcbbbboo
39

---

1st Layer
2nd Layer
3rd Layer
Adam ups
Activation Kartosis
203
1e416
1e-05
1e-04
Training
Training 5th Layer
Training Step
Layer
6th Layer
10a
Activation Kurtodis
20-1
10-4
2K
4K
Training
Figure 37: Kurtosis plot with different Adam with MLP on CIFAR-10. The model uses Pre-Norm structure with SRMSNorm normalisation. Like in Fig 29 we see that larger generally leads to smaller OFs.
MLP on CIFAR-10
Training Accuracy
70
Adam eps
1e-10
1c-08
1c-07
1e-06
1e-05
1e-04
10
1K
2K
3K
4K
Training Step
Figure 38: Train accuracy plot with different Adam of MLP on CIFAR-10, equivalent to Fig 37 this experiment, milder values of c E (le 5, le 6) converge fastest.
In
40

---

RMS of input-wiie wise correl ations
1st Layer
2nd Layer
3rd Layer
Adam, Pro-SRMS
Adam, Normless
Pre-SRMS
0.3
SGD, Nomiles
0.2
0.1
Training Step
Training Step 5th Layer
Training Step
RMS of Inpur.wiii wise correlations
4th Layer
6th Layer
0.3
0.2
Training Step
Training Step
Taaining Step
Figure 39: Effect of SGD vs Adam on Signal Prop, for models plotted in Fig 35
VÎT-B Average Kurtosis on ImageNet-1K
20.0
Model
OP
17.5
Pre-LN
15.0
Pre-LN LayerScale
Average Kurtosis
12.5
10.0
7.5
5.0
0.0
20
40
60
80
100
120
140
Epoch
Figure 40: Average Kurtosis trajectories of ViT's trained on ImageNet-1K: Y-axis is average kurtosis across the 12 residual stream layers.
41

---

ViT-B Test Accuracy (top1%) ImageNet-1K
Test Accuracy [top1%)
50
30
20
Model
OP
10
Pre-LN
Pre- LN- LayerScale
20
40
60
80
100
120
140
Epoch
Figure 41: Test accuracies of VITs trained on ImageNet-1K K. The final test accuracies are 79.5% for OP, 79.3% for Pre-LN + LayerScale, and 78.7% 7% for Pre-LN.
42

---

E.2 Ablating the components of the OP block
In Tabs 3 and and Fig 42 we ablate the components of our OP block. Tab assesses the impact of not having an EntReg mechanism on training stability and convergence speed on the Languini dataset 39 at 320M scale. 42g confirms the loss of EntReg causes entropy collapse on CodeParrot at 130M scale, which is shown to lead to unstable training in Fig 43 In these experiments, we also try the tanh thresholding as an alternative EntReg mechanism to QK-Norm. Tab goes from Pre-LN to OP one step at time, assessing the impact of different norms and downweighted residuals, in terms of OFE.
Table Ablating the convergence and training benefits the OP block. The asterisk denotes that training failed without Flash Attention [109], which centres pre-softmax logits based their max value and therefore more stable. This highlights the training instability of not having some entropy regulating (EntReg) mechanism, where smaller LRs required for stability. At smaller (but stable) LR, the naive unnormalised model without EntReg converges much slower (17.4 vs 16.2 ppl) in this example. Even with larger LR, the EntReg mechanism in the OP block improves convergence (16.6 vs 16.2 ppl for QK-RMSNorm) compared to the naive unnormalised model. Tanh thresholding (from Grok-1) also works an example of alternative EntReg mechanism to QK-Norm. Because Pre-Norms appear before Query/Key weights, they already provide an implicit EntReg mechanism. As result, adding EntReg Pre-Norm models results in only minor changes convergence speed in this experiment (though ViT-22B shows in other settings Pre-Norm alone is not enough (59))
Models are 320M parameters, trained also for 3.3B tokens on Languini [39
as Tab !
Model
MLP/Attn Pre-Norm
EntReg
Scaled Residual
LR
Eval PPL
Pre-LN Pre-RMSNorm
LN
None
Implicit
le-3
16.2
RMS
None
Implicit
le-3
16.3
Pre-LN+QK-Norm
LN
QK-RMS
Implicit
le-3
16.0
Pre-LN+Tannh
LN
Tanh
Implicit
le-3
16.2
Naive unnormalised
None
None
Yes
3e-4
17.4
Naive unnormalised
None
None
Yes
le-3
16.6*
OP (QK-Norm)
None
QK-RMS
Yes
le-3
16.2
OP (Tanh)
None
Tanh
Yes
le-3
16.4
1st Block
2nd Block
3rd Block
Attn entropy
20K
60K
BOK
20K
40K
60K
40K
BOK
Training Step
Training Step
Step
4th Block
5th Block
6th Block
Normless
Atti entropy
GP (OK (oK-Nomm)
20K
60K
40K
60K
40K
8OK
Training Step
Training Step
Training Step
Figure 42: No EntReg leads to entropy collapse without Pre-Norms, which means training fails (as
seen in Fig 43
43

---

14
12
Pre-LN
Normless
Train Loss
10
OP (Tanh)
OP (QK-Norm)
6
4
2
20K
40K 60K
80K
Training Step
Figure 43: Entropy collapse leads failed training. Experiment 130M scale CodeParrot. OP with tanh does not fail but does converge slower in this setting; compared to Tab 3 we use learnt positional encodings in the input embedding layer, difference. We tuned a few values of the mar _att__tttt_ooo r_al hyperparameter with tanh thresholding: f(z) = mar max_attn _attt_ttnaa settnvall-- tanh(z/ /aa_ maar_att__uue _attn_val),, which is set by default 30 in Grok-1 but they did not close the convergence speed loss.
RoPE, which may account for this
1st Block
2nd Block
3rd Block
Activation Kurtosi
10a
10a
10-1
10-2
20K
40K
60K
B9K
20K
40K
60K
00K
20K
40K
60K
BOK
Trai ning Step
Training Step
Training Step
4th Block
Block
6th Block
Activation Kurtoss
102
10P
Normless
(Tanh)
(OK Norm)
10-2
20K
40K
60K
BOK
20K
40K
60K
BOK
20K
40K
60K
8OK
Trai ning Step
Training Step
Training Step
Figure 44: OP with Tanh still has reduced peak OFs compared Pre-LN. This plot corresponds the models shown in Fig 42
44

---

RMS of Input WiGe correiatins
1st Block
2nd Block
Block
0.6
0.4
0.2
20K
40K
60K
BOK
20K
40K
60K
BOK
20K
40K
60X
BOK
Training Step
Training Step
Training Step
RMS of Imputt-wp - Se correiatti ctrrelations
4th Block
5th Block
Block
Pre-LN
Normless
0.6
OP (Tank)
(OK-Nom)
0.4
0.2
20K
40K
60K
B0K
0
20K
40K
60K
BOK
20K
40K
60X
BOK
Training Step
Training Step
Training Step
Figure 45: Signal Prop plot with OP Tanh. This plot corresponds to the models shown in Fig
Table Going from Pre-Norm OP step by step. We remove or add Norms one by one, with different Norm locations depicted in Fig 47 All models trained well (at similar speeds), they all have some form of entropy regulation (either explicit or implicit) and downweighted residuals. We present the peak Kurtosis (Eq (T). Signal Propagation (RMS of input-wise correlations), and activation RMS (|X F) over the training run, with mean and standard deviation over three seeds. We present results where activations X are the input to the second transformer block. We see that that preventing attention entropy collapse through QK-Norm helps reduce OFs (which we see coincides with improved signal propagation). On the other hand, peak activation RMS does not correlate well as a metric with peak kurtosis, across the different models. In addition, the 2 best models in terms of OFs (our OP and also the third last row, which has Pre-V Pre-MLP Norms) (at least at initialisation), which implies that the fact that Pre-V V or Pre- Norms make the residual stream scale independent detrimental for OFE. This is corroborated by Fig 46 which plots the trajectories for the three models (1. Post-QK+Pre-VV 2. QK Norms only and 3. OP) that achieved peak kurtosis lower than 10. Fig [6] shows that the non-homogeneity (due Pre-V Norm) leads to a large initial increase in kurtosis and signal propagation in this setting, like we consistently see with Models are 130M scale on CodeParrot.
1 1-homogeneous
Pre-Norm blocks e.g. Fig 5
Model
Norm
Scaled Resid
Homog.?
Act RMS
Signal Prop
Kurtosis
Post-QK
Pre-OK
Pre-V
Pre-MLP
Pie-RMS
None
RMS
RMS
RMS
Implicit
No
5,45
0.72 co
131.8 131.84m2
Scaled Resids
None
RMS
RMS
RMS
Yes
No
3.97
0.47
46.4415.
Norms
RMS
RMS
RMS
Yes
No
3.92.a.07
0.24_o.co
12.7410.
Attn Norms only
RMS
RMS
RMS
None
Yes
No
4.38
0.29
11.8
Post- QK+PreVV V
RMS
None
RMS
None
Yes
No
4.40gg.06
0.27
6.441.32
QK Norms only
RMS
RMS
None
None
Yes
Yes
1.32+0.06
0.15
2.540.95
QK- only
None
RMS
None
None
Yes
Yes
4.38
0.37
61.04.05
(ours)
RMS
None
None
None
Yes
4.4640.00
0.17=0.0z
4.341.00
45

---

RM5 of Ingut Conrel attons
0.215
oPfcano Pom-oketree
0.250
0.225
0.200
Artivation scale
0:355
0.25
0.225
(Pun
0.20
Fre @k+imel Worne
1.5
0.755
coC
20C
60K
S0K
20K
Joc
Training Step
Trainingg Step
Training Step
Figure 46: Training trajectories of kurtosis, signal propagation and activation scales for the three best configurations in Tab 4 The setting with Pre-V Norm (which is not I-homogeneous) sees a large initial increase in all metrics, with kurtosis and input correlations peaking within 10K steps before reducing during training.
MLP Out
NonLin
MLP In
Pre-MLP Norm
Proj
Hx
Attention
Post-QK Norm
Pre-QK Norm
Pre-V Norm
Figure 47: A transformer block with many different Norm layers depicted, help parse the ablations we consider in Tab Note we break down the standard attention Pre-Norm into Pre-QK Norm and Pre-V Norm because removal of Pre-V Norm makes the attention sub-block homogeneous (i.e. f()) homogeneous if fkkz) = kf(x) for some scalar & > 0), hence acts differently to Pre-QK Norm, which as an implicit regulator for attention entropy.
46

---

F Orders of Activation Updates for Kurtosis
To better appreciate the effect of different optimiser hyperparameters on OFs, now consider how the updates that arise during training to a representation matrix X E RRnd can lead to increasing kurtosis (and OFs). In general, training step (e.g. with a gradient/Adam update on trainable parameters earlier in the forward pass than X) will lead to an update X + X+
Recall that Kurt(X) is an defined through comparing the fourth m4 (X) and second mx moments of neuron RMS
+
- X3, for different j. As such, it is natural to ask how updating X <-
X + A updates these moment statisties. We first study the second moment update tg:
U2 def
m2 X + )x (X)
1
Σ
X
Xa
(14)
d
,=1
1
1u2,1 +
(2222,
with
(15)
nd
d n
d
n
22
def
X
Xa.i
12,2
aaj
(16)
j=la=1
j=1a=1
Likewise for the fourth moment update 144
1
2
U4
def
Ma
X + A7)
(X)
d
X
A
x2
ond
(17)
j=l
1
11,1
+ u1 + u1 + U1,1),
with
(18)
n²d
n
n
U4,1
def
a.j a,B=1
X
X3
U4
def
Σ
2(
32
+4 a.j X..
43,
Xa.i
(19)
j=1
j=1a,9=1
1
d
n Σ
U4,3
del
ΣΣ j=1o,fml
4x.
AX
UsA4
del Σ j=1a,8-1
A ( ( )3
(20)
Above, we have broken down the tth moment update up into (up,))s. where up. denotes the contri- bution to Up that is order / in X The reason for this is that, typically, a learning rate parameter n) is used such A is linear in 7) and so Up is order / in n/ 2T Usually, 1) is chosen to be small such that A is small element wise relative to X. Note that the quadratic update terms Up, are always positive whereas the linear terms Up are not necessarily positive, so we might expect quadratic terms to drive any increase in the pd moment map
of these (up.1) terms, for our OP block, default Pre-LN block, and also three modifications that reduce OFs in Pre-LN (increasing Adam epsilon from 1e - 8 to le reducing maximum LR from 1e 3 to 3e 4, and using non-diagonal preconditioner e.g. SOAP (69) trained on CodeParrot at 130M scale. We see indeed that the cumulative 44 quadratic term dominates the update to u4 and drives the increase in m4 in the default Pre-LN model. Rducing LR, increasing Adam e, and using SOAP, reduce this term, which also reduces the growth in fourth moment and kurtosis. For the choice of LR this is intuitive: in the small LR 77 - limit the linear first order term 14,1 will dominate and the effect of quadratic U422 can be ignored. The impact of sub-leading order terms like U122 in OFE related to the discretisation drift between discrete- discrete-time gradient descent and continuous-time gradient flow [10.. 49 49 plots the non-cumulative version of Fig 48)
In Fig 48 plot the cumulative
On the other hand, in Fig 48 the OP block has a relatively large cumulative increase from 44,2 that is matched by decrease in 14,1 and large increase in 12, which means the kurtosis (which is the
For example, if we have X HW for previous layer H that is fixed (e.g. embedding layer in transformer). Then we usually update weights W + AW linearly in v, and so Λx HAW is also linear in 1)- For other layers we need to consider the change in H too, but this will also be linear in to leading order. This is straightforward
for w2.2. For u4.2 the second summand can be factorised as Σ, (E. Aa.) which is positive.
47

---

ratio m4/m3) does not increase as much as Pre-LN. Fig 50 shows that u422 dominates the cubic u4133 and quartic u4.4 update terms to the fourth moment, so we can focus on studying u422- We plot the moment updates for the input to the second attention block (out of six).
The models presented in Figs 48 to 50
trained using Adam SOAP without momentum, akin and 32 0.95 in Adam 23 The reason for this was to separate the contribution of individual training steps on the kurtosis updates. If instead we re-introduce momentum with 31 = 0.9, then the different update steps become mixed and the leading order 41,1 dominates the updates to the kurtosis for the Pre-LN model, seen in Fig 51 The choice of 32 = 0.95 highlights that OFs are not specific to the standard choice of B2 0.999 in Adam W.
to RMSProp [1II: we set B1
Recall SOAP 69 is just Adam in rotated basis (obtained via Shampoo), and still uses B,, 92 hyperparam-
eters
48

---

Kurtosis
Hnrpaaaneeee Model
Adamel
Pre-LN, Adami
Pre-LN, AdamiN small UB
Rocated A0APW
Cumulative 42,1
-11
344
Cumulative Us,2
Cumulative u2
$.90
1000
8.00
Cumulalive u
5.00
-4.0
200
Cumulative 2
6.055
Cumulative de
tetn har www. Mww. www.://wwwwwwny......... Axy
8,000
8.003
10k
20k
300
Training Step
Figure 48: Cumulative metrics to track kurtosis updates. Models were trained without momentum. We see that the quadratic u4,2 term dominates updates to the fourth moment.
49

---

Kurtosis
304
Model
AdamW
Fre UN, AdamW
Pre- UN, AdamIW small UR
Pre-IN, Adamty ope-1e-4
UN Rotated SOAPPI--OAP)
c0005
200000
Ut.7
2.0
CO002
COOOS
-0.000
20000
U4..
8.000
8000
8900
3.6
03
1/4
1000
1.00
00004
1302
Training Step
Figure 49: Non-cumulative metrics to track kurtosis updates. Models were trained without momen- tum.
50

---

[U, M compared [u
Model
Pre-1Ni, Adammi
Pro-LN, AdammV, small
Lt, AdamW, aps-1e44
S0APP)
1.00
106
J0E
Training Step
Figure 50: Sub-leading order terms are dominated by u4.2-
51

---

Kurtosis
004
Model
CP. AdamW
AdamW
104
Sm
Fre-LN, AdamW. eps=1e-
Pre-LN, Rotated AdenWW-SSAPI
104
Cumulative V1.1
4.001
-.00
Cumulative us.2
0.04
Cumulative U3
8,000
Cumulative
$.00%
8,000
8.009
Cumulaaivv M4,2
1.0000
1.00%
6.0000
Cumulative
10005
8.00%
8.000
8.905
80000
80005
8300
10K
19K
286
40C
Traning Step
Figure 51: Cumulative metrics to track kurtosis updates. Models trained with momentum. The leading order u4,1 term now dominates the updates to the fourth moment.
52

---

G Worse Signal Prop Means Higher Activation Kurtosis in Gaussian Features
Proposition G.1 (Bad Signal Propagation implies higher kurtosis for Gaussian features). Suppose we have X E RRad zero-mean Gaussian distributed with all inputs uniformly correlated with some P > and independent features (across columns). That is: [ 0 and E x...x.XX... p. 1
k} + (1 p) 1 k} 1fa 3}
Then, if we consider the feature-wise Gram matrix Exp = X we have that the expected squared diagonal entry of E5 is [E((E=),,11 1+ 2p² + (1) increases as increases, whereas the expected diagonal entry is (E((E(( = is independent of p.
Proof. As Gaussians determined by their first two moments, let suppose that X VT- p + yoej. where (Wa.j)o.i and (1z) are independent standard Gaussians. Then, for two neuron indices k, I d we have:
(X"X)a) hdl
=(1
Σ asn
Wa,k Ma,
(21)
+ pm
(22)
+ p(1 p)
Σ
Uo + Ua,
(23)
asn
We are interested in the diagonal elements of Dip = X X, when k = above. In this case, we have (u2 aaak la and of are all independent chi-squared x² distributed with 1 degree of freedom. For Z ~ xi. have E E[Z]- 1 and E [ = 3.
For the first moment, we take the expectation above and note that the summands of Eq 23 are products of independent zero-mean Gaussians (so zero mean). This gives E Xx, = n and hence
1[ 1, as required.
For the second moment, we note that all cross products in (x x)
2- will disappear in expectation
when we square besides the one involving Eqs (1 and (2). as both terms will be x² distributed (hence not zero-mean). On the other hand, all cross products involving Eq (23) will be odd order in at least one zero-mean independent Gaussian (hence zero-mean).
The square of Eq 21 is (1 - p)²n( 2) in expectation, which can be seen by the fact that Ensn " actually a X distribution, with mean n and variance 2n. Hence for Z X we have E|Z²| 2 + Var( 2) = n² 2n.
The square of Eq 22 3p2n² in expectation, again by properties of xi random variables.
The square Eq 23 is (n) (in fact 4p p)n) in expectation and will be dominated by the O(n²) terms. To see this, we note that Eq (23) is sum of n zero mean i.i.d. random variables, so one can use the additive property of variances for independent random variables.
Finally, the cross term between Eqs (1 and 22 is 2p(1 - p)n² in mean. One factor of n comes from the sum of inputs a n and the other comes from Eq (22) already. The product of two independent x random variables is in expectation.
Putting this all together, we have
E E x
=(1
+ 4p(1 +3p
p)n 2p(1 p)n?
(24)
k
=(11 p)2 +3. 3p² + 2p 2p²)n² + O())
(25)
=(1 +222))n n² + (())
(26)
As Ef 12 X we divide Eq (4 by n², and obtain our desired result.
Nooo this covariance gives a "uniform" correlation structure E Xx 1 (1 - p)L,, + p1,,1. which has been studied before in Noci et al. (27 He et al. 28 as way to study signal propagation in sequences. Rank collapse [60] is when p 1.
53

---

Above, we note that [[(Ep),,1] is equivalent to the fourth moment m4 in our feature-wist kurto- sis definition Eq D while [[[())),.1] corresponds to the second moment m2. Hence, Prop G.] demonstrates that worse signal propagation (in terms of higher p) leads to higher kurtosis.
We note that the result is restricted to a Gaussian setting with independent features. This is an accurate description of large- width NN initialisation [19, but does not capture training dynamics discuss in the main paper. Indeed, the maximum kurtosis (1 + 2p²) is 3 when 1, whereas in our experiments obtain much higher values during training (and the maximum the width d, which considerably larger than 3 in practice). This represents a gap in our theoretical understanding and practice, which we leave for future study.
54

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justificcton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided the appendix. If you answer [Yes] question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims made in the abstract and introduction are supported by our theoreti- cal analysis and experimental results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
55

---

Justification:
Directions for future work are discussed in the Discussion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model vell-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification:
Assumptions and proofs provided, among others, in Sections 2 and
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We have provided code for our main experimental setting.
Guidelines:
The answer NA means that the paper does not include experiments.
56

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We have released our code. The data is already open-source. Instructions for reproducing results are provided in the experiments sections and appendix.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips.cc/publcc/eeecc http:///ww...eeeeeeuuubieeeeeeee ssi ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
57

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification:
We report the training details, including hyper-parameet and optimizers.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We do report error bars for subset of experiments. Other experiments were performed at a scale that makes it expensive to run multiple times.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split. initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еттог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We give precise information on the computational resources and training time of our models.
Guidelines:
The answer NA means that the paper does not include experiments.
58

---

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http:p//aaaurrg..ccpaaaallaccaaa
Answer: [Yes]
Justification: Our work in line with code of ethies.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations in their jurisdiction).
Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We have a discussion of broader impacts in our introduction in terms of improved efficiency and accessibility of AI.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
functioning correctly, harms that could arise when the
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
59

---

Justification:
Not applicable
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We use publicly available datasets and give adequate credit.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, haperswithcode com/ datasets has curated licenses for some datascts. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [NA]
Justification:
We do not release new assets.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
60

---

Answer: [NA]
Justification: Not applicable
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Not applicable
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
61


Reviews:
[{"review_id": "UYp5skRvVX", "reviewer": "Reviewer_Lb1w", "summary": "This paper addresses the issue of Outlier Features (OFs), which are neurons with activation magnitudes significantly exceeding the average during neural network training, particularly in transformers. These OFs are undesirable to model quantization, leading to high quantization errors. The authors propose quantitative metrics, like kurtosis, to measure OFs and study how architectural and optimization choices affect them. They also introduce the Outlier Protected transformer block, which removes standard Pre-Norm layers to reduce OFs without affecting convergence speed or training stability. Their findings highlight the importance of controlling signal propagation and suggest practical interventions to minimize OFs during training.", "questions": "See weakness.", "limitations": "Yes, in Section 6.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 1, "strengths": "1. The definition of Outlier Features (OFs) is clear and quantitative with kurtosis score. It allows for a more objective and precise analysis of OFs across different neural network architectures.\n2. The paper presents non-trivial amount of empirical studies that analyze the impact of various architectural and optimization choices on OFs.", "weaknesses": "1. Writing Quality: The paper is poorly written, with comments and claims listed in a piecemeal fashion rather than being unified into a structured story. Glitches: on line 59, the phrase \"we matrix multiply\" is unclear. The writing overall lacks fluidity and coherence across sections.\n\n2. Insufficient Evidence for Outlier Protected Block (OP): In Section 3, the authors claim that the Outlier Protected Block (OP) reduces Outlier Features. However, the evidence provided is insufficient to support this contribution. The figures show a rough trend, but more thorough investigation is needed:\na) Is this observation consistent across all layers? (Figure 4 shows different trends in \"final layers\" compared to \"early layers.\")\nb) Does the claim hold under different training schemes?\nc) Are there existing baselines that achieve similar effects?\nd) What are the gains of using the OP block? How much does it improve network quantization efficiency?\n\n3. Relation Between OF and Signal Propagation: In Section 4, the authors assert that the OF phenomenon is related to Signal Propagation. However, both metrics are mathematically defined by matrix X in different ways, so it is unsurprising that they are related. This diminishes the novelty of the finding.\n\n4. Lack of Depth in Training Design Analysis: In Section 5, while the authors provide several training designs and analyze their effects on OFE, the analysis lacks depth. The paper should offer deeper insights into why these effects occur rather than presenting rough findings and claims.\n\nOverall, the paper feels more like a technical report than a conference paper ready for publication. It lacks a solid contribution and fails to adequately answer the research question of \"why OFs emerge during training.\" The authors acknowledge this in Section 6: \"Perhaps the most important is understanding how signal propagation evolves during training, and which factors affect this.\" This is the core question that the paper should address comprehensively.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**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”.\n\n**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”.\n\nMore 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.\n\nAlso, in Figures 1* and 2* of the additional pdf page we show the effectiveness of the OP block at 7B scale.\n\n**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.\n\n**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.\n\n**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.\n\nMoreover, 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.\n\nIn 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.\n\n**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”. \n\nWe 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.", "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": "", "extra_scores": {}}, {"review_id": "tCaetoWRD8", "reviewer": "Reviewer_Yyp1", "summary": "The paper focuses on Outlier Features (OF) in neural networks, particularly transformers, where certain neurons exhibit significantly higher activation levels than others. OFs hinder model quantisation and their emergence during training is poorly understood. The study introduces quantitative metrics to measure OFs, such as kurtosis of neuron activation norms, and investigates how network architecture and optimization choices affect their occurrence. Practical insights are provided to mitigate OFs, emphasizing the importance of managing signal propagation during training. The Outlier Protected transformer block is proposed as a solution, removing standard Pre-Norm layers to reduce OFs without compromising convergence speed or stability. Overall, the paper advances our understanding of OF dynamics and offers strategies to address this challenging aspect of neural network training.", "questions": "see Cons", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper introduces metrics such as kurtosis of neuron activation norms to quantify Outlier Features (OFs) in neural networks.\n2. It provides insights into how architectural and optimization choices influence the emergence of OFs during transformer training.\n3. The study proposes the Outlier Protected transformer block, which removes standard Pre-Norm layers to mitigate OFs without impacting training stability or convergence speed.\n4. Findings are supported by empirical validation across various datasets, demonstrating the effectiveness of the proposed metrics and strategies.", "weaknesses": "1.\tFor the OP module proposed in this paper, all the parameters $\\beta$  are added before the residual connection, but not after the residual connection. Why is this done? Will this result in the inability to eliminate outliers in the input X of the first block?\n2.\tRegarding the selection of parameter $\\beta$, this paper proposes in the appendix experimental details that it is selected in advance as a fixed parameter in the \"CodeParrot\" experiment, and as a trainable parameter in the \"Languini\" experiment. So how should we choose for other scenarios? In the second experiment, what are the update details of parameter  $\\beta$? Can $\\beta=\\mathcal{O}(1 / \\sqrt{\\text { depth }})$  be guaranteed all the time?\n3.\tIn the comparative experiment of this paper on how the OP module can eliminate OFE, such as Figure 4, the OFs metric using the OP module has a trend of increasing with “Token seen”. How to explain this phenomenon? What is the effect for larger “Token seen”?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "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:\n\n**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. \n\nIn 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.\n\n**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. \n\nEmpirically, 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.\n\n**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. \n\nFor 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.\n\nWe 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.", "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": "", "extra_scores": {}}, {"review_id": "W7REW0fBEy", "reviewer": "Reviewer_7V85", "summary": "The paper tackles outlier features (OF), i.e. neurons with activation magnitudes that are significantly larger than average which can cause issues with quantization and low-precision training. The paper introduces several metrics for quantifying the existence of OFs and uses them to explore which design choices lead to the emergence of OFs. The authors then introduce an Outlier Protected Block (OP), similar to a standard transformer block, that works without explicit normalization and does not cause a loss in performance and reduces OFs. They then show the connection between signal propagation and OF and demonstrate that OP  improves signal prop and therefore reduces OFs, as predicted by their theory. Lastly, the paper discusses how the optimization, in particular learning rate and adaptivity influence  OFs.", "questions": "- The reference to Fig 28 in 314 seems off", "limitations": "The authors discuss the lack of large-scale experiments and give experimental evidence for the validity of the conclusions obtained from eq (4).", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The topic seems relevant and I feel like the authors made a sufficiently large contribution to the research of OFs. \n- I liked the connection between signal prop and OFs in Section 4. The section was formal enough while still being easy to understand. I also think it's good that the authors included a discussion about potential issues with the theory (the off-diagonals on the left-hand side of eq 4) and added experiments to show that in practice, bad signal prop seems to lead to OFEs. \n- While maybe not quite sufficient in width, the experimental design itself seems good and makes the results look plausible.  \n- The paper seems well written and is relatively easy to understand considering its topic and while some figures are a bit too small, I overall like their design.", "weaknesses": "- The building blocks for the OP block should be explained in the appendix since the paper seems relevant for practitioners who might not be familiar with the literature. For example: formal definition of pre and post-norm, entropy regularization signal prop theory. \n- My only significant concern is the width of the experiments. I do understand that training LLMs with 10s of billions of parameters is expensive, but it would clearly strengthen the message of the paper if OP could keep up in performance with standard normalized transformers at a much larger layer count, since it could be that this is only true for relatively \"small\" transformers. If training larger language models is impossible, maybe the authors could add experiments for ViTs. Transformers have obviously become very relevant in the vision domain and it would be very interesting to see if OFs are a problem in vision and, if so, if OP can help reduce them while offering similar performance. \n- Section 5 felt more like an after-thought. While I think it is still interesting, I would personally much rather see a larger experimental discussion (as mentioned in the previous point) and have some of this content be moved to the appendix if necessary.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**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.\n\n**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”. \n\nMoreover, 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.\n\nHaving 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. \n\nDespite 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.\n\n7B 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. \n\nWe 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.\n\n**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.\n\nThe 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.\n\nBesides 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.\n\nTable 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.\n\nWe 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.", "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": "", "extra_scores": {}}, {"review_id": "NYtDUbWpjF", "reviewer": "Reviewer_cUUo", "summary": "This paper is doing a lot and is a rare case of the abstract/title really underselling what the paper contains. Basically, the paper investigates outlier feature emergence (OFE) in LLMs and some potential fixes for it. The paper argues that such a study is important for both practical reasons (preventing OFE aids quantization) and theory reasons (understanding training dynamics). It clearly succeeds on both fronts, developing a transformer block that empirically reduces OFE and producing theoretical insights that touch on Signal Propagation, entropy collapse, optimization, and more.", "questions": "Line 118 has a typo, I think “untrainable” was intended. \n\nLine 193: isn't it normalized by the square of the second moment?\n\nLine 314 has a typo in a figure reference.\n\nFigure 9: make the solid and dashed lines represent the same thing for both SGD and Adam.\n\nConsider moving more discussion of appendix F to the main text.", "limitations": "The limitations are well covered.", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Originality: This work clarifies that some prior hypotheses about causes of OFE are inconsistent with experiments. It makes strides towards understanding OFE by combining ideas from Signal Prop and entropy collapse (e.g.), in ways that are insightful, novel, and excellently motivated. The related mathematical analyses are also very well done.\n\nQuality: This is a thorough and careful study, making claims that are well supported by the experiments.\n\nClarity: This paper is extremely well written and contextualized relative to prior work.\n\nSignificance: This work will be appreciated by multiple communities for its wealth of insights. For example, the augmentation of norm-free Signal Prop-focused networks with QK norm to address entropy collapse (making norm-free transformer training competitive) -- which was inspired by the idea that vanilla norm layers help prevent entropy collapse -- is likely to have effects on transformer training procedures going forward.", "weaknesses": "A lot of experimental details are not in the main text, which should be addressable with an extra page. \n\nExperiments with other LLMs and at larger scales are unnecessary but would further increase the significance of this work. The authors addressed this in their limitations section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\nOverall 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.”\n\n*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.\n\nGiven 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:\n\n**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*.\n\nIn 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.\n\n**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.\n\n**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.\n\n**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. \n\nWe 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.", "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": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their time in reviewing our work. We are encouraged that three out of four reviewers gave scores leaning towards accept, with an average score of 6 across all four reviews. In particular, reviewer cUUo gave the highest possible score of 10, commenting: “this work will be appreciated by multiple communities for its wealth of insights”. 


There is, however, a disappointingly large variance in the scores, with Reviewer Lb1w giving a score of 3 and writing that the paper feels “more like a technical report than a conference paper ready for publication”. We believe this, and several of Reviewer Lb1w’s other criticisms, are unsubstantiated with no evidence, which we have discussed at length in our individual response to Reviewer Lb1w.  Moreover, many of Reviewer Lb1w’s criticisms are in direct opposition to praise received from all three other reviews. We urge the reviewers to reach a more consistent appraisal of our work during the discussion period. Through our author response, we hope to convince Reviewers Lb1w, 7V85 and Yyp1, to raise their scores.

In the individual responses, we believe we have addressed all concerns of each individual reviewer, so in the remainder of the global response we would like to draw attention to additional results in the rebuttal that may be of shared interest to the reviewers. 

*Notation for additional results* Whenever we write an asterisk after a Figure or Table, e.g. Table 1*, this means we refer to a Figure/Table from the additional pdf page. Unasterisked Figures/Tables refer to the original submission.


**Quantisation Experiment and Section 5**
In Table 1* we present an experiment testing the combined effects of our different architecture and optimisation choices in affecting quantisation errors. This goes back to our original motivation in studying OFs, and how we can minimise their emergence with different design choices. 

In Table 1*, we take the 125m scale OPT setting of Bondarenko et al. [14], which trains models in standard mixed FP16/FP32 precision on BookCorpus+Wikipedia for around 12B tokens seen. Post training, we quantise to int8 weight-and-activations (W8A8). We use the same  quantisation recipe and reproduce their results. We present results with standard deviations from 3 seeds in the post training quantisation step. We note that the GPTQ paper https://arxiv.org/abs/2210.17323 shows that post training quantisation is harder at smaller scales like this 125m setting than at larger scales.

Table 1* compares both the standard precision performance and also quantisation performance across architecture and optimiser choices. We compare 3 different architectures: 1) standard Pre-LN, 2) the Gated Attention baseline of [14] (which was their best baseline), and 3) our OP block, as well 4 different optimisation setups that are added one after another on top of each other: 1) the default hyperparameters of [14], 2) removing the dropout regularisation, 3) increasing the maximum LR from 4e-4 to 1e-3, and 4) increasing Adam Epsilon from 1e-8 to 1e-5. 

Optimiser choices 2) and 3) were designed to improve standard precision performance, albeit potentially at the detriment of quantisation performance due to OFs (as suggested by our findings with increase LRs in Section 5). Optimiser choice 4) was chosen to improve and reduce OFs following our contributions in Section 5, and thereby hopefully improve quantisation errors.

As seen in Table 1*, our findings throughout the rest of our paper are validated. Firstly, the kurtosis is seen to be highly correlated with quantisation error, for example the Pre-LN model has consistently high kurtosis, and also consistently catastrophic performance at W8A8 (>45 perplexity increase in W8A8 across all optimiser settings). Secondly, our OP block has consistently low kurtosis (around or below 10), and this directly translates to low quantisation error (the difference in perplexity between standard mixed precision at W8A8 is around or below 0.7 across all optimiser settings). This low kurtosis/quantisation error of the OP block holds true even for aggressive optimiser hyperparameters like large LRs that increase kurtosis, but also improve standard mixed precision performance. Finally, the baseline of [14] struggles when dropout is removed and a large learning rate is used, with kurtosis of ~30 and large quantisation error around 2-3 perplexity but our suggestion of increasing Adam epsilon in Section 5 reduces the kurtosis back down to 15 which translates to a quantisation error of 0.77 perplexity. 

Indeed, the best quantised W8A8 model in this case is actually the Gated Attention baseline of [14] with our optimisation choices from Section 5, which is very slightly better than our OP block (15.54 vs 15.62 perplexity). In relation to comments from reviewers 7V85 and Lb1w, this highlights that Section 5 is a fundamental and key section of the paper, and that our contributions to understanding the impact of optimisation choices on OFs are also important, in addition to our architectural contributions. 

**Scaling experiment**
We agree with Reviewer cUUo that scaling further beyond 1.2B parameters is “unnecessary but would further increase the significance of this work”. To address scale, over the weekend of the response period, we were able to access resources to run limited experiments at 7B, which we present in the additional pdf page. In Figure 1* we see that the OP block closely matches standard Pre-Norm in terms of loss curves, and in Figure 2* we see that the Pre-Norm block suffers from significantly worse OFs, in the sense that its kurtosis is much larger. These results mirror our findings at smaller scales in the submission. Further details and context can be found in our response to Reviewer 7V85.

Meta-review: This paper presents a study on reducing outlier features in neural network training. The paper has sparked considerable discussion regarding its novelty, experimental rigor, and practical implications for neural network quantization and training dynamics. The major concern highlighted is the paper's ability to convincingly demonstrate the efficacy of the proposed techniques across varying network architectures and larger datasets. In their rebuttal, the authors have addressed these issues comprehensively, providing additional data and clarifications that enhance the original submission. While one reviewer remains critical, focusing particularly on the paper's writing style and depth of technical contribution, the other reviewers acknowledge the paper’s potential impact and the rigor of the experimental setup.

Given the mixed nature of the reviews but considering the strong rebuttal and the additional experiments provided, the decision leans towards acceptance, contingent upon the successful incorporation of suggested improvements in its final version.
Decision: Accept (poster)