Summary
This work shows a regularization condition for SGD in the interpolation regime which allows it to have same fast linear convergence rate as deterministic gradient descent. Hence, the theory presented in this paper supports the practical observation that with the same (large) learning rate, mini-batch SGD has almost the same convergence rate of SGD. This goes against the traditional approaches for SGD convergence under PL condition that require SGD to have smaller step-size than GD, hence making SGD's convergence slower in theory. The authors here present the conditions in which SGD have similar iteration complexity as GD.
Strengths
1) The main strength of the work lies in it's novelty. The motivation of the problem seems clear from the introduction.
2) This is an important problem since it is crucial to have a large learning rate for SGD as it helps improving generalization. Hence, the current theories that analyze SGD under PL-condition and requires it to have a smaller l.r, may suffer from having good generalization. The theory presented by the authors allows the same large learning rate to be used as GD. Hence, the variance of SGD coupled with the alowable large learning rate can boost generalization.
Weaknesses
1) Looking at the theorems, I still feel the assumptions made are too strong too hold in a non-convex landscape especially Aiming. For theorem-1.2, which consdiers a non-convex loss landscape on $w$, it is unclear how aiming can hold on a ball of radius $r$ and it's relation to the curvature at that point (given by minimal eigenvalue of NTK). The authors make sure that the iterates remain inside a ball $B_{r}(w_{0})$ near the minima but the radius depends on the eigenvalue of NTK at initializaiton! This makes a very strong assumption that the initial point is set to be very close to the true minima $w_{0}$. More so, if initialization is made far away from the minima, there is a high probability that aiming does not hold.
2) From the theorems, it is unclear that how the stochasticity from the mini-batch gradients effect the convergence or what is it's effect in ensuring that the iterates don't escpae $B_{r}(w_{0})$. For theorem-2.3, I assume that it will take longer iterations if the variance in mini-batch gradient is high. It is not very clear from the inequality as how the relation is.
A minor correction:
3) On line-42, it might be better to cite a different source such as [1] to refer to the generalization effect of large learning rate, as the authors of "Edge of stability" themselves mention that they don't claim the generalization effect of large learning rate but rather focus on the stability effect. See https://youtu.be/6xeh6gfESuc?t=3016
[1] Li, Yuanzhi, Colin Wei, and Tengyu Ma. "Towards explaining the regularization effect of initial large learning rate in training neural networks." Advances in Neural Information Processing Systems 32 (2019).
Questions
See points 1 and 2.
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
The main limitation of this work is a strong assumption on the locality of the SGD analysis. Some further clarification would be helpful.