Neural Collapse Inspired Attraction-Repulsion-Balanced Loss for Imbalanced Learning

TL;DR

This paper introduces ARB-Loss, a novel loss function inspired by neural collapse, designed to address class imbalance in deep learning by balancing gradient components, leading to improved accuracy on imbalanced datasets.

Contribution

The paper proposes a new Attraction-Repulsion-Balanced Loss that effectively balances gradient components to improve classification in imbalanced learning scenarios.

Findings

Achieves state-of-the-art performance on large-scale classification datasets.

Enables one-stage training to outperform traditional multi-stage methods.

Effectively handles extreme class imbalance in classification and segmentation tasks.

Abstract

Class imbalance distribution widely exists in real-world engineering. However, the mainstream optimization algorithms that seek to minimize error will trap the deep learning model in sub-optimums when facing extreme class imbalance. It seriously harms the classification precision, especially on the minor classes. The essential reason is that the gradients of the classifier weights are imbalanced among the components from different classes. In this paper, we propose Attraction-Repulsion-Balanced Loss (ARB-Loss) to balance the different components of the gradients. We perform experiments on the large-scale classification and segmentation datasets and our ARB-Loss can achieve state-of-the-art performance via only one-stage training instead of 2-stage learning like nowadays SOTA works.