Stochastic Gradient Descent (SGD) with Nesterov's momentum is a widely used optimizer in deep learning, which is observed to have excellent generalization performance. However, due to the large stochasticity, SGD with Nesterov's momentum is not robust, i.e., its performance may deviate significantly from the expectation. In this work, we propose Amortized Nesterov's Momentum, a special variant of Nesterov's momentum which has more robust iterates, faster convergence in the early stage and higher efficiency. Our experimental results show that this new momentum achieves similar (sometimes better) generalization performance with little-to-no tuning. In the convex case, we provide optimal convergence rates for our new methods and discuss how the theorems explain the empirical results. In recent years, Gradient Descent (GD) (Cauchy, 1847) and its variants have been widely used to solve large scale machine learning problems. Among them, Stochastic Gradient Descent (SGD) (Robbins & Monro, 1951) , which replaces gradient with an unbiased stochastic gradient estimator, is a popular choice of optimizer especially for neural network training which requires lower precision. Sutskever et al. (2013) found that using SGD with Nesterov's momentum (Nesterov, 1983; 2013b) , which was originally designed to accelerate deterministic convex optimization, achieves substantial speedups for training neural networks. This finding essentially turns SGD with Nesterov's momentum into the benchmarking method of neural network design, especially for classification tasks (He et al., 2016b; a; Zagoruyko & Komodakis, 2016; Huang et al., 2017) . It is observed that in these tasks, the momentum technique plays a key role in achieving good generalization performance. Adaptive methods (Duchi et al., 2011; Kingma & Ba, 2015; Tieleman & Hinton, 2012; Reddi et al., 2018) , which are also becoming increasingly popular in the deep learning community, diagonally scale the gradient to speed up training. However, Wilson et al. (2017) show that these methods always generalize poorly compared with SGD with momentum (both classical momentum (Polyak, 1964 ) and Nesterov's momentum). In this work, we introduce Amortized Nesterov's Momentum, which is a special variant of Nesterov's momentum. From users' perspective, the new momentum has only one additional integer hyper-parameter m to choose, which we call the amortization length. Learning rate and momentum parameter of this variant are strictly aligned with Nesterov's momentum and by choosing m = 1, it recovers Nesterov's momentum. This paper conducts an extensive study based on both empirical evaluation and convex analysis to identify the benefits of the new variant (or from users' angle, to set m apart from 1). We list the advantages of Amortized Nesterov's Momentum as follows: • Increasing m improves robustness 1 . This is an interesting property since the new momentum not only provides acceleration, but also enhances the robustness. We provide an understanding of this property by analyzing the relation between convergence rate and m in the convex setting. • Increasing m reduces (amortized) iteration complexity. • A suitably chosen m boosts the convergence rate in the early stage of training and produces comparable final generalization performance. • It is easy to tune m. The performances of the methods are stable for a wide range of m and we prove that the methods converge for any valid choice of m in the convex setting. • If m is not too large, the methods obtain the optimal convergence rate in general convex setting, just like Nesterov's method. The new variant does have some minor drawbacks: it requires one more memory buffer, which is acceptable in most cases, and it shows some undesired behaviors when working with learning rate schedulers, which can be addressed by a small modification. Considering these pros and cons, we believe that the proposed variant can benefit many large-scale deep learning tasks. Our high level idea is simple: the stochastic Nesterov's momentum can be unreliable since it is provided only by the previous stochastic iterate. The iterate potentially has large variance, which may lead to a false momentum that perturbs the training process. We thus propose to use the stochastic Nesterov's momentum based on several past iterates, which provides robust acceleration. In other words, instead of immediately using an iterate to provide momentum, we put the iterate into an "amortization plan" and use it later. We presented Amortized Nesterov's Momentum, which is a special variant of Nesterov's momentum that utilizes several past iterates to provide the momentum. Based on this idea, we designed two different realizations, namely, AM1-SGD and AM2-SGD. Both of them are simple to implement with little-to-no additional tuning overhead over M-SGD. Our empirical results demonstrate that switching to AM1-SGD and AM2-SGD produces faster early convergence and comparable final generalization performance. AM1-SGD is lightweight and has more robust iterates than M-SGD, and thus can serve as a favorable alternative to M-SGD in large-scale deep learning tasks. AM2-SGD could be favorable for more restrictive tasks (e.g., asynchronous training) due to its extensibility and good performance. Both the methods are proved optimal in the convex case, just like M-SGD. Based on the intuition from convex analysis, the proposed methods are trading acceleration for variance control, which provides hints for the hyper-parameter tuning. We discuss the issues with learning rate schedulers in Appendix A.4. We report the test accuracy results of the ResNet18 experiment (in Section 4) in Appendix A.5. A CIFAR-100 experiment is provided in Appendix A.6. We also provide a sanity check for our implementation in Appendix A.7. Table 4 : Final test accuracy and average accuracy STD of training ResNet34 on CIFAR-10 over 5 runs (including the detailed data of the curves in Figure 1 and Figure 2a) . For all the methods, η 0 = 0.1, β = 0.9. Multiple runs start with the same x 0 . We show in Figure 6 how m affects the convergence of test accuracy. The results show that increasing m speeds up the convergence in the early stage. While for AM1-SGD the convergences of Option I and Option II are similar, AM2-SGD with Option II is consistently better than with Option I in this experiment. It seems that AM2-SGD with Option I does not benefit from increasing m and the algorithm is not robust. Thus, we do not recommend using Option I for AM2-SGD. Table 4 . Labels are formatted as 'AM1/2-SGD-{Option}-{m}' . Best viewed in color.
Amortizing Nesterov's momentum for more robust, lightweight and fast deep learning training.
{"src": "scitldr-AIC_train", "paper_id": "S1xJFREKvB"}