Fisher SAM: Information Geometry and Sharpness Aware Minimisation
Recent sharpness-aware minimisation (SAM) is known to find flat minima which
is beneficial for better generalisation with improved robustness. SAM
essentially modifies the loss function by reporting the maximum loss value
within the small neighborhood around the current iterate. However, it uses the
Euclidean ball to define the neighborhood, which can be inaccurate since loss
functions for neural networks are typically defined over probability
distributions (e.g., class predictive probabilities), rendering the parameter
space non Euclidean. In this paper we consider the information geometry of the
model parameter space when defining the neighborhood, namely replacing SAM's
Euclidean balls with ellipsoids induced by the Fisher information. Our
approach, dubbed Fisher SAM, defines more accurate neighborhood structures that
conform to the intrinsic metric of the underlying statistical manifold. For
instance, SAM may probe the worst-case loss value at either a too nearby or
inappropriately distant point due to the ignorance of the parameter space
geometry, which is avoided by our Fisher SAM. Another recent Adaptive SAM
approach stretches/shrinks the Euclidean ball in accordance with the scale of
the parameter magnitudes. This might be dangerous, potentially destroying the
neighborhood structure. We demonstrate improved performance of the proposed
Fisher SAM on several benchmark datasets/tasks.