The Effect of SGD Batch Size on Autoencoder Learning: Sparsity, Sharpness, and Feature Learning

TL;DR

This paper explores how batch size in SGD affects autoencoder training, revealing that smaller batches lead to sparse, feature-selective solutions, while full batches produce dense, less feature-focused minima, with implications for understanding generalization.

Contribution

The study provides a detailed analysis of SGD dynamics on autoencoders, demonstrating how batch size influences sparsity, feature learning, and the sharpness of minima, supported by new convergence proof techniques.

Findings

Smaller batch sizes induce sparsity and feature selection.

Full batch gradient descent finds dense, less feature-specific minima.

Minima from full batch are flatter than those from smaller batches.

Abstract

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.