Stochastic Gradient Descent Revisited: Variance Reduction

Stochastic Gradient Descent (SGD) is a workhorse for training models on large datasets. Its strength lies in using small data subsets (mini-batches) for gradient estimation, drastically reducing the computational cost per iteration compared to calculating the gradient over the entire dataset (Batch Gradient Descent). However, this efficiency comes at a cost: variance.

Each stochastic gradient estimate, $g_{i_t}(\theta_t) = \nabla f_{i_t}(\theta_t)$ (where $f_{i_t}$ is the loss on a mini-batch $i_t$ at iteration $t$), is an unbiased estimate of the true gradient $\nabla f(\theta_t) = \frac{1}{N} \sum_{i=1}^N \nabla f_i(\theta_t)$. This means, in expectation, the stochastic gradient points in the right direction: $\mathbb{E}[g_{i_t}(\theta_t)] = \nabla f(\theta_t)$. But for any single iteration, $g_{i_t}(\theta_t)$ can deviate significantly from $\nabla f(\theta_t)$. This difference introduces noise or variance into the parameter updates.

Why is Gradient Variance a Problem?

This inherent variance in SGD has several practical consequences:

1. Noisy Trajectory: The optimization path doesn't follow the smooth trajectory towards the minimum seen in Batch GD. Instead, it tends to zigzag, making progress less direct.
2. Slower Convergence: While each step is cheap, the noisy updates often mean that more steps are required to reach a good solution compared to methods using the full gradient. The convergence rate for SGD is typically sublinear, even for strongly convex problems where Batch GD achieves linear convergence.
3. Difficulty with Fine-Tuning: As the parameters approach a minimum, the true gradient magnitude becomes small. However, the variance of the stochastic gradient doesn't necessarily decrease at the same rate. This means the noise can dominate the signal, causing the parameters to oscillate around the minimum instead of converging precisely to it. This necessitates careful tuning of learning rate schedules, often requiring the learning rate to decay over time to dampen the oscillations.

Comparison of optimization paths on a simple quadratic loss surface. Batch GD takes smooth steps. SGD exhibits noisy steps due to gradient variance. Variance Reduced methods aim for faster, less noisy convergence than SGD.

The Principle of Variance Reduction

The core idea behind variance reduction techniques is to modify the stochastic gradient estimate to have lower variance, while retaining computational efficiency. We want an estimate $\tilde{g}_t(\theta_t)$ such that:

1. It remains computationally cheap, ideally requiring only a small mini-batch gradient calculation per iteration, similar to SGD.
2. It has significantly lower variance than the standard SGD gradient $g_{i_t}(\theta_t)$.
3. It is still an unbiased (or nearly unbiased in some variants) estimate of the true gradient $\nabla f(\theta_t)$.

Achieving this typically involves incorporating additional information into the gradient estimate. For instance, methods might periodically compute the full gradient or maintain moving averages of past gradients evaluated at different points. The goal is to construct a gradient estimate that corrects the noisy single-batch gradient towards the true gradient direction.

By reducing the variance, these methods often achieve faster theoretical convergence rates (approaching the linear rates of Batch GD for certain problem classes) and demonstrate better practical performance compared to standard SGD, especially in the later stages of optimization.

The subsequent sections will examine specific algorithms like Stochastic Average Gradient (SAG) and Stochastic Variance Reduced Gradient (SVRG), which implement different strategies to achieve this variance reduction, making large-scale optimization more efficient and effective.