DNT: a Deeply Normalized Transformer that can be trained by Momentum SGD
Authors: Xianbiao Qi
Marco Chen
Wenjie Xiao
Jiaquan Ye
Yelin He
Chun-Guang Li
Zhouchen Lin
Venue: ICLR 2026 Poster (2026)
Keywords: Transformer
Deep Normalization
SGD
Primary area: optimization
TL;DR: We have introduced a novel architecture, Deeply Normalized Transformer (DNT), which enables efficient training with vanilla momentum SGDW (mSGDW), achieving performance on par with AdamW-optimized Transformers.

Abstract: Transformers have become the de facto backbone of modern deep learning, yet their training typically demands an advanced optimizer with adaptive learning rate like AdamW, rather than a momentum SGDW (mSGDW). 
Previous works show that it is mainly due to a heavy-tailed distribution of the gradients. 
In this paper, we introduce a Deeply Normalized Transformer (DNT), that is meticulously engineered to overcome 
the heavy-tailed gradients issue, enabling seamless training with vanilla mSGDW while yielding comparable performance to the Transformers trained via AdamW. 
Specifically, in DNT, we strategically integrate normalization techniques at proper positions in the Transformers to effectively modulate the Jacobian matrices of each layer, balance the influence of weights, activations, and their interactions, and thus enable the distributions of gradients concentrated. 
We provide both theoretical justifications of the normalization technique used in our DNT and extensive empirical evaluation on two popular Transformer architectures (\ie, ViT and GPT), validating that: a) DNT can be effectively trained with a vanilla mSGDW; and b) DNT outperforms its counterparts.

Published as a conference paper at ICLR 2026
DNT: A DEEPLY NORMALIZED TRANSFORMER THAT CAN BE TRAINED BY MOMENTUM SGD
Xianbiao Qi', Marco Chen', Wenjie Xiao JJiauun Ye', Yelin He', Chun-Guang Li
Zhouchen Lin +5,6
Intellifusion Inc.
"Tsinghua University JJhns Hopkins University
*Beijing University of Posts and Telecommunications
"State Key Lab of General AI, School of Intelligence Science and Technology, Peking University Institute for Artificial Intelligence, Peking University
ABSTRACT
Transformers have become the de facto backbone of modern deep learning, yet their training typically demands advanced optimizer with adaptive learning rate like AdamW, W, rather than a momentum SGDW (mSGDW). Previous works show that is mainly due heavy-tailed distribution of the gradients. In this paper, we introduce Deeply Normalized Transformer (DNT), that is meticulously engineered to overcome the heavy-tailed gradients issue, enabling seamless training with vanilla mSGDW while yielding comparable performance to the Transformers trained via AdamW. Specifically, in DNT, we strategically integrate normalization techniques at proper positions in the Transformers to effectively modulate the Jacobian matrices of each layer, balance the influence of weights, activations, and their interactions, and thus enable the distributions of gradients concentrated. We provide both theoretical justifications of the normalization technique used in our DNT and extensive empirical evaluation two popular Transformer architectures (i.e., ViT and GPT), validating that: a) DNT can be effectively trained with a vanilla mSGDW; and b) DNT outperforms its counterparts.
INTRODUCTION
Transformer Vaswani et al. 2017) has revolutionized numerous domains in artificial intelligence,
demonstrated remarkable capabilities
natural language processing (Radford et al. 2018 2019
Brown et al. 2020 Dubey et al. 2024 Team
2023 Liu et al. 2024), computer vision Dosoviiskiy
al. 2020 Liu et 2022 Dehghani et 2023), AIGC (Ramesh al., [002] Peebles & Xie
2023 and multi modal applications (Li et al. 2022 Liu et al. 2023a) and become the de facto
backbone of modern deep learning.
Nowadays, it is widely accepted that Adam (Kingma & Ba, 2014) or its descendant AdamW (Loshchilov Hutter (209)) are the standard optimizer for training Transformers; whereas the classical SGD (Robbins & Monro 1951) and its variants (Nesterov) 1983; 1998; Johnson & Zhang 2013). e-g., momentum SGD (mSGD), usually under-perform when training Transformers. Despite of its heavier load on GPU memory than mSGD, Adam is used as the optimizer in most
recent studies on Large Language Models (LLMs) (Dubey et al., 22024] Team 2023 Liu et al,, (024)
and multi-modal models (Li al. 2022
Liu et al. 2023a). Naturally, interesting question arises:
Can Transformers be trained via mSGD yield performance matched to that is trained via Adam? Or, under what conditions?
To answer these questions, we need to understand why mSGD typically underperforms Adam when training Transformer. Previous studies (Simsekli et al. 2019 Zhang et al. 2020) reveal that the fundamental reason lies the statistical property of the stochastic gradients in Transformer
Corresponding authors

---

Published as a conference paper at ICLR 2026
600
600
Gf
9. est 001101
Distributions of the absolute values of the entries in gradients for VIT with PreNorm (marked in blue) and our V-DNT (marked in orange), where V-DNT denotes the vision variant of our DNT. We observe that the amplitude of the gradients in our V-DNT are typically quite small and well concentrated; whereas the gradient distributions of the standard ViT have a long tail.
FIGURE 1:
architectures. Unlike Convolutional Neural Networks (CNNs) (LeCun al. 1998, He et al. 2016) that are trained on tasks like ImageNet, where the entries of the gradients are typically small and well-concentrated, the gradients of Transformer typically exhibit heavy-tailed distributions, as shown in blue in Figure | This heavy-tailed distribution means that the amplitudes of the gradient entries span a wide range and thus is hard to keep step with each others when updating weights. Thus, Adam uses a normalized term between the first-order term (i.e. gradients) and the square-root of the second-order term. Owing to the normalization, Adam is robust to the heavy-tail distribution of the gradients. This explains why Adam has become the standard optimizer for training Transformer in practice. On contrary, mSGD directly uses the first-order gradient with momentum to update the weights, and thus the weights updating suffers from the difficulty to keep pace with each others. Consequently, an interesting question turns out to be: can we help mSGD to relieve the issue of heavy-tail gradients in training Transformers? And how?
To mitigate the issue of heavy-tail gradients training Transformers with mSGD, we propose to add adjust the positions of normalization operations in Transformers. motivated by analyzing the Jacobian matrix of different modules. Roughly speaking, we proposed to use the properly positioned normalization operator to amend the Jacobian matrix of 8i the activations, or the joint influence of both weights and activations.
8y less affected by the weights,
As illustrated in orange in Figure D we observe that our designed architecture, a Deeply Normalized Transformer (termed as DNT)), exhibits more concentrated gradient distribution than its counterpart which has a heavy-tailed distribution. In this paper, we provide not only theoretical justification for the properly positioned normalization operator our DNT, but also empirical evaluations further validate that our DNT outperforms its counterparts, i.e., ViT and GPT, on ImageNet classification and Open WebText tasks. Since that the distributions of the gradients of DNT are more concentrated, training it with the vanilla mSGD can yield performance on par with that by an Adam optimizer.
To the best of our knowledge, this is the first work show that using a vanilla mSGD can train Trans- former to achieve performance comparable to that of using Adam provided that the Transformer architecture properly modified to mitigate the issue of heavy-tail gradients.
PRELIMINARIES
This section will provide some preliminaries on nigh-dimensional random vectors, which enjoy many nice properties that are different from their ciw-dimensional counterparts. Two simple yet useful theorems are introduced below. Proofs can be found in Lemma 3.2.4 of (Vershynin) 2018
2

---

Published as a conference paper at ICLR 2026
Theorem 1 (Concentration of norm). Let be an isotropic random vector in R Then, have 2 d. Moreover, if x and are two independent isotropic random vectors, then [xx, y)² = d.
Theorem (Almost orthogonality of high-dimensional independent vectors). Let us normalize the random vectors x and y in Theorem 1, setting a ||æl|| and y 2! in high ligh-dimensional space, the independent and isotropic random vectors and y. tend to be almost orthogonal,
Theorem establishes that | Vd, lyy vd and (x, v) = v with high probability, which implies that the cosine the angle 9 between two random vectors x and y satisfies cos(f)) = va Theorem 2 implies that in high-dimensional space (i.e., d is very large), two random vectors are almost orthogonal. Thus, given x = x + y where a and y are two high-dimensional random vectors,
we have || 2 x V 1 1
Jacobian of normalization. Normalization (Toffe & Szegedy, 2015 Ba et al. 2016 Zhang & Sennrich 2019) is a widely used technique in deep learning. For example, LayerNorm (Ba et al. 2016 is defined as LN) =
Vily
+ B, and y
1 æ; where e is
smoothing factor, y and are two learnable Fd vectors which are usually initialized
and
Most recently, some recent LLMs Touvrro et al. 2023 Chowdhery et al. 2023)] Team 20233 Liu et al. 2024) uses RMSNorm (Zhang & Sennrich 2019) to replace LayerNorm, where RMSNorm is
defined RMSN(() y
vdx
without the centering term and the bias term. The Jacobian
matrix of RMSNorm with respect to æ is calculated as follows:
(R
Vd
diag(y) (I
22
ax
| +e
a
We RMSNorm default normalization technique when mentioning of normalization, but analysis can be generalized to the other normalization techniques. Here we use a numerator layout for all our gradients derivation throughout this paper.
Stochastic Gradient Descent (SGD). SGD (Robbins & Monro 1951) is a classical and fundamental optimization algorithm in machine learning for training models by minimizing their cost functions. However, the vanilla SGD often suffers from slow convergence, especially in complex optimization landscapes with ravines, saddle points, local minima. To address these limitations, momentum
SGD [Nesterov] 1983, 2013, Sutskever et al. 2013) was introduced as extension of the basic SGD algorithm. Momentum SGD Nesterov 1983 2013 Sutskever et (203)) introduces velocity term m that accumulates gradients over time, i.e., m4+1 μim,, V [(w,), Wt+1 we 0, where μ E [0, 1) is the momentum coefficien that determines how much of the previous velocity is retained and og is the learning rate for the time step Unlike in the vanilla SGD, mSGD allows the optimization to build up a "momentum" in the direction of persistent gradient descent, which can effectively dampen the oscillations in high-curvature directions.
3 THEORETICAL JUSTIFICATION ON WHY DNT CAN BE TRAINED WITH MOMENTUM SGD
3.1 PROBLEM 1: WHAT IS THE ROOT CAUSE HEAVY-TAIL DISTRIBUTION OF GRADIENTS?
Previous works (Zhang et al.] 2020 Simsekli et al. 2019) have pointed out that heavy-tailed distribution of the stochastic gradients is a root cause of SGD's poor performance. Here, we will investigate this issue by analyzing the backpropagation of Transformers.
Suppose that l++ f(x') and we have obtained calculate dz!
FT backpropagation process, then we f
ac using a numerator layout as dul an ox! where 0z! called as the Jacobian ac
Oc
matrix. Having had ac for any forward layer with x² W' we compute aw? as aha!
Oc Oc 1
I-TT Oc 0x I-IT
(1)
Owl
dr
dzl+1 dx
Typical values for y range from 0.9 0.99. In default, for all experiments, l 0.90.
3

---

Published as a conference paper at ICLR 2026
From Equation 1 observe that the heavy-tail problem in gradients is indeed closely related to 0w" when its singular values are very diverse. The Jacobian matrix suffers from highly diverse singular values for several reasons: 1) the weight matrix contains very diverse singular values; 2) the activations span widely, leading to Jacobians with very uneven singular value distributions. When matrix has wide range of singular values (i.e., very large condition number), it means that the transformation stretches the input very differently along different directions. Thus, during the backpropagation phase, the heavy- satiy-taillaa issue in the gradients occurs. Therefore, a reasonable solution to relieve the heavy-tail issue is to constrain the uneven singular values of the Jacobian matrix by controlling the weight matrix and the activations. This is the basic idea in this paper.
the Jacobian matrix 0z*+1
3.2 PROBLEM 2: MITIGATE THE HEAVY-TAIL GRADIENT PROBLEM BY ANALYZING THE JACOBIAN MATRIX
In this subsection, we will describe how we use different normalizations, adding or adjusting the position of the normalizations- constrain the Jacobian matrix to relieve the heavy-tail gradient issue. Note that we do not claim that we discover any new normalization methods, instead, we provide our understanding to (Loshchilov et al.] 2025; Zhu et al.] 2025)] Qi et al., 2025b] (023) for more discussions about the normalizations. We will use red, green, blue, purple, magenta to associate with "InputNorm", "PreNorm", "MidNorm", "PostNorm" and "QKNorm", individually.
how each normalization affects the Jacobian matrix. We refer the readers
WERPE
InputNorm
Blocks
PreNorm
SAFFN
SAPFAN
MiaNarm
(a) InputNorm
(b) PreNorm
(c) MidNorm
W
QKNNrrm
OKNorm
SA/FFN
PoctNorm
W.
MaiMuil
(d) PostNorm
(e) QKNorm
FIGURE 2: Five different normalization methods. The only difference between them is the position of normalization. In (A), "WE/PE" indicates Word Embedding (WE) and Patch Embedding (PE).
3.2.1 INPUTNORM
Definition of InputNorm. InputNorm in Transformer defined as the normalization that is applied after the first word embedding in NLP or the first patch embedding in vision Transformer. As shown Figure (d), InputNorm is defined as
a Input Norm(hh. where h = Embeddingg()),
(2)
where is the input and Embedding(() denotes the word embedding or the patch embedding. For standard residual block in Transformer, we have that:
2+1 =aa x ++(x(() z'- f(='=1) f(z') x f(x²)) )+ +(.'))....4 +....+(x'-- + f(x').
Each x' will be the input into some modules. such as normalization, self-attention and feed-forward layers. The Jacobian matrices of some modules, such as LayerNorm and the dot- dot-product self-attention, are sensitive to the norm of the input.
Under the assumption that random vectors are almost orthogonal in high dimension, we have that:
|
| + 1 [l(()))
Proposition 1 (Effect of the norm of the input embedding the gradients). In dimensional settings when all parameters and activations are dimessimmss,,nal, if the term | [2 is very large, will lead gradient vanishing problem in all subsequent layers, provided that InputNorm is not used.
4

---

Published as a conference paper at ICLR 2026
It means that if the term ||>'12 is large, then the norm ||al+||| in each layer will also be large. If || 'Il2 is the imput into a normalization layer, according to the Jacobian equation normalization OHMSN(a"*)) 0a'
the gradient flow each layer will be significantly affected by the term ||=0 Thus, need to constrain | before it is used as the input into the following layer.
VIa=++113+ diag(7)
va
Remark 1. The ||<0|| has a large inftuence of the gradient flow of the subsequent layers. If is very large, it will lead to gradient vanishing. and if it is very small, it may lead to gradient exploding. Meanwhile, the network is also sensitive the changes of ||x011
3.2.2 PRENORM
Definition of PreNorm. PreNorm in Transformer is defined as the normalization applied before the self-attention (or the feed-forward) module. As shown in Figure PreNorm is defined as
Y
Self-Attention(XX) where X'
[a' s.n), atj PreNorm(zz;).
single-head self-attention is defined as
Y= WAX
where A = softmax(
P Ruxa is called as the attention matrix, P = X W W,X in which Vdg
P is called as the logit, and X E Rdxn Wa E Rdaxd W. E Rd,xd W, E Rd.xd Vd,
Herein, our goal is to calculate @vec( x By vectorization of Y W.XAA A, we have
@vec(Y)) = (A @ W.)dvec(X + (In @ W. (.xX)ovcc(A)
Putting together all these terms, we have that
@vec (A & (W) X
Vd (XX W. W, 0 I.CC+ &
X Wg (W))
Övec X)
(5)
where @ denotes the Kronecker product, Cdn is the commutation matrix, and
J blockdiag (diag( (A.,1) A diag (A. A.
The detailed derivation process can also be found in prior work et et al., 2025a). Nevertheless, we note that rather than analyzing W W. in the self-attention module, here analyze the influence of X.
According to the Jacobian matrix in Equation 5 we have the following proposition.
Proposition 2 (PreNorm can stabilize the gradient in self-attention module). If. for each column æ; have x; Q.j where O; R is normalization scalar, according Equation 4 with the same Wg. Wx and W.. and let Y = Self-Attentionn X) and Y Self-Attention (X Then have that the Jacobian matrices at X and X' are the same, i.e., dvecl Bvec(Y')
Fver X
Zivert X
According to Proposition] 2 have the following remark.
Remark 2. PreNorm will guarantee that the norms of the column vectors of X (which are the input the self-attention layers) are relatively stable range. According Equation we see that if these norms are relatively stable, then the Jacobian matrix will also be stable. Meanwhile, since the gradients with respect W,, WE and W, are directly involving X, a stable X will guarantee that the gradients with respect to W W. and W. are relatively stable.
3.2.3 MIDNORM
Definition of MidNorm. MidNorm in Transformer is defined as the normalization applied after the self-attention feed-forward module and meanwhile before the residual shorteut, As shown in Figure 2 (b), MidNorm is defined as
y MidNorm(z): where
FFN(a; Wi. W2) W2 ReLU(W,,)))
(6)

---

Published as a conference paper at ICLR 2026
In self-attention, the matrices W. and W. can be viewed similar function as W1 and W2 in FFN. we only use a single-head attention, then we have 2 = W. W.x
The Jacobian matrix of an FFN can be computed as: J2(x) W2 diag (1 (W,æ 0) The Jacobian matrix of RMSNorm
Oz
OF W.)
Bre
di
layer is
dy
va diag(r) |=
||=1|] The Jacobian matrix of an FFN followed by an RMSNorm is:
Dy
Oy 0z vddiag(>) (I [I-
zz
diai (1 (Wi > O))W
(7)
dæ
Oz dx
W2 ReLU(
[W, [Wa)ll in high-dimensional settings when W,, W2 and x are high-dimensional and random, the singular values of M will be only related the shape of W1 and W2. and will be independent to the magnitude of W1 and W2.
Proposition (Effect of MidNorm). Let M W2 diag (I(W,>0)):W
According to Proposition 3 we have the following remark.
Remark 3. MidNorm can effectively guarantee that the norms of Wi, W2. W., and W. will not affect the Jacobian matrix shown in Equation 7 It means that even the magnitudes of these weight matrices are very large, it will not magnify the gradients due to the normalization.
3.2.4 POSTNORM
Definition of PostNorm. PostNorm in Transformer is defined as the normalization applied after the residual block. As shown in Figure (d), PostNorm is defined as
x I+ PostNormmz"* where z = x' + f(æ'; W
(8)
Proposition 4 (PostNorm is sensitive to the vector norm of activation). If g1+1 in Equation is very large, then it will significantly decrease the gradient.
Proof. From Equation 8 we have Ow
Va I| I ll2
diag(n)
2 1 1 ). If is very
0z!+1
0n large, according to 0= FT 0zt-r, we have that the gradient decreased.
ag will be significantly
In a classical Transformer /Vawwani et al. 2017). if fa: W,, W2) = W. ReLU(W, w)), along with the training process, the largest singular values of W and W iee, a (W1) and a (W2) will usually become too large (e.g., around 1000). In this way, (= W ][ will be very large. As a sequence, z/+1 in Equation 8 will be very large. Therefore. have that PostNorm under this circumstance will lead to gradient vanishing.
Remark We need to be very careful when using PostNorm. That is, we must ensure that the norm of the input vector PostNorm is within reasonable range; otherwise the network is likely to cause a gradient vanishing when z being very large or a gradient exploding z when z' being very small.
3.2.5 QKNORM
Definition of QKNorm. QKNorm (Henry et al. 2020) Transformer is defined as the normalization applied on the queries and keys in the self-attention block. As shown in Figure (e), the self-attention with QKNorm (Dehghani et al. 2023) defined as
(
P Ph, P Q K' Vdn
Y W, XA, where A' softmax
where of and ky are the i-th column and the j-th column in Q' and K', respectively, in which
/da W /da diag(T))
a,
qi (W.Noorm,,,,.)
I W,
|W,
(10)
yd, Wayy diag(m))
We
k/ QKNorr(W..,))=
W
[

---

Published as a conference paper at ICLR 2026
where d, is the head dimension.
To facilitate the derivation, we denote Q = WgX and = W&X as before, and use qi and kj to denote the i-th column and the j-th column in Q and K, respectively. Thus, have that df = QKNorm( (q:) and kj = QKNorm (k;)) Moreover, we have that Pij = q'; k',, where Ps is
scalar, and the gradient is computed as follows:
iP
Ogk
ok
d'
dæ
dx
dæ
g/g W
W. Waa;
= <d didglngaa jk (I
WWa::ll2 + dd (iag)()))
2
(1I)
Proposition 5 (Effect of QKNorm). In figh-dimensional settings, i.e., when all W We and 8
is independent of the are high-ddieessonal and random, in Equation 11 the gradient term of Oz magnitudes Wa and Ws.
According to the Proposition, we have the following remark.
Remark S. QKNorm can mitigate the joint effect of W Wい on the gradient of the self-attention g Wy has been revealed to be root reason
layer. The fast increase of the singular values of W T leading to model crash. Our analysis shows that QKNorm effectively mitigate the reason to cause model crash brought by W Wg-
Though QKNorm can mitigate the problem brought by W. WR. it cannot fully replace the role of PreNorm, because PreNorm can jointly deal with the problem of Wy. Wi and W., and the gradient of W. is also affected by the value of X.
3.3
DNT: A TRANSFORMER THAT CAN RELIEVE THE ISSUE OF HEAVY-TAIL GRADIENTS
WE/PE
InputNorm
PreNorm
SA with GiKNorm
MidNorm
My PreNorm
FFN
MidNorm
FIGURE DNT architecture. The second PreNorm marked with dashed and rounded corners optional. By default, we do not use the second PreNorm.
Having analyzed the effects of different normalizations, we use four types of normalizations, including InputNorm, PreNorm, MidNorm and QKNomm, except for PostNorm. The reason why do not use PostNorm is that it may cause some training instability issue. For clarity, we illustrate the architecture of our DNT in Figure
Our DNT model commences with Word Embeddings (WE) or Patch Encodings (PE). Then, the initial representations undergo an InputNorm processing and the normalized embeddings are used for the subsequent operations. The core transformer module consists of N blocks. In each block, a PreNorm applied first, followed by self-attention which is augmented with query-key normalization (i.e., QKNorm). Then, MidNorm is applied to process the attention outputs before integrating with the residual connection. In the second sub-block. before the ceed-forward network (FFN), the PreNorm is optional; and after FFN, a MidNorm is applied before integrating with the residual connection. This entire structure is replicated N times to form the whole architecture.
We visualize the effects of each normalization in our DNT in Figure According to the analysis mentioned above, we summarize the advantages of our DNT as following: a) the magnitude of x will significantly affect the gradient of each layer in the Transformer, but we introduce InputNorm to resolve the influence of 0;; b) PreNorm can constrain the norm of each column in activations X in each timestep, and thus amend the Jacobian matrix of self-attention to not be significantly affected by the magnitude of X; c) MidNorm will amend the Jacobian matrix of each sub-block (i.e., the sub-block with self-attention and the sub-block with FFN) in our DNT not be affected by
7

---

Published as a conference paper at ICLR 2026
x .x.....x'
x x .....! X
InputNorm
W.. Wv. W.. W.. W1. W2
PreNorm
W. W.. W.
(a) Influence InputNorm.
(b) Influence of PreNorm.
x P'.'''.. , x
Q
MidNorm
W.. Wo. WT.
OKNorm
W. W,
(c) Influence of MidNorm.
(d) Influence of QKNorm.
FIGURE 4: Influence of different normalizations.
For instance, InputNorm stabilizes x4. ,
W, Wt. W., W., Wi, W2 by constraining X
the magnitude of W1, W2. W, and W.; d) QKNorm can relieve or even remove the influence of the magnitude W and W on the Jacobian matrix of self-attention, and thus reduce the risk of problems, such as rank collapse (Noci et al. 2022), entropy collapse (Zhai et al. (023), spectral energy concentration Qi al. 2025a caused by W Wk-
In DNT, we use four different types of normalizations. We observe that nGPT (Loshchilov et al. 2025) also uses some of the normalizations mentioned above. Here, we would like to emphasize the differences between DNT and nGPT that: a) DNT provides theoretical justifications for each normalization different position; b) DNT uses InputNorm rather than PostNorm, whereas nGPT many PostNorms but not InputNorm; e) nGPT normalizes the activations the weights into spheres, whereas DNT only normalizes the activations but does not requires activations on spheres.
We term model Deeply Normalized Transformer (DNT for short), because it is designed by properly adding or positioning normalization operators in the conventional Transformer. For vision problem, term it as V-DNT, and for language problem, term it as L-DNT. The key difference between V-DNT and L-DNT is that V-DNT uses patch embedding, but L-DNT uses word embedding and mask for attention computation.
Open Vill Loos
ImageMiet Va on Accuracy
3.30
325
AdaniM--AM (NT-5mml
AdamiM
50.0
mSs2WW Small
60.0
g 3.10
Accuracy
3304
20.0
234
AsamtiVM T- T-aage
2.9
ONT
MScBI y CNT Large
175k
80
150
Step
Esoch
FIGURE 5: Validation loss (Left) on Open WebText and recognition accuracy (Right) on ImageNet. We compare L-DNT-Small (124M) to GPT2-Small (124M), and V-DNT-Large (307M) ViT- Large (307M). By effectively relieving the heavy-tail gradient issue, our DNT network trained with naive mSGDW can achieve competitive performance to AdamW (Val loss 2.849 2.863 OpenWeeText, Acc 8/. 5% vs. 82. 1% on ImageNet). However, in classical Transformer with PreNorm, the performance of mSGDW under-performs AdamW significantly (Val loss 2. 2.906 2.867 on OpenWeebeext, Ace 78.2% vs 81.7% on ImageNet). See Appendix C for the training parameters.
4 EXPERIMENTS
We conduct experiments with two popular Transformer architectures: Vision Transformer (ViT) and Generative Pretrained Transformer (GPT). Our implementation leverages established repositories:
8

---

Published as a conference paper at ICLR 2026
timm (Wightman 2019) for VIT and nanoGPT (Karpathy 2022) for GPT models. For experiments with VIT. we utilized two model scales: ViT-Large (307M parameters) and ViT-Huge (632M),
following the configurations described in (Dosovitskiy al. 2020). The data augmentation strategy aligns with (Xie et al. 2024) to ensure fair comparison with previously reported results. For experiments with GPT, employed the nanoGPT implementation focusing GPT2-Small (124M)
and GPTe-Large (774M) variants due to computational constraints. The results of our baselines align with previous work, including Sophia (Liu et al. 2023b) on OpenWebText and MAE (He al. (022) on ImageNet. Training was conducted using Py lorch (Paszke et al. 2019) with 6 precision on A800 GPUs, employing a cosine learning rate schedule.
4.1 VISUALIZATION OF GRADIENTS OF DNT AND TRANSFORMER WITH PRENORM
To visually compare the standard Transformer with our DNT, we visualize the gradients with respect to different weights, including W.. Wk. W., W. W1 and W2- We chose the early checkpoints of the model training for visualization, but we found that the same phenomenon also occurs in the middle and later stages of the model training. The visualization results are shown in Figure We can see that our DNT network can well relieve the issue of heavy tail distribution of the gradients. For instance, in the Transformer, the amplitude of the entries in the gradients almost spreads in the range [0, 10 10= whereas the amplitude of the entries in the gradients in our DNT concentrates around [0, 10 10-9).
4.2 MSGDW ACHIEVES PERFORMANCE ON PAR WITH ADAMW.
We also give a quantitative comparison of the standard Transformer and our DNT both trained with Adam and mSGD Open WeenWetText and ImageNet in Table We can see that training our DNT via mSGDW achieves a similar result to that trained with Adam W. We can also see that using mSGDW to train DNT model greatly outperforms the performance of using mSGDW to train the standard Transformer. In Figure 5, we display the validation loss on Open WebText and the training accuracy on ImageNet along with the training process. Note that we did not the learning rate too much. We just followed the learning rate settings in the previous works Karpathy (2022); Liu et al. (2023b). We believe tuning learning rate will bring in some differences. But overall, DNT network can enable mSGDW compete with AdamW.
TABLE 1: Quantitative comparison of standard ViT/GPT2 and V-DNT/L-DNT trained with AdamW and mSGDW on Open WebTeet and ImageNet. Results on ImageNet are based on 150 epochs, and results Open Web'eet are based 200K steps.
ImageNet (Acc. 1)
Open Web (Val Loss. }
Optimizer
Types Model
307M
632M
124M
774M
1436M
AdamW
VITIGPT2
81.7
80.8
2.867
2,492
2.435
AdamW
V-DNT/L-DNT
82.1
81.9
2.863
2.481
2.396
mSGDW
ViTIGPT2
78.2
73.5
2,906
2,544
2.472
mSGDW
V-DNT/L-DNT
81.5
81.2
2.849
2.503
2.408
4.3 HOW MUCH MEMORY MSGDW SAVES RATHER THAN ADAMW?
We compare the memory usage by mSGDW and AdamW. The results are shown in Table 2 Theoreti- cally, we can calculate that the memory taken by Adam (only the optimizer part) is 1T,5GB and the memory costed by mSGDW (only the optimizer part) is 5.7GB. In the experiment, ANT+A (model plus optimizer) costs 67GB, and DNT+mSGDW (model plus optimizer) costs 61GB. Using mSGDW instead of AdamW 1. 1.4B model can save around 6GB memory.
obtained

---

Published as a conference paper at ICLR 2026
TABLE Comparision of GPU memory used by mSGDW and AdamW trained on 1.4B DNT model. Adam means the network usage and the optimizer usage of GPU memory. i denotes the theoretically calculated values, and & denotes the observed values practice.
AdamW
mSGDW
DNT+AdamW
DNT+mSGDW
Memory 11.5t GB
5.7 GB
< 67 GB
61 GB
4.4 ABLATION STUDY
Comparison of different normalization methods. We conduct ablation study of five different normalization methods. Figure in the Dnjeutttiiiluusttratess these five different network settings. Let us brief introduce these five settings below: 1) Setting Standard transformer with PreNorm, for which we abbreviate as $1; 2) Setting 2: S1 + QKNorm; 3) Setting 3: $2 InputNorm; 4) Setting 4: 2PreNorms MidNorm QKNorm + InputNorm: 5) Setting 5: only PreNorm before self-attention + MidNorm + QKNorm + InputNorm. We use momentum mSGDW for all the training in this subsection. All models are trained with the same hyper- ppper-parameters. The results are shown in Figure We have the following observations.
On WebTee the original PreNorm setting (S1) shows the worst performance. The perfor- mance of S2 is similar to that of S1. The setting of S3 with input obtained better performance. Finally, S4 and S5 obtained the best performance. Meanwhile, the performance of S4 and S5 is similar.
On ImageNet, the original PreNorm setting (SI) is significantly worse than the other four settings. S3 achieves the best performance, and S4 and S5 also obtain excellent performance.
OpeniWebT Vill Loos
IrageMitt Va ldst on Arcuzzcy
3:1
51
50.0
51
3.10
55
50.0
3.00
Rccuracy
40.0
299
30.0
2.94
Step
Esoch
FIGURE 6: Ablation study of different settings using mSGD optimizer on ImageNet and OpenWeb Text. Left side shows accuracy curve of Huge vision model (632M) on ImageNet. Right side shows the validation loss of language model (124M) on OebTeet.ext.
5 CONCLUSION
We introduced novel architecture, named Deeply Normalized Transformer (DNT), which enables efficient training with the vanilla momentum SGDW (mSGDW), achieving performance on par with Adam W-optimized Transformers. Unlike traditional approaches that rely on sophisticated optimizers to address the challenges of heavy-tailed gradient distributions, our DNT properly integrated normal- ization techniques into the architecture of Transformer to effectively regulate the Jacobian matrices of each block, the contributions of the weights, activations, and their interactions, and thus make the gradient distribution concentrated. Our findings demonstrated that properly designed architecture can make a simple optimizers like mSGDW just as effective as sophisticated ones. This opened new opportunities for creating efficient, scalable, and accessible Transformer models.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
Zhouchen Lin is supported by the National Natural Science Foundation of China (NSFC) under Grant No. 62276004. Chun-Guang Li is supported by NSFC under Grant No. 62576048.
ETHICS STATEMENT
This work presents novel deeply normalized Transformer architecture. It does involve human subjects and poses no potential risks. The study is free from conflicts of interest, sponsorship issues, or concerns related to discrimination, bias, fairness. All data used adhere legal and ethical standards, and privacy and security considerations have been addressed. Our work fully adheres to research integrity principles, and no ethical concerns have arisen during the course of this study.
REPRODUCIBILITY STATEMENT
To facilitate reproducibility, provide comprehensive experimental details in the Appendices, including theoretical proofs, experimental settings, and configurations. Our implementation builds on nanoGPT and timm. The ImageNet and Open WebText datasets publicly available.
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint vrXivv1607.064500 2016.
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Arizzadenesheli, and Animashree Anandkumar. signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning, pp. 560-569. PMLR, 2018.
Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 33:1877-1901, 2020.
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems,
2023.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. In International Conference on Machine Learning, pp. 7480-7512. PMLR, 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image worth 16x16 words: Transformers for image recognition at scale. International Conference on Learning Representations, 2020.
Abhimanyu Dubey, Abhinay Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The llama herd of models. arXiv preprini arXiv:2407.21788. 2024.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal machine learning research, 12(7), 2011.
11

---

Published as a conference paper at ICLR 2026
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 16000-16009, 2022.
Alex Henry, Prudhvi Raj Dachapally. Shubham Shantaram Pawar, and Yuxuan Chen. Query-key normalization for transformers. In Findings of the Association for Computational Linguistics: EMNIP 2020, pp. 2246-4255 2020.
Hinton. Rmsprop: divide the gradient by running average of its recent magnitude. COURSERA:
Neural Networks for Machine Learning, 13, 2012.
Roger A Horn and Charles R Johnson. Matrix analysis. Cambridge university press, 2012.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. PMLR, 2015.
Rie Johnson and Tong Zhang, Accelerating stochastic gradient descent using predictive variance reduction. Advances in neural information processing systems, 26, 2013.
Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https //kellerjordan, qithub hctp:/sss/oostsooo//////oooooo
Andrej Karpathy, NanoGPT.
2022.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 141.6994 2014.
Yann LeCun, I. eon Bottou, Yoshua Bengio, al. Gradient-based learning applied to document recognition. PROCEEDINGS OF THE IEEE, pp. 1, 1998.
Junnan Li. Dongxu Li, Caiming Xiong, and Steven Hoi. Blip: Bootstrapping language- image pre- training for unified vision-language understanding and generation. In International conference on PMLR, 2022.
machine learning, pp.
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437, 2024.
Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. Advances in neural information processing systems, 36:34892-3499 -39916, 2023a.
Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. arXiv preprint arXiv:2305. 14342, 2023b.
7e Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, al. Swin transformer v2: Scaling up capacity and resolution. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12009-12019, 2022.
Ilya Loshchilov and Frank Hutter. Fixing weight decay regularization in adam. In International Conference on Learning Representations, 2019.
Ilya Loshchilov, Cheng-Ping Hsieh, Simeng Sun, and Boris Ginsburg. ngpt: Normalized transformer with representation learning on the hypersphere. The Thirteenth International Conference on Learning Representations, 2025.
Yurii Nesterov. method for unconstrained convex minimization problem with the rate of conver- gence (1/k* 2). In Doklady ussr, volume 269, pp. 543-547, 1983.
12

---

Published as a conference paper at ICLR 2026
Yurii Nesterov. Introductory lectures convex optimization: A basic course, volume 87. Springer Science & Business Media, 2013.
Yurri Nesterov. Introductory lectures on convex programming volume i: Basic course. 1998.
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 35:27198-22 12221199-2211,,, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer. James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, perhoperformance deep learning library. Advances in neural information processing systems, 32, 2019.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4195 4205, 2023.
Xianbiao Qi, Jianan Wang, Yihao Chen, Yukai Shi, and Lei Zhang. Lipsformer: Introducing lipschitz continuity to vision transformers. In The Eleventh International Conference Learning Representations, 2023.
Xianbiao Qi, Yelin He, Jiaquan Ye, Chun-Guang Li, Bojia Zi, Xili Dai, Qin Zou, and Rong Xiao. Taming transformer without using learning rate warmup. The Thirteenth International Conference on Learning Representations, 2025a.
Xianbiao Qi, Jiaquan Ye. Yelin He, Chun-Guang Li, Bojia Zi, Xili Dai, Qin Zou, and Rong Xiao. Stable-transfrmmer: Towards a stable transformer training, 2025b. URL https
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, al. Improving language understanding by generative pre-training. 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International conference on machine learning, pp. 8821-8831. Pmir, 2021.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pp. 400-407, 1951.
Noam Shazeer and Mitchell Stern Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4604 PMLR, 2018.
Umut Simsekli, Levent Sagun, and Mert Gurbuzbalaban. A tail-index analysis of stochastic gradient noise in deep neural networks. In International Conference on Machine Learning, pp. 5827-5837. PMLR, 2019.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In International conference on machine learning, pp. 1139-1147.
PMLR, 2013.
Terence Tao. Topics in random matrix theory, volume 132. American Mathematical Soc., 2012.
Qwen Team. Qwen technical report. arXiv preprint arXiv:2309.16609 2023.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Ilama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.00288 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017
13

---

Published as a conference paper at ICLR 2026
Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018.
Jing Wang and Anna Choromanska. A survey of optimization methods for training di models:
Theoretical perspective on convergence and generalization. arXiv preprint arXiv:2501.14458. 2025.
Ross
Wightman. Pytorch image models. 2019.
https://github.c http:://gitthu..commmw com rwttmaangtttmn///
Xingyu Xie, Pan Zhou, Huan Li, Zhouchen Lin, and Shuicheng Yan. Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models. IEEE Transactions Pattern Analysis and Machine Intelligence, 2024.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv:1708. 038888 2017.
Huizhuo Yuan, Yifeng Liu, Shuang Wu, Xun Zhou, and Quanquan Gu. Mars: Unleashing the power of variance reduction for training large models. arXiv preprint arXiv:2411, (tt::2411.1048,, 2024.
Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Joshua M Susskind. Stabilizing transformer training by preventing attention entropy collapse. In International Conference Machine Learning, pp. 40770 77770-40803. PMLR, 2023.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 13:13-33-5,393 2020.
Jiachen Zhu, Xinlei Chen, Kaiming He, Yann LeCun, and Zhuang Liu. Transformers without normalization. arXiv preprint arXiv:2503. /0622, 2025.
14

---

Published as a conference paper at ICLR 2026
A
LLM USAGE
During the preparation of this work, the authors used ChatGPT for language editing and to assist in the creation of TikZ diagrams. The models were firstly prompted with draft text rough sketches to improve clarity and fluency of language and to generate code snippets for figures. Then, we carefully reviewed and modified all generated content. The core ideas, research, analysis, and conclusions remain entirely the work of the authors, and the LLMs were not involved in any intellectual contribution.
B PROOF OF PROPOSITION 2, 3 AND
Proposition and Proposition 4 very easy prove, we have given brief proof in the main body. Therefore, in the appendix, we only provide the proof of Proposition 2, 3 and 5.
B.1 PROOF OF PROPOSITION 2 ON PRENORM
Proof. If for each column Z; we have 2 ay with the same W, We and W,, v. we have that Y Y' given Y Soaf-Attention(2 () and Y Sclf-Attention(X X') because we will obtain the same input the self-attention after the PreNorm.
That is, after the PreNorm, we have that PreNorm () PreNorm (X)),, according to Equation 5 thun we have that: @ven( @vec(Y")
X Xv)
A single-head self-attention can be defined as
== W.XA,
where A = softmax P dd
and P x W W.X, in which A is called the attention matrix
and P is called the logit, and A Ruxn X E Rdx= W. E Rd.xd, Here, our goal is to Vd.
calculate åvee (Y) (x) By vectorization of W XA, we have
Ōvec(Y)= (AT & W.) W.)ovec( X) (In & (o)e)oe)(AA).
Bringing in all the terms, we get the following formula
Övec( Y) (A (W.)(L
ddg (IXTWx WWKTWoO 8 Cn + 8 & Tw)).
J
Ōvec(X)
where Can the commutation matrix[and
J blockdiag (diadg A.,)) (diag(A..1) A., A. ........ diag(A A..)) A.n AT,.).))
Note that J is function of A, and a function of X associated with softmax function. Obviously, @veo(Y) is a high-order function of where the softmax makes the analysis more complicated. [vee[X]
Here, we drop the softmax operation and give an analysis of the Jacobian matrix of the linear attention module, i.e., A
and P = X W, W.X. For the linear attention, the Jacobian matrix is:
dy
Ōvec(Y YY Ovec(X)
(A W.) + (L. 2 W, W,X)
(W Tw, I.) (I, 8 SX W.))
ddg
(12)
Obviously, if the norm of each feature vector for each token is large, the magnitude of each element @veet Y
Ovee[ will have large probability be large, and the singular value of @vee(Y)) [ve((
may be magnified
second-orderly by the norm of each column in X.
https://en wik wikipedieiiee org /ttp://aooooo/o///oooooootoooooo
15

---

Published as a conference paper at ICLR 2026
Furthermore, would like to conduct deeper analysis of the gradient of the loss with respect to the weights. the backpropagation, since that we have obtained ōveo we would like to further
analyze avee(W,) (ve) and W.e)
For the weight matrix Wg. we have
OC Oc @vec Y) Ōvec(A) (A) @vec P)
(
@vec(Y) Av) θvec P) @vec Wg)
(13)
dC
(In @ W,X) ((wx)) 1 ((xx))&xx))
=
@vec(Y)
For the weight matrix WR. we have
UC dC @vec(Y) Övec(A) Övec(P)
@vec (Wx)
@vec(Y) dvec(A) δvec (P) Ovec(WW))
(14)
dC @vec(Y)
J
(In @ W.X)
(x @ (W
ddg
For the weight matrix W., we know that vec(Y) = ((XA) T @ 1) vec(W,)) thus we have
dc 0C dvec(Y)
Ōvec W.)
Yve Ōvec(
(15)
OC
((AT))81)
evec(Y)
We can see that in Equations 13 15 the gradients of loss with respect to W., We and W, are all related X. After PreNorm, in relatively stable range. thus the gradients of W, go We and Wo are all in a stable range.
In conclusion, the normalization of X can help stablize the gradient of the loss function with respect to Wg. W and We, and meanwhile help make @vec( X) more stable.
ovec(Y)
B.2 PROOF OF PROPOSITION ON MIDNORM
Proof. Starting with the definition W, M where y Wa and W E Rmxx, let's derive the relationship between singular values:
For random matrix W with i.idd entries (i.e., zero mean, variance ofy) and random vector with i..d. entries (i.e., zero mean, variance o²)):
E| y
2
E
W
E|x w Wæ].
Using the trace property:
E E|
W W σ] = E[tr(x:" w Wx)) = E Wr x.. WT).
With and W independent, and E Exxx"] = o In:
Eltr(Wara' w -3 J, E. [ir W w [u(
For W with i.i.d. entries, E|u W W E||l w | Therefore, we have:
|| 0²
oy
m
ow 0
16

---

Published as a conference paper at ICLR 2026
According to the measure concentration property, we have that ||yll2 concentrates around its expecta- tion with high probability:
| s Bllbll n ow o
By taking the square root, we have:
1yll vm n ow or with high probability.
Considering the SVD of W, say W UEVT, where E contains singular values or o, (W). We have that the singular values of W1 are:
W
0 (
(W = oi
||yl|2
13112
Substituting it into our concentration result, we have that:
(W)
(w1) vm n ow oz
For large random matrices with i.i.d. entries, the random matrix theory (Hor & Johnson, 2012 2012) tells that the largest singular value follows:
Tao
01 (W) = (Vm + Vn)ow.
Substituting this into our result above, we have:
T1 (Wi) ~ (vm Vnlow mm- n ow 02 mm n Oz
vm + /n
(16)
This derivation result in Equation 16 shows that in high dimension, the largest singular value of W, becomes essentially deterministic, depending only on the dimensions of W and the statistical property o. (which is the standard variance of each entry in æ) of the random vector æ. If m then we have that 01 (W()
Vm.oz"
B.3 PROOF OF PROPOSITION ON QKNORM
Proof. Self-attention with QKNorm (Dehghani al. 2023) defined as:
Y = W.
K', where A' = softmax( P' P' QTKK, and qf and ky are the i-th column and the j-th column Vd, in Q' and K' respectively, and we define
dp W
Wa
d. = RMSN W v,,;) = -4
ddn diag(7a)
[Waa:li
|,,e;ll
dr Waay
Wa
KG RMSN Wax,) = -Y
da diag(Ts)
W ll2
W l
To facilitate derivation, will use Q W W,X X and K W.X before, qi and k, to denote the i-th column and the j-th column in Q and K, respectively. Thus, we denote q. RMSS(gg)) and ky RMSN(Ky).
17

---

Published as a conference paper at ICLR 2026
Therefore, according to the product rule and chain rule, we denote the Jacobian matrix of Y with respect X as follows:
(v)
(A @ W.) (L, @ W. X)
Ovec (A))
(vec(
@vec x)
@vec, (A
Ovec(
(A W.) + (In @ W.
Ovec P
(v)
Ovec(A')
Ōvec @vee[PP)
Qvec(Q))
@vec (P)
Ovec K')
(A" & W.) + (I, 00 W.
Uvec P!)
Ovee (Q") Ovec ()) Ovec K')
Sheeo (vee((X))
(A" W.) + (La W.
Ovec
Uvec P') (vve( (ive (Q) Üvee! PP) (heeo(K') Ovec K)
Ovec(P")
@vec(Q()) (vec(Q) @vee(X)
@veet(K") Ovec (K) (vec(X))
we need to derive out each term in the above equation. To derive out Ovee(
Bvec
Since we know p Q"T K', then we have,
Ōvec PP) Övec K') = Q"T
Similarly, we have
Övec( Ōvec(Q)
(K"T @ 1) CaN.
where CaN is the communication matrix.
Then, have
@vec(Q') blockdiag
iq 0q4
0q'N
J
Övec
0qi 0qp
OqN
K' J Ovec(K') blockdiag
Okl 0k!
0k's
K
(K)
0k: ak2
OkN
where
0qh
ddh diag(?g)
9.9.
0q:
|
|
Moreover, have
@vec( (Q)
1 X W
Ōvec( X)
Ōvec K 1 2 W..
Ōvec X)
Thus. have
0g 0g. 0q:
/ dr diag(?))
9:g.
qig.
W
W, v da diag(?))
Oa:
0q. Ow:
|| 12
|/qa 1
I WW
ok? okl, ōkj
ddn diag(7))
k,k]
k,,k]
W
W₁ v /d diag(Tk)
dx, ok, Oay 1
kk
|| II3
W
2|
In Proposition 3, we have proved that in high-dimensional setting, the singular values of independent to the magnitude of W. Thus, until now, we have proved the Proposition 5.
Further, would like to discuss the Jacobian matrix
Ovee Y Aves after QKNorm. We have that
@vec(Y) (A &W.)+(L J (KK" @ 8) Cax + JJ (T Wh)
(v)
Vah
(17)
It should be noted that:
18

---

Published as a conference paper at ICLR 2026
K' and Q' are two normalized terms that have relatively stable range of values.
the Jacobian matrix of J is relatively independent to the magnitude Q and J K relatively
independent to the magnitude K in a high-dimensional setting.
QKNorm cannot fully replace the value of PreNorm because @vec [ in Equation 17 directly
affected by X.
We note that QKNorm will elliviate the influence of the magnitude of W and W. on the Jacobian
@vec(Y) In the traditional self-attention, @veo[XX
@vec(Y) is largely affected by W W However, after
(vec(
QKNorm, @vec(Y) will only affected by Wa or Wt. independently. rather than their product (i.e., (x)
W. W&.. This is important for the training stability because the singular values of W Wa will
increase extremely fast when both singular values of W,, and We are increasing.
C EXPERIMENTAL DETAILS
TABLE 3: Model configurations, peak learning rate and weight decay for different optimizers.
Adam
mSGDW
Acronym
Size
model
head
depth
LR
WD
LR
WD
L-DNT-Small
124M
768
12
12
6c-4
1.0
lc-4
L-DNT-LLarg
774M
1280
20
36
6e-4
0.1
1.0
1c-4
L-DNT-XL
1436M
1536
24
48
6c-4
0.1
1.0
1c-4
V-DNT-Large
307M
1024
16
24
Ic-3
0.1
0.5
2c-4
V-DNT-Huge
632M
1280
16
32
1e-3
1e-3
We conduct experiments on two popular architectures: Vision Transformer (ViT) and Generative Pre- trained Transformer (GPT). Our implementation leverages established repositories: timm (Wightman) 2019) for ViT and nanoGPT (Karpathy (2022) for GPT models. We utilized five model configurations: L-DNTT-Small (124M parameters), L-DNT--Larg Large (774M parameters), L-DNT-XL (1436M parame- ters), V-DNT-Large (307M parameters), and V-DNT-Huge (632M parameters). Model specifications including hidden dimension (d_ (d_model), number of attention heads (n_head), and network depth detailed in Table 3 The training is conducted using PyTorch (Paszke et al.) 2019) with bfloat 16 precision GPUs, employing cosine learning rate schedule.
All language models are trained on Open WebTet, using GPT-2 tokenizer. The training dataset contains 9B tokens. with validation set of 4.4M tokens, following the train-validation split from nanoGPT. We employed distributed data parallel training with gradient accumulation. All models are trained using bfloat precision. The 124M models are trained on machines with 8 GPUs, 774M models are trained with 16 A800 GPUs, while 1436M models are trained with 32 GPUs. Our global batch sizes for 125M, 770M and 1436M models are 480, 512 and 512 individually. et al. 2023b), they use 480 global batch size for all models. For all language models, we used 2000 steps or learning rate warmup to the maximum learning rate, and then used a cosine learning rate decay. takes around four days train 200K steps for the 1.4B model on 32 GPUs.
Sophia (Liu)
All vision models are trained on ImageNet dataset. We trained each 150 epoches as (Xie et al. 2024). We used learning rate warmup of 60 epochs to the maximum learning rate, and then used cosine learning rate decay.
For experiments, focus comparing AdamW and mSGDW optimizers. The hyperparameters for AdamW are carefully tuned, with 3 = 0.9 and /3 = 0.95, following the dominant configuration in LLM pre-training literature. For weight decay, used 0.1 for AdamW as (Karpathy) 2022 Liu
19

---

Published as a conference paper at ICLR 2026
al. 2023b). We used the recommended learning rate by nanoGPT for Adam in GPT. Since DNT is robust to large learning rate, we use 6 x 10 for all our L-DNT models and 10 for all V-DNT models for AdamW. For mSGDW,
rough grid search for the learning rate, we cannot search fine-grained learning rate and weight decay due to shortage in computation resources. For the momentum in mSGDW, we used default 0.9 for all experiments. We the implementation of mSGDW from timm (Wightman] 2019). This implementation is a decoupled weight decay regularization used in Adam W (Loshchilov Hutter 2019) Note that mSGDW not directly to add a weight decay in the original implementation of mSGD in the official PyTorch, it will have performance problem.
simply
TABLE 4: Training configurations for ViT and V-DNT.
training config
ViT-L/H (224")
ViT-L/H (224")
V-DNT-L/H (2249))
V-DNT-L/H (224²)
optimizer
AdamW
mSGDW
AdamW
mSGDW
learning rate schedule
cosine decay
peak learning rate
le-3
0.5/0.1
le-3
0,5/0.1
minimum learning
1c-8
1e-8
1c-8
1e-8
weight decay
0.1
2c-4/1c-3
0.1
2c-4/1c-3
optimizer momentum
81 32
0.9, 0.99
/
Br 82 0.9, 0.99
# 0.9
warmup epoches
60
60
60
60
weight init
Truncated Xavier
batch size
1024
training epochs
randaugment
(9, 0.5)
mixup
0.8
cutmix
1.0
random erasing
label smoothing
0.1
stochastic depth
0.1/0.5
gradient clip
None
exp. mov. avg. (EMA)
no
https://githui: http:///thhaaaaaeeen////////sssr optim/ py http://00ohaaaoooraaaaaaaaaaaaaa
20

---

Published as a conference paper at ICLR 2026
TABLE 5: Training configurations for GPT and L-DNT.
training config
GPT2-S/L/XL
GPT2-S/L/XL
L-DNT-S/LXXL
L-DNT-S/L/XL
optimizer
AdamW
mSGDW
AdamW
mSGDW
learning rate schedule peak learning rate
cosine decay
6--4/2.5e-441.5-4
1.0
6e-4
1.0
minimum learning rate
60-5
6c-5
6c-5
6c-5
weight decay
0.1
le-4
0.1
le.4
optimizer momentum
B,
0.9, 0.95
# 0.9
81 82
0.9, 0.95
J 0.9
warmup steps
2000
0
2000
0
weight init
Xavier
tokens seen each update
480K/512K/512K
iters
200K
batch size
480/512/512
sequence length
1024
dropout
0.0
6
Truc
gradient clipping
1.0
WE/PE
PreNorm
SA
PreNorm
FFN
N
(a) Setting
SA with OKNorm
PreNorm
N
(b) Setting 2
WEPE
InputNorm
PreNorm
OKNorm
PreNorm
FEN
(c) Setting 3
WE/PE
InputNorm
PreNorm
SA with QKNNrrm
MidNorm
PreNorm
FFN
MidNorm
N
(d) Setting 4
WE/F
InputNorm
ProNorm
SA with OKKNrrm
MidNNorm
MidNorm
N
(e) Setting 5
FIGURE 7: Five differrntt settings of normalizations evaluated in the ablation study. (a): the standard Transformer with PreNorm. (b): (a) + QKNorm. (c): (b) + InputNorm. (d) and (c): two versions of DNTs
21

---

Published as a conference paper at ICLR 2026
D FIVE DIFFERENT NETWORK SETTINGS
In Figure 7 we illustrate five different network settings and we have conducted of ablation studies for the five settings in the main body part.
E EXPERIMENTS ON LARGER MODEL
We further compare larger V-DNT and L-DNT models and original VIT and GPT2 models ImageNet and OpenWebTeer using mSGDW and AdamW. The results are shown in Figures and 10
We see that on OpenWebText L-DNT-large with mSGDW achieves a comparable performance with L-DNNT-arge with AdamW and achieves much better performance than GPT2-large with mSGDW. Meanwhile, we find out that V-DNTT-huge with mSGDW achieves a comparable performance with V-DNT-huge with AdamW and obtains much better performance than ViT-huge with mSGDW.
OpenWebText Val Loss
3.50
AdamW-GPTZ-Large
3.39
AdamWAL-DNT-Large
mSGDW-GPTZ-Large
mSGDW-L-DNT-Large
3.27
3.16
88 Loss
3.04
2.93
2.81
2.70
2.58
2.47
25k
50k
75k
100k
125k
150k
175k
200k
Step
FIGURE Comparison of GPT2-Large and I--DNT-Large (774M) on OpenWebText. All models trained with 200K in total. GPT2-Large training mSGDW under-performe GPT2-Large with AdamW significantly, but L-DNT-Large with mSGDW can achieve comparable performane with L-DNT-LLrre with AdamW.
22

---

Published as a conference paper at ICLR 2026
OpenWebText Val Loss
3.30
AdamW-GPT2-114
3,20
AdamW-L-NNT-14B
mSGDW-GPT2-144
mSGDWAL-DNT-1.48
3.10
2.99
Loss
2.89
- 2.79
2.69
2.58
2.48
2.38
0k
25k
50k
75k
100k
125k
150k
175k
200k
Step
FIGURE Comparison of GPT2-Large and L-DNT-Large (1436M) on OpenWebText All models are trained with 200K in total. GPTe-Large training mSGDW under-performs GPT2-Large with AdamW significantly, but L-DNT-Large with mSGDW can achieve comparable performane with L-DNT-Large with AdamW.
ImageNet Validation Accuracy
80.0
70.0
60.0
Accuracy
50.0
40.0
30.0
20.0
AdamW-VIT -iT-Hgg
10.0
AdamW V DNT Huge
mSGDW-VIT Huge
0.0
mSGDW-V-DNT
20
60
80
100
120
140
160
Epoch
FIGURE 10: Comparison of ViT-Huge and V-DNT-Huge (632M) on ImageNet. All models are trained with 160 epochs in total. ViT-Huge training mSGDW under-perform: ViT-Huge with AdamW significantly, but V-DNT-Huge with mSGDW can achieve a comparable performane with V-DNT- Huge with A.
23

---

Published as a conference paper at ICLR 2026
OpenWebTex Text Val Loss
1000
Val Loss (Log Scale)
100
DNT
Remove QKNorm
Remove PreNorm
Remove MidNorm
Remove InputNorm
0k 5k
10k
15k
20k
25k
30k
Step
FIGURE 11: Ablation study on the influence of different normalization methods training stability.
F MORE ABLATION STUDY
In this section, we conducted more ablation studies. In these experiments, we started with complete L-DNT model (we use setting in Figure that includes one InputNorm, two PreNorm layers, two MidNorm layers, and one QKNmm.. We removed InputNorm, PreNorm (both layers), MidNorm (both layers), and QKNorm individually. To observe the instability issue, we trained each ablated model with a relatively large learning rate (Adam optimizer with learning rate 0.3, without warmup).
According to the experimental results, we observed that:
When only removing QKNorm, the model collapsed after just few training steps;
When only removing PreNorm, the model began to collapse around 4000 steps;
When only removing MidNorm, the model exhibited significant oscillations between 10K and 30K steps;
When only removing InputNorm, the model was able to converge normally with the 0.3 learning rate.
Therefore, from the perspective of training stability, we conclude that the importance of these normalization layers can be ranked in the following order: QKNorm PreNorm MidNorm InputNorm.
Meanwhile, we also conducted an ablation study on the V-DNT model using AdamW, and observed similar phenomenon. QKNorm is crucial, while PreNorm and MidNorm exhibit comparable stability. InputNorm has relatively minor impact on stability.
F.1 EXPERIMENTS USING MUON OPTIMIZER
In this section, we conducted experiments with using the Muon optimizer to compare GPT-2 and DNT. First of all, we give brief introduction to Muon: Given an update and its SVD decomposition, say G UΣV Muon current Muon implementation, for matrices, we use Muon to approximate the matrices, while for vectors, directly employ the AdamW optimizer.
Newton-Schula method to approximate the UVT matrix. In the
24

---

Published as a conference paper at ICLR 2026
In our experiments with Muon,
learning rate of 10-4 for matrices and 6 x 10 for vectors.
We conducted four sets of experiments: a) GPT-2 with AdamW, b) DNT with Adam W, c) GPT-2 with Muon, and d) DNT with Muon. Same previous experiments, we use 2000 steps warmup for all four experiments. The results are presented in Figure 12
OpenWebText Val Loss
3.40
AdamW-GPT2
3.34
AdamW-L-DNT
Muon-GPT2
3.28
Muon-L-DNT
3.21
Losg
3.15
- 3.09
3.03
2.96
2.90
2.84
25k
75k
100k
125k
175k
200k
Step
FIGURE 12: Performance comparison of DNT and GPT using different optimizers.
Based on the observations in experiments, we note the following two points.
When using the learning rates (10- 6 10 10_4). GPT-2 begins show an increase in validation loss after 50K training steps, indicating training instability. For DNT, we did not observe any instability issues, and the training curve exhibits almost no fluctuations. Additionally, when using Muon, our DNT demonstrates significantly better convergence speed compared to DNT with AdamW.
When both DNT and GPT-2 use AdamW, our DNT shows slightly better performance than GPT-2.
G A BRIEF INTRODUCTION TO SOME EXISTING OPTIMIZERS
Optimization methods in deep learning can be broadly categorized into first-order methods and second- order methods, each with distinct characteristics and applications. First-order optimization algorithms dominate deep learning due their computational efficiency, particularly for high-dimensional and large-scale problems. First-order methods rely primarily on gradient information to find the minimum or maximum of function. Based on learning rate selection strategies. These methods can be divided into optimizers with fixed step size and optimizers with adaptive learning rate.
Stochastic Gradient Descent (SGD) (Robbins & Monro, 1951) serves as the foundational algorithm for neural network optimization. It updates parameters in the opposite direction of the gradient of the objective function. While simple and effective, vanilla SGD can struggle with navigating ravines and saddle points in the loss landscape. Momentum SGD (mSGD) (Nesterov 1983) addresses the limitations vanilla SGD by accelerating gradient descent in relevant directions while dampening oscillations. This method augments the gradient direction with fraction of the update vector from the previous step, allowing faster convergence and helping escape local minima. Other notable variants include signSGD (Bernstein al. 2018), which only the sign of gradients for updates;
25

---

Published as a conference paper at ICLR 2026
SVRG (Johnson & Zhang 2013), which reduces variance in stochastic gradients; LARS 2017). which adjusts learning rates layer-wise.
You et al.
Adaptive methods revolutionized gradient-based optimization by incorporating two key innovations. First, they implement grrameter-speciffie learning rate adaptation, performing smaller updates for frequently occurring features and larger updates for less frequent features. Second, they incorporate historical gradient information, often approximating second-order properties of the loss landscape. AdaGrad (Duchi al. 2011)) I) adapts learning rates based historical gradient information and is particularly effective for sparse data. RMSprop (Hinton) (012)) addresses AdaGrad's 's radically diminishing learning rates by using an exponentially weighted moving average. Adam (Kingma
Ba (014) combines momentum with adaptive learning rates, incorporating both first and second moments of gradients. Adam (Loshchilov & Hutter 2019) modifies Adam with more effective weight decay regularization, while Adafactor (Shazeer & Stern 2018) provides memory-effcient adaptive method. Défossez et provides a unified formulation for adaptive methods like AdaGrad, Adam, and AdaDelta.
The optimization field continues to evolve with recent innovations including MUON (Jordan et al. 2024), LION (Chen et al. 2023), Sophia (Liu et al. (023b), and Mars (Yuan et al. 2024). These methods represent the cutting edge of adaptive optimization techniques, further advancing efficiency and performance in training deep learning models.
Wang & Choromanskal (2025) give a detailed analysis and survey of optimization methods, we would like to recommend the reader to read for full reference.
This paper is orthogonal to these works discussed in this section. Our primary contribution is to demonstrate that vanilla mSGD achieve strong performance Transformer architecture when it does not have a heavy-tail problem in gradients. Notably, the optimizers discussed here can also be effectively applied to proposed DNT network.
Remark.
26


Reviews:
[{"review_id": "HvjPDIvZ5j", "reviewer": "Reviewer_QNDm", "summary": "This paper proposes the Deeply Normalized Transformer (DNT), a Transformer variant designed to address the heavy-tailed gradient problem that hinders the performance of vanilla momentum SGD (mSGDW) in training Transformers. The authors provide theoretical analysis connecting the heavy-tail issue to the Jacobian matrices of different Transformer components and argue that strategically introducing normalization at specific positions (InputNorm, PreNorm, MidNorm, and QKNorm) can stabilize gradients and mitigate this issue. DNT is evaluated on both Vision Transformers (ViT) and GPT architectures, showing that it can be trained effectively with mSGDW to reach performance comparable to AdamW. The paper includes gradient distribution visualizations and empirical results on ImageNet and OpenWebText benchmarks.", "questions": "- How does DNT scale when applied to very large LLMs (e.g., tens of billions of parameters)? Are there architectural or stability issues?\n- Could the authors clarify whether the performance parity with AdamW holds when hyperparameters (e.g., learning rates, momentum) are more extensively tuned for mSGDW?\n- How does DNT interact with modern optimizers such as Muon—does normalization reduce or amplify their benefits?\n- What are the computational overheads introduced by the additional normalization layers in large-scale training?\n- It appears that mSGDW with DNT generally performs worse than AdamW during the early stages of training—could the authors provide an explanation for this behavior?\n\nTypo:\n- line 158: any a forward layer", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper tackles a practical and important problem—reducing dependency on adaptive optimizers like Adam—by improving Transformer architectures to work with simpler optimizers. \n\n- The theoretical analysis provides a clear connection between normalization placement and Jacobian conditioning, offering intuition for the design of DNT. They also did a comprehensive analysis of different normalization techniques.\n\n- Experimental results across both vision and language models demonstrate that DNT narrows the performance gap between mSGDW and AdamW, suggesting potential for simpler and more efficient training pipelines.", "weaknesses": "- The experiments, while promising, are limited in scale (e.g., GPT2-Small/Large, ViT-Large) and lack validation on larger models or diverse datasets to confirm robustness.\n- The empirical novelty is modest, as the approach primarily reorganizes existing normalization techniques rather than introducing new mechanisms.\n- The paper does not include detailed ablation studies to isolate the contribution of each normalization type, which would strengthen the empirical validation.", "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": "Qu2Uk4mXUD", "reviewer": "Reviewer_rT6F", "summary": "The paper proposes Deeply Normalized Transformer (DNT), a Transformer architecture designed to be effectively trained with momentum SGD (mSGDW) instead of adaptive optimizers like AdamW. The authors identify that Transformers exhibit heavy-tailed gradient distributions, which make SGD-based optimizers unstable. DNT addresses this by inserting or repositioning normalization layers (InputNorm, PreNorm, MidNorm, and QKNorm) to modulate the Jacobian of each block, ensuring more concentrated gradient distributions and reducing training instability.", "questions": "1. Could you include ablations isolating the impact of each normalization (InputNorm, PreNorm, MidNorm, QKNorm) individually on gradient statistics and performance?\n2. How much GPU memory and wall-clock time are saved when training with mSGDW compared to AdamW?\n3. How does DNT differ conceptually from “Transformers without normalization” (Zhu et al., 2025)? Could these approaches be unified?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Theoretical justification on how each normalization position affects the Jacobian and stabilizes gradient magnitudes.\n2. Empirical results showing that DNT trained with mSGDW performs comparably to standard Transformers trained with AdamW, both on ImageNet (ViT) and OpenWebText (GPT2).", "weaknesses": "1. Theoretical assumptions are idealized: The high-dimensional isotropy and orthogonality assumptions may not hold exactly for real Transformer activations.\n2. Similar ideas appear in nGPT, StableTransformer, and Lipsformer (which are cited), but the novelty claim is modest—it’s mostly a systematic integration and justification rather than a new normalization method. \n3. The comparison is primarily between mSGD and AdamW. It would be more compelling to see how it performs against other optimizers like Sophia or Lion.\n4. While comparing mSGDW and AdamW, the paper does not clarify if both optimizers were optimally tuned (learning rates, weight decay, warmup). The authors state \"we did not tune the learning rate too much\". However, the hyperparameter tables (e.g., Table 2, 4) show that the settings for mSGDW and AdamW are vastly different. For instance, L-DNT-Small uses LR=1.0 for mSGDW versus LR=6e-4 for AdamW , and V-DNT-Large uses LR=0.5 for mSGDW versus LR=1e-3 for AdamW.\n5. In addition, There's also no evidence that DNT maintains benefits under fine-tuning, transfer, or longer training schedules.", "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": "J4Kr78jnMQ", "reviewer": "Reviewer_epVN", "summary": "This paper addresses a key limitation in training Transformers: the poor performance of momentum SGD (mSGD) compared to adaptive optimizers like AdamW. The root cause is identified as the heavy-tailed distribution of gradients in Transformers, which causes uneven updates across parameters.\n\nTo resolve this, the authors propose Deeply Normalized Transformers (DNT), which strategically apply normalization at specific positions in the architecture to modulate Jacobians, balance weight and activation contributions, and reduce the heavy-tail behavior of gradients. Theoretical justifications are provided for why these normalization choices improve training with mSGD. Empirically, DNT achieves performance comparable to AdamW on ImageNet (Vision Transformers) and OpenWebText (GPT) while requiring less memory and computation, thanks to the ability to use mSGDW.", "questions": "See weaknesses.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper addresses a well-known problem in training Transformers, namely why momentum SGD (mSGD) tends to fail compared to adaptive optimizers like AdamW. It provides a detailed theoretical analysis showing how the placement of normalization layers affects the conditioning of Jacobian matrices and the variance of gradients, explaining why these adjustments are crucial for stable training. \n\n2. The approach is practical, leveraging existing normalization techniques in new positions without introducing additional components. Empirically, the proposed Deeply Normalized Transformer (DNT) matches AdamW performance on ImageNet and OpenWebText, with gradients that are more concentrated and stable under mSGDW. \n\n3. Using mSGDW also reduces memory and computational requirements compared to AdamW, which is a significant practical advantage. \n\n4. Finally, the paper offers analytical insights linking the structure of normalization to optimizer behavior, providing a deeper understanding of the interplay between architecture and training dynamics.", "weaknesses": "1. The paper has several limitations regarding the scope and scale of its experiments. It evaluates the proposed Deeply Normalized Transformer (DNT) only on two benchmarks, ImageNet and OpenWebText, and does not include large-scale or multimodal tasks. This narrow evaluation makes it difficult to assess how well the method generalizes to other domains or to the training of state-of-the-art large models.\n\n2. Another limitation is the increased complexity introduced by multiple normalization placements. While these placements are key to stabilizing mSGD, they also add implementation overhead and require careful hyperparameter tuning.\n\n3. The paper also lacks comparisons with other recent approaches designed to improve stability in Transformers, such as nGPT, Stable Transformer, or LipsFormer. Similarly, there is no evaluation against newer optimizers like Muon, which could provide important context for the relative benefits of DNT and mSGDW (Muon is closer to SGD than Adam).\n\n4. Finally, the evaluation on GPT2-small is somewhat limited in scale and may not reflect the challenges of training modern large language models. The optimizer shows instabilities in some experiments, and the loss gap between mSGDW and AdamW remains non-negligible, which could be a critical concern for training larger or more complex architectures.", "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": "SLsFJfRvCR", "reviewer": "Reviewer_gReu", "summary": "This manuscript targets the demand for advanced optimizers with adaptive learning rates in the training of Transformers. A deep normalization transformer is proposed, in which the heavy-tail gradient problem is overcome by strategically integrating different normalization techniques into the appropriate positions of the Transformer.", "questions": "see above section for detailed information", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The idea is timely, clearly presented, theoretically analyzed, and empirically evaluated on two popular Transformer architectures.", "weaknesses": "However, key clarifications and stronger validation are necessary. The detailed comments are as follows:\n\n1. Based on the normalization techniques mentioned in Figure 4, is the contribution of this manuscript merely an engineering improvement regarding the strategic integration of different normalization techniques? What is the logic of cooperation between them? And how to determine the appropriate positions of each technology in the Transformer?\n\n2. In Tables 2-4 of Appendix C, is the selection of parameters specified by the authors or the relatively optimal settings obtained through search? In the experiment, key hyperparameters and design choices were lacking in ablation.\n\n3. There are inconsistencies or errors in the reference writing, such as:\n\"Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems, 36, 2024.\"\nshould be:\n\"Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems, 36:49205-49233, 2023.\"\nIt is suggested that authors should carefully check, revise and improve.\n\n4. In line 426, \"See Appendix H for the training parameters.\" In fact, \"Appendix H\" does not exist in the manuscript.\n\n5. In the manuscript, \"the norm of * is very large\", \"the norm of * is large\" and \"* become too large\" are mentioned times. How to define or quantitatively describe \"very large\", \"large\", or \"too large\"?  \n\t\n6. What is the definition of σ(∙) in lines 303-304 σ(W_1 ) and σ(W_2 )? In Equation (4), \"Y = Self-Attention(X′), where x′ = PreNorm(x) \". However, what kind of variable is X′? Or is there any connection between X' and x'? Furthermore, in line 256, \"Y = Self-Attention(X) and Y' = Self-Attention(X')\". Is it accurate that it is different from the \"Y = Self-Attention(X')\" in equation (4)?\n\t\n7. In order to facilitate readers' understanding, the authors should introduce and explain the alphabetic symbols and various operation symbols that appear in the manuscript. For example: The definition of \"⨂\" in the equation on lines 241-242 and in Equation (5) needs to be declared. What is \"C\" in Equation (5)? The \"⨀\" in line 254 needs to declare its definition. Furthermore, x_0, which appears in lines 372-373, seems not to have been used. So, where does it play any role? What parameter is \"C_dn\" in the formula on lines 730-731? What is \"C\" in Equation (13)?\n\n8. There is a cross-reference error in \"according to Equation ??\" in line 754. Please check and confirm whether other citations are standardized.\n\t\n9. In lines 759-758, there are ∂L/(∂vec(Y)) and  ∂L/(∂vec(W_q)),   ∂L/(∂vec(W_k)),   ∂L/(∂vec(W_v)), where L and L should be the same?\n\t\n10. The captions in Figures 7-10 are not consistent with those of the other figures mentioned earlier in the manuscript.\n\t\n11. What are the dimensions W_1 and W_2 that appear in Equation (6) respectively? In line 274, does the W_1 x>0 in \"W_2 diag(1(W_1 x>0)) W_1\" represent the relationship between the vector W_1 x and 0? How exactly is it defined? Furthermore, what kind of calculation is 1(W_1 x>0)? The authors should supplement the corresponding definitions or calculation methods.\n\t\n12. The \"if\" in line 754 should be \"If\".", "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)