Understanding Imbalanced Semantic Segmentation Through Neural Collapse

TL;DR

This paper investigates the neural collapse phenomenon in semantic segmentation, identifies how class imbalance affects this structure, and proposes a regularizer to improve segmentation performance, achieving state-of-the-art results.

Contribution

It introduces a regularizer based on neural collapse principles to enhance feature learning in imbalanced semantic segmentation tasks.

Findings

Significant performance improvements on 2D and 3D benchmarks.

Achieved 1st place with a +6.8% mIoU on ScanNet200.

Neural collapse structures are disrupted by class imbalance in segmentation.

Abstract

A recent study has shown a phenomenon called neural collapse in that the within-class means of features and the classifier weight vectors converge to the vertices of a simplex equiangular tight frame at the terminal phase of training for classification. In this paper, we explore the corresponding structures of the last-layer feature centers and classifiers in semantic segmentation. Based on our empirical and theoretical analysis, we point out that semantic segmentation naturally brings contextual correlation and imbalanced distribution among classes, which breaks the equiangular and maximally separated structure of neural collapse for both feature centers and classifiers. However, such a symmetric structure is beneficial to discrimination for the minor classes. To preserve these advantages, we introduce a regularizer on feature centers to encourage the network to learn features closer to the appealing structure in imbalanced semantic segmentation. Experimental results show that our method can bring significant improvements on both 2D and 3D semantic segmentation benchmarks. Moreover, our method ranks 1st and sets a new record (+6.8% mIoU) on the ScanNet200 test leaderboard. Code will be available at https://github.com/dvlab-research/Imbalanced-Learning.