Inducing Neural Collapse in Imbalanced Learning: Do We Really Need a Learnable Classifier at the End of Deep Neural Network?

TL;DR

This paper investigates whether fixing a neural network classifier to an ETF structure can induce neural collapse and improve training efficiency in imbalanced datasets, challenging the necessity of learnable classifiers.

Contribution

The study demonstrates that a fixed ETF classifier can induce neural collapse and enhance training in imbalanced data, replacing the need for a learnable classifier and cross entropy loss.

Findings

Fixed ETF classifier induces neural collapse in imbalanced data

Replacing cross entropy with squared loss improves convergence

Method achieves faster training and better performance on multiple datasets

Abstract

Modern deep neural networks for classification usually jointly learn a backbone for representation and a linear classifier to output the logit of each class. A recent study has shown a phenomenon called neural collapse that the within-class means of features and the classifier vectors converge to the vertices of a simplex equiangular tight frame (ETF) at the terminal phase of training on a balanced dataset. Since the ETF geometric structure maximally separates the pair-wise angles of all classes in the classifier, it is natural to raise the question, why do we spend an effort to learn a classifier when we know its optimal geometric structure? In this paper, we study the potential of learning a neural network for classification with the classifier randomly initialized as an ETF and fixed during training. Our analytical work based on the layer-peeled model indicates that the feature learning with a fixed ETF classifier naturally leads to the neural collapse state even when the dataset is imbalanced among classes. We further show that in this case the cross entropy (CE) loss is not necessary and can be replaced by a simple squared loss that shares the same global optimality but enjoys a better convergence property. Our experimental results show that our method is able to bring significant improvements with faster convergence on multiple imbalanced datasets.