Enhancing Sharpness-Aware Optimization Through Variance Suppression

TL;DR

This paper introduces VaSSO, a variance suppression technique that stabilizes sharpness-aware optimization, improving generalization and robustness of deep neural networks beyond existing methods like SAM.

Contribution

It proposes a novel variance suppression approach to enhance the stability and effectiveness of sharpness-aware minimization in deep learning.

Findings

VaSSO improves model generalization over SAM.

VaSSO enhances robustness against label noise.

Experiments show numerical improvements in image classification and translation.

Abstract

Sharpness-aware minimization (SAM) has well documented merits in enhancing generalization of deep neural networks, even without sizable data augmentation. Embracing the geometry of the loss function, where neighborhoods of 'flat minima' heighten generalization ability, SAM seeks 'flat valleys' by minimizing the maximum loss caused by an adversary perturbing parameters within the neighborhood. Although critical to account for sharpness of the loss function, such an 'over-friendly adversary' can curtail the outmost level of generalization. The novel approach of this contribution fosters stabilization of adversaries through variance suppression (VaSSO) to avoid such friendliness. VaSSO's provable stability safeguards its numerical improvement over SAM in model-agnostic tasks, including image classification and machine translation. In addition, experiments confirm that VaSSO endows SAM with robustness against high levels of label noise.