On the Training Instability of Shuffling SGD with Batch Normalization

TL;DR

This paper investigates how different SGD variants, Single Shuffle and Random Reshuffle, interact with batch normalization, revealing that RR offers more stable training and avoids divergence and distorted optima compared to SS.

Contribution

It provides a theoretical and empirical analysis of the contrasting behaviors of SS and RR with batch normalization, highlighting RR's stability advantages.

Findings

RR leads to more stable training loss evolution than SS.

SS can cause divergence and distorted optima in training.

Empirical validation confirms the theoretical differences in practical settings.

Abstract

We uncover how SGD interacts with batch normalization and can exhibit undesirable training dynamics such as divergence. More precisely, we study how Single Shuffle (SS) and Random Reshuffle (RR) -- two widely used variants of SGD -- interact surprisingly differently in the presence of batch normalization: RR leads to much more stable evolution of training loss than SS. As a concrete example, for regression using a linear network with batch normalization, we prove that SS and RR converge to distinct global optima that are "distorted" away from gradient descent. Thereafter, for classification we characterize conditions under which training divergence for SS and RR can, and cannot occur. We present explicit constructions to show how SS leads to distorted optima in regression and divergence for classification, whereas RR avoids both distortion and divergence. We validate our results by confirming them empirically in realistic settings, and conclude that the separation between SS and RR used with batch normalization is relevant in practice.