Optimal Risk Bounds of Stochastic Gradient Descent for Shallow ReLU Networks

Recent work shows that gradient descent (GD) can achieve almost optimal risk bounds for shallow ReLU networks with a logarithmic width. However, these discussions require $O(n^2)$ gradient computations to achieve this optimality, which is not appealing for modern machine learning problems with large sample size $n$. In this paper, we significantly improve the existing gradient complexity by showing that stochastic gradient descent (SGD) can achieve similar risk bounds with $O(n)$ gradient computations for shallow ReLU networks with a logarithmic width. As compared to GD, the analysis with SGD is more challenging as there are additional fluctuations incurred by the stochastic gradient noise along the optimization process. With concentration inequalities to handle these fluctuations, we show that the entire SGD trajectory stays around its initialization point with high probability. Under an NTK separability condition with margin $\gamma$, we show SGD achieves near-optimal risk bounds $\widetilde{O}(1/(n\gamma^2))$ with a logarithmic width.