Hands-on Practical: Implementing SVRG

Stochastic Variance Reduced Gradient (SVRG) is a potent technique designed to combat the high variance inherent in standard Stochastic Gradient Descent (SGD). An implementation guide for SVRG covers observing its behavior and understanding the practical considerations involved. The focus is on a clear implementation to highlight the main mechanics of the algorithm.

Recall the core idea from our discussion on SVRG: periodically compute a full gradient over the entire dataset to serve as a reference point. Then, in subsequent stochastic steps, adjust the standard stochastic gradient by comparing the gradient of the same data point evaluated at the current parameters versus the reference parameters. This adjustment significantly reduces the variance of the gradient estimate, often leading to faster convergence compared to SGD, especially when measured in terms of passes over the data.

The SVRG Algorithm Structure

Let's outline the steps involved in a typical SVRG implementation for minimizing a function $f(w) = \frac{1}{n} \sum_{i=1}^n f_i(w)$, where $f_i(w)$ is the loss for data point $i$:

1. Initialization: Choose an initial parameter vector $w_0$, a learning rate $\eta > 0$, and the number of inner loop iterations $m$. The value $m$ typically corresponds to the effective number of stochastic updates performed between full gradient computations, often set to be proportional to $n$ (e.g., $m=n$ or $m=2n$).
2. Outer Loop: Iterate for $S$ epochs (passes over the full gradient computation). a. Snapshot Parameters: Store the current parameter vector as the snapshot $\tilde{w} = w_{s-1}$ (where $s$ is the outer loop index). b. Compute Full Gradient: Calculate the average gradient over the entire dataset using the snapshot parameters: $\mu = \nabla f(\tilde{w}) = \frac{1}{n} \sum_{i=1}^n \nabla f_i(\tilde{w})$. This is the computationally expensive step performed once per outer loop. c. Initialize Inner Loop Parameters: Set the starting point for the inner loop: $w_0 = \tilde{w}$. d. Inner Loop: Iterate $t$ from $0$ to $m-1$. i. Sample Data Point: Randomly select an index $i_t$ from $\{1, ..., n\}$. ii. Compute Stochastic Gradients: Calculate the gradient for the selected data point at the current parameters $w_t$ and at the snapshot parameters $\tilde{w}$: * $\nabla f_{i_t}(w_t)$ * $\nabla f_{i_t}(\tilde{w})$ iii. Update Parameters: Apply the SVRG update rule: $w_{t+1} = w_t - \eta \left( \nabla f_{i_t}(w_t) - \nabla f_{i_t}(\tilde{w}) + \mu \right)$ e. Update Outer Loop Parameters: Set the result of the inner loop as the parameters for the next outer loop iteration: $w_s = w_m$. (Alternatively, one could choose $w_s$ randomly from $\{w_0, ..., w_m\}$ or use the average). Using $w_s = w_m$ is common.

Implementing SVRG for Linear Regression

Let's implement SVRG to solve a simple least-squares linear regression problem. Our objective function for data $(X, y)$ where $X \in \mathbb{R}^{n \times d}$ and $y \in \mathbb{R}^n$ is:

$$f(w) = \frac{1}{2n} \| Xw - y \|_2^2 = \frac{1}{n} \sum_{i=1}^n \frac{1}{2} (x_i^T w - y_i)^2$$

Here, $f_i(w) = \frac{1}{2} (x_i^T w - y_i)^2$. The gradient for a single data point $(x_i, y_i)$ is:

$$\nabla f_i(w) = (x_i^T w - y_i) x_i$$

And the full gradient is:

$$\nabla f(w) = \frac{1}{n} X^T (Xw - y)$$

Below is a Python implementation using NumPy.

import numpy as np

def linear_loss(w, X, y):
  """Calculates the mean squared error loss."""
  n_samples = X.shape[0]
  predictions = X.dot(w)
  loss = (1 / (2 * n_samples)) * np.sum((predictions - y)**2)
  return loss

def full_gradient(w, X, y):
  """Calculates the full gradient for linear regression."""
  n_samples = X.shape[0]
  predictions = X.dot(w)
  gradient = (1 / n_samples) * X.T.dot(predictions - y)
  return gradient

def stochastic_gradient(w, X, y, index):
  """Calculates the stochastic gradient for a single data point."""
  x_i = X[index]
  y_i = y[index]
  prediction_i = x_i.dot(w)
  # Reshape x_i to ensure correct dot product dimensions if needed
  if x_i.ndim == 1:
      x_i = x_i[:, np.newaxis] # Make it a column vector if it's flat

  # Calculate gradient: (prediction - actual) * feature_vector
  # Ensure the scalar (prediction_i - y_i) multiplies the vector x_i correctly
  # The result should have the same shape as w
  gradient_i = (prediction_i - y_i) * x_i.flatten() # Flatten x_i back if necessary or ensure result shape matches w

  # Ensure gradient_i has the same shape as w (e.g., (d,))
  if w.ndim == 1 and gradient_i.shape != w.shape:
       gradient_i = gradient_i.reshape(w.shape)
  elif w.ndim > 1 and gradient_i.shape != w.shape: # e.g. w is (d, 1)
       gradient_i = gradient_i.reshape(w.shape)

  return gradient_i

def svrg_optimizer(X, y, n_outer_epochs, m, learning_rate):
  """Implements the SVRG algorithm for linear regression."""
  n_samples, n_features = X.shape

  # Ensure y is a 1D array or a column vector consistent with predictions
  if y.ndim == 1:
      y = y # Keep as 1D array
  elif y.ndim == 2 and y.shape[1] == 1:
       y = y.flatten() # Make it 1D for easier calculations with predictions
  else:
       raise ValueError("y must be a 1D array or a column vector")

  # Initialize parameters (e.g., zeros)
  w = np.zeros(n_features) 
  # If you prefer column vector parameters:
  # w = np.zeros((n_features, 1))

  loss_history = []

  print(f"Starting SVRG: {n_outer_epochs} outer epochs, m={m}, lr={learning_rate}")

  for s in range(n_outer_epochs):
    # 1. Snapshot and Full Gradient Calculation
    w_snapshot = np.copy(w)
    mu = full_gradient(w_snapshot, X, y)

    # Store current loss (calculated with the snapshot before inner loop)
    current_loss = linear_loss(w_snapshot, X, y)
    loss_history.append(current_loss)
    if s % 10 == 0 or s == n_outer_epochs - 1:
        print(f"Outer Epoch {s+1}/{n_outer_epochs}, Loss: {current_loss:.6f}")

    # 2. Inner Loop
    w_inner = np.copy(w_snapshot) # Start inner loop from snapshot
    for t in range(m):
      # Select random sample
      i_t = np.random.randint(0, n_samples)

      # Calculate gradients
      grad_current = stochastic_gradient(w_inner, X, y, i_t)
      grad_snapshot = stochastic_gradient(w_snapshot, X, y, i_t)

      # SVRG update
      w_inner = w_inner - learning_rate * (grad_current - grad_snapshot + mu)

    # Update outer loop parameters
    w = np.copy(w_inner) # Use the result of the inner loop

  print("SVRG Finished.")
  return w, loss_history

# --- Example Usage ---
# Generate synthetic data
np.random.seed(42)
n_samples, n_features = 1000, 10
X = np.random.randn(n_samples, n_features)
true_w = np.random.randn(n_features) * 5 
y = X.dot(true_w) + np.random.randn(n_samples) * 0.5 # Add some noise

# SVRG parameters
n_outer_epochs = 50  # Number of times we compute the full gradient
m = n_samples        # Number of inner loop steps (1 effective pass per outer epoch)
learning_rate = 0.1 

# Run SVRG
w_svrg, svrg_loss_history = svrg_optimizer(X, y, n_outer_epochs, m, learning_rate)

print("\nLearned weights (SVRG):", w_svrg)

Analyzing the Results

To appreciate the behavior of SVRG, it's instructive to compare its convergence profile against standard SGD. You would typically implement a similar loop for SGD (updating using only $\nabla f_{i_t}(w_t)$) and run it for an equivalent number of gradient computations or effective data passes.

Let's visualize the potential difference in convergence speed. We'll plot the loss against the number of effective passes through the data. For SVRG, each outer epoch involves one full gradient calculation (equivalent to one pass) plus $m$ stochastic gradient calculations. If $m=n$, each outer epoch corresponds to roughly two effective passes over the data (one for the full gradient, one effective pass for the $n$ stochastic updates). For standard SGD, $n$ stochastic updates constitute one effective pass.

Illustrative comparison of loss reduction per effective data pass for SVRG and a typical SGD run. SVRG often exhibits faster convergence, especially on a log scale for the loss. Note that the SGD curve is manually created for illustration. Run your own SGD implementation for a direct comparison on the same data.

Implementation Details and Tuning

• Choosing m: The number of inner loop iterations, $m$, is a significant hyperparameter.
  • Small $m$: More frequent full gradient calculations. This keeps the variance low but increases the overhead per effective data pass.
  • Large $m$: Less frequent full gradient calculations. Cheaper outer epochs, but the parameters $w_t$ might drift further from the snapshot $\tilde{w}$, potentially slowing convergence if $m$ is too large.
  • Common choices are $m=n$ or $m=2n$, effectively performing 1 or 2 passes worth of stochastic updates between full gradient computations. Theoretical analysis often suggests $m$ should be proportional to $n$. Experimentation is crucial.
• Learning Rate $\eta$: SVRG can often tolerate, and may even require, larger learning rates than SGD to converge quickly. The variance reduction allows for more aggressive steps. However, too large a learning rate can still lead to instability. Tuning $\eta$ is necessary, often via grid search or based on theoretical guidelines if the problem structure allows (e.g., knowledge of smoothness constants).
• Computational Cost: The main cost is the full gradient calculation ($\mathcal{O}(nd)$ for many standard ML models) performed once per outer epoch. The inner loop involves $m$ stochastic gradient computations, which are typically much cheaper ($\mathcal{O}(d)$ per step). SVRG is most effective when the cost of the full gradient is manageable periodically, and the benefit of faster convergence (fewer total effective passes) outweighs this periodic cost compared to the slower, high-variance progress of SGD.

This practical implementation should provide a solid foundation for applying SVRG to larger, more complex problems. Remember that careful tuning of $m$ and $\eta$ is often required to achieve the best performance. Experimenting with SVRG on your own datasets and comparing it to SGD is the best way to gain intuition about its strengths and practical behavior.