Feature Learning Beyond the Edge of Stability

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autore principale: Terjék, Dávid
Natura: Preprint
Pubblicazione: 2025
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866908368127393792
author Terjék, Dávid
author_facet Terjék, Dávid
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2502_13110
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Feature Learning Beyond the Edge of Stability
Terjék, Dávid
Machine Learning
68T07
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.
title Feature Learning Beyond the Edge of Stability
topic Machine Learning
68T07
url https://arxiv.org/abs/2502.13110