From Attention to Activation: Unraveling the Enigmas of Large Language Models
Authors: Prannay Kaul
Chengcheng Ma
Ismail Elezi
Jiankang Deng
Venue: ICLR 2025 Poster (2025)
Keywords: Transformers
Adam
Optimizer
Outliers
Attention
Quantization
Primary area: foundation or frontier models, including LLMs
TL;DR: Popular transformer models have strange behaviours including large outlier activations and attention dominance by specific tokens---we fix these with a change to softmax and a suitable change of optimizer.

Abstract: We study two strange phenomena in auto-regressive Transformers: (1) the dominance of the ﬁrst token in attention heads; (2) the occurrence of large outlier activations in the hidden states. We ﬁnd that popular large language models, such as Llama attend maximally to the first token in 98% of attention heads, a behaviour we attribute to the softmax function. To mitigate this issue, we propose a reformulation of softmax to softmax-1. Furthermore, we identify adaptive optimisers, e.g. Adam, as the primary contributor to the large outlier activations and introduce OrthoAdam, a novel optimiser that utilises orthogonal matrices to transform gradients, to address this issue. Finally, not only do our methods prevent these phenomena from occurring, but additionally, they enable Transformers to sustain their performance when quantised using basic algorithms, something that standard methods are unable to do. In summary, our methods reduce the attention proportion on the first token from 65% to 3.3%, the activation kurtosis in the hidden states from 1657 to 3.1, and perplexity penalty under 4-bit weight quantisation from 3565 to 0.3. Code is available at https://github.com/prannaykaul/OrthoAdam

Published as a conference paper at ICLR 2025
FROM ATTENTION TO ACTIVATION: UNRAVELING THE ENIGMAS OF LARGE LANGUAGE MODELS
Prannay Kaul'- Chengeheng Ma²
Ismail Elezi't
Jiankang Deng'
'Huawei Noah's s Ark Lab, London, UK
Institute of Automation, Chinese Academy of Sciences (CASIA)
CURRENT TRANSFORMER MODELS
Key Pagtion 66
Chaggel
124
142
320
314
446
0.8
0.4
114
118
124
0.0
126
10a
OUR TRANSFORMER MODELS
Key Phottoon
Charnel ID
0.8
0.6
0.2
132
10
11a
10
1.0
126
10t1
Figure I: (top) (a) The mean attention map across all heads and layers of a GPT2- Medium model the first ken strangely dominates attention (boxed in red). (b) The mean hidden outlier activations emerge specific feature dimensions (boxed red). first token position exhibits most extreme outlier activations-(cird))) with our proposed sofimax-1 function eliminates the first token dominance. (b) Using our proposed optimiser, OrthoAdam, removes outlier activations without any reduction in model performance.
layers of the model-
red). (bottom) (a) Replacing the canonical softmax function
ABSTRACT
We study two strange phenomena in auto-regressive Transformers: (1) the dom- inance of the first token in attention heads; (2) the occurrence of large outlier activations in the hidden states. We find that popular large language models, the first token 98% of attention heads, a behaviour we attribute to the softmax function. To mitigate this issue, we pro- pose a reformulation of softmax to softmax-1. Furthermore, we identify adap- tive optimisers, e.g., Adam, as the primary contributor to the large outlier ac- tivations and introduce OrthoAdam, a novel optimiser that utilises orthogonal matrices to transform gradients, address this issue. Finally, not only do our methods prevent these phenomena from occurring, but additionally, they en- able Transformers sustain their performance when quantised using basic al- gorithms, something that standard methods are unable to do. In summary, our methods reduce the attention proportion the first token from 65% to 3.3%, the activation kurtosis the hidden states from 1657 to 3.1, and perplexity penalty under 4-bit weight quantisation from 3565 0.3. Code available https: http::/ggttt...tt/ com wtw.h//w..co..ppaaaaaaaau//au//a
such Llama attend maximally
"Work conducted during internship
1 "oorrsponndecce sma e ez @huawei

---

Published as a conference paper at ICLR 2025
INTRODUCTION
Transformers have revolutionised machine learning, achieving state-of-the-art performance across
diverse domains, including language, vision and protein structure prediction (OpenAT 2023 Carion
et al., 2020, Jumper et al., 2021). However, the inner workings of -uto-reeressive Transformer remain enigmatic. Recent studies (Elhage et al. 2022] Olsson et al. 2022 Bansal et al., 2023)
unravelled some of their complexities, yet we find that two surprising phenomena remain pervasive:
The strong, consistent dominance of the first token in attention maps-see see top of Figure la The presence of outlier activation values, across sequence position, in specific feature channels of the hidden states (the intermediate features of each layer after the residual connection) that are orders of magnitude larger than other values-see see top of Figure Tb
2.
We ask: What causes these phenomena? Are they essential to performant models? And, if not, how can we mitigate them?
These two phenomena are aesthetically curious, but also have important practical implications. For instance, Llama models (Touvron et al., 2023bb Dubey et al. 2024) exhibit the aforementioned first token dominance of attention, and so requiring complicated attention masking schemes to extend Llama models to tasks with long sequences (Xiao et al. 2024). i.e., increase the maximum context length used during training. This is particularly crucial for instruction-ture
models where long conversations are desirable (Wei et al. 2022; Ouyang et al. 2022). y. Similarly, the presence of outlier activations leads to challenges in quantising large language models (LLMs). Large outlier activa- tions increase the required quantisation range (to capture the outliers), resulting in low effective bits for the non-outlier activations, causing severe performance degradation post-quantisation. To ad- dress this issue, prior work has proposed mixed-pra sixee-precision decomposition of LLMs (Dettmers et al., 2022 complex scaling of the weights and activations which must be learnt for each model (Xiao) ct al. 2023). Therefore, our additional motivation is to understand and mitigate these phenomena in general manner, such that these issues resolved during training.
We begin by examining the attention mechanism, and surprisingly find, across numerous input se- quences, query tokens attend most the first key token up 98% of the time. This is striking considering the limited semantic information the first token typically contains-it is often a special token indicating the start of sequence, such as <bos>, We explore explanations for this, ruling out positional encodings, non-linearity choice, or feature normalisation. Ultimately, we identify the softmax function used in the attention mechanism combined with causal masking as the root cause- cause-excessive attention on the first key token demonstrates an attention head effectively doing nothing (Bondarenko et al.) 2023 Clark et al. 2019). The first token is privileged due to causal masking; it is the only key token to which all query tokens can attend. We propose an adjustment to softmax as a solution, softmax-1, removing first token dominance in attention (bottom of Figure Ia
Despite removing first token dominance in attention, us- ing softmax-1, we find that the problem of outlier acti- vations in the hidden states persists. Once again, in- vestigate potential causes of this issue and discover the outliers are primarily caused by the use of adaptive op- timisers, e.g., Adam (Kingma & Ba 2015). Specifically, our experiments show the exponential decaying averages of first and second moments of gradients result in outlier activations. To tackle this, we propose novel optimiser, OrthoAdam, which transforms computed gradients using orthogonal matrices, thus storing gradients in an alterna- tive basis to the model parameters. Our results demon- strate this optimiser eliminates the outliers in the hidden states of Transformers (bottom of Figure lb).
Model
#Parameters
PPL
FP16
4-bil Quant
GPT2-Small
137M
4456.1
GPT2-Medium
350M
28.8
2435.3
GPT2-Large
812M
25.2
571.0
GPT2-XL
166
23.2
7981.8
Llamu2-7B
6.7B
191477.5
Llama3.1-4 1-8B
8B
10.2
2087638.0
GPT2 (Ours)
350M
17.1
GPT2 (Ours)
1.4B
13.3
13.6
Table 1: Duc surprising phenomena in Transformer models, basic zeropoint bib weight quantisation leads to catastrophic performance degradation.
Our models trained with softmax-1 and OrthoAdam ex- hibit improved robustness to quantisation.
Our research extends beyond aesthetic curiosities. While LLMs perform well despite first token dominance and outlier activations, they lead to practical challenges. Although advanced schemes have been developed to enable quantised LLMs to maintain their performance, we show our ap- proach enables LLMs to maintain their performance with the most basic quantisation methods, such as per-tensor 8-bit absmax weight/activatiin quantisation and 4-bit zeropoint weight quantisation. Thus, our investigation helps to better understand Transformers, while offering practical benefits.
2

---

Published as a conference paper at ICLR 2025
In summary, contributions follows:
We identify the dominance of the first token in attention and the occurrence of outliers in the activations of the hidden states as significant issues in auto-regressive Transformers.
We propose two simple, effective solutions: reformulation of the softmax function, softmax-1, address the former issue, and novel optimiser, OrthoAdam, to tackle the latter. Our methods reduce first token attention from 65% to 3.3% and activation kurtosis from 1657 to 3.1.
We demonstrate that these proposals not only resolve the identified problems but also lead to practical improvements in the performance of Transformers under 8-bit weight/activation and 4-bit weight quantisation. Our method reduces the perplexity penalty under 4-bit weight quanti- sation from 3565 0.3.
PROBLEM DEFINITION
This work investigates the two most prominent and strange phenomena of auto-regressive Trans- former models: (1) strong, consistent dominance the first token the attention maps; (2) strong, consistent outlier activations in specific feature channels of the hidden states (the intermediate fea- tures computed immediately after the residual connections) see top of Figure I We aim to under- stand the cause of these phenomena and propose individual solutions for each of them. They have been investigated or commented on previously (Bondarenko et al., 2023 Dettmers et al. 2022 Xiao ee al. 2023), but work reaches different conclusions the causes and suggests novel solutions. We start by describing these two anomalies in detail.
2.1 FIRST TOKEN DOMINANCE IN ATTENTION MAPS
The top of Figure Ta shows the attention map, averaged across all layers and heads, of a Transformer model. specifically a pretrained BPT2-Medium model (Radford et al.] 2019), for a single real natural language sequence. Strangely, in this average attention map the key corresponding receives the highest attention across all queries. Quantitatively, we find the first key token is the most attended key in 76% of (query, head) pairs and receives 52% of all attention, when evaluating on the en validation split of the C4 dataset (Raffel et al., 2020, Dodge et al. (021). l). This behaviour is consistent across different LLMs, including the Llama series (Tovrron et al. 2023b; Dubey et al.)
the first token
2024), DeepSeek (Liu et al. 2024). and the GPT2 series (Radford et al. 2019). See Appendix K for detailed examples of attention maps for these models.
Attention is key component of the Transformer architecture, and work on the interpretability of LLMs often focuses on analysing attention (Elhage et al.) (022)). Moreover, many models, such as Llama2, special token for the beginning of sequence (the <bos> token), which is always the first token in an input sequence. This makes first token dominance particularly puzzling, as such models should learn the initial input structure easily. We hypothesise that this phenomenon in the attention mechanism is a symptom of a fundamental problem in the Transformer architecture and is necessary for a performant auto-regressive Transformer.
2.2
OUTLIER ACTIVATIONS IN THE HIDDEN STATES
The top of Figure lb shows the activation magnitude in the hidden states of a pretrained GPT2- Medium model. We observe the hidden states of the Transformer model exhibit consistent outlier activations in specific feature channels across all token positions (boxed red), with the most extreme outliers occurring in the first token position (circled red). Once again, this behaviour consistent across different L.Mss and is invariant to the input sequence, i.e., the same feature channels always exhibit outlier activations. See Appendix [ for examples of hidden states in pretrained models.
From practical perspective, these outlier activations are problematic with regards to quantising models for deployment (Lin et al. [22 Dettmers et al. 2022). However, from a theoretical perspec- tive, the cause of these outlier activations not well understood. Previous works, have suggested these outliers are related to first token domination in attention maps (Xiao et al. 2023 Bondarenko et al. 2023). This is plausible for the most extreme outliers observed in the first token position, but it does not explain the outlier activations observed across all token positions. In this work, we show the two phenomena unrelated and separate solutions are required to address each.
3

---

Published as a conference paper at ICLR 2025
3 METHOD: FIRST TOKEN DOMINANCE OF ATTENTION MAPS
We start by climinating plausible causes of the first phenomenon of interest: first token dominance of attention maps. We mainly consider GPT2 representative auto-regressive Transformer, because of its simplicity, but also consider the more recent Llama2 model to narrow down possible causes of this phenomenon. For all experiments, unless mentioned otherwise, we use GPT2 model with 130M parameters, trained on the en split of the C4 dataset.
3.1 ELIMINATING CERTAIN CAUSES OF FIRST TOKEN DOMINANCE OF ATTENTION MAPS
Both GPT2 and Llama exhibit first token dominance in attention maps. Thus, rule out parts of their architecture that are different:
Positional encoding. Llama models use Rotary Positional Encodings (RoPE) (Su et al. 2024). while GPT2 models uses learnt absolute positional encodings (Vaswani et al. 2017).
Initial token. Llama models use a <bos> token to denote the beginning of a sequence, while GPT2 models do not.
Activation function. Llama models use SiLU (Elfwing et al. 2018 in the feedforward layers, while GPT2 models use GeLU (Hendrycks & Gimpel| (016)).
Feature Normalisation. Llama models use RMSNorm (Zhang & Sennrich, 2019). while GPT2
models use LayerNorm (Ba et al. 2016).
Note that Llama and GPT2 use different positional encoding, but possible that any form of positional encoding might be cause of first token dominance. To test this possibility, we train GPT2 model without any positional encodings and observe the attention maps. We find equivalently trained GPT2 models with/without positional encodings exhibit first token dominance in 33%/20% of (query, head) pairs and allocate 17%/10% of all attention to the first token. Thus, we conclude that positional encodings are not the cause of these anomalies. The models mentioned here trained for relatively few steps and first token dominance is more pronounced in our longer-trained models and in publicly available pretrained models.
3.2 REMOVING FIRST TOKEN DOMINANCE OF ATTENTION MAPS
After eliminating the above causes, we have two aspects of Transformers that could cause first token dominance: (1) causal masking in self-attention; and (2) softmax normalisation in attention heads.
Consider the self-attention mechanism on the initial token in a causal Transformer. The first query token can only attend its own key token and therefore it receives an attention score of 1, due to softmax normalisation. Similarly, the second query can only attend to the first two keys, whose attention scores must sum to 1. Prior work establishes attention heads specialise to concepts or concept groups (Bansal et al.) 2023] |Elhage et al.) (022). However, given a query irrelevant to the specialisation of an attention head, it must still allocate attention across the keys summing up to 1. Moreover, causal masking privileges the first key token above all others; it is the only key token to which all tokens can attend. This explains why the first token specifically dominates attention maps.
Clearly, particular attention head should be able to atlend nowhere if no relevant information is present. Thus, we modify the softmax function to the following:
exp(x;)
softmax- (xx;) -aaaxx-1(x))
softmax-1(x;) <
(1)
E exp(x,)
This modification removes the strict enforcement of attention scores summing to 1, allowing the model to allocate attention it sees fit, including having low attention scores everywhere. From registers/attention sink perspective (Darcet et al. 2024 Xiao et al. 2024), the in the denominator is equivalent to register/attention sink key token which has dol product with any query token.
Validating the hypothesis. We train two GPT2 models, one with canonical softmax and one with softmax-1, keeping all other variables the same. The model trained with canon- ical softmax attention exhibits first token dominance; the first key token is the most at- tended to key in 53% of (query, head) pairs. However, the model trained with softmax- lowers this to just 2%. Furthermore, with canonical softmax 46% of all attention is re- ceived by the first key, while using softmax-1 lowers this to 4%, thereby validating our idea.
4

---

Published as a conference paper at ICLR 2025
Sey Postion
12
The difference in attention maps between canonical softmax and
softmax-1 is shown in Figure Ia] which compares the attention maps of two models on the same input sequence. Furthermore, we find using softmax-1 I has no effect on training stability, convergence model performance (see Appendix foor the training curves of all our trained models).
0.
Query Roston
What if causal masking is relaxed? To verify the first token is privileged by causal masking, causing first token dominance, we train a GPT2 model with canonical softmax in which causal mask- ing is removed for the first 10 tokens. (the loss function is appropri- ately modified). This way, all queries can attend to the first 10 keys. Figure@ shows one of these tokens (this happens with uniform dis- tribution) still dominates the attention map.
96
Figure 2: Relaxing causal mask- ing leads attention domination by a different token
METHOD: OUTLIER ACTIVATIONS
To quantitatively establish the extent of outliers in the hidden states, we use kurtosis. Kurtosis, in this case, is measure of tail heaviness of a set of activation values. Activations which are normally dis- tributed have a kurtosis of ~-3, while higher kurtosis indicates a heavier-taie tailed distribution (e.g., the exponential distribution) and lower kurtosis indicates a lighter-tailed distribution (e.g., the uniform distribution). Given hidden states X E R M XD ber of layers, L. is the number of tokens and D is the number of feature channels, compute the per-layer, per-position kurtosis of the hidden states as:
Transformer model, where M the num-
Kl Kurtm, [am, = Ex((Xm...d [mm,))
Eal(Xm,ddd )2j2 where Pm.1 = (m(,
(2)
where X Xm, d is the hidden state at layer m at position for feature d, and pmm. is the mean hidden state value at layer m at position I.
4.1 ELIMINATING CERTAIN CAUSES OF OUTLIER ACTIVATIONS
We start by eliminating certain causes which could lead to the presence of outlier activations.
Feedforward Layer Biases. GPT2 uses biases feedforward layers, while Llama uses none, therefore it is unlikely feedforward layer biases cause of outlier activations.
Normalisation Layers. GPT2 uses LayerNorm (Ba et al. 2016) while LLama uses RM- SNorm (Zhang & Sennrich 2019), which both learn individual scaling parameters for each feature channel, potentially causing the outlier activations. To remove such an effect, we replace Layer- Norm in trained GPT2 models with RMSNorm version which applies single global scale instead of per-channel scaling, and call it "RMMSNorSSingee"-SS similar to "Simple RMSNrrm from Qin al. (2023) which has learned parameters. We find outlier activations persist in the hidden states of GPT2 model with RMSNormSingle. In Table 5e show kurtosis remains high in models trained without biases and/or with RMSNormSingle.
Optimiser. Most Transformer models are trained with Adam (Kingma & Ba 2015) or a variant. These optimisers track the first and second moments of the computed gradients using exponen- tial moving averages, tracking these moments at parameter level. The main hyperparameters of Adam-liks like optimisers are 81 and 82, which control the decay rates of the first and second moments, respectively. If B2 O, only the first moment of the gradients is tracked, resembling stochastic gradient descent (SGD) with momentum. Conversely, 31 = O, only the second moment of the gradients is tracked, resembling RMSProp. We suspect that given the optimiser tracks moments in the same basis as the model parameters, it is the most likely cause of the outlier activations in the hidden states auto-regressive Transformer models.
Validating the hypothesis. We train a series of GPT2 models using Adam, RMSProp. SGD with and without momentum, tuning the learning rate and training schedule to encourage convergence. The model trained with SGD has the slowest convergence and highest validation perplexity, while the model trained with Adam converges the fastest and has the lowest perplexity. However, find
5

---

Published as a conference paper at ICLR 2025
models trained with Adam and R.MSProp have high kurtosis, 140 and 70, respectively, while training with SGD gives a kurtosis of ~3.0. We provide these results in our ablation study (Section (.3).
4.2 ORTHOADAM
The previous section leaves important question for training Transformer models: "How can we train a model with an optimiser which has the speed and convergence properties of Adam, but produces activations properties sim- ilar to SGD"?
Optimisers which track exponential decaying averages of the first and/or sec- ond moments of the gradients lead to outlier activations in the hidden states of Transformer models. Moreover, in the models trained above, the largest absolute parameter values correspond to the features which exhibit outlier activations in the hidden states, i.e., if outlier activations occur in feature chan- nel i of the hidden states, the largest model parameter values correspond to specific weights which act on feature channel i of the hidden states, e.g., the thh output channel of the output projection weights of the attention/MLP lay- ers. Therefore, to arrive at these large model parameter values, the optimiser (e.g., Adam) must provide relatively large updates to these specific parameters and not others. We note here that Adam and similar optimisers calculate gra- dient moments in the same basis the model parameters. Additionally, given the channels which contain outlier activations appear invariant to the input se- quence, we hypothesise that these channels are an artefact of the optimiser and do not correspond to any meaningful feature in the input sequence-see Ap- pendix M for plots of the hidden states of pretrained models with different input sequences. Given these observations, we discuss an idealised case of observed hidden states below, and show how orthogonal transformations can be used to reduce outlier activations.
Figure Rotating vectors with dominant components leads to reduction in the maxi- mum absolute value.
Consider D-dimensional vector, x = ae; + z, where e; the 7th unit vector in the standard basis, P. 1 a >> 1 and z (N, I). The first term represents the single outlier activation specific to the 7th channel and the second term represents the "informative" activations. The vector x represents the hidden states of a Transformer model with high kurtosis. This simplified model makes two assumptions: (1) there is a single outlier activation channel; and (2) the informative activations normally distributed.
For values of similar to that of Transformer models, i.e., D ss [10*, 10 Kurt[xy] = O(D). Therefore, we expect larger Transformer models of given architecture to have larger kurtosis in their hidden states. Moreover, the ratio of the -nrm to the fn-norm of the hidden states in our simplified model, ||x|| is close This ratio is another proxy for the extent of outliers.
||x||
Now we consider the effect of an appropriate orthogonal transformation on the vector x. Let Q DxD be an orthogonal matrix, and y =Qx. Under a particular orthogonal transforma- tion,
? 1 1 and Kurtly,] 3. The orthogonal transformation which achieves this is one which
rotates the vector such that Qe J 1. Figure illustrates this rotation process in 2D and 3D. The kurtosis and norm ratio results quoted in this section are derived Appendix G and Ap- pendix H respectively, and are shown to be empirically valid for models we train from the plots in Appendix 1.2 and Appendix I.3 respectively.
One option is to apply orthogonal transformations directly to the hidden states of the model, ie., make Q part of the model parameters that are kept fixed during training. Instead, we propose novel optimizer, OrthoAdam, which applies orthogonal transformations to incoming gradients such that the moment calculations (which our experiments in Table 3 show are the key factor in produc- ing outlier activations) are performed in different basis to the model parameters to prevent gradient updates to any particular set of parameters which lead to outlier activations. We provide the full algorithm in Algorithm
In our experiments, we randomly sample the orthogonal matrix for each parameter (which remains fixed during the training of the model). We find that using OrthoAdam leads to significant

---

Published as a conference paper at ICLR 2025
Algorithm OrthoAdam, our proposed optimiser for redueing activation outliers. gi the element- square g. C go. With 8i and 83 we mean 31 and 33 taken to the power of t.
given learning rate: 0.001, first moment decay rate: 01 0.9, sccond moment decay rate: 92 0.999, numerical epsilon: = 10- initialise time step: t- 0, parameter vector: Ot- E IR"," first moment vector: mu-u +- 0, second moment
vector: V==o < 0, schedule multiplier: λ==o E IR, unique orthogonal matrix: Q E O
repeat
t
(f( Selec(Batch(O.-.)
select batch and calculate gradient D- store the gradient in model parameter basis transform gradient into unique optimiser basis update biased first moment estimate update biased second raw moment estimate compute bias-corrected first moment estimate
g
< 1)((
ge <- MatMul ).g.
-1tur-1 +1
V; B2V,-1 + 82)g:
m f mm./(1 8i)
V. + v./(1 82)
D compute bias-correcter second raw moment estimate
S, -- mi +c)
P calculate the update step unique optimizer basis transform the update step back model parameter basis b can fixed, decay, or also be used warm restarts b apply parameter update
R, MatMul
A - SetScheluleMutiipiecchh
0. + 01-1 Xersi
until stopping criterion is met return optimised parameters Os
duction in the kurtosis of hidden states in Transformer models, effectively eliminating the outlier activations. This is shown qualitatively at the top of Figure [b where feature channels with high absolute activation values in the hidden states are no longer present across all token positions, and quantitatively in Table Saolwngg the kurtosis of hidden states in models trained with OrthoAdam is close to 3, with no performance penally.
EXPERIMENTS
Datasets. We train all models on the en training split of the C4 dataset (Dodge et al. 2021
Raffel
al. 2020) and evaluate on 100000 samples from the validation en split.
Models. We train GPT2 models with (600M, 130M, 350M, 1.4B) parameters and Llama2 mod- els with ~100M parameters. Apart from changing the softmax function, the only other changes we make to the model architectures are the use of RMSNormSingle and we do not use biases in feedforward layers. We ablate these changes in the ablation study at the end of this section.
Training. Unless stated otherwise, we use batch size of 512 and a cosine learning rate schedule with linear warmup for {1000, 2000, 6000, 10000) steps for models with {60M, 130M, 350M, 1.4B) parameters respectively, with a maximum learning rate of 10-3, We train models with {60M, 130M 350M, 1.4B) parameters for { (160k, 320k, 960k, 600k} steps respectively. Note that we use reduced number of steps for the 1.4B parameter model due to computational constraints. In the ablation study, we train GPT2 models with 130M parameters for 40k steps only.
Metrics. We evaluate experiments in the following metrics: (1) the perplexity (PPL) of models on the validation set; (2) the mean kurtosis across all layers of the model (evaluated separately for the first token and the remaining tokens); (3) the maximum absolute activation across all layers of the model (again evaluated separately); (4) the percentage of (query, head) pairs in which the first key token is the most attended to key token. We calculate (1) to ensure
method least maintains the vanilla language model performance, f.e., to ensure the model is not harmed by softmax- I or Adam (2) and (3) show quantitatively the extent to which outlier activations are present the hidden states. Finally, (4) shows the extent to which the first token dominates attention in the model.
5.1
MAIN RESULTS
We show the results of softmax-1 I and Ortmo used to train GPT2 and Llama2 models in Table We observe that across both model architectures and all sizes, the evaluated PPL is the same or slightly lower when comparing a model with softmax-1 and trained with OrthoAdam to the vanilla model with neither, indicating that method does not change model performance. Despite no
7

---

Published as a conference paper at ICLR 2025
Model
AParameters
Softme++1
OrthoAdam?
PPL
En [Ke.a]
Kurtosis
3m [Krcaa]
Em [[Xma..ll
Activation Value
En [AXmatte
Fini Alln
X
31,9
313.8
77.9
1856.11
266.6
0,489
60M
32.4
31.6
2503
105.6
S1.4
10.6
1419,9
304,9
259.0
114.7
0.021
0,365
7.6
0.019
514.0
7018.1
1014.8
0.527
GPT2
130M
X
23
22.7
446.4
175.4
148.2
20.2
4285.0
1134,3
433.4
967.5
0.424
0.024
22.8
10.1
7.3
318.1
261.6
0.019
350M
82033
3.1
161.
40196.0
388.1
3801.1
0.021
0,579
1.48
13.3
1656.5
3.1
351.9
3.0
56798.1
181.9
7051.2
132.1
0.648
0.033
X
X
17.4
435.0
170.0
1622.77
1627.4 4
0.105
Llama 130M
174
172
208
435.8
181.2
4685.9
1340.4
1629.1
1229.5 5
0,105
0,016
4.2
157.0
0.017
Table 2: Main results showing the impact of softmax-/ and OrthoAdam on trained GPT2 and Llama2 models. Utilising softmax-1 and OrthoAdam, significantly reduces the kurtosis and the max activation values of hidden states. Using sofimax-I only is sufficient to reduce first token dominance in attention. We generally find that all combinations of seftmax-I and/or OrthoAdam at a given model size lead to similar performance. Em [8m.a] mean kurtosis of the first token; Em [Km,>1]: mean kurtosis of all other tokens; Em [| mean max absolute activation value of the first token; Em [I mean max absolute activation value of all other tokens. All values are averaged across all layers.
significant change in PPL, each of our proposed methods lead to a significant reduction outlier activations in the hidden states (shown by considerably lower mean layer kurtosis and maximum absolute activation), with the largest reduction observed when both softmax-1 and Oram are used. In particular, for GPT-2 models with 60M, 130M, 350M and 1.4B parameters, the kurtosis without our modifications were 77.9. 141.5, 161.8 and 351.0, while after our modification they drop to 7, 7.3, 3.1, and 3.0. We observe similar results for Llama2-1300 the same as the original version, but kurtosis reduced from 170 to 6.9. Similar to kurtosis, in all cases we see significant reduction the mean activation value. Furthermore, we also observe the drastic drop in first token attention. While the vanilla versions of the model have maximal first token attention of up to 64.8%, after our modification, it is reduced to 1-3%.
where the perplexity is around
5.2 QUANTISATION
We quantise trained models using Absmax and Zeropoini quantisation. Absmax quantisation scales a given tensor (weight or activation) using the absolute maximum absolute value. On the other hand, Zeropoint quantisation shifts the quantised tensor such that the minimum tensor value is the mini- mum representable value. See Dettmers et al. (2022) for exact details on the quantisation schemes.
Experimental Setup. We quantise the trained models using Absmax quantisation using 8-bit in- tegers and the more powerful Zeropoint quantisation using 4-bit integers. In the case of Absmax quantisation, we use 3 different configurations: (1) fine quantisation, where "per-channel" scaling is used for imput activations and weights; (2) moderate quantisation, with "per-tensor" scaling for input activations and weights; and (3) coarse quantisation, with "per-tensor" scaling for input and output activations and weights. In the case of Zeropoint quantisation, we use a single configuration where "per-channel" scaling is used for weights only. We only quantize the linear layers, while the embeddings, normalisation layers and softmax activations are not quantised.
Results. In Table 3 we show the results of quantising the trained models using Absmax and Zero- point quantisation. We experimentally confirm that in all cases, models trained with softmax-1 1 and Adam are more robust to Absmax quantisation schemes than models trained with the canonical softmax function and Adam. The difference in performance is most pronounced when using mod- erate and coarse quantisation maintain performance while models trained with canonical softmax and Adam suffer significant degradation performance. In particular, in the coarse setting, our method outperforms the baseline by up to 36.12 points. For Zeropoint quantisation, we observe that all GPT2 models trained with canonical softmax and Adam become unusable when using 4-bit integer weight quantisation, while models trained with softmax-1 1 and OrthoAdam suffer only small drop in performance. Llama2 models in both cases remain usable after quantisation, but the performance drop is more pronounced when using the canonical softmax function and Adam.
trained with softmax-1 and OrthoAdam are able to
8

---

Published as a conference paper at ICLR 2025
Model
#Panameters
OA S1P
PPL
Foll
moderate
A
line
4-bil
X
31.88
43.53
11.65
34.87
2,59
32.15
0.27
36.6
31.83
32.30
0.47
32.18
0.35
31.89
0.06
33.9
2.1
5.42
130M
22.75
23.21
0.43
23:10
0.32
22.83
0.05
24.0
X
16:37
52.49
19.92
3:55
16.50
0.13
118507.1
118490.7
<
16.31
16.50
0.19
16.46
0.15
16.33
0.02
17.1
0.8
45.05
1.75
13.68
3564.3
<
13.33
13:45
0.12
13:43
0.10
13.34
0.01
13.6
0.2
Lluma2
130M
17.31
20.85
3.54
24.16
20.11
7.07
2.80
17.38
0.07
19.7
2.4
Table 3: Performance of our trained models under various quantisation settings. When using OrthoAdam and softmax- (OA + S1), the performance penalty due to quantisation significantly reduced. The benefits of our proposed changes are more pronounced under more aggressive quantisation settings, i.e., 4-bit weight and coarse 8-bit weight/activation quantisation (vanilla models exhibit catastrophic performance degradation).
5.3 ABLATION STUDY
TableB]showe the results of an ablation study on GPT2 models with 130M parameters. As expected from the discussion in Section we find removing biases from linear layers and varying the position encodings does not prevent first token domination see a small reduction in first token domina- tion when positional encodings are removed. Using softmax-1, first token dominance is mitigated with only ~2% of (query, head) pairs having the first key token as the most attended to key token.
Switching from LayerNorm to RMSNorm with learnt scale for each channel (RMSNorm-M. the normalisation used in Llama2) does not reduce the prevalence of outlier activations in the hidden states. However, switching to RMSNorm with a single learnt scale (RMSNorm-S reduces the mean layer kurtosis and max absolute activation by 400% which remains high. In all of the above cases in which Adam is used as the optimiser, we observe similar perplexity to the initial model (top row). Slight exceptions being the use of rotary and no positional encodings, in which perplexity reduces and increases by 1.3 and 0.5, respectively.
Changing the optimiser to RMSProp leads to increased perplexity (0.5 compared to the initial model), reduced mean layer kurtosis and max absolute activation, by ~500 and ~30%, respectively, when comparing to the equivalent model trained with Adam. In contrast to all previous cases, using SGD with/without momentum (on longer schedule encourage convergence), leads to significant decrease in mean layer kurtosis and max abso- lute activation, by up to 98% and 97%, respectively, when comparing to the equivalent model trained with Adam. However, using SGD requires significantly longer training schedule approach initial model per- formance. Using SGD without momentum leads to significantly higher perplexity (6.8 compared to the initial model). This finding confirms the importance of the optimiser causing outlier acti- vations in the hidden states.
Model
Speed
VRAM
60m-vanilla
14 iter/sec
16.4GB
60m-SI +OA
iter/sec
16.8GB
130m vanilla
7.5 iter/sec
22.6GB
130m-S1+OA
6.0 iter/sec
23.3GB
350m- vanilla
3.3 iter/sec
46.6GB
350m-S1+OA
3.0 iter/sec
47.3GB
1.4B-vanilla
1.0 iter/sec
61.9GB
1.4--S1+AA
1.1 iter/sec
65.0GB
Table 4: Time and m performance.
Using Odtm yields the desirable results from SGD without momentum- -namlyy significant decrease in mean layer kurtosis (140 to 3.0) and max absolute activation (432 to 43.5) and the desirable results from Adam-namely similar perplexity to a model trained with Adam and therefore much faster and better convergence than SGD without momentum.
The final three rows of Table 5 show that using OrthoAdam with softmax-1 and RMSNorm-S leads to the most desirable results, and critically the removal of softmax-1 and the use of LayerNorm or RMSNorm-M reintroduces first token attention dominance and outlier activations, respectively.
Time and memory increase. In Table show that our modifications come with small and tolerable increase in time and memory.
Increasing the sequence length. In Table δ of Appendix, we show that our method is robust to increasing the training sequence length. We show results with models trained in 512 and 1024 sequence length, getting similar results to those of Table

---

Published as a conference paper at ICLR 2025
Biseos
Positio Eacading
Nomolisation
Optimiser
Suftracc
Kanosis
9iiii
Max
Absouue Absouse
LazerNoum
Adun
x
26.9
291.7
0.333
1675.9
LagerNoom
Adun
x
26.9
263.7
0.305
1106.0
X
None Rotary
LayerKorm LayerSorm
Adum
x
27.4
283.3
9.197
1478.7
Adum
x
25.6
341.9
9.336
2577.4
Layerhiom RMSNomm M RMSNom
Adun
244.7
0
Abselute
230.4
0102
628
1400
0.020
RMSProp
Abenttte
92
0.017
Abanini
RMSNom-S
Artbo
30
0,022
Abavini
RMSNim-S
OrthoAdlm
27.9
123.0
0.231
726-4
Absnlnic
RMSNirm-M
Orthosdien
393.9
0.025
737.2
Absolune
LayerVorm
Ortbo/doo
26.6
185.4
00025
514.5
Table 5: Ablation study the impact of various architectural choices on the performance of GPT2 model with sim130M parameter model. "SGD models are trained 8x longer than the others to encourage convergence.
RELATED WORK
Language Models. Language models are based on Transformers (Vaswani et al. 2017 7 ). While there
Transformer-based LLMSs that used the original encoder-decoder architecture such as T5 (Raffel
iet al. 2020). researchers developed models such as BERT (Devlin et al. 2019) and RoBERTa (Liu al.
2019). which are encoder-only. However, most current LLMs such as the GPT (Radford al.
2018 2019; Brown et al. 2020) and Llama series (Touvron et al.,
2023a b Dubey et 20.44 use
a oncoder-only architecture. In our work, we focus on this variant using GPT2 and Llama.
Attention Dominance. BBondarenko et al. (2023) identify the dominance of bland tokens in the attention maps of the BERT Transformer, and suggest complex clipping schemes, additional hy- perparameters, and gating mechanism to to mitigate this issue. Other researchers found the same issue in long-range attention (Xiao et al., (024) and found a workaround using "attention sinks" and discontinuous attention masking. In vision Transformers, Darcet et al. (2024) made the same obser- vation and proposed a solution using "registers". In contrast, we find the root cause of this issue, the softmax in attention, and reformulate it to prevent the first token dominance happening.
Outlier Activations. Previous works have shown that in certain Transformer models which use post-normalisation the norm of the weights of the learnt model must increase (Arora et al. 2019 Soudry et al.) 2018). However the same reasoning does not apply for most recent decoder-only Transformers which
pre-normalisation (Xiong al.] (i.0),i.e. normalisation before the resid- ual connection). A blog-post by Elhage et al. (2023) discusses the presence of outlier activations in the hidden states of Transformer models and rules out numerical precision as the cause. Another blog-post by Miller (2023) posits the activation outliers are caused by the attention mechanism, however, we find outliers and attention dominance are disjoint phenomena. He et al. (2024) identify the presence of outliers and propose an "Outlier Protected Transformer Block" which makes many architectural changes such as removing normalisation layers and severely downscaling the activa- tions at the residual connection. In our contrast, similar to first token dominance, we first find the root cause of this strange behaviour, and then fix it without doing architecture changes.
Outlier-Aware Quantisation. The presence of outliers in the activations of the hidden states has led to a number of works, such as LLM. int (Dettmers et al. (022), per-embedding group quanti- sation (Bondarenko et al.) 2021), and SmoothQuant (Xiao et al. (023) propose varying quantisation schemes to handle the presence of outliers, which require calibration. In contrast, we climinate the presence of outliers in trained models thus enabling the use of the most basic quantisation schemes such as Absmax and Zeropoint quantisation.
7 CONCLUSION
In this work, we study two surprising phenomena in large auto-regressive Transformers: (1) the strong, consistent dominance of the first token in attention maps; and (2) the presence of outlier activations in the hidden states. We propose novel solutions: (1) the softmax-1 1 function remove first token dominance; and (2) the OrthoAdam optimiser which mitigates outlier activations. By doing so, we reduce first token dominance of attention maps by up 95% and the activation kurtosis by up to 99.89 99.8%. Furthermore, our work improves our understanding of Transformers but also offer practical benefits in model quantisation, reducing the quantisation penalty by up to 99.99999 99.9%.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Sanjeev Arora, Zhiyuan Li, and Kaifeng Lyu. Theoretical analysis of auto rate-tuning by batch normalization. In The Seventh International Conference on Learning Representations, 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. stat, 1050:21, 2016.
Hritik Bansal, Karthik Gopalakrishnan, Saket Dingliwal, Sravan Bodapati, Katrin Kirchhoff, and Dan Roth. Rethinking the role of scale for in-context learning: An interpretability-bass case study at 66 billion scale. In Proceedings of the 61 Annual Meeting of the Association for Com- putatiooal Linguistics (Volume I: Long Papers), 2023.
Yelysei Bondarenko, Markus Nagel, and Tijmen Blankevoort. Understanding and overcoming the challenges of efficient transformer quantization. In Proceedings of the 2021 Conference on Em- pirical Methods in Natural Language Processing, 2021.
Yelysei Bondarenko, Markus Nagel, and Tijmen Blankevoort. Quantizable transformers: Removing outliers by helping attention heads do nothing. In Advances Neural Information Processing Systems, volume 36, 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Merbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. Advances in Neural Information Processing Systems, volume 33, 2020.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European Conference on Computer Vision, 2020.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. What does BERT look at? analysis of BERT's s attention. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks NLP, 2019. URL htRp:ttps://att paatps://aclanthooo.gg
org Feg/W19 4828
Harald Cramér. Mathematical Methods of Statistics. Princeton University Press, 1946.
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers. In The Twelfth International Conference on Learning Representations, 2024.
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm.int80:: 8-bit matrix mul- tiplication for transformers at scale. In Advances Neural Information Processing Systems, volume 35, 2022.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics, 2019.
Jesse Dodge, Maarten Sap. Ana Marasović, William Agnew, Gabriel Ilharco, Dirk Groeneveld,
Margaret Mitchell, and Matt Gardner. Documenting large webtext corpora: A case study on the colossal clean crawled corpus. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing. 2021.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, al. The llama 3 herd of models. arXiv preprint 2:Xi::2407.21783 2024.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid- linear units for neural network function approximation in reinforcement learning. Neural networks, 107, 2018.

---

Published as a conference paper at ICLR 2025
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodde Danny Hernandez, Andy Jones, Jackson Kermion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits, 2021. URL https: 7/transformer circuits
Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-Dodss Robert Lasenby, Dawn Drain, Carol Chen, et al. Toy models of superposi- tion. arXiv preprint arXiv:2209. 2022.
Nelson Elhage, Chris Olah, Robert Lasenby, and Shan Carter. Privileged bases in the trans- former residual stream, 2023. URL
privileeed-bassis/ iineeiiiiiiiiiiiiieemm html
Bobby He, Lorenzo Noci, Daniele Paliotta, Imanol Schlag, and Thomas Hofmann. Understanding and minimising outlier features in neural network training. In Workshop on Efficient Systems for Foundation Models II @ ICML2024, 2024.
Dan Hendrycks and Kevin Gimpel. arXiv:1606. 08415, 2016.
Gaussian error
linear units (gelus). arXiv preprint
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In The Third International Conference Learning Representations, 2015.
Ye Lin, Yanyang Li, Tengbo Liu, Tong Xiao, Tongran Liu, and Jingbo Zhu. Towards fully 8-bit integer inference for the transformer model. In Proceedings of the Twenty-Nithh International Joint Conference on Artificial Intelligence, 2021.
Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, et al. Deepseek-v2: A strong, economical, and efficient mixture- of-experts language model. arXiv preprini arXiv:2405.0444 044444 2024.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907. 1692, 2019.
Evan Miller. Attention is off by one, 2023. attention-is-off-bb--oo- html
org
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint 27X1v:2209.11885, 2022.
OpenAl. ChatGPT.
comm 2023.
Long Ouyang. Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kel- ton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul F Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. In Ad- vances Neural Information Processing Systems, volume 35, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems, volume 32, 2019.
12

---

Published as a conference paper at ICLR 2025
Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Bao- hong Lv, Fei Yuan, Xiao Luo, et al. Scaling transnormer to 175 billion parameters. arXiv preprint 11Xi::2307.11995. 2023.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing with unsupervised learning. 2018.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21, 2020.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The im- plicit bias of gradient descent on separable data, Journal of Machine Learning Research, 19, 2018.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 568, 2024.
James Sylvester. Thoughts on inverse orthogonal matrices, simultaneous signsuccessions, and tes- sellated pavements in two or more colours, with applications to newton's 's rule, ornamental tile- work, and the theory of numbers. Philosophical Magazine Series 1, 34, 1867.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint 2rXiv:2302.13977 2023,
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.09288. 2023b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems, volume 30, 2017.
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoc V Le. Finetuned language models are zero-shot learners. In The Tenth International Conference on Learning Representations, 2022.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowiez, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, 2020.
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 Proceedings of the 40th International Conference on Machine Learning, volume 202, 2023.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference Learning Rep- resentations, 2024.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture In International Conference on Machine Learning, 2020.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. In Advances in Neural Information Processing Systems, volume 32, 2019.
13

---

Published as a conference paper at ICLR 2025
A TRAINING DETAILS
In this section, we provide details on the training of our models.
batch size of 512 and in all experiments using Adam OrthoAdam the optimiser, we use peak learning rate of 10-3 This excludes the experiments in Section 5.3 which SGD the optimiser, which peak learning learning rate schedule with linear warmup for { {1000, 2000, 6000, 10000} steps for models with (60M, 130M, 350M, 1.4B} parameters respectively. Note that we use reduced number of steps for the 1.4B model due to computational constraints. For the main experimental results in Tables 2 and 3 we train the models with {60M, 130M, 350M, 1.4B) parameters for {160k, 320k, 960k, 600K} steps respectively. For the ablation study in Section [.3 we train GPT2 models with 130M parameters for 40k steps with 2000 warmup steps. We use maximum sequence length of 256 tokens, which we find is sufficient to observe the anomalies of first token attention dominance and large outlier activations found in popular pretrained models such as GPT2 (Radford et al. 2019) and Llama (Touvron al., [203a b Dubey et al. 2024). The result of our training setup that models trained for the main experimental results with [60M, 130M, 350M, 1.4B) parameters are trained on trained on 5B tokens. We train models on 8 NVIDIA 32GB V100 GPUs using the Pytorch deep-learning framework (Paszke)
In all experiments
of 0.2. In all experiments we use cosine
(21B, 42B, 126B, 79B) tokens respectively. The ablation experiments
al. 2019) and the HuggingFace Transformers library (Wolf al. 2020).
14

---

Published as a conference paper at ICLR 2025
B LONGER SEQUENCE TRAINING
Model Size
Setup
Full
Coarse
(Coarse)
Moderate
(Moderate)
Fine
(Fine)
60M-256
Vanilla
31.88
43.53
11.65
34.87
2.99
32.15
0.27
60M-512
Vanilla
32.66
48.55
15.89
37.24
4.58
33.07
0,41
60M-1024
Vanilla
33.52
57.68
24.16
38.22
4.70
33.80
0.28
60M-256
S1+OA
31.93
32.46
0.53
32.32
0.39
32.00
0.07
60M-512
S1+OA
31.83
32.30
0.06
32.18
0.35
31.89
0.47
60M-1024
S1+OA
32.25
32.85
0.60
32.73
0.48
32.32
0.07
130M-256
Vanilla
22.89
46.49
23.60
28.31
5.42
23.07
0.18
130M-512
Vanilla
22.80
42.34
19.54
28.14
5.34
22.98
0.18
130M-1024
Vanilla
22.93
38.78
15.85
29.04
6.11
23.16
0.23
130M-256
ST+OA
22.78
23.21
0.43
23.10
0.32
22.83
0.05
130M-512
S1+OA
23.16
0.43
23.04
0.31
22.79
0.06
130M-1024
S1+OA
23.87
23.28
0.41
23.19
0.32
22.94
0.07
Table 6: Performance results for various model sizes and setups under longer sequence length.
In Table o we provide results when trained with sequence length of 512 and 1024, and compare with sequence length 256. As can be seen, our model is very robust when we increase the sequence length, showing no noticeable performance drop in perplexity be it under the general setting, or when quantized. In all cases, especially under quantization schemes, our method outperforms the vanilla one when trained with longer sequences.
C LARGER LLMS
We show that the first token attention and the outliers happen also in large modern LLMs such as Llama-3.1-88 Furthermore, these issues happen regardless if the training is done unsu- pervised manner (next-token prediction) supervised manner (intruction tuning). We down- loaded Llama-3.1-8B http:///wggggaaaa...oom///aa/aaa
and Llama3. 1-mma3.--B-Insttreee
In Table we show that these large models, the first token attention increases (at over 95% compared to the results shown in the main paper. Furthermore,
also checked the cumulative sum of attention to the first token and found it out to be at 73.49%. In other words, 73.49% of the entire attention in Llama-3.1 1-8B is the first token. This can be interpreted that the larger the network, the more specialized the heads are, and most of the heads will simply do nothing. Attending on the first token the mechanism the Transformer has developed to learn to do nothing. We also check the kurtosis of Llama-3.1 1-8B, showing that the method has very high kurtosis for both the first token and on average.
Finally, we show that these results remain very similar if the model is finetuned in instruction data. We observe that the first token attention and first token kurtosis virtually the same in Instruct model as in the original one, while the average kurtosis actually increases in the Instruct model. Thus, we conclude that our findings stand for modern large LLMs, regardless they are finetuned in instruction data or not.
Method
1 atteetttnoonn ferfffrttttokrrr
1st kurtosis
Average kurtosis
Llama-3. 3.1BBB
95.45
73.49
1227
55
Llama-3. 1-----IIsstucce
95.53
70.13
1228
69
Table 7: Comparison of attention and kurtosis metrics for I.Lama-3.1 1-BBBBBBBBBBBBBB and I.lmm-3.11 1-8-88-----
D INSTRUCTION TUNING
To complement the previous experiment, we run a new experiment, doing instruction finetuning (supervised learning) in Alpaca dataset. We compare the results of model trained with canonical
15

---

Published as a conference paper at ICLR 2025
Method
P/uti
Courrs
A
P
Amul
P/nnc
A
(26)
(
Egcloads 364.75
Akartari
GPT-2-130M
18.35
31.03
12.07
21.09
2.76
18.49
0.16
63.57
41.49
69.31
GPT-2-130M + S1 OA
18.29
18.31
0.02
18.31
0.02
18.29
0.0
2,4
0.8
2.97
2.96
Table 8: Results instruction tuning using GPT-2-130M as baseline and comparing with our approach. Prull represents perplexity without any quantization. Proarees Pmod and Pfine represent perplexity under coarse, moderate and fine quantization. Similarly, Acearee: Aodd and A fine represent the increase perplexity under these three quantization schemes compared to the not quantized method. Late (%)) represent the percentage of tokens that attend to the first one, lat (sum) represent the cumulative sum of the attention in the first token, Nurtosss represent the first token kurtosis, while Aburtosis represents the average kurtosis in the network. As we can observe, our method reaches best results under every setting, while reducing the attention on the first token and kurtosis.
softmax and Adam, compared to our method with softmax-1 and OrthoAdam. We present the results GPT-2-130m models.
We show that while both models reach roughly the same perplexity, only our model keeps the same perplexity under all three quantization schemes. On the other hand, the vanilla model drops for 12.07 points in coarse quantization, 2.76 in moderate quantization, and 0.16 in fine quantization.
To evaluate that this is a direct effect of the first token attention and outliers, we also present the results in the percentage of maximum attention in the first token, the cumulative attention score of the first token, the 1st token kurtosis, and the average kurtosis. We show that in the vanilla model, the maximum attention is in 65.57% of cases in the first token, and the cumulative attention on the first token is 41.49% 1.... In contrast, our method has maximum attention on the first token is 2.4% and the first token contributes only 0.8% of the attention. Furthermore, while the vanilla model has kurtosis of 564.75 for the first token, and 69.31 for the average token, our method has 2.97 kurtosis for the first token, and 2.96 kurtosis for average token, very similar the kurtosis of normal distribution (3).
In this way. empirically show that our findings of the main paper stand also for instruction tuning training. Models trained with instruction tuning drop in accuracy under quantization schemes because of their outliers, while our method remains stable and as we showed, has the same kurtosis as normal distribution and not high attention in the first token.
16

---

Published as a conference paper at ICLR 2025
E
OPTIMISER BASIS
Outout Weight Layer:
Moments Parameter Layer:
4x 13
2x
1017
MLP Weight Layer:
10-4
10.0
-.5
-s.0
10
MLP Weight Layer:
Layer:
10-
10-
MLP Weight
Layer:
Fa
355
Weight Layer:
Layer:
15
10-
10
16
Outout MLP Weight
Second Moments Parmmteer
:10
10-4
16-
Outout MLP Weight
Second Moments Parmeter
10t
:00
167
Outout MLP Weight Norm
Second Moments Parmeter
15
100
Outpout MIP Weight Norm
Sum of Second Moments of Paramette Layw:
20
:0
Dutpuu MLP Weight Norm Layer: 10
Stam of Serond Moments of Parameter Layen:
J5
10
10t
Dutput MLP Weightt Norm Layer: 11
Stam Serond Moments of Parameter Layen: 11
30
20
4x
2x
10
Output MIP Weightt Norm Layer: 12
Sum of Second Moments of Parameter Layer. 12
-20
Ex13
4x10
10
Figure 4: Left: MLP output feature weight euclidean norm each layer in our GPT2- 130M model training with canonical softmax and Adam. Right: Sum of gradient second moments for the corresponding MLP output feature weight when training with canonical softmax and Adam. Training with softmax and Adam leads to small outlier gradient second moment moving averages with lead to disproportionately large gradient steps for the feature dimensions containing these small outliers. This in turn leads to large outliers the model weights.
17

---

Published as a conference paper at ICLR 2025
Outputt Weight Nerm Layer
Second Nomwets Parameter Layer
-00
4x2007
Output MIP Weight Norm Layer
Sum Second Noments Parmmeter Layer
3.0
24xx10
24x10-
4.5
22
2x.00-
160
Output MLP Weight Norm Layer
Sum of Second Noments of Parameter Layer
33.00
22-10-1
4.72
2x30
14x
10
Output Weight Norm Layer
Sum Second Morwents Parameter Layer
17-1e-7
5.25
5.00
1.5 130
14x10-
4.75
Norm Layer.
Norwets Layer
10-
20-0-7
5.0
14x
Output M&P Weight Norm Layer
Sum of Second Norwets of Parameter Layer
6.00
13 xip
5.75
14 14-
1.3 32x0
1.3 x300
5:23
Outtput M.P Weight Nor Layer?
Sum of Second Nomeets of Parameter Layer
15 25-1
Outsus Weight Norm Layer
Sum Second Moments Pa rameter Layer
+.75
1.25 300
1.06 40-7
8.25
20
6.00
Outout Norm Layer
Second Moments Pa rember Layer
6.8
6.6
$x
-64
Output MIL Weightt Norm Layer 10
Sum of Second Nom of Parameter Layer: 10
7:25
9x-00)
7.00
6.75
Ouiput MIP We ght Norm Layer. 11
Sum of Second Mom ants of Parameter Layer: 11
Rx30-)
7.0
7x
Output Weight Layer
Sum of Second ants Parameter Layer:
10
9x
73
8x 30-
7x-3--
20-1
Figure 5: Left: MIP output feature weight euclidean norm for each layer in our GPT2-130M model training with softmax-1 and OrthoAdam. Right: Sum of gradient second moments for the corresponding MLP output feature weight when training with softmax- I and OrthoAdam. Training with softmax-1 and OrthoAdam leads gradient second moment moving averages be considerably more uniform than training with softmax and Adam. This prevents large gradient steps from being taken removing outlier weights in the model.
18

---

Published as a conference paper at ICLR 2025
Our observation of outlier activations in the hidden states of language models leads us the clusion that some aspect of the training process is "basis- messis-dependent..
indiiidedeeendent" func- tions those which equivariant under orthogonal transformations. A function f(x) is basis- independent if f(@x) (z( for any orthogonal matrix Q. We remove biases from our linear layers and introduce single-scale version of RMSNorm to remove "basis-dependen" effects from the model itself. It is straightforward to show that linear layers with biases (i.e., afline transfor- mations) and multi-scale RMSNorm (as is standard) are dasis-independent (i.e., they are basis- dependent). SGD and SGD with momentum are basis- iesss-independent: We can show that the standard SGD update rule is basis-indeepnnee
Given the standard SGD update rule ltt1 = 0. (.-1)V(((.) transforms the parameters to - 0, = Q9. where 0g is an orthogonal matrix:
1c+1 = 0, -
(E.1)
By chain rule:
Vi(0,) vL())) 00:
((,) Q" QV)( 00,
(E.2)
0c+1 = 0, - (VVV((() = Q(, LVV)) = Q0++1-
SGD with momentum follows similarly.
However, Adam and RMSProp not basis-independee because of tracking the second-order mo- ments. Let U/ be the second-order moment the gradicnts at step t. then
vt+1 = + - (1 @ VL()))
(E.3)
B2b; + (1- vi B20 + (1- o QVL(O.)
(E.4)
* Qu++1
The above shows that SGD and SGD with momentum provide basis-independent update rule which proportional to the gradient, on the other hand, Adam and R.MSProp are not basis-inde and provide updates that depend on the element-wiss root of the second-order moment of the gradient. We believe that basisedep functions (whether in the model or the optimizer) are the cause of the outlier activations we observe in the hidden states of language models. The second-order moment tracking in Adam and RMSProp allows disproportionately large gradient updates in certain weights of the model especially in the early steps of training where the moment moving averages not well- calibrated (a result of adaptive per parameter learning rate scaling). Therefore we expect features for which the second-order moment is small to have disproportionately large weights in the model. Disproportionately large weights in particular dimensions of the model cause the outlier activations observe. Using ArthoAdam, the moving average moments of the gradients are computed in unique random orthogonal basis for each parameter, which prevents small values of the second-order moment from causing disproportionately large updates in the model parameters (outlier gradients large or small in the model basis are transformed to be likely of similar magnitude in the orthogonal basis).
The outlier weights in the model are generally contained in the output linear layer of each MLP block. To verify our intuition above, we have computed the norm of the output weights in output linear layer of each MLP block in GPT2-130M model trained with/without softmax-1 and Or- thoAdam. We additionally plot the sum of the second-order moments of the gradients of the output weights in each MLP block. We observe that in all cases of large outlier weights in the model, the sum of the second-orded moments of the gradients of the output weights in the corresponding MLP block small outliers. This is consistent with intuition that the small outlier second-orde moments of the gradients are causing the large outlier weights in the model. We show these plots in Figure and Figure 5 respectively
19

---

Published as a conference paper at ICLR 2025
F TRANSFORMING THE OUTPUT INTO ORTHOGONAL BASIS
Normalized Depth Kurtosis
10a
Layer Output Activation Kurtosis
102
Vanilla (X)
S10A (X)
S10A (XQ)
S10A (XQt)
10
Normalized Depth
Figure Transforming the output into orthogonal basis
We investigated transforming the output of each layer (i.e., the added activations to the hidden states for said layer) into the orthogonal basis used by that output layer in Ortho We 350M model trained with Softmax-I and OrthoAdam. Given layer i has output X; E RLxD and the orthogonal basis for the output layer in Odmm is Q; E OD, where L is the sequence length, D is the hidden dimension and O is the set of D x D orthogonal matrices. We plot the activation kurtosis of X,, Q,Q X, and X, V, where V; is the right singular vectors X,. Using V as a transformed basis gives baseline for how large one could increase the kurtosis of the activations by transforming them into an orthogonal basis. Additionally we give the activation kurtosis for our GPT2-350M vanilla model.
GPT2-
20

---

Published as a conference paper at ICLR 2025
G KURTOSIS GROWS WITH THE NUMBER OF DIMENSIONS IN TRANSFORMERS
In this section, how the kurtosis of the hidden states grows with the number of dimensions in the hidden states. This is something we observe empirically in the hidden states of transformer models and is a key motivation for our work. Table 2 shows the kurtosis of the hidden states of transformer models trained without softmax- or OrthoAdam grows as the model size increases, as does the maximum activation value in the hidden states.
observations from the hidden states of transformer models to illustrate
To make this mathematically rigorous, we shall consider a simple example, in which we shall ap- proximate the hidden states of a transformer model at single token position as a D-dimensional vector comprising of the sum of a scaled one-hot vector and a standard normal vector.
Consider a D-dimensional vector x which is the sum of two D-dimensional vectors and z, where c; is the ith unit vector in the standard basis, x E RR,, a E and N(O,Ip), Therefore the elements of x are given by:
zj = ===
+ 2g
for j = .,,,...,
where 6zg is the Kronecker delta function. The mean is given by:
E [ry]]
= E [adi + zn] = aB [Øt]] + E [zy]
a
a as [zy] = 0 by definition of the standard normal distribution
+
0
=
D
D
(G.1)
D
The variance is given by:
o² Var [x,]] = Var (ad:s zj]
Var [6.y] Var [zy] ad:; and 2y are independent in our model
Var [ōt]]
Var [zy] by definition of the standard normal distribution
(EE (6 = T 1 1
Var [6g]
E [8t]]
Therefore:
2 a²
D
+1
(G.2)
D
The kurtosis of the elements of x is given by:
(
47
E
(Iz 11)
Kurt[æ,] E
0²
When + i:
E (1g 1))]]
E
2j D
E
4 4
+ 623
a
- 42j (0) 3
+
a
D
D
D
D
= E [29]
4E
a D
+ 6E
a0
D
4E
1zy]
(0)
+
a
D
As E [>]]]
0 and E [=/]]
= 3:
a 2
3+6 6
+
21

---

Published as a conference paper at ICLR 2025
When d:
E (1g 1))² = E
a + 2j ]
D )22 }
E
α
1 11-11 =
+ 4
=a(a((
1 (a(1-
] ) 2j
1
3
E
a
D
+
a
1
1
z + (a (a(1-
(1-
1 +)
=(a
1
a (1 | ) +3
Therefore, the overall fourth moment of the elements of x is given by:
(
) + +6 a 1 - )) +3
E [00 (105-1)33 -))
D
D --1
) + (B)*)
+
3 3+6
And the kurtosis of the elements of is given by:
b (aa (1- + 6 (1 - b))² +3) + D-1
(3+6 (z)) + +z)))
D
Kurt [xtz]
02 (1 } + 1)²
3 +
4r
6a
6
Kurt [zy]]
17
17
17
172
133
1f
(G.3)
1 2a
20
26
D
D
D2
D
At this point, we can see that Kurtosis is function of and D, however if we consider the limit as D + we can see that Kurt[zr]] - 3, i.e., the kurtosis of Gaussian distribution. However, this neglects the importance of the scaling factor a which we know empirically larger than the dimensionality of the hidden states. The table below summarises the maximum activation values (analogous to a) and the dimension of the hidden states for the models we trained. Given this
Model
#Parameters
Model Size (D)
Max Activation (a)
60M
512
1856
130M
768
7018
GPT2
350M
1024
40196
1.4B
2048
56798
Llama
130M
768
4623
Table 9: Model sizes and maximum activation values for the models used in our experiments.
empirical information, make the conservative assumption that D. Under this assumption
22

---

Published as a conference paper at ICLR 2025
which is supported by empirical observations, Equation
G.3 simplifies to:
3+
6D
4D'
6D
+ 6D'
Kurt [zz]]
D
D
D2
D2
D3
174
1+ 2D2
D
20²
132
+
13
2D'
128
+ 124
D
3+ D + 6D - 4D2 -6+62 6D - 3
1 2D - 2 D
2D +1
D 4D + 12D 6
D
12
Kurt [x,]]
4+ D D2
O(D)
(G.4)
Using our conservative assumption that D 1, we can see that the kurtosis of the hidden states grows linearly with the dimensionality of the hidden states when D is in the region of 10 10 as is the case for transformer models.
This simple example serves as mathematical illustration of the empirical observations make the hidden states of transformer models. have shown that the kurtosis the hidden states is expected to grow linearly with the dimensionality of the hidden states, and the issue of outlier activations is expected to grow as the hidden states of transformer models grow size.
23

---

Published as a conference paper at ICLR 2025
H ORTHOGONAL TRANSFORMATIONS AND REDUCTION IN -NORM AND KURTOSIS
From simple model in Appendix we have simplified model of Transformer hidden states, E R., where the first element is a and the rest are standard normal random variables.
X ac; z where 2j ~ (0)
From this model, we can compute the expected no-m:
D = Σ == Σ
D
|
2) =²
DVa [25] = a²
(H.1)
j=1
Using the triangle inequality, compute a range for the too -orm:
[||
E
max USiSD
la Ia+ + zzl. max
j7t
Given a > 1, we can drop the terms for j # and compute the expected 6.0 norm using the 7th element:
x xo. = E ||a + =|]
la + zul S |a |
Using folded normal distribution properties, E [|=i]] = V << a, therefore:
E nol
Given that a > 1, we can safely assume that ||x|
2 a. Therefore:
x
a
E
x
VD + a
Note from Table that the maximum activation value, a, is generally much larger than the model size, D.
x
E
1
(H.2)
x
We find this empirically be the case in the middle layers the Transformer models we study (see plots in Appendix
||x]| ||x|
The co-norm of x can be thought of as proxy for the extent of outliers in vector. If then vector has at least one large outlier and consequently high kurtosis.
1,
We will now show that applying orthogonal transformation vector can reduce the l norm constrained to a fixed fo-nomm. Using the same definition of x as above, let Q E Dx be an orthogonal matrix and let y
Qx
ll 2 y"y x Qx=""x==""" Txx |x12
Elly| = E||||| = a
(H.3)
This standard proof shows that applying an orthogonal transformation to a vector does not change the & &-mmr of the vector. It can however lead dramatic reduction in the form of the vector. We will now show that for a vector, y E R constrained to have a fixed E--nomm, va + D. the -orm of vector can be reduced significantly by applying an orthogonal transformation such that y
D + z', Vj E [1, D], where z ~ 1).
y Qx Q (ac; 22) aQc, Qz
24

---

Published as a conference paper at ICLR 2025
Select Q such that Qe
=
( 1
VD vD
) given Q orthogonal, Qz ID).
Illy
2
a
+ 22 D,
using extreme value theory Cramer 1946)
VD
V2D
VD
The expected ratio of &o norm to fg-norm is:
|ly
1
(a+V2DIn D 1)
y
E
|
2
E
Iby
Using the same conservative assumption in Appendix G Ghat D:
D, Table 9
shows empirically
y &
D + In D 2DD22 D 2 In D 2 D
E
1 D
D +
As D grows, the last term of the numerator and the 1 in the denominator become negligible:
yy ly
D 2+222
O (
D
D
Therefore, under an orthogonal transformation, the lo norm to l2- norn ratio can be reduced sig- nificantly. It is trivial to show that Kurt[ys] and see many of experiments which OrthoAdam and softmax-1 exhibit this behaviour (see plots in Appendix [.2]
x
y &
1
ae; +2,
E
|
1
y
Qx,
E
|
~
D
2
12
6
Kurt [rz]] D 4
O(D)
Kurt [y,]
=
D
D
The exact form of Q can be computed numerically or constructed using appropriately normalised Hadamard matrices (Sylvester 1867).
25

---

Published as a conference paper at ICLR 2025
I LAYER PROGRESSION OF FIRST TOKEN ATTENTION DOMINANCE, KURTOSIS, & -NRM TO f2-NORM RATIO AND MAXIMUM ABSOLUTE ACTIVATION
For brevity, we give metrics for the first token attention dominance, hidden state kurtosis and ab- solute maximum activation averaged over all layers in Table 2 which gives the results of the main experiments in our work.
However, the layer-wise progression of these metrics is also of interest and can provide insights into the behaviour of the model. Additionally, we provide the same metrics for popular pretrained GPT2 and Llama models to show the similarity to models trained without softmax-1 and OrthoAdam.
Finally, to establish a relationship between activation kurtosis and the -nrm to & 2--ormm ratio, we calculate the Pearson's 's correlation coefficients between per-layer kurtosis and per-layer -oom fg-norm ratio for all models our main experimental results from Table
All metrics computed the same validation set of the C4 dataset (Raffel al. 2020) in the main paper (Section 5).
I.1 FIRST TOKEN ATTENTION DOMINANCE
We begin by examining the progression of first token attention dominance across layers. We calcu- late the percentage (head, query) pairs where the query token attends most the first (key) token. Given different models have different number of layers, we normalise the layer index to the range [0, 1] for each model.
We find general trend across our trained models which use the canonical softmax function where the first token attention dominance begins low in the initial layers where models do initial processing of all input tokens. The dominance rises to peak in the middle layers where heads specialise to specific sub-tasks and so the first token is attended to as a default "no-op" (Bondarenko et al.) 2023 Clark et al., 2019). Finally, the dominance decreases in the final layers where the model "detokenises" the features back into token space.
1.1.1 GPT2-60M
go12: 50m
Propertin of First lokens with Maximum Arsention
1.0
go02: 50m, OA
go02: 60m,
gal2: 80m, 51, DA
8.n
0.6
0.2
0.0
2.0
0.2
0.4
5.6
0.8
2.0
Normalised Depth
Figure 7: Layer-wise progression of first token attention dominance for GPT2-60M. The x-axis is normalised the range 1]. S1/OA denote models trained with softmax-1 and/or OrthoAdam.
26

---

Published as a conference paper at ICLR 2025
1.1.2 GPT2-130M 130M
3012: 130m
2012 130m,
Proporiion of First Tokens with Maximuum Attention
150m,
130mm
6.6
0.0
022
0.4
0.6
0.8
10
Normalised Depth
Figure 8: Layer-wiss wise progression of first token attention dominance GPT2-130M, The x x-axis normalised the range (0, 1]. S1/OA denote models trained with softmax-1 and/or Adam..
1.1.3 GPT2-350M AND GPT2-1.4B
3012 350m
2012 350m, $1,
Propertin of First lokens with Maximum Arsention
2012 1o
9p12: 15, S1, DA
8.n.
0.6
0.2
0.0
3.0
02
0.4
96
6.8
2.0
Normalised Depth
Figure 9: Layer-wise progression of first token attention dominance for GPT2-350M and GPT2-1.4B The
x-axis normalised the range [0, 1]. S1/OA denote models trained with softmax-1 and/or OrthoAdam.
27

---

Published as a conference paper at ICLR 2025
1.1.4
LLAMA-130M 130M
Proportion of First Tokens with Maximum Artention
1.0
Farme: 130m,
0.8
0.4
0.2
Normalised Depth
Figure 10: Layer-wise progression of first token attention dominance for Llama- 130M. The x-axis is normalised
the range (0, 1]. S1/OA denote models trained with softmax-1 and/or Ortmo
1.1.5 POPULAR PRETRAINED MODELS GPT2 AND LLAMA
Proportion of First Tokons with Maximum Attention
0.4
op32
ge/2-al
0012 modium
Libna 2:76
Meta:llmm--31-00
Normalssed Depth
Figure 11: Layer-wise progression of first token attention dominance for popular pretrained GPT2 and Llama models. The x-axis is normalised to the range [0, 1).
1.2 ACTIVATION KURTOSIS
Next, we examine the progression of activation kurtosis across layers. As observable in Table the kurtosis of the first hidden state is significantly higher than the other hidden states and plot the kurtosis the first hidden state only for brevity.
We observe in the plots below that models trained without OrthoAdam exhibit general trend of increasing kurtosis as the hidden states progress through the layers. Demonstrating that multiple layers of the model contribute to the emergence of large activation values. Models trained with
28

---

Published as a conference paper at ICLR 2025
OrthoAdam but not softmax-1 exhibit similar trend, but with lower kurtosis values initially. Fi- nally, models trained with both Ort and softmax-1 exhibit consistent small kurtosis across layers-areen around the value of 3 which is the kurtosis of Gaussian distribution. Interestingly, GPT2- 60M and GPT2-130M show small rises in the final layers-the the cause this is left for future work.
We find that models show reduction in kurtosis in the final layers, again attribute this to the "detokenisation" of the features back into token space.
1.2.1 GPT2-60M
go22 60m
po22 60m,
go52:
50m, DA
gel2: 20m, 51, PA
Hidden State Activation Kurtosis
3.0
02
DM
64
6.8
1.0
Normallsed Depth
Figure 12: aayer-wise progression of activation kurtosis of the first token position for GPT2-60M. The x-axis normalised to range (0, 1]. S1/OA denote models trained with softmax-1 and/or OrthoAdam.
I.2.2 GPT2-130M
3p12 130m
2012 130m,
2012 130m 0A
9p12: 130m, 51, DA
Hidden State Activation Kirtosis
3.0
02
DA
646
6.6
1.0
Normalised Depth
Figure 13: Layer-wise progression of activation kurtosis of the first token position for GPT2- 130M. The x-axis normalised the range (0, 1). S1/OA denote models trained with softmax-1 and/or OrthoAdam.
29

---

Published as a conference paper at ICLR 2025
1,2.3
GPT2-3501 AND GPT2-1.44 1.4B
2012: 350m
2012 250m, 51, DA
3012: 1a
3p12: 15, S1, DA
Hidden State Activation Kurtosis
3.0
02
D4
64
6.8
1.0
Normalised Depth
Figure 14: Layer wise progression of activation kurtosis of the first token position for GPT2 350M and GPT2- 1.4B. The x x-axis is normalised to the range [0, 1]. S1/OA denote models trained with softmax-1 and/or Or- Adam.
1.2.4
LLAMA-130M
tama: 130m
tama: 130m,
lama: 130m,
Bamm: 130m, 51, DA
Activation Kirtosis
Hidden State
3.0
02
D4
646
6.6
1.0
Normalised Depth
Figure 15: Layer-wise progression of activation kurtosis of the first token position for Llama- 130M. The x-axis normalised to the range (0, 1). S1/OA denote models trained with softmax-1 and/or OrthoAdam.
30

---

Published as a conference paper at ICLR 2025
1.2.5
POPULAR PRETRAINED MODELS- GPT2 AND LLAMA
-+-
1piz
9012 medium
Usmo2 76-hf Nota-Jama-3.340
sp12-large
10r
Activation
204
10a
Normalised
Figure 16: Layer-wise progression of activation kurtosis of the first token position for popular pretrained GPT2 and Llama models. The x-axis normalised to range [0, 1].
1.3 -NRR TO 2-- NNRM RATIO
The plots below show the progression of the -nom to 22-norm ratio across layers. We observe that models trained without OrthoAdam exhibit general trend of inereasing ratio as the hidden states progress through the layers. Once again as this ratio is maximal in the first hidden state, we plot the ratio of the first hidden state only for brevity (as done for kurtosis).
The trends are similar to the kurtosis plots and so the same commentary applies.
1.3.1 GPT2-60M
0o2::
1.0
gol2: 80m, DA
Hidden Infirity Norm Norr
0.4
0.2
Normalised Depth
Figure 17: Layer-wise progression of the loo norm to form ratio in the hidden states of the first token posi- tion for GPT2-60N. The x-axis is normalised to the range [0, 1). S1/OA denote models trained with softmax-1 and/or OrthoAdam.
31

---

Published as a conference paper at ICLR 2025
1.3.2
GPT2- 130M 130M
DDG2: 130m
130m.
1.0
gol:: 130m,
992. 130m, $1. 04
0.6
Hidden Infinity
04
0.2
0.0
0.2
1.0
Normalised Depth
Figure 18: Layer-wise progression of the -oomo to 4--norm ratio in the hidden states of the first token position for GPT2-130M. The x-axis is normalised the range [0, 1]. S1/OA denote models trained with softmax-1 and/or OrthoAdam.
I.3.3 GPT2-350M AND GPT2-1.4B
350m
1.0
gol2: 36,51,
Hidden Infinity Norm
0.4
0.2
Normalised Depth
Figure 19; Layer-wise progression of the Ls to fo-norm ratio in the hidden states of the first token position for GPT2-350M and GPT2-1.4B. The x-axis is normalised to the range [0, 1). S1/OA denote models trained with softmax-1 and/or OrthoAdam.
32

---

Published as a conference paper at ICLR 2025
1.3.4 LLAMA-130M
lama: 130m
Hame: 150m,
Nama: 150m,
Iamma: 1.0m,
Hidden Infinity Norm 2-Norm
08
0.0 3.0
522
0.4
1.0
Normalised Depth
Figure 20: Layer-wiss progression ls norm 2--nrrm ratio the hidden states first token position for Llama-1300 130M. The x- x-axis is normalised to the range [0, 1]. S1/OA denote models trained with softmax- and/or OrthoAdam.
1.3.5 POPULAR PRETRAINED MODELS GPT2 AND LLAMA
Hidden Infinity Norm 2 Norm
0.6-
84
gue2
ge
Liame-2-Tb-hf
large
Mota U.1b331448
0.
0.2
Normalised Depth
Figure 21: Layer-wise progression -oom to 2z-norm ratio the hidden states first token position for popular pretrained GPT2 and Llama models. The x-axis is normalised to the range (0, 1].
To further clarify that in Transformer models the l to fnrm ratio is proxy for activation kurtosis, we calculate the Pearson's correlation coefficients between the two metries for models in our main experimental results from Table 2 and public GPT2 and Llama models. The results are shown in Table I0 We find a strong positive correlation between the two metrics across all models which reinforces intuition that using orthogonal matrices transform the gradients in the optimiser an effective way to mitigate the emergence of large activation values, as an orthogonal transformation reduce the norm of vector substantially for given fornr..
33

---

Published as a conference paper at ICLR 2025
Correlation Kurtosis to Norm Ratio
Model
iParameters
Softmax+1?
OrthoAdam?
Other Tokens
0.961
0.932
0.932
0.934
60M
0.986
0.972
0.968
0.970
0.988
0.932
GPT2 (Ours)
0,927
0.924
130M
0,992
0,962
0.935
0.953
0,990
0.929
350M
0.998
0.997
0.988
0.952
1.4B
0.994
0.995
0.931
0.903
0.864
0.877
Llama2 (Ours)
130M
0,931
0.905
0.560
0.975
GPT2 (Public)
137M
0.985
0,944
GPT2-Modium (Public)
350M
0.969
0,846
GPT2-Large (Public)
812M
0.985
0.896
GPT2-XI. (Public
1.6B
0.956 0.987 0,928
0.939
I.lama-7B (Public)
6.7B
0.902
Llama3.1-8B (Public)
SB
0.915
Table 10: Correlation of the kurtosis and norm-ratio of the hidden states of our trained models and popular pretrained models.
1.4 MAXIMUM ABSOLUTE ACTIVATION
Finally, we examine the progression of the maximum absolute activation across layers.
1.4.1
GPT2-60M
Max Absalute Hidden State Activation
20"
207.
30
Normalised Depth
Figure 22: Layer-wise progression of the maximum absolute activation in the hidden states of the first token po- sition for GPT2-60M. The x-axis is normaliscd to the range 1]. S1/OA denote models trained with softmax-1 and/or OrthoAdam.
34

---

Published as a conference paper at ICLR 2025
1.4.2
GPT2-130M 130M
Max Absolute Hidden State Activation
100-
201
get2:
9.0
02
0.4
0.6
0.8
1.0
Normalised Depth
Figure 23: Layer-wise progression of the maximum absolute activation in the hidden states of the first token position for GPT2-130M. The x-axis is normalised the range [0, 1]. S1/OA denote models trained with softmax-1 and/or OrthoAdam.
1.4.3 GPT2-350M AND GPT2-1.4B
Max Absalute Hidden State Activation
20"
207.
0012:
Normalised Depth
Figure 24: Layer-wise progression of the maximum absolute activation in the hidden states of the first token position for GPT2-350M and GPT2-1.4B. The x-axis is normalised to the range [0, 1). S1/OA denote models trained with softmax-1 and/or OrthoAdam.
35

---

Published as a conference paper at ICLR 2025
1.4.4 LLAMA-130M
Max Absolute Hidden State Activation
30%
207
201-
Itama:
5.0
Normalised Depth
Figure 25: Layer-wise progression of the maximum absolute activation in the hidden states of the first token position for Llama-130M. The x-axis is normalised to the range [0, S1/OA denote models trained with softmax-1 and/or OrthoAdam.
1.4.5 POPULAR PRETRAINED MODELS GGP2 AND LLAMA
101
Aboolute
Meto
Normalised
Figure Layer-wise progression the maximum absolute activation the hidden states of the first token position for popular pretrained GPT2 and Llama models. The x-axis is normalised to the range (0, 1].
36

---

Published as a conference paper at ICLR 2025
J HIDDEN STATES OF PRETRAINED MODELS
In this section, we present the progression of hidden states of popular pretrained models. This shows how models establish outlier activations and how they persist in the same feature dimensions across layers. For each model we show the absolute activation values in the features containing the largest activations. We show the mean across layers, the first layer, and & of the layers.
GPT2 Hidden States
Input Sequence: Overall Mean Hidden States
Input Sequence: Overall Mean Hidden States
Feature
10a
Taken Position
10
12
10
116
114
10
118
122
10-1
124
10-2
Layer: Feature Index
Layer:
Feature Index
256
320
384
448
128
192
256
320
448
Token Position
10
10
12
14
112
10
118
120
10
124
126
10-2
Layer:
Layer:
Feature Index
Feature Index
64
128
192
256
320
384
448
64
128
192
256
320
384
448
Taken Position
10
102
112
114
120
100
122
Layer:
Layer:
256
10
Token Position
10
10a
12
10
114
116
10
122.
124
10
Figure 27: Example hidden state plots for GPT2-Smas model.
37

---

Published as a conference paper at ICLR 2025
GPT2-Medium Hidden States
Input Sequence:
Input Sequence: 2
Overall Mean Hidden States
Overall Mean Hidden States
Feature Index
Feature Index
120
192
256
320
384
448
128
192
256
320
304
448
10a
Taken Position
14
112
10
118
120
10
124
126
10
Layer:
Layer:
Feature Index
Feature Index
64
128
192
256
320
384
448
64
128
192
256
320
384
448
Token Position
10
10
12
114
112
10
120
10
122
10
Lay Feature Index
Layer: Feature Index
256
320
320
Taken Position
12
10
102
10
116
122
124
126
10-4
Layer: 18
Layer: 18 Feature Index
Feature Index
120
192
256
320
384
448
128
192
256
320
304
448
103
Token Position
10
14
112
118
120
10
124
126
10
Figure 28: Example hidden state plots for GPT2-Medium model.
38

---

Published as a conference paper at ICLR 2025
GPT2-Large Hidden States
Input Sequence:
Input Sequence: 2
Overall Mean Hidden States
Overall Mean Hidden States
Feature Index
Feature Index
120
192
256
320
384
448
128
192
256
320
304
448
10a
Taken Position
14
10
112
118
120
10-1
124
126
10
Layer:
Layer:
Feature Index
Feature Index
64
128
192
256
320
384
448
64
128
192
256
320
384
448
Token Position
10
10
10
12
112
10
114
120
10
122
10
Lay La: Feature Index
Layer: Feature Index 320
256 320
Token Position
10
10
12:
14.
10
114
116
10
122
124
126
10-1
Layer: 27
Layer: 27 Feature Index
Feature Index
120
192
256
320
384
448
128
192
256
320
304
448
103
Token Position
14
112
10
118
120
10a
124
126
10
Figure 29: Example hidden state plots for a GPT2-Large model.
39

---

Published as a conference paper at ICLR 2025
GPT2-XL Hidden States
Input Sequence: 1 Overall Mean Hidden States
Input Sequence: 2
Overall Mean Hidden States
Feature Index
Feature Index
1024
1088
1152
1216
1280 1344
1408
1472
1024
1080
1152
1215 1280 1344
1408
1472
10a
102
Taken Position
10
14
10
112
118
120
10-1
124
126
10
Layer: 0 Feature Index
Layer: 0
Feature Index
1024
1088
1152
1216 1280 1344
1408
1472
1024
1088
1152
1236
1280
1344
1408
1472
Token Position
10²
10
10
12
112
114
120
10-1
122
10
Layer: 12
Layer: 12
Feature Index
Feature Index
1216 1280 1344
1024
1216 1280 1344
10a
laken Position
10
12
14.
10
114
116
122
124
126
10-4
Layer: 36
Lay 36
Feature Index
Feature index 1280
1024
1088
1152
1216 1280 1344
1408
1472
1024
1080
1152
1216
1344
1408
1472
10a
Token Position
10a
14
112
10
118
120
10
124
126
10
Figure 30: Example hidden state plots for a GPT2-XL model.
40

---

Published as a conference paper at ICLR 2025
Llama2-7B Hidden States
Input Sequence: 1
Input Sequence: 2
Overall Mean Hidden States
Overall Mean Hidden States
Feature Index
Feature Index
2048
2112
2176
2240
2304 2368
2432
2496
2048
2112
2176
2240 2304 2368
2432
2496
10a
Taken Position
10
14
10
112
10c
118
120
10-1
124
126
10
Layer: 0
Layer: 0
Feature Index
Feature Index
2048
2112
2176
2240 2304 2368
2432
2496
2048
2112
2176
22:40
2304
2368
2432
2496
Token Position
10
10
10
12
112
114
120
10
122
10
Layer: 8
Layer: 8 Feature Index
Feature Index
2240 2304 2368
2240 2304
107
Taken Position
10
12
14.
10
114
116
10
122
124
126
10
Layer: 24
Layer: 24
Feature Index
Feature index 2240
2048
2112
2176
2240 2304 2368
2432
2496
2048
2112 2176
2304
2368
2432
2496
10a
Token Position
14
112
10
118
120
10a
124
126
10
Figure 31: Example hidden state plots for Llama2-7B model.
41

---

Published as a conference paper at ICLR 2025
Llama3.1-88 Hidden States
Input Sequence: 1 Overall Mean Hidden States
Input Sequence: 2
Overall Mean Hidden States
Feature Index
Feature Index 2240 2304
2048
2112
2176
2240 2304 2368
2432
2496
2048
2112
2176
2368
2432
2496
10a
Taken Position
10
14
112
10
118
120
10
124
126
10 2
Layer: O
Layer: 0
Feature Index
Feature Index
2048
2112
2176
2240 2304
2368
2432
2496
2048
2112
2176
22:40
2304
2368
2432
2496
Token Position
10
12
10
112
114
10
120
122
10
Layer: 8
Layer: 8 Feature Index
Feature Index
2240 2304 2368
2240 2304
102
Taken Position
12:
10
116
122
124
126
10
Layer: 24
Layer: 24
Feature Index
Feature index 2240
2048
2112
2176
2240
2304 2368
2432
2496
2048
2112 2176
2304
2368
2432
2496
102
Token Position
10
14
112
118
10c
120
124
126
10
Figure 32: Example hidden state plots for Llama3.1-88 1-BB model.
42

---

Published as a conference paper at ICLR 2025
DeepSeek2-Lite Hidden States
Input Sequence: 1
Input Sequence: 2
Overall Mean Hidden States
Overall Mean Hidden States
Feature Index
Feature Index 1728 1792 1056
1536
1600
1664
1720
1792
1856
1920
1984
1536
1600
1664
1920
1984
10
102
Taken Position
10
14
112
10c
118
120
124
126
10 2
Layer: 0 Feature Index
Layer: 0
Feature Index
1536
1600
1664
1728 1792 1856
1920
1984
1536
1600
1664
1728
1792
1856
1920
1984
10a
Token Position
10
10
12
10
112
114
120
10
122
10
Layer: 7
Layer: 7
Feature Index
Feature Index 1792
1664
1728 1792 1856
1664
1856
10a
Taken Position
10
10
12
116
10
122
124
126
10
Layer: 21
Layer: 21
Feature Index
Feature index 1728
1536
1600
1664
1720 1792 1856
1920
1984
1536
1600
1664
1792
1056
1920
1984
10a
Token Position
14
10
112
118
10
120
124
126
10
Figure 33: Example hidden state plots for a DeepSeekv2-Lice model.
43

---

Published as a conference paper at ICLR 2025
K ATTENTION MAPS OF PRETRAINED MODELS
In this section, we present the attention maps of popular pretrained models. This shows how models establish attention patterns and how they persist after initial layers. This shows that generally after the first or second layer, first token attention dominance is highly established and persists across layers.
We show the mean across layers, the first layer, 1 and of the layers over all heads in each case.
GPT2 Attention Maps
Imputt 3 Kay hagiian
Input Sequence: Map Key Fog ton
1.0
08
24
Query
0.0
Layer
hugiian
Key Pog
32
Quary
Layer 3
Layer
Pogion
Pigtion
0.0
Ouary hatton
Layer 9
Layer
0.4
Ouryy Nottin
Figure 34: Example attention maps for GPP22-Smal model.
44

---

Published as a conference paper at ICLR 2025
GPTZ-Medium Attention Maps
Input Sequence: 1 Overall Mean Actention Map Key Pasition
Input Sequence:
Overall Mean Attention Map Key Position
96
0.8
0.6
Query Position
02
Laver Key Pasition
Key Position
100
0.6
0.4
Query Positoon
Laver
Laver
Pasition
Key Pocttion
10
0.4
Query Positoo
0.0
Laver 18
Laver 18
Key Papitio
Key Posttio
32
96
10
0.8
32
0.4
Query Prsition
0.2
0.0
Figure 35: Example attention maps for GPT2-Medium model.
45

---

Published as a conference paper at ICLR 2025
GPTZ-Large Attention Maps
Input Sequence: 1 Overall Mean Actention Map Key Pasition
Input Sequence:
Overall Mean Attention Map Key Position
96
0.6
Query Position
02
Laver Key Pasition
Key Position
1.0
0.6
0.4
Query Positton
Laver
Laver
Pagition
Pocttion
10
0.4
Query Rosition
0.0
Laver 27 Key Pacison
Laver 27
Key Position
96
10
0.8
32
Query Prsitton
0.2
0.0
Figure 36: Example attention maps for GPT2-Large model.
46

---

Published as a conference paper at ICLR 2025
GPT2-XL Attention Maps
Input Sequence: Overall Mean Actention Map Key Pasition
Input Sequence:
Overall Mean Attention Map Key Position
96
0.6
Query Position
02
Laver Key Pasition
Key Position
1.0
0.6
0.4
Query Positoon
Layer
Layer
PagiBon
Posttion
10
0.4
Query Positon
0.0
Laver 36
Laver 36
Key Pacition
Key Posttion
32
96
10
0.8
32
0.4
Query Prsition
0.2
0.0
Figure 37: Example attention maps for GPT2-XL. model.
47

---

Published as a conference paper at ICLR 2025
Llama2-7B Attention Maps
Input Sequence: 1 Overall Mean Actention Map Key Pasition
Input Sequence:
Overall Mean Attention Map Key Position
96
0.6
Query Position
02
36
Laver Key Pasition
Key Position
1.0
0.6
0.4
Query Positton
Laver&
Laver
Pagition
Pocttion
96
10
0.4
Query Positoon
0.0
Laver 24 Key Pacison
Laver 24
Key Posttion
96
10
0.8
32
0.4
Query Position
0.2
0.0
Figure 38: Example attention maps for a Llama2-7B model.
48

---

Published as a conference paper at ICLR 2025
Llama3.1-88 Attention Maps
Input Sequence: 1 Overall Mean Actention Map Key Pasition
Input Sequence:
Overall Mean Attention Map Key Position
96
0.8
0.6
Query Position
02
36
Laver Key Pasition
Key Position
1.0
0.6
0.4
Query Positoon
Laver&
Laver
Pasition
Pocttion
96
10
0.4
Query Posiion
44
0.0
Laver 24 Key Papison
Laver 24
Key Posttion
32
96
10
0.8
32
0.4
Query Prsition
0.2
0.0
Figure 39: Example attention maps for Llama3. 1-8B model.
49

---

Published as a conference paper at ICLR 2025
DeepSeek2-Lite Attention Maps
Input Sequence: 1 Overall Mean Actention Map Key Pasition
Input Sequence:
Overall Mean Attention Map Key Positin
96
0.6
query Position
Key Pasition
Key Position
1.0
0.6
Query Position
Laver
Laver
Pasition
Pocttion
10
0.4
Query Rosition
0.0
Layer 20
Laver 20
Pagiton
Key Position
32
96
0.8
Query Prsition
0.2
0.0
Figure 40: Example attention maps for DeepSeekv2-Liit model.
50

---

Published as a conference paper at ICLR 2025
L TRAINING CURVES
To demonstrate that our proposed methods, L.e. replacing the canonical softmax function with softmax-/ and using our proposed optimiser, OrthoAdam, do not negatively impact the training of large language models, we provide the training curves for our models here. One can observe that training curves for models using either or both our proposed changes are stable and converge to similar loss value the baseline models.
GPT2-60M
gpt2:
gpt2: 60m, 3.56
gpt2: 60m, QA Final Loss: 3.59
gpt2: 60m, S1, OA Final Loss: 3.57
Training Loss
20000
40000
60000
anbco
100000
1200n0
100000
180000
Training Step
Figure 41: Training for GPT2-60M models with different optimisers and softmax functions. The models using OrthoAdam and softmax-1 are stable and converge to similar loss value as baseline models. S1/OA denotes the model using softmax-1 and/or OrthoAdam.
L.2 GPT2-130M
9912: 130m Final Loss: 3.20
gpt2: 130m, S1 Final Loss: 3.19
gpt2: 130m, OA Final Loss: 3.21
gpt2: 130m, S1, OA Final Loss: 3.19
Training Loss
3.1
53020
150000
250000
Training Step
Figure 42: Training curves for GPT2-130M models with different optimisers and softmax functions. The models using OrthoAdam and softmax-1 stable and converge to similar loss value the baseline models. S1/OA denotes model using softmax-1 and/or OrthoAdam.
51

---

Published as a conference paper at ICLR 2025
L.3
GPT2-350M AND GPT2-1.4B
gpt2: 350m Final Loss: 2.86
gpt2: 350m,
Loss: 2.86
gpt2: Final Loss: 2.64
gpt2: 1b, OA Final Loss: 2.63
Training Loss
2.8
2.6
102000
200090
300300
400000
503000
503000
Training Step
Figure 43: Training curves GPT2-350M and GPT2-1.4B models with different optimisers and softmax functions. The models using OrthoAdan and softmax- / are stable and converge to a similar loss value as the baseline models. S1/OA denotes the model using softmax-1 and/or OrthoAdam.
L.4 LLAMA-130M
llama: 130m Final Loss: 2.84 130m,
llama: Final
liama: 130m, S1, OA Final Loss: 2.83
Training Loss
3.0
2.9
2.8
50690
100000
L50000
200000
250000
300000
Training Step
Figure 44: Training curves for Llama-130M models with different optimisers and softmax functions. The models using OrthoAdam and safimax-1 are stable and converge to similar loss value as baseline models. S1/OA denotes the model using softmax-1 and/or OrthoAdam.
52


Reviews:
[{"review_id": "h9JchJP2gr", "reviewer": "Reviewer_SMxg", "summary": "This paper studies two phenomena in transformer-based models. The first is the strong dominance of the first token in the attention maps, and the second is the presence of outlier activations in the hidden states. To address these two issues, the paper proposes a new softmax philosophy and a new optimizer. The experiments show the advantages of the proposed methods under the quantization scenario.", "questions": "1. Why do the authors train the LLaMA2-130M model? It is strange to me.\n\n2. Can the authors explain the how the case would change where the model is used to generate following some instruction. Will the LM attend strongly to the first generated token?\n\n\n---\nAfter rebuttal\nThanks for the authors' detailed responses. My concerns are below.\n\n1. My first concern is the **originality of the softmax-1**. However, I failed to find the original paper when I first read the paper until I see the comments.\n\n2. The second concern is the **contribution of the paper**. In addition, the first token attention issue was also discussed in [1]. I am afraid the story of the paper on 1 and 2 is overstated in the submission. That is why I give a lower contribution score.\n\n3. The authors state \"large language models\" in their submission title, which is confusing. While the authors provide the LLaMA-8B results in the rebuttal, only a 8B LLM can not represent the entire group of LLMs. The situation can be very different when the model size comes to 70B or even 130B. A more suitable statement for the current submission would be **\"autoregressive language models\"**.\n\nTherefore, I decide to keep my score.\n\n[1] Efficient Streaming Language Models with Attention Sinks", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "1. The paper conducts a series of experiments to verify the phenomena, which are convincing.\n\n2. The proposed new softmax and optimizer are simple and effective.\n\n3. The experiments under the quantization scenario show the usefulness of the study, which potentially enlightens future research.", "weaknesses": "1. The paper states that the study is on popular large language models (LLMs). However, the experiments in the paper exclude large language models, considering that the largest-sized one in the experiments is GPT2-1.4B, not reaching the bar of a regular sized LLM (>7B). It makes the results of the paper less convincing.\n\n2. I am not convinced that the highlighted two issues in the paper are critical for recent LMs. From the experiments in the paper, one can only see the promising usage under 4-bit/8-bit quant; in other scenarios, still not clear.\n\n3. The paper reports PPL and other metrics to access the LM performances. However, it is not clear whether the resultant models perform well on downstream tasks and instruction following, which are important aspects to evaluate an LM/LLM. Therefore, it is not positive that the study will greatly contribute to the community.\n\n4. Keep up with the last point, the paper does not consider (or explain) whether the study still holds under the instruction-following scenario. This situation is different from doing language modeling; or it is still the same, since the first generated token is not the start of the input sentence.\n\n5. Minorly, it is also not mentioned in the paper the relationship of the two phenomena. It seems that they belong to two independent research aspects.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "Bqmd9KTmH1", "reviewer": "Reviewer_rxLH", "summary": "This work studies two unexpected phenomena that occur in Transformers-based language models: the dominance of first tokens in the self-attention maps, and the presence of outlier dimensions throughout the hidden representations of these models. The authors propose a different mitigation technique for each phenomenon, and successfully mitigate both first-token dominance and outlier dimensions. They finally show that combining both mitigation techniques yields language models that can be quantized easily, with minimal performance loss.\n\nThe first mitigation technique relies on a softmax-1 activation function, that adds a bias to the denominator of the softmax function to allow for null rows in the attention map. The authors conduct experiments that show that this technique solves the first-token dominance problem but not outlier dimensions. This second problem is mitigated by OrthoAdam, a slight modification to the Adam optimizer that compute gradients in a given orthonormal basis to avoid specific canonical dimensions being used for massive coefficients.", "questions": "- The idea of using an orthogonal basis $Q$ in Adam is very elegant, especially in the context of quantization. Nevertheless, from a modeling/theoretical viewpoint, it would be interesting to check whether the models are still using outlier \"dimensions\" in this orthogonal basis, that is if the projection of their activations in such an orthogonal basis still have outlier coefficients. Did you conduct such experiments? In other words, did you verify that the models are not still relying on outlier \"dimensions\" in a different basis than the canonical one?\n- It is common to decrease the learning rates for larger models. Why did you choose a single learning rate for the different model sizes? \n\n**Typos / Remarks**\n\nL419 : this sentence (*Note that only linear layers...*) seems a bit unformal.\n\nL459: *but still remain high* -> *which remains high* ?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 4, "strengths": "This paper successfully addresses two issues that have been identified by extensive literature in the language modeling field when training Transformers-based models. It provides relevant methods to mitigate these issues and show that such methods can help with performance \n (slightly) and significantly eases the quantization process.\nThe paper is well-written and pedagogical. The proposed solutions are elegant and simple to implement.\nOverall, this paper paves the way for a better understanding of observed self-attention maps and for more expressive inductive biases for language models. It also provides immediate and substantial benefits for the field of quantization.", "weaknesses": "The main weakness of the paper is its failure to support its main theoretical claim. From the abstract on, the authors argue that they identify \"Adam as the primary contributor to the large outlier activations\", but they fail to properly discuss the theoretical background of this claim. They support this claim with the following arguments, which all have flaws:\n- *the optimiser tracks moments in the same basis as the model parameters* : this reason does not explain why the optimiser is pushing for these outlier activations in gradients and their moments, but only why these outlier dimensions are passed on to the model's weights;\n- *Adam and RMSProp have high[er] kurtosis [than SGD]* : this indeed shows that Adam and RMSProp are particularly accelerating the emergence of outlier activations, but it should be noted that they are also accelerating convergence. Hence, this could simply imply that Adam and RMSProp are efficiently facilitating the emergence of outlier activations that could still **be caused by another element of the system** that would rely on outlier activations to increase performance. In other words, the \"outlier activation\" state could be a favorable state for the model in terms of performance because of another component, and this state could be reached more quickly/easily using Adam and RMSProp.\n\nIt could be argued that the fact that OrthoAdam mitigates the outlier dimension problem proves that Adam is the culprit for such a phenomenon. Nevertheless, OrthoAdam is almost explicitly designed to avoid outlier dimensions *in the canonical basis*, as discussed in section 4.2, and could be seen as a variation of Adam that enforces the absence of outlier dimensions. Hence, this just shows that it is possible to train a model under this constraint without substantial performance decrease, but it does not show that Adam is the key **causal** element in the emergence of outlier dimensions.\n\nAnother limitation of the paper lies in some parameters and models that were used to validate the approaches:\n- The models were trained using a sequence length of 256, which questions the scalability of softmax-1 to longer sequences, as sequence length may be a crucial factor for this activation function;\n- All resulting models seem to have a noticeably high perplexity on the validation set. As a comparison, the 70M Pythia model reaches approximately a 3.2 loss (=24 perplexity) at half-training (~150B tokens), and the 410M version reaches 2.2 for cross-entropy (=9 perplexity)(https://wandb.ai/eleutherai/pythia-extra-seeds/reports/Some-loss-curves-for-smaller-Pythia-models--Vmlldzo2NTkxNDIw). This is significantly smaller than what is described in the paper for equivalent models, which raises questions about implementation, hyperparameter choices and evaluation.\n\nFinally, the authors do not include a discussion on training dynamics, showing how the models converge during training, and they do not provide an analysis of OrthoAdam in terms of memory requirements and training latency, which may question the scalability and technical relevance of the method. My intuition is that the cost should be tolerable, but I think this should be mentioned in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "LXgyfaBRaW", "reviewer": "Reviewer_hVNT", "summary": "This paper studies two interesting phenomena in transformer-based LLMs. First, the large amount of attention the first token receives. Second, the emergence of outlier neurons in the transformer layer activation. They attribute the first phenomenon to the softmax function, and the second to the Adam optimizer. They propose variants of these two methods (softmax-1 and OrthoAdam), which resolve these two issues. They show that using these methods allows models to quantize far better than without them.", "questions": "- As the authors note, Llama2 remains usable after quantization even in the vanilla version. Do the authors have some intuition regarding this?\n\n- I had a hard time understanding the discussion in #222 around causal masking. What does it mean to relax causal masking? train with an MLM loss? and why would the fifth token dominate in this case?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Interesting and important research questions\n- A well executed study. In particular, several hypotheses are considered for the phenomena in question, and the authors are able to pin-down on what seems to be their root cause \n- Very interesting findings\n- Practical results that allow simpler quantization", "weaknesses": "- softmax-1 has been introduced before in the streamingLLM paper (https://arxiv.org/abs/2309.17453, referred to as $Softmax_1$ or Zero-sink in that work). The authors in that work showed that \"_while the zero sink alleviates the attention sink problem to some extent, the model still relies on other initial tokens as attention sinks_\". This work seems to argue differently. It seems important to pin down the differences between both setups that lead to the different behaviors.\n\nIn addition, I found some of the major claims to require further discussion.\n\n- From table 4, it seems RMSNorm-s is also important for mitigating the outlier effect. This is currently not really \ndiscussed, and the authors only mention softamx-1 and OrthoAdam as the important parts.\n\n- The softmax-1 method requires some discussion and intuition. Why is the specific method appropriate here? Specifically, the authors say (#208) that it allows \"having low attention scores everywhere.\" But standard softmax also allows that, e.g., by assigning uniform attention scores to all tokens. Is it about a better inductive bias? It would be helpful to see a more detailed comparison between the two methods. E.g., an explanation of why softmax-1 is more effective than standard softmax for this task, and/or a discussion of any potential trade-offs or downsides to using softmax-1.\n\n- If I understand the discussion around OrthoAdam correctly, the authors say the reason for the outliers is (a) the appearance of high values of features, and (b) that optimizers like Adam lead to such high values. Is my understanding correct? If so, this should be stated more explicitly.\n\n\n\nMinor:\n\n- Using *x*s to contrast with *v*s in table 2 (as used in tables 3 and 4) would make it clearer than leaving it blank.\n\nAlso, there are several typos and such in the paper:\n\n- #195: \"... therefore it *in* receives ...\" (drop \"in\")\n- #306: \"... under *an* particular ...\" (should be \"a\")\n- #313: \" the model *during* i.e.  ...\" (something's strange here, either drop during or add something after it)\n- #420: \"Note that only linear layers are *quantised the* embeddings, ...\" (Missing period/, while)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Poster)