Demystifying SGD with Doubly Stochastic Gradients

Optimization objectives in the form of a sum of intractable expectations are rising in importance (e.g., diffusion models, variational autoencoders, and many more), a setting also known as "finite sum with infinite data." For these problems, a popular strategy is to employ SGD with doubly stochastic gradients (doubly SGD): the expectations are estimated using the gradient estimator of each component, while the sum is estimated by subsampling over these estimators. Despite its popularity, little is known about the convergence properties of doubly SGD, except under strong assumptions such as bounded variance. In this work, we establish the convergence of doubly SGD with independent minibatching and random reshuffling under general conditions, which encompasses dependent component gradient estimators. In particular, for dependent estimators, our analysis allows fined-grained analysis of the effect correlations. As a result, under a per-iteration computational budget of × , where is the minibatch size and is the number of Monte Carlo samples, our analysis suggests where one should invest most of the budget in general. Furthermore, we prove that random reshuffling (RR) improves the complexity dependence on the subsampling noise. ef f Effective sample size of Eq. ( 5 ) ( ) Unbiased stochastic estimator of ∇ Eq. ( 7 ) ( ; ) Integrand of estimator ( ) Eq. ( 7 ) ( ) Doubly stochastic estimator of ∇ Eq. ( 8 ) ℒ sub ER constant (Definition 1) of Assu. ℒ ER constant (Definition 1) of Assu. 2 BV constant (Definition 2) of Assu. 2 BV constant (Definition 2) of Assu. Stochastic Gradient Descent on Finite-Sums Stochastic gradient descent (SGD) is an optimization algorithm that repeats the steps where, Π is a projection operator onto , ( ) =0 is some stepsize schedule, ( ) is an unbiased estimate of ∇ ( ).