Week 9: Variance Reduction

Finite-Sum Structure#

The ERMEmpirical Risk Minimization objective has a special structure explicitly exploited by variance reduction:

$f(x) = \frac{1}{n}\sum_{i=1}^n f_i(x)$

Each $f_i$ is the loss on a single training example (or a minibatch, for batched VR methods). The gradient is:

$\nabla f(x) = \frac{1}{n}\sum_{i=1}^n \nabla f_i(x)$

Computing the full gradient costs $O(n)$ per iteration — cheap for $n = 10^4$, expensive for $n = 10^7$. SGD samples one index $i$ uniformly and uses $\nabla f_i(x)$ as a stochastic gradient:

$\mathbb{E}_i[\nabla f_i(x)] = \nabla f(x)$

The estimator is unbiased, but its variance is the root problem.

Variance of SGD#

Even at the optimum $x^*$ where $\nabla f(x^*) = 0$, a stochastic gradient $\nabla f_i(x^*)$ is generally nonzero — different data points pull in different directions, and only their average cancels. The gradient variance at $x$ is:

$\sigma^2(x) = \frac{1}{n}\sum_{i=1}^n \|\nabla f_i(x) - \nabla f(x)\|^2$

SGD's convergence rate under strong convexity is:

$\mathbb{E}[\|x_k - x^*\|^2] \leq (1 - \eta\mu)^k \|x_0 - x^*\|^2 + \frac{\eta \sigma^2}{\mu}$

where $\sigma^2$ bounds the variance. The second term is an asymptotic noise floor that does not vanish as $k \to \infty$. To converge to the exact minimiser, SGD must decay the learning rate $\eta_k \to 0$. A constant learning rate converges only to a neighbourhood of the optimum, with radius proportional to $\sqrt{\eta \sigma^2}$.

The core tension: SGD is cheap per iteration ($O(1)$ vs $O(n)$ for full-batch) but requires decaying learning rates and many epochs for high accuracy. Variance reduction retains the $O(1)$ per-iteration cost (amortised) while achieving linear convergence to the exact solution — no learning rate decay needed.

SVRG#

Stochastic Variance Reduced Gradient (SVRG) uses a control variate — a correlated, zero-mean correction that reduces the effective variance of the stochastic gradient.

SVRG operates in outer loops (epochs). At the start of each outer loop $s$, a full gradient is computed at a snapshot $\tilde{x}_s$:

$\tilde{\mu}_s = \nabla f(\tilde{x}_s) = \frac{1}{n}\sum_{i=1}^n \nabla f_i(\tilde{x}_s)$

Then for $m$ inner iterations ($m \approx 2n$ in the original paper), the stochastic gradient is:

$v_k = \nabla f_i(x_k) - \nabla f_i(\tilde{x}_s) + \tilde{\mu}_s$

where $i$ is sampled uniformly at random. The key insight: $\mathbb{E}_i[\nabla f_i(\tilde{x}_s)] = \tilde{\mu}_s$, so the term $-\nabla f_i(\tilde{x}_s) + \tilde{\mu}_s$ has zero mean but is correlated with $\nabla f_i(x_k)$. When $x_k \approx \tilde{x}_s$, the two gradients are similar, and their difference is small — the variance is dramatically reduced.

After $m$ inner steps, the snapshot is updated: $\tilde{x}_{s+1} = x_m$ (or a random/average of the inner iterates), and a new full gradient is computed.

Convergence. Under strong convexity and smoothness, SVRG converges linearly (geometric rate) to the exact minimiser with a constant step size — no decay needed. The amortised cost per iteration is roughly $O(1 + n/m)$, balancing the occasional $O(n)$ full gradient against $m$ cheap $O(1)$ inner steps.

SAGA#

SAGA is a close cousin that avoids the outer-loop structure by maintaining a table of historical gradients:

At all times, the algorithm stores $\nabla f_i(\hat{x}_i)$ for each $i$, where $\hat{x}_i$ is the most recent iterate at which $\nabla f_i$ was evaluated. Each iteration:

1. Sample index $i$ uniformly.
2. Form the SAGA gradient: $v_k = \nabla f_i(x_k) - \nabla f_i(\hat{x}_i) + \frac{1}{n}\sum_{j=1}^n \nabla f_j(\hat{x}_j)$
3. Update the table entry for $i$: $\hat{x}_i = x_k$, and update the running average of stored gradients.

The running average acts as an approximation of the full gradient, updated incrementally (no full-batch pass needed). SAGA requires $O(n)$ memory to store the gradient table (one vector per training example), which limits its use to moderate $n$. SVRG's memory is $O(1)$ beyond storing a single snapshot gradient.

Practice Notes#

Where VR methods shine:

• Moderate $n$ (thousands to low millions of examples) where computing full gradients periodically is affordable.
• Convex or strongly convex problems where linear convergence matters.
• Logistic regression, SVMs, and classical ML with explicit finite-sum structure.

Where VR methods are rarely used:

• Deep learning with $n \gg 10^7$: storing per-example gradients (SAGA) or computing periodic full gradients (SVRG) is impractical.
• Nonconvex objectives: the control variate argument relies on convexity; the theory is weaker in the nonconvex case, though some extensions exist.
• When the computational bottleneck is $O(d)$ per gradient computation rather than the number of gradient evaluations — VR reduces the number of gradient calls, not the cost per call.

Legacy and influence. While classical SVRG/SAGA are not standard in deep learning pipelines, the control-variate idea has influenced modern methods: momentum-based variance reduction, stochastic Newton sketches, and the analysis of large-batch training all draw on the same variance decomposition.

Browser lab#

On synthetic finite-sum least squares ($n = 1000$, $d = 20$), compare the gradient error norm $\|\nabla f(x_k)\|$ for SGD vs SVRG over epochs. Demonstrate that SVRG drives the gradient to zero while SGD stagnates at the noise floor.

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)
n, d = 1000, 20
X = np.random.randn(n, d)
w_opt = np.random.randn(d)
y = X @ w_opt + 0.1 * np.random.randn(n)

# Precompute gradients for speed
grads = np.array([X[i] * (X[i] @ w - y[i]) for i, w in enumerate([np.zeros(d)] * n)])

def compute_gradient(w):
    return X.T @ (X @ w - y) / n

def compute_gradient_i(w, i):
    return X[i] * (X[i] @ w - y[i])

# SGD with constant step size
eta = 0.05
n_epochs = 60
w_sgd = np.zeros(d)
sgd_grad_norms = []

for epoch in range(n_epochs):
    perm = np.random.permutation(n)
    for i in perm:
        g = compute_gradient_i(w_sgd, i)
        w_sgd = w_sgd - eta * g
    sgd_grad_norms.append(np.linalg.norm(compute_gradient(w_sgd)))

# SVRG
m = 2 * n  # inner iterations per outer loop
w_svrg = np.zeros(d)
svrg_grad_norms = []

for outer in range(n_epochs // 2):
    w_snap = w_svrg.copy()
    full_grad_snap = compute_gradient(w_snap)

    for inner in range(m):
        i = np.random.randint(n)
        g = compute_gradient_i(w_svrg, i) - compute_gradient_i(w_snap, i) + full_grad_snap
        w_svrg = w_svrg - eta * g

        if inner % n == 0:
            svrg_grad_norms.append(np.linalg.norm(compute_gradient(w_svrg)))

fig, ax = plt.subplots(figsize=(8, 5))
ax.semilogy(sgd_grad_norms, "b-", linewidth=1.2, alpha=0.8, label="SGD (constant $\eta$)")
ax.semilogy(svrg_grad_norms, "r-", linewidth=1.5, label="SVRG")
ax.axhline(0.1, color="gray", linestyle="--", alpha=0.5, label="Noise floor (SGD)")
ax.set_xlabel("Epoch")
ax.set_ylabel("$\|\|\\nabla f(w)\|\|$ (log scale)")
ax.set_title("Gradient norm: SGD vs SVRG on finite-sum least squares")
ax.legend()
ax.grid(True, alpha=0.3)
plt.show()

print(f"Final SGD gradient norm: {sgd_grad_norms[-1]:.6f}")
print(f"Final SVRG gradient norm: {svrg_grad_norms[-1]:.6f}")