Quantitative Convergence Analysis of Projected Stochastic Gradient Descent for Non-Convex Losses via the Goldstein Subdifferential
This work provides theoretical guarantees for a widely used optimization algorithm in machine learning, addressing a gap in convergence analysis for constrained non-convex problems, though it is incremental as it builds on existing SGD frameworks.
The paper tackles the problem of analyzing convergence for projected stochastic gradient descent (SGD) with non-convex losses, proposing a new criterion based on the Goldstein subdifferential that allows direct comparison with unconstrained SGD and avoids the need for variance reduction. It obtains asymptotic convergence and non-asymptotic bounds, including O(N^{-1/3}) in expectation and O(N^{-1/5}) with high probability for IID sub-Gaussian data.
Stochastic gradient descent (SGD) is the main algorithm behind a large body of work in machine learning. In many cases, constraints are enforced via projections, leading to projected stochastic gradient algorithms. In recent years, a large body of work has examined the convergence properties of projected SGD for non-convex losses in asymptotic and non-asymptotic settings. Strong quantitative guarantees are available for convergence measured via Moreau envelopes. However, these results cannot be compared directly with work on unconstrained SGD, since the Moreau envelope construction changes the gradient. Other common measures based on gradient mappings have the limitation that convergence can only be guaranteed if variance reduction methods, such as mini-batching, are employed. This paper presents an analysis of projected SGD for non-convex losses over compact convex sets. Convergence is measured via the distance of the gradient to the Goldstein subdifferential generated by the constraints. Our proposed convergence criterion directly reduces to commonly used criteria in the unconstrained case, and we obtain convergence without requiring variance reduction. We obtain results for data that are independent, identically distributed (IID) or satisfy mixing conditions ($L$-mixing). In these cases, we derive asymptotic convergence and $O(N^{-1/3})$ non-asymptotic bounds in expectation, where $N$ is the number of steps. In the case of IID sub-Gaussian data, we obtain almost-sure asymptotic convergence and high-probability non-asymptotic $O(N^{-1/5})$ bounds. In particular, these are the first non-asymptotic high-probability bounds for projected SGD with non-convex losses.