Feature Learning Beyond the Edge of Stability

TL;DR

This paper introduces a new neural network training method that leverages a polynomial width pattern and depthwise gradient scaling to enhance feature learning and stability beyond traditional limits, supported by theoretical formulas and empirical results.

Contribution

It presents a novel homogeneous multilayer perceptron parameterization with a specific width pattern and gradient scaling scheme, enabling stable training beyond the edge of stability.

Findings

Improved feature learning demonstrated empirically.

Gradient scaling scheme enables stable training beyond stability edge.

Theoretical formulas connect sharpness and feature quality.

Abstract

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.