SSE-SAM: Balancing Head and Tail Classes Gradually through Stage-Wise SAM

TL;DR

This paper introduces SSE-SAM, a stage-wise method that balances head and tail class learning in long-tailed datasets by combining the strengths of SAM and ImbSAM to improve generalization and saddle point escape.

Contribution

The paper proposes SSE-SAM, a novel stage-wise approach that effectively balances head and tail class training by leveraging complementary strengths of SAM and ImbSAM.

Findings

SSE-SAM outperforms existing methods in escaping saddle points for both head and tail classes.

SSE-SAM achieves improved classification performance on long-tailed datasets.

The staged approach effectively balances class representation and enhances generalization.

Abstract

Real-world datasets often exhibit a long-tailed distribution, where vast majority of classes known as tail classes have only few samples. Traditional methods tend to overfit on these tail classes. Recently, a new approach called Imbalanced SAM (ImbSAM) is proposed to leverage the generalization benefits of Sharpness-Aware Minimization (SAM) for long-tailed distributions. The main strategy is to merely enhance the smoothness of the loss function for tail classes. However, we argue that improving generalization in long-tail scenarios requires a careful balance between head and tail classes. We show that neither SAM nor ImbSAM alone can fully achieve this balance. For SAM, we prove that although it enhances the model's generalization ability by escaping saddle point in the overall loss landscape, it does not effectively address this for tail-class losses. Conversely, while ImbSAM is more effective at avoiding saddle points in tail classes, the head classes are trained insufficiently, resulting in significant performance drops. Based on these insights, we propose Stage-wise Saddle Escaping SAM (SSE-SAM), which uses complementary strengths of ImbSAM and SAM in a phased approach. Initially, SSE-SAM follows the majority sample to avoid saddle points of the head-class loss. During the later phase, it focuses on tail-classes to help them escape saddle points. Our experiments confirm that SSE-SAM has better ability in escaping saddles both on head and tail classes, and shows performance improvements.