In a world where data can be collected continuously and storage costs are cheap, issues related to the growing size of interesting datasets can pose a problem unless we have the right tools for the task. Indeed, in the event where we have streaming data it might be impossible to wait until the "end" before fitting our model since it may never come. Alternatively it might be problematic to even store all of the data, scattered across many different servers, in memory before using it. Instead it would be preferable to do an update each time some new data (or a small batch of it) arrives. Similarly we might find ourselves in an offline situation where the number of training examples is very large and traditional approaches, such as gradient descent, start to become too slow for our needs.

Stochastic gradient descent (SGD) offers an easy solution to all of these problems.

In this post we explore the convergence properties of stochastic gradient descent and a few of its variants, namely

• Polyak-Ruppert averaged stochastic gradient;
• Adaptive Gradient, also known as AdaGrad;
• Stochastic Average Gradient, also known as SAG.

 

Stochastic Gradient Descent for Logistic Regression¶

In this post we will consider the classification dataset quantum.mat that I used during an assignment for the course CPSC 540 - Machine Learning at UBC. It contains a matrix $X$ with dimensions $50000 \times 79$, as well as a vector $y$ taking values $\pm 1$, which we will classify using logistic regression. The true minimum of our model's cost function lies at $f(w) = 2.7068 \times 10^4$, but various implementations of SGD have drastically different performances as we will see.

The cost function for logistic regression with L2-regularization can be written in a form well-suited to our minimization objective:

\begin{equation} f(w) = \sum_{i=1}^n f_i(w), \;\;\;\;\; \text{with} \;\; f_i(w) = \log \Big( 1 + \exp \left( - y_i X_{ia} w_a \right) \Big) + \frac{\lambda}{2 n} w^2. \end{equation}

The key idea behind SGD is to approximate the true gradient $\nabla f(w)$ by the successive gradients at individual training examples $\nabla f_i(w)$, which naturally befits online learning. The rationale for doing so is that it is very costly to compute the exact direction of the gradient, whereas a good but noisy estimate can be obtained by looking at a few examples only. This noise has the added benefit of preventing SGD from getting stuck in the shallow local minima that might be present for non-convex optimization objectives (such as neural networks), at the cost of never truly converging to a minimum but rather in a neighborhood around it.

We now proceed to build a LogisticRegressionSGD() class that implements the logistic regression cost function written above. We will then include different stochastic optimization methods and see how close they can get to the true minimum after 10 epochs.

Regular Stochastic Gradient Descent¶

Schematically the standard SGD algorithm takes the form

1. Initialize weights $w$ and learning rate $\eta_t$.

2. • Randomly permute training examples.

   • For $i = 1 : n$ do $w \leftarrow w - \eta_t \nabla f_i(w)$.

3. Repeat step 2 until convergence.

Many theory papers use the step size $\eta_t = 1/\lambda t$, which offers the best convergence rate in the worst case scenario. However choosing this learning rate typically leads the SGD algorithm in regions of parameter space afflicted by numerical overflow of the cost function before it ultimately "converges" to $f(w) \geq 4.5 \times 10^4$, far away from the global minimum.

Alternatives include choosing a constant learning rate (we found that $\eta = 10^{-4}$ gave reasonable results) or an iteration-dependent rate that slowly converges to 0, such as

\begin{equation} \eta_t = \frac{\sqrt{n}}{\sqrt{n} + t}. \end{equation}

We find that the last two approaches yield similar results, although the latter requires the fine-tuning of both the numerator and denominator in order to work optimally, and also makes it hard to decide when to stop since later iterations move very slowly.

Polyak-Ruppert Averaged Stochastic Gradient¶

Rather than use the information contained in the weights $w$ at iteration $t$ to determine the descent direction, it is often an improvement to use a running average instead, which keeps a memory of previous iterations

\begin{equation} \overline{w}_t = \overline{w}_{t-1} - \frac{1}{t} \left( \overline{w}_{t-1} - w_t \right). \end{equation}

Doing so results in a slight improvement over regular SGD. Note that convergence improves if we start averaging after nAvg $\geq 2$ passes in order to smooth out the initial irregularities.

Adaptive Gradient (AdaGrad)¶

One of the main drawbacks of the stochastic optimization methods outlined above is the need to manually choose the optimal learning rate for the problem at hand. AdaGrad, an algorithm proposed in 2011, eschews this problem by computing an appropriate learning rate for each direction $\hat{w_a} \in \mathbb{R}^d$.

AdaGrad automatically assigns a higher learning rate to rare/sparse features, which typically have a higher predictive power than common ones. We can understand this intuitively by thinking about words in a story: rare words like Daenerys and dragons provide significantly more information and context for the audience of Game of Thrones than common ones such as the or a. Therefore AdaGrad ensures that the most predictive features have larger updates (i.e. the associated weights increase/decrease proportionally to their importance) than the ones providing irrelevant information.

The weight update for AdaGrad is given by

\begin{equation} w_{t+1} = w_t - \eta_t D_t \nabla f(w_t), \end{equation}

where the diagonal matrix $D_t$ has elements

\begin{equation} (D_t)_{jj} = \frac{1}{\sqrt{\delta + \sum_{k=0}^t \nabla_j f_{i_k}(w_k)^2}}. \end{equation}

Here $i_k$ denotes example $i$ chosen randomly on iteration $k$, $\nabla_j$ is the $j$th element of the gradient, and $\delta$ is a small number to prevent division by 0. All we need to do is fiddle with the constant learning rate $\eta_t = \eta$ since $D_t$ automatically takes care of assigning higher importance to sparse features.

Stochastic Adaptive Gradient (SAG)¶

Last but not least, we now discuss the SAG algorithm, a variant on batching in SGD that was published in 2015. The basic implementation of this method can be explained schematically as follows:

1. Randomly select $i_t$ and compute the gradient $\nabla f_{i_t} (w_t)$.

2. Update the weights by taking a step towards the average of all the gradients computed so far

   $$ w_{t+1} = w_t - \eta_t \left( \frac{1}{n} \sum_{i=1}^n G_i^t \right), $$

   where $G_i^t$ keeps in memory all the gradients $\nabla f_{i_t} (w)$ computed before iteration $t$ (with replacement if training example $i_t$ is visited repeatedly).

3. Repeat.

Additionally, in contrast the the methods outlined before, SAG also leverages a property we have not used so far: Lipschitz continuity in the gradient of convex cost functions $f$

$$ \lVert \nabla f(x) - \nabla f(y) \lVert \; \leq \; L \lVert x - y \lVert. $$

By choosing the learning rate to be inversely proportional to the maximal Lipschitz constant over all training examples

$$ L = \frac{1}{4} \max_{1 \leq i \leq n} \left( \lVert x^i \lVert^2 \right) + \lambda, \;\;\; \eta_t = 1/L, $$

(here $x^i$ denotes a row of $X$), SAG achieves vastly superior convergence than all of the methods discussed above. In fact it is the only method of the ones outlined in this post that converges to the global minimum to 5 significant figures.

A caveat on randomness¶

The implementation of SAG below approaches the random updating of the gradients in two different ways, with surprising consequences.

1. In case = 1, the index $i_t$ is sampled randomly with replacement, meaning that not all training examples are necessarily visited after an epoch has been completed. This choice of sampling leads to the best convergence properties.

2. In case = 2, the index $i_t$ is sampled randomly without replacement, such that all training examples are cycled through exactly once during each pass. It turns out that simply reshuffling the cycle after each pass, the method of choice for all the methods above, actually yields a much worse performance for SAG.

It can be verified that random sampling with replacement barely affects the other SGD algorithms, but it remains somewhat of a mystery to me why this choice in randomness affects convergence so drastically.

Whereas regular stochastic gradient descent has been known for a long time, some of the variants discussed in this post are quite recent: AdaGrad and SAG were first described in peer-reviewed publications in 2011 and 2015 respectively. It is interesting to note that the latter is already one of the main solvers available for Logistic and Ridge regressions in scikit-learn because of its efficiency with large datasets. The rapid rise of online learning and artificial neural networks, where stochastic optimization shines brightest, are sure to stimulate research for even better methods in the near future.