SketchySGD: Reliable Stochastic Optimization via Randomized Curvature Estimates
Zachary Frangella, Pratik Rathore, Shipu Zhao, Madeleine Udell · SIAM Journal on Mathematics of Data Science · 2024
Abstract. We introduce SketchySGD, a stochastic second-order method that uses sketching to approximate the curvature of the loss function. SketchySGD improves on existing stochastic gradient methods in machine learning by using randomized low-rank approximations to the subsampled Hessian and by introducing an automated step size that works well across a wide range of convex machine learning problems. We show theoretically that SketchySGD with a fixed step size converges linearly to a small ball around the minimum. Further, in the ill-conditioned setting, we show that SketchySGD converges at a faster rate than stochastic gradient descent for least-squares problems. We validate this improvement empirically with ridge regression experiments on real data. Numerical experiments on both ridge and logistic regression problems with dense and sparse data show that SketchySGD equipped with its default hyperparameters can achieve comparable or better results than popular stochastic gradient methods and preconditioned conjugate gradients, even when they have been tuned to yield their best performance. In particular, SketchySGD is able to solve an ill-conditioned logistic regression problem with a data matrix that takes more than 840 GB of RAM to store, while its competitors, even when tuned, are unable to make any progress. SketchySGD’s ability to work out of the box with its default hyperparameters and excel on ill-conditioned problems is an advantage over other stochastic gradient methods, most of which require careful hyperparameter tuning (especially of the learning rate) to obtain good performance and degrade in the presence of ill-conditioning.