Flatter, Faster: Scaling Momentum for Optimal Speedup of SGD

We thank the referee for their positive feedback on our paper, and for the constructive questions that helped expanding and improving our presentation. We believe we addressed all the referee's questions. We detail our reply and updates to the draft below:

1. The main assumption is that the neural network is in the overparameterized regime. There has been extensive literature (as discussed in Sec. 2) supporting the statement that, in overparameterized neural networks, global minima tend to form connected manifolds, what we call ``valley''. A separate but related assumption, is that the closer the weights are to the valley, the more applicable we expect our prediction to be. Given the relationship between overparameterization and existence of zero-loss valleys have been studied across many models, we believe these assumptions are architecture-independent and we expect they should apply to transformer-based models as well, in particular in the context of fine-tuning where models are often overparameterized (see e.g. https://arxiv.org/abs/1908.05620). We also point out that, in our experiments, even when training is initialized far from the valley (e.g. due to random initialization), we still observe the scaling law we proposed. In fact, we now added in Sec. 4.3.1 and Fig. 2 (bottom right) an experiment training a ResNet-18 model on CIFAR-10 starting with random initialization, and observed the optimal scaling relationship we proposed.

2. Extending this approach to Nesterov acceleration involves evaluating the gradients at $w+\beta \pi$ instead of $w$, where $w$ are the network weights and $\pi$ is momentum. Technically, while it will be slightly more involved to derive the limit diffusion equation of Theorem 3.4, we expect it to be still doable as the latter result is derived through a Taylor expansion in $\pi$, since the Theorem is valid in the limiting case in which the configuration $(w,\pi)$ is close to the valley, which is located inside the subspace $\pi=0$. We also note that it is in principle possible that Nesterov-like acceleration schemes will lead to a different power from $2/3$. The important point is that there will still be a power law since the steps involved in the derivation will be qualitatively similar.

3. We included an error for our numerical estimates (see e.g. Fig. 2), which was derived using the $\operatorname{curve\_fit}$ function from the package $\operatorname{scipy.optimize}$. This function computes errors of the estimates of the parameters using the "linear approximation method" from Vugrin, Kay White, et al. "Confidence region estimation techniques for nonlinear regression in groundwater flow: Three case studies." Water Resources Research 43.3 (2007). In our case we assume that the true data do lie on a line with noise which is i.i.d. distributed. While this assumption may not be strictly satisfied (e.g. the noise of the data points to fit may depend on $\eta$, or there might be deviations from our scaling law due to finite strength of the noise we use in experiments) the data show a clear linear relationship and such deviations are not apparent. Responding specifically to the question about adding more points -- we do not expect the estimate to change significantly because new points will be independent from old points, and if they lie inside the extensive range of learning rate values we have explored, will be approximately identically distributed as well. This means that our error estimate of one standard deviation about $0.05$ provides a good measure for the expected deviation of the estimator of the slope.

4. As we mentioned in point 1., the scaling law is related to the presence of a zero-loss valley. We do not expect the specific power of $2/3$ to hold for a generic "rough" loss surface. An additional observation is that the power we found holds empirically for random initial conditions, i.e. even when we initialize far from any valley. We observed this across various datasets (artificial, FashionMNIST, and CIFAR10). In terms of the loss surface, this could mean that a large fraction of training is spent near a zero-loss valley.

Following the referee's suggestions, we compressed part of Secs. 3.1 and 3.2, leaving only the necessary definitions and the most important results. We extended the experiments section (Sec. 4) by including the results we obtained after the original submission.

We thank the referee for their constructive questions that helped expanding and improving our presentation, and led us to design further experiments. We believe we addressed all the referee's questions. We detail our reply and updates to the draft below:

1. We appreciate the referee's feedback on the definition of $\tau_1$ and $\tau_2$. Regarding $\tau_1$, in the introduction we replaced "mixing time" with "the inverse of the smallest nonzero eigenvalue of the damping term" which is more self-contained and should completely clarify what we mean. Regarding $\tau_2$, we added appendix B to provide a formal definition of this quantity and included two more sentences at the end of the first paragraph of Sec. 1.2 to make the reasoning around tau2 more explicit. We hope these elements are now clarified.

2. We thank the referee for this question, as it prompted us to further investigate empirically what happens with random initialization. We already implemented random initialization for part of our experiments (Matrix sensing and MLP on artificial dataset), showing a consistent behavior with the scaling law we predicted without the need to initialize near the manifold of the global minima. To answer the referee's question, we performed a new set of experiments in which we trained ResNet-18 on CIFAR10 with Kaiming's initialization. We found clear evidence of the optimality of the 2/3 scaling law in this setting, as illustrated in Fig. 2 (bottom right) and Sec. 4.3.1. From a theoretical standpoint, the properties of the loss surface are known to be very challenging to characterize in generic deep-learning models, in particular away from the zero-loss valley, which on one hand makes it hard to rigorously predict this last empirical finding, and on the other hand we believe that, because of this reason, the empirical evidence we found is even more valuable.

3. As mentioned in point 2., stimulated by the referee's question, we took the opportunity to look at the standard setting of training with Kaiming initialization and SGD noise over CIFAR 10, finding strong empirical evidence that our prediction holds in realistic scenarios. We respectfully remind the referee that we also performed experiments on FashionMNIST using a deep MLP, which also provides positive evidence of our prediction. The aim of this paper is to develop a rigorous theory to ground our predictions and to initiate and display its verification across progressively complex models.

We thank the referee for their positive feedback on our paper, and for the constructive questions that helped expanding and improving our presentation. We believe we addressed all the referee's questions. We detail our reply and updates to the draft below:

1. The referee is correct that in our analysis so far we focused on SGD label noise. We carried out a new set of simulations training ResNet-18 on CIFAR10 using SGD instead of label noise. We find a clear signal of the gamma=2/3 scaling relationship. Strictly speaking, our mathematical framework currently assumes that the noise is constant, which is not quite the case for SGD noise and is the reason why we did not consider SGD noise in our experiments in the first place. Yet, this result shows that in practice the differences with SGD noise might not be minor in terms of the power law. It is gratifying to see empirically that the optimal scaling law still applies, and we are grateful to the referee for bringing up this question. We included the new results in Sec. 4.3.1 and Fig. 2 (bottom right).

2. The number of timesteps it takes for the model to converge to a low-curvature point of the zero-loss valley is controlled by the slowest between $\tau_1$ and $\tau_2$. By ``controlled'' we mean that the longer is $\max(\tau_1, \tau_2)$, the longer it takes to converge to such low-curvature point. More specifically, optimality is reached when $\tau_1 = \tau_2$ because, as illustrated in Fig. 1(a), if $\tau_2$ becomes faster, $\tau_1$ becomes slower and viceversa. The transverse timescale $\tau_1$ is important because if this is too slow, even though longitudinally convergence has happened, there might still be residual transverse fluctuations (similar to the top left of Fig. 2).

3. Regarding our previous experiment on CIFAR10 (where we explicitly separated phase 1 and phase 2 in the training schedule), label noise would have continued to cause a drift unless SGD would have caused the weights to converge to the widest minimum by the end of phase 1. Please let us know if this answers your question, we are happy to discuss more.

4. Fig 1 is obtained using various values of $\eta$ for a fixed value of $\gamma$. As explained in the last paragraph before App. A.2, the way we extracted the exponent $\alpha$ was to fit the number of timesteps to convergence $T_c$ following the relationship $T_c = T_0 \eta^{\alpha}$, where we performed several runs with different values of $\eta$ and the associated $T_c$, and extracted $T_0$ and the exponent $\alpha$ through a fit. Because of this, the plots of Fig. 1 are defined only when running across several values of $\eta$, it is not possible to plot them for a single value of $\eta$.

Additionally, we have two remarks about the weaknesses mentioned by the referee:

• Prefactor: We appreciate the referee's critique that the theory does not give a prescription for the prefactor $C$. In principle, this could mean that we have not reduced the search space of hyperparameters, but we notice two things. First, given the optimal power is fixed to 2/3 by the theory, to know $C$ one only needs to perform a few runs to find the optimal momentum hyperparameter for a given fixed value of the learning rate. In particular, from this information it is possible to extract the best momentum hyperparameter for other values of the learning rate thanks to the knowledge of the exponent. This effectively reduces the hyperparameter search to one parameter instead of two. Secondly, and intriguingly, we noted that in all our experiments, $C\approx 0.1-0.5$ independently of the models and tasks we considered.

• Simple Extension: We would also like to note that, while (Li et al. 2022) has been a main result inspiring our research, our framework has substantial differences. In order to treat the momentum hyperparameter $\beta$ as an additional small parameter (i.e. in addition to the learning rate $\eta$), we had to develop a completely different scheme of limits to derive our final equation. In that derivation, we relied on the small-noise limit rather than small-learning rate.