Phase diagram of early training dynamics in deep neural networks: effect of the learning rate, depth, and width

We systematically analyze optimization dynamics in deep neural networks (DNNs) trained with stochastic gradient descent (SGD) and study the effect of learning rate $η$, depth $d$, and width $w$ of the neural network. By analyzing the maximum eigenvalue $λ^H_t$ of the Hessian of the loss, which is a measure of sharpness of the loss landscape, we find that the dynamics can show four distinct regimes: (i) an early time transient regime, (ii) an intermediate saturation regime, (iii) a progressive sharpening regime, and (iv) a late time ``edge of stability" regime. The early and intermediate regimes (i) and (ii) exhibit a rich phase diagram depending on $η\equiv c / λ_0^H $, $d$, and $w$. We identify several critical values of $c$, which separate qualitatively distinct phenomena in the early time dynamics of training loss and sharpness. Notably, we discover the opening up of a ``sharpness reduction" phase, where sharpness decreases at early times, as $d$ and $1/w$ are increased.

Paper

Similar papers

Peer review

Reviewer bhX57/10 · confidence 4/52023-06-16

Summary

The paper studies the effect of depth, width and learning rate on the early training dynamics of DNNs. The authors first describe four phases throughout, and then focus on the first two phases. In the first phase, they identify three ranges of learning rates, according to how they affect the evolution of the training loss and the sharpness of the loss in the first few training steps: - for small learning rates both loss and sharpness decrease monotonically. Surprisingly, this regime extends to learning rates that are larger than the traditional threshold of $2 / \lambda_{max}$ which guarantees a close match of GD to GF. - For larger learning rates, the loss is non-monotone but the sharpness still decreases. - For yet larger learning rates, both the loss and sharpness are non-monotone. Any larger learning rates lead to divergence. In general larger depths $d$ and smaller widths $w$ shifts these regime to larger values of the learning rate. In the second phase, the sharpness settles to some value. For small learning rates $\eta$, this value is roughly constant w.r.t. $\eta$ (but which depends on the width), but for large $\eta$ the sharpness decreases and settles to a value which seems to be independent of the width. The simple model of a shallow linear network with one datapoint and one dimensional inputs and outputs is studied, and the qualitative features described above are recovered.

Strengths

Understanding the different regimes of training as a function of width, depth and learning rates is an important problem, and any work that studies these questions either mathematically or empirically is welcome. The analysis and discussions are clear and insightful. A range of different architectures are tested. The regimes proposed and studied are to my knowledge new.

Weaknesses

The authors define a large number of regimes acoording to some rather subtle changes in behavior sometimes, and it is not always clear whether these regimes are also correlated with other quantities that we actually care about (e.g. generalization/feature learning). The training phase of neural networks features many strange/non-monotonic behavior and it is not clear that all of them have to be identified, since they might have a rather small impact.

Questions

It is shortly mentioned that setting the outputs of the network at initialization to zero changes the ranges of learning rates corresponding to the regimes of phase 1. In the appendix, it appears that the transitions between ranges become almost independent of the width in this case, and the smallest range (where both loss and sharpness decrease) stops much closer to the expected value of $2 / \lambda_{max}$. This suggests that the dependence of these values on the width $w$ which is quite thoroughly studied in the main paper, might only be related to the variance size of the outputs at initialization, and not some deeper property of the network. Ideally I think that the plots of the appendix describing this phenomenon should be put in the main, as it greatly impacted my personal interpretation of the results of the paper. Did the authors also check the behavior of the output function norm in the first few timesteps? I was wondering whether the non-monotonic behavior at the begining of training might be related to the function norm having a similar non-monotonic behavior (scaling up then down, possibly to some value close to zero), since the function norm has an obvious impact on both the loss and sharpness of the loss.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

The limitations of the results presented are discussed well.

Reviewer BcD37/10 · confidence 4/52023-07-06

Summary

The paper identifies new phenomena involving the dynamics of the sharpness (largest eigenvalue of the Hessian) in the early stages of neural net training, for ReLU nets of various architectures that are initialized using He initialization. Previously, it was known that for very wide networks, there are two 'phases' in hyperparameter space: on the one hand, if the step size exceeds 2/(initial sharpness), then a "catapult" occurs (https://arxiv.org/abs/2003.02218) and the sharpness drops; on the other hand, if the step size is below this threshold, then the sharpness remains at its initial value. This paper fills in the picture for _narrow_ nets. For these nets, the paper shows that the sharpness decreases at the beginning of training even when the step size is less than 2/(initial sharpness). As a consequence, the train loss does not spike even when the step size is bigger than 2/(initial sharpness), since the sharpness drops first. Only for larger learning rates does the train loss initially spike. Then, for still larger learning rates, the paper shows the _sharpness_ also spikes (which the paper calls a "sharpness catapult"), which has not been observed previously. For several architectures, the paper draws phase diagrams depicting which learning rates cause which behavior. The paper then investigates connections to the loss connectivity literature, finding that only very large learning rates result in a barrier in the loss landscape (and that not every catapult results in a barrier). The paper then investigates the role of initial network output, finding that initial sharpness reduction goes away if the network is made to output zero at initialization. The paper finally reproduces all of these phenomena analytically in a simplified model (the u-v model from the catapult paper).

Strengths

- The paper identifies and carefully studies a novel phenomenon involving the dynamics of the sharpness at the early stage of training. - The experiments are very thorough (though limited to ReLU nets initialized with He initialization) - There is an analytical explanation in a simplified model - The findings shed some light on mode connectivity

Weaknesses

- The paper doesn't explain the early-stage sharpness reduction for general neural nets (i.e. beyond the simple u-v model)

Questions

- Have you looked at loss connectivity between the final parameters (at the end of training) and the initialization?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

2 fair

Contribution

3 good

Limitations

discussed above

Reviewer qjGV5/10 · confidence 3/52023-07-07

Summary

In this paper, the authors studied the effects of learning rate, width and depth of neural networks for early training dynamics. Specifically, the authors focused on the value of loss and the maximum eigenvalue of Hessian of loss and found that there are 4 different possible training dynamics in the early training. These different regimes are determined by the value of learning rate and are related with depth and width. Experiment results are provided to support the claim. A simple example with the same phenomena is also given and analyzed.

Strengths

1. Understanding the effect of large learning rate and its effect on the nonlinear training dynamics is an important problem. 2. The paper presents many experiment results to support their findings. 3. The finding that there exist different regimes depending on the learning rate ($c_{loss}$, $c_{sharp}$, $c_{max}$) in the early training seems to be interesting and new.

Weaknesses

1. The current paper seems to have lots of results and experiments in the main text. As a reader, it is not very easy for me to get the main conclusion for each section/set of experiments. It might be good to highlight the conclusions so that the readers can understand the point easier. 2. The current paper tries to study the effect of depth and width on those $c$ values ($c_{loss}$, $c_{sharp}$, $c_{max}$) related with stepsize. However, it seems to me that only a few choices of width $w$ are tried in the experiments (Figure 3) and no results showing the effect of depth $d$ in the main text. It might be good to include more results on different width and depth so that the fitted curves could use more than only 3 or 4 data points and the support of the conclusion could be stronger.

Questions

1. I was wondering for these $c$ values ($c_{loss}$, $c_{sharp}$, $c_{max}$), how many times do the experiments repeat? 2. For the definition of these $c$ values, I was wondering how time $T_1$ is chosen. 3. In Definition 4, should $\chi_\tau$ be $\chi_\tau^\prime$? 4. In Figure 4, I was wondering if the networks have the same depth and width, or different point represents network that may have different depth and width.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

The limitations are discussed in the paper. This is a theoretical work and therefore does not seem to have negative societal impact.

Reviewer bXcV6/10 · confidence 3/52023-07-08

Summary

This work mainly investigates the learning dynamics depending on the learning rate and categorizes its qualitative behaviors in some phases. In particular, the authors identify new phases; sharpness reduction, and loss-sharpness catapult. They investigate the dependence of the critical learning rates (mainly, $c_{class}, c_{sharp}, c_{max}$ and $c_{barrier}$) on the width, depth & loss type, and try to identify a universal relationship between them. Their empirical observations in some deep neural networks are also confirmed in a simple linear network and the obtained phase diagram is expected to hold in various situations.

Strengths

This work attacks the challenging problem to find any universal behavior of early learning dynamics in finite-size networks (I guess that the terminology of "phase diagram" comes from physics or chemistry but such a dynamic and finite-sized diagram seems not known). Despite this difficulty, the authors verified new phases (characterized by $c_{loss},c_{sharp},c_{max}$) in various settings and these new phases seem to have the originality. Their careful and elaborated experiments are persuasive and enhance the quality of the paper.

Weaknesses

(a) The whole findings are empirical. Even in a simple model of $uvx$, the authors did not solve the dynamics in an analytical form. It is quite hard to obtain any intuition on why the three phases characterized by $c_{loss}<c_{sharp}<c_{max}$ universally appear. (b)They are several unclear points that the authors should more clearly mention. See below.

Questions

**Connection to the infinite width limit** While the current work focuses on finite-sized networks, some previous work investigated the infinite-width limit [64]. In particular, [64] claims that for the infinite-width limit, the learning dynamics *must be* the kernel regime for SP initialization (nevertheless, the learning dynamics does not progress or explodes). How is this infinite-width study related to your results? Since the critical learning rate of the kernel regime is given by $2/\lambda_{0}^{NTK}$, does this suggest that all $c_{loss,sharp,max}$ converge to 2 in the infinite width limit? **Phase diagram in previous work** Related to [64], there are several studies on the phase diagram of convergence in the infinite width case (e.g. https://arxiv.org/abs/2012.15110 and https://arxiv.org/abs/2007.07497). It would be better to mention them for giving a rich overview of the phase-diagram studies in the literature. More related to the current work, the phase diagram for width v.s. learning rate is also investigated in the literature (https://arxiv.org/abs/1806.01316). This work implies that 2/$\lambda_{0}^{NTK}$ works as a critical learning rate of the convergence in the infinite width. **The effect of early learning dynamics on eventual minima** I understand that the main focus of the current work is on the early training regime, but it is quite curious how the sharpness reduction, loss catapult, and loss-sharpness catapult phases determine the eventual generalization performance and sharpness of minima in the trained models. If this point would be clarified, it would increase the significance of this paper much more.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

**No normalization layers** The architectures in this work have no batch/layer normalization layer. Since the batch norm is known to strongly affect the sharpness (e.g. https://arxiv.org/abs/1901.10159), it is curious how the normalization layer potentially changes the results. **Difference between MSE and cross-entropy cases** It seems that the cross-entropy loss shows a different behavior of the critical learning rates. In particular, in most cases, $c_{loss}$ and $c_{sharp}$ decrease close to 2 as $1/w$ decreases in the MSE case while they keep far from 2 in the cross-ent case. In the revised version, I believe it would be better to give a more elaborated explanation of the difference in the main text.

Reviewer qjGV2023-08-13

Thanks for the response to address my question. It makes things clear and help me have a better understanding of the paper. I will increase my score.

Reviewer bXcV2023-08-14

Thank you for your kind reply. I am looking forward to seeing the accepted version, which will briefly include the above answers. >a complete understanding of the catapult phenomenon in the context of cross-entropy loss requires a more detailed examination. I agree. I hope that this point will be clarified in your next work or other subsequent works.

Area Chair kLa12023-08-18

Thanks for authors' rebuttal!

Reviewer BcD3, did the authors address your concerns on early-stage sharpness reduction for general neural nets? Thanks.

Authorsrebuttal2023-08-21

Further clarifications on early sharpness reduction and loss catapult

Even though this question was directed at Reviewer BcD3, we would like to take the opportunity to clarify this point. We have observed two separate phenomena in general neural networks with standard parameterizations: (i) the initial reduction in sharpness (i.e. the "sharpness reduction phase") and (ii) the systematic increase of $c_{loss}$ (marking the onset of the loss catapult) with $1/w$, leading to the opening up of the sharpness reduction phase with $1/w$. To further understand these phenomena, we performed systematic experiments as discussed in Section 4 and Appendices G and H. These results indicate that the initial sharpness reduction depends on the function output scale at initialization (see Figures 29 and 31). This provides some explanation, albeit incomplete, for the origin of (i). However, the reasons behind the scaling of $c_{loss}$ with $1/w$ remain more unclear. This scaling only ceases when the function output is set to zero at initialization (Appendix G, Figure 30) and persists even when the network output is made small by scaling it by a constant (Appendix H, Figure 32). We made some progress in developing a theoretical understanding of these phenomena in a toy model (the uv model), but developing a complete theoretical understanding for general neural nets is beyond the scope of this work. Given the richness and universality of these findings across architectures and datasets, we believe it deserves the attention of the wider community to arrive at a complete understanding.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC