When Do Flat Minima Optimizers Work? (2202.00661v5)

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.

• The paper compares flat minima optimizers SWA and SAM, revealing task-specific strengths and weaknesses through comprehensive empirical evaluation.
• It employs linear interpolation in weight space to analyze basin connectivity and introduces the novel WASAM method that consistently outperforms individual optimizers.
• The study highlights that optimizer performance is strongly influenced by the choice of tasks, network architecture, and dataset characteristics.

When Do Flat Minima Optimizers Work?

The paper "When Do Flat Minima Optimizers Work?" (2202.00661) focuses on the analysis and empirical evaluation of flat-minima optimizers, specifically the Stochastic Weight Averaging (SWA) and Sharpness-Aware Minimization (SAM) methods. It provides a comprehensive investigation into their effectiveness across various tasks, shedding light on the circumstances where these optimizers might excel or underperform.

Introduction and Background

Flat-minima optimizers have gained attention due to their potential to enhance generalization in neural networks by locating parameter spaces with large neighborhoods yielding low loss values. SWA and SAM are the primary methods examined. SWA averages weights over iterations to converge towards flat regions, while SAM explicitly minimizes sharpness by optimizing maximum loss within a local neighborhood. Despite widespread use, there's been limited comprehensive benchmarking to understand their comparative performance and underlying properties.

Empirical Investigation

The paper conducts a broad benchmarking paper across domains such as computer vision (CV), NLP, and graph representation learning (GRL). It employs linear interpolation in weight space to understand the landscape between non-flat and flat minima, revealing differences in basin connectivity and susceptibility to sharp directions. Notably, it introduces Weight-Averaged SAM (WASAM) by averaging SAM iterates, which shows consistent improvement over SWA or SAM alone across multiple tasks.

(Figure 5) *Figure 5: GraphSAGE on OGB-Proteins: Adam's (bullet) solution performs about equally well as SAM (%%%%0%%%%), and better than SWA (%%%%1%%%%).*

Key Findings

1. Task-Specific Performance: SWA tends to underperform in NLP tasks compared to SAM, which improves generalization significantly in several instances.
2. Architectural Influence: Different architectures exhibit varied benefits from flat-minima optimizers; for instance, SWA is less effective with Transformers.
3. Dataset Impact: The efficacy of these optimizers is heavily influenced by datasets, with SAM particularly excelling in NLP and failing in specific GRL tasks.
4. WASAM Advantage: WASAM frequently outperforms individual optimizers by combining the robustness of SWA with the sharper minimization of SAM.
5. Linear Interpolation Insights: Interpolation paths show that while SWA and SAM often lie in different basins, SAM explores hidden sharp minima compared to SWA.

(Figure 10) *Figure 10: WRN-28-10: Changing SAM's rho result in different basins.*

Discussion of Failures

The investigation identifies scenarios where neither SWA nor SAM outperforms standard baselines, emphasizing the misalignment between training and test loss surfaces. For instance, node property prediction tasks on the OGB-Proteins dataset revealed that SAM doesn't necessarily achieve lower generalization error despite flatter training loss landscapes.

Conclusion

The analysis establishes that while flat minima optimizers like SWA and SAM often enhance generalization, their effectiveness is nuanced and task-dependent. Future work could focus on improving optimizer robustness across different architectures and datasets, as well as exploring dynamic adaptations of hyperparameters like SAM's neighborhood size $\rho$.

Through extensive experimental evaluation, the paper provides practical insights into the conditions under which flat-minima optimization techniques are advantageous, while also highlighting the need for continued development and nuanced application of these methods in diverse machine learning tasks.