Better Theory for SGD in the Nonconvex World (2002.03329v3)

Abstract: Large-scale nonconvex optimization problems are ubiquitous in modern machine learning, and among practitioners interested in solving them, Stochastic Gradient Descent (SGD) reigns supreme. We revisit the analysis of SGD in the nonconvex setting and propose a new variant of the recently introduced expected smoothness assumption which governs the behaviour of the second moment of the stochastic gradient. We show that our assumption is both more general and more reasonable than assumptions made in all prior work. Moreover, our results yield the optimal $\mathcal{O}(\varepsilon^{-4})$ rate for finding a stationary point of nonconvex smooth functions, and recover the optimal $\mathcal{O}(\varepsilon^{-1})$ rate for finding a global solution if the Polyak-{\L}ojasiewicz condition is satisfied. We compare against convergence rates under convexity and prove a theorem on the convergence of SGD under Quadratic Functional Growth and convexity, which might be of independent interest. Moreover, we perform our analysis in a framework which allows for a detailed study of the effects of a wide array of sampling strategies and minibatch sizes for finite-sum optimization problems. We corroborate our theoretical results with experiments on real and synthetic data.

A Theoretical Perspective on Stochastic Gradient Descent in Nonconvex Optimization

The paper, "Better Theory for SGD in the Nonconvex World," by Khaled and Richtárik, addresses the substantial challenge of nonconvex optimization in machine learning, focusing on stochastic gradient descent (SGD), particularly in the field of nonconvex functions. The authors propose a novel variant of expected smoothness (ES) to facilitate more accurate theoretical analysis and practical performance assessments of SGD, aligning better with the realistic dynamics of nonconvex loss landscapes encountered in contemporary machine learning.

Overview and Contributions

The paper introduces a generalization of the expected smoothness assumption, offering a more flexible and realistic theoretical framework for analyzing SGD in nonconvex settings. This assumption pertains to the behavior of the second moment of the stochastic gradient. The authors argue that previous assumptions used in literature are often overly restrictive and fail to capture the stochastic dynamics accurately as experienced in empirical practice. In contrast, the proposed ES assumption allows for a more comprehensive modeling of nonconvex optimization problems, accommodating various sampling strategies and minibatch sizes without the stringent conditions found in traditional analyses.

The key contributions of the paper are as follows:

1. Expected Smoothness Assumption: The introduction of the ES assumption, positing that the second moment of the stochastic gradient can be bounded in terms of function values, gradients, and a constant, is pivotal. This provides a more nuanced view than traditional convex-oriented assumptions, including bounded variance or strong growth conditions.
2. Rate of Convergence: The analysis yields optimal rates for the convergence of SGD to stationary points in nonconvex optimization at $O(\epsilon^{-4})$. Moreover, if the Polyak-Łojasiewicz (PL) condition holds, a $O(\epsilon^{-1})$ rate can be attained for global solutions, which is significant in bridging the gap between nonconvex scenarios and tractable convex optimization results.
3. Sampling and Minibatch Strategies: The framework allows for a detailed exploration of sampling strategies, demonstrating that the ES assumption holds under both these sampling methods and stochastic gradient compression techniques. Furthermore, it provides theoretical guidance for selecting optimal minibatch sizes and importance sampling probabilities, advancing both the theoretical understanding and practical implementation of SGD in nonconvex settings.
4. Empirical Verification: The theoretical results are supported by rigorous experiments using synthetic and real-world datasets. These demonstrate not only the practical applicability of the ES model but also validate the theoretical convergence rates and effectiveness of importance sampling.

Implications and Future Directions

This paper insists on reconsidering classic assumptions in analysis, urging the research community to adopt more realistic models in stochastic optimization that capture the intricacies encountered in large-scale nonconvex problems. The implications are manifold:

• Theoretical Advancements: Broadening the lens through which stochastic gradients are analyzed will inspire future work to refine optimization theory in nonconvex contexts, potentially leading to new algorithms that leverage these theoretical insights for faster convergence and enhanced robustness.
• Practical Applications: For practitioners, the findings underscore the importance of model-relevant assumptions that accurately capture stochasticity. The framework enables more efficient utilization of computational resources by optimizing sampling and minibatching strategies, enhancing applicability in machine learning tasks like deep learning.
• Extensions to Other Algorithms: Future research could explore extensions of the ES assumption to other randomized optimization methods, comparing efficacy against SGD in various nonconvex scenarios and potentially unveiling new avenues in optimization strategies.

In conclusion, Khaled and Richtárik's paper makes significant strides in understanding SGD in nonconvex environments, proving both theoretically robust and empirically valid, and establishing a foundation for continued exploration within stochastic optimization paradigms.