Stochastic Optimization with Bandit Sampling
This work addresses a bottleneck in stochastic optimization for machine learning practitioners, offering a general method to improve efficiency, though it is incremental as it builds on existing sampling techniques.
The paper tackles the problem of high variance in gradient estimators for stochastic optimization by proposing a non-uniform sampling approach using multi-armed bandits, resulting in asymptotic approximation of optimal variance within a factor of 3 and significant reductions in convergence time and variance for algorithms like SGD, SVRG, and SAGA.
Many stochastic optimization algorithms work by estimating the gradient of the cost function on the fly by sampling datapoints uniformly at random from a training set. However, the estimator might have a large variance, which inadvertently slows down the convergence rate of the algorithms. One way to reduce this variance is to sample the datapoints from a carefully selected non-uniform distribution. In this work, we propose a novel non-uniform sampling approach that uses the multi-armed bandit framework. Theoretically, we show that our algorithm asymptotically approximates the optimal variance within a factor of 3. Empirically, we show that using this datapoint-selection technique results in a significant reduction in the convergence time and variance of several stochastic optimization algorithms such as SGD, SVRG and SAGA. This approach for sampling datapoints is general, and can be used in conjunction with any algorithm that uses an unbiased gradient estimation -- we expect it to have broad applicability beyond the specific examples explored in this work.