Variance Reduction in Stochastic Gradient Descent: SVRG and SAGA

Stochastic gradient descent pays for cheap iterations with a rate that depends inversely on the target accuracy. The cause is the variance of the sampled gradient, which does not shrink as the iterate approaches the optimum, so a constant step size leaves the iterate jittering inside a ball and a decaying step size forces a sublinear rate. For the finite-sum problems that dominate supervised learning, where the objective is an average of per-example losses, the variance can be removed instead of tolerated. The trick is to estimate the gradient from a correlated pair of quantities whose difference has small variance but still averages to the true gradient, an idea imported from Monte Carlo simulation under the name control variate.

Why the noise survives at the optimum

The variance of the sampled gradient at a point x is the variance of the per-example gradients ∇fᵢ(x). At the minimiser of the average loss, the average gradient is zero, but individual examples generally have non-zero gradients that sum to zero: one example wants the parameters moved one way, another wants them moved back. That cancellation happens only in the average. The per-example variance therefore stays bounded away from zero at the optimum, and any method that samples examples without correcting for the disagreement retains the corresponding noise.

The consequence is a hard limit on the rate. To reduce the expected function-value gap of a strongly convex objective to ε, fixed-step stochastic gradient descent requires on the order of 1/ε iterations; full-batch gradient descent requires O(κ log(1/ε)) iterations but each costs n gradient evaluations. For small accuracy targets the second number wins by a wide margin once κ is modest, and it is exactly the small-accuracy regime that matters at the end of training. Variance reduction is the attempt to have the cheap iteration and the logarithmic rate at the same time.

Control variates: subtract an old gradient, add it back exactly

Let ḡ = (1/n) Σ ∇fᵢ(w̃) be the full gradient at a reference point w̃. The estimator

ĝᵢ = ∇fᵢ(w) − ∇fᵢ(w̃) + ḡ

is unbiased for ∇f(w) for any fixed w̃, because the expectation of the first two terms is ∇f(w) − ḡ. Its variance is the variance of the random variable ∇fᵢ(w) − ∇fᵢ(w̃), which vanishes as w → w̃. If the per-example gradient mapping is continuous, the differences shrink with the step, so the noise decreases as the iterate approaches the reference point. That is the entire mechanism: the estimator is made unbiased by adding the full gradient, and made low-variance by differencing two nearby evaluations of the same example.

The requirement is that a fresh reference point can be produced at a bounded cost. Producing one costs a full pass over the data, so the method must take many cheap steps per full pass, and that ratio is the parameter that governs everything.

The two loops of SVRG

Stochastic variance-reduced gradient descent organises the computation as an outer loop and an inner loop. The outer loop computes the snapshot gradient ḡ at the current w̃, costing n per-example gradient evaluations. The inner loop runs m steps with a constant step size, each using the control-variate estimator at the fixed snapshot. After m inner steps, the current point becomes the new snapshot and the process repeats.

Within an outer round the step size can be constant rather than decaying, because the variance is bounded by the movement of the iterate away from the snapshot rather than by the problem’s intrinsic noise. For a μ-strongly convex objective with L-smooth per-example losses, the analysis requires the step to be a sufficiently small multiple of 1/L and the inner-loop length m to exceed a multiple of the condition number κ = L/μ. Under those conditions the expected gap contracts geometrically across outer rounds, and the total number of per-example gradient evaluations to reach accuracy ε is O((n + κ) log(1/ε)). The n term is the unavoidable cost of the snapshots; the κ term is what one epoch costs. Compare with full-batch gradient descent, which needs O(κ log(1/ε)) iterations at n evaluations each, and the factor of n is gone.

import random

def svrg(params, grads, n, m, eta, rounds):
    """Outer rounds of snapshot gradients followed by m corrected inner steps."""
    x = list(params)
    for _ in range(rounds):
        snap = list(x)                                  # reference point for this round
        g_snap = [sum(grads(snap, i)[j] for i in range(n)) / n
                  for j in range(len(x))]               # one full pass
        for _ in range(m):
            i = random.randrange(n)
            gi, gj = grads(x, i), grads(snap, i)
            # unbiased for grad f(x) because E[grad f_i(snap) - g_snap] = 0
            g = [gi[j] - gj[j] + g_snap[j] for j in range(len(x))]
            x = [x[j] - eta * g[j] for j in range(len(x))]
    return x

One detail of the loop is worth stating explicitly: the snapshot gradient is computed once per round from a single full pass and is reused unchanged by every iteration of the inner loop, so the cost is n evaluations per round regardless of m. Its length is the number of parameters, so storing it costs one extra array. The inner loop’s memory is therefore the parameters plus that array, which is what makes SVRG attractive on large models where the per-coordinate state of adaptive methods is also counted in arrays of parameter size.

SAGA and the stored gradient table

The alternative to a snapshot is to keep the last gradient seen for every example. SAGA maintains a table Gᵢ holding ∇fᵢ evaluated at the last point where example i was sampled, and updates a single row per iteration:

ĝ = ∇fᵢ(w) − Gᵢ + (1/n)Σⱼ Gⱼ, then Gᵢ ← ∇fᵢ(w).

The running average of the table is maintained incrementally, in one array of parameter length, so no full pass is needed after initialization. The bias of each stored gradient decays as the parameters move, so the estimator has the same control-variate structure as SVRG without the outer loop. The cost is the table: n vectors of length d, so n·d floating-point values. For a problem with 10⁶ examples and 10³ features that is 10⁹ values, on the order of 4 GB in single precision, which is more memory than many models’ parameters. SAGA also handles the composite case, applying a proximal operator to the result of the step, which places it in the family described in projected and proximal gradient methods.

The epoch length is not a free parameter

Both methods have a regime where they stop being better than plain stochastic gradient descent, and the boundary is the point worth remembering. SVRG’s inner-loop length must grow with the condition number; a schedule with m fixed while κ grows violates the analysis and the method degenerates to noisy gradient descent with an expensive full pass attached. The requirement is intuitive once the control-variate variance is written down: the estimator’s variance is proportional to how far the iterate has moved from the reference, and in an ill-conditioned problem the iterate moves a long way along the flat directions while making little progress in function value. If the inner loop is shorter than the time it takes to move appreciably from the snapshot, the variance reduction barely applies; if the loop is much longer, the later inner iterations are doing almost uncorrected stochastic gradient descent at constant step and the iterate is driven out to the noise ball.

The condition number also governs the step size, and the step size in turn bounds the noise floor from stochastic gradient descent and step decay: the same ησ²/μ expression appears, with σ² now the variance of the corrected estimator rather than the raw per-example variance. Variance reduction shrinks σ²; it does not change the curvature.

Where the method stops paying

Four conditions have to hold together for the gains to be real. The objective must be a finite sum with a moderate number of terms, so that the snapshot or the table is affordable. The individual losses must be smooth, since the control variate differences gradients and a non-smooth loss gives no useful difference at the kink. The scaling with the condition number must be tolerable, since for κ in the thousands, the epoch length itself becomes thousands of steps, which is fine, but the step size must shrink correspondingly and each outer round then costs a full pass for little progress. And the analysis applies to convex or strongly convex objectives; for non-convex problems the results available characterise convergence to a stationary point with a rate that no longer separates cleanly from plain stochastic gradient descent, and the practical methods that use variance reduction inside deep networks do so with schemes such as periodic snapshots and modified inner loops that no longer match the published conditions. Mini-batching composes with variance reduction in a straightforward way, reducing the estimator’s variance further by the batch factor at a proportional cost per iteration, which is often a better use of memory than a longer history in the table.

A final structural limitation is that the variance reduction is tied to the sample. Sampling examples independently with replacement wastes part of the gain that random reshuffling would preserve, and combining variance reduction with an epoch-based permutation requires care, because the control variate depends on each example’s gradient being evaluated at a well-defined reference point rather than on the order in which examples appear. The methods that combine both properties exist, but their epoch structure is more intricate than either ingredient alone.