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