# Feature Learning Beyond the Edge of Stability

arXiv: https://arxiv.org/abs/2502.13110 · HTML: https://mlsanity.com/paper/2502.13110/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 12 of its 14 references that are in the map.
- First posted on arXiv 18 Feb 2025.

## Builds on (most cited first, 12 of 12)

- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [The large learning rate phase of deep learning: the catapult mechanism](https://mlsanity.com/paper/2003.02218/) (2020)
- [Towards Explaining the Regularization Effect of Initial Large Learning Rate in Training Neural Networks](https://mlsanity.com/paper/1907.04595/) (2019)
- [Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability](https://mlsanity.com/paper/2209.15594/) (2022)
- [Understanding Gradient Descent on Edge of Stability in Deep Learning](https://mlsanity.com/paper/2205.09745/) (2022)
- [On the Impact of the Activation Function on Deep Neural Networks Training](https://mlsanity.com/paper/1902.06853/) (2019)
- [Understanding Optimization in Deep Learning with Central Flows](https://mlsanity.com/paper/2410.24206/) (2024)
- [The Implicit Regularization of Dynamical Stability in Stochastic Gradient Descent](https://mlsanity.com/paper/2305.17490/) (2023)
- [Catapults in SGD: spikes in the training loss and their impact on generalization through feature learning](https://mlsanity.com/paper/2306.04815/) (2023)
- [Edge of Stochastic Stability: Revisiting the Edge of Stability for SGD](https://mlsanity.com/paper/2412.20553/) (2024)
- [Super Consistency of Neural Network Landscapes and Learning Rate Transfer](https://mlsanity.com/paper/2402.17457/) (2024)
- [Special Properties of Gradient Descent with Large Learning Rates](https://mlsanity.com/paper/2205.15142/) (2022)

## Abstract (arXiv)

We propose a homogeneous multilayer perceptron parameterization with polynomial hidden layer width pattern and analyze its training dynamics under stochastic gradient descent with depthwise gradient scaling in a general supervised learning scenario. We obtain formulas for the first three Taylor coefficients of the minibatch loss during training that illuminate the connection between sharpness and feature learning, providing in particular a soft rank variant that quantifies the quality of learned hidden layer features. Based on our theory, we design a gradient scaling scheme that in tandem with a quadratic width pattern enables training beyond the edge of stability without loss explosions or numerical errors, resulting in improved feature learning and implicit sharpness regularization as demonstrated empirically.