(S)GD over Diagonal Linear Networks: Implicit Regularisation, Large Stepsizes and Edge of Stability
This work addresses optimization dynamics in overparameterized models, providing insights into stepsize selection for practitioners, but it is incremental as it builds on existing regularization theories.
The paper investigates how stochasticity and large stepsizes affect the implicit regularization of gradient descent and stochastic gradient descent in diagonal linear networks, showing that large stepsizes consistently benefit SGD for sparse regression but can hinder sparse solution recovery for GD, especially near the divergence threshold.
In this paper, we investigate the impact of stochasticity and large stepsizes on the implicit regularisation of gradient descent (GD) and stochastic gradient descent (SGD) over diagonal linear networks. We prove the convergence of GD and SGD with macroscopic stepsizes in an overparametrised regression setting and characterise their solutions through an implicit regularisation problem. Our crisp characterisation leads to qualitative insights about the impact of stochasticity and stepsizes on the recovered solution. Specifically, we show that large stepsizes consistently benefit SGD for sparse regression problems, while they can hinder the recovery of sparse solutions for GD. These effects are magnified for stepsizes in a tight window just below the divergence threshold, in the "edge of stability" regime. Our findings are supported by experimental results.