Rethinking Sharpness-Aware Minimization as Variational Inference

Understand

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.