Intra-Class Uncertainty Loss Function for Classification

TL;DR

This paper introduces an intra-class uncertainty loss function modeled by Gaussian distributions to improve classification accuracy, especially in unbalanced datasets, by capturing class variability and reducing feature imbalance.

Contribution

The paper proposes a novel loss function that models intra-class uncertainty with Gaussian distributions, enhancing class representation and classification performance.

Findings

Improved accuracy on MNIST, CIFAR, ImageNet, and Long-tailed CIFAR datasets.

Effective modeling of intra-class variability with Gaussian distributions.

Enhanced class compactness and reduced feature imbalance.

Abstract

Most classification models can be considered as the process of matching templates. However, when intra-class uncertainty/variability is not considered, especially for datasets containing unbalanced classes, this may lead to classification errors. To address this issue, we propose a loss function with intra-class uncertainty following Gaussian distribution. Specifically, in our framework, the features extracted by deep networks of each class are characterized by independent Gaussian distribution. The parameters of distribution are learned with a likelihood regularization along with other network parameters. The means of the Gaussian play a similar role as the center anchor in existing methods, and the variance describes the uncertainty of different classes. In addition, similar to the inter-class margin in traditional loss functions, we introduce a margin to intra-class uncertainty to make each cluster more compact and reduce the imbalance of feature distribution from different categories. Based on MNIST, CIFAR, ImageNet, and Long-tailed CIFAR analyses, the proposed approach shows improved classification performance, through learning a better class representation.