The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
Authors: Jinbo Wang
Mingze Wang
Zhanpeng Zhou
Junchi Yan
Weinan E
Lei Wu
Venue: ICML 2025 poster (2025)
Keywords: Sharpness
Optimization
Transformer
LLM pre-training
Primary area: deep_learning->algorithms
TL;DR: We identify a clear sharpness disparity across transformer blocks and introduce a novel Blockwise Learning Rate (LR) strategy that speeds up language model pre-training by up to 2x.

Abstract: Transformers have become the cornerstone of modern AI. Unlike traditional architectures, transformers exhibit a distinctive characteristic: diverse types of building blocks, such as embedding layers, normalization layers, self-attention mechanisms, and point-wise feed-forward networks, work collaboratively. Understanding the disparities and interactions among these blocks is therefore important.
In this paper, we uncover a clear **sharpness disparity** across these blocks, which  intriguingly emerges early in training and persists throughout the training process.
Building on this insight, we propose a novel **Blockwise Learning Rate (LR)** strategy to accelerate large language model (LLM) pre-training. Specifically, by integrating Blockwise LR into AdamW, we consistently achieve lower terminal loss and nearly $2\times$ speedup compared to vanilla AdamW. This improvement is demonstrated across GPT-2 and LLaMA models, with model sizes ranging from 0.12B to 1.1B and datasets including OpenWebText and MiniPile.
Finally, we incorporate Blockwise LR into Adam-mini (Zhang et al., 2024), a recently proposed memory-efficient variant of Adam, achieving a combined $2\times$ speedup and $2\times$ memory savings. These results underscore the potential of leveraging the sharpness disparity principle to improve LLM training.

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
Jinbo Wang" Mingze Wang"
I Zhanpeng Zhou 22
Junchi Yan? 2 Weinan E 134 Lei Wu 134
Abstract
They have become the de facto design in modern Al mod- els (Team et al., 2023; Achiam et al., 2023; Liu et al., 2024a).
Transformers consist of diverse building blocks, such as embedding layers, normalization layers, self-attention mechanisms, and point-wise feed- forward networks. Thus, understanding the dif- ferences and interactions among these blocks is important. In this paper, we uncover clear sharpness disparity across these blocks, which emerges early in training and intriguingly per- sists throughout the training process. Motivated by this finding, we propose Blockwise Learning Rate (LR), a strategy that tailors the LR to each block's sharpness, accelerating large language model (LLM) pre-training. By integrating Block- wise LR into AdamW, we consistently achieve lower terminal loss and nearly 2x speedup com- pared to vanilla AdamW. We demonstrate this ac- celeration across GPT-2 and LLaMA, with model sizes ranging from 0.12B to 2B and datasets of WpenWeTText, MiniPile, and C4. Finally, we incorporate Blockwise LR into other optimiz- such Adam-mini (Zhang et al., 2024c), recently proposed memory-efficient variant of Adam, achieving combined 2 speedup and 2x memory saving. These results underscore the potential of exploiting the sharpness disparity to improve LLM training.
Compared to traditional architectures, e.g. multilayer per- ceptrons (MLPs), convolutional neural networks (CNNs), and recurrent neural networks (RNNs), transformers exhibit distinctive alloy-like characteristies, where diverse types of blocks synergistically combine to achieve superior perfor- mance. A transformer at minimum includes self-attention (further broken down into query-key (QK) and value-output (VO)) blocks, point-wise feedforward networks (FFN), nor- malization layers (Norm), and embedding layers (Emb). Uncovering the distinct properties these blocks, as well as the differences and interactions among them, is thus crucial for gaining deeper insight into transformer models (Wang & E, 2024).
In practice, transformers are typically trained using the Adam optimizer (Kingma & Ba, 2014; Loshchilov & Hutter, 2017). Dissecting the alloy-like characteristics of transformers can provide insights into why Adam outper- forms stochastic gradient descent (SGD) for transformer training (Devlin, 2018; Zhang et al., 2020; Pesme & Flam- marion, 2023; Kunstner et al., 2024; Zhang et al., 2024b) and even holds promise for unlocking further improvements in training efficiency (Popel & Bojar, 2018; Xiong et al., 2020; Zhang et al., 2024c). Particularly, Zhang et al. (2024b) and Zhang et al. (2024c) observed that unlike ML.s and CNNs, the Hessian (aka sharpness or curvature) of transformers exhibits distinct blockwise heterogeneity. Building this insight, Zhang et al. (2024c) successfully reduced Adam's 's memory footprint nearly by half without sacrificing training efficiency for a variety of LLM and non-LLM training tasks.
1. Introduction
Transformers (Vaswani et al., 2017) have achieved remark- able success across fields, including natural language pro- cessing (Brown et al., 2020), vision (Dosovitskiy et al., 2020), and scientific computing (Jumper et al., 2021).
Our Contribution. In this work, we aim to explore how we can leverage the aforementioned alloy-like characteris- ties of transformers to improve training efficiency. Specifi- cally, our contributions can be summarized as follows:
"Equal contribution 'School of Mathematical Sciences, Peking University 'Shanghai Jiao Tong University 'Center for Machine Learning Research, Peking University *AI for Science Institute, Beijing, China. Correspondence to: Lei Wu <leiwuu @ math, m.th.pku.edu.com pku.edu.en>. Mingze Wang <mingze- wang@ stu, wtt.pku.cc...c.....
The sharpness disparity principle. Motivated by the alloy-like characteristics, we examine the sharpness of transformers at the level of block type. Surprisingly, we discover a distinct disparity in sharpness across different block types, summarized as follows:
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre- Pre-Training
Itaration 40000
LLaMA (1.18) OpanWebTes Text
S(Emb) <((K <
< ) <( (1)
Embedd
Adari (501
Adari (108-)
Here S))) denotes the average sharpness of block type
veliiaatoon Rook
Slucwi
FFN
(see (4) for the calculation details). See Figure (left) for an illustration of this principle. Intriguingly, this principle emerges in the early training stage and persists throughout the subsequent training process, as shown in Figure 3. These findings are validated through exten- sive experiments on the training of GPT-2 (Radford et al., 2019) and LLaMA models (Touvron et al., 2023), span- ning various model sizes and datasets. We also provide preliminary theoretical explanations to complement these empirical observations.
Nom
10
12 10
10
1217
1014
10
10
Avg. Sharpness (Lage)
num stops
Figure 1: (left) Sharpness disparity among block types in pre-trained GPT-2 (small) on OpenWebText, exhibiting clear order relationship as characterized by Principle (1), (right) For the pre-training of LLaMA (1.1B) on Open Web- Text, incorporating our Blockwise LR strategy into Adam results a lower terminal loss and a 1.92x speedup com- pared to the well-tuned vanilla AdamW.
The Blockwise LR strategy. Inspired by Principle (I), we propose tuning LRs by block type to accelerate LLM pre-training. Specifically, we adjust the LRs of blocks within the same type in proportion to their sharpness, while keeping the LR of the block type with the high- est sharpness unchanged. This strategy accelerates the dynamics along low-sharpness directions without com- promising training stability, as the latter is governed by the high-sharpness directions.
2. Related Works
Sharpness structures in transformers. Recent work has started to investigate blockwise sharpness patterns in trans- former models through Hessian-based analyses. For exam- ple, Zhang et al. (2024b) empirically observed the sharpness' blockwise heterogeneity but did not establish a clear princi- ple regarding the sharpness disparity among different blocks. Meanwhile, Ormaniec et al. (2024) provided Hessian anal- ysis for a single self-attention (SA) layer, focusing only on the sharpness disparity between the query-key (QK) and value-output (VO) blocks within the same layer.
The effectiveness of Blockwise LR is extensively vali- dated in LLM pre-training across both GPT-2 and LLaMA models, with model sizes ranging from 0. 1.12B to 2B pa- rameters, and datasets including Webeet (Gokaslan & Cohen, 2019), MiniPile (Kaddour, 2023), and C4 (Raf- fel et al., 2020). The results can be summarized as follows:
In contrast, we examine sharpness at the block-type level across the entire transformer architecture, rather than focus- ing on individual blocks (as in Zhang et al. (2024b)) or a sin- gle layer (as in Ormaniec et al. (2024)). This coarse graseed perspective reveals a consistent disparity, as formalized by Principle (1), which persists throughout most of the training during the initial steps.
AdamW with Blockwise LR achieves lower terminal loss and is nearly 2 faster than vanilla AdamW.
See Figure 1 (right) for a quick view of the accelcration effect achieved by Blockwise LR. Furthermore, we ex- plore the compatibility of Blockwise LR with other Adam- based optimizers. Specifically, we integrate our Block- wise LR into Adam-mini (Zhang et al., 2024c), achieving both 2x speedup and 2x memory saving.
Efficient optimizers for LLM pre-training, AdamW (Adam with decoupled weight decay) (Loshchilov & Hutter, 2017) has become the default optimizer in LLM pre-training. Efforts to design more efficient optimizers generally fall into two main categories: accelerating convergence and reduc- ing memory footprint. Accelerations have been developed using techniques such as Nesterov momentum (Xie et al., 2022), diagonal second-order estimates (Liu et al., 2024b; Wang et al., 2024), variance reduction (Yuan et al., 2024), and matrix- preconditioners (Keller et al., 2024; Vyas et al., 2024). Memory-efficient optimizers utilize sign-based methods (Chen et al., 2024), reduced usage of second mo- ments in Adam (Zhang et al., 2024c), and gradient low-rank projection (Zhao et al., 2024a). The closest work to our Blockwise LR is Wang et al. (2024), which also increases the LR along low-sharpness directions. A detailed compari- son is deferred to Section 5.
Remark 1.1. There has been long- long-standing effort in deep learning to accelerate neural network training by adapting layerwise learning rates, a strategy that has proven effective in architectures such as MLPs and CNNs (Yang, 2019; Yang et al., 2022; Everett et al., 2024; Shin et al., 2024). However, these approaches have not been successfully transferred to the training of deep transformers. We hypothesize that this gup stems from trunsformers' distinctive alloy-like charac- teristics: the inherent block-level diversity makes layerwise learning rate strategies inadequate. To investigate this fur- ther, we examine layer-level sharpness in Figure and no clear trends emerge across layers. This suggests that while sharpness disparity exists at the block-type level, it does not exhibit consistent pattern at the layer level.
The edge of stability (FoS) phenomenon. Neural network
2

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
training typically occurs at the EoS stage (Wu et al., 2018; Jastrzebski et al., 2020; Cohen et al., 2021; 2022), where the optimizer exhibits oscillatory behavior along sharp di- rections without diverging, while steadily progressing along flat directions, leading to loss reduction. Several works (Wen et al., 2024; Song et al., 2024; Cohen et al., 2024; Wang et al., 2024) have highlighted the crucial role of the dynamics along flat directions (referred to as river directions by Wen et al. (2024), bulk directions by Song et al. (2024), and stable direction in Wang et al. (2024)) in reducing total loss. Notably, Wen et al. (2024) further demonstrated that this picture is essential for understanding LLM pre- -re-training. Building on these insights, our Blockwise LR approach is designed to accelerate training by amplifying the dynamics particularly along the flat river directions.
TF(X) x R For each layer E [4], the
computations are as follows:
X (1- = x (1-1) + SA
(111/2) (-10 ());
(2)
x = x (1 - 1) FFN (1) (Norm* (1) (x(1-
Norm blocks. Here, Norm (v) (v {i 1)/ denote normalization layers (e.g., LayerNorm (Lei Ba et al., 2016)
and RMSNNrm (Zhang & Sennrich, 2019)) with learnable parameters y ,,(() E RIIDD For LayerNorm, the com- putation for a token x E 1x is:
x E [c]
Norm Norm (x)
+ 0 + B(0)
where the bias B is omitted in LLMs such as nanoGPT.
FFN blocks. FFN denotes a (token-wise) two-layer
3. Preliminaries
FFN of width M, comprising parameters w(O RDxM W E RMxD, and using activation function a()) such as ReLU. For any token æ E RIxD, the operation is:
Notations. Let I1-112- I1-Ilp- and Tr(-) denote the spectral norm, Frobenius norm and trace for matrices, respectively. Given A E IR m its row-wise vectorization is defined as
veo( A) = (01,1.-.. 01,u: ... 0m,1,"' ". amm,n) E Rm The Kronecker product and Hadamard product are denoted by @ and @ respectively. The row-wise mean and variance of A E _=*x" are denoted by [A] E R2xx and [A[A] E Rmxx, respectively. Specifically, they are defined as: for all i E [m],j E [ra]. (A)(s =
FFN (a) a (x ?
SA blocks. SA(), multi-head self-attention, has parame- ters wi w:O w w! E RDD When applied to a Q K
sequence Z E RaxD it operates as:
1 E Arx)² We will standard big-O notations like 0(), (()), and ((-)
"E=1 A,,: (Vr[A]),s (A,, - n
H
SA 11) (2) = Σ
SA
(L))
(Z)
(A)
to hide
constants.
h=1
Zwwll (L.b)
ZW(Lh) (L.b)
Zw
+ +M
Jacobian matrix. Given a vector-valued function: b +
K
(zW(lN))
softmax
(b) with b E R² and a(b) E R", the Jacobian is de- (8a.)
VD/H
fined as aa ab = (bb ja.j E mm Analogously, for a
matrix-valued function: B + A(B) where B 1 and A(B) E Rmxx to avoid directly working with tensors,
(1,)) w:ll.)) willh)) where H is the head number, and W DD//HH) w-(kk)
(1.h)
(D/H))D are split from K W w! by heads, respectively. The oper- ator softmax(() represents the row-wise softmax normaliza- tion. For the next-token prediction, the mask M E B satisfies Ma, = oo if < j and M.,- = 0 otherwise.
the Jacobian is defined as
0A_ Öven(A) (A) E menpp Bvec(B)
w w w
OB
3.1. The Transformer Architecture
Given n-token input sequence X x= [ "T.TTT E Raxd, where d refers to the vocabulary size in LLM and each 2, corresponds to the token's s one-hot encoding, an L-layer transformer TF processes it as follows.
3.2. Blockwise Sharpness and the Efficient Estimation
Mcasuring sharpness requires accessing the Hessian matrix, which is computationally expensive due to the high dimen- sionality of the parameter space. Consequently, approximate methods are needed to reduce computational complexity.
Embedding layer. First, each input token is embedded into the latent space through an embedding layer with parameters WE E Rdx D bg IRD
x
(0) x. WE + bg, s E [n].
Let ((.-)) denote the cross-entropy loss. For an input data x E Rdd and label y E Rd,, let the model's prediction be I (x:; 0) E Rdy. The Fisher (Gauss-Newton) matrix F F(0) is widely recognized approximation of the Hessian, particu- larly near minima. Thus, the diagonal Hessian can be esti- mated as h(0) = diag( (())). a popular technique in deep learning optimization (Martens Grosse, 2015; Grosse
where the bias by is omitted in LLMs such as nanoGPT (Karpathy, 2022).
L-layer SA-FFN blocks. Then the embedded sequence (() is processed by L-layer SA-FFN blocks, and the output of the final layer is taken as the output sequence
3

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
Martens, 2016; George al., 2018; Mi al.. 2022; Liu et al., 2024b; Wang et al., 2024). Moreover, given in- put batch {(T6, yb)] Jmm the empirical diagonal Fisher can
the best of our knowledge, provide the first comprehensive comparison of sharpness across block types in transformers.
b E Ve(s(>>;; f); ÿb) O (((f(< 0); 2)), where yb voftmax(((0;z))) How- ever, noted by Liu et al. (2024b), implementing this estimator is computationally expensive due to the need to calculate B single-batch gradients. Liu et al. (2024b) pro- posed a more convenient estimator diag( Par (0)), which only requires the computation of the mini-batch gradi- ent (O)
be estimated: diag( (()))
lterationn 40000
Iteration = 40000
Embed
0.5
Emoed
0K
Density
0.3
FIN-
Norm
0.2
Norm-
0.0
1 -E Vi(f( (s) B
with y)
20-11
10
10-23
10-2 10-14
16-13
10
Avg. Sharpess (Log Scale)
Shappees (Log Scale)
softmax( (z6; 0)):
Figure 2: (left) The average sharpness for the five typical block types in pre-trained LLaMA model (0.25B); (right) the sharpness distribution across different blocks in a pre- trained GPT-2 (small) model.
h0)) = diag( Parr (0)) = B (Ln @ Vin(O)- (3)
According to Liu et al. (2024b, Section 2), this estimator is unbiased, i.e., EEy [diagg For [0)] = Egg ((tiag(((())
Figure 2 (right) plots the full sharpness distribution for each block type, whereas Figures (left) and 2 (left) only re- port mean sharpness values. Evidently, even at the dis- tribution level, Principle (1) remains valid. Interestingly, the Emb block exhibits much higher variance compared to other blocks. This behavior likely stems from the embed- ding layer's ss direct interaction with the entire vocabulary, where rare tokens result in the wide spread of small sharp- ness and frequent tokens contribute to large sharpness. A similar insight has been utilized by Kunstner et al. (2024) to explain the necessity of Adam in training NLP models.
Given a block type * E {Emb, QK, VO, FFN, Norm), let O[o] represent the parameters associated with all blocks of that type, and let ()))) denote the corresponding diagonal Hessian. The average sharpness for each block type can then be approximated as follows:
B Vo. (100
((0 = (1()))
(4)
#((
where En corresponds to (3) and #(O)) denotes the num- ber of parameters associated with the block type . For brevity, f in (4) will be omitted when there is no ambiguity.
Furthermore, Figure 3 illustrates the evolution of block- wise sharpness during the training process. We can see that Principle (1) is not exclusive to well-trained trans- formers; instead, it emerges in the early stages of training and persists consistently throughout the subsequent training process. This observation underscores the potential of lever- aging Principle (1) to enhance LLM pre-training we refer to Section for further explorations.
Remark 3.1. It worth noting that in (4), the sharpness averaged over all blocks of the same type, which may be distributed across different layers, rather than being calculated within each individual block.
4. The Sharpness Disparity Principle
4.1. Main Findings
Comparison with existing works. Our findings build on prior work, extending key observations. Zhang et al. (2024b) noted the block heterogeneity in the Hessian of transformers but did not establish a clear principle for sharpness distine- tions across blocks, as we do with Principle (1). The work of Ormaniec et al. (2024) is more closely related but focuses solely on a single self-attenioon layer (SA), reporting the relationship S(QK) < S(VO). In contrast, we analyze all major block types in transformers, including Emb, FFN, and Norm, thereby offering a more comprehensive principle that captures the full scope of sharpness disparity.
We first investigate the sharpness disparity across differ- ent types of building blocks (Emb, QK, VO, FFN, Norm) in transformer-based LLMs. Specifically, we pre-trained GPT-2 (Radford et al., 2019) and LLaMA (Touvron et al., 2023) models on the OpenWebText dataset using default configurations. Blockwise diagonal Hessians are analyzed at various checkpoints using the Hessian estimator (3). The experimental details can be found in Appendix A.
In Figures (left) and (left), we report the average sharp- ness, estimated using (4), of the five typical types of blocks for GPT-2 and LLaMA, respectively. The results reveal clear and consistent sharpness disparity among different block types, as summarized in Principle (1). Specifically, Norm layers consistently exhibit the highest sharpness, the Emb layers are the flattest, and QK layers are relatively flatter compared FFN and VO layers. These findings, to
4.2. Theoretical Insights
To provide theoretical insights into explaining Principle (I), we derive analytic expressions of S(+) and analyze their dependence on parameter magnitudes and numbers of each block. For simplicity, we denote ((0) == LB()), where
4

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre- Pre-Training
Itrntionn:
llmaiionn
100
Itmatin 1000
limiion 10000
Itmition 20000
Eribdd
Entec
Emred
Droed
Embed
VG-
V0
V0
V0
12 20 20 10
10
10
10 12"h
11
10
20th 20"
20 1o 10
Avy, Sharpness (Log Scale)
Aug Sharpness (Lor Soale)
Avg. Shonpness Scale)
Aug. Sharpness (Log))
Ang. Sharyness (Log Scale)
(a) Evolution of the average sharpness across different blocks during pre pre-training GPT-2 (small) on TeenWeeText.
llrriiont
Reseiion 100
Desdion 1000
Rerdioor 10000
Iestiion 20000
Drbed
Embed
Embed
Dmbed
Entedd
F+M
H00-
PPN
PPM
V0
28
13t4
28
Aug. Sharpresss (Log Stadie)
Aug. Shorpnees (Log Stode)
Aug. Shorppasss (Log Scalle
Avg. Shoppess (take)
Aug. Sharpeess (Log Stode))
(b) Evolution of the average sharpness across different blocks during pre-training LLaMA (0.25B) on OpenWecText.
WehText.
Figure 3: In these experiments, the total training steps are both 50k. Principle (1) emerges during the initial phase (from iteration to iteration Ik), which accounts for only approximately 2% of the total steps, and persists throughout the subsequent training process.
(0) is defined in (3). Then from (4), we have S(+) =
where the denominators (D² or D) reflect the number of parameters each group.
B PIV. 21F (+). Without loss of generality, we set =
1. Our calculations for QQ apply to general @
Theorem 4.1 provides theoretical support for our main find- ing: S(FFN) is substantially smaller than S(Norm). As illustrated in Figure (a), during training, 11 gradually decreases, and ||WlllF ( E {1,2}) in FFN layers remains larger than l/ resulting in D2 [[ 1 >>>>>>lyll Dl
Considering blocks across different layers is complicated. Therefore, focus comparisons within the same layer. Specifically, we examine the following sharpness compar- isons: (i) FFN vs. Norm within the same layer; (ii) SA (comprising QK and VO) vs. Norm within the same layer: and (iii) Emb vs. the adjacent Norm.
Theorem 4.2 (QK, VO vs. Norm). Consider the (1- 2/-- th layer in (2). Omitting the layer index for simplicity,
Theorem 4.1 (FFN vs. Norm). Consider the I-th layer in a transformer (2). Omitting the layer index for simplicity, let Y = X + FFN (Norm (X; ; W1, W2), where FFN utilizes the (Leaky) ReLU activation a. Then, the gradients of Q w.r.l. W1, W2. and y are:
let Y = X + SA SAASSANNNn Norm (X:7); WK, Wo, Wv, Wv.Wo) Consider a single-head attention (i.e., H = 1) for sim- plicity. Then, the gradients of Q w.r.t. different blocks (WK, Wo. Wv, Wo, are provided in Appendix B.2. Furthermore, there exist two
Wo)
constants 4, v > 0 (detailed in Appendix B.2), such that:
0Q 0Q|
0A bM
N
W1
6 Ta;
OW2 2 OY
42
0Q = 0Q
0A
S(W.) =0
E (K,Q);
In ④
W
Noorm @ IM):
D2 W. 2
ow; OY
am
0Q 8Q
In
W2 1,2
(I, W, 7) 0M
0A
S(W.) o
122
E (V,0);
dY OY
D2 WP|lW.l 2
diag( vcc( X staa)) (1nx1 Ia),
( +
2
S(y) O
where Xstd 1= X XNorm Norm(X:
D 1| 12
/v.[x]
where the denominators (D² or D) reflect the number of parameters in each group.
X ② (1nx1 X 7), A == (M). M == XNorm W1. Let
== DVD Il SM ll Will ||aall ||>l Then,
the blockwise average sharpness can be bounded
Theorem 4.2 provides theoretical support for our main find- ing that both S(QK) and S(VO) are significantly smaller than S(Norm). The inclusion of the softmax operation in attention layers introduces additional complexity in the cal- culations. Detailed derivations are given in the appendix. As shown in Figure 8 (b), during training. ||>||p gradually
2
S(W.) = 0
E (1,2);
D2 W.
2
S(7) = 0
|3
5

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
decreases, and |/ (. {K, QVQ.V, in SA blocks re- mains larger than ||> resulting in D2 [.. > D l
Along high-sharpness directions, the optimizer exhibits significant fluctuations without con- verging or diverging. These components of dynamies govern training stability, as further increasing the LR in these directions can lead to instability, while con- tributing little to loss reduction.
Fast dynamics:
This theorem does not explicitly establish that S(QK) < S(VO). Studying this relation requires deeper analysis of the constants & and V, as well as the magnitudes of W. llp- Ormaniec et al. (2024) has demonstrated S(QK) < S(VO) both theoretically and experimentally, and we defer to that analysis instead of repeating it here.
Along low-sharpness directions, the optimizer progresses steadily, making the primary con- tribution to loss reduction, albeit at a slow rate.
Slow dynamics:
Theorem 4.3 (Emb v.s. Norm). Consider the embed- ding layer and its adjoint normalization layer of a trans- former (2). Omitting the layer index for simplicity, let: Y := Norm(X y)- The gradients of Q w.r. W and "y are derived in Appendix B.3. Moreover, there ex- ists a
Inspired by the above picture, a promising approach to ac- celerating training is as follows: given a base optimizer, increase the LRs along low-sharpness directions while keep- ing the LR of high-sharpness directions unchanged. This strategy aims to speed up loss reduction without compro- mising training stability.
constant V > (also detailed in Appendix B.3), such that:
Wang et al. (2024) has implemented this idea by adjusting the LR of each parameter based on its sharpness. How- ever, this approach faces two key challenges: 1) it requires frequent diagonal Hessian estimation, which imposes sig- nificant computational and memory overhead; 2) sharpness estimates at the individual parameter level can be unreliable. := WW [WE] e). The
TT2
S(WE) =
Dd min | iE|d]
T/2
S(7) 0
D
l
T where WF = (w E denominators (Dd or D) represent the number of parame- ters in each group.
W (E,)
The Blockwise LR. Unlike Wang et al. (2024), we propose adjusting LRs at the block-type level, as our Principle (1) reveals a consistent sharpness disparity at this granularity. Specifically, let mass denote the I.R for base optimizers such as AdamW, the LR for each block type is then adjusted as follows:
Theorem 4.3 provides theoretical justification for our main finding that S(Emb) is much smaller than S (Norm). As shown in Figure 8(c), during training. Dd f l12 >
D mel (Notice that the vocabulary size d is very large in practice, e.g., 50304 for the GPT tokenizer.)
Norm blocks (the sharpest directions): we still use the base LR, 'orrm = Thase. to keep training stability;
Recalling the definition of average sharpness (4), the key step in deriving Theorem 4.1 and 4.2, and 4.3 is establish- ing QI 011/1 [I). This relationship is highly intuitive given the compound multiplicative nature of trans- former blocks, where the norm of the derivatives is inversely proportional to the norm of associated parameters, even with weak non- -ineariiss...s. For example, if = [[= 24 and O <py), then [8Q/0z 00/0x: 10 (f())// mad x 1/|2r/ for all i [n].
Other blocks (low-sharpness directions): we adjust the LRs of these blocks by 7 & (())Mbase where (Emb, QK, FFN, VO}, where r()) denotes the adjusting ratio for the block type
Naturally, we can set r()) x ((Nor))S(( However, in practice, we find that manually tuning r())' inivvolvin only four hyperparameters-whii following the qualitative trend described by Principle (1) is more effective. Further details are provided in Section 6.
5. The Blockwise LR Strategy
It is also worth noting that due to its simplicity, Blockwise LR can be seamlessly integrated into modern LLM training frameworks such as Megatron (Shoeybi et al., 2019).
Recalling Figure 3, the sharpness disparity across different blocks, as described in (1), emerges early in training and persists until convergence. This insight can be leveraged to accelerate LLM pre-training, as elaborated later.
6. Experiments
Models and datasets. We evaluate our proposed Blockwise LR in the pre-training of decoder-only LLMS across various model types, model sizes, and datasets'. Specifically, we
Fast-slow dynamics at EoS. As discussed in Section 2, re- cent studies (Wen et al., 2024; Song et al., 2024; Wang et al., 2024) have highlighted the distinct roles of the dynamics along high- and sow-sharpness directions during EoS. The main picture summarized as follows:
'The code is available https: actps://github.com Wangboo/H BBlooo/Blookwwr wtwp://eraiiiiiiieiiiiiinnn selieannnnnaaa
6

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
consider two widely-used LLMs: LLaMA and GPT-2; we experiment with model sizes ranging from 0.12B to 2B parameters; the datasets includes Open WeenWe (Gokaslan & Cohen, 2019)* ?, MiniPile (Kaddour, 2023)', and Colossal Clean Crawled Corpus (C4) (Raffel et al., 2020)*, providing a highly diverse text corpus.
6.1. Main Results
Main findings. In Figure 4 and Figures (rright), we com- pare the performance of AdamW with Blockwise LR against vanilla AdamW across various settings. Our observations, which consistently hold across all experiments inccluding both GPT-2 and LLaMA models with sizes ranging from 02B to and datasets including Open WebText and MiniPile, are as follows:
Baselines. As a baseline, we use the default Adam opti- mizer, configured with the hyperparameters B1 = 0.9, 82 = 0.95 and weight decay À = 0.1. To ensure training stability, gradient clipping is applied with 1.0. These settings align with the training protocols used in nanoGPT and LLaMA models (Touvron et al., 2023). The LR strategy includes a linear warm-up up phase followed by a cosine decay scheduler, capped at r max. And the terminal L.R 1 min n is set to 1r_max/20. For each experiment, we first tune the 1r _max to be optimal for Adam W, and the baselines are trained us- ing these optimal 1 1r_ax'' s. Details of the tuned r mmax values can be found in Appendix A.
Given the same total number of training steps, Block- wise LR enables AdamW to reach a lower terminal loss than vanilla AdamW.
Across different total training steps, AdamW with Blockwise LR achieves nearly 2x speedup com- pared to vanilla AdamW.
An intriguing observation in Figure 4 is that Adam with Blockwise LR often starts to outperform vanilla Adam from the mid-to-late stages of training. This behavior re- sembles the WSD scheduler (Wen et al., 2024; Hu et al., 2024), which typically surpasses cosine or linear decay LR schedulers in the late stage (during the decay phase). Under- standing the underlying cause of this phenomenon requires further investigation, which we leave for future work.
Adjusting ratio tuning and its transferability. To incor- porate the Blockwise LR into A we simply use the Ir_max (tuned for vanilla A) for Norm blocks. Then, we only tuned the four adjusting ratios in a single small- scale experiment specifically the pre-training of LLaMA (0.25B) on Minipile following the rule: r()) is adjusted according to the trend of S(Norm) guided by Principle (1).
Scaling law is in favor of Blockwise LR. To further exam- ine scaling behavior, Figure 5 (right) visualizes the scaling laws of Adam W with Blockwise LR versus AdamW during LLLMA pre-training. For MiniPile (left) and Open OpenWebTText (middle), the performance gaps between the two optimiz- ers get larger models size grows. For C4 (right), the performance gap remains stable across model scales, with the corresponding scaling curves remaining nearly parallel. These results suggest that the gains offered by Blockwise LR may persist at larger model scales.
S(+)
The tuned hyperparameters are:
r(Emb) = 10,100, r(KK) 8, r(FFN) 6, r(VO) 4. (5)
Notably, the adjusting ratios are highly robust hyperparame- ters, as demonstrated in the following ways:
First, as shown in Figure 10, in the experiments for tun- ing the adjusting ratios, Blockwise LR demonstrates ro- bustness these hyperparameters, consistently acceler- ating pre-training across a range of r()) ss The configu- ration in (5) achieves the largest improvements among those tested. Notably, even with suboptimal ratios, Blockwise IR still delivers significant performance gains. Further details are provided in Appendix A.2.
Evaluation on downstream tasks. Furthermore, as ob- served in Table 1, the improvement in validation loss trans- fers to an improvement in downstream task accuracy. Within the same number of pre-training steps, the LLaMA (L.IB) trained with Blockwise LR shows better downstream perfor- mance than Adam among all evaluated tasks.
Second, the configuration in (5), tuned from single experiment, transfers perfectly across all Adam W experiments conducted in this paper. Consequently, we adopt (5) as the default adjusting ratios for all Adam V experiments. This robustness aligns with the consistency of Principle (1), which holds across GPT and LLaMA models, various model sizes, and datasets.
6.2. Ablation Studies
In the preceding experiments, Blockwise LR is applied to all major blocks simultaneously. Here, we conduct ablation studies to assess the contribution of each block type individ- ually. Specifically, we pre-train a LLaMA model (0.25B) on Open Tex focusing three comparisons: (i) applying Blockwise LR exclusively to Emb; (ii) applying Blockwise LR to both Emb and FFN; (iii) applying Blockwise LR to blocks of all the four types (Emb, FFN, QK, and VO). The adjusting ratios follow Eq. (5) and the results are shown in Table 2.
2 An opensource recreation of the WebText corpus, widely used for LLM pre-training such as RoBERTa (Liu et al., 2019) and GPT-2.
AA 6GB subset of the deduplicated Pile (825GB) (Gao et al.,
2020)
'A large-scale public language datasets, widely used for LLM pre-training such as T5 (Raffel et al., 2020)
7

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre- Pre-Training
GPT-2 (small) on OpenWec Text
LLaMA (0.13B) on MiniPlic
LLaMA (0.25B) on MiniPile
LLaMA (0. on MiniPlic
Aduni (506)
Adaaeiyy (38k)
Alleet (3h)
Adbini (108))
2.50
A
validaaiio loss
Blackwisc (506)
sccllwee (R()
validaztion loss.
(B(30)
Brocloss (R(3))
volidaation
2.
2:35
244
260
2:45
spondup
45k
steps
steps
steps
steps
LLaMA (0. 258) OpenWeb Text
LLaMA (0.58) Ween Text
LLaMA (0.758) OpenWeb Text
LLaMA (2B) Open WexText
2.50
Adlent/ (500k)
Adar MileeW//o (220))
Adar (200k)
Al (LO0K)
2.00
Blockinee Pntcciniie LR (504)
velidatoo low
Blociieiue LR (500)
Bioclivise (S(5)))
validation lows
Blodowiie LR (50k)
valddation
validation
2.00-
2,45
2.01
0,001
2.75
1.51x spcooup
1,51%
speodup
1.55x
75
95x
1004
75
95%
1004
num steps.
num stops
num stags
num steps
Figure 4: AdamW with Blockwise LR consistently outperforms AdamW in LLM pre-training tasks across different model types, varying model sizes, and datasets.
Scaling Laws (LLaMA MiniPile)
Scaling Lavs (LLaMA OpenWebTess Text)
Scaling Laws (LLaMA on C4)
2.50
termiral valleation loss
Accomity
torminal validation loss
Adamtt
terminal soliclation loss
3.0
AdomwW
2.45
Blockwlse
Blockwiss
2.0
Blockwise
2:40-
2.0
2.7
2.35
2.0
2.30
2.5
200
330
D00
350
1000
1503
2000
2xx 3x 109t
parameters
M
parameters
M
(noe-embedddnng) parameters, M
Figure 5: Scaling-law comparison of AdamW with Blockwise LR and AdamW various datasets for LLaMA models.
Table 1: Evaluation results on downstream tasks (0-shot with innluuationhharess, of LLaMA models (I.1B) pre-trained on OpenWebText using AdamW W or Blockwise LR for 50K steps. The best scores in each column are bolded.
Method
ARC_E
ARC_C
PIQA
Hel Halag
OBQA
WinoGrande
SCIQ
Adam
52.69
22.87
68.71
36.13
19.40
55.17
77.60
Blockwise LR
54.29
25.34
69.53
38.00
22.60
59.83
81.60
First, the results show that applying Blockwise LR any block consistently improves performance, supporting the hypothesis that dynamics along low-sharpness directions are crucial for loss reduction. Among all blocks, apply- ing Blockwise LR to FFN yields the largest improvement (0.043 - 0.016 = 0.027), likely because FFN blocks com- prise the majority of model parameters, offering the greatest potential for optimization gains.
Norm with other blocks.
In summary, these ablation studies further validate the ef- fectiveness of Blockwise LR and confirm the rationale of selecting specific types of blocks for LR amplification, as guided by the sharpness disparity principle.
Table 2: Ablation results for the effectiveness of Blockwise LRR in pre-training LLaMA (0.25B) on
Second, we conduct an additional experiment to assess the impact of increasing the LR for Norm blocks. Specifically, the Norm LR is doubled, while the LR for other blocks remains unchanged from the baseline. As shown in the last row of Table 2, this leads to deterioration in performance, contrasting with the improvements seen when increasing the LRs for other blocks by far more than double. This result underscores a fundamental difference in the dynamics of
Blockwise LR
terminal loss (50k steps)
w/o
2.834
Emb
2.818 (-0.016 /)
Emb & FFN
2.791 (-0.043 /)
Emb & FFN QK & VO
2.784 (-0.050 //
Norm
2.837 (+0.003 X)
8

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
6.3. Integration into Other Optimization Schemes
closely resemble those in Figure 3(b), which uses AdamW. Our Principle (Eq. (1)) emerges during the initial phase, and persists throughout the subsequent training process.
In practice, there are two popular directions for improving LLM pre-training: acceleration and reducing memory con- sumption. While Blockwise LR has demonstrated remark- able success in accelerating pre-training, a natural ques- tion arises: Can Blockwise LR be combined with memory- efficient optimizers to achieve both faster training and fewer memory consumption?
Blockwise LR with wsd scheduler. The preceding exper- iments employ the cosine decayed LR scheduler. In this section, we evaluate Blockwise LR under an alternative and increasingly popular scheduler: warmup-stable-decay (wsd) (Hu al., 2024), which includes a linear warm-up LR to peak 1r max, followed by stable phase where LR remains at 1 j_max, and then a linear decay to Ir.min. We extend our experiments to incorporate the WSD scheduler. Experimental details are provided in Appendix A.3. As shown in Figure 7 (right), Blockwise LR still achieves a 2x speedup when used with the wsd scheduler.
Blockwise LR on Adam-mini. Without loss of generality, we choose the Adam-mini (Zhang et al., 2024c) optimizer, an Adam variant that reduces memory consumption by ap- proximately 2x compared to AdamW. Here, we conduct experiments to explore whether Blockwise LR can also ac- celerate Adam-mini. Following Zhang et al. (2024c), we adopt the 1r_max that tuned for AdamW as the the 1r _max of Adam-mini. However, since Adam-min employs SGD within each block, its dynamics differs significantly from A Consequently, for Adam-mini with Blockwise LR, we re-tune the ratios r()) for o E {Emb, QK, FFN, VO). More experimental details are provided in Appendix A.3.
LLaMA (0.258) Open
LLaMA (0.258) OpenWeb Text: wed scheduler
Adl-elt (S3)
3310
Lin (UI--)
Adamltt (234)
2:50
Bluckvins LR(C8k)
2.55
Bluckudes (53))
2:00-
2m
2.5
apocckup
LLaMA (0.25B) on OpenWebText
LLaMA (0.5B) on MiniPile
254
mum of staps
50+
12x
3308
25%
num steps.
50c
75c
2004
500-
Adbuss wi (574)
Adoee (86))
validtion bos
3000
3:00
(Bockwies LIT (204)
Arbivs vili (1204)
validation 1.2
Adbee (eci) iF (20))
Figure In pre-training tasks, (left) Lion with Blockwise LR outperforms Lion; (right) when using wsd LR scheduler, AdamW with Blockwise LR outperforms AdamW.
3:00
2.05
2:80
9:75
These experiments demonstrate that Blockwise LR is not limited to accelerating Adam but can also be effectively combined with other optimizers such as Adam-mini and Lion, and LR scheduler such as wsd, while preserving their unique advantages. This finding paves the way for future research exploring the integration of Blockwise LR with other optimization algorithms.
25%
10%
75
10ok
15%
706
4tA
num of steps
num of steps
Figure 6: Adam-mini with Blockwise LR outperforms Adam-mini in pre- trainia tasks.
The results, presented in Figure 6, demonstrate that Block- wise LR achieves 2x speedup Adam-mini. Since vanilla Adam-mini already achieves 2 memory saving compared to AdamW while maintaining nearly the same convergence speed, Adam-mini combined with Blockwise LR achieves both a 2x speedup and 2x memory saving compared to vanilla AdamW. We leave more ablation stud- ies with other optimizers for future work.
7. Conclusion and Outlook
In this paper, we uncovered a sharpness disparity principle among different types of blocks in transformers, as formal- ized in Eq. (1). Notably, this blockwise sharpness disparity persists throughout the entire training process, except during the initial few steps. Building on this discovery, we pro- posed a novel Blockwise LR adjustment principle, which effectively accelerates base optimizers such as Adam W and Adam-mini in L.LM pre-training tasks.
Blockwise LR on Lion. Another memory-efficient opti- mizer is Lion (Chen et al., 2024), which eliminates second- order moments in Adam W. We conduct experiments to ex- plore whether Blockwise LR can also accelerate Lion. We begin by tuning the 1 for Lion baseline, as detailed in Appendix A.3. For Lion with Blockwise LR, we directly apply the ratios r()) in Eq. (5) (note that this originally tuned for AdamW with Blockwise LR). The results, pre- sented in Figure 7 (left), demonstrate that Blockwise LR yields 2x speedup on Lion.
Future works It would be valuable to investigate the ap- plicability of our Blockwise LR to non-LLM tasks, such as computer vision, and its compatibility with other optimizers. such as Muon (Keller et al., 2024) and other like archi- tectures such as Mamba (Gu& & Dao, 2023). Furthermore, our findings open up opportunities to develop other block- adaptive optimization strategies, such as blockwise weight decay and gradient clipping, which could further enhance training efficiency and performance.
Additionally, we evaluate the evolution of the average sharp- ness across different blocks when trained using Lion opti- mizer. As shown in Figure 11 in Appendix A.3, the results

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
Acknowledgements
Cohen, J. M., Damian, A., Talwalkar, A., Kolter, Z., and Lee,
J. D. Understanding optimization in deep learning with central flows. arXiv preprint arXiv:2410.24206, 2024. 3
Lei Wu was supported by the National Key R&D Program of China (No. 2022YFA1008200) and National Natural Science Foundation of China (No. 12288101). Mingze Wang supported by Young Scientists (PhD) Fund of the National Natural Science Foundation of China (No. 124B2028). Junchi Yan and Zhanpeng Zhou were partly supported by NSFC (72342023).
Devlin, J. Bert: Pre-training of deep bidirectional trans-
formers for language understanding. arXiv preprint arXiv: 101.:488000 2018.
Dosovitskiy, A., Beyer, L.,, Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et An image worth 16x16 words: Transformers for image recognition at scale. arXiv 2020.
Impact Statement
preprint arXiv:
This paper contributes advancing the field of deep learn- ing, with focus understanding and improving the pre- training of LLMs. While our work has the potential to impact society in various ways, we do not identify any spe- cific societal consequences that require particular emphasis at this time.
Everett, K., Xiao, L., Wortsman, M., Alemi, A. A., Novak, R., Liu, P. J., Gur, I., Sohl-Dickstein, J., Kaelbling, P., Lee, J., et al. Scaling exponents across parameterizations and optimizers. arXiv preprint arXiv:2447.00882,.20 2024.
2
References
Gao, I.., Biderman, S., Black, S., Golding, I.., Hoppe, T. Foster, C., Phang, J., He, H. Thite, A., Nabeshima, N.. et al. The Pile: An 800GB dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00022 2020. 7, 13
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. GPT-4 technical report. arXiv preprint wrXiv:2303.08774. 2023. 1
George. T., Laurent, C., Bouthillier, X., Ballas, N., and Vincent, P. Fast approximate natural gradient descent in a Kronecker-factored eigenbasis. Advances in Neural Information Processing Systems, 31, 2018. 4
Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebron, F., and Sanghai, S. GQA: Training generalized multi-query transformer models from multi-head check- points. In Bouamor, H., Pino, J., and Bali, K. (eds.), Pro- ceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 4895-4901, Sin- gapore. December 2023. Association for Computational Linguistics. doi: 10. 18653/v1/2025 wwt.:/0w....mmmmmmmmmm..m 298. 13
Gokaslan, and Cohen, V. Openwebtext cor- pus.
http://Skyl 1 (ann07.giithu..it. OpenWebTextCCoppus 2019. 2, 13
Grosse, R. and Martens, J. A Kronecker- cronecker-factored approx- imate Fisher matrix for convolution layers. In Interna- tional Conference Machine Learning, pp. 573-582. PMLR, 2016. 3
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G.,
Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXhv:231..00522, 2023.
Chen, X., Liang, C., Huang, D., Real, E., Wang, K., Pham, H., Dong, X., Luong, T., Hsieh, C.-J., Lu, Y., et al. Sym- bolic discovery optimization algorithms. Advances in Neural Information Processing Systems, 36, 2024. 2, 9,
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford. E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15530 2022. 14
16
Cohen, M., Kaur, S., Li, Y., Kolter, J. Z., and Talwalkar, A. Gradient descent neural networks typically occurs at the edge of stability. International Conference Learning Representations, 2021.
Hu, S., Tu, Y., Han, X., He, C., Cui, G., Long, X., Zheng, Z., Fang, Y., Huang, Y., Zhao, W., et al. Minicpm: Unveiling the potential of small language models with scalable train- ing strategies. arXiv preprint arXiv:2404.06395 2024. 7, 9
Cohen, J. M., Ghorbani, B., Krishnan, S., Agarwal, N., Medapati, S., Badura, M., Suo, D., Cardoze, D., Nado, Z., Dahl, G. E., et al. Adaptive gradient methods at the edge of stability. arXiv preprint arXiv:2207.14484 2022.
Jastrzebski, S., Szymczak, M., Fort, S., Arpit, D., Tabor, J., Cho, K., and Geras, K. The break-even point on
10

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre- Pre-Training
optimization trajectories of deep neural networks. In International Conference Learning Representations, 2020. 3
Information Processing Systems, 35:30950-30962, 2022.
4
Ormaniec, W., Dangel, F., and Singh, S. P. What does it mean to be a transformer? insights from a theoretical hessian analysis. arXiv preprint arXiv:2410. 200::440..008,, 2024. 2, 6, 6
Jumper, J., Evans, R., Pritzel, A., Green, T.,, Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Zidek, A., Potapenko, A., et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):5833-88, 2021. 1
Pesme, S. and Flammarion, N. Saddle-to-saddle dynamics in diagonal linear networks. Advances in Neural Informa- lion Processing Systems, 2023.
Kaddour, The MiniPile challenge for data-efficient lan- guage models. arXiv preprint arXiv:2304.00 008424 2023.
7. 13
Popel, M. and Bojar, O. Training tips for the transformer model. arXiv preprint arXiv:/804.0000 002004000027..4 2018.
Karpathy, A. NanoGPT. https://github: karpathy/ Satiyy/nnnna /wwnonnnnnnPP.. 2022. 3, 13
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019. 2, 4, 13
Keller, et al. Muon optimizer.
https:
//github. , www./eeeeeeeeeaeeeaee readme-ov- 2024. 2.9
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, Exploring the limits of transfer learning with unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-552 2020. 2, 7, 13
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. XrXiv preprint 1412: 2014.
Kunstner, F. Yadav, R., Milligan, A., Schmidt, M., and Bietti, A. Heavy-tailed class imbalance and why adam outperforms gradient descent on language models. arXiv preprint arXiv:2402.1 19449, 2024. 1,
Shin, K. Y. Kim, S., and Moon, S.-M. Initializing the layer-wise learning rate, 2024. URL https:// openreview net/ worurorrmm/mm feerummid=mSs.... 100.mmi002YE.,...
Lei Ba, J., Kiros, J. R., and Hinton, G. E. Layer normaliza- tion. ArXiv e-prints, pp. arXiv-1607, 2016.
Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-Im: Training multi- billion parameter language models using model paral- lelism. arXiv preprint /909, 2019. 6
Liu, A., Feng, B., Xue, B., Wang, B., Wu, B., Lu, C.,, Zhao, C., Deng, C., Zhang, C., Ruan, C., et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437, 2024a.
preprine
Song, M., Ahn, K., and Yun, C. Does sgd really happen in tiny subspaces? arXiv preprint 2rXiv:2405.1600 2024. 3, 6
Liu, H., Li, Z., Hall, D., Liang. P., and Ma, T. Sophia: A scalable stochastic second-order optimizer for language model pre-training. International Conference on Learn- ing Representations, 2024b. 2. 4, 14
Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Roformer: Enhanced transformer with rotary position embedding. Neurocomputingg 568:12706: 2024. 13
Liu, Y., Ott, M., Goyal, N., Du J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., and Stoyanov, V. Roberta: A robustly optimized bert pretraining approach. arXiv preprint 1907. /692, 2019. 7, 13
Team, G., Anil, R., Borgeaud, S., Alayrac, J.-.., Yu, Sori- cut, R., Schalkwyk, J., Dai, A. M., Hauth, A., Millican, K., et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11880 2023. 1
Loshchilov, and Hutter, F. Decoupled weight decay reg- ularization. arXiv preprint arXiv: 2017. 1,
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation language models. arXiv preprint a3X7v:2302.13977, 2023.
Martens, J. and Grosse, R. Optimizing neural networks with Kronecker- factorrd approximate curvature. Interna- tional conference on machine learning, pp. 2408-2417. PMLR. 2015.
2, 2,4,7 7,
Vaswani. A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, E.,, and Polosukhin, I. At- tention is all you need. Advances in neural information processing systems, 30, 2017.
Mi, P., Shen, L., Ren, T., Zhou, Y., Sun, X., Ji, R., and Tao, D. Make sharpness aware minimization stronger: sparsified perturbation approach. Advances in Neural

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
Vyas, N., Morwani, D., Zhao, R., Shapira, Brandfon- brener, D., Janson, L.., and Kakade, S. Soap: Improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409. 11211, 2024. 2
Zhang, B. and Sennrich, R. Root mean square layer nor- malization. Advances in Neural Information Processing Systems, 32, 2019.
Zhang, J., Karimireddy, S. P., Veit, A., Kim, S., Reddi, S.,
Wang, M. and E, W. Understanding the expressive power and mechanisms of transformer for sequence model- ing. Advances in Neural Information Processing Systems, 2024.
Kumar. S., and Sra, S. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33:15383-15393 2020.
Zhang, P., Zeng, G., Wang, T., and Lu, W. Tinyllama: An open-source small language model, 2024a. 13
Wang. M., Wang, J., He, H., Wang. Z., Huang, G., Xiong, F., 7., E, W., and Wu, I.. Improving generalization and convergence by enhancing implicit regularization. arXiv
Zhang, Y., Chen, C., Ding, T., Li, Z., Sun, R., and Luo, Z.- Q. Why transformers need adam: hessian perspective. arXiv preprint arXiv:2402. 16788, 2024b. 1, 2,
preprint arXiv:2405. 20763, 2024. 2, 3, 4, 6, 14
Wen, K., Li, Z., Wang. J., Hall, D., Liang, P., and Ma, T. Understanding warmup-stable-decay learning rates: A river valley loss landscape perspective. arXiv preprint arXiv:24 10.05 2024. 3, 6,
Zhang, Y., Chen, C., Li, Z., Ding, T., Wu, C., Ye, Y., Luo, Z. Z.-Q.. and Sun, R. Adam-mini: Use fewer learning rates to gain more. arXiv preprint arXiv:2406. 26t.:23,6..6693,, 2024c. 1, 9, 16
Wolf, T. Debut, L., Sanh, V.. Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., Davison, J., Shleifer, S., von Platen, P., Ma, C., Jernite, Y., Plu, J., Xu, C., Scao, T. L., Gugger, S., Drame, M., Lhoest, Q., and Rush, A. M. Transformers: State-of- the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38- 45, Online, October 2020. Association for Computational Linguistics. 13
Zhao, J., Zhang, Z., Chen, B., Wang, Z., Anandkumar, A., and Tian, Y. Galore: Memory-efficient Ilm train- ing by gradient low-rank projection. arXiv preprint arXiv:2403.03507 2024a. 2, 13, 14
Zhao, R., Morwani, D., Brandfonbrener, D., Vyas, N., and Kakade, S. Deconstructing what makes good optimizer for language models. arXiv preprint arXiv:2407.07972 2024b. 13, 14
Wu, L., Ma, C., and E. W. How SGD selects the global minima in over- parameterizedd learning: A dynamical stability perspective. Advances in Neural Information Processing Systems, 11:8279-8288, 2018. 3
Xie, X., Zhou, P., Li, H., Lin, Z., and Yan, S. Adan: Adap- tive nesterov momentum algorithm for faster optimizing deep models. arXiv preprint arXiv:2208.0660
2022. 2
Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C.. Zhang, H., Lan, Y., Wang. L., and Liu, T. On layer normalization in the transformer architecture. In Inter- national Conference on Machine Learning, pp. 10524- 10533. PMLR, 2020. 2,
Yang, G., Hu, E. J., Babuschkin, I., Sidor, S., Liu, X., Farhi, D., Ryder, N., Pachocki, J., Chen, W., and Gao, Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.0344 0346000..66,.... 2022.
Yang, Z. Xlnet: Generalized autoregressive pretrain- ing for language understanding. arXiv:/906.0080 1906.08237 2019. 2
arXiv preprint
Yuan, H., Liu, Y., Wu, S., Zhou, X., and Gu, Q. Mars: Unleashing the power of variance reduction for training large models. arXiv preprint arXiv:241 1.14388. 2024.
12

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre- Pre-Training
A. Experimental Details
Models. utilize two popular classes LLM models for our pre-training experiments:
GPT-2. We use GPT-2 (small) model (Radford et al., 2019), implemented via the nanoGPT code base (Karpathy, 2022). Following nanoGPT, the model employs Gaussian Error Linear Unit (GELU) activations and standard Layer Normalization (LayerNNrm). Detailed model configurations are provided in 3.
LLaMA. LLaMA (Touvron et al., 2023) is another popular decoder-only Transformer architecture, incorporating Rotary Positional Encoding (RoPE) (Su et al., 2024), Swish-Gated Linear Unit (SwiGLU), and Root mean square layer normalization (RMSNorm). We pre-train LLaMA models of sizes ranging from 0.13B to 2B parameters. For implementation, for the 1.1B model configuration, we follow TinyLlama (Zhang al.. 2024a), which employs grouped-query attention (Ainslie al., 2023); for other model sizes, utilize the LLaMA code from HuggingFace Transformers Library (Wolf et al., 2020). Additional model configurations are detailed in Table 3 and 4.
Datasets. Models pre-trained the following datasets:
OpenWehText (Gokaslan & Cohen, 2019). It is an opensource recreation of the WebText corpus, is extensively utilized for LLM pre-training such as RoBERTa (Liu al., 2019) and GPT-2.
MiniPile. (Kaddour, 2023). It is 6GB subset of the deduplicated Pile (825GB) (Gao et al., 2020) presents a highly diverse text corpus. Given its diversity, training on minipile poses challenges and potential instabilities.
Colossal Clean Crawled Corpus (C4) (Raffel al., 2020). It is large-scale public language dataset, widely used for LLM pre-training such as T5 (Raffel et al., 2020), and prior pre- trainiiinng studies (Zhao et al., 2024a:b).
All experiments are conducted on 4 A800/H800 80G GPUs
A.1. Training Configurations for Adam Baselines
Table 3: Model configurations and optimally-tuned peak learning rates on OpenWebText and MiniPile.
Acronym
Size
dinode)
dyF
head
depth
1r.max on Open
1r max MiniPile
GPT-2 (small)
124M
768
3072
12
12
6e-4
6e-4
LLaMA (3.)
I34M
768
3072
12
6
1.2e-3
LLaMA (0.25B)
237M
1024
4096
16
8
8e-4
7.5e-4
LLaMA (0.5B)
522M
1280
5120
20
15
8e-4
4.5e-4
LLaMA (0.75B)
743M
1664
6656
26
13
6e-4
LLaMA (L.1B)
1175M
2048
5632
32
22
4e-4
LLaMA (2B)
2025M
2048
8192
32
22
2c-4
Table 4: Model configurations and optimally-tunde peak learning rates on C4.
Acronym
Size
dmadel
dFF
n head
depth
1 maax
LLaMA (66M)
66M
512
2048
le-3
LLaMA (0.2B)
200M
768
3072
le-3
LLaMA (0.4B)
400M
1280
5120
16
12
6e-4
ILaMA (IB)
1004M
1600
6400
25
22
3e-4
As baseline optimizer, use the default AdamW for LLM pre-training, configured with the hyperparameters 3 0.9, Bg = 0.95 and weight decay À = 0.1. To ensure training stability, gradient clipping is applied by norm with threshold 1.0. These settings align with the training protocols used nanoGPT and LLaMA models (Touvron et al., 2023). The default LR strategy integrates a linear warm-up phase, followed by a cosine decay scheduler with the peak learning rate max and the final learning rate r n=l max/20. Additionally,
13

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
OpenWeeText pre-training. The (max) sequence length is 1024, and the batch size is set to 480, following nanoGPT and Liu al. (2024b). The total training duration is 50,000 or 100,000 steps, including 1,000 warm-up steps. The grid search for 1r _r_ is performed over {2e-6 4, 4e 4, 6e 4, 8e 4, 1e-3). Optimal learning rates for each model are detailed in Table 3.
MiniPile pre-training. The (max) sequence length is set to 512, and the batch size set to 300, following Wang et al. (2024). The total training duration is 30,000 or 60,000 steps, including 600 warm-up steps. The grid search for 1r _max is performed over {3e- 4, 4 5e 4, 6e 7..e-..... 4, 9e 4, 1. 2. 1. 1.....5220.. 3}. 5e-3). Optimal learning rates for each model are detailed in Table 3.
C4 pre-training We follow the setup of Zhao et al. (2024a:b), using a sequence length of 256 and batch size of 512. Following the Chinchilla scaling law (Hoffmann et al., 2022), the total number training tokens is set to be approximately 20 times the number model parameters. The training includes I, 1,000 warm-up steps. The grid search for 1r ma is performed over {le 4, 2e-4, 4, 3e-4, 4, 6e 4, 1e 3, 1. 5e-3). Optimal learning rates for each model are detailed in Tables 4. We use the T5 tokenizer, with the vocabulary size 32100.
Baselines: models are pre-trained using AdamW with the respective tuned 1 r__ax for each dataset and model configuration.
Related Experiments.
Blockwise LR Experiments. The baseline results in Figure 4, Figure (right), and Table 2 (the w/o line) are trained following the configurations above.
Sharpness Principle Experiments. Models for Figure 1 (left), Figure 2, Figure 3, are trained using the baseline configurations for GPT-2 (small) or LLaMA (0.25B) on WebTe,, with a total training duration 50,000 steps. In these experiments, the sharpness estimated using (0)) in Eq. (3), with set to 1024. The sharpness distributions and average sharpness values for different blocks are calculated on a logarithmic scale, i.c. log h(@) (l]))).
Additionally, the experiment in Figure 9 employs the same model and sharpness estimator.
Theoretical Analysis Support. To support theoretical insights in Section 4.2, Figure 8 shows the evaluation of the parameter norms across different blocks during training. The model used is LlaMa (0.25B), trained OpenWebTer. The model is LLaMA (0.25B), trained on WpenWebText following the baseline configurations.
25
20-
W
15-
125
Diuines löxir
W
W c
152
130
216
20cc
30k
50k
234
2Mc
30%
40k.
50
12k
233c
300c
num of steps
num of stops
num of stops
illustrate Thcorem 4.1) Norms of input/output weight pa- rameters in FFN and the weight parameters of Norm before FFN, averaged by the number of layers.
(a)
(b) To illustrate Theorem 4.2)
(c) (To illustrate Theorem 4.3) Norms of weight parameters in Emb and the weight parameters in the adjoint Norm layer after Emb.
Norms querykcy/valuc/outpuu parameters in SA and the weight parameters of Norm before SA,
averaged by the number of layers.
Figure 8: Evolution of parameter norms across different blocks during pre-training LLaMA (0.25B) WeenWeTText.
A.2. Experimental Details for Blockwise LR on AdamW
Switching Time. The principle of blockwise sharpness heterogeneity emerges clearly after the initial training phase, as shown in Figure 3. To leverage this principle, in our experiments of AdamW using Blockwise LR, we switch from standard AdamW W to Adam W with Blockwise LR at the end of LRR warmup phase.
14

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre- Pre-Training
10-1
FFN
VD
10-3
10-7
10-6
10-12
layer index
layer index
(a) Average sharpness across different lay- ers. Layer 0 corresponds to the Emb layer. Layers 1, 1,- correspond to the SA-FFN layers.
(b) Average sharpness of the blocks {QK, FFN, VO, Norm across different layers (1 = 1, 8).
Figure 9: In pre-trained LLaMA (0.25B) (with I. 8 layer), there is no clear disparity for the average sharpness across the layers. This is in stark contrast to our our sharpness disparity Principle (1) across the blocks.
Experiments in Figure 4. We adopt the adjusting ratios (5) as the default adjusting ratios for all experiments of Adam W with Blockwise LR.
Experiment on Hyper-pane pere-parrmeter Tuning. We only tune the four adjusting ratios r(.) (+ {Emb, QK, VO, FFN)) in a single small-scale experiment: pre-training LLaMA (0.25B) on Minipile. Specifically, we compare the results under the following configurations of ratios:
r (Emb) 6, (QK) 4, (FF(N 3, r(VO) 2;
rVO) 3; r (Emb) 10, (QK) - 8, (FFFN) rVVO) 4.
(Emb) 8, r(K) -6, 6(F()
The results for the tuning experiments are presented in Figure 10. One can see that the configuration r(Emb) 10, r(KK) 8, r(F(NN 6, r(VO) 4 (Eq. (5)) achieves the largest improvement in terminal loss. Additionally, Blockwise LR demonstrates robustness to these ratios, consistently accelerating pre-training across all tested configurations.
LLaMA (0.25B) on OpenWebTen
3.05
3.00
validation loss
2.50
AdamW
Bleckwise LR (6. 3.
2.80
Blockwise LR 6, 3)
2.75
Blockwise LR (10, 4)
25k num of steps
50k
Figure 10: Pre-training LLaMA (0.25B) on Minipile using AdamW with Blockwise LR across three configurations of adjusting ratios.
Experiments Table We pre-train LLaMA (0.25B) on OpenWeeTex. with a focusing on the three comparisons: (i) applying Blockwise LR exclusively Emb; (ii) applying Blockwise LR both Emb and FFN; (iii) applying Blockwise LR blocks of all the four types (Emb, FFN, QK, and VO). The adjusting ratios are maintained per the tuned in Eq. (5).
A.3. Experimental details for Adam-mini, Lion, and wsd Experiments for Adam-mini.
15

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
Baseline. In the baseline experiments in Figure 6, following Zhang al. (2024c), we adopt the same peak learning rate 1 max tuned for AdamW as the 1 max of Adam-mini.
Hyperparameter tuning. Since Adam-mini uses SGD within each blocks, its dynamics differs significantly from those of AdamW. Thus, for Adam-mini with Blockwise LR, we re-tune the ratios r(o) E (1,2.4) for 1, r(FFN) 4, rvO)) - 4, which
{Emb, QK, FFN, VOJ. The tuned ratios are r(Emb) 4, 4,K(
S(Norm) This discrepancy may stem from the unique dynamics of Adam-mini, particularly its SGD-like behavior within blocks. We leave further investigation for future work.
are used in the experiments in Figure 6. Note that these ratios do not satisfy r(.) x Si.)
Experiments for Lion In the baseline experiments in Figure 7 (left), following Chen et al. (2024), search for the optimal maximum learning rate Ir_max for Lion within the set {1/3, 1/5, 1/10} of the Irumax values tuned for AdamW, using corresponding weight decay values of {0.3, 0.5, 1.0]. Ultimately, we adopt 1/5 of the AdamW-tuned 1r max with weight decay of 0.5. For Lion with Blockwise LR, we directly apply the ratios in Eq. (5) (note that this is originally tuned for AdamW with Blockwise LR).
hanation
Iteration - 100
lterstion I 1000
Bentioo 2 10000
lleration 1 20000
Debed
Inbed
Droed
010e:
Embec
04-
CC
F1N
Figure 11: Evolution of the average sharpness across different blocks during pre-training LLaMA (0.25B) on WpenWeTText using the Lion optimizer. The total training duration is 50k steps. The results closely resemble those in our Figure 3(b), which uses AdamW. Our Principle (Eq. (1)) emerges during the initial phase (from iteration 0 to iteration Ik), which accounts for only approximately 2% of the total steps, and persists throughout the subsequent training process.
Experiments for wsd sheduler. In the experiments in Figure (right), followed by a stable phase where LR remains at 1 r_max (up to 66.7% of the total training steps), and then linear decay to 0. The 1r max is identical to that used in the cosine decay scheduler, as reported in Table 3.
linear warm-up LR to peak 1 Ir_max,
B. Proofs in Section 4
B.1. Proof of Theorem 4.1
We focus on the transformation from X to (11-1/2);
x = (1-1/2) + (F (Norm' (x1- ( w(
From the chain rule, it follows that:
0Q aQ ax
E [1,2}:
! (1) Ox ow!")
aQ aq u
0y(() ( 0-(()
Thus, it suffices to compute
and @ dyt For simplicity, we define:
X E, [X]
X = 1/2)
X atd
X := Norm(X;;)))) = Xxsd O (Iax1 87),
[V
M XNorm W1,
A:= (M),
F = AW2, Y (()) X FF.
16

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre- Pre-Training
where (-) represents the ReLU or Leacky ReLU activation function. We now compute
0Y JW.
and 0Y
Bry
It is straightforward that:
0Y
OF
OF
DA am
aA 8M
(In & W2)
X Norm @ IM);
OW
OW,
DA ΠM OW,
0Y
OF
Nor
OF DA 2M No
0
θy
aA aM e Norm
(In & W
0A
[In
(vec) Wi (diag( vee( X sud)) (Inx1 @ ID)).
2M
For the (Leaky) ReLUU it holds that (2) = (z). Thus, for 2Y TW,
we have:
OY
OF
AQID [Norm W1 o
0A
@ ID-
OW
2M
Now we derive the upper bounds. First, notice that:
))
2 1/2
1/2
X 2
XstalFF
Σ D
VnD:
X Normll |F = Xatd @ (lax1 @ >1)lF S stalle |nxt 7llF S VnD I Ielllllll S nVD IylF
Consequently, we have the following estimates:
iQ
iQ
dY
iQ OY
aA GM
<
=
[I,
W
(N & Ins)
OW
OY
OW
F
F
0Q
0A
aQ
aA
In @ W2 ll || |Noomm & IM <
| [ W 1l [I l|
0Y
oY
aM
0Q
0A
0Q
0A
12llF X |F 3 nVD
W2lle ||yle;
0Y
2M
0Y
aM
0Q
0 0Y
0Y
aQ
aA 6M
X Norm
W
+
s I
aw
aw
TY
FF
F
|
00
aA am
iQ
X Norm
W1
X Norm
Will
0Y
0Y
aM
02
0A
eQ
0A
| |Ip
X Norm Mp - nVD
D
Wi ll lrllp:
0Y
aM
0Y
@M
1:
00
0Q
0Y Oy
07
0Y
F
|F
0Q
aA
Inn
W
(In
W,T) (diag(veo(Xx.) (1nx1 2 Ib))
0Y
aM
00
0A
I X W || I, 8 W [[diag (1nx1 @ ID) l
0Y
8Q
aA
I 12 W 1. 1l Will ||diag |1 1 & Ipll
0Y
17

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre- Pre-Training
aQ
0A
<
Wii
FF
W2
Jp Xsta IF
| 1
0Y
aM
iQ
0A
SnvD
[ ll
[2
0Y
aM
F
Thus, if we define
iQ
aA
V = nVD
Will 1/2ll lyll
0Y
?M
then it holds that:
0Q
V
aQ
V
iQ
v
Ow
[Wil FF
Ow2
W2
0y
|
Therefore,
0Q = 0
V
S(W.) #(W)
E (1,2);
OW. FF
D2
iQ
Tp2
S(>)
#(7)
0y
=
D |
B.2. Proof of Theorem 4.2
We focus on the transformation from (1-1) x(1-1/2);
x x
(1-1) + SA (Norm (t-1/2) (x ((-111) ( w w w
From the chain rule, follows that:
0Q
0Q (X
. (K,
dw!!) (1)
(1-11-22) Ow!!)
0Q
0Q
0X
(y12-1/2)
(1-11122)
0/(1-1/2)*
(-1 Sw('-1/3) and ax For simplicity, we define:
Thus, it suffices to compute
(-(111/2)
X E-
X X 11-1) X std
X Norm Norm(X; X O 7
VV,
X
M N Wo W K Norm
softmax (M), S A. Wr Wo:
VD
Y 11 1/2) X SS
We now compute
8Y aw, and NY
OY
as
as aA 5M
dA
N @ X WK
(I, & W W. X Norm
OW
Q
OW Q
AA aM
OW
2M
VD
OY
0S
8S AA aM
aA
N ? X W Q
(L, W. W. X Norm)
OWk
oWk K
0A 6M OW
2M
VD
18

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre- Pre-Training
0Y
äs
Norr & Wo:
dwv
dwv
OY
A. Wv ID.
awo
dwo
Moreover,
0Y
0Y
Norm
0S NX θr
0
NNom
By
NXNorm
In & Wo WT x T
ÖM (In @ NNrm WxW3 + Knn, (In @ XNom NNomWo WoWK)
0A
VD
Norm
+ A Woww (diag(vec(Xstaa)) @ Ia)).
where Kn is the commutation matrix².
Recalling the proof in Appendix B.1, we have:
[ = VnD,
X 1 S nvDDhyl
Then, similar to the proof in Appendix B.1, we have the following upper bounds:
aQ
0Q
0A
[ |11 [W 11e Woll:
[NNom l1
OWQ
VD
0Y
aM
FF
nIV
D) 3
aQ
aA
[Kll Wvll Wolls
|
VD
0Y
aM
0Q
8Q
aA
[ Wvll Wollp |Norm
OW
VD
YY
FF
F
(nvD) D
0A
[Wolle [W Wolls ||
VD
0Y
aM
0Q
iQ
aQ
|A]]: [Wollr || Normll nVD
||A]| WWoll l
OWv
0Y
0Y
P
F
aQ
IQ
iQ
||A|| [Wv XNomlll S nvD
|||| IWvlllWW
@Wo
0Y
0Y |
|F
|
0Q
iQ
0A aM
+ Www
<
Vn
Xata|Fr
In &
W W X Norm
(In @ XNorm WK W
07
0Y
D
F
00
2
èA
SnVD
D
|Kll Wollo Woll X 11: + ||A||| Wvll |
0Y
VD
@M
iQ
2(nv'D
3
0A 7M
|l Woll, [w Wolp Ihl +n |A nn
0Y
VD
TThe commutation matrix K mm transforms column-wise vectorization into row-wise vectorization.
19

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre- Pre-Training
Therefore, if we define:
(nv 3
IQ
0A 3M
[ Woll, [Wll Wolle lrlp
VD
0Y
V n>
| W [oll Ihlle
0Y
then it holds that:
0
+
0Q
$
S
OWk
llF
dWQ F
|Wall
AQ
V
aQ
V
dWv
W
@Wo
WollF
iQ
2 V
0y
||
F
Therefore,
0Q
$
S(W.) = #(W)
E [K,QH
OW.
D2 I. 1
S(W.) 1
0Q
(V.0);
#WW)
aW.
D2 W. l
F
2
x22 + 422
s(>)
#(>)
Qy
D |>ll
F
B.3. Proof Theorem 4.3
We focus on the transformation from X to = Norm X WWe (1/2)). For simplicity, we define:
Z E, Z
Z = XWE, Zetd ==
Y Norm Norm(Z:7) = Zeed @ (1ax1 @ 77).
Zz Z
straightforward that:
oY
(ieag (veo( (Ztaa) & ID).
Recalling the proof Appendix B.1, we have:
0Y 8y F
S n.DD
aY Then we calculate 8Wg
For simplicity, we denote
31
7 = 2 E.
E Rd
Zd
W
WE & WE
E WE),
WE =
E RR
WE E
E dxD
WF.
E,
20

---

The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
By the proof in (Xiong et al., 2020), for a vector E RIxD, denoted by x E[ac]
then
vD (I 2 1ID)) Thus, we have:
I
= |
2
0Y
0Y 0Zzd oz
oWE
3Zdd dz OWE
VD |اا1|2
& 2
1 ) iE(ml
=In @ diag (vec(>))) diag
I
1
(x @ ID).
|
Recalling the relationship zij = -hhm Wk,j we have E[z,] = Ex F E which implies
2
Fisk Lkk.
Combining this property with the that are one-hot fact of the inputs X, we have:
min /E|n]
12l12 min 11021
ke|d
Additionally, the one-hot encoding ensures:
1/2
1lp
Σ im
= vn.
Now we have the following bound:
OY
OWE E
$ In @ diag (vec(y))llr diag
VD
2T 2; || 12,1
I (1
1 1 likobro))
|x @ Ipllz
||
iE(n)
VD
||>l|
|
SVn (hrlF min
|Xl1 S nvDD D
S nv D
1
min
||
minge(d)
|
If we choose V = nVD Ipl then we have:
OY 0y |F
0Y
V min
<
<
|
ow
[d] |
Therefore,
8Q
222
S(WE) #(WB)
=
OW
Dd min od
|.
(Q
1/2
S(7) #(>)
O
04
D ||>l
21


Reviews:
[{"review_id": "OYVC5Nx30Z", "reviewer": "Reviewer_umnV", "summary": "This paper presents the Sharpness Disparity Principle, which identifies a systematic difference in sharpness across different transformer components. Specifically, the authors find that normalization layers exhibit the highest sharpness, while embedding layers have the lowest, with other blocks lying in between. This pattern emerges early in training and persists throughout.\n\nLeveraging this observation, they propose Blockwise LR, a novel learning rate scaling strategy that assigns different LRs to different block types based on their sharpness. By integrating this strategy into AdamW, they achieve nearly 2x training speed-ups while maintaining (or improving) final loss values across different models (from 0.12B to 1.1B).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your great efforts on the review of this paper and your appreciation. We'll try our best to address your questions.\n\n**Q1: Suggestion for larger-scale experiments.**\n\n**A1**: Thanks for this question.\n- We clarify that our largest model (1.1B) is large enough for current dataset scale (OpenWebText and MiniPile), consistent with standard practices in the community (e.g., [Liu et al, 2023] trains models up to 0.77B on OpenWebText).\n- Our method scales well from 0.12B to 1.1B models, suggesting its potential scalability (see [Fig.R1](https://anonymous.4open.science/api/repo/ICML-2025-Sharpness-Optimization/file/R1_scaling_web.pdf?v=5c61cfe0). We are currently extending our experiments to **2B models**; the results are on the way.\n- To further support the effectiveness, we conducted **additional experiments** on a larger dataset, C4. In [Fig.R2](https://anonymous.4open.science/api/repo/ICML-2025-Sharpness-Optimization/file/R2_scaling_c4.pdf?v=f2e12859), our algorithm consistently achieves lower terminal loss than AdamW across model sizes from 0.2B to 1B, consistent with previous findings.\n\n**Q2: Concerns about the deep cause of the sharpness disparity principle.**\n\n**A2**: Thanks for this question.\n- We clarify that we have provided a first attempt to analyze this sharpness disparity principle from the parameter norms, supported by experiments and discussion following the theorems. Moreover, the discussion in Lines 298–306 offers an intuitive explanation from the multiplicative structure of Transformer blocks.\n- A more rigorous theoretical analysis is beyond this paper’s scope. Even for simpler two-layer neural networks, the analysis of the parameter norms remains challenging [Du et al, 2018; Allen-Zhu et al, 2018]. We leave this as important future work.\n\n**Q3: Suggestion for training steps.**\n\n**A3**: Thanks for this question.\n- We clarify that our training steps are sufficiently large given current model and dataset sizes. For example, 100k steps on OpenWebText yields 480x1024x100k~**50 billion tokens**. According to Tab. 3 in [Hoffmann et al, 2022], most of our experiments *exceed the recommended token budget*.\n- Moreover, this training setup aligns with standard practice in the community, e.g., [Liu et al, 2023; D’Angelo, 2023].\n- Following your suggestion, we conducted a **new experiment** by extending the training of 0.25B LLaMA from 50k/100k to 150k/300k steps. In [Fig.R3](https://anonymous.4open.science/api/repo/ICML-2025-Sharpness-Optimization/file/R3_025B_long.pdf?v=16de8e07), Blockwise LR still achieves lower terminal loss and is 2x faster than AdamW even at longer durations.\n\n**Q4: Concerns about downstream performance.**\n\n**A4**: Thanks for raising this concern.\n- We clarify that In LLM pretraining, downstream performance is widely known to correlate strongly with final pretraining loss and less so with factors like architecture or optimizer choice. See [Du et al, 2024] for a detailed analysis. Thus, the primary focus in LLM pre-training is to reduce the final loss as much as possible, under specific compute and data budgets.\n- In response to your suggestion, we **evaluated downstream performance**. In [Tab.R4](https://anonymous.4open.science/api/repo/SharpnessTable-Downstream/file/R4_evaluate.pdf?v=8dadc1ea), LLaMA trained with our algorithm outperforms the one trained with AdamW across all evaluated tasks.\n\n**Q5: Concerns about comparison with other optimization strategies or optimizers.**\n\n**A5**: Thanks for the constructive question.\n- We clarify that our Blockwise LR is the first successful blockwise LR method for Transformers. As noted in Remark 1.1, traditional layerwise LR methods--originally developed for MLPs and CNNs--have not translated successfully to Transformers.\n- Notably, Blockwise LR is compatible with various optimizers, e.g., AdamW and Adam-mini. Thus, instead of comparing with other optimizers, we follow your suggestions (in Q7) to evaluate Blockwise LR in combination with other optimizers (details in A7). \n- Following your suggestion, we also conducted a **new experiment** using another popular strategy, WSD scheduler [Hu et al, 2024]. In [Fig.R5](https://anonymous.4open.science/api/repo/ICML-2025-Sharpness-Optimization/file/R5_wsd_web.pdf?v=f591fc57), Blockwise LR still achieves a 2x speedup over AdamW in this setting.\n\n**Q6: Question about the precision settings.**\n\n**A6**: All the models are trained with BFloat16. We will clarify it in the revised version.\n\n**Q7: Question about the compatibility with other optimizers.**\n\n**A7**: Thanks for the interesting question. We followed your suggestion and combined Blockwise LR with another popular optimizer, Lion [Chen et al, 2023]. In [Fig.R6](https://anonymous.4open.science/api/repo/ICML-2025-Sharpness-Optimization/file/R6_lion_web.pdf?v=6932b09d), this combination achieves lower terminal loss and 2x speedup over well-tuned Lion.\n\n### Reference\nDue to space limit, see Reference in *Rebuttal to Reviewer w3eV*.", "claims_and_evidence": "Most calims are supported clearly. But I have some concerns:\n - The results focus on smaller-scale pretraining (1.1B max). It remains unclear whether this strategy generalizes to massive-scale models.\n- The cause of this disparity is not deeply analyzed. While the paper suggests that parameter norms influence sharpness, a more in-depth theoretical analysis would be helpful.\n- The number of training steps is not that large. It is better to observe from a larger time scale.", "methods_and_evaluation_criteria": "Yes. Blockwise LR is well-motivated by empirical findings and aligns with existing work on learning rate adaptation. Ablation studies clearly show that increasing LR for norm layers is harmful, supporting the rationale behind the method.\n\nCons:\n- Limited discussion on transferability to larger models. Will Blockwise LR still provide speedups at large scale?\n- Lack of comparisons to alternative optimization strategies. It would be useful to compare against layerwise LR tuning methods.", "theoretical_claims": "Yes. The paper provides some theoretical justification by analyzing blockwise sharpness via the Fisher Information Matrix.", "experimental_designs_or_analyses": "Pros:\n- Results are consistent across different datasets and model sizes.\n- Ablation studies show that modifying LR for Norm layers degrades performance.\n\nCons:\n- Most experiments stop at 100K steps. Does the speedup persist in billion-step training?\n- The study focuses on pretraining loss but does not analyze downstream performance on reasoning/QA tasks.", "relation_to_broader_scientific_literature": "The work is well-grounded.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Would benefit from larger-scale testing and comparisons to other adaptive optimizers.", "other_comments_or_suggestions": "NA", "questions_for_authors": "- What is the precision setting? Is it trained based on mix-precision?\n- Can this be combined with second-order optimizers like Sophia or Shampoo?", "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": "V2qH5TvAGp", "reviewer": "Reviewer_qPBA", "summary": "The paper proposed a blockwise learning rate method to accelerate training. The blockwise learning rate is designed based on blockwise sharpness estimation.  The writing is clear. The principle is reasonable and makes sense to me. The experiments are mostly convincing. Overally speaking, this is a good paper and might motivate more future algorithm designs.  The paper is worth sharing with community.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your great efforts on the review of this paper and your appreciation. We will try our best to address your questions.\n\n**Q1: Questions of the meaning of \"(50k)\" and \"(100k)\".** \"In Figure 4 and 5, what do \"(50k)\" and \"(100k)\" mean? I cannot find the description anywhere around the figure or in Section 6.\"\n\n**A1**: Thank you for pointing this out. 50k and 100k refers to the total training steps. We will provide more explanations in future revision.\n\n**Q2: Questions about the faster convergence of Blockwise LR in the final steps.** \"According to Figure 4, blockwise lr (50k) converges slower than AdamW (50k) in the early stage, and then blockwise lr (50k) converges faster in the final steps. Why would this happen? Any explanation or intuition?\"\n\n**A2**: Very interesting question. This behavior resembles that of WSD schedulers, which often outperform cosine decay in later stages. A preliminary explanation draws on the **river-valley** loss landscape [Wen et al., 2024], which splits the loss into two components: a *river component* (the primary loss along the river at the bottom of the hills) and a *hill component* (additional loss from deviations in height from the river’s course).\n- *Early training*: Blockwise LR boosts LR in river (low-sharpness) directions, enabling faster core progress minimizing the river component,  but also inevitably in some hill (high-sharpness) directions due to noise in data or Hessian estimates—causing larger oscillations and higher loss.\n- *Late training*: As LR decays, the oscillations along the hill component diminish and iterates settle close to the river path [Wen et al., 2024]. Since Blockwise LR made more progress along the river early on, it achieves a lower terminal loss.\n\nWe will add this discussion in the revision.\n\n**Q3: Questions about the total number of training tokens.** \"In Figure 4, the training seems far from convergence. What is the total number of tokens T? Does the advantage of blockwise lr maintains if we train more tokens?\"\n\n**A3**: Thanks for this question.\n- We clarify that our training steps are sufficiently large given current model and dataset sizes. For example, 100k steps on OpenWebText yields 480x1024x100k~*50 billion tokens*. According to Tab. 3 in [Hoffmann et al., 2022], most of our experiments *exceed the recommended token budget*. \n- Moreover, this training setup aligns with standard practice in the community, e.g., [Liu et al., 2023; D'Angelo et al., 2023].\n- Following your suggestion, we conducted a **new experiment** by extending the training of 0.25B LLaMA from 50k/100k to 150k/300k steps. As shown in [Fig.R3](https://anonymous.4open.science/api/repo/ICML-2025-Sharpness-Optimization/file/R3_025B_long.pdf?v=16de8e07), Blockwise LR still achieves lower terminal loss and is 2x faster than AdamW even at longer durations.\n\n\n**Q4: Suggestions for WSD experiments.** \"The current blockwise lr is tested on the top of cosine schedule. How does the proposed blockwise lr reconcile with other lr schedules like WSD? Does the acceleration maintains?\"\n\n**A4**: Thanks for the constructive suggestion. Following your recommendation, we conducted a **new experiment** using the WSD scheduler. As shown in [Fig.R5](https://anonymous.4open.science/api/repo/ICML-2025-Sharpness-Optimization/file/R5_wsd_web.pdf?v=f591fc57), Blockwise LR still achieves a 2x speedup over AdamW under this setting. \n\n\n### Reference\nDue to space limit, see Reference in *Rebuttal to Reviewer w3eV*.", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "yes", "experimental_designs_or_analyses": "yes", "relation_to_broader_scientific_literature": "no", "essential_references_not_discussed": "no", "other_strengths_and_weaknesses": "see below", "other_comments_or_suggestions": "see below", "questions_for_authors": "Q1: In Figure 4 and 5, what do \"(50k)\" and \"(100k)\" mean? I cannot find the description anywhere around the figure or in Section 6.\n\nQ2: According to Figure 4, blockwise lr (50k) converges slower than AdamW (50k) in the early stage, and then  blockwise lr (50k) converges faster in the final steps. Why would this happen? Any explanation or intuition?\n\nQ3: In Figure 4, the training seems far from convergence. What is the total number of tokens T?   Does the advantage of blockwise lr maintains if we train more tokens?\n\nQ4: The current blockwise lr is tested on the top of cosine schedule. How does the proposed blockwise lr reconcile with other lr schedules like WSD? Does the acceleration maintains?", "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": "h2Qst9UZSJ", "reviewer": "Reviewer_jDEo", "summary": "The authors demonstrate that there is a sharpness disparity between the different transformer blocks, which appears early in training and persists throughout. Based on their observation, the authors introduce a novel approach called Blockwise Learning Rate, which adjusts the learning rate of each transformer block based on its sharpness. Blockwise Learning Rate with AdamW achieves lower loss with the same number of gradient steps, and reaches the same loss with nearly half the steps on various sizes of GPT-2 and LLaMA across two datasets. Additionally, the authors demonstrate that Blockwise Learning Rate is compatible with Adam-mini, a memory-efficient variant of Adam.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your great efforts on the review of this paper and your appreciation. We will try our best to address your questions.\n\n**Q1: Concerns about the downstream performance.** \"However, the authors only report the loss and do not evaluate the downstream performance, which would strengthen their evidence.\" \"The performance is only evaluated in terms of loss, which may not necessarily translate to better performance on downstream tasks, so I would like to see some downstream evaluation to be added.\"\n\n**A1**: Thanks for raising this concern.\n- We clarify that In LLM pretraining, downstream performance is widely known to correlate strongly with final pretraining loss and less so with factors like architecture or optimizer choice. See [Du et al., 2024] for a detailed analysis. Thus, the primary focus in LLM pre-training is to reduce the final loss as much as possible, under specific compute and data budgets.\n- In response to your suggestion, we **evaluated downstream performance**. As shown in [Tab.R4](https://anonymous.4open.science/api/repo/SharpnessTable-Downstream/file/R4_evaluate.pdf?v=8dadc1ea), LLaMA trained with our algorithm outperforms the one trained with AdamW across all evaluated tasks.\n\n**Q2: Concerns about the final learning rate.** \"In the literature, it is more common for the final learning rate to be 10% of the peak learning rate instead of 5%, but I do not believe that this invalidates their observations.\"\n\n**A2**: Thank you for the helpful comment.\n- In our experiments, we followed the setup in Sophia  (Sec. 3.1 in [Liu et al., 2023]), where the final LR is set to 5% of the peak LR.\n- As noted by the reviewer, it is more common in the literature for the final LR to be 10% of the peak LR. In our **new experiments** conducted on C4 dataset, we adopt this setting, and the corresponding results are shown in [Fig.R2](https://anonymous.4open.science/api/repo/ICML-2025-Sharpness-Optimization/file/R2_scaling_c4.pdf?v=f2e12859).\n\n**Q3: Concerns about the number of training steps.** \"The number of gradient steps is relatively small (30K to 100K) compared to modern LLMs, but I believe it is just enough to draw conclusions.\" \"Can you extend the training of one of the small models, preferably Llama 0.25B?\"\n\n**A3**: Thanks for this question.\n- We clarify that our training steps are sufficiently large given current model and dataset sizes. For example, 100k steps on OpenWebText yields 480x1024x100k~**50 billion tokens**. According to Tab. 3 in [Hoffmann et al., 2022], most of our experiments *exceed the recommended token budget*. \n- Moreover, this training setup aligns with standard practice in the community, e.g., [Liu et al., 2023; D'Angelo et al., 2023].\n- Following your suggestion, we conducted a **new experiment** by extending the training of 0.25B LLaMA from 50k/100k to 150k/300k steps. As shown in [Fig.R3](https://anonymous.4open.science/api/repo/ICML-2025-Sharpness-Optimization/file/R3_025B_long.pdf?v=16de8e07), Blockwise LR still achieves lower terminal loss and is 2x faster than AdamW even at longer durations.\n\n\n**Q4: Suggestions for \"point-wise feed-forward network\".** \"I suggest replacing \"point-wise feed-forward network\" with either \"position-wise\" or \"token-wise\" as they are more common.\"\n\n**A4**: Thank you for your suggestion. We will revise the terminology in the revised version.\n\n\n**Q5: Suggestions for WSD experiments.** \"Do you know if it also works across learning rate schedulers? Can you apply your Blockwise Learning Rate with warm-stable-decay (WSD) instead of cosine decay?\"\n\n**A5**: Thanks for the constructive suggestion. Following your recommendation, we conducted a **new experiment** using the WSD scheduler. As shown in [Fig.R5](https://anonymous.4open.science/api/repo/ICML-2025-Sharpness-Optimization/file/R5_wsd_web.pdf?v=f591fc57), Blockwise LR still achieves a 2x speedup over AdamW under this setting. \n\n\n### Reference\nDue to space limit, see Reference in *Rebuttal to Reviewer w3eV*.", "claims_and_evidence": "The two claims are (1) there is a sharpness disparity between the different transformer blocks, and (2) adjusting the learning independently for each block based on its sharpness results in faster training speed with respect to the number of gradient steps. The theoretical proofs and experimental evaluation support their claims.", "methods_and_evaluation_criteria": "The evaluation methodology is sound: the models, datasets, and training procedure are standard, and they consider various model sizes, optimizers, and hyperparameters. However, the authors only report the loss and do not evaluate the downstream performance, which would strengthen their evidence.", "theoretical_claims": "I have not reviewed the proofs in detail due to a lack of time. However, the demonstrations in the main paper *appear* to be sound and supported by empirical evidence.", "experimental_designs_or_analyses": "As mentioned above, the experimental design is sound. The authors consider two widely popular models, GPT-2 and LLAMA, with various sizes. The two datasets considered are well established too, although they are small in comparison to modern datasets such as RedPajama or RefineWeb. The training procedure is modern, relying on AdamW with tuned β=(0.9,0.95) and a warmup phase followed by a cosine decay. In the literature, it is more common for the final learning rate to be 10% of the peak learning rate instead of 5%, but I do not believe that this invalidates their observations. The number of gradient steps is relatively small (30K to 100K) compared to modern LLMs, but I believe it is just enough to draw conclusions. The performance is only evaluated in terms of loss, which may not necessarily translate to better performance on downstream tasks, so I would like to see some downstream evaluation to be added.", "relation_to_broader_scientific_literature": "This paper contributes to the field of sharpness analysis and efficient optimizers for transformers. As far as I am aware, this is the firs study to consider the sharpness of blocks rather than layers and to suggest modifying the learning rate per-block across the entire model.", "essential_references_not_discussed": "I am not aware of missing relevant works.", "other_strengths_and_weaknesses": "The figures are clear, especially how the sharpness of each block is depicted in Figure 3. I appreciate that the authors validated their proposed Blockwise Learning Rate across two models, two datasets, two optimizers, and various sizes. Additionally, I appreciate that they tuned the learning rate to ensure optimal performance for AdamW.\n\nThe main weaknesses are the lack of downstream evaluation, which may reveal a closer performance than the loss suggests, and the number of gradient steps, which may close the gap.", "other_comments_or_suggestions": "I suggest replacing \"point-wise feed-forward network\" with either \"position-wise\" or \"token-wise\" as they are more common.", "questions_for_authors": "- Can you extend the training of one of the small models, preferably Llama 0.25B?\n- Can you evaluate a few models on downstream benchmarks?\n- You have shown that the Blockwise Learning Rate works across optimizers. Do you know if it also works across learning rate schedulers? Can you apply your Blockwise Learning Rate with warm-stable-decay (WSD) instead of cosine decay?", "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": "EFvtoD6fiq", "reviewer": "Reviewer_w3eV", "summary": "This paper uncovers a sharpness disparity across different blocks in Transformers, which persists throughout the training process. The authors propose a novel Blockwise Learning Rate strategy to accelerate LLM (e.g., GPT and LLaMA) pre-training. Furthermore, the proposed method consistently achieves lower loss.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your great efforts on the review of this paper and your appreciation. We will try our best to address your questions.\n\n**Q1: Concerns about the gap between estimated digonal Hessian and the true Hessian.** \"The approximation is only performed for the diagonal Hessian matrix, and the gap between the estimated Hessian matrix and the true Hessian matrix is not been controlled.\"\n\n**A1**: Thanks for this question.\n- First, we’d like to clarify that our sharpness measure (Eq. (4)) is based on the trace of blockwise Hessians, for which *only the diagonal Hessians are needed*. This follows from the identity ${\\rm Tr}(H)={\\rm Tr}({\\rm diag}(H))$, which we will clarify in the revised version.\n- To approximate the diagonal Hessian, we adopt the diagonal Fisher matrix. The Fisher is widely regarded as a reasonable approximation of the Hessian in optimization and deep learning [LeCun et al., 1998; Amari, 1998; Martens, 2014].\n\n**Q2: Suggestions for experiments on non-LLM tasks and other optimizers.** \"For non-LLM tasks and other optimizers, the sharpness principle may not hold and need many computations to check.\"\n\n\n**A2**: Thanks for the constructive suggestions. To address them, we conducted two **new experiments**.\n- **Other optimizers.** We evaluated the sharpness disparity principle on LLaMA trained with Lion [Chen et al., 2023]. As shown in [Fig.R7](https://anonymous.4open.science/api/repo/ICML-2025-Sharpness-OPT-V2/file/R7_law_lion.pdf?v=c52ad373), it exhibits *almost the same* principle as with AdamW (Eq. (1) and Fig 3(b)). Then we integrated Blockwise LR into Lion. Remarkably, as shown in [Fig.R6](https://anonymous.4open.science/api/repo/ICML-2025-Sharpness-Optimization/file/R6_lion_web.pdf?v=6932b09d), Lion with Blockwise LR achieves lower terminal loss and 2x speedup over well-tuned Lion.\n- **Non-LLM tasks.** While our primary focus is on language models, as indicated by the title, we followed your suggestion and evaluated the sharpness principle on a ViT-B trained on ImageNet-1k. Surprisingly, as shown in [Fig. R8](https://anonymous.4open.science/api/repo/ICML-2025-Sharpness-OPT-V2/file/R8_ViT.pdf?v=e754a2fa), ViT exhibits a similar sharpness ordering as LLMs: S(QK)<S(FFN)<S(VO)<S(Norm). The only difference is that the embedding layer is no longer the flattest, likely due to structural differences between image and language inputs. These results suggest Blockwise LR can be extended to vision models using this revised principle.\n\n\n**Q3: Questions on the design of Blockwise LR.** \"Why do you not decrease the LRs along high-sharpness directions?\"\n\n**A3**: Thanks for this insightful question. \n- Our design follows the view in [Wen et al., 2024; Wang et al., 2024; Song et al., 2024]: low-sharpness directions primarily drive loss descent, while high-sharpness directions determine training stability. To maintain stability, we keep the learning rate in high-sharpness directions unchanged.\n- Reducing LR in high-sharpness directions may suppress oscillations in these directions but alters the stability condition, and its long-term impact remains unclear. [Wen et al., 2024] shows that using relatively large LR in high-sharpness (\"hill\") directions early in training can result in lower final loss.\n\n### Reference\n\nAllen-Zhu et al. A Convergence Theory for Deep Learning via Over-Parameterization. 2018.\n\nAmari, S. Natural Gradient Works Efficiently in Learning. Neural Computation. 1998.\n\nChen et al. Symbolic Discovery of Optimization Algorithms. 2023.\n\nD'Angelo et al. Why Do We Need Weight Decay in Modern Deep Learning? 2023.\n\nDu et al. Understanding Emergent Abilities of Language Models from the Loss Perspective. 2024.\n\nDu et al. Gradient Descent Finds Global Minima of Deep Neural Networks. 2018.\n\nHoffmann et al. Training Compute-Optimal Large Language Models. 2022.\n\nHu et al. MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies. 2024.\n\nLeCun et al. Efficient Backprop.1998. \n\nLiu et al. Sophia: A scalable stochastic second-order optimizer for language model pre-training. 2024.\n\nMartens, J. New insights and perspectives on the natural gradient method. 2014.\n\nSong et al. Does SGD really happen in tiny subspaces? 2024.\n\nWang et al. Improving Generalization and Convergence by Enhancing Implicit Regularization. 2024. \n\nWen et al. Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective. 2024.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "I did not check the proof carefully, but the theoretical results seem faithful and supported by the empirical results.", "experimental_designs_or_analyses": "Yes, I think the experiments are solid enough to support the authors' claim.", "relation_to_broader_scientific_literature": "NA.", "essential_references_not_discussed": "NA.", "other_strengths_and_weaknesses": "Strengths:\n1. The paper is written clearly.\n2. The paper designs a systematic study on the impact of sharpness in LLM pretraining, and the experiments support their claims well.\n3. The experimental results are promising.\n\nWeaknesses:\n1. The approximation is only performed for the digonal Hessian matrix, and the gap between the estimated Hessian matrix and the true Hessian matrix is not been controlled.\n2. For non-LLM tasks and other optimizers, the sharpness principle may not hold and need many computations to check.", "other_comments_or_suggestions": "NA.", "questions_for_authors": "1. Why do you not decrease the LRs along high-sharpness directions?", "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:

Meta-review: Reviewers agree that the idea on measure sharpness of different layers of the Transformer, and use it to change LR is interesting. The experiments are also quite thorough. Moreover, the training dynamics (slower than Adam first and faster later) also triggers discussions. Overall the work is worth publishing and may draw interest in the community.
Decision: Accept (poster)