How does one do Stochastic Gradient Descent (SGD) on an objective function that has a regularizer?

Question

I know that for Stochastic Gradient Descent, one picks a data point $(x_n, y_n)$ at random from the training set $S_N$ and then updates the parameter of the model in question.

If the cost function is:

$$J(w; S_N) = \frac{1}{N} \sum^{N}_{n = 1} J(w;x,y) = \frac{1}{N} \sum^{N}_{n = 1} Loss(w;x,y) = \frac{1}{N} \sum^{N}_{n = 1} cost(w;x,y) $$

A typical SGD update would look as follows:

$$ w := w - \gamma \nabla J(w; x,y)$$

However, if there is a regularizer present in the cost function, then its is not 100% clear to me how to stochastic gradient descent (SGD) would actually work.

Consider a regularized optimization problem that we want to optimize via SGD:

$$ H[w] = J(w; S_N) + \lambda R(w) = \frac{1}{N} \sum^{N}_{n = 1} J(w;x,y) + \lambda R(w) $$

My conjecture is that the correct way to do stochastic gradient descent is by doing:

$$ w := w - \gamma \nabla H[w] = w - \gamma ( \nabla J(w; x_n,y_n) + \lambda \nabla R(w) )$$

by choosing a random data point $(x_n, y_n)$.

The justification I have for this is the following (and wanted to check it with the community). The random direction that we move the parameters in depends on the random update that we do. However, what is the expected descent that we might do?

i.e. what is:

$$ \mathbb{E}_{n}[\nabla H(w)] = \mathbb{E}_{n}[ \nabla J(w; x,y) + \lambda \nabla R(w) ] $$

where the expectation is with respect to choosing a random data point at random from $S_N$ uniformly.

Since we are picking a data point at random uniformly the expectation of our SGD becomes:

$$ \mathbb{E}_{n}[ \nabla J(w; x,y) + \lambda \nabla R(w) ] = \frac{1}{N} \sum^{N}_{n = 1} \nabla J(w;x,y) + \mathbb{E}_{n}[\nabla \lambda R(w) ] $$

$$ \frac{1}{N} \sum^{N}_{n = 1} \nabla J(w;x,y) + \lambda \nabla R(w) \mathbb{E}_{n}[1] = \frac{1}{N} \sum^{N}_{n = 1} \nabla J(w;x,y) + \lambda \nabla R(w) $$

which is the same as if we tried to optimize the objective function in a batch way. Using this logic, this seams to me to be the way to use SGD with regularization. What do people think?

Answer 1

This is a good question, which annoyingly I haven't come across in many places. You've essentially got the right idea, but I'll flesh it out a little and describe the minibatch SGD scenario, which is more common.

Say you have training data $\mathcal{D} = \{x_i,y_i\}_{i=1}^N$, a loss function $J(w;\mathcal{D})$ and regulariser $\mathcal{R}(w)$, then your complete objective to minimise in parameters $w$ is, $$ H(w) = \lambda\mathcal{R}(w) + J(w;\mathcal{D}) = \lambda\mathcal{R}(w) + \sum_{i=1}^N J(w;x_i,y_i) $$ where $\lambda$ controls the influence of the regulariser. So your gradient is quite simply $$ \nabla_w H(w) = \lambda\nabla_w\mathcal{R}(w) + \sum_{i=1}^N \nabla_wJ(w;x_i,y_i) $$ In gradient descent your update equation would involve evaluating this expression on all the data and then updating at time $t$ as $$ w_t = w_{t-1} - \gamma \nabla_w H(w_{t-1}) $$ where $\gamma$ is the step size. It turns out to be much faster to sample the data and update $w$ based on the gradient evaluated on just that sample. This is fine, but it is prone to high variance. A much better idea is to sample a minibatch of data of size $M$ and use the average gradient on that for the update so $$ w_t = w_{t-1} - \gamma \left(\lambda\frac{M}{N}\nabla_w\mathcal{R}(w) + \sum_{i=1}^M \nabla_wJ(w;x_i,y_i)\right) $$ Note that we have had to scale the regulariser by $\frac{M}{N}$, since we are evaluating the gradient of the loss on only $M$ points, instead of all $N$ (I think this was the main part of the question). Furthermore, instead of sampling the data, people tend to simply cycle through it sequentially, because then you can guarantee it is all 'seen'.

One final intriguing comment. In one run through the data, we have $\frac{N}{M}$ minibatches. We can share the regulariser across those minibatches in any way we like as long as the sum of all the shared parts equals $\nabla_w \mathcal{R}(w)$ for instance in this paper the authors advocate multiplying the regulariser gradient by $\frac{2^{M-i}}{2^M - 1}$ for the $i^{\text{th}}$ minibatch in the first run through the data.