Understanding Sharpness Dynamics in NN Training with a Minimalist Example: The Effects of Dataset Difficulty, Depth, Stochasticity, and More

We deeply appreciate your constructive feedback. Below are our responses to your questions and suggestions.

Extending our work to more realistic scenarios

Thank 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.

For 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:

While 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.

How does $Q$ generalize to non-linear networks & Lack of nonlinear dynamics analysis

We 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.

We 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.

The 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.

The gradient flow dynamics for $o_i$ and $v_1$ are given by:

$$\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] = - \frac{1}{N}v_1 \sigma_i \left[(v_1h(\sigma_i o_i) - d_i) \cdot h'(\sigma_i o_i) \right].$$

and

$$\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] = - \frac{1}{N} \sum_{i=1}^r h(\sigma_i o_i) \left[ \left(v_1h(\sigma_i o_i) - d_i \right) \right] .$$

We can derive the following equation:

$$v_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)} .$$

Let $g$ be an antiderivative of $\frac{h}{h'}$. Integrating both sides gives:

$$\frac{1}{2} v_1^2 = C+ \sum_{i=1}^r \frac{1}{\sigma_i^2} g(\sigma_i o_i).$$

Assuming the convergence of the loss, we have:

$$o_i = \frac{1}{\sigma_i} h^{-1}\left( \frac{d_i}{v_1} \right).$$

Therefore, at convergence, the following equation holds:

$$\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).$$

This 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.

Why does $\hat S_D$ correlate with empirical sharpness under non-balanced initialization?

For 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.

Our minimalist model is a marginal improvement

We 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.

Why does precision matter so much?

Please refer to our answer to Reviewer smTr.

RK4 in Figure 4

Please refer to our answer to Reviewer 2e3Z.

We appreciate your valuable feedback. Below, we address your comments.

Balanced initializations for D>2 & GD/SGD analysis only on D=2

While it is true that the assumption of balanced initialization for $D > 2$ may appear restrictive, we believe it is justified for the following reasons:

• 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.
• Balanced initialization is also a widely adopted assumption in the literature of deep linear networks [2, 3].
• As shown in Figure 7, experimental results that don’t satisfy balanced initialization align with our theoretical predictions.

In 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.

Effect of precision on EoS in typical theory

Few 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.

To 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:

• $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$.
• $y_t = \nabla S(\theta^\star) \cdot (\theta_t - \theta^\star)$ quantifies the change in sharpness from its threshold value $2/\eta$.

We 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:

• Progressive Sharpening:

  For small $x_t$ and $y_t$, the sharpness increases linearly: $y_{t+1} - y_t \approx \eta \alpha \, (\alpha > 0).$

  The update for $x_t$ is: $x_{t+1} \approx -(1+\eta y_t) x_t .$

  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.

• Blowup:

  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.

• Self-Stabilization:

  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|$.

Note 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.

Why is RK4 step size proportional to the inverse sharpness?

The 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.

Further rationale on our choice of RK4 can be found in our response to Reviewer 2e3Z.

Sharpness dynamics in Figure 2

To 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.

---

[1] Damian et al., 2023, Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability, ICLR.

[2] Arora et al., 2018, On the Optimization of Deep Networks: Implicit Acceleration by Overparameterization, ICML.

[3] Arora et al., 2019, A Convergence Analysis of Gradient Descent for Deep Linear Neural Networks, ICLR.

We appreciate your insightful feedback. We have provided figures for additional experiments in the supplementary PDF file [Link]. Below, we address your comments.

Applicability of results across datasets

We 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.

Additional Experiments on Precision (Appendix C)

We 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.

• 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).
• 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.

For further intuition, please refer to our response to Reviewer smTr.

Ratio of batch size and learning rate as a key parameter

We 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.

In 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.

Use of Runge-Kutta in Figures 1 & 4

Our 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.

Choice of sharpness measure

Our 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.

Minor concerns

Cite the LOBPCG paper

Thank you for pointing this out. We will add the citation.

Modify Figure 3 to also show the edge of stability. Figure 4 already does this fairly well.

We 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$.

$o_i$ is used without definition

In 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$.

---

[1] Cohen et al., Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability, ICLR 2021.

[2] Damian et al., Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability, ICLR 2023.

[3] Warden P., Speech Commands: A Dataset for Limited-Vocabulary Speech Recognition, 2018.

We appreciate your thoughtful comments. Below, we address your recommendations and questions.

Comprehensive Literature Review & Implications

Thank you for highlighting relevant literature that we initially omitted. We agree that incorporating a more comprehensive review will strengthen our work.

• 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.
• 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.
• 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.

The 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.

Increase of $C(\theta)$ when it is positive

We 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.

In Appendix A.5, instead of applying the Cauchy-Schwarz inequality in line 847 (Eq. 23), we derive the following exact characterization:

For SGD,

\ $ &\mathbb{E}\_P[C(\theta_\text{SGD}^+)] - C(\theta) \\\\ &= \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))] \ $

We 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.

Then, $C(\theta)$ increases for GD ($C(\theta_\text{GD}^+) - C(\theta) \geq 0$) if and only if

$$C(\theta) \leq \frac{\Omega_1(\theta)}{\Psi_1(\theta)}=: T_1(\theta).$$

Moreover, $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

$$C(\theta) \leq \frac{\Omega_2(\theta)- \Omega_1(\theta)}{\Psi_2(\theta) - \Psi_1(\theta)}=: T_2(\theta).$$

We 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] 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.

We will incorporate this discussion into the revised manuscript.

Underlying rationale for balanced condition

Please refer to our response to Reviewer smTr.

---

[1] Implicit Bias of SGD for Diagonal Linear Networks: a Provable Benefit of Stochasticity, NeurIPS 2021

[2] Implicit Bias of the Step Size in Linear Diagonal Neural Networks, ICML 2022

[3] Sharpness-Aware Minimization for Efficiently Improving Generalization, ICLR 2021

[4] Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability, ICLR 2021

[5] A Modern Look at the Relationship between Sharpness and Generalization, ICML 2023

[6] Catapults in SGD: spikes in the training loss and their impact on generalization through feature learning, ICML 2024