From Information to Generative Exponent: Learning Rate Induces Phase Transitions in SGD
Konstantinos Tsiolis, Alireza Mousavi-Hosseini, Murat Erdogdu
To understand feature learning dynamics in neural networks, recent theoretical works have focused on gradient-based learning of Gaussian single-index models, where the label is a nonlinear function of a latent one-dimensional projection of the input. While the sample complexity of online SGD is determined by the information exponent of the nonlinear link, recent works improved this by reusing samples or modifying the loss function --- transformations which introduce non-correlational updates --- and instead are limited by the (potentially much smaller) generative exponent. However, this picture is only valid if the learning rate is sufficiently large. In this paper, we characterize the relationship between learning rate and sample complexity for a broad class of gradient-based algorithms that encapsulates both correlational and non-correlational updates, and demonstrate a phase transition from an "information exponent regime" with small learning rate to a "generative exponent regime" with large learning rate. Our framework covers prior analyses of one-pass SGD and SGD with batch reuse, while also introducing a new layer-wise training algorithm that leverages a two-timescales approach to go beyond correlational queries without reusing samples or modifying the loss from squared error. Our theoretical study demonstrates that the choice of learning rate is as important as the design of the algorithm in achieving statistical and computational efficiency.