This project was submitted to the ICLR Reproducibility Challenge 2019. The paper we reproduced is “A Resizable Mini-batch Gradient Descent based on a Multi-Armed Bandit”. Its aim is to combine SGD and Multi-Armed Bandit to chose the size of the mini-batch used. Throughout the optimization procedure, the current estimate oscilates between zones where the loss surface is rather flat, and zones where the loss surface is much more complicated. This adaptive size of the mini-batch combines fast descent rates when possible and variance reduction when needed. This is the Exploration-Exploitation Dilemma. Below, you can see a heamap showing the chosen batch size at each epoch for MNIST optimization.
It is beautiful to see that large batch sizes are chosen at the end of the optimization since variance reduction is needed approaching the local minimum.