Balanced Sharpness-Aware Minimization for Imbalanced Regression

TL;DR

This paper introduces Balanced Sharpness-Aware Minimization (BSAM), a novel method to improve regression models' generalization in imbalanced data scenarios, achieving better performance on tasks like age and depth estimation.

Contribution

The paper proposes BSAM, a new loss optimization technique that enforces uniform generalization across the observation space in imbalanced regression tasks, with theoretical guarantees.

Findings

BSAM outperforms existing methods on multiple vision regression tasks.

The approach provides a theoretical generalization bound.

Extensive experiments validate the effectiveness of BSAM.

Abstract

Regression is fundamental in computer vision and is widely used in various tasks including age estimation, depth estimation, target localization, \etc However, real-world data often exhibits imbalanced distribution, making regression models perform poorly especially for target values with rare observations~(known as the imbalanced regression problem). In this paper, we reframe imbalanced regression as an imbalanced generalization problem. To tackle that, we look into the loss sharpness property for measuring the generalization ability of regression models in the observation space. Namely, given a certain perturbation on the model parameters, we check how model performance changes according to the loss values of different target observations. We propose a simple yet effective approach called Balanced Sharpness-Aware Minimization~(BSAM) to enforce the uniform generalization ability of regression models for the entire observation space. In particular, we start from the traditional sharpness-aware minimization and then introduce a novel targeted reweighting strategy to homogenize the generalization ability across the observation space, which guarantees a theoretical generalization bound. Extensive experiments on multiple vision regression tasks, including age and depth estimation, demonstrate that our BSAM method consistently outperforms existing approaches. The code is available \href{https://github.com/manmanjun/BSAM_for_Imbalanced_Regression}{here}.