Is it possible to use stochastic gradient descent at the beginning, then switch to batch gradient descent with only a few training examples?

Question

Batch gradient descent is extremely slow for large datasets, but it can find the lowest possible value for the cost function. Stochastic gradient descent is relatively fast, but it kind of finds the general area where convergence happens and it kind of oscillates around that area.

Is it possible to use stochastic gradient descent at the beginning and find the way to a general convergence and then use batch gradient descent on only a few training examples out of the huge dataset to get even closer to the exact point of convergence?

I know that a model with a cost function that's a bit away from the lowest value for the cost function performs well in stochastic gradient descent, but assuming you want better results, will this work well?

Answer 1

There is a trade-off between the:

• Memory capacity of computation device
• Quality of gradient approximation
• Generalization ability of the network

Memory capacity

I would say, that it is possible to process the whole dataset at once only for small enough dataset and image resolution (or any other measure of the data sample size - text sequence length, number of points in point cloud, whatever). Indeed, one can compute the update by accumulation several gradients (this functionality is provided in Pytorch Lightning) and only then update the parameters, but this would be rather slow.

For dataset of ImageNet size you will need to wait several hours (on single of few GPU's) to make a single update, which is prohibitive.

Usually, one is tempted to take large batches in order to traverse the training dataset as fast as possible, due to the large parallelization ability of modern computation accelerators (GPUs, TPUs, e.t.c). Typical batch size is of order 1k-2k on ImageNet.

Generalization

There is quite a lot of research devoted to the optimal choice of the batch size. A smaller batch size is said to be beneficial in the initial stage of training since it allows to find better and wider optima. Wider optima lead to more robust behavior of the training procedure and have less train-test error discrepancy, as mentioned, for instance, in this paper.

In the consequent stage of training, when one is close to the optimum, the typical strategy is to decrease the amount of "noise" in SGD. This objective can be achieved in several ways:

• Decrease of the learning rate (making smaller steps)
• Increase of the batch size

Assuming that the individual gradients are uncorrelated (which may be a good or bad assumption), the standard deviation of the mini-batch gradients will scale roughly as: $$ \mathcal{O}\left(\frac{1}{\sqrt{N}}\right) $$

The answer

Provided your computational resource allow, or you can allow for the accumulation of a large number of batches - you can try to update the weights, based on the gradients on the whole dataset. But the learning rate decay is a simpler strategy.

One step on the dataset with 1M samples with learning rate $\eta$ is roughly equivalent to 1K steps with batch size 1K and learning rate $\eta / 1000$.