[Paper Review] SGD with Large Step Sizes Learns Sparse Features

Maksym Andriushchenko, Aditya Varre+2 more|arXiv (Cornell University)|Oct 11, 2022

Stochastic Gradient Optimization Techniques4 citations

TL;DR

This paper demonstrates that large step sizes in SGD induce a loss stabilization phase that drives implicit bias toward sparse feature representations in neural networks, even without explicit regularization. Through empirical and theoretical analysis, it shows that prolonged stabilization amplifies noise-driven dynamics, enabling SGD to learn sparse features via an SDE model linked to neural tangent kernel features.

ABSTRACT

We showcase important features of the dynamics of the Stochastic Gradient Descent (SGD) in the training of neural networks. We present empirical observations that commonly used large step sizes (i) lead the iterates to jump from one side of a valley to the other causing loss stabilization, and (ii) this stabilization induces a hidden stochastic dynamics orthogonal to the bouncing directions that biases it implicitly toward sparse predictors. Furthermore, we show empirically that the longer large step sizes keep SGD high in the loss landscape valleys, the better the implicit regularization can operate and find sparse representations. Notably, no explicit regularization is used so that the regularization effect comes solely from the SGD training dynamics influenced by the step size schedule. Therefore, these observations unveil how, through the step size schedules, both gradient and noise drive together the SGD dynamics through the loss landscape of neural networks. We justify these findings theoretically through the study of simple neural network models as well as qualitative arguments inspired from stochastic processes. Finally, this analysis allows us to shed a new light on some common practice and observed phenomena when training neural networks. The code of our experiments is available at https://github.com/tml-epfl/sgd-sparse-features.

Motivation & Objective

• Understand the implicit regularization effects of large step sizes in SGD during training of deep neural networks.
• Investigate why loss stabilization—commonly observed in practice—leads to improved generalization despite no further loss reduction.
• Characterize the hidden stochastic dynamics during the stabilization phase that favor sparse feature learning.
• Link the observed behavior to stochastic differential equation (SDE) models with multiplicative noise tied to neural tangent kernel features.
• Provide theoretical and empirical justification for how large step sizes enable transition from lazy training to rich feature learning regimes.

Proposed method

• Empirically analyze training dynamics of ResNet-18 on CIFAR-10 using piecewise constant step size schedules with and without decay.
• Observe loss stabilization patterns where training loss remains approximately constant despite no further descent.
• Model the slow dynamics during stabilization using a stochastic differential equation (SDE) with multiplicative noise proportional to the neural tangent kernel (NTK) features.
• Validate the SDE model by comparing predicted feature sparsity with actual network behavior across architectures (ResNet-18, ResNet-34) on CIFAR-10 and CIFAR-100.
• Use a two-layer teacher-student setup to confirm that small step sizes prevent feature learning, while large step sizes enable it.
• Analyze the role of batch size by showing that replacing step size with effective step size scaled by batch size $ B $ preserves the dynamics ($ \gamma \leftarrow \gamma / B $).

Experimental results

Research questions

• RQ1 How do large step sizes in SGD influence the dynamics of training beyond loss descent?
• RQ2 What hidden stochastic dynamics occur during the loss stabilization phase that lead to improved generalization?
• RQ3 Can the implicit bias toward sparse features during stabilization be modeled via an SDE with NTK-dependent noise?
• RQ4 To what extent does the duration of the stabilization phase affect the sparsity of learned features and generalization performance?
• RQ5 How do common practices like weight decay, BatchNorm, and SAM relate to the implicit regularization induced by large step sizes?

Key findings

• Large step sizes induce a loss stabilization phase where training loss remains approximately constant, yet generalization improves with longer duration.
• During stabilization, SGD exhibits a slow drift driven by the interplay of gradient and non-vanishing multiplicative noise, which biases the model toward sparse feature representations.
• The effective dynamics during stabilization are well-modeled by a stochastic differential equation (SDE) whose noise term is proportional to the neural tangent kernel (NTK) features.
• Empirical results confirm that increasing the duration of the stabilization phase leads to higher feature sparsity and lower test error, even without explicit regularization.
• The sparsity of features correlates strongly with improved generalization, with test error differences of up to 12–35% depending on step size schedule.
• Normalization techniques like BatchNorm and weight decay can be interpreted as enabling larger step sizes that enhance the implicit bias toward sparsity, preventing divergence while amplifying regularization.