SGD for Structured Nonconvex Functions: Learning Rates, Minibatching and Interpolation

TL;DR

This paper establishes new convergence guarantees for SGD on structured non-convex functions, including quasar convex and Polyak-Lojasiewicz functions, with insights into minibatching and interpolation scenarios.

Contribution

It introduces weaker residual conditions for convergence analysis and provides optimal minibatch size insights for structured non-convex functions.

Findings

SGD converges to a global minimum under structural assumptions.

Expected Residual condition is weaker than previous assumptions.

Optimal minibatch size is characterized for efficient training.

Abstract

Stochastic Gradient Descent (SGD) is being used routinely for optimizing non-convex functions. Yet, the standard convergence theory for SGD in the smooth non-convex setting gives a slow sublinear convergence to a stationary point. In this work, we provide several convergence theorems for SGD showing convergence to a global minimum for non-convex problems satisfying some extra structural assumptions. In particular, we focus on two large classes of structured non-convex functions: (i) Quasar (Strongly) Convex functions (a generalization of convex functions) and (ii) functions satisfying the Polyak-Lojasiewicz condition (a generalization of strongly-convex functions). Our analysis relies on an Expected Residual condition which we show is a strictly weaker assumption than previously used growth conditions, expected smoothness or bounded variance assumptions. We provide theoretical guarantees for the convergence of SGD for different step-size selections including constant, decreasing and the recently proposed stochastic Polyak step-size. In addition, all of our analysis holds for the arbitrary sampling paradigm, and as such, we give insights into the complexity of minibatching and determine an optimal minibatch size. Finally, we show that for models that interpolate the training data, we can dispense of our Expected Residual condition and give state-of-the-art results in this setting.