SGD with large step sizes learns sparse features
Authors: Maksym Andriushchenko
Aditya Vardhan Varre
Loucas Pillaud-Vivien
Nicolas Flammarion
Venue: Submitted to ICLR 2023 (2023)
Keywords: SGD
large step sizes
implicit regularization
feature learning

Abstract: We showcase important features of the dynamics of the Stochastic Gradient Descent (SGD) in the training of neural networks. We present empirical observations that the commonly used large step sizes (i) lead the iterates to jump from one side of a valley to the other causing \textit{loss stabilisation} (ii) this stabilisation induces a hidden stochastic dynamics orthogonal to the bouncing directions that \textit{biases it implicitly} toward simple predictors. Furthermore, we show empirically that the longer large step sizes keep SGD high in the loss landscape valleys, the better the implicit regularization can operate and find sparse representations. Notably, no explicit regularization is used so that the regularization effect comes solely from the SGD training dynamics influenced by the step size schedule. Therefore, these observations unveil how, through the step size schedules, both gradient and noise drive together the SGD dynamics through the loss landscape of neural networks. We justify these findings theoretically through the study of simple neural network models. Finally, we shed a new light on some common practice and observed phenomena when training neural networks.

Under review as a conference paper at ICLR 2023
SGD WITH LARGE STEP SIZES LEARNS SPARSE FEATURES
Anonymous authors Paper under double-blind review
ABSTRACT
We showcase important features of the dynamics of the Stochastic Gradient Descent (SGD) in the training of neural networks. We present empirical observations that commonly used large step sizes (i) lead the iterates to jump from one side of a valley to the other causing loss stabilization, and (ii) this stabilization induces a hidden stochastic dynamies that biases implicitly toward simple predictors. Furthermore, we show empirically that the longer large step sizes keep SGD high in the loss landscape valleys, the better the implicit regularization can operate and find sparse representations. Notably, no explicit regularization is used so that the regularization effect comes solely from the SGD dynamics influenced by the step size schedule. Therefore, these observations unveil how, through the step size schedules, both gradient and noise drive together the SGD dynamics through the loss landscape of neural networks. We justify these findings theoretically through the study of simple neural network models as well as qualitative arguments inspired from stochastic processes. Finally, this analysis allows to shed a new light on some common practice and observed phenomena when training deep networks.
INTRODUCTION
Deep neural networks have accomplished remarkable achievements wide variety of tasks. Yet, the understanding of their remarkable effectiveness remains incomplete. From an optimization perspective, stochastic training procedures challenge many insights drawn from convex models. E.g.. large step-size schedules used in practice lead to unexpected patterns of stabilizations and sudden drops the training loss, see e.g. He al. (2016). From generalization perspective, overparametrized deep nets generalize well while fitting perfectly the data and without any explicit regularizers (Zhang et al., 2017). This suggests that optimization and generalization are tightly intertwined: neural networks find solutions that generalize well thanks to the optimization procedure used to train them. This property, known as implicit bias or algorithmic regularization, has been studied recently both for regression (Li et al., 2018; 88; Woodworth et al., 2020) and classification (Soudry et al., 2018; Lyu and Li, 2020; Chizat and Bach, 2020). However, for all these theoretical results, it is also shown that typical timescales needed to enter the beneficial feature learning regimes are prohibitively long Woodworth et al., 2020; Moroshko et al., 2020).
In this paper, we aim staying closer to the experimental practice and consider the SGD schedules from the ResNet paper (He et al., 2016) where the large step size is first kept constant and then decayed, potentially multiple times. We illustrate this behavior in Fig. 1 where we reproduce minimal setting without data augmentation or momentum, and with only one step size decrease. We draw attention to two key observations regarding the large step-size phase: (a) quickly after the start of training, the loss remains approximately constant on average and (b) despite no progress on the training loss, running this phase for longer leads to better generalization. We refer such large step-size phase as loss stabilization. The better generalization hints at some hidden dynamics in the parameter space not captured by the loss curves in Fig. 1. Our main contribution is to unveil the hidden dynamics behind this phase: loss stabilization helps to amplify the noise of SGD that drives the network towards a solution with sparser features (see Appendix, Figure for 2D-visualzzation).
1.1
OUR CONTRIBUTIONS
The effective dynamics behind loss stabilization. We characterize two main components of the SGD dynamics with large step sizes: (i) fast movement determined by the bouncing directions

---

Under review as a conference paper at ICLR 2023
40%
Thaining loss
decay
35%
30-1
U. decay at 100 spocte
(i 5. decay at 20% spache
22%
di desay at 50% spocts
1014
30%
15%
Epoch
Epoch
Figure 1: A typical training dynamics for a ResNet-18 trained on CIFAR-10. We use weight decay but momentum or data augmentation for this experiment. We a substantial difference in generalization (as large as 12% vs. 35% test error) depending on the step size n and its schedule. When the training loss stabilizes, there is hidden progress occurring which aim to characterize.
causing loss stabilization, (ii) a slow dynamics driven by the combination of the gradient and the multiplicative noise-which is non-vanishing due to the loss stabilization.
SDE model and sparse feature learning. We model the effective slow dynamics during loss stabilization by a stochastic differential equation (SDE) whose multiplicative noise is related to the neural tangent kernel features, and validate this modeling experimentally, Building on the existing theory on diagonal linear networks, which shows that this noise structure leads to sparse predictors, we conjecture a similar "sparsifying" effect on the features of more complex architectures. We experimentally confirm this neural networks of increasing complexity.
Insights from our understanding. We draw a clear general picture: the hidden optimization dynamics induced by large step sizes and loss stabilization enable the transition to a sparse feature learning regime. We argue that after a short initial phase of training, SGD first identifies sparse features of the training data and eventually fits the data when the step size is decreased. Finally, we discuss informally how many deep learning regularization methods (weight decay, BatchNorm, SAM) may also fit into the same picture.
1.2 RELATED WORK
He et al. (2016) popularized the piecewise constant step-size schedule which often exhibits clear loss stabilization pattern. However, they did not provide any explanations for such training dynamics and its implicit regularization effect. Non-monotonic patterns of the training loss have been explored in recent works. However, the loss stabilization regime we consider is different (i) from the catapult mechanism (Lewkowycz et al., 2020) where the training loss shows only one spike at the start of training and then monotonically converges without stabilization, and (ii) from the edge of stability regime of full-bated GD (Cohen et al., 2021) where the training loss shows many regular spikes after some point in training but again without stabilization.
Past works conjectured that large step sizes induce the minimization of some hidden complexity measures related to flatness of minima (Keskar et al., 2016; Smith and Le, 2018). Notably, Xing al. (2018) point out that SGD moves through the loss landscape bouncing between the walls of a valley where the role of the step size is to guide the noisy iterates of SGD towards a flatter minimum. However, many typically used flatness definitions are questionable for this purpose since (1) they are not invariant under reparametrizations that lead to an equivalent neural network (Dinh et al., 2017), and (2) even for naturally trained networks, full-batch gradient descent with large step sizes (unlike SGD) can lead to flat solutions which are not eelergenraliiing (Kaur et al., 2022). Note that it is possible to bridge the gap between GD and SGD by using explicit regularization as Geiping et al. (2022). We instead focus on the implicit regularization of SGD which remains the most practical approach for training deep networks.
The importance of large step sizes has been investigated with diverse motivations. However, we believe that existing approaches do not sufficiently capture the hidden stochastic dynamics behind the loss stabilization phenomenon observed for deep networks. Attempts to explain it on strongly convex models (Nakkiran, 2020; Wu et al., 2021; Beugnot et al., 2022) are inherently incomplete since is phenomenon related to the existence of many zero solutions with very different generalization
2

---

Under review as a conference paper at ICLR 2023
properties. Li et al. (2019b) analyzed the role of loss stabilization for synthetic distribution containing different patterns, but it is not clear how this analysis can be extended to general problems. Works based stability analysis characterize the properties of the minimum that SGD or GD can potentially converge depending on the step size (Wu et al., 2018; Mulayoff et al., 2021; Ma and Ying, 2021; Nacson al., 2022). However, these approaches do not capture the entire training dynamics such as the large step size phase that we consider where SGD converges only after the step size is decayed. SGD with label noise has been studied because of its beneficial regularization effect and its resemblance to SGD's 's standard noise. Its implicit bias has been first characterized by Blanc et al. (2020) and extended by Li al. (2022). However, their analysis only holds in the final phase of the training, close to zero-loss manifold. Our work instead is closer in spirit to Pillaud- Vivien et al. (2022) where the label noise dynamics analyzed in the central phase of the training, i.e., when the training loss is still substantially above zero.
2 THE EFFECTIVE DYNAMICS OF SGD WITH LARGE STEP-SIZE: SPARSE FEATURE LEARNING
In this section, we show that large step sizes lead the loss stabilize by making SGD bounce above a valley. We then unveil the effective dynamics induced by this loss stabilization. To clarify our exposition showease our results for the mean square error but other losses like the cross-entropy carry the same key properties in terms of the noise covariance (Wojtowytsch, 2021b, Lemma 2.14). We consider a generic parameterized family of prediction functions setting which encompasses neural networks. In this case, the training loss on input/output samples (xs, Vilisign Rd IR reads
{z + h1((z), 0 E R²),
1 n
[()) = 2n Σ (ha y1)²
(1)
We consider the overparameterized setting, i.e. p > n, hence, there shall exists many parameters l that lead to zero loss, i.e., perfectly interpolate the dataset. Therefore, the question of which interpolator the algorithm converges to is of paramount importance in terms of generalization. We
focus on the SGD recursion with step size 17 O, initialized at do E IP: for all t E N,
0.+1 = 0, n7 (zx,) ys, VV Voho, (xi,),
(2)
where i, ~ U n]) is the uniform distribution over the sample indexes. In the following, note that SGD with mini batches of size B 1 would lead to similar analysis but with // B instead of 17.
2.1 BACKGROUND: SGD IS GD WITH SPECIFIC LABEL NOISE
To emphasize the combined roles of gradient and noise, we highlight the connection between the SGD dynamics and that of full-batch GD plus a specific label noise. Such manner of reformulating the dynamics has already been used in previous works attempting to understand the specificity of the
SGD noise (HaoChen et al., 2021: Ziyin et al., 2022). We formalize it in the following proposition.
Proposition 1. Let (0)) follow the SGD dynamics (2) with sampling function (iz) Let 15=i, be indicator function, define for t 2 0, the random vector & E RR" such that for all i E [1, n].
[e], : (he, (xx) y1)(1 - n1;=))).
(3)
Then follows the full- gradient dynamics on C with label noise (Ee)12o that is
Oct1 0.
7
(x:) - Vahn,
n
where we define the random labels 3/ := + & Furthermore, &e is a mean zero random vector with
variance such that m(n-1) E ||E 2C(c.).
This reformulation shows two crucial aspects of the SGD noise: (i) the noisy part at state 0 always belongs to the linear space spanned by (Vohg (z1).... Vohe(w,)) loss. Concerning (ii), we highlight in the following section that the loss
(x,)), and (ii) it scales as the training
stabilize because of large step sizes which makes the effective scale of label noise constant. These features are of paramount
importance when modelling the effective dynamics that take place during loss stabilization.
3

---

Under review as a conference paper at ICLR 2023
2.2 THE EFFECTIVE DYNAMICS BEHIND LOSS STABILIZATION
On loss stabilization. For generic quadratic costs, e.g., F(3) = X - yl2.. gradient descent with step size 1/ is convergent for " 2/ 2//max- divergent for 1/ > 2/Amax and converges bouncing 2-periodic dynamics for 2/ Xmaxx where Amax is the largest eigenvalue of the Hessian. However, the practitioner is not likely hit perfectly this unstable step size and. almost surely, the dynamics shall either converge diverge. Yet, non-quadratic costs bring to this picture particular complexity: it has been shown that, even for non-convex toy models, there exist an open interval of step sizes for which the gradient descent neither converge nor diverge (Ma et al., 2022; Chen and Bruna, 2022). As we are interested in SGD, we complement this result by presenting a toy example in which loss stabilization occurs almost surely in the case of stochustic updates. Indeed, consider a regression problem with quadratic parameterization on one-dimensional data inputs x; 's, coming from distribution P, and outputs generated by the linear model y, 2,6?. The loss writes F(0)= == Ep (y- =x²)²² and the SGD iterates with step size rj > 0 follow, for any E N,
0c+1 = 0, + + 0z Ii, (v.. Ti, 2i,07)
where
Ii, ~ P.
(5)
For the sake of concreteness and clarity, suppose that 0. and supp(p) = (a, b), we have the following proposition (a more general result can be found in Proposition 3 of the Appendix).
Proposition 2. For any 1) (a=2, 1.2.5 b b-2) and initialization 0g E (0, for t > 0.
81 < F(O,)= < 82
almost surely, and almost surely.
(6)
3 > Vk T
O142x < < O2+28+1
where ,1, 82 T > O are constant given in the Appendix.
The proposition is divided in two parts: if the step size is large enough, (q.(6) the loss stabilizes in between level sets óN and 82 and (7) shows that after some initial phase, the iterates bounce from one side of the loss valley to the other one. Note that despite the stochasticity of the procedure, the results hold almost surely.
The effective dynamics. As observed in the prototypical SGD training dynamics of Fig. 1 and proved in the non-convex toy model of Proposition 2, large step sizes lead the loss to stabilize around some level set. To further understand the effect of this loss stabilization in parameter space, we shall assume perfect stabilization. Then, from Proposition 1, we conjecture the following behaviour
During loss stabilization, SGD is well modelled by GD with constant label noise.
Label noise dynamics have been studied recently (Blanc et al., 2020; Damian et al., 2021; Li et al., 2022) thanks their connection with Stochastic Differential Equations (SDEs). To properly write a SDE model, the drift should match the gradient descent and the noise should have the correct covariance structure (Li al., 2019a; Wojtowytsch, 2021a). Proposition I implies that the noise at state 0 is spanned by the gradient vectors (Voho( (..)...., Voho (x..) and has a constant intensity corresponding to the loss stabilization at a level > Hence, we propose the following SDE model
....,
d0, (o + Vnă de, ddB,
(8)
where (Bt):20 ):20 is standard Brownian motion in R." and (x( == [Vho (x:) T |==1 E Raax referred to the Neural Tangent Kernel (NTK) feature matrix (Jacot et al., 2018). This SDE can be seen as the effective slow dynamics that drives the iterates while they bounce rapidly in some directions at the level set d (fast dynamics). It highlights the combination of the deterministic part of the full-batch gradient and the noise induced by SGD at level set which depends on the step size of SGD. We confirm the validity of this SDE modeling empirically in Sec. C showing that the SDE captures the dynamics of large step size SGD even for non-linear networks. In the next section, we leverage the SDE (8) understand the implicit bias of such learning dynamics.
2.3 SPARSE FEATURE LEARNING
In this section, we give insights the effective dynamics given by Eq)( We begin with simple model of diagonal linear networks that showease a sparsity inducing dynamics and further disclose our general message about the overall implicit bias promoted by the effective dynamics.
4

---

Under review as a conference paper at ICLR 2023
2.3.1 A WARM-UP: DIAGONAL LINEAR NETWORKS
An appealing example of simple non-linear networks that help in forging an intuition for more complicated architectures is diagonal linear networks (Vaskevicius et al., 2019; Woodworth et al., 2020; HaoChen et al., 2021; Pesme et al., 2021). They are two-layer linear networks with only diagonal connections: the prediction function writes (z)
(u, v C z) (u + e, e,z) where @ denotes elementwise multiplication. Even though the loss is convex in the associated linear predictor B = u @ v E TR", it is not in (u, v),, hence the training of such simple models already exhibit a rich non-convex dynamics. In this case, (x)) = 0 ④ x, and the SDE model Eq.(8) writes
du
X (X @ vn) - v) @ Us dt + V 18 @ [X xd
(9)
where (Bt)c20 is a standard Brownian motion in I.. Equations are symmetric for (v:)c20-
What is the behaviour of this effective dynamics? Pillaud- Vivien et al. (2022) answered this question by analyzing a similar stochastic dynamics and unveiled the sparse nature of the resulting solutions. Indeed, under sparse recovery assumptions, denoting 3 the sparsest linear predictor that interpolates the data, it is shown that the associated linear predictor St = u, ) 0y: (i) converges exponentially fast zero outside of the support of 3 (ii) is with high probability in O(v1)) neighborhood of 3 in its support after a time O(8-1).
Overall conclusion on the model. During a first phase, SGD with large step sizes 7) decreases the training loss until stabilization at some level set 6 > 0. During this loss stabilization, an effective noise-driven dynamics takes place. It shrinks the coordinates outside of the support of the sparsest signal and oscillates in parameter space at level Opvipō) on its support. Hence, decreasing later the step size leads to perfect recovery of the sparsest predictor. This behaviour is illustrated in our experiments in Figure 2.
2.3.2 THE SPARSE FEATURE LEARNING CONJECTURE FOR MORE GENERAL MODELS
Results for diagonal linear nets recalled in the previous paragraph show that the noisy dynamics (9) induce a sparsity bias. As emphasized in HaoChen et al. (2021), this effect is largely due to the multiplicative structure of the noise v O [X cBB.] that, in this case, has a shrinking effect on the coordinates (because of the coordinate-wel wise multiplication with r). In the general case, we see, thanks to (8), that the same multiplicative structure of the noise still happens but this time with respect to the NTK feature matrix do (X). Hence, this suggests that similarly to the diagonal linear network case, the implicit bias of the noise can lead to shrinkage effect applied to do X) which depends on the noise intensity 8 and the step size of SGD. Indeed, an interesting property of Brownian motion is that, for U E Rº, (v, B+) = [| where the equality is valid in law and (Wi) is a -indimenniooal Brownian motion. Hence, the process (8) is equivalent to a process whose i-th coordinate is driven by a noise proportional to | |[WW. where di is the i-th column of Φo (X) and (W))r2o is a one dimensional Brownian motion. This SDE structure, similar the geometric Brownian motion, is expected to induce the shrinkage of each multiplicative factor (Oksendal, 2013, Section 5.1), i.e., in our case (II (lVoh)() Il Thus, we conjecture:
The noise part of (8) seeks minimize the εg-norm of the columns of (X)).
Note that the fitting part of the dynamics prevents the NTK feature matrix to collapse totally to zero, but as soon they not needed to fit the signal, columns be reduced to zero. Remarkably, from a stability perspective, Blanc et al. (2020) showed a similar bias: locally around a minimum, the SGD dynamics implicitly tries minimize the Frobenius norm (d)(X)]lF - (Voha(x))] We provide below a specification of this implicit bias for different architectures:
Diagonal linear networks: For (2) = (u O U, x), we have (u, = [0 2 x, u z]. Thus, for generic data matrix X, minimizing the norm of each column of pu,, (X) amounts to put the maximum of zero coordinates and hence to minimize I|u @ viloo
ReLU networks: We take the prototypical one hidden layer to exhibit the sparsification effect. Let ha.w (xx) = (a, o (Wr)), then V.ha, w (x) = o(Wx) and V ha,w(x) (x) = a, (1,=\(,,,=))000 Note that the form of the column corresponding to the neuron is reduced when it is activated at a minimal number of training points, hence the implicit bias enables the learning of sparse data-active features. Finally, when some directions are needed to fit the data, similarly activated neurons align fit, allowing the rank of de (X) to be also a good proxy for this feature sparsity.
5

---

Under review as a conference paper at ICLR 2023
Toin loss
Iteatiin
Iteratio
Iteration
Iteration
SGD 1- deccy
SCD -7. rway at 305 Receiiu
sm0 0.% decp 52% beattees
Figure 2: Diagonal linear networks. We observe loss stabilization, better generalization for longer schedules, minimization of the rank of dh (X) and sparsity of the predictor u O V.
Overall, fully understanding theoretically the structural implications of the implicit bias described above remains an exciting avenue for future work. We show next that the conjectured sparsity is indeed observed empirically for variety of models, well as that the rank reduction of (x( can be used as a good proxy of the hidden progress of the loss stabilization phase. This is confirmed both for SGD and its SDE modeling (the latter we show in App. C).
EMPIRICAL EVIDENCE OF SPARSE FEATURE LEARNING DRIVEN BY SGD
Here we present empirical resultsfor neural networks of increasing complexity: from diagonal linear networks to deep residual networks on CIFAR-10 and CIFAR-100. We make the following common observations for all these networks trained using SGD schedules with large step sizes:
(01) Loss stabilization: training loss stabilizes around a high level set until step size is decayed,
(02) Generalization benefit: longer loss stabilization leads better generalization,
(03) Sparse feature learning: longer loss stabilization leads to sparser features.
explicit regularization in our experiments that the training dynamics is driven purely by SGD and the step size schedule. Additionally, in some cases, we cannot find a single large step size that would lead to loss stabilization. In such cases, whenever explicitly mentioned, we use a warmup step size schedule-i.e. increasing step sizes according to some schedule- to make sure that the training loss stabilizes around some level set. Such warmup schedules used in practice (He et al., 2016; Devlin et al., 2018). Warmup is often motivated purely from the optimization perspective
Importantly,
commonly
way accelerate training (Agarwal et al., 2021) but we suggest that, more importantly, it is also a way to amplify the regularization effect of the SGD noise which is proportional the step size.
Measuring sparse feature learning, Our main insight is that the NTK feature matrix is significantly
simplified the loss stabilization phase, and that the rank of (X) (i.e., the sparsity of its singular values) is a good proxy to track this dynamics. We compute it over iterations for each model (except deep networks where it is not feasible) by using a fixed threshold on the singular values of dip (y((X (X) normalized by the largest singular value. In this way, we ensure that the difference in the rank that we detect is not simply due to different scales (X). Moreover, we always compute (x(X) (X) on the number of fresh samples equal to the number of parameters |0] to make that rank deficiency is not coming from n < which is the case in the overparametrized settings we consider.
Furthermore, we also want to track more direct and interpretable notion of feature sparsity. This motivates us to count the average number of distinct (i.e., counting a group highly correlated activations one), non-zero activations at some layer over the training set which we refer to as the feature sparsity coefficient. We count pair of activations i and j as highly correlated if their Pearson's 's correlation coefficient is at least 0.95. Unlike ank( (X)), the feature sparsity coefficient scales to deep networks and has an easy-to-graspp meaning.
3.1 SPARSE FEATURE LEARNING IN DIAGONAL LINEAR NETWORKS
Setup. The inputs ,,...., Fr with n = 80 are sampled from (0,, Ia) where I is an identity matrix with d = 200, and the outputs generated as yı (3., where B. Rd is 20 sparse. We

---

Under review as a conference paper at ICLR 2023
Iterwtion
0.000 Peradim
5123, 0-2000
Racstion
Figure 4: Two-layer ReLU networks for ID regression. We observe loss stabilization, simplifica- tion of the model trained with longer schedule, lower rank of (X),,, and much sparser features.
consider four different SGD runs (started from u4 0.1, U, 0 for each i): one with small step size and three other with initial large step size decayed after 10,, 30%, 50% iterations, respectively.
Observations. We show the results in Fig. 2 and note that (01)-(03) hold even in this simple model trained with vanilla SGD without any explicit regularization or layer normalization schemes. We observe that the training loss stabilizes around 10-1.3, the test loss improves for longer schedules, both rank(og (X)) and ||u O ol1o decrease during the loss stabilization phase leading to sparse final predictor. While the training loss has seemingly converged to 10 1.5 hidden dynamics suggested by (9) occurs which slowly drifts the iterates to sparse solution. This implicit sparsification explains the dependence of the final test loss on the time when the large step size is decayed, similarly to what has been observed for deep networks in Fig. 1. Interestingly, we also note that SGD with large step-size schedules encounters saddle points after we decay the step size (see the training loss curves in Fig. 2) which resembles the saddle-to-saddit regime described in Jacot et al. (2021) which does not occur in the large-initializzn initialization lazy training regime.
Training loss 2D subspace
SGD and GD have different implicit biases. Since observe from Fig. 2 that for loss stabilization, stochasticity alone does not suf- fice and large step sizes are necessary, one may wonder if conversely only large step sizes can be sufficient to have a sparsifying effect. Even if special instances can be found for which large step sizes are sufficient (such as for non-centered input features as in Nacson et al. (2022)), we answer this negatively showing that gradient descent in general does not go to the sparsest solution as demonstrated in Fig. 10 in the Appendix. Moreover, in Fig. we visualize the dif- ference in trajectory between the two methods taken with large step sizes over a 2D subspace spanned by w* Winit and W flow Winits where w* is the ground truth, flow is the result of gradient flow, and Winit is the initialization. This example provides an important intuition that loss stabilization alone is not sufficient for sparsification and that the role of noise described earlier is crucial.
GD
SGD
Figure 3: GD and SGD take different trajectories.
3.2 SPARSE FEATURE L.AARNING IN SIMPLE RELU NETWORKS
Two-layer ReLU network in ID. We consider the one-dimensional regression task from Blanc et al. (2020) with 12 points, where label noise SGD has been shown to learn a simple model. We show that similar results can be achieved with large-step-size SGD via loss stabilization. We train a ReLU network with 100 neurons with SGD with a linear warmup (otherwise, we were unable achieve approximate loss stabilization). directly followed by step-size decay. The two plots correspond to warmup/decay transition 2% and 50% of iterations, respectively. The results shown in Fig. confirm that (01)-(3) hold: the training loss stabilizes around 10 0.5 the predictor becomes much simpler and is expected to generalize better, and both rank (o (X)) and the feature sparsity coefficient substantially decrease during the loss stabilization phase. For this one-dimensional task, we can directly observe the final predictor which is sparse in terms of the number of distinct ReLU kinks as captured by the feature sparsity coefficient and the rank of the NTK feature matrix. Interestingly, we also observed overregularization for even larger step sizes when we cannot fit all the training points (see Fig. 11 in Appendix). This phenomenon clearly illustrates how the capacity control is induced by the optimization algorithm: the function class over which we optimize depends on the step size schedule. Additionally, Fig. 12 in Appendix shows the evolution of the predictor over iterations. The
7

---

Under review as a conference paper at ICLR 2023
Ierracin
Iteration
Tuoration om
Berraton
Figure 5: Thyee-layer ReLU networks in a teacher-student setup. We observe loss stabilization, lower rank of the NTK feature matrix and lower feature sparsity coefficient on both hidden layers.
ResNet-18 on R-10, no explicit regularization
Epoch
Spoch
Epoch
Epoch
4 0.03
6.075 warmup, decay 10%
n- 0.75, wammup. decay at 30%
4= 575, warmun, decay at 50%
18 on 100,R-10, statc aht setting
300
Training
Epoch
Epoch
Epoch
Epoch
9- D.0GZ
D.05, decay 10%
n- 0.25, daciyy 30%
0.55 dacay 50%
Figure 6: ResNet-18 trained on CIFAR-10. Both without explicit regularization and in the state- of ah- setting, the training loss stabilizes, the test loss noticeably depends the length of the schedule, and the feature sparsity coefficient is minimized over iterations.
general picture is confirmed: first the model is simplified during the loss stabilization phase and only then fits the training data.
We use a teacher-student setup with a random three-layer teacher ReLU network having neurons on each hidden layer. The student network is overparametrized with 10 neurons on each layer and is trained on 50 examples. Such teacher-student setup is useful since we know that the student network can implement the ground truth function but might not find it due to the small sample size. We train models using SGD with a medium constant step size and large step size with warmup decayed after 10%, 30%, 50% iterations, respectively, The results shown in Fig. confirm that (01)-(Q (O3) hold: the training loss stabilizes around 10 f.5 the test loss is smaller for longer schedules, and both rank(ow (X)) and the feature sparsity coefficient substantially decrease during the loss stabilization phase. All methods have the same value of the training loss (10-3) after 1 iterations but different generalization. Moreover, we see that the feature sparsity coefficient decreases on each layer which makes this metric a promising one to consider for deeper networks.
Deeper ReLU networks.
3.3 SPARSE FEATURE LEARNING DEEP RELU NETWORKS
Setup. We consider here an image classification task and train a ResNet-18 18 and ResNet-34 on CIFAR-10 and CIFAR-100 using SGD with batch size 256 and different step size schedules. We use an exponentially increasing warmup schedule with exponent 1.05 to stabilize the training loss. We cannot measure the rank of ((X) here since this matrix is too large (Rz 50 000 x 20 000 000) so measure only the feature sparsity coefficient taken at two layers: at the end of super-blocd block 3 (i.e., in the middle of the network) and super-block 4 (i.e., right before global average pooling the end of the network) of ResNets. We test two settings: a basic setting without explicit regularizers and a Hate-offthe-ars setting with weight decay, momentum, and standard augmentations.
8

---

Under review as a conference paper at ICLR 2023
Observations. The results on CIFAR-I0 shown in Fig. 6 confirm that our main findings still hold also in this setting: the training loss stabilizes either slightly below 10-1 or above 10 10-1, the test error is becoming progressively better for longer schedules, as well as the feature sparsity coefficient. Small step sizes lead to bad generalization, especially without explicit regularization: 35% test error compared to 15% for large step sizes. This poor performance confirms that it is crucial to leverage the implicit bias of large step sizes. The difference in the feature sparsity coefficient is also substantial with the final model having 70% instead of 24% block 4 without explicit regularization. The observations are similar for the -tate-of- offe-of-the--ane setting as well where even with explicit regularization, we still see a noticeable difference in generalization and feature sparsity depending on the step size and schedule. We further note that feature sparsity coefficient is gradually minimized over iterations in this case (similarly to Figures 2, 4. 5) while without explicit regularization we observe a different pattern: a very quick drop down to almost zero at the very first epoch and then a gradual increase.
We show the results with similar findings on CIFAR-100 in Fig. 15 in Appendix. Additionally, Fig. 14 illustrates that for small step sizes, the early and middle layers stay very close to their random initialization which indicates the absence of feature learning similarly to what is suggested by the neuron movement plot in Fig. 9 in the Appendix for two-layer network in teacher-student setup.
4 INSIGHTS FROM OUR UNDERSTANDING OF THE TRAINING DYNAMICS
Here we provide an extended discussion on the implications of our theoretical and empirical findings.
The multiple stages of the SGD training dynamics. As analyzed and shown empirically, the training dynamics we considered can be split onto three distinct phases: (i) initial phase of reducing the loss down to some level where stabilization can occur, (ii) a loss stabilization phase where noise and gradient directions combine to find sechittctuur-deeeedeee
sparse representations of the data, (iii) a final phase when the step size is decreased to fit the training data. This typology allows to clearly disentangle the effect of the stabilization phase (ii) which relies on the implicit bias of SGD to simplify the model. Note that phases (ii) and (iii) can be repeated a few times until final convergence (He et al., 2016). Moreover, in some training schedules, (ii) does not explicitly occur, and the effect of loss stabilization (ii) and data fitting (iii) can occur simultancously (Nakkiran et al., 2019).
From lazy training to feature learning. Similar sparse implicit biases have been shown for regression with infinitely small initialization (Boursier et al., 2022) and for classification (Chizat and Bach, 2020; Lyu and Li, 2020). However, both approaches are not practical from the computational point of view since (i) the origin is a saddle point for regression leading to the vanishing gradient problem (especially, for deep networks), and (ii) max-margin bias for classification is only expected to happen in the asymptotic phase (Moroshko et al., 2020). On the contrary, large step sizes enable to initialize far from the origin, while allowing to efficiently transition from a regime close to the lazy NTK regime (Jacot et al., 2018) to the rich feature learning regime.
Common patterns in the existing techniques. Tuning the step size to obtain loss stabilization can be diffieult. To prevent early divergence caused by too large step sizes, we sometimes had to rely on an increasing step size schedule (known as warmup). Interpreting such schedules as a tool to favor implicit regularization provides a new explanation their success and popularity. Additionally, normalization schemes like batch normalization or weight decay, beyond carrying their own implicit or explicit regularization properties, can be analyzed from a similar lens: they allow to use larger step sizes that boost further the implicit bias effect of SGD while preventing divergence (Bjorck et al., 2018; Zhang et al., 2018). Note also that we derived our analysis with batch size equal to one for the sake of clarity, but an arbitrary batch size B would simply be equivalent to replacing Y + x Similarly to the consequence of large step sizes, preferring smaller batch sizes (Keskar et al., 2016) while avoiding divergence seem key to benefit from the implicit bias of SGD. Finally, the effect of large step sizes or small batches is often connected to measures of flatness of the loss surface via stability analysis (Wu et al., 2018) and some methods like the Hessian regularization (Damian et al., 2021) or SAM (Foret et al., 2021) explicitly optimize it. Such methods resemble the implicit bias of SGD with loss stabilization implied by the label noise equation (Eq. where matrix (o(X) is the key component of the Hessian. However, an important practical difference is that the regularization strength in these methods is explicit and decoupled from the step size schedule which may be harder to properly tune since it is simultancously responsible for optimization and generalization.

---

Under review as a conference paper at ICLR 2023
REFERENCES
N. Agarwal, S. Goel, and C. Zhang. Acceleration via fractal learning rate schedules. In International Conference on Machine Learning. pages 87-99. PMLR, 2021.
G. Beugnot, Mairal, and Rudi. On the benefits of large learning rates for kernel methods. arXiv preprint arXiv:2202.1377 2022.
N. Bjorck, C. P. Gomes, B. Selman, and K. Q. Weinberger. Understanding batch normalization. Advances in neural information processing systems, 31, 2018.
G. Blanc, N. Gupta, G. Valiant, and P. Valiant. Implicit regularization for deep neural networks driven by an ornstein-uble -hheein-uuhlebeckk like process. In Conference on learning theory, pages 483-513, PMLR, 2020.
E. Boursier, L. Hillaud-iivien and N. Flammarion. Gradient flow dynamics of shallow relu networks for square loss and orthogonal inputs. arXiv preprint arXiv:2206.0099 00039,, 2022.
I. Chen and J. Bruna. On gradient descent convergence beyond the edge of stability, arXiv preprint arXiv:2206.00 04172, 2022.
L. Chizat and F. Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on Learning Theory, pages 1305-1338. PMLR, 2020.
L. Chizat, E. Oyallon, and F. Bach. On lazy training in differentiable programming. Neural Information Processing Systems, 32, 2019.
Advances in
J. M. Cohen, S. Kaur, Li, J. Z. Kolter, and A. Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.00065, 2021.
A. Damian, T. Ma, and D. Lee. Label noise sgd provably prefers flat global minimizers. Advances in Neural Information Processing Systems, (1227449-27666 2021.
J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.0488) 2018.
L. Dinh, R. Pascanu, S. Bengio, and Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, pages 1019-1028. PMLR, 2017.
Foret, A. Kleiner, H. Mobahi, and B. Neyshabur. Sharpness- aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021.
J. Geiping, M. Goldblum, P. E. Pope, M. Moeller, and T. Goldstein. Stochastic training is not necessary for generalization. In International Conference on Learning Representations, 2022.
J. Z. HaoChen, C. Wei, J. D. Lee, and T. Ma. Shape matters: Understanding the implicit bias of the noise covariance. In Conference Learning Theory, 2021.
K. He, X. Zhang, Ren, and J. Sun. Deep residual learning for image recognition. CVPR, 2016.
A. Jacot, Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization neural networks. Advances in neural information processing systems, 31, 2018.
A. Jacot, F. Ged, B. Şimşek, C. Hongler, and F. Gabriel. Saddle-to-sadde dynamics in deep linear networks: Small initialization training, symmetry, and sparsity. In International Conference on Machine Learning, 2021.
Kaur, Cohen, and Z. C. Lipton. On the maximum Hessian eigenvalue and generalization. arXiv preprint arXiv:2206. 10654, 2022.
N. S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: 1699.068996 2016.
A. Lewkowycz, Y. Bahri, E. Dyer, J. Sohl-Dickstein, and G. Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint 2xXiv:2003.02211 2020.
10

---

Under review as a conference paper at ICLR 2023
Q. Li, C. Tai, and E. Weinan. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. The Journal of Machine Learning Research, 20(1): 1474-1520, 2019a.
Y. Li, T. Ma, and H. Zhang. Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory, pages 2-47. PMLR, 2018.
Y. Li, C. Wei, and T. Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. NeurlPS, 2019b.
Z. Li, T. Wang, and S. Arora. What happens after sgd reaches zero loss?-a mathematical framework. In International Conference on Learning Representations, 2022.
K. Lyu and J. Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations, 2020.
C. Ma and L. Ying. The Sobolev regularization effect of stochastic gradient descent. arXiv preprint 1xXiv:2105.1462, 2021.
C. Ma, L. Wu, and L. Ying. The multiscale structure of neural network loss functions: The effect on optimization and origin. arXiv preprint arXiv:2204. 1326, 2022.
E, Moroshko, B. Woodworth, S. Gunasekar, J. D. Lee, N. Srebro, and D. Soudry. Implicit bias in deep linear classification: Initialization scale vs training accuracy. In Advances in Neural Information Processing Systems, volume 33, pages 22182-22193, 2020.
R. Mulayoff, T. Michaeli, and D. Soudry. The implicit bias of minima stability: A view from function space. Advances Neural Information Processing Systems, 34:17749-17761, 2021.
M. S. Nacson, K. Ravichandran, N. Srebro, and D. Soudry. Implicit bias of the step size in linear diagonal neural networks. International Conference on Machine Learning, pages 16270-16295. PMLR, 2022.
Nakkiran. Learning rate annealing can provably help generalization, even for convex problems. arXiv preprint 20Xv:2005.07360 2020.
Nakkiran, G. Kaplun, Y. Bansal, Yang, B. Barak, and I. Sutskever. Deep double descent: Where bigger models and more data hurt. International Conference on Learning Representations, 2019.
B. Oksendal. Stochastic differential equations: introduction with applications. Springer Science & Business Media, 2013.
S. Pesme, L. Pillaud-Vivve Vivien, and N. Flammarion. Implicit bias of sgd for diagonal linear networks: provable benefit of stochasticity. In Advances Neural Information Processing Systems, 2021.
L. Fillaud-iiveen J. Reygner, and N. Flammarion. Label noise (stochastic) gradient descent implicitly solves the lasso for quadratic parametrisation. In Conference on Learning Theory, 2022.
S. L. Smith and Q Le. A Bayesian perspective on generalization and stochastic gradient descent. In International Conference on Learning Representations, 2018.
D. Soudry, F. Hoffer, M. S. Nacson, S. Gunasekar, and N. Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research, 19(1):282222875 2018.
T. Vaskevicius, V. Kanade, and P. Rebeschini. Implicit regularization for optimal sparse recovery. Advances in Neural Information Processing Systems, 32, 2019.
S. Wojtowytsch. Stochastic gradient descent with noise machine learning type. Part II: Continuous time analysis. arXiv preprint arXiv:2106.02588 2021a.
S. Wojtowytsch. Stochastic gradient descent with noise machine learning type. Part Discrete time analysis. arXiv preprint arXiv:2105.01650, 2021b.

---

Under review as a conference paper at ICLR 2023
B. Woodworth, S. Gunasekar, J. D. Lee, E. Moroshko, P. Savarese, Golan, D. Soudry, and N. Srebro. Kernel and rich regimes in overparametrized models. Proceedings Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 3635-3673. PMLR, 09-12 Jul 2020.
J. Wu, D. Zou, V. Braverman, and Q. Gu. Direction matters: On the implicit bias of stochastic gradient descent with moderate learning rate. International Conference on Learning Representations, 2021.
L. Wu, C. Ma, al. How sgd selects the global minima in over wer-pparameterized learning: dynamical stability perspective. Advances in Neural Information Processing Systems. 31, 2018.
C. Xing, D. Arpit, C. Tsirigotis, and Y. Bengio. walk with sgd. arXiv preprint arXiv: 1802.08770 2018.
C. Zhang. S. Bengio, M. Hardt, B. Recht, and O Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017.
G. Zhang, C. Wang, B. Xu, and R. Grosse. Three mechanisms of weight decay regularization. In International Conference on Learning Representations, 2018.
L. Ziyin, K. Liu, T. Mori, and M. Ueda. Strength of minibatch noise in sgd. In International Conference on Learning Representations, 2022.
12

---

Under review as a conference paper at ICLR 2023
A PPENDIX
In Section A, we show Proposition on the equivalence between SGD and GD with added noise. In Section B. we provide the proof that loss stabilization occurs as written in Proposition 2. In Section C, we show experimentally that the proposed SDE model matches well the SDE dynamics. Finally, we present additional experiments in Section
Figure 7: dirre-dimensional visualisation of the SGD dynamics in a non-convex loss landscape. The SGD dynamics (blue points) is bouncing side-to-side to the bottom of the valley (the dotted green line). A slow movement occurs pushing the iterates in the direction given by the green arrows.
To begin this appendix, we provide in Figure 7 a toy visualization in which we showcase a typical SGD dynamics when loss stabilization occurs. We run SGD on the diagonal linear network with one sample in two dimensions (n = 1, 2) adding label noise of the shape given by equation (9),)), with balanced layers u = v. The blue points corresponds to iterates of the dynamics (that are linked with the orange dotted lines). The green line corresponds to the global minimum of the loss, what can be called the "bottom of the valley". This hopefully will serve the reader forge a visual intuition on (i) the bouncing dynamics side-to-side the bottom of the valley (in green), and (ii) the slow stochastic movement (in the direction of the green arrows).
A SGD AND LABEL NOISE GD
For the sake of clarity we recall below the statement of the Proposition which we prove in this section.
Proposition 1. Let (0i) follow the SGD dynamics E2) with sampling function (ft)c2o Let
1i=i, be indicator function, define for 2 0, the random vector & E RR" such that for all i E [1, n].
[Gt] == (I;) 11)(1 -1)m))))
(10)
Then (ilioffollows the full-batch gradient dynamics on L with label noise (Ee) that is
0.t 0
7
(ho. Voha,
(11)
T2
where we define the random labels 3/¹ = 21 &s. Furthermore, &e a mean zero random vector with variance such that min-11 E 1|&. -C(()).
Proof. Note that
(ho, (xx)) Voho,
(x:) 34 -lEl) Voho, (2:).
(12)
1=1
13

---

Under review as a conference paper at ICLR 2023
Using [Erla = (ho, (24) y1)(1 mll=)))
1
he, (244) y1 - (20) - -v:) nl 1) Voho,
(13)
n
Σ 1 (z,) - y4) Vi)Vohe = (ho, (xi.) - Voho,
(14)
which is exactly the stochastic gradient wrt to sample (Ii,, Mir))
Now we prove the latter part of the proposition regarding the scale of the noise. Recall that, for all & n, we have [gt]; = (ho, (x:) - y1)(1 nl;=i,,, where it ~ 24 ([], Now taking the expectation,
ElE,] E (ho,())) y1 m1)= = (ho, - 21)(1 nE [1==4,] 0, (15)
as E [1== 1/n. Coming to the variance,
EE ||E.|1 E
[E 2
B[Ga]:
(16)
(he, (xx;) - -y,) [(1 - n n1:==1)9 )
(17)
Σ (ho, (2:) - =3 [(1 - 2n ;-i, + 122 1==))]
(18)
(ho, (2:)) =1))²11 2 n)
(19)
(n- (ao,(x;) 1s)² = 2n(n - 1))(
(20)
and this concludes the proof of the proposition.
B QUADRATIC PARAMETERIZATION IN ONE DIMENSION
Again, for the Appendix to be self-contained, we recall the setup of the Proposition 2 on loss stabilization. We consider a regression problem with quadratic parameterization on one-dimensional data imputs z,'s, coming from a distribution P. and outputs generated by the linear model y = 2,02.
The loss writes F(0) = 1Rp (y- - =0 and the SGD iterates with step size vj > 0 follow, for any t N.
0,+1 = 0, + n 02 Fi, (Vi, Fi, 07)
where Ii, ~ p.
(21)
We rewrite the proposition here. Proposition 3. (Extended version of Proposition 2) Assume 3 Fmin Tmax > 0 such that supp(p) C [rmin Tmaxx] Then for any v) E ((0..min)) 1..5 -2), any initialization in Og (0, 0.,) for t E N, we have almost surely
F(O.) E (e. 02,0.77
(22)
where min { (1((0 - 1)/3,0.02}. Also, almost surely, there exists t, > 0 such that 0142 E (0.65 0,, (1 -6) 0.) and 01-24+1 E ((1 + c.) 8., 1.162 0.).
Proof. Consider SGD recursion (21 and note that x02.
0c+1 0, + 17 02 x xx02 x0?)
(23)
0 = 0, + n 0 2 (0 0?)
(24)
14

---

Under review as a conference paper at ICLR 2023
For the clarity of exposition, consider the rescaled recursion of the original SGD recursion.
0c+1/0. = 01/0. + 0- x² 01/0.
(1 (2//e)))
(25)
and, by making the benign change 8, +- 1-0/ we focus on the stochastic recursion instead,
0++1 0, (1 0,),
(26)
where 7 ~ fyy the pushforward á under the application + n 0² x². Let r ;= supp((,). the support of the distribution of y. From the range of 17, it can be verified that I S (1.1.25). Now the proof the theorem follows from Lemma 5.
Lemma 4 (Bounded Region). Consider the recursion Eq. I S (1,1.25) and 0 < then all t > 0, 0, E (0, L.162).
Proof. Consider a single step of (2,( for some 7 E 1.25),
0+ 0 + 70 6
The aim is to show that 0+ stays in the interval (0, 1.162). In order to show this, we do casewise analysis.
For 0 E 10: Since 0 < 0 s 1, we have 04 2 0 > 0. To prove the bound above, consider the following quantity,
mmaz
max max TE(1,1.25) (E(1)
70(1 62)
(27)
Say (0) 0 + +0(11 - 82), note that (0) + 02 and (0) -670 Hence, for
any in our domain, the maximum is attained at 0,, 1 V 1 + and (0,)) 3/37
2(1+-)
2(1 + 3/32 7E(55.1.25
max max + 7 02) = max 7E(1.1.25) dE(U,1]
(28)
3/37
It can be verified that 2(1+>)*/²
is increasing with gamma in the interval (1, 1.25). Hence,
3/3y
2(1 7
2(1 + 7)
max yE(1,1255)
< 1.162
(29)
3/
3/3
y=1.25
Combining them, we get,
04 < max max 0 + 70(1 02) 1.162
(30)
7 26) 8E
For E 1111:: Since 8 1, we have, 0+ < A < 1.162. For lower bound, note that for 04 to be less than 0, we need 7 702 < 0. But for y E (1, 1.25) and 0 E (1,1.162),
7002 1) < 1.25 (1.112)) 1) < 1.
(31)
Hence, it never goes below 0.
Lemma S. Consider the recursion (26) with I' C 1.25) and 0g initialized uniformly in (0, 1). Then, there exists Co > such that for all cu there exists t such that for any 0,
Oz12k E (0.65, 1 c) and 0 (1 + c,
(32)
almost surely.
Proof. Define Ymin as the infimum of the support I'. Let min (Tmon 11/3, 0.023. Note that
co > 0 as "ynin > 1. Now for any 0 < c < 6o, we have Ymiu (2 <) c) >
Divide the interval (0.1.162) into 4 regions, Io (0, 0.65,, I (0.65, <). L [1 1). I3 = (1, 1.12). The strategy of the proof is that the iterates will eventually end up in In and that once it ends up in I,, it comes back to I in steps.
Let f be initialized uniformly random in (0, 1). Consider the sequence (fx) generated by 0+1 = h,, (0.) = 0, + 0x 0?) where "t ~ P..
(33)
We prove the following facts (P1)-(P4):
15

---

Under review as a conference paper at ICLR 2023
(P1) There exists t such that the 0g I UlUUs.
(P2) Let 0g I3, then 0t+1 E 1 Ul2.
(P3) Let 0g E 22, there exists & > 0 such that for k' < k. 0ct2x E 12 and 0442 E h
(P4) When 0, E 11, then for all & 2 0, 01+22 1 and 01+26+1 E (1 +1++0 e, 1,1662)
Proof of (P1)-(P4): Let E N, note first that the event (0: = 1} Usgr (0z = 1 * 1} and hence a finite union of zero measure sets. Hence (0, = 1 is a zero measure set and therefore we do not consider it below. For any other sequence, from the above four properties, we conclude that the lemma holds.
Proof of P1: Assume that until time > 0, the iterates are all in To. then have
8, = 0 7(1 1)) 12-1(2 62-1) 1.5 0.-1 > 1.5 8g
(34)
Hence, the sequence eventually exits Ig. We know that it will stay bounded from Lemma 4, hence will end up in 1 UI2 Us.
Proof of P2: For any 0, E 1.162,, 1< < < 1.25, since h, (.) is decreasing in (1,1.162), we have (1.162) < (0.) < (1). Also (0) is linear in gamma with negative coefficient for 0 > 1. Hence decreases as increases. Using this,
652 = h1 25 (1.162) < (1.162) < (0,) < h1(
(35)
Hence, 0.11 E 1 Ul.
Proof of P3: The proof of this follows from Lemma
Proof of P4: The proof of this follows from Lemma 10.
Lemma 6. For any E I U12 and any a, I, (a)) UI2
hymns (h-ymax (0) < ha (h) S hymin (hrmii (0)).
(36)
Proof. For any 7 I, recall
(0) = 0 7011
= -
-0--11
(37)
Note that for 0 E T U TT, 0(1 0) 1, Hence 70(1 0) > 1. This gives us that (0) > 1. Now we will track where T, T2 can end up after two stochastic gradient steps.
For any b r, as 0 e I U 12, have
hymaax (0) (0) hymin (0) > 1,
note hyma (0) (₂) h hym (0) holds since 8 1.
Now for any E I' and x > 1, (x(x) is decreasing function in 2. Hence
ha (hymar (0)) < ha (ha())) I ha (h (0)).
Using Tmin a. ho (h.yain (0)) < h (hymii (0)), Similarly using Tmax > a, we have, h (hyamx (8) < ha yuuas (0)). Combining them we get,
hyaaas (hrmax (0)) < ha (hah())) I h (h (0)).
(38)
Similar argument can extend it to,
h 25 (h125 < ha (O) < (O),
(39)
Lemma 7. Let 0g T2 there exists k such that 0442k E l
16

---

Under review as a conference paper at ICLR 2023
Proof.
For any y H, let 0+ = (0), then we have
h, - 0 = h0, - 10 + +0+
(40)
Furthermore,
0 = 0 + 70 0 = f(1 + 7(1 -))),
(41)
1 0 = 1 + +-0+00(11 0-) = (1 0) + 70 - 0)),
(42)
1 1 0+ 1 - -0 0²) (1 0) -0 + 0)).
(43)
And multiplying the above three terms and adding 0(1 - 0²), we get,
0+ -
+ 0 -
0²) = 0(1 0² (1+ )²) + - -0100000-000101000000000
P(0)
(44)
For E I2, using Yunin (2 c)(1 c) > we have the inequalities
(1 (1+7(1 7(1 (1 + + 0)) 1,
(45)
(1 (1-70(1+0)) < Mmin e) < -1,
(46)
P(O) < -1.
(47)
Hence,
(h, - 0 = >11 (2) P(O)) <
(48)
Therefore, for [1 e, 1). for any y E I, h, (0) < 0. Hence for any two stochastic gradient step with a, b I, from Eq. 0e+2 = ha (0)) S hymm (h..... (0:) < 0r. From any point in 12, we have 10++2 1| > 00 1/, for any a, b E Т. Intutively this means that in
gradient steps the iterates move further away from / until it eventually leaves the interval 12 as the sequence {0142 1k20 strictly decreasing with no limit point in l2. From Lemma 9, we know that in steps the iterates will never leave I U I2. Hence they will eventually end up in I leaving I2.
Property 8. Define (0) == (h, (0)) for the sake of brevity. The followings properties hold for 0 E I U 12, E r and O,, the root of by (0):
Q1 (0) 2 g,
Q2 The function (-(.) is decreasing in [0.65, 07) and increasing in (6y, 1].
Proof. Note (0) 11 -² has most root 0, E (0, 1). Note that for all y F,
0, E In U 12. For any 7. (0) = (h,(h, For any 0 E T U 12, we have, h²) >
1 h² (0) 0. Therefore, (0) has only one root in U T2. Since 0, I T2, note 8, (0,)) = hú (h, (o,)), > 0. Therefore, (.) attains its minimum at 6y and this shows the desired properties.
Lemma 9. For any 0 E 1 U and any a, a E I, ha (h) Ul
Proof. Lower Bound: From (39)),), we know
h (5 < ha (O)
(49)
We know that from property Q1 that g, (0) 2 (0,).. Hence
81.25 (5).2) < 21) < h (h,
(50)
It can be quickly checked that .65 < 81.25 (0 25). Hence the lower bound holds.
Upper Bound: From (39), we know
ha (h)( h1 (hz(9))
(51)
We know that from property Q2 that gi (9) < max(gi (1), gi (0.65)) : It can be easily verified that gi (0.65) < 0.98. Hence gi (0) < 1.
17

---

Under review as a conference paper at ICLR 2023
Lemma 10. For any 0 1 and any a, b r, (h)( E I1 and (a) (1 + c. 1.162).
Proof. The lower bound in Lemma holds here. For the upper bound, from and Eq.(36),
ha (nhh:(() I hymi yh (0)).
(52)
Using property Q2,
hrmmin (h yomie (0)) < max yomt (1 c), 7mmi (0.65)
(53)
From (48, Noni (1 -() 6) < 1 - c. From E39( 8 (0.65) < (0.65) < 0.98 < 1 - 6. In
1t the function (.) first increases reaches maximum and decreases. Hence for 0 E 41, (0) 2 m0n)6.(0.65) ha e)}
ha(1 2 1 - e + a(1 (1 (1) -),, c),
(54)
1 + ale - (1-1
(55)
2 + Mmmin (2c () - -))
(56)
= 1 + + (Tmin (2 - -))(1 () - > 1+
(57)
Also ha (0.65) > hi (0.65) > 1.02 > 1+c, therefore h, (0) > 1+ and this completes the proof. 0
C EMPIRICAL VALIDATION OF THE SDE MODELING
In this section, we experimentally check the validity of the SDE modeling of SGD in (8) in terms of the key metries: training loss, test loss, rank of the NTK feature matrix, and feature sparsity.
SDE discretization. Let "/e be the SDE discretization step size, t) the step size of the corresponding SGD that we aim validate, & the noise intensity level, and Z, ~ (O 4n). Then we discretize the SDE from (8) as follows:
0++1= (V + VTV no 6e, (x)
(58)
To approximate continuous time, we use small discretization step size "t : nn/10 and run the discretization for 10x longer than the corresponding SGD run. We use 7) == [[//10] and
SGD
:= (t/10]) ) where is a constant that select for each setting separately to match the training dynamics of the corresponding SGD run. In addition, we also evaluate discretization of gradient flow (i.c., (58) without the noise term) which helps to draw conclusions about the role of the noise term.
Experimental results. We present the discretization results in Fig. 8 for all models considered in the paper except deep networks for which computing the NTK matrix de, on each iteration the SDE discretization is too costly. In all cases, the dynamics of the SDE discretization qualitatively matches the dynamics of the corresponding SGD run. In particular, we observe similar levels of decrease in the rank of the NTK matrix and feature sparsity coefficient. We note that the match between SDE and SGD curves is not expected to be precise due to the inherent randomness of the process. Finally, observe that gradient flow discretization exhibits rank minimization or feature sparsity which suggests that the presence of the noise (either from the original SGD or its SDE discretization) plays key role in learning sparse features.
18

---

Under review as a conference paper at ICLR 2023
Diagonal linear networks
Iteration
Iteration
Iteration
Iteration
Geadiar flx dsciviiation
9P% britione
from
Two-layer ReLU networks ID regression
Sind X
1000 2000
2000
2008
2000
Iterrtioo
Two-layer ReLU networks a teacher-student setup
Aigurds aine
To0 loos
Iterations (rescaled)
Iterations (rescaled)
Iorrtions (rescaled)
Iterations (rescaled)
Goadeet Bow diccstiantion 0.00]
SGD decay 50% teubooo
S0E dicncication with and from the 5CD 516)
Three-layer ReLU networks in teacher-student setup
Reratio
Tteration
Tosaaiion
Iteration
Sric Bety Derst
deorttarion
Figure 8: Empirical validation of the SDE modeling. In all cases, the dynamics of the SDE discretization qualitatively matches the dynamics of the corresponding SGD run. Moreover, gradient flow discretization exhibits no rank minimization feature sparsity which suggests that the presence of the noise plays key role learning sparse features.
19

---

Under review as a conference paper at ICLR 2023
D ADDITIONAL EXPERIMENTAL RESULTS
This section of the appendix presents additional experiments complementing the ones presented in the main text.
SGD 6-0.15
SGD, n-0.66. decoy 50% iterationn
Illustration of neuron dynamics. We illus- trate the change of neurons during training of two-layer ReLU networks in the teacher- student setup of Chizat al. (2019) (see Fig. therein) using a large initialization scale for which small step sizes of GD or SGD lead lazy training. We postpone the illustration of (01)-(3) to Fig. 13 Appendix as our inter- est is on showing neuron dynamics (Fig. 9). We see that for SGD with a small step size, the neu- rons Wy stay close to their initialization, while for a large step size, there is clear clustering of directions wi along the teacher directions w.. The overall picture is very similar to Fig. of Chizat et al. (2019) where the same feature learning effect achieved via gradient flow from small initialization which is, however, much more computationally expensive due to the saddle point at zero. Finally, we note that the clustering phenomenon of neurons w; motivates the removal of highly correlated activations in the feature sparsity coefficient: although the corresponding activations are often non-zero, many of them in fact implement the same feature and thus should counted only once.
randorts intalized
randomly intialiadd
Figure 9: Only for large step size, the neurons cluster along the teacher neurons w leading model that uses sparse set of features.
Further results. We give short overview of additional figures referred to in the main text. More details can found the captions.
Figure 10 shows that even if loss stabilization occurs in diagonal linear networks, the implicit bias towards sparsity is largely weaker than that of SGD and generalization is poor.
Figures 11 and 12 demonstrate that the implicit bias resulting from high-loss stabilization makes the neural nets learn first a simple model then eventually fits the data.
Figure presents the sparsifying effect corresponding to the neurons' movements exhibited in Figure 9.
Figure 14 showcases the features learning induced by large step sizes for different layers of ResNets-18 18 when trained on CIFAR-10.
Figure 15 exhibits the feature sparsity in ResNets architecture CIFAR-100 without any regularization (plain SGD) and the state-of-the-art setup.
Itoration
Iteration
Iteration
GD 2.00
60 355 eady 1010
GO 233 G4d0y 24 Ter01
side 21 50h, Parzioe
5GD 027 dew tansioo
Figure 10: Diagonal linear networks. Loss stabilization also occurs for full- gradient descent but does not lead to similar level of sparsity as SGD and also does not improve the test loss.
20

---

Under review as a conference paper at ICLR 2023
Mvil
Tral
Iteration
Iteration
Iterstin
SG2, v-60000
Figure 11: Two-layer ReLU networks for ID regression. Unlike for Fig. here we larger warmup coefficient (500x vs. x00x) which leads to overregularization such that the -cheucce run fails to fit all the training points and gets stuck at too high value of the training loss (s= 10 0.5.
SGD, doeay nouaices
SGD decay 50% ixration
Figure 12: Two-layer ReLU networks for 1D regression. Illustration the resulting models from Fig. 4 over training iterations. We can see that first the model simplified and only then it fits the training data.
Festure sparsity coof
12.00
deryy
Figure 13: Two-layer ReLU networks in teacher-student setup. Loss stabilization for two-layer ReLU nets the teacher-studee setup with input dimension 2. We observe loss stabilization, better test loss for longer schedules and sparser features due simplification of (xX).
Early layer Small 7)
Middle layer Small v
Last layer Small n
Initial
Large
Initial
Large n
Initial
Large rj
Figure 14: Visualization on four sets of convolutional filters taken from different layers of ResNets- 18 trained on CIFAR-10 10 with small vs. large step size (the 50% decay schedule). For small step sizes, the early and middle layers stay very close randomly initialized ones which indicates the absence of feature learning.
21

---

Under review as a conference paper at ICLR 2023
ResNet-34 on CIFAR- 1000 no explicit regularization
Freher reet sere
Epoch
Epoch
Epoch
Epoch
8.02
0-7.5 warmup, decay at 30%%
0.75, wammpp decay ot 30%
4- 075, warmus, decay at 50%
ResNet-34 on CIFAR-100, state-of-the-art setting
Traiiing losx
Apraateeaeyyy Jny
Epoch
Epoch
o- decay
docay
decay
Figure 15: ResNet-34 trained on CIFAR-100. Both without explicit regularization and in the setting, the training loss stabilizes, the test loss significantly depends on the length of the schedule, and feature sparsity minimized over iterations. However, differently from the plots on CIFAR-10, here without explicit regularization we observe oscillating behavior after the step size decay (although at very low level between 10 4 and 10 10-2).
22


Reviews:
[{"review_id": "kp5CN9cgNp", "reviewer": "Reviewer_FkNE", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper empirically shows that SGD with large step size learns sparse features and presented a few minimal theoretical examples that corroborate the empirical observation\n", "summary_of_the_review": "The paper has great promise but is hindered by the disclarity, missing of discussion of important previous works, overclaim, and potential incorrectness of the theory\n\nI cannot recommend the paper for acceptance in its current state. I might reconsider my position if the authors make extensive updates to improve the manuscript; otherwise, I suggest resubmitting to a different conference", "strength_and_weaknesses": "Strength: some pieces of insights are novel. For example, the authors suggest the mechanism of loss stabilization and suggest its connection to sparsity.\n\nthe point that GD and SGD have different inductive biases is also an important insight\n\nThe following are the weaknesses that the authors need to address.\n1. the title is too strong: the title implies that using a large learning rate always leads to sparsity, but this is not what the paper has shown. The paper only showed that *empirically, a large learning rate leads to sparsity in certain restricted scenarios*. The title is a overclaim that needs to be changed\n\n2. the first contribution has already been studied extensively in quite a few previous works. the authors need to discuss these works and modify the claim accordingly: (1) https://openreview.net/forum?id=uorVGbWV5sw; (2) https://proceedings.mlr.press/v162/mori22a.html\n\n3. contribution 3: it is not exploration that leads to sparse features; it is the multiplicative noise. After all, \"exploration\" here is a very ambiguous word here and needs revision\n\n4. page3: \"we showcase our results for the mean square error, but other losses like the cross-entropy carry the same properties (W2021b).\" I don't see which theorem in W2021b shows that cross-entropy carries the same properties. The authors should be more specific about which theorem in W2021b and which \"property\" they are referring to\n\n5. prop 1: the notation $\\mathbf{1}_{i=i_t}$ is undefined\n\n6. prop 1 feels incorrect (it is, at least, unclear). The theorem statement seems to suggest that the SGD on $L(y)$ is equivalent to GD on $L(y + \\xi)$, but this is not the case: the second equation of the proof does not compute the gradient with respect to $\\xi_i$, which is also a function of $\\theta$. The proof thus feels incorrect\n- Thus, the message that SGD = GD + label noise is correct is also a misstatement.\n- The authors needs to either correct the proof or update the proposition statement\n\n7. the word \"stabilize\" is used many times but I feel is inappropriate. Look at figure 1: at a large learning rate, the loss does not stabilize to a fixed value but fluctuates drastically around a mean. I think the word \"stabilize\" is confusing and should be changed\n\n8. the fonts in Figures 3, 4, 6, 7 are not visible to human eyes. The authors need to use a font size similar to the font size of the main text in the figures\n\n9. the word \"feature sparsity\" is undefined throughout in the paper even though it plays an important role. this leads to a lot of potential confusion. For example, see the third figure in figure 7; the figure tells me that a \"larger learning rate leads to smaller feature sparsity,\" but I guess this is not what the authors are trying to convey\n\n10. I cannot agree with a main (empirical) claim of the paper: \"longer stabilization leads to sparser features.\" The empirical results are only sufficient to show that longer stabilization is positively correlated with sparser features (because they happen simultaneously), but I do not think the empirical results allow one to conclude that one leads to the other. I stress the following point: *correlation cannot imply causality*\n\n \n\n\nMinor:\n1. perhaps the authors should point out that eq 2 is essentially SGD with batch size 1 \n2. the authors should give equation numbers to all the equations in the appendix so they are easier to refer to during the review", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "d6D_yTAIyWJ", "reviewer": "Reviewer_USkr", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "The authors claim that in SGD-trained neural networks, during the phases where the loss is constant, the stochastic noise in the optimization serves to allow the model to learn sparse features, which endows it with a benevolent inductive bias in terms of generalization. The authors demonstrate their ideas in experiments with small to medium sized neural networks.", "summary_of_the_review": "I think the authors' work provides some interesting findings that would be helpful for a more nuanced understanding of the generalization behavior of deep neural networks. At points, more justification for their modeling choices and/or their inferences seem to be needed, as well as more concrete comparions with the existing body of work.", "strength_and_weaknesses": "The authors' argumentation seems to include the following steps:\n\n1- SGD can be formulated as GD with a specific random noise (Prop 1.).\n2- Variance of the norm of the specific random noise they introduced is proportional to the batch loss (Lemma 2.).\n3- It is possible for an SGD trained algorithm to maintain an equilibrium where the loss is lower and upper bounded all the while the parameters take different values in the basin of the minimum but at some distance away from it (Prop 3.).\n4- In the phase of the learning where the loss is constant, SGD can be modelled by a SDE whose noise term is characterized by the stochastic gradients (Eq 7).\n5- This SDE implies that when the loss is constant, the optimization procedure will be inclined to reduce the l_2 norm of the stochastic gradients, leading to a sparse representation.\n\nI find the authors' work valuable in investigating how the training dynamics and sparse representations evolve during training of neural networks. How efficient representations are achieved throughout the training procedure is a topic that is promising for understanding how neural networks generalize. The authors have a more or less self-consistent argumentation (albeit with the caveats below) and provide interesting experimental results to complement their claims.\n\nAlthough I believe that the conference audience would benefit from their work, I believe answering the questions below would improve my and possibly other conference audience's understanding of their work:\n\n- The authors' comparison of their approach to the rest of the literature is cursory at best. Although they mention a lot of other works, they usually shy away from making comparisons that illuminate why their results differed from those previously published. I would appreciate even less citations in exchange for meaningful comparison. For example, (Pg. 2 Par 5) What makes authors experimental regime different than that of Lewkowycz et al. 2020.? What does flatness definitions being \"questionable\" mean (Pg. 2)? Why are their findings regarding batch GD is different than those of Nacson et al 2022? (Pg 7.) In which aspects are authors' demonstrations are \"closer to practice\" (Pg. 1) and compared to whom?\n- Pg. 4 Par. 4: How would the authors argue that Proposition 3 supports the importance of large learning rates per se in leading to the loss stabiliziation? Moreover, how does Prop. 1 serves the authors' exposition exactly? What in framing SGD as a GD + noise enable Claim 4 to be made, that would not be possible if it was not framed as GD?\n- What is the justification for the choice of Brownian motion in the authors' proposed SDE? Large learning rates are assumed to produce heavy-tailed behavior in neural network parameters, would these results change the interpretation of the authors' results?\n- Throughout the introduction and exposition the authors refer to label noise being a key in understanding the generalization behavior of SGD. I think this might be confusing for some audience since the type of label noise the authors are proposing is simply a theoretical device for their derivation, as opposed to researchers working with uncertainty in supervision signal.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "OaNIASlquq", "reviewer": "Reviewer_gDbe", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper studies the implicit bias of SGD with large step sizes. The authors present empirical observations and some theoretical understanding. They show that large step sizes SGD enables exploration by loss stabilization, and this exploration leads SGD to learn sparse features. The authors also provide some insights into the deep neural networks' training dynamics and understanding of some existing tricks.", "summary_of_the_review": "This paper contains very interesting observations and analyses. However, I feel that the results are overclaimed and analysis support is not enough. Consequently,  I am a bit skeptical about the applicability of the findings.\n", "strength_and_weaknesses": "The paper is generally well-written. Understanding the implicit bias of SGD with large step sizes is a significant theoretical and practical issue. This paper provides a rich discussion and offers a sparse feature learning perspective.  However, I feel that neither the analysis is well-designed nor the claim is well-supported. Specifically, I have two primary concerns from the main text of this work: the persuasiveness of theories in Section 2 and the sparse feature bias of deep networks in Section 3.\nThese two main points also limit the significance or rather limit my ability to determine this work's significance. \n\n\n- The persuasiveness of theories in Section 2.\n    - **The reasonableness of the SDE in (Eqn. 7).**  SDE is good modeling of SGD only in small step size regime, but it is generally unclear how this SDE modeling is relevant for understanding SGD with large step sizes.  As the authors mentioned in Section 2.3.1, this SDE may be reasonable for the diagonal linear model. However, for nonlinear networks such as ReLU nets, the reasonableness of the SDE modeling (eqn 7) seems questionable (especially the noise covariance). This SDE is used to motivate the bias of learning sparse features but is never used in experiments, where vanilla SGD is applied. It is thus unclear why do not use vanilla SGD to motivate the sparse feature learning.   The authors can address this issue using experiments. \n    \n    - **The hidden stochastic dynamics orthogonal to the bouncing directions.** In the abstract, the authors mention ''this loss stabilization induces a hidden stochastic dynamics orthogonal to the bouncing directions''. However, it seems that this paper does not point out the bouncing direction as well as explain the reason for orthogonality. The 1-D example (Eqn. 4) gives us insights into loss stabilization, but it cannot help us understand the authors' claim about the bouncing direction and the orthogonal hidden stochastic dynamics. The authors can construct a 3-D toy example to illustrate the loss stabilization, the bouncing direction, and the orthogonal effective dynamics.\n    \n    - **The claim that large step sizes learn sparse feature solutions.**\n    I do not follow the explanation in Section 2.3.2 of why the eqn (7) implies that SGD implicitly minimizes the $\\ell_2$ norm of each column of the feature matrix. This claim is so big and concrete but I do not see too much support for this claim. In addition, all of the authors' analysis holds for any step size in the interval of step sizes that satisfy loss stabilization. It means that SGD with different step sizes in this interval has the same sparse-feature-learning bias. Hence, it seems that the authors' viewpoint can not distinguish the different implicit biases of two sizes in this interval. If this is the case I do not view it as grounds for rejection but it is important for a reader to know. \n \n    \n- The sparse feature bias of deep networks in Section 3.\n    \n    - **The clarification of the feature sparsity coefficient.** What is the feature sparsity coefficient for deep neural networks mentioned on Page 6? As an essential metric, there is neither a clear explanation nor a mathematical formulation of this in the main text, which is not very clear.\n    \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "I-mTwgAYezV", "reviewer": "Reviewer_Jpgc", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper provides some insights on the potential reason that the initial large step sizes can learn the model that can generalize well in practice. Center to the claim is an SDE with specific noise covariance structure, and the authors argue that, if we use the large step sizes at the beginning, the loss should stabilize at a certain level without further decreasing. During this phase, the unnecessary feature will be optimized towards 0, which will eventually lead to a solution with sparse features.\n", "summary_of_the_review": "The authors provide a quite good intuition to interpret the effectiveness of large initial step sizes. Although the example is a little bit toy, I believe this paper will make impacts in the related areas. I feel the authors should be happy to include some discussions for [1]. Although it appears after the ICLR submission deadline, I feel this paper provides some theoretical analysis that potentially demonstrates the intuition discussed in this paper.\n\n[1] Mousavi-Hosseini, Alireza, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A. Erdogdu. \"Neural Networks Efficiently Learn Low-Dimensional Representations with SGD.\" arXiv preprint arXiv:2209.14863 (2022).", "strength_and_weaknesses": "### Strength:\n* The idea is novel and interesting.\n* The authors provide some synthetic examples to interpret the intuition.\n* The authors conduct lots of experiments to demonstrate the conjecture.\n\n### Weakness:\n* It’s a little bit hard to follow as the authors don’t provide clear logic in the presentation of Section 2.\n* The examples are a little bit toy.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject