Semi-Discrete Optimal Transport for Long-Tailed Classification

The long-tailed data distribution poses an enormous challenge for training neural networks in classification. A classification network can be decoupled into a feature extractor and a classifier. This paper takes a semi-discrete optimal transport (OT) perspective to analyze the long-tailed classification problem, where the feature space is viewed as a continuous source domain, and the classifier weights are viewed as a discrete target domain. The classifier is indeed to find a cell decomposition of the feature space with each cell corresponding to one class. An imbalanced training set causes the more frequent classes to have larger volume cells, which means that the classifier’s decision boundary is biased towards less frequent classes, resulting in reduced classification performance in the inference phase. Therefore, we propose a novel OT-dynamic softmax loss, which dynamically adjusts the decision boundary in the training phase to avoid overfitting in the tail classes. In addition, our method incorporates the supervised contrastive loss so that the feature space can satisfy the uniform distribution condition. Extensive and comprehensive experiments demonstrate that our method achieves state-of-the-art performance on multiple long-tailed recognition benchmarks, including CIFAR-LT, ImageNet-LT, iNaturalist 2018, and Places-LT.