Rethinking Sharpness-Aware Minimization as Variational Inference

TL;DR

This paper connects sharpness-aware minimization (SAM) with mean-field variational inference (MFVI), showing they share similar interpretations and proposing algorithms that combine both for improved neural network training.

Contribution

It establishes a theoretical link between SAM and MFVI, and introduces algorithms that interpolate between them, with empirical evaluation on benchmark datasets.

Findings

SAM and MFVI share similar flatness optimization interpretations

Proposed algorithms outperform or match existing methods on benchmarks

SAM-like updates can serve as a drop-in replacement for the reparametrisation trick

Abstract

Sharpness-aware minimization (SAM) aims to improve the generalisation of gradient-based learning by seeking out flat minima. In this work, we establish connections between SAM and Mean-Field Variational Inference (MFVI) of neural network parameters. We show that both these methods have interpretations as optimizing notions of flatness, and when using the reparametrisation trick, they both boil down to calculating the gradient at a perturbed version of the current mean parameter. This thinking motivates our study of algorithms that combine or interpolate between SAM and MFVI. We evaluate the proposed variational algorithms on several benchmark datasets, and compare their performance to variants of SAM. Taking a broader perspective, our work suggests that SAM-like updates can be used as a drop-in replacement for the reparametrisation trick.