Understanding Sharpness Dynamics in NN Training with a Minimalist Example: The Effects of Dataset Difficulty, Depth, Stochasticity, and More
Authors: Geonhui Yoo
Minhak Song
Chulhee Yun
Venue: ICML 2025 poster (2025)
Keywords: progressive sharpening
sharpness
Primary area: deep_learning->theory
TL;DR: We propose a minimalist model that successfully replicates the progressive sharpening and edge of stability phenomena, and empirically and theoretically analyze the effect of problem parameters in progressive sharpening

Abstract: When training deep neural networks with gradient descent, sharpness often increases---a phenomenon known as *progressive sharpening*---before saturating at the *edge of stability*. Although commonly observed in practice, the underlying mechanisms behind progressive sharpening remain poorly understood. In this work, we study this phenomenon using a minimalist model: a deep linear network with a single neuron per layer. We show that this simple model effectively captures the sharpness dynamics observed in recent empirical studies, offering a simple testbed to better understand neural network training. Moreover, we theoretically analyze how dataset properties, network depth, stochasticity of optimizers, and step size affect the degree of progressive sharpening in the minimalist model. We then empirically demonstrate how these theoretical insights extend to practical scenarios. This study offers a deeper understanding of sharpness dynamics in neural network training, highlighting the interplay between depth, training data, and optimizers.

Understanding Sharpness Dynamics in NN Training with a Minimalist Example: The Effects of Dataset Difficulty, Depth, Stochasticity, and More
Geonhui Yoo
Minhak Song? ? Chulhee Yun
Abstract
ing mechanisms to explain its behavior (Ahn et al., 2022; 2023; Arora et al., 2022; Damian et al., 2023; Song & Yun, 2023; Wang et al., 2022; Zhu et al., 2023). Notably, Damian et al. (2023) introduce the self-staailizaaioon mechanism, at- tributing the edge of stability to a negative feedback loop formed by the third-order term in the Taylor expansion of the loss. However, their analysis relies on the assumption that progressive sharpening occurs, i.e., that the sharpness tends to increase along the negative gradient direction.
When training deep neural networks with gradi- ent descent, sharpness often increases-a phe- nomenon known as progressive sharpening- before saturating at the edge of stability. Although commonly observed in practice, the underlying mechanisms behind progressive sharpening re- main poorly understood. In this work, we study this phenomenon using a minimalist model: a deep linear network with single neuron per layer. We show that this simple model effectively cap- tures the sharpness dynamics observed in recent empirical studies, offering simple testbed to bet- ter understand neural network training. Moreover, we theoretically analyze how dataset properties, network depth, stochasticity of optimizers, and step size affect the degree of progressive sharpen- ing in the minimalist model. We then empirically demonstrate how these theoretical insights extend to practical scenarios. This study offers deeper understanding of sharpness dynamics in neural network training. highlighting the interplay be- tween depth, training data, and optimizers.
Although commonly observed in practice, the underlying mechanisms behind progressive sharpening remain poorly understood. Analyzing this phenomenon is particularly chal- lenging due to its dependence on various factors, including network architecture, training data, and optimizers. Cohen et al. (2021) conduct systematic experiments to investigate how these factors influence the degree of progressive sharp- ening, and we summarize their observations in Section 2.
In this work, we study progressive sharpening using a mini- malist model: a deep linear network with a single neuron per layer. We show that this simple model effectively cap- tures the sharpness dynamics observed in recent empirical studies. Furthermore, we empirically demonstrate that our theoretical findings from the minimalist model extend to practical scenarios. Our main contributions are summarized below:
1. Introduction
Understanding the learning dynamics of neural network training challenging due the non-convex nature of its loss landscape. Recent empirical studies have highlighted that when training deep neural networks using gradient de- scent with step size 7). sharpness often increases a phe- nomenon known as progressive sharpening- eventually hovers near 2/q, a regime known as the edge of stability (Jas-
In Section 2, we identify key factors influencing progres- sive sharpening (dataset size, network depth, batch size, and learning rate) based on Cohen et al. (2021), and sum- marize them in Phenomenon 1.
In Section 3, we propose a minimalist model that suc- cessfully replicates sharpness dynamics in deep learning, including the effects of key factors on progressive sharpen- ing (Phenomenon 1) and behavior at the edge of stability.
trzębski et al., 2019; 2020: Cohen et al., 2021).
The sharpness dynamics in deep learning have garnered sig- nificant interest in recent years, with several studies propos-
In Section 4 and Section 5, we provide rigorous theo- retical analysis of the minimalist model. We introduce the concept of dataset difficulty and derive bounds on the sharpness using this quantity. Furthermore, we show that the predicted sharpness from these bounds aligns well with empirical observations, even beyond our theoretical setup.
KAIST AI KAIST Math. Correspondence to: Chulhee Yun <chulhee, yun @ krz>.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
1.1. Related Works
600
500
GF
10-
GF
The dynamics of sharpness during neural network training have gameeeed significant interest in recent years. Jastrzębski et al. (2019; 2020) observe that sharpness increases during the initial phase of training and step size influences sharp- ness along the optimization trajectory. Cohen et al. (2021) formalize the phenomena of progressive sharpening and the edge of stability through extensive controlled experiments, laying the groundwork for subsequent theoretical and em- pirical studies. In this subsection, we briefly mention some of the related works. Refer to Appendix A for more related works.
Sharpees
400-
10
4-2000
0-3850
200
m00000
10-3
100
10"
50
100
150
200
250
50
100
150
200
250
time
tima
(a) Sharpness
(b) Loss
Figure I. GF closely tracks GD dynamics before EoS. Sharpness of GF saturates as loss converges zero. Sharpness of GD saturates as it enters the EoS regime. For experimental details, refer to Appendix C.I.
Progressive Sharpening. Several recent works have explored the mechanisms behind progressive sharpening. Wang et al. (2022) employ the output-layer norm as a proxy for sharpness to explain progressive sharpening. In a two- layer linear network, they prove progressive sharpening under certain conditions, but their characterization is lim- ited to a certain interval and does not specify the limit behavior of sharpness. Agarwala et al. (2023) analyze a quadratic regression model and showed progressive sharp- ening occurs at initialization. Rosenfeld & Risteski (2024) empirically observe that the training dynamics of neural networks are heavily influenced by outliers with opposing signals, suggesting that progressive sharpening arises due to these outliers. However, they do not quantify the degree of progressive sharpening or analyze its correlation with data properties. In addition, they restrict their theoretical analy- sis to synthetic data. Closely related to our work, Marion & Chizat (2024) study sharpness dynamics in deep linear networks and characterize the sharpness of solutions found by gradient flow. However, their analysis focus primarily establishing an upper bound on sharpness, whereas our work provides both lower and upper bounds. Furthermore, we introduce the concept of dataset difficulty and demonstrate its correlation with these bounds.
sharpening depends on problem parameters, based on the observations of Cohen et al. (2021). Their experiments ex- amine the influence of factors such as network architecture and training data on the degree of progressive sharpening, which they quantify as the maximum sharpness along the gradient flow trajectory (GD with infinitesimal step size).
Notably, Cohen et al. (2021) observe that GD with fixed step size 1) closely follows the gradient flow trajectory un- til the sharpness approaches 2/1) (see Figure 1). Once the sharpness reaches 2/1,, the training switches to the Edge of Stability (EoS) regime. In this regime, the GD trajectory deviates from the gradient flow trajectory and instead oscil- lates along the central flow trajectory (Cohen et al., 2025). Thus, gradient flow effectively represents what GD would do if GD didn't 't have to worry about instability (Cohen et al., 2021). Therefore, if we know the sharpness at convergence of GF, we can predict whether the same training of GD enters the EoS regime, based on the step size.
In Phenomenon 1, we summarize how problem parameters influence the degree of progressive sharpening, as observed in Cohen et al. (2021). We focus the mean squared loss setting since, with cross-entropy loss, sharpness decreases at the end of training due to margin maximization, making comparisons less straightforward.
Sharpness Dynamies of SGD. While existing works on progressive sharpening and the edge of stability primarily focus GD dynamics, several recent studies have analyzed how SGD differs from GD. For instance, SGD with a large step size has been observed to operate at stochastic edge of stability, where sharpness stabilizes at a threshold smaller than that of GD (Lee& & Jang, 2023; A Agarwala & Pennington, 2024). Agarwala & Pennington (2024) analyze sharpness dynamics in SGD for quadratic regression model, showing that SGD noise reduces sharpness increase compared to GD.
Phenomenon 1 (Key factors of progressive sharpening). The degree of progressive sharpening depends on the fol- lowing problem parameters:
Progressive sharpening occurs to a greater degree as the size of the training dataset increases. For example, when training on different-sized subsets of CI- FAR 10 using gradient flow, progressive sharpening is more pronounced with larger datasets (see Figure 2a).
Dataset size:
Network depth: Progressive sharpening occurs to
2. Key Factors of Progressive Sharpening
greater degree as network depth increases. For instance, when training fully-connected networks with fixed width and varying depth using gradient flow, deeper networks exhibit more pronounced progressive sharpening (see Fig- ure 2h).
Progressive sharpening refers to the phenomenon where sharpness increases during gradient descent (GD) training. In this section, we discuss how the degree of progressive
2

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
1400
1900
2000
Sharpness
1200
1000
N=10000
N-5000
1200
D=5
(nd
Sharpness
W-1004
W-=248B
N-2000
N-1000
gha
800
800
D-3
1000
W-256
W4-12
200
400
400
200
500
W-66
0
150
time
time
time
(a) Dataset size
(b) Network depth
(c) Network width
Figure 2. Effect of dataset size, network depth, and network width of tanh NN, for experimental details, refer to Appendix C.I.
== 125
= 250
B 500
BCN
220
Learning Rate
200
-1=2/1000
Sharpness
180
1]228000
160
7F22600
140
7F2/400
120
1)-2/000
100
20
40 60
80
100
0
20
40 60 80
100
20 40 60 time
80
100
20 40 60 80
100
time 11" iterations
time
n iterations
iterations
n iterations
Figure 3. Effect of batch size, and learning rate in SGD and GD, for experimental details, refer to Appendix C.I.
Batch size: Progressive sharpening occurs to a greater degree as the SGD batch size increases. When training with SGD, larger batch sizes lead to more pronounced progressive sharpening (see Figure 3).
Rd we use S to denote its orthogonal complement.
Minimalist Model. We consider a simple deep linear network f Rdd - IR, where each hidden layer consists of a single neuron with the identity activation function and D > 2 is the depth. The network is defined
Learning rate of SGD: Progressive sharpening occurs to a lesser degree the SGD learning rate increases, especially when the batch size is small (see Figure 3).
f(z;0) :- (x I
(I)
Cohen et al. (2021) also observe that wider networks ex- hibit less progressive sharpening with cross-entropy loss. However, our experiments suggest that this trend may not hold for squared loss, as shown in Figure 2c, Due to this discrepancy, exclude the network width from the key factors of progressive sharpening.
where 0 = (u, 01, UD-1) represents the collection of all model parameters. Here, E R is the weight of the first layer, and U/ E R is the weight of the (i +i+1 1)-th layer. We use psep= = d+D to denote the total number parameters.
Task. We study the empirical risk minimization problem under the squared loss. Let the training dataset be defined by the data matrix X = [<]
xx]] E RRXxd and
T
3. Minimalist Model for Sharpness Dynamics
the label vector y = [yn ... yN] E N which satisfy y * 0. The loss function L IRP - given by
In this section, we introduce a minimalist model-a deep lin- car network with single neuron per layer-thar that effectively captures the key characteristics of progressive sharpening (Phenomenon 1) as well the edge of stability phenomenon observed in practical setups.
N Σ (f(=;))) -
D-1 Xu
L(0) = 2N
02
- y
2N
==1
For convenience, we abbreviate the network output for the entire dataset as
3.1. Problem Setup
Notation. For positive integer n, use [n] to denote {1..... [I.....n). For a vector ||v|| denotes its Euclidean norm. For a matrix A, let ||Al]2 detnoes its spectral norm, and (A) and (A)(A) denote its column space and row space, respectively. For a symmetric matrix M, let (MM de- note its maximum eigenvalue. For linear subspace of
f(x;0)==(Xu)
1/
and the residual vector as
=(0) = f(X; 0) y
so we can write I(O) = 2/ (X;() 21 - 11)))
3

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
B
B
200
N
3000
Learning Rate
2500
2/18000
sharpness
2/12000
2000
2/8000
1500
2/3000
1000
500
0
10
20
30
10
20
30
10
20
30
10
20
30
time 1n
iterations
time = n) iterations
time = iterations
time = iterations
Figure 4. Effect of batch size and learning rate in our minimalist model (D = 2). We use a 2-label subset of N 1000 from CIFARIO.
sharpness dynamics of neural network training. Here, we empirically showcase that the model not only reproduces the key observations made in Phenomenon 1, but also the self-stabilization dynamics in the edge of stability regime. Our observations with minimalist model suggest that this simple deep linear network could offer useful testbed for understanding the sharpness dynamics in deep learning.
5000
Dm0
Sharpness
4020
N=1000
N-5500
Degg
Nu-300
2000
N=1000
2k
1020
15
100
time
time
(a) Depth
(b) Dataset size
Progressive Sharpening. For Figure 4 and Figure 5a, we trained our minimalist model on 2-label (cat vs dog) subset of CIFARIO (d = 3072) with N 1000, where labels are set to #1 depending on the class. For Figure 5b, we changed N for different runs. In Figure 4, used random-reshufflngg SGD for B N cases, and GD for B = N cases. For Figure 5a and Figure 5b, we used Runge- Kutta 4 algorithm with adaptive step size of 1.0 where S()) measured once every 50 iterations. In Figure and Figure 5, can observe the same trend described in Phenomenon 1 with our minimalist model: progressive sharpening happens to a greater degree with larger datasets, deeper networks, larger batch size, and smaller learning rate (when B is small). While preserving the essential properties of progressive sharpening, the simplicity of our model makes it amenable rigorous theoretical analyses. which present in the subsequent sections.
Figure 5. Effeets of depth and dataset size in minimalist models. All experiments used 2-label subset of N = 1000 from CIFARIO. In (b), except for N 1000, runs terminated after L(0) 10-T
Optimizers. We consider three optimization algorithms:
Gradient Flow (GF); 0(1) = L(((().
S(B)>
Gradient Descent (GD): ((++1) + g(0) - -)V
Mini-batch Stochastic Gradient Descent (SGD):
@(t+1) t- ((t) 17 p(t)),
where we use to denote the learning rate (a.k.a. step size). For SGD, (0; p(O) is the mini-batch loss at step t:
L(,; (()) 2B
(f(:)) -
1
p(0 (0)]?
Edge of Stability. Before we present further investigations into progressive sharpening, we discuss our empirical find- ings on the edge of stability regime here. To see whether our minimalist model also captures the characteristics of sharpness dynamics in this regime, we chose X E R2x2 and y E 2² (N d 2) whose entries were sampled from the standard normal distribution. We trained depth D = 2 minimalist model and observed the evolution of sharpness and training loss throughout the run of GD. For comparison, we also trained a Transformer on SST2 dataset (Socher et al., 2013). For more details, refer to Appendix C.3.
2B
defined by an independently sampled random diagonal ma- trix p(t) E N with exactly diagonal entries chosen uniformly at random and set to 1 and the rest set to O.
Sharpness. The primary goal of this paper is to understand how much the sharpness of the loss increases along the training trajectory. The quantity sharpness S(Ø) at 0 is defined as the maximum eigenvalue of the loss Hessian at 0:
S(0) == Amax 2 (())
In Figure 6a and Figure 6b, we observe that the sharpness dynamics exhibited by the minimalist model closely resem- bles the typical sharpness curve at the edge of stability. The model first goes through progressive sharpening until the
3.2. Minimalist Model Replicates Sharpness Dynamies
We now demonstrate that our minimalist model success- fully replicates most of the interesting phenomena in the

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
1.85 S
700
For simplicity of exposition, assume without loss of generality that y E col(X), which allows the model to attain zero training loss L(0*) = 0 at global minima." The label vector can then be decomposed as
600-
Sharpees
8.6
Sharpees
8
40D-
203-
102
2000
2000
3000
4100
5000
8000
2030
400
0001
1000
2003
Σ dee..
(3)
Stpp
stop
(a) Sharpness of minimalist
(b) Sharpness of Transformer
where di,...,d, ., are scalars. Let W = row(X) span(w1,.. wy). and define II as the projection onto W, and IT as the projection onto its orthogonal comple- ment W Then, the first-layer weight E Rd can be decomposed as
10
10
1
loss
10
10"1
10-4
10
0
1000 2000
3000
4000
DO00
200
420
600
800
step
Step
o + I
(c) Loss of minimalist
(d) Loss of Transformer
Then, we can express the GF, GD, and SGD dynamics in terms of og, di, o, for i E [r] and v..... ................................ UD-1 A detailed derivation is provided in Appendix B. It is worth noting that the updates to u occur only within the subspace W.
Figure 6. Our model captures typical "edge of stability" behaviors. sharpness reaches 2/1,, and then the sharpness oscillates around the threshold 2/1z. Figure 6C and Figure 6d show the loss curve for the two models. As usually seen in the loss curves of practical models, the loss of our minimalist model decreases mon-monotonically with occasional spikes during the edge of stability phase. Replicating this convergence be- havior with minimalist model is intriguing, because many existing studies based on other minimalist models fail to do so (Zhu et al., 2023; Kreisler et al., 2023; Kalra et al., 2025).
Now we introduce a key concept that determines the degree of progressive sharpening.
Definition 4.1 (Dataset Difficulty). The difficulty of a dataset (X, y) is defined as
Q
In practical models, the magnitude of the sharpness oscil- lation around 2/n often decays over time, as can be seen in Figure 6b and also Figure of Damian al. (2023). To the best of our knowledge, the reason for this attenu- ated oscillation is not well-understood. Interestingly, we find that our minimalist model captures this characteristic as well (Figure 6a), suggesting its potential usefulness for deeper theoretical understanding. However, it looks quite challenging analyze this decay theoretically, because we discovered that the sharpness which the rapid sharpness drop starts in fact depends on machine precision, and high precision can even make the loss blow up instead of decay- ing with occasional spikes. See Appendix D for details.
Intuitively, the quantity Q captures the overall difficulty for model to perfectly fit the dataset. For illustration, let us
temporarily consider the linear model: we would like to learn vector B E I2d that satisfies X3 y. From the decompositions in (2) and (3), it is straightforward to check that any solution 8 must satisfy 3* = Σx= di dww, + Hence, each di can be thought of as "distance to travel" for d, the model to achieve e, (XB - v) = 0. The squared sum of such distances thus captures the total amount of effort required to fit the entire dataset.
Notice that Q is only dependent on the dataset, independent of any architecture or optimization algorithm. We will show later that not only Q is useful in our theoretical analysis of minimalist models, but also Q can be used to predict the degree of progressive sharpening in larger fully-connected networks with nonlinear activations. Table shows aver- age values of Q computed for 2-label subsets of CIFARIO, 10, SVHN and Google speech commands datasets. We see that Q increases as N increases.
4. Sharpness at Minimizer of Training Loss
In this section, we provide a theoretical characterization of sharpness S(O*) for any given minimizer 6* of the minimal- ist model. We start by introducing some necessary notation.
We let r == rank(X), and denote the singular value decom- position (SVD) of the data matrix X E TR. Nxd as
x Σ T.e
(2)
'This is without loss of generality, because any y can be de- composed into = W] + y1. with yl E col(X) and yı. E cal(X)*, This then allows decomposing the loss into: T.(0) = z
where 01 .... 2 o, > 0 are the singular values, e; E N are the left singular vectors, and w; E Rd are the right singular vectors.
xIx ITp'' --- 112 s lly l|2. The second constant term becomes the global minimum value of T.(0).
5

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
4.1. Two-layer Linear Networks
can be seen later in Figure 9, we observe in practice that C(f()) is substantially smaller than C throughout training.
We first consider the two-layer case, i.e. when D = 2. Here, we prove upper and lower bounds on the sharpness S(0*) at given solution L(0-) 0 as function of the imbalance between two layers of 6*
Thus, under the assumption C(0*) S Č, we establish the following relationship:
Remark 4.4. The sharpness at a global minimizer S(0*) increases with the largest singular value oi of the data ma- trix, increases with dataset difficulty Q, and decreases with layer imbalance C(0*).
Definition 4.2 (Layer Imbalance). The layer imbalance of the two-layer minimalist model (D = 2) with parameter 0 = (u, 01) is defined by
This observation will play a crucial role in Section 5.2.
((0)= == [[ u| of = 1=1 of.
4.2. Deep Linear Networks
As will be seen in Section 5, layer imbalance is preserved quantity for the trajectory of GF, and slowly increasing quantity for GD and SGD.
Now we extend our analysis to the general case, i.e., min- imalist model with arbitrary depth D 2 2. For the deeper case, we focus on the case where all layers are balanced, for the sake of simplicity. To this end, we first introduce the following assumption.
In the next theorem, we show that the sharpness of a global minimum 0* can characterized using C(O*) and Q.
Theorem 4.3 (Sharpness at minimizer, two-layer case). For a two-layer minimalist model (1) trained on a dataset (X, y) with difficulty Q, the sharpness at a global minimizer 8* (u* vi) of [(0) is bounded by
Assumption 4.5 (Balanced Layers), For a minimalist model of depth 1) 2 2 with parameter 0 = (u, U1 - (D-1), we say the parameter @ is balanced if
Hwwll = |eal =... |UD-I|-
of(1))++ + df
(17)2 SS(0's 1 1/ N
| (f(())2+ Ei= df (ut)²
N
The balancedness assumption is widely adopted in the deep linear network literature (Arora et al., 2018; 2019). In our setting, the key difference is that we consider the norm of the first-layer weight u projected onto the subspace W.
where of satisfies
(ut)² VC 4Q -((0))
Theorem 4.6 (Sharpness at minimizer, general case). For the minimalist model (1) of depth D 2 2 trained on a dataset (X. with difficulty Q let 8 be a global mini- mizer of (()) that is balanced (Assumption 4.5). Then, the sharpness at 6* is bounded by
2
Specifically, if the layers are balanced, i.e., C(0*) = O, then
1 N
02Q +dQQ 1/2] S S(0')
1 N 0/02/22 + (En Q-)
040 D-1
+D--1)dQ b) S(())
N
$ N 170 D=t +
10 11- 1> dl e
Theorem 4.3 provides bounds on the sharpness at global minimizer 6* when the layer imbalance Ci(0*) is known. We can easily observe that if (vi)² 2 o7 Vi d. both the lower and upper bounds of 5(0*) increase with (of)?. Furthermore, (ut)² is an increasing function of the dataset difficulty Q and a decreasing function of the layer imbalance C(f*). Therefore, if C(#*) increases while (ot)² 2 oT ' VI= d. then both the lower and upper bounds decreases. The condition (ut)² 2 oi ' V. di holds if and only if
In Theorem 4.6, the effect of depth is dependent on the scale of Q. If Q > 1, the increase of depth will result in higher sharpness, and vice versa if Q < 1. In datasets of practical size, Q is usually much larger than 1, as Table I suggests.
5. Optimizers of Minimalist Models
So far, we have discussed how dataset difficulty and layer imbalance determine the sharpness at global minimum of the training loss L(l) in our minimalist model. In this section, we consider the trajectory of GF, GD, and SGD on L(O). For GF, using the fact that the layer imbalance is a conserved quantity, we can characterize the amount of pro- gressive sharpening until convergence. Through numerical experiments we show that this prediction aligns well with
Q ((0)ss: VEL d
E T1
dc
Note that C is again only dependent on the dataset, and Table shows its average numerical value for CIFARIO (see Table 8 for SVHN and Google speech commands). As
6

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
Combining Lemma 5.2 with Theorem 4.3 and Theorem 4.6:
Table 1. Average Q for different dataset size N. Samples were taken from two labels (cat vs dog. vs 5, yes vs no) of each dataset, with standard deviation in parenthesis.
Corollary 5.3 (two-layer case). Consider a two-layer min- imalist model (1) trained with the gradient flow. Under Assumption 5.1, the sharpness at convergence S(((cc)) sat- isfies the bounds of Theorem 4.3 with C(0*) - C(
Name CIFARIO
= 100
300
N 1000
0.22(0.04)
44.44(3.32) 1.70(0.16)
SVHN
1.16(0.26)
21.13(2.90)
859.4(67.2)
Corollary 5.4 (general case). Consider a minimalist model (1) trained with the gradient flow using balanced initialization (Assumption 4.5). Under Assumption 5.1, the sharpness at convergence (00())) satisfies the bounds of Theorem 4.6.
Google speech
0.26(0.04)
1.67(0.18)
26.34(2.35)
Table 2. Average O that minimizes upper bound of Theorem 4.3, computed from 2-label subsets of CIFARIO.
N 100
N 300
N 500
N = 1000
These corollaries show that, for GF, we can get the bounds of sharpness convergence in Theorem 4.3 and Theorem 4.6 only based on the information of initialization, without the need to run GF. The bounds obtained for ((0()))) align with the observation in Phenomenon that larger datasets (larger Q) and deeper models (larger D when Q > 1) induce stronger progressive sharpening.
6.0(1.2)
47.4(5.3)
155.9(15.4)
(246.1(100.2)
Table Average Sn of CIFARIO.
= o 1 computed from 2-lahel subsets
D
--100
N-300
= 500
= 1000
2
365(61)
1017(91)
1851(128)
5243(296)
3
283(53)
1111(114)
2465(202)
9876(662)
Our theoretical results so far only characterize the sharpness at convergence. We now present theoretical results based on the initial weight assumption specifically for our two-layer minimalist model, thereby deriving theoretical bounds for both the initial sharpness and the converged sharpness.
4
249(50)
1161(127)
2845(252)
13555(983)
5
231(48)
119((136)
3101(287)
(1391(1244)
Table Average D-1 (E Q) --
t computed from 2-label
subsets of CIFAR1Ó.
Assumption 5.5 (a.3 Initialization). For a two-layer min- imalist model with parameter 0 (u, v1). we say the pa- rameter f uses aB initialization if u ()) ~ N (0, (0,a2T.) and
D
= 100
N = 300
N = 500
N 1000
2
0.77(0.04) 2.18(..18)
(.43()
0.15(0.01)
3.36(0.19)
1.68(0.05)
1.13(0.03)
0.56(0.01)
(0) N (o, (2, where a E R+ and B E R
4
4.42(0.18)
2.63(0.06)
1.96(0.04)
1.16(0.02)
5
5.46(0.18)
3.60(0.07)
2.84(0.05)
1.87(0.03)
Under Assumption 5.5, our analysis of the minimalist model can encompass many widely used weight-initiallzation tech- niques, such as those of Glorot & Bengio (2010), He et al. (2015), and LeCun et al. (2002).
the actual post-training sharpness value in both linear and nonlinear neural networks. For GD and SGD, we study how the discrete and stochastic nature of the algorithms affects the evolution of layer imbalance C(B(+),, which has impli- cations on the sharpness of final solutions (Remark 4.4).
We now present two theorems, one characterizing expected sharpness initialization and the other at convergence.
Theorem 5.6 (Initial Sharpness Bound with ag Initializa- tion). Under the Assumption f.r a two-layer minimalist model trained on a dataset (X, y). the expected sharpness at the initialization (00) (u v1 (0) is bounded by
5.1. Gradient Flow
We introduce following assumption for theoretical results. Assumption 5.1 (Gradient flow converges to global mini- mum). The gradient flow dynamics O(1) converges to the solution ((x): == limy-poc (()), where (.())) = O. We introduce the following useful property of gradient flow:
(a² + 18 S [(000))] [}]{))((00))
N En az0? + fog + +VEi of 222 + d))
Lemma 5.2 (Conservation of layer imbalance under gradi- ent flow). Let O(t) be a gradient flow trajectory trained on the minimalist model (1). Then,
Theorem 5.6 provides bounds of the sharpness at initial- ization under the Assumption 5.5. We can easily observe that both the lower and upper bounds of E[S()())) increase with a² and B², meaning that the larger variance of the initialization scheme results in sharper initialization.
* For 2, the layer imbulance remains constant along the gradient flow trajectory: C(((t)) C(B(0)) for all
t t 0.
* For D > 2 and balanced initialization (00). 0(t) re- mains balanced for all 2 O.
7

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
CIFAR N=100
CIFAR N=:00 (Q=1)
Emplical Sharpness
2000
2400-
2000
1800-
33100-
1300-
1500
1E00-
14000-
10000
0 500
430
330
300
590
400
2000
1430
1600
1830
1900
-W-25461
w
1000
200-
230
600
1400
5500-
3
2000
1200
1400
2500
2000
2500
Depth
Depth
Prodicted Sharpness
Prodicted Sharpnoss
(a) identity activation, depth 5, (b) tanh activation, depth 4, width 2048 in CIFAR 2 lahel width 1024 in SVHN label (IV= 300, Correlation 0.99) (N 100, Correlation 0.81)
Figure 8. Depth vs Sharpness plot. "I" denotes identity activation. "T" denotes tanh activation, and "w" means width.
Figure? 7. Correlation of SD vs empirical ((0(o))), for more re- sults, refer to Appendix C.S.
same experiments on SVHN and Google speech commands; see Appendix C.6. Based on Table and Table 4, we can observe that the term the gap between the upper and lower bounds should be orders of magnitude smaller than ei Q Pot. Therefore, we give a name for the quantity:
of N Q D
dominates both bounds, and
Theorem 5.7 (Sharpness Bound at Convergence with a.3 Ini- tialization). Under the Assumption 5.1 and Assumption 5.5,
for a two-layer minimalist model trained on a dataset (X.y) with difficulty Q and gradient flow, the expected sharpness at convergence 4 is lower bounded by:
Definition 5.8 (Predicted Sharpness). We define the pre- dicted sharpness as the following: Sp == 2 "
E[S(0*)) 2N [of )) (F(C(())))* 4Q
For fully-connected networks of varying width, depth, acti- vation, and dataset size, we trained the model 50 times using different random 2-label subsets of size N from CIFAR1O, SVHN, and Google speech commands. For each run, GF randomly initialized using the default Pytorch initial- ization scheme. All training runs were terminated when L(((t)) < 10-6, and we treated the iterate at termination as ((x)).
+ +( - oî) B(C(O)]]
and upper bounded by:
[[S(0')] s 2N
(09 + Ei
(E(C((*))) + 2ra + 2,3 + 4QQ
Figure 7 shows the scatter plots of final sharpness ((0()))) vs our predicted sharpness Sp on a 5-layer linear network of width 2048 trained on CIFARIO, and tanb-activated 4-layer network of width 1024 trained on SVHN. We emphasize that the predicted sharpness provides a reasonable estimate of the post- post-training sharpness, even though the experiment settings were different from our theory in number of ways?: width was not fixed to 1, activation was nonlinear, and initialization was not balanced.
+ (a -oi) BIC(O+)]]
where E[C((*)] = B[C(((0))] 2-2-33 B2
These bounds are analogous Theorem 4.3. The lower bound of Theorem 4.3 can be reparameterized as follows:
1 N
+
de [op
= IN [0] + ( (C ($$=+4++++00770 -03 -
In Figure 8, we show that our theory captures the effect of dataset difficulty Q and depth D. As expected from predicted sharpness, we observe different effects of depth D on ((((c))), depending on Q > 1 (when N = 300). and Q < (when N = 100). More detailed results are deferred to Table 5 and Table Appendix C.5.
The only difference between the lower bound of Theo- rem 5.7 and Theorem 4.3 is whether we consider the expec- tation or specific value for C(O*). For the upper bound, reparameterization shows similar results, except for the term 2raa + 2,31. Therefore, the same characterization we stated in Remark 4.4 can also be applied to Theorem 5.7.
We detail experiments for Theorem 5.6 and 5.7 in Ap- pendix C.7.
Experiments. To verify if our theory provides a good prediction of the post-training sharpness, we numerically calculate the terms that appear in Theorem 4.6. We used 2-label (cat vs dog) subset of CIFARIO (d = 3072) in Table 3 and Table 4. We randomly selected 50 datasets of size N, while ensuring a balanced distribution between labels. For each dataset, we computed the quantities and report the average and standard deviation. We also run the
5.2. Gradient Descent and Stochastic Gradient Descent
For GD and SGD, we specifically focus on the case of D 2 and address the last two parts of Phenomenon 1: how batch size and step size affect progressive sharpening.
Due to this reason, the scale of empirical sharpness and pre- dicted sharpness in Figure 7 doesn't 't match exactly.
8

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
B
1
Learning Rate
0.8
2/18000
2/12000
0.4
2/8000
0.2
2/3000
-0.2
-0.4
10
20
10
20
30
20
10
20
time = 17 iteratiions
time = 1] iterations
time = 1] iterations
time = n iterations
Figure 9. Effect of batch size and learning rate in minimal models in the dynamics of layer imbalance C(0()), in D D 2 N 1000..
By Theorem 4.3, knowing C(0*) at convergence determines the bound of sharpness. Therefore, we focus on the change of C' after an update of GD and SGD. Starting at the same point 0, let OtD and fSoD be the parameters after step of GD and SGD, respectively. In the theorem below, we show that a step of SGD incurs a greater increase of C' compared to GD. Together with Remark 4.4, Theorem 5.9 sheds light on why SGD induces less progressive sharpening.
Note that B[C(0San) - - is more pronounced for smaller batch size B, highlighting the role of stochastic- ity. Smaller B results on greater increase of C in SGD, resulting in less progressive sharpening. Larger step size y also amplifies this mechanism and leads to even smaller increase of sharpness. In Figure 23-26, we numerically calculate 1 (0(*), 22((()), 22 (0(*) and (2(()) values of GD and SGD, which offers helpful insights on how to interpret Theorem 5.9 and Figure 9. As for GD, although the increase of C is proportional to 172 its L 1 (1(()) and 2 (0(9) of are relatively small, so combined with the factors N² the increase becomes very small and largely unaffected by the step size. For SGD, (2(((() is far larger than V (8(0), and likewise (2(((()) greatly exceeds 22 (0(9). Consequently, C increases much more noticeably when B is small and rj is large. This dynamics of C nicely correlates with the degree of progressive sharpening in Figure 4.
Theorem 5.9 (Increase of Layer Imbalance). For the min- imalist model (1) with D 2, the following holds for the change of C for an update of GD and SGD:
C(OTD) C(0) 222 (2 (0)
B[C((0t()) -C(OG)) -)(N-B) IN-(NII]- v (O))C(()
+ (12( - (
where
(0) = Σ Te)))
6. Conclusion
(0) = N of1= (0)
(0) (x,(0(
e,)oo]² (== N. 11. 0 (e)) 00 (=)) o e,
Throughout this paper, we studied how problem parame- ters influence the sharpness dynamics of neural networks. We introduced a minimalist model that reproduces progres- sive sharpening and edge of stability behavior observed in practice. Theoretically, both initialization and convergence as functions of problem parameters. Empirically, we showed these bounds are nu- merically tight and can predict convergence sharpness under gradient flow. For GD and SGD, we established theorem that offers insights into how batch size and learning rate influence sharpness dynamics.
Also, (0) > 1 (0) 2 0 and (0) > (1 0.
derived sharpness bounds at
The symbol @ in the theorem statement denotes the element- wise product. Theorem 5.9 shows that whenever (0) 5 Ql100) =: (f), GD is guaranteed to increase C and SGD (P())
increases C even more whenever C(#) < (0, (0). In Figure 9, we visualize how C(o(!) evolves in our minimal models with 2-label subset of CIFAR10 dataset. The plots are obtained from the same runs as Figure 4. First, we observe that C()(0) keeps increasing. Also in Figure 21 and 22, we observe that C(((0)) < Ti (0(9)) and C(0(0) 5 (2((()) hold across all settings.
Vi
As step toward broader generalization, Appendix E presents a preliminary theoretical analysis of networks with nonlinear activations. We believe that further extending the problem setup, such as varying the architecture, width, or optimizer, can help bridge the gap between theoretical understanding and empirical observations of progressive sharpening.
There are few important implications of the theorem. In- deed, this theorem does not fully prove that SGD has larger value of C(0*) at convergence and hence smaller final sharp- ness 5(8*) (Remark 4.4) than GD. Nevertheless, the theo- rem offers useful insights on how the batch size B and step size 1i affects the degree of progressive sharpening.
9

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
Acknowledgements
Learning Research, pp- 244-253. PMLR, 10-15 Jul 2018. URL hwt.hhtsssppocceeeeiiiiee ml r. www.ppresseeee..ee vv80 aroral8a.htme
This work was supported by two National Research Foun- dation of Korea (NRF) grants funded by the Korean gov- ernment (MSIT) (No. RS-2023-0021 No. RS-2024- 00421203).
Arora, S., Cohen, N., Golowich, N., and Hu, W. A conver- gence analysis of gradient descent for deep linear neu- ral networks. In International Conference on Learning Representations, 2019. URL https::/p///.. WtRLttttt////enteeeee openreviee.. net net/forum?id=Sae SkM0gg355....
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which feel must be specifically highlighted here.
Arora, S., Li, Z., and Panigrahi. A. Understanding gra- dient descent on the edge of stability in deep learn ing. In Chaudhuri, K., Jegelka, S., Song, L., Szepes- vari, C., Niu, G., and Sabato, S. (eds.), Proceed- ings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 948-1024. PMLR, 17-23 2022. httpt v1 v162/aronee 2022a2rra22...... html.
References
Agarwala, A. and Pennington, High dimensional analysis reveals conservative sharpening and stochastic edge of stability. arXiv preprint arXiv:2404. 22222444449.2.... 19261, 2024.
/ http///poeeeenggeeeeeeeeee mlr. pre..ress/
Agarwala, A., Pedregosa, F,, and Pennington, J. Second- order regression models exhibit progressive sharpening to the edge of stability. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Ma- chine Learning Research, pp. 169-195. PMLR, 23-29 Jul 2023. URL https::::// /ttptttss::/p/ooee http://ooceedeeieeeee mwr. .rrsss/ v202 http//agaaaaaal... a.3b. html.
Cohen, J., Kaur, S., Li, Y., Kolter, J. Z., and Talwalkar, A. Gradient descent neural networks typically curs at the edge of stability. In International Confer- ence on Learning Representations, 2021. URL https: / //openreview net www.mrirrrrrrorr -hw.rfhkGra.....t.--.
Cohen, J., Damian, A., Talwalkar, A., Kolter, J. Z., and Lee, J. Understanding optimization in deep learning with central flows. In The Thirteenth International Conference Learning Representations, 2025. URL https:// openreview net forum i IIErII3ZP.
Ahn, K., Zhang, J., and Sra, S. Understanding the unsta- ble convergence of gradient descent. In International Conference on Machine Learning, pp. 247-257, PMLR, 2022.
Damian, A., Nichani, E., and Lee, J. D. Self-stabilizzation: The implicit bias of gradient descent at the edge of stability. In The Eleventh International Conference on Learning Representations, 2023. URL https:// openreview net / SagrumiiddnnKKAA5550
Ahn, K. Bubeck, S., Chewi, S., Lee, Y. T., Suarez, F,, and Zhang, Y. Learning threshold neurons via edge of stability. In Seirth-eeventh Conference on Neural Information Processing Systems, 2023. URL https: //openreview net forum? wtt.mmdddd:::.ooammmm.
Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, Sharpness-awan minimization for efficiently improving generalization. In International Conference Learning Representations, 2021. URL https: /owpr/vieeeeeeeeeeeeeeeee net wawp:mmmiidmmmmmmmss111...
Andriushchenko, M., Croce, F., Müller, M., Hein, M., and Flammarion, N. A modern look at the relationship be- tween sharpness and generalization. In Krause, A., Brun- skill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, ), Proceedings of the 40th International Confer- ence on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 840-902, PMLR, 23- mlr.
Glorot, X. and Bengio, Y. Understanding the diffi- culty of training deep feedforward neural networks. Teh, Y. W. and Titterington, M. (eds.), Proceed- ings of the Thirteenth International Conference on Ar- tificial Intelligence and Statistics, volume of Pro- ceedings of Machine Learning Research, pp. 249- 256, Chia Laguna Resort, Sardinia, Italy, 13-15 May 2010. PMLR. URL https: h http///pooeeeeeeeee//ee//e/eesee .mlr.
29 Jul 2023. URL httpttt/pppooceeeiiii press/v205 wwwraanscheeenooooooooo htmml
Arora, S., Cohen, N., and Hazan, E. On the optimiza- tion of deep networks: Implicit acceleration by overpa- rameterization. In Dy, J. and Krause, A. (eds.), Pro- ceedings of the 35th International Conference on Ma- chine Learning, volume 80 of Proceedings of Machine
html. htmm.
He, K., Zhang, X., Ren, S., and Sun, J. Delving deep into rectifiers: Surpassing human-level performance on
10

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
imagenet classification. In Proceedings of the IEEE inter- national conference on computer vision, pp- 1026-1034, 2015.
LeCun, Y., Bottou, L., Orr, G. B., and Müller, K.-R. Effi- cient backprop. In Neural networks: Tricks of the trade, pp. 9-50. Springer, 2002.
Holoborodko, P. Mpfr c++, 2010. URL https:// github. com/ advanpix/mpreal.
Lee, and Jang, C. A new characterization of the edge of stability based sharpness measure aware of batch gra- dient distribution. In The Eleventh International Confer- ence on Learning Representations, 2023. URL https: //oopnneview.reeeeee www.eeettooreee forum?id--i wtw.um//------cXX66.K------
Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks. In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R. (eds.), Advances in Neural Information Process- ing Systems, volume 31. Curran Associates, Inc., 2018. URL https://proceeddnngs neurips. cc cc/paper siaapps_ffless/ppeee///////e Sa4bel (eebaaa..oeeee20006c666662222220 pdf.
Li, Y., Wei, C.,, and Ma. T. Towards explaining the regularization effect of initial large learning rate in training neural networks. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Gamett, R. (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https::/ nsgu cc cc/paper cc/paper_files/ /tt/speer_Lieeepppppe///////LLL (oc90oo222ffaaee7701447666666222 pdf.
Jastrzębski, S., Kenton, Z., Ballas, N., Fischer, A., Bengio, and Storkey, A. On the relation between the sharpest directions of DNN loss and the SGD step length. In International Conference on Learning Representations, 2019. URL httttp:://ppereviiii wwwppeeeeewe.eeeeeeeeeeeeeeeeeee net/ forum? id=SkgEa
Marion, P. and Chizat, L. Deep linear networks for re- gression are implicitly regularized towards flat minima. In The Thirty-eighth Annual Conference on Neural In- formation Processing Systems, 2024. URL https: / //operrevi ew. net. forum? d=F738MY 3m4.
Jastrzębski, S., Szymezak, M., Fort, S., Arpit, D., Tabor, J.. Cho*, K., and Geras*, K. The break-even point optimization trajectories of deep neural networks. In International Conference on Learning Representations, 2020. URL https::////....... wwwpp/pwwwwwwweeeeeeooe forum? Beerig87C4KK.
Nacson, M. S., Ravichandran, K., Srebro, N., and Soudry, D. Implicit bias of the step size in linear diagonal neural networks.
Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S. (eds.), Pro- ceedings of the 39th International Conference on Ma- chine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 16270-16295 PMLR, 17-23 Jul 2022. URL /ttps:: http:s::/ppcceeeeeeeeeee mlr press/ v1 62 / ht/p//acccn2aaaa html.
Kalra, D. S., He, T., and Barkeshli, M. Universal sharpness dynamics in neural network training: Fixed point analysis, edge of stability, and route to chaos. In The Thirteenth International Conference on Learning Representations, www.eeiiiiweeee/oounnn
2025. URL https:///// id=v2NOirKn10
Pesme, S., Pillaud-Vivien, I.., and Flammarion, N. Implicit bias of SGD for diagonal linear networks: a provable benefit of stochasticity. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances Neural Information Processing Systems, 2021. URL https: / // //opnreeview net/ forumforumm/mm iidivvii 7ttriikKKHHHA...
Knyazev, A. V. Toward the optimal preconditioned cigensolver: Locally optimal block preconditioned con- jugate gradient method. SIAM Journal on Scien- rific Computiing, 23(2):517-541, 2001. doi: 10.1137/ $1064827500036624 URL https::tt//.o...... http:///o/.oo///
Rosenfeld, E. and Risteski, A. Outliers with opposing signals have an outsized effect on neural network op- timization. In The Twelfth International Conference Learning Representations, 2024. URL https:// openreview net
Kreisler, I., Nacson, M. S., Soudry, D., and Carmon, Y. Gradient descent monotonically decreases the sharp- ness of gradient flow solutions in scalar networks and beyond. In Krause, A., Brunskill, E., Cho, K., En- gelhardt, B., Sabato, S., and Scarlett, J. (eds.), Pro- ceedings of the 40th International Conference on Ma- chine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 17684-17744. PMLR, 23-29 Jul 2023. URL https: /ttpaa/:///////eeeeeeeeee http///pceeenng.eee diggs. mlr.a. .pess// v2.2 http:/keisseeeerr html.
Socher, R., Perelygin, A., Wu, J., Chuang, J., Manning. C. D., Ng, A., and Potts, C. Recursive deep models for semantic compositionality over a sentiment treebank. In Yarowsky, D., Baldwin, T., Korhonen, A., Livescu, K., and Bethard, S. (eds.), Proceedings of the 2013 Confer- ence on Empirical Methods Natural Language Process- ing, pp. 1631-1642. Seattle, Washington, USA, October

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
2013. Association for Computational Linguisties. URL https: / a/cl acllanoolooyy org/D13-1170/.
Song, M. and Yun, C. Trajectory aligmment: Understanding the edge of stability phenomenon via bifurcation the- ory. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems, 2023. URL https: //openreview net. wwr.moormmiii.... dd=-nnaAAA.....
Wang. Z., Li, Z., and Li, J. Analyzing sharpness along GD trajectory: Progressive sharpening and edge of sta- bility. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022. URL http::://ppeerevieeee / Corum?/d=thhIito0xxS
Warden, P. Speech Commands: A Dataset for Limited- Vocabulary Speech Recognition. ArXiv e-prints, April 2018. URL https://arxiv. wttp://wwwr///////////// 03209.
Wu, J., Braverman, V., and Lee, J. D. Implicit bias of gradient descent for logistic regression at the edge of stability. In seirth- Conference on Neural In- formation Processing Systems, 2023. URL ht hts: //openreview. net.ieew.eeeeeeeeeeeeeeeeeeeeeee forum? iidiiidii IT 9MOLYNpO.
Zhu, L., Liu, C., Radhakrishnan, A., and Belkin, M. Cata- pults in SGD: spikes in the training loss and their impact on generalization through feature learning. In Forty- first International Conference on Machine Learning, 2024. URL htttpss///pprreve// id- id-jumG101441.
net/ forum?
Zhu, X., Wang, Z., Wang, X., Zhou, M., and Ge, R. Under- standing edge-of-stability training dynamics with a min- imalist example. In The Eleventh International Confer- ence on Learning Representations, 2023. URL https: // wet.eewww....eeeeeeeeeeeeeeeeeee tttet//oormmmidp7722agg....
12

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
A. More Related Works
Edge of Stability. The edge of stability regime, where sharpness stabilizes near 2/1,, has been extensively studied in recent years (Ahn al., 2022; Arora et al., 2022; Damian et al., 2023; Wu al., 2023). Damian et (2023) propose the self-stabilization mechanism, attributing this phenomenon a negative feedback loop arising from the third-order term in the loss's Taylor expansion. Recent theoretical works have further analyzed training dynamics under simplified models. For instance, Ahn et al. (2023) study loss function of the form (z, v) - (wry), and Song & Yun (2023) extend these results to 2-layer linear networks. Similarly, Zhu et al. (2023) characterize the edge of stability in a 4-layer scalar network, and Kreisler et al. (2023) generalize this analysis to deep scalar networks.
Specifically, Kreisler et al. (2023) consider a scalar linear network with loss ((w), for depth E N and weights w IRD. Their Theorem 3.2 shows that gradient descent does not increase the sharpness of the gradient flow solution initialized at GD iterates (referred to GFS sharpness). Similarly, our Theorem 5.9 shows that C(O) increases over time when mild conditions on C(0) satisfied. Together with our Remark 4.4, these results imply that GFS sharpness decreases training progresses under GD/SGD in our minimalist model.
While these works provide valuable insights, they focus on training with single data point, limiting their applicability to more general settings. In contrast, our minimalist model considers general training data, enabling us to capture how dataset properties influence sharpness dynamics.
Linear Diagonal Neural Networks. Our minimalist model shares similarities with diagonal linear networks in a sparse regression setting. Pesme et al. (2021) show that SGD leads solutions with better generalization than GD. Similarly, our Theorem 5.9 shows that SGD induces less progressive sharpening than GD, leading lower sharpness at convergence. Considering that lower sharpness correlates with improved generalization in diagonal linear networks (Nacson et al., 2022), they both unveil how stochasticity can help generalization.
Connection between sharpness and generalization. SAM (Foret et al., 2021) introduce the hypothesis that minimizing sharpness improves generalization, and its benefit demonstrated practical training. Moreover, GD with large learning rate shown to implicitly find flatter solutions (Cohen et al., 2021), which often generalize better than those obtained with small learning rates (Li et al., 2019). While these works suggest correlation between sharpness and generalization, Andriushchenko al. (2023) show that this relationship data-dependent.
Potential practical implications on learning rate scheduling. The study by Zhu et al. (2024) highlights that the catapult mechanism contributes positively to model generalization, and catapults can be induced by designing proper learning rate schedule. In light of this, predicting sharpness evolution can offer practical value when designing such schedulers.
13

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
B. Technical Details
this section, we present proof' of the main theorems and derivations of key formulas. For simplicity, we use the following notation: for function mapping model parameters scalar, vector, or matrix, we write Z(t) ((((t)) for gradient flow (GF) and (() 2(0lt) for gradient descent (GD) and stochastic gradient descent (SGD).
B.1. Reparameterization the Minimalist Model
In this subsection, we reparameterize the gradient flow (GF), gradient descent (GD), and stochastic gradient descent (SGD) dynamics for the minimalist model (1) in terms of og d,, o, for i [r]. and U1..... PD-1, as introduced in Section This reparameterization serves the foundation for subsequent theoretical proofs.
Network output and residual.
We can decompose network output f(X; 0) into e1, 21.....f Er components by
D-1
D-1 0
f(x:0) = Σ f(x)()e) =
Xu
Bj
6 =
II
V
Ei1
(4)
j=1
where the last equality is obtained by replacing with Di=1 o (eew))) ). Similarly, we can decompose residual =(0) by
=(0) f(;;0)
Σ
0:0;
1/ di Ci.
(5)
Hence, network output (4) and residual (5) can be reparameterized in terms os di, or for i E [r]. and Uj...., Up-1-
Gradient Flow. Recall that the loss at is L(0) =
+
- -11 |f((;
2 (0)1. The GF dynamics is given by
2N
AL
D-1
i(t)
Du (t) =
T =(t)
II (t(t),
(6)
N
J=1
0L
(1) dv (1) =1)) Xu(t) II (t),,, eje 1].
(7)
We can replace X with E os (e:w))) ) in (7) and obtain
(t())
1 N Σ a,(x(t)) (w) u(t))
II
(t) = N Σ (x(((() (:)o:(t) II (t),
4#3
for each [D [D-1]. 1]. Hence, we have
(t) Σ σ. (e
1
=(t) (t
II (t).
(8)
973
Similarly, inner product with w; to both hand sides of (6) and replacing X with Er az(eww))) gives
D-1
ot(t) = -
n. i(t) = - o (e_ x(t))
(t), Vielj.
(9)
Therefore, (8) and (9) together give the reparameterizztioo of the GF dyanmics.
Gradient Descent.
Similar to the GF dynamics, the GD dynamics with step size y can be reparameterized as
((+1)
n Σ
(e 20))
VieW- 1].
(10)
N
1=1
14

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
and
D-1 (cF²(()) ((+1) 0i No oi
(t)
II
(11)
Mini-batch Stochastic Gradient Descent.
Recall that the mini-batch loss L(0; p(t)) at step t given by
(((:))))) L(0; ((0)) 2B
-)
1 1000000 2B 20) 20),
where p(t) E RN-NN is an independently sampled random diagonal matrix with exactly B diagonal entries chosen uniformly at random and set to 1 and the rest set to O.
The update rule of SGD is given by
D-1
u (e+1)
(t
7 x p())(0)
II
(12)
B
==1
(0+1)
(t)
d (=(00)] p(t) Xu
1
11)
vje D - 1].
(13)
B
We can replace X with Ex= oz(erw!!) and rewrite (13) as
(++)) = vg/
11
ai
e p(0))(0)
II 10 0#1
vje
- D-1
1].
(14)
B 1=1
and similarly rewrite (12)
D-1
0f
(++1) (t)
B os (e, p(t)((1)
II o/
[re
(15)
Note that GD is special case of SGD when B = N and P(O) = T, where is an N-by-N identity matrix.
B.2. Proof of Theorem 4.3
Let 0* (u",e]] be global minimizer of L(0) for two-layer minimalist model (1) trained on dataset (X, y) with difficulty Q. We denote of = w. u² for each E [r]. Since L(O*) ===))) the residual =(0*) is zero vector. Combining with (5), we have
=T=((') = u d: =0 Vi [r].
Moreover, we have
C(0*) = [i10)3 - (11)2
de by the definition of the layer imbalance. Substituting न.o. gives
(17)² Q
C(0*)
(wt))
(ut)²
(uq)²
which can be rewritten quadratic equation in (u
C(0)) - 0.
15

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
This quadratic equation has unique positive solution, given by
(vi) - +
+ +
(16)
2
is worth noting that, for a given dataset difficulty and layer imbalance at global minimum 0 = (u", vi), the second-layer d:
weight of and of for each i [r] are uniquely determined up to sign, as given by (16) and of d.es
Based on these facts, we will bound the sharpness S(0*). First, note that the loss Hessian matrix exactly matches with the (normalized) Gauss-Newton (GN) matrix at a global minimum:
V²L(0")= N J(0))
where J(0') 38 -0*)= [Xut Xu*] E RRx a Jacobian matrix of the minimalist model. Hence, the sharpness 0*
matches (up to scaling) the spectral norm of the NTK matrix ((0)) (Jacot al., 2018):
1 [[((0))((0))) 112
S(0*) Amaax V²I(0))) N [[( J(0 11
N
The NTK matrix can be written as
J(0) = (xx(()) + Xu"(") = (wf)²
Σ
ofee +
Σ
o1, 01 en
i=1
$1.82=1
Hence, the NTK matrix at 0 is uniquely determined by the dataset and C(00)) as
J(0)) ( -0(0) + (C(0))² + 4Q
2 Lose + (-c(0))+ (-C0))++CCC)+++
Σ di di i
2
VC00.)+
(17)
Based on (17), we can lower bound the spectral norm of the NTK matrix by
= of((t)² + df
[((*)) [
(of)²
and upper bound by
Ei of (u+)2
[(00)) 112 S ( (v1)²l2 + [Xx" x < of(1)) + 1x of(0))2
Therefore, obtain the desired bound on the sharpness:
N (+))) +
di (1))2
1 s S(())
N (ott)
Σi d (0
B.3. Proof of Theorem 4.6
The proof of Theorem 4.6 is analogous to the proof of Theorem 4.3 presented Appendix B.2. Let 0 (u', of....,... ")b be a global minimizer of L()) for a minimalist model (1) of depth D trained on dataset (X, y) with difficulty Q. We assume that 6* is balanced, i.e., [[ = |el
-Ip-il-- We denote of w. u" for each i [r]. Since L(#*) z (12 the residual z(0)) is zero vector. Combining with (5), we have
D-1
2T²) = 0;0; II
un d, = =0, Wi rr).
j=1
16

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
Moreover, since 0 is balanced,
-0.b..""-1>> 112- o u| (uf)²
Substituting of = dd./(o; [[o; D-1 0f) gives
Σ of 11 D- T (0z)2 (of)² =.. (oh-1)²
Multiplying both sides by (uf)2D-2 gives
120 1)20 Σ d
(ut)2D =.. (1
= Q.
i=1 0
Hence, we can observe that
(u7)² =...= -(1カ-1) = Q
and
d o
(o))² of I (1))
Q' Vi [r]
is worth noting that 0....,,.__.. and 07 ...,, of are uniquely determined up to sign.
Based on these facts, will bound the sharpness S((*). First, note that the loss Hessian matrix exactly matches with the (normalized) Gauss-Newton (GN) matrix at global minimum:
LOL( J(O))
where J(Ø-)= = of (0r) = (II uz) x _u"
__u"²
E Nxp is a Jacobian matrix of the minimalist model.
"I
Hence, the sharpness at 0 matches (up to scaling) the spectral norm of the NTK matrix ((0)) (Jacot et al., 2018):
S(A*) Amax
> (V*L(("))
1 N
1 J(0)))
(((0)))(0~)T
N
The NTK matrix can be written as
D-1
D-1 x (xg))=
D-1
J(O') = x
II j=l
(ut)² +
Σ Xu'(u") x
II (et)2
(18)
5=1
= Q
o E T + +(D - (D-1)QQ
D-2
oi 01g Che
(19)
i=1
f1
= Q
ofee T + +(D
-D--11 - 1)Q D-2
d d Q p
ene 6/2
(20)
$1.2
Q P-1
Σ deee T - 1D-11)Q y
Σ
di dig Ci e
(21)
$1.82
Note that the NTK matrix is also uniquely determined. Based on (18)-(21), we can lower bound the spectral norm of the NTK matrix by
[0(0"))(0) TTOOO) = -7 +(D (D-1
17

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
and upper bound by
D-1
S ||x 1/ II (oi)² +
II ==1
(ot)
qQ "F - 1)Q " [uu xT112
q0 D=t
(D 1)Q 0= |X 12
q " - 1)- 2-2
o((2))2
q D + 1)Q D
QQ
0i " 5 + D 1)Q 1
df
A=
Therefore, we can obtain the desired bound on the sharpness:
00p00 D=t
-DD-1) b s S(00) S
1
D-1
+ (D 1)
Σ c
r
Q b
N
B.4. Proof of Lemma 5.2
We provide proof for D 2 and D) > 2 separately.
(1) 2. It suffices to prove that C(O(L)) = 0 for all / 2 0. Recall that the layer imbalance is defined
Σ ot(()) (t)2.
C(0(t))
Diffcrentiating both sides with respect to gives
c(((t)) ())((())()) (t)
2 Σ σi (c] 2()) (t)) a
2
Σ σ 2()) (,(t)o)(t)
N
N
0,
where we used (8) and (9) for the second equality.
(2) D>>2 2. It suffices to prove that
a
a 5y (vi(1)²) =
i (1)))
at
holds for any
O. Using (8) and (9), we have
0
r 2
r Σ 2o,(t)
2 r
D-1
[I (t))l))))))
Σ
tt (o,(t)?)
(e (()))o,(t)
II v,((),
j=1
and
D-1
-
2 N
Σ
asleT x(t))o,(t) II (t).
Dt
f'=1
for any jE [D- 1]. This completes the proof.
18

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
B.5. Auxiliary Lemmas and Proofs
In this subsection, we provide auxiliary lemmas for the proof of Theorem 5.6 and 5.7.
Oxxk v
vl 0]
Lemma B.1. Given a matrix A
and k E N dimensional arbitrary vector E I
||||2 |||||
Proof. Since A is symmetric, its spectral norm is equal to the maximum absolute eigenvalue of A. Let (x,8) E IR++1 be an cigenvector corresponding to an eigenvalue X, where x E R and y E R. Then the eigenvalue equation
A
A
y
y
becomes
Okxk
x
y
0
x
y
This yields the system:
vy Xx,
v T x Ay-
Case 1: of O. In this case, from the first equation obtain
V.
assuming # 0. Substituting this expression into the second equation gives:
v T
U
Xy
||ll Ay-
Since y * canceling y we get:
||/l22 X2
#||0|
Case 2: 0. If = 0, the first equation becomes:
Ar.
For nontrivial eigenvector (i.e., x + 0), we must have À =
Thus, the eigenvalues are:
A = |ol|| 1=-100] and A = 0 (with multiplicity at least k - 1).
Therefore, the spectral norm of
||||2 = max( ||0||
Lemma B.2. Define
g(x) V 4Q,
where Q 2 0. Then g is convex on IR.
Proof. Compute the first and second derivatives:
x2 + 4Q
4Q
g'(x) Vx² +22+QQ
g"(x)
(o² + 4Q)3/2 2 0.
x² + +QQ
Since g0(x) 2 for all x E R, g convex.
19

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
Lemma B.3 (Sharpness Bounds for Two-Layer Minimalist Model with Arbitrary 0). For two-layer minimalist model trained on dataset (X,y). y), the sharpness at (u,r1) v1) is bounded by
1
2070j00 uf of
207d1 01 o9 of
1 S ((0) S N
u of N
2 2 1
0f
Σ (a,0vu - dd:)²
i=1
Consider a two-layer minimalist model with parameters 0 = (u, v1) and Jacobian J(0). By differentiating the loss twice,
Proof.
(z(0)==(0) = - ((0))((O)) N
2N
8² 20)
where we define (0) E gaxpxp as (.), = 20,00% 2022
for E [NV, J. k E ip] and
(1H0), (0)) j.k lvl.
Therefore, sharpness at 0 is given by
S(0) [J() ((0) ( (H(O), ((M))
We may write J(O)](0) = XXT of Xwu T XT. Using the singular value decomposition XXT =Ei= of Gge and representing Xu Ei=1 0j es, it follows that
((O)((O) of D of ete +
Σ
0,0, 0,0, e.e_
1
Lower bound For any unit vector Rd++1 have
| J() + (H(O), 2 wT (J(0) ((0)))(()) + (H()), z(0)))w. (()))
J(0 [ie; ](
Choosing
straightforward calculation on the term J(O) J(0) shows that
le 100 eT (J0)J(07)² (wi)a + (0 )2+ El ofo (v)) o
Since (ox( 2 (01) )2, follows that
lel J(0)p2 et (JJTT)) 2 (f 0
1
Next, noting that
Oaxd
xT)(0)
(11(0);(0))
xx
0
similar calculation shows that
e =(0) = 20 01 un (o 01 - d1).
Ti 01
where J(0)== =
Dividing by |le 2(0)] yields the term
(o) Po + (z (oi) "f
20 01
Thus, combining these results, we obtain the following inequality:
1 N
(o)
20 la (ur) + 2
(x)) + o(
< S(())
(01) + of
20

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
Upper bound
To provide upper bound, use the triangular inequality as follows:
1
[00+ - N [[((0) N
J(0) + =
By standard norm inequalities.
(()))l = |](0) 112 < |XX (11)2 + 1
Since ||XXT| 1|2 = of and Xule = E (o:)) we have
(J(0) J(0)) S sf((1)²+ + + (o,)).
im
Similarly,
T
(()0) =|==))))==00)>>>>>>> xi| =
Σ o. =0))
a(a000011 - dijuss
(as (a; 0:11 - d.))
The first equality is based Lemma B.I. Therefore, the upper bound for the sharpness is
1 S(0) < N
(11))² + Σ (o:))
Σ
oe a5 (
0000 ocur - -d:))²
B.6. Proof of Theorem 5.6
Note that the initialization follows
1(0)~~ (0, 1 ~
Then,
o (0, E [0,00)}] = =0,
E [ = =32
From the lower bound of Lemma B.3,
N ( +
20 (()) of + oño
2a7
0 (0)_(0) (0) 01
$ S(
(M) 10) + ( (0)
Hence, the expectation of the lower bound becomes
B[S(0())] 2 of (a² + +), N
where the contribution of do (0) (0)
tof (o)p term vanishes by symmetry, and we may drop 201 ((.
2 0 for
(o tof
(0)p
brevity.
From the upper bound of Lemma B.3,
21

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
r
5(0 S
(1) + Σ (
11 - d.) d)) 0 o
(0)
(22)
Applying Jensen's 's inequality to the square-root term in Equation (22) and taking the expectation yields
Σ
σi
(0) 00) 03 0 11 - d:) ):
5
E
o
0i
(0) (0) 0;
V1 -
di) 2
r
E
o
(11²)²
20,070, + to?
(a232ot + offf)
Therefore, taking the expectation to Equation (22) results in the following inequality.
s(
S N
Σ a² + oo +
Σ 1=1
o (o fo? + +d))
Thus, we obtain the overall bounds
of (a
2 + 8 =
s N
Σ
a o + 8 o
Σ of (a (2
N
:=1
This finishes the proof.
B.7. Proof of Theorem 5.7
Our analysis aims to derive the expectation of both the lower and upper bounds of Theorem 4.3. By substituting the expression for (ut and simplifying the resulting terms for the lower bound, we obtain:
S(0)) N (((t))2+ +
d
C(O2)2 +4Q C(0*)
2d
(uT)2
N
2
V C( + 4Q C(O+)
di
(C(+)) + +4Q + C(O))
C(O+)² + 4Q - C(O*)
N
2
2Q
d
( 0/ c(()))
(d
+
(C +4
(23)
2N
Q
For the upper bound, we obtain the following in the same manner:
S(0)) < 1/ N (f(01)2 Σi df
[ (f
Ei
VO(0-)++ 1Q+
[izg df
of
c(0y)
(24)
(ut)2
2N
Q
Q
22

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
To derive expectations for both the lower bound and the upper bound, first characterize [[C((")):
B[C(O*)] - - = E
(o. [o) (0 =ra² -B²
(25)
where the first equality can be directly derived by Corollary 5.3
We characterize E] /C(-))) + 4Q] by introducing two-sided bounds, using the following auxiliary lemma.
We now apply Jensen's inequality to bound B[ /C((*)> + 4Q] from below and above.
Lower bound via convexity. By Lemma B.2, g(x) = v + +4Q is convex. Hence
(E(C(())) < [[a(C((")) = (B(C(O+)))² +4Q s E x[VC(0+}++40)
Upper bound via concavity. The function h(y) V is concave on [0, oo). Let Y = C(07)2 4Q > 0. Jensen's 's inequality then gives
E hYY) & (E[Y)
E 4Q s +
Combining the two yields the desired two-sided bound:
VIEIC(O)))* 4Q = [CC(0+)+
+ [
s <v +
(26)
For numerical bounds for both sides, need to derive E[C((*)?] as follows:
B(C((") =
(0) 2
0,
Σ
2
Σ
2)
(n1²) 2
+ (ufo0)
Σ 0:
(0)
+ ΣΣ Σ
(
2 (000))
-2 2 Σ (d0)) 2)
(v19) + (1f00)
=r (3a*) +
a 2ra² 2 32 + 3,3
= ((2 + 2r)a 4
2ra² + 3
(27)
Combining Equation (26) with Equation (25) and Equation (27), we can derive bounds as follows:
V (ra² (33)- + 4Q E [C(0+)+ s -
+ 2r)a - 2ra2 3 + 4Q.
(28)
Finally, to derive our desired lower bound for the expectation of the sharpness at convergence, we combine the expectation of Equation (23) with Equation (25) and the lower bound of Equation (28) as follows:
23

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
1
df
( -oi c(O))
B[S(0))] 2 2N E
0 +
C(
1
of +
d
B
|C(0+)2 + 10
+
di of B|C(O')]
2N
1
o
df
V
(ra2 82)² + 1Q
df
of (ra² - 330]
2N
Q
Q
We derive the desired upper bound in the same manner.
EisI
df
Ei - of c(Oy
E[S(0') -
E
01 +
v ((0+)2+400
2N
Q
Ei df
+ (21q df
01
2
II
VC(0-)2 + 44
oo B[C(OU]]
2N
[
Di= de
(12 + 2r 2ra2 3/4 1Q
-oq
of (ra? -33)
2N
Q
Q
This finishes the proof.
B.8. Proof of Theorem 5.9
For the proof, we consider the update equations SGD, as it can unify the analysis for both GD and SGD. Starting from and its corresponding values of 01....,0, Oy, the update equations under one SGD update can be written as
Do; 11 0z Pej)1 for E [r].
11 E o= Pe:)os.
Ari
where we write == 2(0) for simplicity. Note that plugging in P = I and B = N recovers the GD update. We now analyze change of C after the update:
C(OSta) - C(0)
(o, + Ao:)² - -(11 + 2v1)²
Σ 2 ot
Σ (20:D;; + (Aoo)))) - (201 + (Avi)
1=1
Σ (2 (2n))².
(29)
Now, substituting the definitions of Ao; and Avi.
27
Σ (Ao,)2 - (Avr)²
n
Σ (2 T Pes)
Eoit 0 pei)o:
(30)
B
For simplification, we denote d a;(2(0) 3 et), and wiP = ==) Peş) for any arbitrary matrix P. RNNN Then,
24

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
we can derive the following exact characterization:
C(05Gn) C(0)
(o.) (Av1)²
27
n
(wF) 2 vi
Σ
vi Oi
B
(w+)** (of
- -2t]
- Σ
Σ
2 p wif oz0s]
B
n²2
Σ (wP)2 11 o/- Σ og
+ Σ
(WP)
ΣΣ
2 P wh
[,0,
Ba
it
n/2
Σ
(vP) 2 C(0) - ΣΣ
(31)
B²
For GD, setting P= I and B = N finishes the proof. For SGD, we need to take the expectation over randomness of P. Recall that P E TR. NRNN a random diagonal matrix with exactly B diagonal entries chosen uniformly at random and set to 1 and the rest set to 0.
To make the expectation easier to calculate, let us define p E 1) satisfying P = diag(p), where diag(-) denotes a diagonal matrix whose diagonal entries are equal to the components of the input vector. With this notation, we can write
Ep [vf [f = Ep [osaz= diag(es)pp diag(e;)=] 0z0,z di) [pp] "] diag(e;)z.
(32)
Now calculate the expectation of the random variable pp Recall that the entries of p are sampled without replacement. For E [N], let (0, 1} denote the i-th entry of p. Then,
B N
and E[P.py] = B(B - 1) N(N 1)
E Elp
Elp,1
for any i # j.
From this, Ep [pp]] is matrix whose diagonal entries are all B and off-diagonals are B(B-1) We can write N(N-1)-
B B(B B(B-1)) ) 1+ ) B(B - 1)
Ep [pp ]
N(N 11 11
N(N -
B(N- B), I+ B(B 1) 11
N(N - 1) I+
(33)
N(N - 1)
where 1 E N is a vector filled with 1's.
Substituting Equation (33) to Equation (32) yields
B(N- B) z T diag(ee))aag(()
B(B 1)
Ep JuP v9 0
N(N 11) diag(e))11T diag(e,);
N(N -1)
B(N B)
N(N -1) (20 et) (x00e)) + B(B 1)
(x Te)(TTe)))
N(N 1)
B(N B)
B(B - 1) N(N 1)
N(N - 1) (159)
40
w v/.
(34)
Combining Equation (31) and Equation (34) together, we get
25

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
(N B) BN -
Ep [C(($an) CC())
3 ((()00) +
+ΣΣ 1100000 - 00oill
131
if -1 1) + BN - 1)
Σ (() ΣΣ
(v/0g - o0:)) 2
131
(2 B) (N -1
N 12 C(0) N ΣΣ 1/ Or - 0i 1
22
(2
+ IP(B-1)
Ew!P + ΣΣ (v - o0)0))
BN - 1)
=.11(0)
=:]](0)
2
(N B)
V1(0)C(0)+ + (0)
(-(V=(0) (1()))
+ (02(0) - (2 (0)))
N
BN - 1)
-(()))-((0)
where the last equality used
N-B + B-1 BN 1) BN(N 1) N
The last thing to show is (22(0) (0) and (2) 2 Q (0). For (2) V (6). it suffices to show 0|zle (z ee.².. If we denote dg : [z 3 es)a. the k-th entry of O 6j, then the inequality boils down to
N
N
N
Σ
aj
i=1
This finishes which is true by Jensen's 's inequality. The same holds for (0) 2 22 (0), we denote a : [6, 00/o,k. 0,1,.. the proof.
C. Experimental Details and More Results
We bootstrapped our own implementation from the work of Damian et al. (2023). We mainly used JAX for all the experiments except for Figure 28. To get the sharpness, we used LOBPCG (Knyazev, 2001) algorithm. If there's 's no specific mention in each experiment for initialization, we replicated pytorch default initialization, which has distribution of
Van
Tan
Also, dataset difficulty Q is calculated as follows to avoid the numerical instability: Q = >E=
de max(o?.10-0)
Code is available at ?
wwwpaaddrrrooooeeesseeeeeee sharpening......g..
C.1. Experimental Details for Figure 1-3
In all cases, we trained tanh activation networks. More details are given as follows:
Figure 1: We trained a width 512, 3-layer NN. For the dataset, we trained a N = 2000 subset of CIFARIO (d = 3072)
dataset. Except for the GF curve in the plot, we trained using GD with learning rates 2/500, 2/350, 2/200.
Figure 2a: We trained width 512, 3-layer NN. For each case we trained N = 10000, N = 5000, N = 2000,
N 1000 subsets of CIFAR10 (d 3072) dataset.
26

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
Figure 2b: We trained width 512, with 3-, 4-, and 5-layer NN. For the dataset, trained N 10000 subset of CIFARIO (d 3072) dataset.
Figure 2c: We trained 3-layer NNs of widths 64, 128, 256, 512, 1024, and 2048. For the dataset, we trained a 10000 subset of CIFARTO (d 3072) dataset.
Figure 3: We trained a width 512, 3-layer NN. For B = N case we used GD, and in other cases, we used random reshufiting SGD (batch size 125, 250, and 500). Both are trained with learning rates 2/1000, 2/800, 2/600, 2/400, and 2/200. For the dataset, we trained on full CIFARIO (d = 3072) dataset.
For Figure and GF curve Figure 1, we used Runge- Mutge-Kutta& 4 algorithm with the adaptive step size of S(Ø) where S(0) measured once every 50 iterations. Note that our choice of experimental setting of GF followed Appendix 1.5 of Cohen et al. (2021). All images have been normalized and standardized on per-channel (RGB) basis.
C.2. Data Processing Details Google Speech Commands
We provide details the preprocessing of the Google Speech Commands dataset (Warden, 2018) in Appendix C.4 and Appendix C.5.
We filtered the audio with equal or more than 1 second duration. These audio files are cropped to second, and then converted to mel-spectrogram with the following code.
mel librosa. feature mewspeet -y=audooaarray sr=16000, n_mels=64) db = librosa power_to_de (mel, ref=np.. max)
Then, mel-spectrograms were normalized and standardized globally.
C.3. Experimental Details for Figure 6
For minimal model training, sampled X, y from pytorch normal random as follows:
1.54099607
0.29342899
1. 1.84452377
X
-2.17878938
0.56843126
y 1. .333599555
for initialization, (0) [0.01 0.01]
1 [0.01]. For GD learning rate we used rj 2/50.
For real-world experiments, we used the same transformer model architecture provided in Damian et al. (2023), which is 2-layer Transformer with hidden dimension 64, and two attention heads. We used GELU activation, and N 100 samples of SST2 with binary MSE loss, using GD learning rate = 2/400.
C.4. Experiments Figure 2-Figure 5 on SVHN and Google Speech Commands Dataset
We present experimental results on SVHN and Google Speech datasets that mirror Figure 2-Figure In particular, Figure 10-Figure 13 reproduce our results in main text, while Figure 14-Figure 17 do so for speech data. All experimental settings are consistent across datasets, with only the dataset itself varying.
These confirm that the observed trends in Figure 2 Figgue 5 hold across CIFAR-10 10, SVHN, and Google speech dataset.
27

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
C.4.1. EXPERIMENTS ON SVHN
7000 N
8000 W
6000
1-000
14k
D
7000
W--248
Sharpness
5000
N=5000
D
6000
W=1024
4000
N-2000
10k-
D=3
5000
WF512
3000
N-00
8k
6k-
3000
4000
256
2000
2000
1000
2k-
1000
W=66
0
100
200
300
400
500
0
50
100
150
200
250
0
50
100
150
200
250
300
time
time
Datasct size
(b) Network depth
(c) Network width
Figure 10. Effect of datasct size, network depth, and network width of tanh NN, SVHN dataset
B = 125
B = 250
B = 500
B=N
1800
Learning Rate
1600
-=22/0000.0
Sharpness
1400
n=2/4000.0
1200
1000
n=2/2000.0
800
n=2/1500.0
600
n=2/1000.0
400
200
100
150
100
150
100
150
100
150
time n iterations
time n iterations
time n iterations
time = n iterations
Figure 11. Effect of batch size, and learning rate in SGD and GD of tanh NN, SVHN dataset.
B 25
200
B=N
20k
Learning Rate -=2/15000
Sharpness
15k
7=2/20000
n=2/25000
10k-
n=2/30000
n=2/40000
5k
04
50
100
50
100
0
50
100
50
100
time = iterations
time iterations
time r iterations
time = iterations
Figure 12, Effect of batch size, and learning rate in SGD and GD of minimal model, SVHN dataset.
250k
15k
N=50C
Sherpness
200k
1.kk N-300
150k
D-3
100k
6kc
40
100
time
time
Dataset size
(b) Network depth
Figure 13, Effect of dataset size, network depth of minimal model, SVHN dataset
28

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
C.4.2. EXPERIMENTS ON GOOGLE SPEECH COMMANDS
2000
700
Sharpness
1500
N=5000
N-10000
Sharpness
1500
2000
5
D=4
Sharpness
500
800
W-1024
W-2048
1000
N-2000
D=3
400
W=512
N=1000
W-256
500
500
200
W=84
0
0-
100
200
400
600 800
1000 1200 1400
200
400
600
800
0
50
100
150
time
time
time
(a) Datasct size
(b) Network depth
(c) Network width
Figure 14, Effect of dataset size, network depth, and network width of tanh NN, Google speech commands dataset
B = 125
B 250
500
B=N
Learning Rate
170
n=2/800.0
Sharpness
160
www.mmm
n=2/600.0
-=2/400.00
140
Q=2/200.0
130
n=2/150.0
110
100
50
100
150
150
0
50
100
150
50
100
150
= n iterations
time = iterations
time n iterations
n iterations
Figure 15. Effect of batch size, and learning rate in SGD and GD of tanh NN, Gongle speech commands dataset
=
B
B=N
1600
Learning Rate
M
-=2/1500
Sharpness
1400
n=2/2000
1200
n=2/3000
1000
7=2/4000
n=2/6000
600
0
10
20
30
0
10
20
30
10
20
30
10
20
30
time = iterations
time = iterations
time = iterations
time = iterations
Figure 16. Effect of batch size, and learning rate in SGD and GD of minimal model, 2-label subsets of Google speech commands dataset (yes no)
3500
10k
3000
1=1000
Sharpness
Dm5
2500
2000
N=300
6A
D=3
N-100
1.00
50
150
time
time
(a) Dataset size
(b) Network depth
Figure Effect of dataset size, network depth of minimal model, 2-label subsets of google speech commands dataset (yes vs no)
29

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
C.5. More Detailed Experimental Results for Figure 7 and Figure
We present more detailed experiments omitted from Figure 7 and Figure
C.5.1. EXPERIMENTS ON 2-I.ABEI. (CAT vs DOG) SUBSETS OF CIFARIO
Table 5 shows values of correlation between empirical ((0())) vs predicted sharpness Sn and average empirical ((0 for 2-label subsets of CIFARIO with N 100 and N 300.
Table 5. Detailed results of GF experiments on 2-label subsets of CIFARIO.
(a) N 100, correlation of empirical ((0(r)))
predicted sharpness Sp
activation
identity
tanh
SiLU
ELU
width
512
1024
2048
512
1024
2048
512
1024
2048
512
1024
2048
depth=3
0.99
0.99
0.99
0.90
0.93
0.95
0.90
0.90
0.94
0.96
0.96
0.97
depth=4
0.99
0.99
0.99
0.89
0.93
0.95
0.84
0.89
0.92
0.95
0.96
0.97
depth=5
1.00
0.99
0.99
0.89
0.93
0.95
0.83
0.87
0.91
0.95
0.96
0.97
(b) . 300, correlation of empirical S(0(co)) predicted sharpness SD
activation
identity
tanh
SiLU
ELU
width
512
1024
2048
512
1024
2048
512
1024
2048
512
1024
2048
depth=3
0.99
0.99
0.99
0.74
0.76
0.83
0.86
0.86
0.87
0.85
0.89
0.91
depth=4
0.99
0.99
0.99
0.73
0.75
0.85
0.85
0.84
0.86
0.85
0.89
0.90
depth=5
0.99
0.99
0.99
0.80
0.78
0.84
0.80
0.82
0.85
0.87
0.89
0.90
(c) N
100, average empirical ((((co)
activation
identity
tanh
SiLU
ELU
width
512
1024
2048
512
1024
2048
512
1024
2048
512
1024
2048
depth=3
495
518
566
248
266
293
97
97
101
266
283
313
depth=4
478
493
527
236
250
270
73
71
72
240
254
273
depth=5
468
479
505
231
243
259
62
59
58
228
239
255
(d)
300, average empirical ((0(co))
activation
identity
tanh
SiLU
ELU
width
512
1024
2048
512
1024
2048
512
1024
2048
512
1024
2048
depth=3
1471
1502
1564
510
552
615
236
234
239
603
640
690
depth=4
1701
1730
1791
573
620
680
215
204
204
647
683
724
depth=5
1871
1897
1961
624
670
733
203
188
183
686
721
763
30

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
C.5.2. EXPERIMENTS ON 2-LABEL (3 vs 5) SUBSETS SVHN
Table 6 shows values of correlation between empirical ((0()))) vs predicted sharpness Sp and average empirical (((()))) for 2-label subsets of SVHN with N 100.
Table Detailed results of GF experiments 2-label subsets of SVHN.
(a) N=100, correlation of empirical ((()))) vs SD
activation
identity
tanh
SiLU
ELU
width
512
1024
2048
512
1024
2048
512
1024
2048
512
1024
2048
depth=3
1.00
1.00
1,00
0.74
0.82
0.86
0.78
0.79
0.83
0.85
0.89
0.92
depth=4
1.00
0.99
1.00
0.78
0.81
0,86
0.72
0.76
0.78
0.86
0.87
0.91
depth=5
0.99
0.99
1.00
0.80
0.84
0.87
0.66
0.73
0.77
0.86
0.88
0.91
(b) N=100, average empirical S(0(oo))
activation
identity
tanh
SiLU
ELU
width
512
1024
2048
512
1024
2048
512
1024
2048
512
1024
2048
3epth=3
2759
2823
2946
1021
1086
1164
497
476
474
1029
1115
1219
depth=4
3099
3153
3263
1113
1172
1241
478
432
1059
1134
1218
depth=5
3352
3395
3493
1179
1246
1300
474
408
368
1089
1166
1240
C.5.3. EXPERIMENTS ON 2-LABEL (YES NO) SUBSETS GOOGLE SPEECH COMMANDS
Table 7 shows values of correlation between empirical (((o)) predicted sharpness Šn and average empirical S(0(co)) for 2-label subsets of Google speech commands with N 100.
Table 7. Detailed results of GF experiments 2-label subsets of Google speech commands dataset(yes vs no).
(a) N=100, correlation of empirical S(O(zc)) vs Sp
activation
identity
tanh
SiLU
ELU
width
512
1024
2048
512
1024
2048
512
1024
2048
512
1024
2048
depth=3
0.98
0.98
0.96
0.75
0.87
0.88
0.60
0.68
0.77
0.78
0.87
0.88
depth=4
0.98
0.98
0.98
0.74
0.86
0.90
0.61
0.64
0.73
0.81
0.85
0.89
depth=5
0.98
0.99
0.98
0.78
0.85
0.89
0.64
0.64
0.71
0.83
0.86
0.89
(b) N=100, average empirical ((0()))
activation
identity
tanh
SiLU
ELU
width
512
1024
2048
512
1024
2048
512
1024
2048
512
1024
2048
depth=3
500
532
598
252
282
328
178
162
154
303
323
366
depth=4
491
514
561
238
265
301
151
133
119
278
295
325
depth=5
487
504
541
230
256
287
136
116
102
261
280
304
31

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
C.6. More Experiments for Table 2-4 2-label Subsets of SVHN (3 vs 5) and Google Speech Commands (yes vs no)
We present more detailed experimental results omitted from Table to Table 4. All experiments are carried out using 2-label (3 5, yes no) subsets of SVHN and Google speech commands dataset.
Table 8 shows average values of Č' computed for 2-label subsets of SVHN and Google speech commands of different N's. omitted from Section 4.1. We observe the same drastic growth of as N increases.
Table Average Ö that minimizes upper bound of Theorem 4.3, computed from 2-label subsets of SVHN and Google speech commands.
Dataset
100
N=300 897.2A(127.75) 43.96(4.86)
N 500
N 1000
SVHN
49.29(11.76) 7.10(1.12)
4167.96(537.27) 119.38(10.01)
(6556.80(30155.00 688.18(60.07)
Google speech commands
Table 9 and Table 10 show values of the two terms that appear in Theorem 4.6, calculated using 2-label subsets of SVHN and Google speech commands. We observe the same trend as Section 5.1: the predicted sharpness Šp dominates the bounds, making it reliable proxy for the final sharpness.
Table Average So
of Q computed from 2-label subsets the following datasets
(a) SVHN, with standard deviation parenthesis
D
N = 100
N = 300
N = 500
N
1000
1936.80(302.71)
8264.30(712.54)
17851.97(1343.15)
12935.36(2685..
3
1988.28(379.45
13748. 10(1470.64)
38350.34(6400447)
(63272.16(10276.97))
2015,60(419.45
17735.64 (0735.64((084.52)
286762.97(19824.5)
2032.51(443.88)
70724.27(7841.86)
402067.19(29295.22
(b) Google speech commands, with standard deviation in parenthesis
= 100
300
= 500
N 1000
362.74 (34.85)
891.27 (57.97)
1468.31 (73.87)
3500.62 (151.26)
291.73 (34.35)
971.45 (77.30)
1890.28 (117.04)
6038.67 (345.34)
4
261.68 (33.80)
1014.34 (88.64)
2144.92 (145.91)
7931.86 (510.36)
5
245.18 (33.38)
1041.02 (96.00)
2313.93 (166.05)
9342.15 (641.56)
Table 10. Average Pwl
EE=
d)) Q t computed from 2-label subsets of the following datasets
(a) SVHN, with standard deviation parenthesis
D
N
300
N 500
N 1000
2
0.95(0.12)
0.22(0.02)
0.10(111
0.03(0.00)
3
1.93(0.16)
0.73(0.03)
0.44(0.02)
0.21(0.01)
4
2.92(0.18)
1.4000005)
0.96(0.03)
0.55(0.01)
5
3.91(0.19)
2.18(0.06)
1.60(0.04)
1. 1.04(0.02)
(b) Google speech commands, with standard deviation in parentheses
W= 100
300
= 500
N 1000
1.94(0.14)
0.78(0.04)
0.47(0.02)
0.20(0.01)
3.11(0.15)
1.69(0.06)
1.21(0.03)
0.67(0.02)
4.18(0.16)
2.64(0.07)
2.06(0.04)
1.33(0.03)
5.21(0.15)
3.61(0.08)
2.96(0.05)
2.08(0.04)
32

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
C.7. Experiments for Theorem 5.6 and Theorem 5.7
To numerically validate the Theorems 5.6 and 5.7 and assess the tightness of our bounds, report Table 11-15, and Figure 18 20. We use a² 37 3 3 that the initialization matches the variance of PyTorch's 's default initialization.
Table 11 reports the lower bound, and Table 12 shows the corresponding gap from Theorem 5.6. Table 12 confirms that our theoretical bounds are tight. In Table 13, we empirically calculated the initial sharpness by sampling 400 random weight initializations for each 2-label subset of datasets, thereby calculating the average and standard deviation of 20000 different sharpness values. A comparison with the values in Table 11 reveals agreement between theory and experiment. We note that the standard deviation in Table 13 is relatively large, which is be expected given the extremely small width of our minimalist model.
Table 14 reports the lower bound, and Table 15 shows the corresponding gap for Theorem 5.7. These numerical results demonstrate that the bound gap decreases, and thus the theoretical bounds become tighter as the dataset size N grows.
To verify that the lower bound in Table 14 predicts the sharpness at convergence, we performed additional experiments following the setup of Figure 5b, but using the a& initialization with a² + 3 = over three random weight seeds. Although individual runs exhibit variability sharpness due to the model's extremely small width, Figures 18-20 show that, for small dataset sizes, the empirical sharpness saturates and its magnitude closely matches the prediction in Table 14. For larger N (i.e., N 500 for SVHN, and N 1000 for CIFAR10 and Google speech commands), convergence to this saturation of sharpness would require more training iterations than are shown in our current plots.
By combining the results in Tables 11 and 14, we verify that the expected sharpness increases from initialization to convergence.
33

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
Table 11. Average lower bound of E[ computed from 50 2-lahel subsets of datasets, with standard deviation in parenthesis,
dataset
N 100
=03
N 500
N = 1000
CIFAR
262.57(29.72)
260.42(15.41)
261.94(12.15)
261.93(9.16)
SVHN
602.25(39.35)
600.50(23.20)
601.97(15.33)
602.17(12.55) 229.77(6.93)
Google speech commands
233.42(12.69)
230.01(10.14))
227.72(3.80)
Table 12. Average gap hetween the upper bound and the lower bound of [R(S( standard deviation parenthesis, a² 8
computed from 50 2-lahel subsets of datasets, with
dataset
= 100
3-300
N = 500
= 1000 7.00(0.27)
CIFAR
8.91(1.03)
7.44(0.47)
7.18(0.36)
SVHN
12.54(1.43)
11.66(0.71)
11.49(0.42)
11.33(0.25) 10.69(0.16)
Google speech commands
11.36(0.58)
10.88(0.37)
(0.81
Table Average of S(O()) over two-label subset each dataset. each subset, S(g(0)) was computed across random weight
initializations (20,000 samples total); the standard deviation is shown in parentheses.
dataset
N 100
=03
N = 500
N = 1000
CIFAR
266.56 (388.70)
264.07 (383.16)
265.37 (384.81)
265.40 (384.68)
SVHN
610.09 (886.94)
608.07 (882.80)
609.07 (883.88)
609.38 (884.33)
Google speech commands
237.50 (342.38)
233.95 (337.19)
233.68 (336.58)
231.59 (333.44)
Table 14. Average lower bound of [[)(())) computed from 50 2-label subsets of datasets, with standard deviation in parenthesis, a²
g²
37
dataset CIFAR SVHN
N 100
N=
N 500
N 1000
514.11(75.15)
1141.41(97.99
1964.56(132.45)
5325.42((98.65
2250.11(316.08)
8539.61(718.93)
(8105.67(1347.25)
53138.75 $3138.75(2688.77) 3559.35(151.9))
Google speech commands
490.39(38.74)
994.79(60.39)
1557.68(75.34)
Table Average gap between upper bound and lower bound of [E[((")] computed from 2-label subsets datasets, with standard deviation in parenthesis, a² 32 B² 3
dataset
100
303
N 500
1000
CIFAR
44.16(4.46)
17.15(1.00)
9.64(0.57)
3.42(0.15)
SVHN Google speech commands
47.14(5.06)
11.15(0.81)
5.19(0.30)
1.75(0.07)
36.04(2.77)
15.36(1.10)
9.39(0.48)
3.89(0.19)
34

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
5000
N
N
N
4000
Sharpness
4000
N=1000
Sharpness
3000
N-1000
Sharpness
4000
N=1000
3000
300
N-000
3000
N-300
N-000
2000
N=100
N=200
2000
2000
1000
1000
0
80
80
20
80
100
20
100
time
time
time
(a) Seed
(b) Seed 2
(c) Seed 3
Figure 18. Minimalist model (D = 2) trained on different random weight seed, normal distribution with a² , 14. = iiiitialized, and the same 2-label subset of CIFARIO.
20k N
N
N
N=1000
15k
N=1000
N=1000
Sharpness
N=500
Sharpness
N=500
Sharpness
15k
N#500
Nr0300
10k
300
NI:300
10kk
10k
40
60
100
40
60
100
20
60
100
time
time
time
(a) Seed
(b) Seed 2
(c) Seed 3
Figure 19. Minimalist model (D trained different random weight seed, normal distribution with a² 3d 3 iinitialized, and the same 2-label subset of SVHN.
3000
2500 N
3000 N
Sharpness
2500
N=500
Sharpness
N=500
Sharpness
2500
N=500
2000
N=800
1500
N=300
N=300
1500
N-100
1000
N-000
1500
N-100
500
1000
500
500
60
80
100
20
80
20
100
time
time
time
(a) Seed
(b) Seed 2
(c) Seed 3
Figure 20. Minimalist model (D 2) trained on different random weight seed, normal distribution with a² 1. 82 iiiitialized, and the same 2-label subset of Google speech commands.
35

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
C.8. Change of T1 T23 1, 22 R2, and 92 for (S)GD
We present our experimental results for the terms we introduced in Theorem 5.9. Experimental settings are the same as Figure and Figure 9.
Learning Rate
3.5
2/18000
2/12000
2/6000
2/3000
1.5
0
10
20
30
0
10
20
30
0
10
20
30
0
10
20
30
time n iterations
time iterations
time iterations
time n iterations
Figure 21. Change of T1
=
B = 50
= 200
Learning Rate
22188000
2/12000
2.5
2/6000
2
2/3000
1.5
0.5
15
20
25
30
15
20
25
30
15
20
25
30
time iterations
time = iterations
time = iterations
Figure 22, Change of
B 25
B 50
B 200
100M
Learning Rate
10M
2/18000
2/12000
1M
2/6000
Ps
100k
2/3000
1000
10
20
30
0
10
20
30
0
10
20
30
0
10
20
30
time r iterations
time r iterations
time iterations
time iterations
Figure 23, Change of V
36

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
Learning Rate
2/18000
2/12000
Psi 2
18
2/6000
2/3000
5
10
15
20
25
30
10
15
20
25
30
10
15
20
25
30
time n iterations
time iterations
time iterations
Figure 24. Change of 2/2
B 25
B 50
B 200
B
1B
Learning Rate
100M
2/18000
Omega 1
2/12000
1M
2/6000
2/3000
100k
10k
1000
10
10
30
20
30
time 17 iterations
time r iterations
time = iterations
time = iterations
Figure Change of Q2
B 25
B 50
B = 200
Learning Rate
2/18000
Omega_2 22
2/12000
2/6000
18
2/3000
10
15
20
25
30
10
15
20
25
30
10
15
20
25
30
time n iterations
time n iterations
time r iterations
Figure 26.
Change of Pz
37

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
D. Precision Dependence of EoS Behaviors
D.I. Experiments Synthetic Dataset
To check the precision dependence of EoS behaviors, we show our experimental results from JAX, and our own framework. For the training data, we generated our synthetic data with the code below.
def minimal (d, common_sizer signal size, alpha, beta):
np. w...ddnd..eedd (0) X = zeros (2, d))
x_common np. ag..ranomm random| (1, )
X_opposite np. randon. random (1,
c_mmmonn /- np www.inallg.nomm///em (X_oommmX____ooe
X_opposite -= X_common (X_opposite flatte @ ._commn..flatten flatten ())
X_opposite /= ipp.iinae www.Inallg.omme (X_opposite) Orthogonalizing common and Xopposite
[[0] common_siie _sze * c_mmmn + signal_size X_opposite
X[]] common_size c_momon signal_size X_opposite
X flat reshae (2, -1)
eigs, vecs np. www.lialgg...m eigh flat @ X X_flat
= vecs I 0] * alpha vecs [;, 1] beta return X,
We used the result of minimal datan100, 5.477, 0.233, 0.3, 1.414) for both cases. We trained our minimal model of D using 7 = 0.02 GD.
200
10
M
sharpness
100
60
250
loss
1
10
10
1
40
00
10-2
1
10:3
10
(a) Sharpness with float32
(b) Sharpness with float64
(c) Loss with float32
(d) Loss with float64
Figure 27. Effects of precision the Edge of Stability, Experiment done in JAX
In Figure 27, we can see that our loss and sharpness plot of minimal model highly depends on the precision. We used the data type as muxp.. for float32 case, and jax fumpy.float64 for float64 case.
To clearly see the impact of precision in training dynamics, for the same settings, we made simple C++ program replicating the Py PyTch training using MPFR C++ (Holoborodko, 2010). In Figure 28, can clearly observe distinct trends across different precisions.
Notably, we find that high precision may lead to blow up of the training loss instead of entering the edge of stability regime and converging non-mnnot.nically
38

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
Precision: 16
Predslon: 32
Procision: 64
Prectsion: 128
134
1017
102
10--
10t1
10-4
1241
10-
10-4
1014
10
10
10
10.0
250
500
750
1000
1250
1500
250
503
750
1903
1250
35:00
250
5:0
750
1000
1250
1900
0
100
200
303
405
Precision: 256
Precision: 512
Precisioo:: 1024
Precision: 2048
1050
2010
10th
12th
10t
1014
12"
36
1.7
10a
307
169
101
16th
182
19a
1.0
10a
10a
100
200
300
400
502
100
200
300
400
500
100
200
300
400
500
0
100
200
300
400
500
(a) Loss plot
Precision:
Precision:
Precision:
Precision:
340
150
80-
100
100
40
25
250
500
150
1000
1250
1500
230
500
750
1205
1250
1500
0
230
520
750
1000
1250
1500
200
302
1.25
Precision: 256
Precision: 512
Precision: 1024
Precision: 2048
0
120
200
300
400
502
0
100
200
300
400
500
0
100
200
300
400
500
100
200
300
400
500
(b) Sharpness plot
Figure 28. D 2 minimal model in C++ framework training. axis given training step numbers
39

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
D.2. Experiments in 2-lahel Subset of CIFARIO Dataset
1000
10c
precision=8
precision=8
800
precision=12
precision=12
Sharpness
precision=16
precision=16
10
600
prectsion=32
loss
precision=32
precision=64
precision=64
400
10a
200
0
20
40
60
80
100
120
140
0
20
60
80
100
120
140
iterations
iterations
(a) Sharpness
(b) Loss
Figure 29. The effect precision minimalist model(D=2). with CIFAR-I0 2 label subset 300, /) 200
precision=32
600
precision=256
precision=64
precision=512
Sharpness
600
precision=128
Sharpness
500
precision=1024
precision=256
400
200
200
100
100
0-
100
200
300
400
500
600
200
300
400
500
600
iterations
iterations
(a) Sharpness low precision
(b) Sharpness high precision
10a
10a
precision=32
precision=256
10a
precision=64
precision=512
precision=128
precision=1024
loss
10
precision=256
loss
10c
10-1
10-2
10-2.
10
10-3
100
200
300
400
500
600
100
200
300
400
500
600
iterations
iterations
(c) Loss low precision
(d) Loss high precision
Figure 30. The effect of precision in realistic model (3 layer SiLU activated NN, width=32), and realistic data CIFAR 2 label subset 300, GD r
200
To check the generality of the observations in Appendix D.1, we show our experimental results from our own framework. For the training data, we used 2-label subset of CIFAR1O dataset.
For the experimental result in Figure 29, trained minimalist model (width=1, depth=2) CIFAR-2 subset (N = 300). The results confirm that higher precision leads to blowup. Additionally, Figure 29a shows that increased precision postpones the sharpness drop (catapult phase).
For the experimental result in Figure 30, we trained 3-layer SiLU-activated NNs (width=32) on a CIFAR-2 subset (N = 300). While settings do not exhibit blowup, higher precision delays the catapult phase until extremely high precision (2 256). Beyond this threshold, the training dynamics become nearly identical, with completely overlapping curves.
40

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
D.3. Hypothesizing Precision Dependence in Seff-Stabilization Framework
To elaborate on the intuition behind this dependency, briefly introduce Section of Damian al. (2023)'. They introduce concise 2-dimensional description of gradient descent dynamics the edge of stability by Taylor-expanding the loss gradient V L around fixed reference point d (which can be understood as the first 0 iterate that reaches sharpness 2/n)..
Two key quantities are defined as follows:
Zt = u (fg- 8*) measures the displacement from l* along the unstable (top eigenvector of V² direction u.
B/+ VS((") (02- 0*) quantifies the change in sharpness from its threshold value 2/17.
While assuming that progressive sharpening happens at scale of o:: (T((0))-VS(0))) - a > 0, the dynamics are described in three stages cycling throughout:
Progressive Sharpening: For small 2, and yo. the sharpness increases linearly: 3t+1 yt rja, (a 0). The update for I, is: 2++1 (1 + my az Thus, Ye becomes positive, the factor (1 my: > 1 causes | to grow. These updates rely on the Ist-order approximation of the gradient.
Blowup: With y > 0, the multiplicative effect in the z-update leads exponential growth in zzl. marking the blowup phase where the Ist-order approximation no longer suffices.
Self-Stabilization: When |aa] is large, the 2nd-order approximation of gradient (which involves L yields: y4+1 y1 ~
xx which provides a negative feedback that reduces y/ and stops further growth of [zz]-
Note from above that when y O, (i.e., before blowup), | shrinks exponentially. We hypothesize that higher numerical precision allows |zz| remain small for longer iterations, postponing blowup and self stbllttabilizaiinn, resulting in abnormally high sharpness.
'e adopt here the notation Damian al. (2023).

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
E. Preliminary Analysis of Nonlinear Activations
present preliminary results extending our minimalist model to networks with nonlinear activation functions. Assumption E.1 (Orthogonal Data Assumption). X is a diagonal matrix.
Note that each e; is given the standard basis vector under the Assumption E.1. Then we define our two-layer non-linear minimalist model and its corresponding layer imbalance as follows:
Definition E.2 (Two-layer Non-linear Minimalist Model), We define our two-layer non- non-linear minimalist model f : R + R
as
[(;; 0) :- h(x h(x'u)o1:
where 0 = (u, 01) represents the collection of all model parameters and h R -> IK is given an activation function that satisfies the following conditions:
h an injective function.
h E C (())) + 0, Vo R.
We define g antiderivate of 8.
Definition E.3 (Layer imbalance in non-linea minimalist model). The layer imbalance of the two-layer non-linear minimalist model with parameter 0 = (u, 01) is defined by
C(0) = g(0,0,)
UT
Definition E.3 can be viewed as the generalized version of Definition 4.2 4. For an arbitrary vector = (011 D2, .........
Inn
(((=)),(((2)).))) RR, we denote h(r)) (h(x1), h(x2), ..... h((rr)))) E TR Then, the loss function is given by T(0; 2N [ (1((u))n v|l2, and w. u. We propose useful lemma analogous to Lemma 5.2.
Lemma E.4 (Conservation of layer imbalance in non-linear minimalist model under gradient flow). Let ((t) be a gradient flow trajectory trained on the two-layer non-linear minimalist model. Then, the layer imbalance remains constant along the gradient flow trajectory: C(O(1)) C(O(0)) for all 2
Proof. It suffices to prove that C(((t)) Recall that the layer imbalance of two-layer non-linear minimalist model defined as
gloros(t))
C(O(t)) ==
0n (t)
Differentiating both sides with respect to gives
C(0(1)) = ox(t) (ato((t))
20, (t)
σ; (((o,o((t)
N (1) (eT =1)) M(oro)))))
2
Σ h(ooos) (c[=(1) (4)
N
0
Proposition E.5 (Solutions non-linear activation). Under Assumption E. 1. Iwo-layer non-linear minimalist model trained dataset (X, 2), the following equality holds at the global minimizer 0" = (u" v²) of L(0):
(0))² = Σ g(
(35)
of
*Given M(x) = J. Var E IR, (x)) = 4 results the same definition the imbalance.
42

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
where o = i x² h
4
Proof. Let 0r = (u", vi) be global minimizer of L(0) for a two-layer non-linear minimalist model trained on a dataset (X, y) that holds Assumption E.1. We denote of = wi u" for each E [r]] Since L((*) = =l=")=(0))1 the residual x(0))
a zero vector. Combining with (37), we have
e, =(0)) h(o,0))0 yuj di =, Vie
Moreover, we have
C(0') =2 > g(a;ot) of
(ut)²
i=1
by the definition of the layer imbalance. Substituting of 1 h 1
de of
finishes the proof.
By solving the equation of Proposition E.5 based on Lemma E.4 under Assumption 5.1, we can specify the solution of the gradient flow dynamics. Based on the selection of h, Equation (35) can be expressed in the following form:
tanh. In this case, g(0,0)) 1 cosh"'ooa (o;0ţ) follows the definition of g. Then,
1 cosh 2 (o,o)) (u7)2
C(O0)
i=1
1
o, 1 - tanh (ann?""000) (vj)²
1
(vi)²
(ot)² df (1T)2,
where the last equality used of = 1 a;
tanh -1
d: 77
sigmoid. In this case, g(050)) = exp(o;0;) + o,o follows the definition of g. Then,
010 + exp(o;o;)
C(0*) = 2 Σ
- (ut)²
o
1=1
In
d:
+
di
Uj
-n-di
11
(ot)²,
o²
where the last equality used of ln
d,
While we consider these two cases be relatively simple compared other nonlinear activation functions, a thorough analysis under highly nonlinear equation remains necessary and is therefore deferred to future work.
E.1. Reparameterization of the Two-layer Non-linear Minimalist Model
In this subsection, we reparameterize the gradient flow (GF) dynamics for the two-layer non-linear minimalist model (Definition E.2) under the Assumption E.1 in terms of o, d, o, for i [r]. and Un.
Network output and residual.
We can decompose network output f(X; 0) into e ei...., Er components by
S(X:)) - Σ (eT ((x;)ee
Σ (eT ((Xu)uz) = Σ (M(a)))))))
(36)
43

---

Understanding Sharpness Dynamics in NN Training with Minimalist Example
where the last equality is obtained by replacing e, ((u)) with liee Xu), and X with Ex=1 (.(ew,)) Similarly, can decompose residual =(0) by
x(0) = ((x;0) - (h(0,00)11- -d))
(37)
Hence, network output (36) and residual (37) can be reparameterized in terms of d,, o for E [r], and v1.
Gradient Flow. Recall that the loss at is L(0)
1 2N (f(;;0) yll 2\ (0)]) The GF dynamics is given by
0L
i(t)
Ou (f) N (1 [(()))(=())))))))))) (X())())),
(38)
aL
(t) dv (t) = h(Xu(t)) )(t))
(39)
N
We can replace X with Ei= o.(e,w?)) in (39) and obtain
d1(t)
1 N
Σ [t)
=(t) 1 (traa(t)){[2(t)
Hence, we have
1
(t)
(o):o(()))TT(() z(t.
(40)
N
Similarly, inner product with w, to both hand sides of (38) and replacing X with E or(e,,, gives
1 No (4) =()) N ((,00))), e [r].
àe - ú =
(41)
Therefore, (40) and (41) together give the reparameterization of the GF dynamics.
44


Reviews:
[{"review_id": "75MhUI9mZd", "reviewer": "Reviewer_47Zn", "summary": "This paper employs a minimalist model to investigate the progressive sharpening phenomenon in the training of deep neural networks. Progressive sharpening is a widely observed phenomenon characterized by the enhancement of sharpness during training using gradient descent or stochastic gradient descent, before reaching a saturation point at the boundary of stability. The primary objective of this study is to explore the relationship between progressive sharpening and various problem parameters, such as dataset size, network depth, batch size, and learning rate. The minimalist model is a regression problem with a single neuron per layer and identity activation.", "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": "We appreciate your thoughtful comments. Below, we address your recommendations and questions.\n\n## Comprehensive Literature Review & Implications\n\nThank you for highlighting relevant literature that we initially omitted. We agree that incorporating a more comprehensive review will strengthen our work.\n\n- Linear diagonal networks: Our minimalist model shares similarities with diagonal linear networks in a sparse regression setting. [1] showed that SGD leads to solutions with better generalization than GD. Similarly, our results show that SGD induces less progressive sharpening than GD, leading to lower sharpness at convergence. Considering that lower sharpness correlates with improved generalization in diagonal linear networks [2], they both unveil how stochasticity can help generalization.\n- Potential practical implications on learning rate scheduling: The study by [6] highlights that the catapult mechanism contributes positively to model generalization, and catapults can be induced by designing a proper learning rate schedule. In light of this, predicting sharpness evolution can offer practical value when designing such schedulers.\n- Connection between sharpness and generalization: SAM [3] was introduced under the hypothesis that minimizing sharpness improves generalization. Moreover, GD with a large learning rate has been shown to implicitly find flatter solutions [4], which often generalize better than those obtained with small learning rates. While these works suggest a correlation between sharpness and generalization, [5] showed that this relationship is data-dependent.\n\nThe link between sharpness and generalization remains an active research area, though our study focuses on optimization dynamics rather than generalization. Nevertheless, our findings on progressive sharpening may relate to generalization, as larger learning rates and smaller batch sizes yield less progressive sharpening, effectively regularizing toward lower-sharpness solutions. This also connects our analysis to the broader literature on the implicit bias of GD and SGD, including the works mentioned by the reviewer.\n\n## Increase of $C(\\theta)$ when it is positive\n\nWe acknowledge that Theorem 5.6 may have caused confusion, especially when $C(\\theta) > 0$. Our simplified version emphasized the effect of batch size and learning rate while avoiding complexity. Below is the complete result.\n\nIn Appendix A.5, instead of applying the Cauchy-Schwarz inequality in line 847 (Eq. 23), we derive the following exact characterization:\n\nFor SGD,\n\n$$\n\\\\begin{align*}\n&\\mathbb{E}\\_P[C(\\theta_\\text{SGD}^+)] - C(\\theta) \\\\\\\\\n&= \\underbrace{\\frac{\\eta^2}{N^2} [ - \\Psi_1(\\theta) C(\\theta) + \\Omega_1(\\theta) ]}\\_{C(\\theta_\\text{GD}^+) - C(\\theta)} + \\frac{\\eta^2(N-B)}{BN^2(N-1)}[- (\\Psi_2(\\theta) - \\Psi_1(\\theta)) C(\\theta) + (\\Omega_2(\\theta) - \\Omega_1(\\theta))]\n\\\\end{align*}\n$$\n\nWe denote, $\\Omega_1(\\theta) \\triangleq \\sum_{i} \\sum_{j>i} \\left[ \\sigma_i (z(\\theta)^\\top  e_i) o_j - \\sigma_j (z(\\theta)^\\top  e_j) o_i \\right ]^2$, $\\Omega_2(\\theta) \\triangleq N\\sum_{i} \\sum_{j>i} \\lVert \\sigma_i(z(\\theta) \\odot e_i) o_j - \\sigma_j(z(\\theta) \\odot e_j) o_i \\rVert^2$, and use $\\Psi_1$, $\\Psi_2$ from our theorem statement.  \n\nThen, $C(\\theta)$  increases for GD ($C(\\theta_\\text{GD}^+) - C(\\theta) \\geq 0$) if and only if\n\n$$\nC(\\theta) \\leq  \\frac{\\Omega_1(\\theta)}{\\Psi_1(\\theta)}=: T_1(\\theta).\n$$\n\nMoreover, $C(\\theta)$  increases for SGD faster than GD ($\\mathbb{E}\\_P [C(\\theta\\_\\text{SGD}^+)] - C(\\theta\\_\\text{GD}^+) \\geq 0$) if and only if\n\n$$\nC(\\theta) \\leq  \\frac{\\Omega_2(\\theta)- \\Omega_1(\\theta)}{\\Psi_2(\\theta) - \\Psi_1(\\theta)}=: T_2(\\theta).\n$$\n\nWe numerically verified that $C(\\theta) \\leq T_1(\\theta)$ and $C(\\theta) \\leq T_2(\\theta)$ holds throughout the GD/SGD training. The experiment results in PDF [[link]](https://anonymous.4open.science/r/understand_progressive_sharpening-E2F5/Experimental_Results_47Zn.pdf) illustrate how $C(\\theta)$ increases even when positive. Moreover, we observe the same dependence on learning rate $\\eta$ and batch size $B$—specifically, $C(\\theta)$ increases more with a larger learning rate or a smaller batch size, as long as $C(\\theta) \\leq T_1(\\theta)$ and $C(\\theta) \\leq T_2(\\theta)$ hold.\n\nWe will incorporate this discussion into the revised manuscript.\n\n## Underlying rationale for balanced condition\n\nPlease refer to our response to Reviewer smTr.\n\n---\n\n[1] Implicit Bias of SGD for Diagonal Linear Networks: a Provable Benefit of Stochasticity, NeurIPS 2021\n\n[2] Implicit Bias of the Step Size in Linear Diagonal Neural Networks, ICML 2022\n\n[3] Sharpness-Aware Minimization for Efficiently Improving Generalization, ICLR 2021\n\n[4] Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability, ICLR 2021\n\n[5] A Modern Look at the Relationship between Sharpness and Generalization, ICML 2023\n\n[6] Catapults in SGD: spikes in the training loss and their impact on generalization through feature learning, ICML 2024", "claims_and_evidence": "The provided paper introduces two crucial quantities: the dataset difficulty, denoted as $Q$, and the layer imbalance, denoted as $C(\\theta)$. It establishes a bound on the sharpness of the minimizer in terms of these two quantities under specific assumptions. The paper substantiates these claims through both proofs and empirical evidence.", "methods_and_evaluation_criteria": "The empirical methods appear to be sound and provide a valid basis for the verification of their theoretical results.", "theoretical_claims": "The proofs are grounded in conventional eigendecompositions of the data matrix, which appears sound to me.", "experimental_designs_or_analyses": "Empirical observations are systematically integrated throughout the paper to substantiate the theoretical assertions.", "relation_to_broader_scientific_literature": "The studied phenomenon is closely related to the ongoing research on the edge-of-stability phenomenon and the potential implicit regularization effect of SGD. The minimalist examples also appear to have drawn inspiration from several previous works [1, 2], which should be included in the literature review (see also the next question).\n\n[1] Gunasekar, Suriya, et al. \"Implicit bias of gradient descent on linear convolutional networks.\" Advances in neural information processing systems 31 (2018).\n[2] Woodworth, Blake, et al. \"Kernel and rich regimes in overparametrized models.\" Conference on Learning Theory. PMLR, 2020.", "essential_references_not_discussed": "I recommend that the paper conduct a more comprehensive literature review on the training dynamics of neural networks and the previous studies of gradient descent (GD) and stochastic gradient descent (SGD). While the authors have made it evident that their study focuses on the influence of problem settings on the progressive sharpening, I believe it would be essential to discuss the significance of this sharpening (implicit regularization) and how it could be utilized to enhance training (sharpness-aware minimization, SGD). Additionally, as previously mentioned, I believe the connection between the proposed minimalist model and previous works should be further elaborated.\n\nTo list a few:\n(1) Implicit Regularization of SGD: [4, 5, 6]\n(2) SAM: [3, 7, 8]\n(3) Connections to linear diagonal networks (and its variants): [1, 2, 5, 6, 9, 10]\n\n[3] Long, Philip M., and Peter L. Bartlett. \"Sharpness-aware minimization and the edge of stability.\" Journal of Machine Learning Research 25.179 (2024): 1-20.\n[4] Wu, Jingfeng, et al. \"Direction matters: On the implicit bias of stochastic gradient descent with moderate learning rate.\" arXiv preprint arXiv:2011.02538 (2020).\n[5] HaoChen, Jeff Z., et al. \"Shape matters: Understanding the implicit bias of the noise covariance.\" Conference on Learning Theory. PMLR, 2021.\n[6] Ren, Yinuo, Chao Ma, and Lexing Ying. \"Understanding the generalization benefits of late learning rate decay.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2024.\n[7] Andriushchenko, Maksym, and Nicolas Flammarion. \"Towards understanding sharpness-aware minimization.\" International conference on machine learning. PMLR, 2022.\n[8] Foret, Pierre, et al. \"Sharpness-aware Minimization for Efficiently Improving Generalization.\" International Conference on Learning Representations.\n[9] Pesme, Scott, Loucas Pillaud-Vivien, and Nicolas Flammarion. \"Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity.\" Advances in Neural Information Processing Systems 34 (2021): 29218-29230.\n[10] Cai, Yuhang, et al. \"Large stepsize gradient descent for non-homogeneous two-layer networks: Margin improvement and fast optimization.\" Advances in Neural Information Processing Systems 37 (2024): 71306-71351.", "other_strengths_and_weaknesses": "The paper is generally well-written and easy to follow. Heuristic claims are frequently validated with empirical observations, and the theorems are rigorously mathematically constructed.", "other_comments_or_suggestions": "See questions.", "questions_for_authors": "- I am somewhat perplexed by the reasoning presented in Section 5.2. It is asserted that whenever $C(\\theta) \\leq 0$, the corresponding layer imbalance of both GD and SGD is inevitably bound to increase. However, as previously mentioned by the authors, $C(\\theta)$ continues to rise even when $C(\\theta)$ is positive, which is clearly not elucidated by Theorem 5.6. Could you provide a possible explanation for this discrepancy? In this context, how does the batch size $B$ and the step size $\\eta$ influence the progressive sharpening process?\n- The balanced condition (Assumption 4.5) appears to be quite stringent. Could you elucidate the underlying rationale for this assumption?\n- As previously mentioned, how does progressive sharpening impact the generalization and the model’s quality after training? Are there any practical implications of the findings to the actual training of neural networks?", "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": "N6tNHKxoTf", "reviewer": "Reviewer_2e3Z", "summary": "This paper first shows that a \"minimalist model\"--one that has a single unit per layer with linear activations--can effectively captures a recently observed phenomenon called \"progressive sharpening\", where the sharpness of the loss increases as training progresses. This sharpness then stabilizes around $2/\\eta$ ($\\eta$ being the learning rate), a regime called the \"edge of stability\". The paper then defines several metrics: dataset difficulty and layer imbalance, which are used in theoretical bounds on the sharpness at the global minimum.\n\n## Update after rebuttal\n\nI am now more confident in recommending acceptance. The rebuttal text should be added to the paper or the supplementary, and would improve its presentation and overall quality.", "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": "We appreciate your insightful feedback. We have provided figures for additional experiments in the supplementary PDF file [[Link]](https://anonymous.4open.science/r/understand_progressive_sharpening-E2F5/Experimental_Results_2e3Z.pdf). Below, we address your comments.\n\n## Applicability of results across datasets\n\nWe have included additional experimental results on SVHN, shown in Figures 1–4 of the supplementary file, confirming that the observed trends hold across both CIFAR-10 and SVHN. To further test our results on a different modality, we conducted experiments on the Google Speech Commands dataset [3]. The results in Figures 7–10 and Tables 1–5 of the supplementary file again exhibit similar trends. These findings suggest that our observations capture general training dynamics rather than being specific to a particular dataset or modality.\n\n## Additional Experiments on Precision (Appendix C)\n\nWe appreciate the reviewer’s concern regarding the generality of our precision experiments in Appendix C. We have conducted additional runs on deep neural networks and a real-world dataset. The updated results are provided in Figures 5 and 6 of the supplementary PDF file.\n\n- **Figure 5:** We trained our minimalist model (width=1, depth=2) on a CIFAR-2 subset ($N=300, \\eta=0.01$). The results confirm that higher precision leads to a blowup. Additionally, Figure 5a shows that increased precision postpones the sharpness drop (catapult phase).\n- **Figure 6:** We trained a 3-layer SiLU-activated NNs (width=32) on a CIFAR-2 subset ($N=300, \\eta=0.01$). While all settings do not exhibit a blowup, higher precision delays the catapult phase until extremely high precision ($\\geq 256$). Beyond this threshold, the training dynamics become nearly identical, with completely overlapping curves. We suspect that a larger width prevents excessive progressive sharpening and blowup.\n\nFor further intuition, please refer to our response to Reviewer smTr.\n\n## Ratio of batch size and learning rate as a key parameter\n\nWe appreciate the suggestion to analyze the ratio $\\eta/B$ as a key parameter. While we did not explicitly frame our experiments in terms of this ratio, our results in Figure 2 already allow for an implicit comparison. For example, the cases with $B=125, \\eta=2/400$ and $B=250, \\eta=2/200$ share the same ratio but exhibit slightly different sharpness trends, suggesting that $\\eta/B$ alone may not fully capture the observed phenomena.\n\nIn large-batch settings ($B \\approx N$), our results align with the GD Edge of Stability (EoS) literature [1], showing that sharpness evolution remains largely unchanged before EoS, implying that the ratio has little effect on sharpness dynamics in this regime. Given these observations, we are uncertain about the additional insights gained by treating the ratio as the primary parameter, but we welcome further discussion.\n\n## Use of Runge-Kutta in Figures 1 & 4\n\nOur choice of the Runge-Kutta (RK4) method follows the setup in [1] (Appendix I.5). As shown in their Figures 29, 31, 33, 35, 37, and 39, RK4 gradient flow closely tracks GD dynamics before EoS. This allows us to analyze training dynamics independently of the learning rate schedule, providing a cleaner theoretical comparison. We will add the relevant citations in the revised manuscript.\n\n## Choice of sharpness measure\n\nOur sharpness definition follows [1] and [2], as our primary goal is to study progressive sharpening observed in [1] and the transition into the EoS regime. While alternative measures offer valuable insights, they primarily focus on generalization rather than training dynamics. Since our study is centered on the optimization perspective, the Hessian maximum eigenvalue remains the most relevant metric.\n\n## Minor concerns\n\n>Cite the LOBPCG paper \n\nThank you for pointing this out. We will add the citation.\n\n>Modify Figure 3 to also show the edge of stability. Figure 4 already does this fairly well.\n\nWe would like to clarify that Figure 4 does not depict the edge of stability. The observed sharpness saturation in Figure 4 is not due to the edge of stability but rather reflects sharpness dynamics under gradient flow, simulated using the Runge-Kutta (RK4) method instead of gradient descent. The edge of stability is explicitly presented in Figure 2 for the case where $B = N$, where sharpness saturates at 200 for a learning rate of $\\eta = 2/200$.\n\n>$o_i$ is used without definition\n\nIn Section 4, we specify that the first-layer weight $u$ can be decomposed as $u=\\sum_{i=1}^r o_i w_i + \\Pi_W^{\\perp}u$, where $w_i$ are the right singular vectors of $X$. Thus, $o_i$ represents the component of $u$ along $w_i$.\n\n---\n\n[1] Cohen et al., Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability, ICLR 2021.\n\n[2] Damian et al., Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability, ICLR 2023.\n\n[3] Warden P., Speech Commands: A Dataset for Limited-Vocabulary Speech Recognition, 2018.", "claims_and_evidence": "The authors test the effect of dataset size, network depth, batch size, and learning rate in progressive sharpening, and the theory derived lines up well with their experiments.\n\nWhile the datasets used in the paper are common and with good reason (for example, SVHN is \"harder\" than CIFAR-10, and this paper's results align with that), I'm concerned that *only* using these two limits the generalizability of the results. Most of this paper's results solely rely on CIFAR-10 (Figure 1-4), so it's unclear how well they translate to other datasets, such as those with different modalities.\n\nAnother issue along these lines is in Appendix C, where the authors seem to have run the experiment for a single seed, on a single, randomly-generated dataset, with one set of parameters passed to the `minimal_data` function, and only using the minimal model with $D=2$. This limits how representative the results are.", "methods_and_evaluation_criteria": "This is an interesting paper. For how surprising some results are, the authors do a good job designing experiments to show claimed effects. \n\nHowever, I believe the paper misses one aspect surrounding the variables it tests. While the paper tests the effect of learning rate and batch size separately on the sharpness, recent work has shown that these two determine the dynamics of SGD *jointly* as a ratio instead, see [1-2] below. Therefore, this paper would also benefit from performing its experiments with the *ratio* as a parameter. \n\n[1] Jastrzębski, S., Kenton, Z., Arpit, D., Ballas, N., Fischer, A., Bengio, Y., & Storkey, A. (2017). Three factors influencing minima in sgd. arXiv preprint arXiv:1711.04623.  \n[2] Smith, S. L., Kindermans, P.-J., Ying, C., & Le, Q. V. (2018). Don’t Decay the Learning Rate, Increase the Batch Size. International Conference on Learning Representations.", "theoretical_claims": "I checked the proofs in the main paper and did not find any issues. I skimmed the proofs of the Appendix, where results seem to follow naturally from prior work.", "experimental_designs_or_analyses": "I have detailed these concerns under Claims and Evidence.", "relation_to_broader_scientific_literature": "The contributions of this paper are significant. Although the paper's theory and experiments are currently limited to a small architecture, the authors effectively argue why it is still relevant. Moreover, this paper is an advancement in the field, and other papers in the future will build on it to larger architectures.", "essential_references_not_discussed": "I did not find any crucially missing references, though the authors should cite the LOBPCG paper since they use and mention it explicitly:\n\n[1] Knyazev, A. V. (2001). Toward the optimal preconditioned eigensolver: Locally optimal block preconditioned conjugate gradient method. SIAM journal on scientific computing, 23(2), 517-541.", "other_strengths_and_weaknesses": "This paper could benefit from better presentation in a couple of instances:\n\n* On page 4, under \"Progressive Sharpening\", the authors state: \"In Figure 3 and Figure 4, we can observe the same trend described in Phenomenon 1 [...]\". This would be a good chance for Figure 3 to also show the edge of stability in the same plots. Figure 4 already does this fairly well.\n* In Section 4, the $o_i$ terms are used without definition.", "other_comments_or_suggestions": "None.", "questions_for_authors": "1. For Figure 4 (and Figure 1, as stated in Appendix B.1), why do you use the Runge-Kutta method? It's possible I misunderstood, but my read of these figures was that they plot sharpness over optimizer iterations; it's curious to not use something like SGD.\n2. For sharpness, why do you specifically choose the maximum Hessian eigenvalue instead of other established measures such as the robustness under adversarial perturbations [1] or the volume of a polytope with a thresholded height [2]? These, especially the former, have been endorsed in the literature [3-4].\n\n[1] Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., & Tang, P. T. P. (2017). On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima. International Conference on Learning Representations.  \n[2] Hochreiter, S., & Schmidhuber, J. (1997). Flat minima. Neural Computation, 9(1), 1–42.  \n[3] Dinh, L., Pascanu, R., Bengio, S., & Bengio, Y. (2017, July). Sharp minima can generalize for deep nets. In International Conference on Machine Learning (pp. 1019-1028). PMLR.  \n[4] Andriushchenko, Maksym, Francesco Croce, Maximilian Müller, Matthias Hein, and Nicolas Flammarion. “A Modern Look at the Relationship between Sharpness and Generalization.” In International Conference on Machine Learning, 840–902. PMLR, 2023.", "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": "MU91nN97ZE", "reviewer": "Reviewer_smTr", "summary": "This paper analyzes sharpness dynamics and characterizes the effect of dataset complexity, depth, batch size (Phenomena 1). Furthermore, they show a simplified model, a deep linear network with unit width trained on multiple examples, captures the Phenomena 1 and analyzes it in detail. By analyzing the sharpness at convergence and properties of the Gradient Flow, they show that the rate of progressive sharpening increases with depth and data complexity. They also argue the effect of batch noise and learning rate in SGD setting.", "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": "We appreciate your valuable feedback. Below, we address your comments.\n\n## Balanced initializations for D>2 & GD/SGD analysis only on D=2\n\nWhile it is true that the assumption of balanced initialization for $D > 2$ may appear restrictive, we believe it is justified for the following reasons:\n\n- In line 728 of our paper, we implicitly use that condition to derive the GF limit point $v_1^\\star$. Without it, the resulting equation becomes a polynomial of degree $D$, whose closed-form solution can’t be derived when $D\\geq 5$. For $2<D<5$ cases, the solutions are also hard to interpret. Although the roots might be solvable numerically, we prioritized presenting interpretable analytical bounds that highlight the effect of depth.\n- Balanced initialization is also a widely adopted assumption in the literature of deep linear networks [2, 3].\n- As shown in Figure 7, experimental results that don’t satisfy balanced initialization align with our theoretical predictions.\n\nIn the case of GD/SGD analysis, the balanced constant $C$ does not remain fixed throughout training. As a result, even if training starts from a balanced initialization for $D > 2$, the balancedness may not be preserved during the optimization process.  For this reason, we limited our analysis to the two-layer case for intuitive understanding.\n\n## Effect of precision on EoS in typical theory\n\nFew works mention the precision as a meaningful parameter for EoS phase. We found one work [1, Appendix D.3], where they also mention that EoS dynamics is sensitive to the precision.\n\nTo elaborate the intuition, we’ll briefly introduce Section 4 of [1], which introduces a concise 2-dimensional description of gradient descent dynamics near the edge of stability by Taylor-expanding the loss gradient $\\nabla L$ around a fixed reference point $\\theta^\\star$ (which can be understood as the first $\\theta$ iterate that reaches sharpness $2/\\eta$). Two key quantities are defined:\n\n- $x_t = u \\cdot (\\theta_t - \\theta^\\star)$ measures the displacement from $\\theta^\\star$ along the unstable (top eigenvector of $\\nabla^2 L(\\theta^\\star)$) direction $u$.\n- $y_t = \\nabla S(\\theta^\\star) \\cdot (\\theta_t - \\theta^\\star)$ quantifies the change in sharpness from its threshold value $2/\\eta$.\n\nWe assume that progressive sharpening happens at scale of $\\alpha$: $-\\nabla L(\\theta^*) \\cdot \\nabla S(\\theta^*) = \\alpha > 0$. The dynamics are described in three stages cycling throughout:\n\n- **Progressive Sharpening:**\n    \n    For small $x_t$ and $y_t$, the sharpness increases linearly: $y_{t+1} - y_t \\approx \\eta \\alpha \\, (\\alpha > 0).$\n    \n    The update for $x_t$ is: $x_{t+1} \\approx -(1+\\eta y_t) x_t .$\n    \n    Thus, once $y_t$ becomes positive, the factor $(1+ \\eta y_t) > 1$ causes $|x_t|$ to grow. These updates rely on the 1st-order approximation of the gradient.\n    \n- **Blowup:**\n    \n    With $y_t > 0$, the multiplicative effect in the $x$-update leads to exponential growth in $|x_t|$, marking the blowup phase where the 1st-order approximation no longer suffices.\n    \n- **Self-Stabilization:**\n    \n    When $|x_t|$ is large, the 2nd-order approximation of gradient (which involves $\\nabla^3 L$) yields: $y_{t+1} - y_t \\approx \\eta \\left( \\alpha - \\frac{\\beta}{2} x_t^2\\right),$ which provides a negative feedback that reduces $y_t$ and stops further growth of $|x_t|$.\n    \n\nNote from above that when $y_t < 0$, (i.e., before blowup), $|x_t|$ shrinks to zero **exponentially**. We hypothesize that higher numerical precision allows $|x_t|$ to remain small for longer iterations, postponing blowup and self-stabilization, resulting in abnormally high sharpness. This is further supported by the experimental results discussed in our response to Reviewer 2e3Z.\n\n## Why is RK4 step size proportional to the inverse sharpness?\n\nThe RK4 step size is chosen based on its linear stability range. For the ODE $\\frac{dy}{dt} = \\lambda y$ with $\\lambda \\in \\mathbb{R}$, RK4 is stable if $-2.785 \\leq \\lambda \\eta \\leq 0$, where $\\eta$ is the step size. Since we already flip the sign in gradient flow, using a step size proportional to the inverse sharpness ensures stable integration.\n\nFurther rationale on our choice of RK4 can be found in our response to Reviewer 2e3Z.\n\n## Sharpness dynamics in Figure 2\n\nTo clarify, Figure 2 illustrates GD and SGD dynamics (Appendix B.1), not RK4. The observed sharpness patterns are thus likely due to problem-specific characteristics of the dynamics. However, our work focuses on the end-to-end behavior, not intermediate dynamics, leaving this as an open question. \n\n---\n\n[1] Damian et al., 2023, Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability, ICLR.\n\n[2] Arora et al., 2018, On the Optimization of Deep Networks: Implicit Acceleration by Overparameterization, ICML.\n\n[3] Arora et al., 2019, A Convergence Analysis of Gradient Descent for Deep Linear Neural Networks, ICLR.", "claims_and_evidence": "Claims:\n\n* (Empirical) Phenomena 1: Sharpness increases with (1) Dataset size, (2) Depth and decreases at smaller batch sizes (Figure 1 and 2)\n* (Empirical) The simplified model empirically captures Phenomena 1\n* (Theoretical) The theoretical analysis (+ empirical analysis Tables) of the simplified model captures the effect of dataset complexity and depth. \n* The effect of batch size and learning rate in SGD case is captured in restricted settings (depth 2)", "methods_and_evaluation_criteria": "The method and evaluations are solid.", "theoretical_claims": "The theoretical claims are sound. I have verified the claims in the main text and also verified dynamical equations for GF, GD and SGD.", "experimental_designs_or_analyses": "The experimental design and analyses are sound.", "relation_to_broader_scientific_literature": "The paper improves our understanding of progressive sharpness by examining the effect of depth, dataset complexity and batch size. The effect of smaller batch size reducing the rate of progressive sharpening has been known in prior work, such as Ref. [1].\n\n[1] Agarwala, A. and Pennington, J. High dimensional analysis\nreveals conservative sharpening and a stochastic edge of\nstability. arXiv preprint arXiv:2404.19261, 2024", "essential_references_not_discussed": "I think the work adequately discusses existing works.", "other_strengths_and_weaknesses": "Strengths\n\n* The paper is clearly written and the claims are justified well.\n\nWeakness\n\n* For D > 2, only balanced initializations are examined\n* The GD/SGD analysis is restricted to D = 2", "other_comments_or_suggestions": "Suggestions:\n\n* For equation 3, it might be better to use some other variable than d as its already used for input dimension", "questions_for_authors": "Questions:\n\n* Effect of precision on EoS: if high precision can cause loss to diverge, then the theoretical analysis should suggest divergences. But typical theoretical analysis does not suggest that. Do the authors have some understanding of this?\n\n* I am not familiar with step size selection for Runge Kutta methods. Why was the step size chosen to be proportional to inverse sharpness and recomputed periodically? \n\n* In Figure 2, do the authors understand why does the sharpness decrease --> increase --> decrease --> finally increase? Typically, GD trajectories are known to show three types of dynamics [1]: (1) increase throughout, (2) decrease throughout, and (3) decrease then decrease. Is it because of the Runge-Kutta step size selection?\n\n[1] Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability\nhttps://arxiv.org/abs/2103.00065\n\n[2] Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos", "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": "DMKVwlYO9c", "reviewer": "Reviewer_boFk", "summary": "The paper investigates sharpness dynamics in neural network training using a minimalist deep linear network with one neuron per layer. It theoretically and empirically shows that this model captures progressive sharpening and edge-of-stability behaviors observed in practice. Key contributions include:  \n1. Identifying dataset difficulty $Q$, depth, batch size, and learning rate as factors influencing sharpness.  \n2. Deriving theoretical bounds on sharpness at minima, dependent on $Q$ and depth.  \n3. Demonstrating that SGD’s stochasticity reduces progressive sharpening compared to GD via layer imbalance dynamics.\n\n\n## update after rebuttal\nThe authors have done a good job explaining their work. My concerns about how to generalize this framework to non-linear models remain. Nevertheless, I believe this paper should be accepted, and I would retain my score.", "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": "We deeply appreciate your constructive feedback. Below are our responses to your questions and suggestions.\n\n## Extending our work to more realistic scenarios\n\nThank you for the insightful suggestion. To explore the applicability of our findings in more realistic settings, we conducted additional experiments training CNNs on the CIFAR-2 subset ($N=100$). The experimental setup follows the description in Section 5.1, with the only difference being that we fixed the random seed for weight initialization across runs. Our CNN architecture of depth $D=3$ consists of two convolutional layers (32 channels) with average pooling, followed by a fully connected layer.\n\nFor predicted sharpness $\\hat S_D$, we set $D=3$ that matches the depth of our CNN. The Pearson correlation between actual and predicted sharpness is summarized below:\n\n| identity | tanh | SiLU | ReLU |\n| --- | --- | --- | --- |\n| 0.7335 | 0.5723 | 0.3561 | 0.4314 |\n\nWhile the correlation decreases with non-linear activations, the identity activation still achieves a correlation above 0.7, suggesting that the intuition behind predicted sharpness may extend to CNNs. However, the behavior under non-linear activations is less clear. A deeper investigation into broader architectures remains an interesting direction for future work.\n\n## How does $Q$ generalize to non-linear networks & Lack of nonlinear dynamics analysis\n\nWe would like to clarify that $Q$ is a quantity that depends solely on the dataset, not on the architecture of the neural network nor the non-linearity of the activation function. Therefore, we interpret the question on $Q$ as asking whether a nonlinear dynamics analysis is possible. If our understanding is incorrect, please let us know.\n\nWe now outline our initial attempt at analyzing non-linear networks. Assuming that $XX^\\top$ is a diagonal matrix, we can analyze our minimal model with $D=2$. For such $X$’s, each $e_i$ is given as the standard basis vector.\n\nThe loss function is given by $L(\\theta; X) = \\frac{1}{2N} \\lVert h(Xu) v_1  - y \\rVert^2$, and $o_i = e_i^\\top h(Xu) = h(e_i^\\top Xu)$. We consider a non-linear activation function $h$ that is injective, continuously differentiable, and has non-vanishing derivatives.\n\nThe gradient flow dynamics for $o_i$ and $v_1$ are given by:\n\n$$\n\\dot o_i = - \\frac{1}{N}v_1 \\sigma_i  \\left[(v_1h(e_i^\\top Xu) - y) \\cdot h'( e_i^\\top Xu) \\right] \n= - \\frac{1}{N}v_1 \\sigma_i  \\left[(v_1h(\\sigma_i o_i) - d_i) \\cdot h'(\\sigma_i o_i) \\right]. \n$$\n\nand\n\n$$\n\\dot v_1 = - \\frac{1}{N} \\sum_{i=1}^r h(e_i^\\top Xu)  \\left[e_i^\\top \\left(v_1h(Xu) - y \\right) \\right] \n= - \\frac{1}{N} \\sum_{i=1}^r h(\\sigma_i o_i)  \\left[ \\left(v_1h(\\sigma_i o_i) - d_i \\right) \\right] .\n$$\n\nWe can derive the following equation:\n\n$$\nv_1 \\dot v_1 = \\sum_{i=1}^r \\frac{\\dot o_i}{\\sigma_i} \\frac{h(\\sigma_i o_i)}{h'(\\sigma_i o_i)} .\n$$\n\nLet $g$ be an antiderivative of $\\frac{h}{h'}$. Integrating both sides gives:\n\n$$\n\\frac{1}{2} v_1^2 =  C+ \\sum_{i=1}^r \\frac{1}{\\sigma_i^2} g(\\sigma_i o_i).\n$$\n\nAssuming the convergence of the loss, we have:\n\n$$\no_i = \\frac{1}{\\sigma_i} h^{-1}\\left( \\frac{d_i}{v_1} \\right).\n$$\n\nTherefore, at convergence, the following equation holds:\n\n$$\n\\frac{1}{2}v_1^2 = C + \\sum_{i=1}^r \\frac{1}{\\sigma_i^2} g\\left(h^{-1}\\left( \\frac{d_i}{v_1} \\right)\\right).\n$$\n\nThis results in a highly non-linear equation of $v_1$ which requires further investigation to fully understand its solutions $v_1^\\star$. As such, we have decided to leave these parts for future work.\n\n## Why does $\\hat S_D$ correlate with empirical sharpness under non-balanced initialization?\n\nFor the sake of theoretical analysis, our analysis employs a balanced initialization. Nevertheless, we anticipate that similar trends would be observed under a non-balanced initialization scheme. Moreover, owing to the network's relatively large width in our non-linear network experiments, we believe that the randomness in initialization is effectively \"averaged out,\" leading to consistent results across different random seeds.\n\n## Our minimalist model is a marginal improvement\n\nWe agree that the minimalist model, on its own, builds upon previous works and is not the core novelty of our study. Instead, our primary contribution lies in rigorously connecting this model to the phenomenon of progressive sharpening—a relationship that has been underexplored in theoretical analyses. Through this, we derive novel insights into how problem parameters influence the degree of progressive sharpening during training. We believe these results advance the understanding of sharpening dynamics in ways that prior studies have not addressed.\n\n## Why does precision matter so much?\n\nPlease refer to our answer to Reviewer smTr.\n\n## RK4 in Figure 4\n\nPlease refer to our answer to Reviewer 2e3Z.", "claims_and_evidence": "Claim i: The minimalist model replicates sharpness dynamics (progressive sharpening, edge of stability).  \n- Evidence: Figures 3–5 show sharpness trajectories in the model matching practical networks (e.g., Transformers).  \n\nClaim ii: Sharpness at minima is bounded by $\\sigma_1^2 Q^{(D-1)/D}/N$ (Theorem 4.6).  \n- Evidence: Table 3 and Figure 7 show strong correlation between predicted and empirical sharpness. \n\nClaim iii: SGD increases layer imbalance $C(\\theta)$, reducing sharpness (Theorem 5.6).  \n- Evidence: Figure 6 shows $C(\\theta)$ grows faster with smaller batch sizes/larger learning rates, aligning with reduced sharpness in Figure 3.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are appropriate and well-justified for studying sharpness dynamics:\n   - Synthetic Data: Validates theoretical bounds under controlled settings (Sec 4.1).  \n   - Real Data: 2-class subsets of CIFAR10/SVHN balance tractability and realism.  \n   - Correlation Analysis: Tests generalizability of $\\hat{S}_D$ to nonlinear networks (Fig 7, Table 5).", "theoretical_claims": "The theoretical claims are solid.\n\nI checked Sharpness Bounds (Theorems 4.3, 4.6):  \n   - Derived via NTK spectral analysis (Appendix A.2–A.3).  \n   - Key Step: Expressing the NTK matrix in terms of $\\sigma_i$, $d_i$, $Q$, then bounding its spectral norm.  \n   \nI did not carefully check Layer Imbalance Dynamics (Lemma 5.2, Theorem 5.6) but it looks reasonable to me.\n\nThe analysis assumes GF converges to global minima (Assumption 5.1). While empirically validated, formal convergence guarantees are not provided.", "experimental_designs_or_analyses": "- GF/GD/SGD trajectories: Figures 3–6 validate sharpness dynamics and layer imbalance. The experiments look reasonable to me. In figure 4, I am puzzled by why RK4 is used. \n- Figure 7 is good way to demonstrate the upper bound. It looks reasonable to me as well.", "relation_to_broader_scientific_literature": "This work connects to:  \n- Progressive sharpening (Cohen et al., 2021) and edge of stability (Damian et al., 2023).  \n- SGD vs GD sharpness tradeoffs (Agarwala & Pennington, 2024).  \n- Similar minimalist example like Zhu et al. 2023, Kalra et al. 2023, where here the model captures convergence behavior.\nExtends prior work by analyzing dataset difficulty $Q$ and depth $D$ in a tractable model.", "essential_references_not_discussed": "I am not aware of.", "other_strengths_and_weaknesses": "Strengths:  \n- Novel insights into dataset difficulty and depth’s role in sharpness.  \n- Clear empirical validation across architectures/activations (Table 5, Fig 7).  \n\nWeaknesses:  \n- Lack of nonlinear dynamics analysis (e.g., ReLU, attention).  \n- It is a good model but feels like a marginal improvement over the \"convex-hull\" of existing works.", "other_comments_or_suggestions": "It would be very interesting if the authors can try to run some experiments that are closer to real-world usages. Such as GPT-2 level small experiments or some proper image classification tasks. I would like to see how far this simple model can go.", "questions_for_authors": "1. How does dataset difficulty $Q$ generalize to non-linear networks? I knew technically there is no difficulty defining it, but I wonder if the intuition still holds.\n2. Why does $\\hat{S}_D$ correlate with empirical sharpness even under non-balanced initializations (Fig 7)?  Could the authors provide some intuition?\n3. Could authors provide any intuition on why the precision matters so much?", "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: This paper investigates the phenomenon of progressive sharpening in deep neural network training using a minimalist model. The work provides theoretical analysis and empirical validation for how sharpness evolves during training, and how problem parameters like depth, data difficulty, learning rate, and batch size influence this evolution. There is strong consensus among the reviewers that the paper makes a valuable contribution to understanding sharpness dynamics — a actively studied topic in optimization and deep learning theory. Some reviewers pointed out that the analysis is primarily limited to linear models or highly simplified settings. While this is acknowledged by the authors, and some preliminary attempts at nonlinear analysis were discussed, further exploration in this direction would strengthen future work. Given the overall reviewer consensus, the quality of the technical contributions, the clarity of the presentation, and the constructive engagement during the rebuttal phase, I recommend acceptance.
Decision: Accept (poster)