SGD with memory: fundamental properties and stochastic acceleration

Thank you for your careful reading and positive evaluation of our work!

The paper seemed quite convoluted...

That's true; we have rewritten one part that was especially convoluted, please see our general response.

The empirical validation of the results remains quite limited ...

We agree that in our empirical validations we covered a limited set of scenarios. We have validated that the heuristically derived AM1 algorithm actually gives the predicted convergence rate while being stable in numerical simulations, and in the MNIST experiments, we showed that AM1 still works in non-convex setting without exact power-law spectrum and on real-world data. See also a detailed summary of our experiments in the response to reviewer LKb1.

A comprehensive test of AM1 in a variety of modern deep-learning scenarios would be an interesting direction for future research. Yet, we are convinced that the provided set of experiments fits well the current manuscript. We would like to keep the focus of this submission on theoretical analysis and derivation of the algorithm, limiting experiments to validation of main theoretical claims and providing proof-of-concept evidence that the algorithm can work beyond the purely quadratic setting considered in the paper.

It is seems like the optimal rate depends on the right choice of hyperparameters, such as $\alpha_0$, $q_0$ and $\delta$. However it is unclear to me how to select or estimate them in practice. This limits its application to real-world problems.

Thank you for this remark! We tried to address this question in the paper, although not very directly.

• Although theorems 6 and 7 provide specific choice of $\delta, q_0, \alpha_\mathrm{eff}$ hyperparameters, it is natural to assume that the precise values will change if we deviate from the setting of quadratic models with exact power-law spectrum. Yet, we believe that the region of hyperparamters $\delta^{-1}\gg \alpha_\mathrm{eff}\gg q_0=const$, identified in theorem 7, will be more robust to the change of the setting (non-power law spectrum, non-quadratic loss, etc). See the remark in the penultimate paragraph on page 9.

• Another possible difficulty for practical application is formulation of the algorithm in terms of variables $\delta, q_0, \alpha_\mathrm{eff}$. These variables appear to be the most convenient parametrization for our theoretical analysis, but not necessarily for the practical implementation of the algorithm. We address this in the beginning of section M, where we rewrite the AM1 update rule as a modification of HB (a.k.a SGD in most deep learning libraries). Specifically, we can just add can an extra ``kick'' term - updating parameters directly with the gradients of the current iteration. Denoting the learning rate of the new kick term as $\alpha_1$, and the learning rate of the classical momentum term as $\alpha_2$, the AM1 schedule becomes $(\delta_t, \alpha_{1,t}, \alpha_{2,t})\sim (t^{-\overline{\delta}}, const, t^{-\overline{\delta} + \overline{\alpha}})$.

• Hyperparameter choice. One can parameterized the whole schedule as $(\delta_t, \alpha_{1,t}, \alpha_{2,t}) = (c_\delta t^{-\overline{\delta}}, c_{\alpha_1}, c_{\alpha_2}t^{-\overline{\delta} + \overline{\alpha}})$ leading to 5 tunable hyperparameters. A quick starting point could be (i) setting reasonable values for constants $c_\delta,c_{\alpha_1},c_{\alpha_2}$ (ii) picking $\overline{\delta}$ close to its theoretically predicted value $\overline{\delta}=1$ to avoid some edge effects, e.g. $\overline{\delta}=0.9$ (iii) doing a sweep on value of $\overline{\alpha}$, since its theoretical optimal prediction $\overline{\alpha} = \overline{\delta}(1-\frac{1}{\nu})$ requires knowledge of the spectral exponent $\nu$ which is typically unknown for real-world problems. We roughly followed this strategy in our experiments, as shown on figures 4,5.

• Finally, we note that there is freedom in choosing schedule exponents $\overline{\delta},\overline{\alpha}$, removing the need to precisely estimate spectrum exponent $\nu$. From our results in sec. 6, the algorithm would be stable as soon as $\overline{\alpha}\leq \overline{\delta}(1-\frac{1}{\nu})$, and acceleration is determined by the value of $\overline{\alpha}$. Thus, a less aggressive choice than theoretically optimal $\overline{\delta}=1, \overline{\alpha}=1-\frac{1}{\nu}$ would still provide accelerated convergence while being stable.

Thank you for your feedback!

I don't fully understand about the experimental part. Is it just not there?

While the paper primarily focuses on theoretical analysis, we do provide experiments to validate the derived results, as presented in figures 1, 4, 5. All the settings are described in the dedicated section M of the appendix. In short, we believe that our current set of experiments covers the following points:

• Figure 1 demonstrates the divergence of Jacobi accelerated HB that was mentioned in the introduction. Importantly, the divergence of Jacobi accelerated HB happens at any batch size, with larger batch size values just delaying the moment of divergence. This serves as a motivation to look into accelerated strategies, leading us to the formulation and analysis of Memory-M algorithms.

• Same figure 1 shows that the proposed AM1 achieves an accelerated power-low convergence rate, although not as high as that of Jacobi accelerated HB in the noiseless setting. Moreover, the AM1 algorithm is stable even at quite small batch sizes ($|B|=1$ and $|B|=10$ in the figure).

• Figure 1 also provides the noisy and noiseless (full-batch) loss curves in otherwise identical settings. This visually illustrates the effect of noise on the loss trajectory in both stable and diverging cases.

• In terms of scope, we consider either synthetic Gaussian data with the manually set power-law spectrum and MNIST classification with a shallow ReLU network. We describe the details of each setting in section M. These two settings serve different goals.

• The Gaussian setting is directly described by the developed theory and was expected to follow our theoretical predictions. Yet, the last step of designing a schedule for AM1 algorithm and obtaining its convergence rate is heuristic and non-rigorous, in contrast to the rest of the paper. Thus it was important for us to validate the intuition used in the heuristic derivation of AM1. Indeed, the predicted convergence rate, depicted by dashed lines, matches the experimental curves.

• The MNIST setting, although basic by modern deep learning standards, is non-convex and thus deviates from the quadratic optimization problem analyzed in the paper. Yet, we see that the main properties of AM1, stability and accelerated convergence, survive deviation from purely quadratic problems. This gives us some hope that AM1 algorithm could have broad practical applications in real-world optimization problems.

• Finally, in figures 4 and 5 we provide more experiments by changing various parameters of both Gaussian and MNIST settings. These experiments show that the conclusions from figure 1 are robust.

Can authors convince me that their work fits the spirit of this conference? This question is not obvious to me at this point...

We are trying to solve a problem -acceleration of SGD - that is relatively hard and does not seem to have a simple intuitive solution. To this end, we develop and present an original theoretical approach that includes many new elements (general memory-$M$ algorithms, a propagator expansion of the loss, generalized effective learning rate, stable memory-1 algorithms with unbounded effective learning rate ...) and eventually reveals a solution. We believe our work to advance the ML science and thus fit the spirit of ICLR quite well.

Many of the works directly connected to our research and cited in our list of refernces have also been published either at ICLR or similar conferences like NeurIPS, ICML, COLT. This further confirms that our work fits the spirit of ICLR well.

Could we add label noise or regularization?

Both label noise and regularization could be introduced in our framework by generalizing some of its parts. Basically, we need to modify second moments update (12),(13), and then add new respective terms into propagator expansion (19).

In case of additive i.i.d. label noise with variance $\sigma^2$, second moments update acquires a term (one can use derivation in sec. C.1 as an example)

The case of $L_2$ regularization is a bit more involved with more terms changed/added. Since it pulls the parameters to $0$, the optimum will shift to a new location $\mathbf{w}^*_{reg}$, modifying both the main evolution term $A_t$ and the noise term $P_t$ that make up the evolution operator $F_t$ (see section D).

However, the biggest changes would be in the loss propagator expansion. In its current version, we heavily rely on the multiplicative structure of the evolution operator, its form $F_t=A_t+P_t$, and $P_t$ being rank-1. Adding extra terms of a different type to $F_t$, or introducing additive terms into the second moment's update as in the case of label noise, would require redesigning the loss expansion. This is an interesting direction for future research.

Thanks very much for your careful reading and a very positive evaluation of our work!

One possible suggestion is that the paper develops the results initially in the $a=0$ as the shifts are not needed for the main contributions and then provide extensions, possibly in the appendix. The other is to hide some of the intermediate steps to provide space for more intuition for the different objects.

Thank you for these suggestions - we seriously thought about them, especially removing the shifts. We didn't get to implement it, because of the lack of time and because we wanted to keep some discussion of shifts in the main text, but we plan to return to this issue in a future update. However, we have rewritten the excessively dense "loss expansion" section, moving intermediate steps to the appendix as you suggested (see our general response).

One thing that threw me off is using capital letters for the propagators as these are scalars.

Thank you for pointing this out. We now mention explicitly in the text that the propagators are scalar, to minimize possible confusion. There are several reasons why we use capital letters for the propagators: 1) they denote direct contribution to the averaged loss $L_t$ which is also consistently denoted by a capital letter (e.g., in the noiseless case $L_t=\tfrac{1}{2}V_t$); 2) propagators $U_t$ are extensively treated as entries of a circulant matrix $\mathbb U$ in various derivations in the appendix; 3) propagators are important objects, and most lowercase letters are already used for other purposes in the paper.

Could such a framework be extended to non-IID batching, that is, batching with replacement (recent work [1,2])?

We expect that yes, it could be extended in this way. Reading these papers, analysis of batching with replacement is certainly technically more involved than of batching without it. But at least for large training sets and small batches, we can expect both kinds of batching to be close, suggesting similar properties of our generalized SGD.

Could you say anything about the generalization error using [3]?

The paper [3] derives a precise asymptotic characterization of the test error under power-law spectrum. Yet, there seem to be several challenges left to obtain (and then analyze) such characterization for mini-batch gradient-based algorithms, including memory-M algorithms:

• The test error in [3] is derived for kernel ridge regression (which is especially convenient for tools from random matrix theory thanks to its connection to the resolvent), but not for gradient descent. The case of gradient descent is, however, addressed to some extent in some other papers, e.g. [4,5].

• Regarding the stochasticity of inputs, in [3] and other related works the training set is randomly drawn once before applying the learning algorithm. The SGD setting adds one more layer of stochasticity due to the random choice of mini-batches from the finite-size training set. The joint averaging over these two sources of randomness might be challenging.

• If we want to characterize the generalization error under a label noise, this also needs to be added to the framework (see another question below for more details).

On a positive side, though, the SE approximation used in our framework holds exactly for Gaussian features used in [3].

[4] Yicheng Li, Weiye Gan, Zuoqiang Shi, and Qian Lin. Generalization error curves for analytic spectral algorithms under power-law decay, 2024. URL https://arxiv.org/abs/2401.01599.

[5] M. Velikanov, M. Panov, and D. Yarotsky. Generalization error of spectral algorithms. In ICLR 2024. https://openreview.net/forum?id=3SJE1WLB4M.

Thank you very much for your careful reading and a very positive evaluation of our work!

...the Loss expansion section with the introduction of propagators seemed too dense - and was lacking intuitions that would facilitate further reading of the paper..
... I wonder if there are any ways of improving the readability of the technical portions ...

Thank you for this feedback and particularly pointing out the loss expansion section as too dense - we fully agree and have rewritten this piece, moving technical details to a new appendix. We plan to return to improving the readability of the paper in the future update.