The Implicit Bias of Gradient Noise: A Symmetry Perspective
Liu Ziyin Mingze Wang Lei Wu
We characterize the learning dynamics of stochastic gradient descent (SGD)
when continuous symmetry exists in the loss function, where the divergence
between SGD and gradient descent is dramatic. We show that depending on how the
symmetry affects the learning dynamics, we can divide a family of symmetry into
two classes. For one class of symmetry, SGD naturally converges to solutions
that have a balanced and aligned gradient noise. For the other class of
symmetry, SGD will almost always diverge. Then, we show that our result remains
applicable and can help us understand the training dynamics even when the
symmetry is not present in the loss function. Our main result is universal in
the sense that it only depends on the existence of the symmetry and is
independent of the details of the loss function. We demonstrate that the
proposed theory offers an explanation of progressive sharpening and flattening
and can be applied to common practical problems such as representation
normalization, matrix factorization, and the use of warmup.