nGPT: Normalized Transformer with Representation Learning on the Hypersphere
Authors: Ilya Loshchilov
Cheng-Ping Hsieh
Simeng Sun
Boris Ginsburg
Venue: ICLR 2025 Poster (2025)
Keywords: transformer
normalization
representation
Primary area: foundation or frontier models, including LLMs
TL;DR: A novel architecture, where all vectors forming the embeddings, MLP, attention matrices and hidden states are normalized to reside on the same hypersphere.

Abstract: We propose a novel neural network architecture, the normalized Transformer (nGPT) with representation learning on the hypersphere. In nGPT, all vectors forming the embeddings, MLP, attention matrices and hidden states are unit norm normalized. The input stream of tokens travels on the surface of a hypersphere, with each layer contributing a displacement towards the target output predictions. These displacements are defined by the MLP and attention blocks, whose vector components also reside on the same hypersphere. Experiments show that nGPT learns much faster, reducing the number of training steps required to achieve the same accuracy by a factor of 4 to 20, depending on the sequence length.

Published as a conference paper at ICLR 2025
NGPT: NORMALIZED TRANSFORMER WITH REPRE- SENTATION LEARNING ON THE HYPERSPHERE
Iya Loshchilov, Cheng-Ping Hsieh, Simeng Sun & Boris Ginsburg
NVIDIA
loshchhill lov, chsieh, simengs, bgi (nginbburg)@nvidi
ABSTRACT
We propose a novel neural network architecture, the normalized Transformer (nGPT) with representation learning on the hypersphere. In nGPT, all vectors forming the embeddings, MLP, attention matrices and hidden states unit norm normalized. The input stream tokens travels on the surface a hypersphere, with each layer contributing displacement towards the target output predictions. These displacements are defined by the MLP and attention blocks, whose vector components also reside on the same hypersphere. Experiments show that nGPT learns much faster, reducing the number of training steps required to achieve the same accuracy by factor of 4 to 20, depending on the sequence length.
INTRODUCTION
The Transformer architecture (Vaswani et al., (017) is the foundation for most of modern language models. An enormous number of modifications to this architecture have been proposed to improve training stability, inference costs, context length, robustness, etc. It has been noted that the ap- plication of various normalization techniques is beneficial (Salimans & Kingma (016), leading to experiments with adding normalization layers such LayerNorm and RMSNNrm in nearly every possible position within the network (Xiong et al., 2020). Another approach to the model nor- malization is through controlling the norm of weights using weight decay (Loshchilov & Hutter) 2019). Recent studies (Andriushchenko et al.] 2023) suggest reevaluating the role of weight decay and taking a closer look rotations rather than focusing solely on vector norms (Kodryan al.) 2022 Kosson et al.] 2023). Franke et al. (2023) suggested to enforce an upper bound on L2 norm of parameter groups There growing evidence that representation learning on the hypersphere is associated with more stable training, greater embedding space separability, and better performance on downstream tasks (Wang Isola (2020). Recent studies also suggest that transformers implicitly perform gradient descent as meta-optimizers (Von Oswald et al. 2023 Dai et al. (022).
We propose to unify (see also Appendix A.3) various findings and observations made in the field under new perspective of the normalized Transformer. Our key contributions are as follows:
Optimization of network parameters the hypersphere We propose to normalize all vectors forming the embedding dimensions of network matrices to lie on a unit norm hypersphere. This allows view matrix-vector multiplications dot products representing cosine similarities bounded in [-1,1]. The normalization renders weight decay unnecessary.
Normalized Transformer as a variable-metric optimizer on the hypersphere The normalized Transformer itself performs multi-step optimization (two steps per layer) on a hyper- sphere, where each step the attention and MLP updates is controlled by eigen learning rates-the diagonal elements of a learnable variable-metric matrix. For each token (; in the input sequence, the optimization path of the normalized Transformer begins at a point on the hypersphere corresponding to its input embedding vector and moves to point the hypersphere that best predicts the embedding vector of the next token f,+1-
Faster convergence We demonstrate that the normalized Transformer (nGPT) reduces the number of training steps required to achieve the same accuracy by a factor of 4 to 20.

---

Published as a conference paper at ICLR 2025
2 EVOLUTION OF THE TRANSFORMER: FROM GPT TO NGPT
This section outlines the baseline Transformer and the modifications necessary derive its nor- malized version. We illustrate these changes for Transformer decoder with self-attention only. The extension to encoder-decoder and cess-attentionn is straightforward. summary of these changes given in Table with details in Section 2.6
2.1 TOKEN EMBEDDINGS OUTPUT LOGITS
The decoder-only Transformer trained to predict token t, using previous tokens input sequence = (f1, f2.,t....... f For each input token &.. we retrieve its corresponding embedding in Rd from learnable embedding matrix Eimpur E R VVxuuss with vocabulary of size V. Similarly, the target output sequence y is represented using a learnable embedding matrix Euutat E R Vxduas Notably, since both Emput and Eoupur are learnable (unless they are tied to be equivalent), any token i, can have different embeddings in the input and output sequences. measure token similarity during model training, the dot product between the corresponding embedding vectors is used. How- ever, the norms of embedding vectors in the original Transformer are unconstrained, which can lead to inaccurate similarity estimation. To improve the accuracy of similarity estimation, propose to normalize the embedding vectors stored in Emput and Eouppu after each step of the training algorithm.
The next token prediction is enforced by causal masking (see Section 2.3 to ensure that future tokens are considered. This allows the model to compute the prediction error for all T tokens in parallel during training, while preserving the autoregressive nature of the task. After the Transformer processes the sequence (æ)...., 1;-1). ) it produces an output vector h; E Rddsss for each i-th position in the predicted sequence. The logits z; R representing the unnormalized probabilities for each token in the vocabulary, are computed using the output embedding matrix Eoutun:
z; Eoupp
The logits z; are passed through softmax function to convert them into probabilities:
exp(Zi.j.
Pl xi-1)
Ev=1 exp(2i.))
Here, Zi.z, is the logit corresponding to the correct token and the denominator normalizes the log- its into probability distribution over the vocabulary. During inference, the prediction Ú; obtained by selecting the token with the highest probability. Since all nGPT embeddings are normalized, the logits E V in equation represent dot products bounded the range [-,, This limits the confidence (temperature) of the probability distribution generated by the softmax in equation [ To adjust this during training, and line with Hoffer et al. (2018), we introduce a trainable scaling parameter s, R that scales the logits element-wiss
f- 28:
(3)
Table 1: Transformer vs. Normalized Transformer.
Transformer
Normalized Transformer
<- ATTN(RMSNOm((h)
hA + Norm(ATTN(h))
h hA
h Norm(h outha h))
hm MLP(RMSNomm(())
hm Norm(MLP(h))
h < h + hM
h Norm(h cxw h))
Final: h + R.MNNrr(h)
After each batch pass, matrices and embeddings are normalized along their embedding dimension. The hidden state updates are controlled by learnable vectors of eigen learning rates and av.
All parameters of matrices and embeddings unconstrained.
2

---

Published as a conference paper at ICLR 2025
2.2
LAYERS AND BLOCKS
2.2. 1
BASELINE TRANSFORMER
L layers of transformations are applied to the hidden state h, consisting of alternating the self- attention (ATTN) and multi-layer perceptron (MLP) blocks:
h h + ATTN(RMSNOm((A) h + h MLP(RMSSNrr((h))
(4)
(5)
where RMSNorm(h) is one of several possible normalizations. embedding to norm of / /dmodel. then scales each dimension by lcarnable vector of dmodel factors, typically initialized to Since the transformation block outputs are added h, the token embedding norms can vary significantly. To address this, normalization is also applied after the final layer.
used first to normalize each
2.2.2 NORMALIZED TRANSFORMER
For any points a and b on the surface of hypersphere in Rdmm, SLERP (Spherical Linear Interpo- lation) by Shoemake (1985) computes an interpolation along the geodesic (shortest path):
SLERP(a, (a, a; sin =a)0)
sin(at),
sin(O) b
(6)
sin
where 0 = arccos(a b) is the angle between the points a and b, and a E [0, 1] is the interpolation parameter, with a returning and = returning b. Our experiments suggest that SLERP can be approximated by simple linear interpolation (LERP):
LERP(a, b; w) = (1 - a)a ab
Let us rewrite this equation as an update equation in nGPT:
a aab-ab
(8)
where a is h, and, b the point suggested by the attention or MLP block. Then, for the gradient g b, more general form involving variable matrix B Rddaaa Xdnand becomes:
- aBg
In quasi-Newton methods, B approximates the inverse Hessian matrix H When B is diagonal with non-negative elements, aB becomes vector R whose elements correspond to the diagonal of B times the learning rate We denote word eigen, meaning "own," referring to the internal structure of the Transformer). We provide some notes in Appendix A.2 Following equation & the update equations for the attention and MILP blocks are as follows:
eigen learning rates (from the German
h Norm + aha - h))
(10)
h - Norm h + +Mm(hM h))
(11)
where as E RR and aM E R-dase are learnable parameters applied to the normalized outputs 20
of the attention and MLP blocks h = Norm(ATTN(h)) and hv = Norm(MLP(h)), respectively. The function Norm(x) normalizes any vector to have unit norm, and, unlike RMSNorn or Lay- erNorm, does not introduce any clement-wise scaling factors. The normalization can be viewed as the retraction step in Riemannian optimization, mapping the updated solution back to the manifold. Appendix A.4 discusses the extension of our update equations the context Riemannian optimiza- tion. In contrast the baseline Transformer, additional normalization is required after the final layer, as the embeddings are already normalized by the proposed scheme.
2.3
SELF-ATTENTION BLOCK
2.3.1 BASELINE TRANSFORMER
The attention mechanism is a key component of the Transformer. It allows each token to attend every other token in the sequence, enabling the model capture long-range dependencies. The
3

---

Published as a conference paper at ICLR 2025
block typically starts with normalization of the input hidden state h using RMSNorm deal with fluctuating norms of embeddings. Then, the normalized h is projected into three separate vectors the query q, the key k, and the value v:
q - h. g. ke t WWW, ky U + hW.
(12)
where W, Wi, W. E Rdddnoed Xd, are learned projection matrices, and de the dimensionality of the query/key vectors. To incorporate positional information, we apply Rotary Position Embeddings (RoPE) by [Su et al. (2024) to both the query and key vectors. The attention scores are computed by taking the dot product of the query and key vectors, scaling them by - then applying softmax function to obtain attention weights, and finally computing a weighted sum of the value vectors w:
gk
Attention( q, k. v) t softmax
+ +M)
(13)
dd
where M is a matrix that prevents attending to future tokens by setting the corresponding entries to -oc. Specifically, M.,- 0 if j and M.,- = -0 jf > i.
In practice, Ytheads attention heads are used where for each i-th head, separate linear projections Wi. W W are applied, and the attention mechanism is computed independently for each head:
gr
hA + -oCoca(hccad,,... head W
(14)
where head; Attention q², k', v²) and Wo R x de dastel is learned projection matrix, where du is typically set to dmodel/ Theads.
2.3.2 NORMALIZED TRANSFORMER
The matrix- vector multiplication of Wa E Amaas of the i-th head and h E R can be viewed as a dot product between the columns of W, and h. In the baseline Transformer, all matrices, including W are unconstrained, leading to unbounded values in q. We propose to normalize W WE. W, and W. along their embedding dimension so that the computed dot products with h can be interpreted as cosine similarity between unit norm vectors bounded in [-1,1] 1]. Thus, all attention matrices can be viewed as collections of normalized embedding vectors to be compared with.
While each element of q and & is now bounded, the norms of these two vectors can still vary. Moreover, injection of positional information by RoPE further distorts q and k. We propose to additionally normalize q and ks ensuring that the dot product of every query and key is under control:
q Norm 8gk
(15)
k +- Norm (k) (k)8
(16)
where 8gk E Rd is a vector of trainable scaling factors for the i-th head.
In the original Transformer, the softmax scaling factor 1/ /dvd; in equation 13 is introduced to account for the expected variance of de in the dot product of non-normalized query and key vectors. In the normalized Transformer, the expected variance of the dot product between normalized query and key vectors is 1/dg. To restore a variance of 1, the softmax scaling factor should instead be /dg. If the softmax scaling factor is set to 1, this is equivalent to initializing the scaling factors 8yk at d% 1/4
2.4
MLP BLOCK
2.4.1
BASELINE TRANSFORMER
The input hidden state h of the MLP block is first normalized using RMSNorm and then passed through two separate linear projections, producing two intermediate vectors (we omit bias terms);
u f h.W. v + hW
(17)
1 the following equations are defined per head hut we omit i-th head index for the sake readability. "hheee is need for separate scaling factors for q and ls the scaling would simply be applied element- wise when computing the dot product between q and k.
4

---

Published as a conference paper at ICLR 2025
where W., W do d are the learned weight matrices. The intermediate vectors and are combined using a gated activation function called SwiGLU defined by Shazeer (2020) as:
SwiGLU(u, < u SiLU(p)
(18)
where Si)U(w) = V. olv), and (Vv) is the sigmoid function. The result of the gated activation is then passed through final linear transformation WoMLP E Rdve Xdaai
hy - SwiGLU(u, v) D) WOMLP
(19)
2.4.2 NORMALIZED TRANSFORMER
We propose to normalize matrices Wa and W along the embedding dimension that the and V vectors represent the cosine similarity between h and vectors stored in W,, and W., respectively. To control their impact, introduce scaling factors 8n E RRdu and 8y Rdbr.
U f Us
(20)
V8 dmoneely
(21)
where the rescaling of by dimodet is needed to benefit from the non-linearity of SiLU (see the Appendix A.I. The output of the MLP block is invariant to rescaling of u by a scalar.
2.5
EFFECTIVE LEARNING RATES IN ADAM
The core of the Adam algorithm by Kingma 2014) as follows:
m < m + (1 - Bi)g
< 320 (1 32)g²
0 am/ (V +e),
(22)
where is the parameter vector, g is the batch gradient, m is the momentum, v is the estimate of the per-element gradient amplitudes, is the scheduled learning rate, is small constant, and 31 < 82 are momentum factors close to 1. We cite the text of the original Adam paper using our notation: more common scenarios, we will have that m Te
since F(g)/ LET2 1. The effective magnitude the steps taken in parameter space each timestep approximately bounded by the stepsize setting a. Thus, a controls the effective step-size in the search space, while the ratio in can temporarily increase (respectively, decrease) the step-size if the current amplitude of per- Ve
parameter momentum greater (respectively, smaller) than its estimated value over longer time horizon. Consider an example where 0; 0.01 and the global learning rate is 0.001. If the gradient amplitude remains stable (i.e., my Rs 1), would take 0.02-0.01
10 steps to double 0,. However,
yor
0.001
0; = 1.0, it would take 2.0-1.0 in the second case, the number of steps would only decrease if m, Vo.
0,001 1000 steps to double. Even if the gradient's amplitude is larger
In nGPT, for any trainable vector of scaling parameters such as 8a, we use two scalars Sa,init and sa.scale- When initializing s, as trainable parameter, its initial value is set to Sa,scale- However, during the forward pass we restore its actual value by multiplying Sa,init/ / sc,eeee- This allows us to control the effective learning rate for 8, by adjusting 80,scates while keeping the global learning rate unchanged. For example, setting 8a,init = and 8a,scatc = 1/Vd ensures that this parameter is updated with the same effective learning rate other normalized parameters in the network.
2.6 SUMMARY OF MODIFICATIONS
The recipe to convert the baseline Transformer into the normalized Transformer is as follows:
1. Remove all normalization layers such as RMSNorm or LayerNorm.
2. After each training step (and, optionally, during the forward pass), normalize matrices Emput Eouiput, W. Wt. W.. W.. W.. W. and W along their embedding dimension.
3. Replace the update equations and by equations [0 II where a, (and also am) is treated with init 0.05 (in order of //Mayers) and @'Ascaate 1/Vdmadee-
5

---

Published as a conference paper at ICLR 2025
4. Change the softmax scaling factor in attention from 1/vd2 vdg. Implement the rescal- ing and normalization (normalization here is optional) of and k as in equations 15 and
16 where 8yk is treated with inki = 1 and Sak,soate 1/ dd
5. Implement the rescaling of the intermediate state of the MLP block using equations 20 and
where 8u (and also 8y) is treated with Su,init = and &u,soate =1
6. Implement the rescaling of logits using equation 3 where 8N is treated with 8z,init = 1 and
8 scule = 1mmdeeo
7. Remove weight decay and learning rate warmup.
3 EXPERIMENTS
We train both the baseline Transformer (GPT) and the normalized Transformer (nGPT) the Open- WebText dataset (Gokaslan & Cohen, 2019) and evaluate them on a set of standard downstream tasks. We experiment with models containing 0.5B and 1B parameters, including the embeddings. For both GPT and nGPT, we report results using the best initial learning rate settings (see Appendix A detailed description of the setup and hyperparameters is in Appendix A.6
A.7
3.1
ACCELERATION OF TRAINING
GPT with 2008 budgel
rOFT with 228 batgel
budget
Figure 1: Validation loss during training of 1B GPT and nGPT with 4k context length.
Figure I presents the validation loss during the training of GPT and nGPT models with billion pa- rameters and sample length of 4k tokens. After 20k iterations, nGPT achieves the same validation loss that GPT reaches only after 200k iterations (approximately 400 billion tokens), demonstrating a 10x speedup in terms of iterations and tokens used
1K context: OpenWebText
4K context: OpenWebText
BK context: OpenWebText
0.58
GPT
GPT
7.00
0.58
2.P00.58
13PT
03PT
- 13PT
Vilitation loos
Validatonn loss
Uaidtation
230
20c
3M
Training sohinns billions
Training
billions
Training takens billionss
Figure 2: Final validation loss (y-axis) for training runs with different computation budgets in tokens (x-axis). The training of 0.5B and IB nGPT models is about 4x, 10x and 20x faster (in terms of tokens) on 1k, 4k and 8k context lengths, respectively.
While the time per step of nGPT is higher (80% for 4k, and 60% for 8k context respectively). it can be reduced after code optimization. Also the overhead less significant for larger networks (see Appendix A.5)

---

Published as a conference paper at ICLR 2025
context: easy
4K context: hellaswag
context: winogrande
Valua
Value
Value
1.00
649
GPT 0.58
GPT 0.58
GPT 0.58
OPT 1B
GPT 18
GPT 1B
MGPT 0.58
MGPT 0.50
nGPT 0.58
nSPT
045
100
400
100
24
308
40
10)
20
301
46
Training tokens in bilions
Training tokens in bilions
Training tokens in billions
4K context: wsc273
4K context: lambada openal segien1024
4K context: average
0.50
0.99
0.60
2.0
0.52
Value
0:40
2
GPT 0.58 1B
GPT 0.5B
GPT 0.58
RGPT 0.58
n3PT 0.58
C40
nGPT 0.58
nGPT 1B
nGPT 18
nGPT 18
025
Training toleens
Training
Training billions
Figure 3: Models trained with 4k context length. Final performance (y-axis) on a set of downstream tasks and their average value (Bottom-Right for different computation budgets in tokens (x-axis).
Figure 2 illustrates how the performance gap between nGPT and GPT scales across three axes: total token budget, context length, and network size. Training the 0.5B and IB nGPT models is approximately 4x, 10x, and 20x faster context lengths of 1k, 4k, and 8k tokens, respectively.
Figure 3 shows a similar pattern across downstream tasks, confirming that the acceleration is not only reflected in perplexity but also in task performance. Figures 8 and I0 in the Appendix provide results for Ik and 8k context lengths. We observe some saturation for the longest runs of nGPT, suggesting that the model capacity is nearly reached for this number of trainable model parameters.
3.2 INSPECTION NETWORK PARAMETERS
Figure shows that, while nGPT maintains fixed norm for embeddings (by design), GPT hibits significant variation. The distribution of eigenvalues, computed from the covariance matrix of embeddings and normalized by their median, reveals that GPT's input embeddings have higher condition number, especially in the IB model. The distribution of pairwise dot products between embeddings indicates that even in nGPT, embeddings persphere (where the dot product would approach 0), but instead form clusters peossillslllyy reflecting natural patterns in language data. Dot products in GPT tend to have higher values due to its em- beddings forming hyper-ellipsoid, suggested by the spread of vector norms. The ill-conditioned nature of GPT's input embeddings could lead to computational issues involving these embeddings.
uniformly distributed across the hy-
Figure shows the median condition numbers (across heads) for attention and MLP matrices at different layer depths-22 -24 layers for the 0.5B model and 36 layers for the IB model. GPT mod- els exhibit significantly higher condition numbers in their attention matrices compared to nGPT. A closer inspection of these matrices (the 3rd and 4th layers are in Figure 12 and Figure [3 of Ap- pendix) suggests that they degenerate into lower-rank matrices, potentially reducing the learning capacity of these blocks. One could argue that the elevated condition numbers are influenced by the norms of the vectors in these matrices. Our post-training normalization of these matrices is depicted

---

Published as a conference paper at ICLR 2025
Distribution of norms of input emboddings
Sorted Eigenvalues input embeddings
Pairwise dol- pol-products of embeddinng
6FT8688 068
05 -PT
OFT 068
GFT
18
Eigemrake (redann EewdannEgarrnaiieon
13PPt 18
1B
Nornvakke
10000
Norma tred rank
Normalzed tark
Norma rark
Distribttion norms ostput embeddings
Sorted Eigenwaluss of output embeddinga
Paineise dol-products output embeddings
2:98
0.58
0:8
8PT
-PTT
KCPT 0,50
KGPT 0.50
Egenosku madeant Sewdeenniuuuaaaaa www.aal.....
MGPT
Nommvab
VILuA
54 Normalzze rank
Norma teed rams
Normallzed TONK
Figure 4: Left: Distribution of norms of vectors from input (Top line) and output (Bottom line) embedding matrices. Middle: Distribution of eigenvalues divided by its median value. Right: Pairwise distribution of dot products between embeddings. Models are trained for 100k iterations.
Condition number of Q each layer
Condition number of each layer
Condition number of at each layer
GPT 0.58
GPT 0
GPT 0.88
GPT
19
GPT
RGPT
NGPT
Condiron number
Condiar number!
Condi mondita mumiter
Morranizeed tnyer index
Normak red byyer inder
Normaltzed layer index
Condition number of MLP11 layer
Condition number MLP12 each layer
Condition number wach layer
0.58
0.58
0.28
HEDT 0.5
1227 3.5
HGPT 0.5
DGPT
Candiion nur
Curoliin FU
Conettion
No normallzzd layer Indox
Nommaized byer Indox
Nommlizeed layer indox
Figure 5: Median condition numbers for attention and MLP matrices at different layer depth (24 and 36 layers for 0.5B and 1B models, respectively). Models are trained for 100k iterations.
by the dotted lines in Figure of they remain higher than those for nGPT, indicating potential rank deficiency. The need for such normalization highlights one of the issues that nGPT is specifically designed to address.
of the Appendix. While the adjusted condition numbers are reduced,
8

---

Published as a conference paper at ICLR 2025
Eigan leeming ratee of Attention
for matris MLP
Scaling
10 0.35
nOPT 0.58
10.0 0.35
meanssts
meannbrsVVuus)
Normallieed indes
Nommlkzed byw index
Normalliedd indes
Eigen leaming
Scaling for matrie MLP
Distribution of scaling for embeddings
r0F1t
nGPP 18
1B
meanaar
50 Normakzzd byee Index
No Normallzed Iayer Index
Viduu
Figure 6: (Left): Eigen learning rates of Attention and MLP blocks. (Middle): Scaling factors ap- plied to the intermediate states of MLP. (Right): Scaling factors applied before the QK dot product; distribution of per-vector scalings applied to logits. Models are trained for 100k iterations.
An important contribution of this work is the decoupling of predictions made by the Attention and MLP blocks from their impact on the hidden state h. These contributions are controlled by the cigen learning rates a and QM- Their interpretation is straightforward: if C' for an embedding dimension i E Rdaasn is 0.2, then the update follows hG < (1 0.2) + 0.2/ Thus, they directly quantify the contribution of hA, into Figure shows the average absolute values hA and hm at each layer. Notably, the network learns to take only modest steps (20%-30%) in the direction suggested by hΛ and hM- The average magnitude of ala decreases from 0.25 the 0.5B network (24 layers) to 0.20 in the 1B network (36 layers). Meanwhile, parameter count per block linked to its maximum predictive capacity (in our setup, MLP blocks have more parameters than Attention blocks), then the greater eigen learning values observed for MLPs could potentially be attributed to the higher quality predictions made by these blocks.
decreases from 0.37 to 0.32. If the
The scaling factors Su. 8, and 8gk remain relatively stable across layers. The value of 8, be interpreted as a measure of the non-lincarity the SiLU function, which behaves like ReLU for large 8, and approximates linear unit for values near 0 (see also Appendix A. The distribution of 82 is primarily characterized by its mean, which influences the temperature of the softmax during cross-entropy calculations. The introduced scaling factors 8yk; 8,, 8, and s to compensate for the removal of magnitude information when normalizing matrices and embeddings.
3.3 ABLATION STUDIES
Appendix A.9 summarizes numerous ablation experiments. An important finding is that having fixed (non-learaalle values for Syk; Su. 8, and a single global learnable value for 8; leads to only a slight degradation in accuracy. Therefore, our presented general case can be simplified and become easier to interpret. Appendix A.8 demonstrates that nGPT can handle longer contexts without requiring any modifications RoPE

---

Published as a conference paper at ICLR 2025
4 RELATED WORK
Wang & Isola 2020)) provides a comprehensive overview of the arguments for representation learn- ing on the hypersphere. Spherical representations are associated with more stable training in the latent space of variational autoencoders (Xu & Durrett 2018) and in embeddings used for face ver- ification (Wang et al. 2017 Notably, when embeddings are well clustered, they tend be linearly separable from the rest of the embedding space (Wang & Isola 2020). Mettes et al. (2019) demon- strated that classification and regression can unified by placing prototype embeddings uniformly on a hypersphere, allowing for separation with large margins a priori. [Wang & Isola (2020) found strong empirical correlation between downstream task performance and both the alignment (close- ness) and uniformity of embeddings the hypersphere.
Since all embeddings in nGPT lie on the hypersphere, any update that causes the hidden state h to deviate from the manifold is followed by normalization step. This normalization can be interpreted as a retraction in the context of Riemannian optimization. One might attempt to approximate nGPT's update in GPT by applying RMSNorm both at the beginning and end of the block (Xiong et al. 2020). However, this approach does not guarantee a fixed norm for the hidden state, nor does ensure that the recombination approximates SLERP LERP.
The normalization in equations [5] and closely resembles the QK normalization by Henry et al. (2020). In nGPT, this process can be viewed restoring q and & of the i-th head a (dhoode//hhuuss Thends dimensional hypersphere after the projection of h by W, and WE. respectively. Since h and the embedding dimensions Wg and W are already normalized, the norms of q and lc are also com- parable, making their normalization potentially unnecessary. We investigated the effect of omitting this normalization our ablation studies (see Appendix A..p. The results indicate only a minor performance degradation with potential computational savings of about 12%. However, the normal- ization helps to maintain performance when extrapolating the context length (see Appendix A.8)
5 DISCUSSION AND CONCLUSION
This work builds on numerous key findings and observations made in the field which directly (Liu
al.
(017//Wap et al.) [2017: Liu et al.] [2018|Xu Xu & Durrett [2018] Wang & Isola, 2020 Liu al.
2021
Karras et al. 20224) and indirectly (Salimans & Kingma 2016 Franke et al. [023 Kodryan
jet al. 2022. Kosson et al. ) support representation learning on the hypersphere. One of our
main contributions is the normalization of the embedding dimensions across all of the transformer's matrices, ensuring they reside on the same hypersphere. Crucially, we observed that such normal- ization alone would constrain the inputs of non-linear units, and, thus, the scaling factors for these units should be introduced.
In line with recent studies suggesting that transformers implicitly perform gradient descent as meta- optimizers (Von Oswald et al.) 2023 Dai et al.) 2022). we explicitly demonstrate how this process occurs in the normalized Transtormer: the transformation blocks provide gradient information, ii) this information is multiplied by eigen learning rates to adjust the hidden state, and iii) the commonly used normalization can be interpreted retraction step in Riemannian optimization, projecting the point back onto the hypersphere. We believe we are the first to decouple the eigen learning rates from the rest of the network, recognizing them trainable parameters that can be interpreted the diagonal elements of a variable-metre matrix. In other words, the normalized Transformer functions as variable-mette optimizer, searching for output solutions using data-driven gradient information estimated in its attention and ML.P blocks.
The spherical representation provides valuable insights into the internals of nGPT, enabling the collection and analysis of statistics about its normalized components. Most importantly, it allows for the application of mathematical techniques specifically designed for dealing with hyperspheres. We believe that the reported acceleration, by factor from 4 to 20, is only the first step towards uncovering new algorithms and architectures that could emerge from nGPT Future work should explore scaling nGPT to larger network sizes, real-world datasets, and broader range of tasks. For instance, the extension of nGPT to encoder-decoder and hybrid architectures (Dao & Gu 2024 De et al. 2024) is straightforward.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Naman Agarwal, Nicolas Boumal, Brian Bullins, and Coralia Cartis. Adaptive regularization with cubics on manifolds. Mathematical Programming. 188, 2021.
Maksym Andriushchenko, Francesco D' Angelo, Aditya Varre, and Nicolas Flammarion. Why do we need weight decay in modern deep learning? arXiv:2310.044111 2023.
Thomas Bachlechner, Huanru Henry Majumder, Bodhisattwa Prasad Mao, Garrison W. Cottrell, and Julian McAuley, Rezero all you need: Fast convergence at large depth. In arXiv, 2020. URL
https://aae /ttp:://aaivvvvo org/ abs 2003. 04887
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can GPT learn in-context? Language models implicitly perform gradient descent as meta-optimizers.
arXiv:2212.10559 2022.
Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. arXiv:240.. 1060,, 2024.
Soham De, Samuel I. Smith, Anushan Fernando, et al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. arXiv:2402. 19427, 2024.
Kalyanmoy Deb, Karthik Sindhya, and Jussi Hakanen. Multi-objective optimization. Decision sciences, pp. 161-200. CRC Press, 2016.
Jörg K. H. Franke, Michael Hefenbrock, Gregor Koehler, and Frank Hutter. Constrained parameter regularization. 1.Xiv:2311.00058. 2023.
Aaron Gokaslan and Vanya Cohen. OpenWebText corpus. http://8kyyion007..gi0z Lio0penn@ebTeettorpuus 2019.
Alex Henry, Prudhvi Raj Dachapally, Shubham Pawar, and Yuxuan Chen. Query-key normalization for Transformers. arXiv:2010.04225. 2020.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Fix your classifier: the marginal value of training the last weight layer. arXiv preprint arXiv:1801.044 08540400444 2018.
Jerry Yao-Chieh Hu, Dennis Wu, and Han Liu. Provably optimal memory capacity for modern hop- field models: Tight analysis for transforner-compatbbe dense associative memories. Advances in Neural Information Processing Systems (NeurIPS), 37, 2024.
Andrej Karpathy. NanoGPT. http:::/gi.tubb..cm//keeeeeee
com
2023.
Tero Karras, Miia Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyz- ing and improving the training dynamics of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24174 24174-24184, 2024.
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv:1412.69980 2014.
Seijin Kobayashi, Yassir Akram, and Johannes von Oswald. Weight decay induces low-rank at- tention layers. Advances Neural Information Processing Systems (NeurIPS), 2024. URL Equal contribution.
https://en bttps:///arxvv wwtp///////poo..o.///..pp/.../pt /2410 .23819
Maxim Kodryan, Ekaterina Lobacheva, Maksim Nakhodnov, and Dmitry P Vetrov. Training scale- invariant neural networks on the sphere can happen in three regimes. NeurIPS, 2022.
Atli Kosson, Bettina Messmer, and Martin Jaggi. Rotational equilibrium: How weight decay bal- ances learning across neural networks. arXiv:2305. 17212. 2023.
Tim Large, Yang Liu, Jacob Huh, Hyojin Bahng, Phillip Isola, and Jeremy Bernstein. Scalable optimization in the modular norm. Advances in Neural Information Processing Systems, 37: 73501-73548, 2025.
Weiyang Liu, Yan-Ming Zhang, Xingguo Li, Zhiding Yu, Bo Dai, Tuo Zhao, and Le Song. Deep hyperspherical learning. Advances neural information processing systems, 30, 2017.

---

Published as a conference paper at ICLR 2025
Weiyang Liu, Zhen Liu, Zhiding Yu, Bo Dai, Rongmei Lin, Yisen Wang, James M Rehg, and Le Song. Decoupled networks. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 2771-2779 2018.
Yang Liu, Jeremy Bernstein, Markus Meister, and Yisong Yue. Learning by turning: Neural archi- tecture aware optimisation. In International Conference on Machine Learning, pp. 6748-6758.
PMLR. 2021.
Ilya Loshchilov. Weight norm control. arXiv preprint arXiv:2311.11440. 2023.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR, 2019.
Pascal Mettes, Elise Van der Pol, and Cees Snoek. Hyperspherical prototype networks. NeurlPS, 2019.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting of the Association for Computational Linguistics, pp. 311-318, 2002.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever.
Language models unsupervised multitask learners. hatpaeeeelaaagaggeoooooele///aaa multitask learners pdf 2018.
bettps://cdn.ope _ttptee_odllss____________e are_unsupervisee
Tim Salimans and Durk Kingma. Weight normalization: simple reparaaeettrriizationnn accelerate training of deep neural networks. NeurlPS, 2016.
Noam Shazeer. Gated linear units (glu). arXiv:2000.0552222222
Sam Shleifer, Jason Weston, and Myle Ott. Normformer: Improved transformer pretraining with extra normalization. arXiv preprint arXiv:2110.09456, 2021.
Ken Shoemake. Animating rotation with quaternion curves. In Proc. of the 12th annual conference on Computer graphics and interactive techniques, 1985.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model par- allelism. arXiv preprint 1xXi::999900053 2019.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputingg 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NewrlPS, 2017.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In ICMI, 2023.
Feng Wang, Xiang Xiang, Jian Cheng, and Alan Loddon Yuille. Normface: L2 hypersphere em- bedding for face verification. In Proc. of the 25th ACM nternational conference on Multimedia, 2017.
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through align- ment and uniformity on the hypersphere. In ICML, 2020.
Dennis Wu, Jerry Yao-Chieh Hu, Teng- Yun Hsiao, and Han Liu. Uniform memory retrieval with larger capacity for modern hopfield models. arXiv preprint arXiv:2404.00 03827, 2024.
Ruibin Xiong, Yunchang Yang. Di He, Kai Zheng, et al. On layer normalization the Transformer architecture. In ICML, 2020.
Jiacheng Xu and Greg Durrett. Spherical latent spaces for stable variational autoencoders. N.Xiv:8808,08000 2018
12

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.
RESCALING IN THE MLP BLOCK OF THE NORMALIZED TRANSFORMER
When computing SwiGLU using equation 18. each element z of the vector an input to SiLU:
SiLU(x) = o = Ite
(23)
where (xx) sigmoid. For a with large magnitude, SiLU(x)) approximates ReLU(x)): when oc, SiI.U(x:)
-oc, SiLU(xx) + 0, and when -0.278 is located at Imin & 1.2778. While the elements of v represent dot products of dmodet" dimensional vectors and are bounded in [-1,1], their expected absolute value (when they are ran- dom) is BiJ cos(
The minimum of (1n. a
2 0.7979 Thus, we should rescale x by V Vomodals otherwise, Vdmadel for very small x, we end up with SiIU((z) R x/2..
Vdmadel
alternative view is to note that since the variance of each of the normalized vectors (h and vector from W.) is 1/dmodel, the variance of 1 (suitable for the sigmoid part of SiLU) can be restored by rescaling of æ by /dmodel- Based on these two views, we rescale v by dmmodet to benefit from the non-linearity of SiLU.
A.2 EIGEN LEARNING RATES
In the main text of the paper, defined eigen learning rates positive (a |aX| is used during the forward pass). However, when they are not constrained to be positive, we obtain experimental results which are the same (up numerical difference). This surprising observation can be explained as follows. Both the attention and MIP blocks have transformation matrices at their outputs. When the search is unconstrained, it is sufficient for Adam to flip (select) the sign of the i-th row of the output transformation matrix Wo to change the sign of the corresponding i-th coordinate in hA Thus, the transformation calculated as æ₁ W. is the same as a' W where a'A;- -AA, and W (i.)) -W.,(..). In other words, when unconstrained, we can arrive exactly the same transformation by flipping the signs in both a and Wo. which cancel each other. For simplicity and clearer interpretation, we suggest constraining a to be positive in the main paper.
There is, however, another interesting scenario when eigen learning rate could become negative. In quasi-Newton methods, B approximates the inverse Hessian matrix H-1 whose diagonal elements are positive values if the function is locally convex which is the assumption of the Newton method (H needs to be positive definite). However, the diagonal of H can have negative values if the objective function is non-convex and has saddle points. While quasi-Newton methods like BFGS aim to ensure (e.g., via regularization) that B is positive-definite even non-convex problem, some Riemannian optimization methods can exploit negative curvature of H H- (Agarwal et al., 2021. When the diagonal of is unconstrained, we perform variable-metric step, acknowledging that the function may be non-convex locally. We did not mention this in the main paper because, as noted earlier, the results with both constrained and unconstrained a are essentially the same.
During the review process of this paper, we were made aware of the work of Bacclechner et al. (020)) which proposes ReZero (residual with zero initialization). In the context of Transformers, ReZero is implemented as:
h 4- h + ahT
(24)
where the hidden state h at each layer updated with a-rescaled output hT of the transformation block. The trainable layer-specific scalar a is the same for both the MLP and Attention blocks, with its initial value suggested to be set to zero. The authors also proposed to remove LayerNorm and suggested that BatchNorm is suitable alternative. We note that the introduction of the scaling factor a does not affect the architecture of the model because the output projection matrices of both MLP and Attention blocks already contain that same scaling factor their parameters (c.g., a simultaneous rescaling of a by k and the output projection matrices by 1/k results in a network with identical behavior). However, the use of single scaling factor can help
multiplicative part of
13

---

Published as a conference paper at ICLR 2025
to more quickly adjust the impact of each transformation block. Indeed, the authors argue that by initializing a to zero, one can achieve a form of curriculum learning in which individual layers of the network can be gradually activated during optimization.
The update of ReZero differs from the update of nGPT which can be written as:
h + Norm h + oer (rr(Nom(hrr) - -h))
(25)
where the original output of the transformation block hy is normalized before being used to compute an update direction, weighted by the eigen learning rate of the block ar. In contrast to ReZero, the update of nGPT is invariant to any rescaling of the output projection matrix, and, thus, the eigen learning rate can be viewed ReZero, where the scale of a cannot be viewed as a measure of the impact of the block without taking into account the scaling of the block's 's matrices (see, e.g., our previous example with rescaling by k). In contrast to ReZero, where a rescales he, the update of nGPT rescales the direction towards hr- Finally, the update of nGPT holds for the spherical representation where the updated hidden state retracted back to the hypersphere.
measure of the contribution of the block. This is not the case for
Another work that we were made aware of during the review process is NormFormer (Sheeferr al. 2021) which updates the hidden state as follows:
h - h + LN( hA)
(26)
h ah IIN (L((O()(()) W2,
(27)
where hA is the output of the Attention block with LayerNorm normalization of inputs. W1 and W₂ matrices are components of the MLP block with GELU activation function, and, are trainable scaling factors defined parameter-wise. If we omit the variance part of LayerNorm in Equation 26 can viewed as eigen
then I.N hA) can be reshaped to aNorm hy ) and one could wonder learning rates. This is not the case because, similarly to ReZero, i) the norm of h is not constrained, and ii) the scaling is applied to the value itself and not to the direction. In Equation 27 not only the norm of h is not constrained, but the output of the MLP is also not normalized. Instead, LayerNorm normalizes the rescaling factors for the vectors stored in W2. While the name NormFormer may suggest some similarities with nGPT, the approach does not normalize network parameters in any way and does not normalize embeddings during the update over layers (i.e., the typically growing norm of h affects its recombination with the outputs of the transformation blocks).
A.3 SUPPLEMENTARY NOTES TO INTRODUCTION
The original Transformer architecture is a solution in some design space which can be measured with respect to a set of metrics. Many known aspects multi-criteria decision-making and multi- applied (Deb et al., 2016). Various Transformer variants exist because
objective optimization depending on the selection of metrics (c.g.. performance on particular tasks) and constraints (c.g., memory, compute, data), the decision-maker prefers and selects (e.g., based non-dominance or aggregation of objectives) different design solutions. The first part of the Introduction section describes various architectural and optimization findings/solutions of the past. Then, is proposed to "unify various findings and observations made in the field under a new perspective of the normalized Transformer." This unification can be viewed as a recombination procedure in the design space of architectures and optimization approaches. The Introduction section concludes with a description of the key aspects of the proposed solution. The experimental results of the paper demonstrate some advantages of this solution w.r.t. the baseline GPT, e.g., improved condition number of attention matrices, faster convergence.
While the proposed solution can be viewed as a recombination of existing techniques, this is not how it was historically designed. The precursor of nGPT is described in Loshchilov (2023), where a modification of AdamW was studied in the training of Transformer models. That work demonstrated that instead of performing weight decay, the norm of all parameters of the network can be controlled and scheduled over the course of training. While this approach enabled acceleration in training, it was noted that the attention matrices are ill-conditioned and the normalization of the whole network or even individual matrices does not resolve it. In contrast, experiments with normalization of
14

---

Published as a conference paper at ICLR 2025
individual vectors within matrices demonstrated improvements of the condition numbers. Then, by combining this finding with the understanding that LayerNorm and RMSNorm normalize the hidden state to hypersphere (when all scaling factors equal), it hypothesized that the spherical representation a viable option. nGPT was designed from first principles, focusing on ensuring that the Transformer architecture respects the spherical representation for all its vector components. SLERP was identified as the proper way to recombine the hidden state and the output of the transformation blocks when working on the hypersphere and using scalar eigen learning rates. Scaling factors, such as those used for logits, were introduced to relax the constraints imposed by the spherical formulation.
During the review process of this paper, we were made aware of a set of closely related works. Karras et al. (2024) introduces a set of normalization techniques (very similar to ours) which greatly im- prove the training of Diffusion Models. [iu et al. (202T) proposed an optimization method for neural networks under hard hyperspherical weight constraints, Large al, (2025) employs a residual block structure which is related to the one of nGPT but without learnable parameters (eigen-learning rates). By connecting transformers to modern Hopfield models, Wu et al. (2024) show that well-separated memory patterns (learned tepresentations) enhance memorization capability (expressiveness) natu- rally achieved by mapping learning onto the hypersphere. The authors conclude that representation learning the unit hypersphere significantly improves expressiveness and memorization. (Hu al. (224)) claim that mapping to the hypersphere makes representation learning provably optimal for any dataset.
A.4 RIEMANNIAN OPTIMIZATION
If h - hA is viewed as the gradient in the Euclidean space, then, aligning with the requirements of Riemannian optimization, the projection of g onto the tangent space of the hypersphere
9pruj h(hT hA
(28)
The projection is equivalent to g when the dot product h hA is 1. Depending on the alignment between h and hx, the projected gradient varies between h hh (when the vectors are aligned) and - hA (when the vectors are orthogonal). The Riemannian variable-metric update then reads as:
h + Norm h Ba(h(h hx) - ha))
(29)
h + Norm h Ba (oahhh"mm))) -mM)
(30)
The normalization by Norm can be viewed as the retraction step in Riemannian optimization, map- ping the updated solution back to the manifold.
Our experimental results suggest that the impact of h hTu is negligible. Therefore, all experiments in the paper are based on equations 10 and
A.5 TIME COST PER STEP
The time cost per step for nGPT is approximately 80% higher with 4k context length, and 60% higher with 8k context length. This overhead is not only due to nGPT having 6 normalization steps (2 of them are applied for q and k) per layer instead of 2, but also because nGPT's normalizations are not yet fully optimized, unlike GPT, where normalization layers are fused with other operations. Training on larger networks is expected to further reduce this performance gap. as the number of layers (and thus the number of normalizations) increases only modestly with the number of network parameters. Appendix A.9 shows that we can remove the normalization of q and with a minor negative impact on results.
A.6 EXPERIMENTAL SETUP
In all experiments, we use OpenWebText (Gokaslan & Cohen 2019) dataset which, according to experiments of Karpathy 2023), represents good approximation of the OpenAl's internal dataset
15

---

Published as a conference paper at ICLR 2025
Table 2: Model Parameters for GPT and nGPT
Model Parameter
0.5B Models
1.0B Models
Number of Layers (Mayyrr)
24
36
Model Dimension (dmodel)
1024
1280
Number of Attention Heads (Theuds)
16
20
Key Dimension (dg)
dmodet / Pheads md
dmadeel/ / Theads
MLP Dimension (dMLP)
4dmodel
Parameters in GPT
468.22 .8.2M
1025.7M
Parameters in nGPT
468.4M
1026.1M
Table 3: Optimization Parameters for GPT and nGPT
Optimization Parameter
GPT AdamW 0.1
nGPT
Optimizer
Adam (AdamW with weight decay 0.0)
Weight Decay
0.0
Number of Warmup Steps
2000
0
Learning Rate Schedule
Cosine Annealing
Cosine Annealing problem-specific
Initial Learning Rate
problem-specifie
Final Learning Rate
used to train GPT-2 models. We are well aware that this dataset is not of the highest quality. How- ever, we believe that it suitable for academic research, and, should improve the comparability of our findings with other research papers.
We trained our models using 64 A100 GPUs distributed across nodes (8 GPUs per node). Global batch size is 512. We use the LLaMA-2 tokenizer with 32k tokens. We use the same setup for the 0.5B and 1.0B parameter models. All parameters of matrices are stored in bfloat16.
All matrix parameters are initialized by sampling from a zero-mean normal distribution with stan- dard deviation of 0.02 for GPT and 1/ Ammdel for nGPT. The standard deviation for the output matrices was scaled by a factor of /2 X Nlayer. as suggested by Radford al. (2018). The ini- tialization of matrix parameters is not important for nGPT because they are normalized afterwards. The base of RoPE is 10000. The initialization of the additional parameters introduced in nGPT is described in Section
experiments described in this paper were performed using an internal library based on Megatron-LM (Shoeybi et al.) (019). In order to illustrate how nGPT works, we re- implemented nGPT using nanoGPT (Karpathy) (023) and published our re-implementation at coD/NVIDAA///gg It should be noted that the latter re-implementation only qualitat- evely replicates our internal experiments and should not be used in production. However, it should be great quickly get familiar with nGPT its core functionality represents only few dozens of lines of codes. It is very important to note that when implementing nGPT in training libraries, one should make sure that not only instantiated model parameters normalized but also the ones which are used by the optimizer. Missing the latter is a common bug that should be avoided.
A.7 SELECTION OF THE INITIAL LEARNING
The initial learning rate the only hyperparameter we tune for both GPT and nGPT. Figure demonstrates our trials to select the most suitable initial learning rate for GPT and nGPT. Our first experiments started with the 0.5B model and Ik context length. After observing the general trend of these curves, we reduced the number of experiments and followed the trends minimize the total compute used. Apart from estimating the optimal settings for the initial learning rate, our general observation is that the optimal learning rates for GPT and nGPT tend to be similar for the same values of the validation loss. Longer runs are usually associated with the possibility of achieving lower values of the validation loss, and this typically requires lower initial learning rates.
16

---

Published as a conference paper at ICLR 2025
One artifact observed is the increasing sensitivity of the IB nGPT model 8k context length. We found that the smoothness of the hyperparameter curves can be restored (see the dotted lines with squares) by increasing inii and MMiiini from their default value of 1/ V mdddel to 0.1. This change decreases the effective learning rate of Adam these variables by a factor of (1/ /dmoel) / /0.1 se 3. As result, the eigen learning rates are learned by Adam at a slower rate.
0:58
0.58
with 5K contect
9-20k
-0kkk
100k
10k
254
100x
25A
B0k
-C-200x
400%
100
50k
2.5
100%
60K
2004
100
lesming with context
leaning
Intted
100K
100%
200k
200%
Fine veltiaton loak
25h
D0k
6004
100k
18k
100%
laaning
learning
Figure 7: Final validation loss values for different initial learning rates for the 0.5B models (Top) and IB models (Bottom). GPT is denoted by solid lines with circles, while nGPT is represented by the dotted lines with stars. specific case with a different setup. denoted by the dotted lines with squares (Bottom Right), is discussed in the text.
17

---

Published as a conference paper at ICLR 2025
context: exsy
context: hellawwag
winogrande
0.0
6.58
c:88 OPT OPT 16
carn
GPT e nPP
HOPT O SPPTI8 18
28
Training
1K contect: wse275
1K context: lambada openal scqjen1024
IK contest: average
0.1
C CPT 6nn O
C-CPT c.58
CIPT 248
O-CPT
SPPTTT
GPTT
11
2B Trining tikene billoon
100 Triining Bokere
Training ning Moinns
Figure Models trained with I context length. Final performance (y-axis) on set of downstream tasks and their average value (Bottom-Right for different computation budgets in tokens (x-axis).
easy
context: hellaswog
winogrande
0:13
0.38
6.0
GGPT 8.58
GPT SPTT 1B
KGPT 66
8 --PTTT 1B
S-IFTT 18
Triining
Training
contest: wsc275
BK context: lambado openal s0g
BK coptext: average
km
610
GPT 0:B
O GPT c.s8 G-GPT -PFT
GPT 2u58
0:0
HIGTT 18
Thining bilkorn
Treining Iokers
Tratning Miiers
Figure 9: Models trained with 8k context length. Final performance (y-axis) on set of downstream tasks and their average value Moottom-Right for different computation budgets in tokens (x-axis).
18

---

Published as a conference paper at ICLR 2025
GPT 0.55
GIPT 0.25 HOPT HOPT 18
O GIFT 0.38
CHOPT GrrdPPT
HORT S0PP 18
Training tocans niblloss
Traiiig tokans biloos
Traiiing Iokend in bilone
Figure 10:
Performance on WMT14-FR-EN task which evaluates the ability (measured by the
BLEU score of Papineni et al. 2002) to translate French sentence to English given five (French, English) example pairs, tollowing the methodology of Radford et al. (2018) where it was used to benchmark GPT-2.
layer
number bayer
layer
3.--GTT. 3.00
-277777 270
GPT 3.50
GFTT 155
OPT
Condito number
Conditon number
Condaton
layer
Condition number MIL Iayar
Condition number of
Inyer
Condition number of each layer
0.50
GPT 5.0
GPT 0
G27 10
GFT
Condbion UUTEOT
CO NUTTOS
Congron
Normel andd Inger Indeer
Norme fred layer
Normllined lnger Indee
Figure 11: Median condition numbers measured for attention and MLP matrices at different layer depth (24 and 36 layers for 0.5B and 1B networks, respectively). The dotted lines are for the case when GPT's matrices are renormalized after training. Models are trained for 100k iterations.
19

---

Published as a conference paper at ICLR 2025
Ostrrbution of noms Q
Borted Singuler Vakuos for heads of
0P7 6.0
Now
OIf c0 (7150
Noma Redirone
Norraked
Datritution of nomrs of
Scrted Singuler Vakues for heads of
Norn eake
GPT 11. send 10 0,0000
Noma Bred PANE
Norraline 1992
Figure 12: Left: Distribution of the norms of vectors forming the complete attention matrices (not per head). Right: Distribution of singular values for all heads. The results are shown for the 3rd layer.
Oxtribution of nome O
Sorted Singuler Vakues heads of
<APT
Nom
(PT 1R www. C94
2ord--4-1100 4.10
Noma Dedrrwe
Norroleed
Osoribution of nome of
Sorted Singsle Values for heads of
OFf
Norm K010
OPT 0.50 30ma- 130 Eul 1CT 33557
Nomea bad PA0E
Nerralodd 1852
Figure 13: Left: Distribution of the norms of vectors forming the complete attention matrices (not per head). Right: Distribution of singular values for all heads. The results are shown for the 4th layer.
20

---

Published as a conference paper at ICLR 2025
A.8 LENGTH EXTRAPOLATION ABILITY
We investigate the length extrapolation ability of nGPT by evaluating its perplexity on the PG19 dataset, shown in Figure [4 In standard GPTs, perplexity tends increase dramatically when tested on sequences longer than pre-training length. In contrast, nGPT maintains stable perplexity range even extrapolated lengths. This result demonstrates that nGPT can handle longer contexts without requiring any modifications to RoPE, providing clear advantage over standard GPTs in terms of length extrapolation.
In the original version of the paper, we explored the impact of gk-nooralization on performance of nGPT. Our conclusion was that the performance is equivalent but removing jk-nommalization can provide computational savings. During the review process of the paper, two reviewers asked ques- tions which made us further investigate the role qk-normalization and context length extrapolation abilities of nGPT. This resulted in noticing that nGPT without qk ok-normalization shows worse extrap- olation abilities (see Figure Right.. This difference was imitially overlooked because the results of nGPT with and without nk-normalization average accuracy on downstream tasks and validation loss of both variants in Table (6) In other words, their in-distribution performance appears the same (i.e., gk-normalization does not explain nGPT's better performance) but their outtof-distributtin performance is different. The difference from the fact that qk-normalization normalizes q and k. and, thus, guarantees that their dot product is bounded in [- [-1,1]. Even without using qk normalizatiinn in nGPT, the hidden state and the vectors forming Wg and W₁ matrices are normalized so that their dot product is already bounded in [-1,1]. However, the results of these individual dot products (based on vectors from W, and W;) will form q and & per head. Moreover, q and are also affected by RoPE. The qk-normalization restores q and k vectors back to hyper-sphere whose dimensionality is defined by the size of each head.
other tasks practically equivalent (see, e.g., the
While the extrapolation abilities of nGPT can be partially attested to the use of qk-normalization. gap between nGPT without qk-normalization and the baseline GPT still substantial. Another contributing factor the difference is the fact that many attention matrices of the baseline GPT are practically low-rank, as discussed in the main paper. This makes them less efficient in dealing with the context. Interestingly, Kobayashi et al. (2024) recently demonstrated that weight decay induces low-rank attention layers, thus confirming our observations. While the low-rank attention matrices can potentially appear due of factors, we note that weight decay is used in nGPT.
PG19 Perplexity
PG19 Perplexity
10
1K GPT 18
8K GPT 18
10
2K nGPT 0.58 (ramove @Knrmm)
1K rGT 18
(log)
8K rGPT 18
PPL (log)
20
204
10a
10a
2K
8K
Length
Length
Figure 14: PG19 perplexity from 1K 32K among different training lengths for the original GPT and nGPT (Left) and also with GPPT when go-normalization is removed (Right).
21

---

Published as a conference paper at ICLR 2025
A.9 ABLATION STUDIES
We perform a series of ablations in nGPT, focusing on the selection of scaling factors 8yk, &y (8w), 82. well as evaluating the necessity of normalization. For the scaling factors, first examine the impact of different initializations of these parameters. Additionally, we explore whether these factors can be simplified to learnable scalar or a fixed value. For normalization, we analyze whether the removal of 2K--ormalizzttion is feasible alternative. To ensure a fair comparison, all our ablation models have size 0.5B and context length 1 and are trained with learning rate 10-3 for 100k iterations (~52B tokens).
Table presents various combinations of Sinst and Sacule for each scaling factor, as well as the downstream task accuracy and validation loss. Additionally, we report the mean value of each scal- ing factor distribution to show the final converged scaling value after training. For 8yk- we observe that the mean value (Mean(s) s)) is relatievly stable with value around I across most initialization set- tings, except when using smaller Sinit and larger sccale, which results in a Mean(s) value less than 1. For su and g., changes in initialization lead significant shifts in Mean(s), accompanied by increases in final validation loss. For 8., we see a substantial degradation in both accuracy and loss under certain initializations. Overall, these results suggest that the baseline hyperparameter settings are robust but additional tuning of 8u, 8u, and 8, could potentially further improve performance.
In Table modify each learnable per-element vector of the scaling factors 8yk- 8u (8w). 82. and the eigen learning rates OA Q&M to single learned sealar or a fixed value. This ablation helps us determine whether these tuning parameters can be simplified in nGPT. From the table. we observe that most changes result in only negligible degradation (S. 0.3%) in validation loss, and some even lead to slight improvements in accuracy. Notably, replacing the per-element vector Syk with single scalar has minimal impact on the mean value, while the mean value of Su, Su, O' and a show a slight increase. Furthermore, even when fixing Syk. Su. and s,, the model still achieves comparable accuracy and validation loss to the baseline using the per-element vectors.
In Table 6 we investigate the impact of removing QK normalization operations and the approxi- mation of using LERP instead of SLERP to potentially reduce training time per step. Specifically, removing the QK normalization terms (Norm) in equations 15 and reduces training time per step by 12%. Conversely, replacing the baseline LERP in equation 7 with SLERP in equation increesess training time by 10%. Despite these changes, both modifications result in comparable accuracy and loss to the baseline, demonstrating their effectiveness in reducing computational overhead without sacrificing performance.
22

---

Published as a conference paper at ICLR 2025
Table 4: Ablations of sinit and Sscale. Our default setup is the first row of each section. Mean(s)
is the mean value of the learned distribution after training. Avg. Acc is the average accuracy of our selected five downstream tasks, Valid. Loss is the final validation loss, and we use dmodet 1024.
iisst
scaate
Mean(s)
Avg. Acc (%) f
Valid. Loss J
Syyt
1/ V dondee
1.51
54.44
2.252
8gk
0.33
1/ Vmodel
1.36
54.67
-0.33%
8yk
0.05
1/ Vmndel
1.01
53.69
-0.09%
Sgt
1,38
54.19
-0.09%
Agt
0.33
1.11
54.89
-0.08%
Byk
0.05
0.29
52.41
+1.94%
& Su
1.24 & 1.12
54.44
2.252
8,, & 8
Cmnnded
0.03 & 0.03
53.51
+0.92%
$n. & Sn
1/v dononel
2.75 & 11.2
53.90
+0.05%
& Sy
1/Vamodel
1/ 1/Vdmodee
0.40 & 0.25
54.39
+0.48%
82
1/ donadee
60.8
54.44
2.252
dimodel
1/ d/noeeee
106.1
52.60
+1.06%
23.6
52.71
+3.12%
8; Amodel
63.6
51.84
+1.66%
Table 5: Ablations of replacing learnable per-element vector (scaling factors and eigen learning rate) with a single learned scalar or fixed value. The number of the learned vector, learned scalar and fixed value the mean values across all layers.
$y & 8y
$;
0 & a.
Avg. Acc (%)
Valid. Loss J
vector 1.47
vector 1.12 & 1.24
vector 60.80
vector = 0.22 & 0.33
54.44
2.252
scalar 1.49
vector 1.13 & 1.23
vector 62.
vector = 0.22 & 0.33
54.05
+0.22%
vector 1.46
scalar 1.46
vector 60.70
vector = 0.22 & 0.33
54.23
+0.07%
1.47
1.12 1.24
scalar 95.65
0.22 & 0.33
53.69
+0.20%
vector 1.40
vector 1.13 & 1.13
vector 60.90
scalar 0.30 & 0.36
54.86
+0.22%
scalar 1.51
scalar 1.477
vector = 16.188
vector = 0.22 & 0.32
54.52
+0.17%
scalar 1.52
scalar 1.68
scalar 96.65
vector 0.22 & 0.30
52.59
+0.30%
scalar 1.49
scalar 1.17
scalar 88.64
scalar 0.30 & 0.37
53.61
+0.62%
value 1.00
vector 1.12 1.15
vector 60.69
vector = 0.24 & 0.35
54.17
+0.11%
vector 1.45
value 1.00
vector 61.53
vector 0.23 & 0.35
55.51
+0.05%
value
value 1.00
vector 61.26
vector 0.24 & 0.36
53.63
+0.20%
value 1.00
value 1.00
scalar 96.15
vector 0.22 & 0.35
53.37
+0.40%
Table 6: Ablations of the design choices of nGPT. The removal of QK norm involves taking out
the normalization terms Norm in equations 15 and 16 The replace of LERP with SLERP is using
equation 6 without approximation of equation 7
Training time per step (s)
Avg. Ace (%)
Valid, Loss J
Baseline nGPT
0.657
54.44
2.252
I) remove QKnor
0.576
54.71
+0.12%
2) replace LERP with SLERP
0.726
54.80
-0.08%
23

---

Published as a conference paper at ICLR 2025
A.10 ANALYSIS SCALING PARAMETERS
In nGPT, we introduce a total six trainable parameters: the eigen learning rates a and a., along with the scaling factors 8yk. Su, 8y and 82. Since these parameters updated using gradients, are interested in understanding their learned distributions after training. In Figure [5, we present the histograms these saved weights combining across layers and analyze their distributions under different context lengths, model sizes, learning rates, and numbers of training tokens.
For O and a., we observe that their distributions remain stable across various learning rates. However, increasing the context length or the number of training tokens tends to shift the distribu- tions the right, indicating that the hidden state requires more transformation from the Attention and MLP blocks to handle the larger inputs. Additionally, the mean values for these parameters tend be smaller in larger models because we have more layers to update the hidden state.
Regarding the scaling factors, 8yk is quite stable under all the conditions, suggesting that fixed value, similar to the softmax scaling factor in original transformer, may be sufficient. Interestingly, we find that 8gk has high density close to zero, indicating the sparsity of attention in nGPT. For Su and Sy, their distributions shift left as the context length increases, but shift right as the model size the number of training tokens increases. Furthermore, these two factors become flatter when using larger learning rate. Lastly, for 82, we see higher mean values for longer context lengths, larger model sizes, and more training tokens, suggesting that the model may learn to lower temperature to make the final word distribution sharper.
24

---

Published as a conference paper at ICLR 2025
4c 203)
200
Loni 2.000
60-
Etffeentt langthe
dlferest model
diffrrrrttttto tearring
differrent tobess
(reaaed 3m7
reaene
240
(0)
0.
0.2
04
0.0
0.
0.2
0.4
0.
1.6
sbets y otferent lengths
abelSy) differest medel shes
16-
W 1a15l1-dPrree diffrreet leeriigg retee
abotSp) diffrrett tohems
199
0.0
1.5
2.
25
68
00
05
1.0
15
200
25
20
08
00
0.5
20
20
0.
0.5
1.0
1.5
20
25
30
2.0
9, diferentt leggtt
9, dtfferett madel Gume
9y Oiftrrstt Teareing
ffferunt Salerns
Winan-L Ik kirmwiit
(w220003 (aaa-11
2.3
00
23
00
05
1.0
1/5
3.0
35
100
0.
2.0
09
03
1.5
1.5
2.5
2.
30
190
007
120
1.09 Prur -08
17 97.4
5354 pou wel
Figure 15: The distribution of trainable eigen learning rates and scaling factors varies under different context lengths, model sizes, learning rates, and number of training tokens. If not specified, context length 1K, model size 0.5B, learning rate 2.0 x 10-3, 3 and training tokens 52B as default.
25


Reviews:
[{"review_id": "HpXPHoyA4R", "reviewer": "Reviewer_Z8Gb", "summary": "Motivated by previous findings on the advantages of representation learning on a hypersphere, the authors propose a novel neural network architecture called the normalized Transformer (nGPT). In nGPT, all vectors—including embeddings, MLP outputs, attention matrices, and hidden states—are normalized to ensure they reside on the same hypersphere. This setup ensures that each token in the input sequence, during optimization, starts at a position on the hypersphere defined by its input embedding and moves on the surface of the hypersphere to a point that best predicts the embedding of the next token. Experiments on the OpenWebText dataset demonstrate that nGPT learns significantly faster, reducing the number of training steps needed to achieve performance comparable to baseline models.", "questions": "**Q1:** The authors comment on the average magnitudes of $\\alpha_A$ and $\\alpha_M$, noting that $\\alpha_M$ has a higher average magnitude than $\\alpha_A$, implying their suggestions more precise. could the authors the authors elaborate on this point in more detail?\n\n**Q2:** Previous research has demonstrated that representation learning on a hypersphere enhances training stability, with experiments and ablations confirming stable activations, gradients, and weights, as well as showing gradient descent acting as a meta-optimizer. The 4-20x speedup achieved by the normalized Transformer is impressive, and further intuition or analysis on the factors contributing to this acceleration would be valuable.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "The paper is well-written and well-structured. The pair-wise comparisons in Section 2 make it easy to understand the architectural modifications, with clear motivation and explanations for each change.\n\nThe comprehensive experiments and ablation studies effectively demonstrate the method's effectiveness. The experimental results are impressive.", "weaknesses": "Minor: The per-step time of nGPT is approximately 80% higher for a 4k context and 60% higher for an 8k context, primarily due to additional normalization steps. However, Table 6 shows that removing the normalization for q and k only slightly affects performance, suggesting a potential way to reduce overhead.", "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": "ephxH6jiB6", "reviewer": "Reviewer_MiZy", "summary": "The paper introduces a novel neural network architecture, the normalized Transformer (nGPT) with representation learning on the hypersphere, and empirically shows that nGPT learns significantly faster, cutting down the training steps needed to reach comparable accuracy levels.", "questions": "I have no question.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written, and the related work section is presented clearly.\n\n- The experimental results are solid and convincing.", "weaknesses": "- I do not find any major weaknesses. My only concern is that, in essence, the method involves normalizing everything to unit, functioning as an embedding layer within architectures, with empirical results demonstrating its effectiveness. Although hypersphere normalization in representation learning is introduced and applied, the paper provides a clear discussion of related papers.", "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": "SliRfNoxNS", "reviewer": "Reviewer_4Taq", "summary": "This paper proposes a novel architecture called the normalized Transformer (nGPT), which performs learning on a hypersphere. By normalizing parameters and output features, such as embeddings, MLPs, and attention layers, and introducing learnable scale parameters, the model enables effective representation learning on the hypersphere. The model demonstrates faster convergence compared to GPT, even without using conventional techniques like weight decay and learning rate warmup, in models with 0.5B and 1B parameters.", "questions": "- There are existing studies, such as Bachlechner et al. (2021), that demonstrate fast convergence speeds in Transformer. It would be helpful to see how nGPT compares with these previous studies.\n  - Bachlechner et al., “ReZero is All You Need: Fast Convergence at Large Depth”, UAI 2021\n- An interesting aspect of nGPT is the linear interpolation between the output of the attention and MLP blocks and the hidden state $h$ in Section 2.2.2. The authors used linear interpolation as shown in Equation (7), but I am curious why they did not consider directly using geodesic interpolation from Equation (6).\n- It seems that the bold $\\alpha$ in Equation (9) should be corrected as $\\alpha$.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Fast convergence in training is practically important for saving computational power, and this paper aims to offer a meaningful solution through the proposed architecture, nGPT.\n- The model is carefully designed to prevent loss of expressiveness by introducing appropriate scaling parameters when applying existing results on the effectiveness of representation learning on the hypersphere on Transformers.\n- Detailed ablation studies are conducted on the newly introduced learnable parameters, $\\mathbf{s}_{qk}$, $\\mathbf{s}_u$, $\\mathbf{s}_v$, and $\\mathbf{s}_z$, demonstrating that nGPT is robust to the introduced scale parameters.", "weaknesses": "- Many variations of the Transformer architecture have been proposed, but they have generally failed to successfully replace the original Transformer. This is because they have not demonstrated the scalability inherent in the original Transformer. In this paper, scalability also cannot be verified as performance has not been tested on models widely used at the 7B parameter scale or above.\n- As the authors mentioned, nGPT requires an additional computation time of 80% for 4k and 60% for 8k context in a single iteration. While the authors argue that this can be reduced through code optimization, such as fused ops, the normalization of all matrices and embeddings at the beginning of each step would not benefit from fused ops, and GPU communication costs would arise with larger models. Verification of actual wall-clock time on larger models is necessary.\n- The authors experimentally show that nGPT has length extrapolation ability in addition to fast convergence, but there is no discussion as to why nGPT possesses this property. Adding insights into this aspect could greatly benefit future research.", "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": "KPzCqb7vAt", "reviewer": "Reviewer_wJPA", "summary": "This paper proposed the normalized Transformer (nGPT), which normalizes all vectors forming the embeddings, MLP, attention matrices, and hidden states into unit norm to travel the tokens on the surface of a hypersphere via optimal scaling factors, trainable scaling parameters, and a variable-metric optimizer, resulting in faster convergence.", "questions": "**Question 1.** Why nGPT converges faster than the original Transformer? Experimentally, the faster convergence of the normalized Transformer was confirmed, but it was not clearly explained why\n\n**Question 2.** The rationale behind the approximation of the SLERP by LERP needs to be clarified. While the motivation and theoretical background are well explained, the purpose and reason behind this approximation still need to be clarified. Is there any reason not to use SLERP with normalization? Or, is there any benefit to using LERP instead of SLERP?\n\n**Question 3.** The benchmarks used in the experiments solely encompass common sense reasoning and word/sentence prediction. Are the authors convinced that the experimental results are enough to indicate the versatility of the normalized Transformer in other NLP tasks? If so, I would like to know the rationale behind this conclusion.\n\n---\n\n**Things to improve the paper that did not impact the score**:\n\n* [line 88] nGPT appeared before defining what it is\n\n* This paper explained how to build it, summarized the modifications, and described hyperparameter settings in detail. Also, the normalized Transformer is easy to be implemented. However, it is recommended to publicly open the code for reproducing the paper's results.\n\n* The structure of the paper -- which divides the proposed method into its constituent components and presents these as separate sections, and then introduces and compares the proposed method in a sequential manner -- is not a common one. This structure facilitates the reader's comprehension of the proposed method as comparing it with the baseline. However, this structure may be perceived as somewhat akin to that of a technical report. It would be advisable to restructure the paper.\n\n* In [4], the use of a single trainable scale parameter on normalized vectors within a matrix has potential to facilitate an examination of the training behavior exhibited by the model. It may be advantageous to consider the implementation of a single trainable scale parameter for each layer or attention within the normalized Transformer. Could this method result in the generation of superior outcomes compared to those observed previously?\n\n---\n\n[1] Salimans, T., & Kingma, D. P. (2016). Weight normalization: A simple reparameterization to accelerate training of deep neural networks. _Advances in neural information processing systems_, _29_.\n\n[2] Yaras, C., Wang, P., Zhu, Z., Balzano, L., & Qu, Q. (2022). Neural collapse with normalized features: A geometric analysis over the riemannian manifold. _Advances in neural information processing systems_, _35_, 11547-11560.\n\n[3] Shleifer, S., Weston, J., & Ott, M. (2021). Normformer: Improved transformer pretraining with extra normalization. _arXiv preprint arXiv:2110.09456_\n\n[4] Hoffer, E., Hubara, I., & Soudry, D. (2018). Fix your classifier: the marginal value of training the last weight layer. _arXiv preprint arXiv:1801.04540_.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The proposed method was well explained, and it is clear how to convert the original Transformer to the normalized Transformer.\n\n* The correctness of each component of the proposed methods was validated through the comparison of experimental results with the original method.\n\n* This paper proposed a variable-metric optimizer that employs LERP and eigen learning rates. This method allows the optimizer to identify an effective learning rate and discard the weight decay.\n\n* This paper identified the appropriate scaling factor to mitigate potential risks associated with applying diverse normalization in Transformer. The proposed approach appears to improve the convergence rate of normalized Transformer.", "weaknesses": "**Weakness 1.** *Section 1. Introduction* provided a general overview of normalization and representation learning on the hypersphere within the context of Transformer. However, the research problem that this paper addresses and the rationale behind its significance are not sufficiently highlighted. It should be considered to emphasize these aspects before introducing the paper's contributions.\n\n**Weakness 2.** This paper presents two novel components: one is LERP and eigen learning rates, and the other is the identification of an appropriate scaling factor. The other components, such as the normalization of all matrices and the trainable scaling factor, could be regarded as an integration of diverse normalization techniques (Weight Normalization [1], Feature Normalization [2], NormFormer [3]) and scaling parameter tuning on normalized vectors [1,4]. To further enhance the novelty of this paper, it is recommended that the novel components be validated theoretically or experimentally (e.g., experimentally validate how much the LERP approximates the SLERP after training) and that the normalized Transformer be compared with previous normalization methods, (i.e., describes the difference with other normalization techniques)\n\n**Weakness 3.** This paper proposed numerous ablation experiments to investigate the impact of scaling factors and learnable scaling parameters on model performance. However, it seems essential to conduct an ablation study of the components of the normalized Transformer to gain a comprehensive understanding of the influence of each component on model performance.", "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)