Layer-wise Adaptive Step-Sizes for Stochastic First-Order Methods for Deep Learning
This addresses the need for automated hyperparameter tuning in deep learning optimization, offering an incremental improvement over existing first-order methods.
The paper tackles the problem of tuning learning rates in deep learning by proposing a per-layer adaptive step-size method that uses stochastic curvature information from the Hessian, resulting in improved performance over fine-tuned baselines on models like Autoencoder, CNN, and GCN.
We propose a new per-layer adaptive step-size procedure for stochastic first-order optimization methods for minimizing empirical loss functions in deep learning, eliminating the need for the user to tune the learning rate (LR). The proposed approach exploits the layer-wise stochastic curvature information contained in the diagonal blocks of the Hessian in deep neural networks (DNNs) to compute adaptive step-sizes (i.e., LRs) for each layer. The method has memory requirements that are comparable to those of first-order methods, while its per-iteration time complexity is only increased by an amount that is roughly equivalent to an additional gradient computation. Numerical experiments show that SGD with momentum and AdamW combined with the proposed per-layer step-sizes are able to choose effective LR schedules and outperform fine-tuned LR versions of these methods as well as popular first-order and second-order algorithms for training DNNs on Autoencoder, Convolutional Neural Network (CNN) and Graph Convolutional Network (GCN) models. Finally, it is proved that an idealized version of SGD with the layer-wise step sizes converges linearly when using full-batch gradients.