Paper: Flatter, Faster: Scaling Momentum for Optimal Speedup of SGD (ICLR-2024)

Reviewer comment (Writing):
It is very hard to read Section 3 for non-experts in the stochastic process theory. I suggest the authors compress it and extend the section with experiment evaluation.

Author response:
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.

Reviewer comment (Evaluation):
Figure 2 presents fitting line results that confirm the estimated dependence rule, however, I do not find the analysis of the variance of the derived estimate. I am not sure that if one adds more points to the plots, then the dependence is changed significantly.

Author response:
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.

Reviewer comment (Theory):
The authors consider simple models from computer vision tasks like ResNet18 and MLP. Could you please list the assumptions on the deep neural network that are necessary for the correctness of the derived relation between learning rate and momentum hyperparameter? These assumptions will be very helpful in the extension of the presented results to other models from different domains, for example, transformer-based LLMs.

Author response:
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.

Reviewer comment (Theory):
Is it possible to extend the proposed approach from Heavy Ball to a Nesterov-like acceleration scheme? If so, please comment what are the potential obstacles to deriving similar results.

Author response:
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.

Reviewer comment (Theory):
The question related to weakness 2: how robust is the derived estimate w.r.t. the new points that may appear in the plots?

Author response:
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.

Reviewer comment (Theory):
How the derived relations between the learning rate and momentum term can be interpreted from the perspective of loss surface properties? What the derived 2/3 power rule can highlight in the loss surface corresponding to the considered models?

Author response:
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.

Reviewer comment (Writing):
The biggest concern I have is regarding the presentation of this paper. The quantities $\tau_1$ and $\tau_2$ are the focus of this paper. But it is only defined through informal descriptions (for example " Define $\tau_2$ to be the number of time steps it takes so that the displacement of $w_k$ along $\Gamma$ becomes a finite number as we take $\epsilon$ → 0 first, and $\eta$ → 0 afterward") in the introduction, and no (even informal) definition elsewhere. This makes these two terms extremely hard to interpret and understand.

Author response:
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.

Reviewer comment (Theory):
This paper only considers optimization around the manifold of the global minima. Although this is inherited from (Li et al., 2022), I wonder whether this framework can characterize the convergence rate along the whole trajectory, which is the one of more interest. For example, in Figure 2, the initialization is picked where a perfect interpolation has been already achieved. What happens if the initialization is chosen as, for instance, Kaiming's initialization?

Author response:
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.

Reviewer comment (Experiments):
The experiments are too few and toy considering this paper aims to provide modification for algorithms: there is only an experiment of 2-layer MLP and an experiment over CIFAR 10 with a very uncommon initialization (as discussed above).

Author response:
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.

Reviewer comment (Novelty):
The work seems a relatively straightforward extension of Li et al. (2022)

Author response:
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.

Reviewer comment (Theory):
It is not clear if the optimality of the gamma=2/3 exponent applies to standard SGD as well.

Author response:
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).

Reviewer comment (Theory):
The theory does not give a prescription for the prefactor C. I think this makes it not so useful in practice.

Author response:
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.

Reviewer comment (Experiments):
What happens when using standard SGD instead of the label noise one? Can the author provide numerical evidence that gamma=2/3 is still a good exponent? My understanding is that all numerical experiments are carried out with SGD label noise.

Author response:
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).

Reviewer comment (Theory):
Could the author clarify the argument for the tau1 = tau2 criterium. In particular, it is not clear to me why we should care about traversal equilibration. What seems relevant is to move fast toward the flat region staying close enough to the zero loss manifold.

Author response:
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).

Reviewer comment (Experiments):
What happens if also phase 1 is trained with SGD instead of GD? The starting point for the following SGD label noise diffusion would be already in a flat region there wouldn't be much drift?

Author response:
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.

Reviewer comment (Presentation):
Is Fig 1 obtained with a single value of \eta and varying \gamma? Could the author show multiple sets of points corresponding to different eta values?

Author response:
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$.