Training Diagonal Linear Networks with Stochastic Sharpness-Aware Minimization

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Clara, Gabriel, Langer, Sophie, Schmidt-Hieber, Johannes
Format: Preprint
Veröffentlicht: 2025
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866929761460158464
author Clara, Gabriel
Langer, Sophie
Schmidt-Hieber, Johannes
author_facet Clara, Gabriel
Langer, Sophie
Schmidt-Hieber, Johannes
contents We analyze the landscape and training dynamics of diagonal linear networks in a linear regression task, with the network parameters being perturbed by small isotropic normal noise. The addition of such noise may be interpreted as a stochastic form of sharpness-aware minimization (SAM) and we prove several results that relate its action on the underlying landscape and training dynamics to the sharpness of the loss. In particular, the noise changes the expected gradient to force balancing of the weight matrices at a fast rate along the descent trajectory. In the diagonal linear model, we show that this equates to minimizing the average sharpness, as well as the trace of the Hessian matrix, among all possible factorizations of the same matrix. Further, the noise forces the gradient descent iterates towards a shrinkage-thresholding of the underlying true parameter, with the noise level explicitly regulating both the shrinkage factor and the threshold.
format Preprint
id arxiv_https___arxiv_org_abs_2503_11891
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Training Diagonal Linear Networks with Stochastic Sharpness-Aware Minimization
Clara, Gabriel
Langer, Sophie
Schmidt-Hieber, Johannes
Machine Learning
Statistics Theory
We analyze the landscape and training dynamics of diagonal linear networks in a linear regression task, with the network parameters being perturbed by small isotropic normal noise. The addition of such noise may be interpreted as a stochastic form of sharpness-aware minimization (SAM) and we prove several results that relate its action on the underlying landscape and training dynamics to the sharpness of the loss. In particular, the noise changes the expected gradient to force balancing of the weight matrices at a fast rate along the descent trajectory. In the diagonal linear model, we show that this equates to minimizing the average sharpness, as well as the trace of the Hessian matrix, among all possible factorizations of the same matrix. Further, the noise forces the gradient descent iterates towards a shrinkage-thresholding of the underlying true parameter, with the noise level explicitly regulating both the shrinkage factor and the threshold.
title Training Diagonal Linear Networks with Stochastic Sharpness-Aware Minimization
topic Machine Learning
Statistics Theory
url https://arxiv.org/abs/2503.11891