The Effect of SGD Batch Size on Autoencoder Learning: Sparsity, Sharpness, and Feature Learning
This work addresses the impact of batch size on optimization dynamics and feature learning in non-convex autoencoders, providing theoretical insights for machine learning practitioners, though it is incremental as it builds on existing SGD analysis.
The paper investigates how SGD batch size affects the training of a single-neuron autoencoder on orthogonal data, finding that full-batch gradient descent yields dense solutions with little feature learning, while smaller batch sizes lead to sparse solutions and feature selection, with full-batch minima being flatter contrary to prior beliefs.
In this work, we investigate the dynamics of stochastic gradient descent (SGD) when training a single-neuron autoencoder with linear or ReLU activation on orthogonal data. We show that for this non-convex problem, randomly initialized SGD with a constant step size successfully finds a global minimum for any batch size choice. However, the particular global minimum found depends upon the batch size. In the full-batch setting, we show that the solution is dense (i.e., not sparse) and is highly aligned with its initialized direction, showing that relatively little feature learning occurs. On the other hand, for any batch size strictly smaller than the number of samples, SGD finds a global minimum which is sparse and nearly orthogonal to its initialization, showing that the randomness of stochastic gradients induces a qualitatively different type of "feature selection" in this setting. Moreover, if we measure the sharpness of the minimum by the trace of the Hessian, the minima found with full batch gradient descent are flatter than those found with strictly smaller batch sizes, in contrast to previous works which suggest that large batches lead to sharper minima. To prove convergence of SGD with a constant step size, we introduce a powerful tool from the theory of non-homogeneous random walks which may be of independent interest.