When Do Flat Minima Optimizers Work?

TL;DR

This work tackles the question of when flat-minima optimizers improve generalization by performing a rigorous, cross-domain comparison of SWA and SAM, including a geometric analysis of the resulting minima. It combines loss-landscape visualizations, Hessian-based flatness metrics, and a broad benchmark across CV, NLP, and GRL tasks, introducing WASAM as an effective amalgamation. The study uncovers task- and architecture-dependent effects, with NLP favoring SAM, GRL leaning toward SWA, and WASAM delivering robust gains across many settings, while also highlighting failure modes and saddle-point tendencies. The findings offer practical guidance for optimizer selection and lay groundwork for designing methods that align flatness with generalization, accompanied by open-source code for reproducibility.

Abstract

Recently, flat-minima optimizers, which seek to find parameters in low-loss neighborhoods, have been shown to improve a neural network's generalization performance over stochastic and adaptive gradient-based optimizers. Two methods have received significant attention due to their scalability: 1. Stochastic Weight Averaging (SWA), and 2. Sharpness-Aware Minimization (SAM). However, there has been limited investigation into their properties and no systematic benchmarking of them across different domains. We fill this gap here by comparing the loss surfaces of the models trained with each method and through broad benchmarking across computer vision, natural language processing, and graph representation learning tasks. We discover several surprising findings from these results, which we hope will help researchers further improve deep learning optimizers, and practitioners identify the right optimizer for their problem.

Figures (12)

• Figure 1: The mechanics behind SWA and SAM, whose solution is denoted by $\bm{+}$ and $\bm{\times}$, respectively. SWA produces a solution $\theta$ that is pulled towards flatter regions, while SAM approximates sharpness within the parameters' neighborhood (arrows).
• Figure 2: Training (blue) and test (red) losses (---) and accuracies (.4pt) of linear interpolations $\boldsymbol{\theta}(\alpha)=(1-\alpha) \boldsymbol{\theta}+\alpha \boldsymbol{\theta}^{\prime}$ (for $\alpha \in [-1,1.5]$) between SWA ($\bm{+}$) and SAM ($\bm{\times}$) solutions ($\alpha=0.0$) and non-flat baseline solutions ($\bullet, \alpha=1.0$).
• Figure 3: Training (blue) / test (red) losses (---) / accuracies (.4pt) between non-flat baseline ($\bullet$) $\leftrightarrow$ SWA ($\bm{+}$), SAM ($\bm{\times}$) $\leftrightarrow$ WASAM ($\star$).
• Figure 4: (a) NLP test results: Open-Domain Question Answering and Natural Language Understanding (GLUE) including paraphrase, sentiment analysis, and textual entailment. (b) GRL test results: Node Property Prediction (NPP), Graph Property Prediction (GPP), Link Property Prediction (LPP).
• Figure 5: GraphSAGE on OGB-Proteins: Adam's ($\bullet$) solution performs about equally well as SAM ($\bm{\times}$), and better than SWA ($\bm{+}$).