Mini-batch Noise Lowers Sharpness via Dominant-Subspace Fluctuations

Jun-Ho So, Dongwook Shin

Keywords: sharpness, mini-batch noise, training dynamics, SGD, Hessian, implicit regularization, neural network optimization

TL;DR: Mini-batch noise lowers sharpness through fluctuations in the dominant Hessian subspace

Abstract:

During SGD training, the gradients often align strongly with the dominant subspace spanned by the top-k eigenvectors of the Hessian of the loss. While this seems to naturally imply that loss reduction mainly occurs within this space, prior work has shown that updates within this dominant subspace make no meaningful progress in reducing the loss.
In this work, we argue that the dominant subspace is better understood not as the main space for loss reduction, but as a key subspace for explaining the sharpness dynamics of mini-batch SGD. To explain the role of the dominant subspace in reducing top-k sharpness, we show how the averaged gradient over fluctuations in the dominant directions produces a sharpness correction term, and derive a sharpness correction term induced by mini-batch noise in the dominant directions. Experimental results show that adding the derived correction term to GD brings the sharpness evolution of GD closer to that of SGD.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 110