Demystifying SGD with Doubly Stochastic Gradients

TL;DR

This paper analyzes the convergence of doubly stochastic gradient descent methods for complex optimization problems involving expectations, providing new insights into their theoretical properties and practical efficiency.

Contribution

It establishes convergence results for doubly SGD under general conditions, including dependent estimators, and analyzes how to optimally allocate computational resources.

Findings

Convergence of doubly SGD is proven under broad conditions.

Random reshuffling improves complexity dependence.

Guidance on resource allocation for minibatch and Monte Carlo samples.

Abstract

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 $b \times m$, where $b$ is the minibatch size and $m$ 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.