Reproduce “Normalization Layers Are All That Sharpness-Aware Minimization Needs” on OpenResearch

Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.

Universal app for Apple silicon and Intel. Requires macOS 11 or later.

About this paper

Maximilian Mueller, Tiffany Vlaar, David Rolnick, Matthias Hein

The paper introduces SAM-OnlyNorm (SAM-ON), an approach that applies Sharpness-Aware Minimization (SAM) exclusively to the affine parameters of normalization layers, achieving better generalization than full SAM with less than 0.1% of parameters perturbed. This method also demonstrates computational efficiency and reveals that SAM's benefits are not solely derived from finding flat minima.

Key insights

  • The affine parameters of normalization layers hold disproportionate importance for SAM, as perturbing only these few parameters consistently yields superior generalization compared to perturbing all parameters.
  • SAM's generalization benefits are not solely dependent on finding flatter minima; SAM-ON models often generalize better despite converging to *sharper* minima than models trained with full SAM.
  • The uniqueness of normalization layers is critical, as random sparse perturbations at equivalent sparsity levels failed to achieve similar performance gains, highlighting a specific interaction.

Summary by alphaXiv.

Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.