A margin-based replacement for cross-entropy loss

TL;DR

Cross-entropy is the de facto loss for classification but can hinder robustness and generalisation. The authors propose High Error Margin (HEM) loss, a margin-based variant that reweights and thresholds per-sample errors to focus on the largest mistakes, and its imbalance-adjusted variant HEM+. Through extensive experiments across 18+ architectures and 8 datasets, HEM/HEM+ generally outperform CE on unknown-class rejection, adversarial robustness, continual learning, and semantic segmentation, while maintaining near-parity on clean accuracy. The results indicate that HEM offers a versatile, general-purpose replacement for CE, with HEM+ providing additional gains for imbalanced and segmentation tasks. This approach reduces over-confident predictions and catastrophic forgetting, enabling safer and more reliable deployment across diverse vision tasks.

Abstract

Cross-entropy (CE) loss is the de-facto standard for training deep neural networks to perform classification. However, CE-trained deep neural networks struggle with robustness and generalisation issues. To alleviate these issues, we propose high error margin (HEM) loss, a variant of multi-class margin loss that overcomes the training issues of other margin-based losses. We evaluate HEM extensively on a range of architectures and datasets. We find that HEM loss is more effective than cross-entropy loss across a wide range of tasks: unknown class rejection, adversarial robustness, learning with imbalanced data, continual learning, and semantic segmentation (a pixel-level classification task). Despite all training hyper-parameters being chosen for CE loss, HEM is inferior to CE only in terms of clean accuracy and this difference is insignificant. We also compare HEM to specialised losses that have previously been proposed to improve performance on specific tasks. LogitNorm, a loss achieving state-of-the-art performance on unknown class rejection, produces similar performance to HEM for this task, but is much poorer for continual learning and semantic segmentation. Logit-adjusted loss, designed for imbalanced data, has superior results to HEM for that task, but performs more poorly on unknown class rejection and semantic segmentation. DICE, a popular loss for semantic segmentation, is inferior to HEM loss on all tasks, including semantic segmentation. Thus, HEM often out-performs specialised losses, and in contrast to them, is a general-purpose replacement for CE loss.

Figures (11)

• Figure 1: Summary results for all the different tasks considered, comparing the average performance of cross-entropy (CE) loss to each of the alternative losses that have been evaluated: LogitNorm (LN), Logit-adjusted (LA), DICE, multi-class margin (MM), high error margin (HEM), and high error margin with adjusted margins (HEM+). Results are averaged using the arithmetic mean over all other factors that were varied in the experiments, such as data-set and network architecture. For all the evaluation metrics used, higher values indicate better performance. The relative performance is calculated by subtracting the performance produced by CE loss from the corresponding metric for each of the other losses. Hence, positive values indicate performance better than that of CE loss. Note that the results for LA are equal to those of CE, and the results for HEM+ are equal to those of HEM when the training data is balanced ( i.e., when using standard training data and performing continual learning).
• Figure 2: Results when learning with standard data-sets and testing with clean and corrupt images. (a) and (b) directly compare the performance produced by HEM and cross-entropy (CE) losses when used to train networks with MNIST, CIFAR10, CIFAR100, TinyImageNet (TIN), and ImageNet1k (IN) using three different network architectures for each data-set. For each data-set the size of the marker used corresponds to the size of the network. Results above the diagonal are conditions where better performance was obtained when training with HEM rather than CE loss. Performance metrics are averaged over multiple trials performed for each condition (data-set and architecture) and the error bars show the standard deviation recorded across the trials in each condition (in the majority of cases these error bars are too small to be visible). (a) Compares the performance of the two losses in terms of the accuracy of classifying the standard test-data. (b) Compares the performance of the two losses in terms of the accuracy of classifying the common-corruptions test-data. (c) Shows results averaged over all the data-sets and network architecture (and multiple trials in each condition) for all relevant losses: cross-entropy (CE), LogitNorm (LN), DICE, multi-class margin (MM) and HEM. Error bars show the mean standard deviation recorded across the trials in each condition. The inset shows the results for CE, LN, and HEM losses plotted on a separate scale to allow the differences between these losses to be visible.
• Figure 3: Results when learning with standard data-sets and testing on unknown and adversarial images. This figure has an identical format to \ref{['fig-standard_training_IND']} except (a) compares the performance of CE and HEM losses in terms of the ability to distinguish samples from known and unknown classes, and (b) compares the performance of CE and HEM losses in terms of the ability to deal correctly with adversarially perturbed samples. In both (a) and (b) Maximum Softmax Probability (MSP) is used as the confidence score. (c) Shows results averaged over all the data-sets and network architecture (and multiple trials in each condition) for all relevant losses. Closed markers indicate that Maximum Softmax Probability (MSP) was used as the confidence score, while open markers plot results when using Maximum Logit Score (MLS).
• Figure 4: Prediction confidence after learning with standard data-sets. Results are for WRN22-10 networks trained on CIFAR10. Each graph shows histograms of the number of samples classified with different levels of prediction confidence (MSP). Separate histograms are shown for the response generated to unseen samples from known classes (the CIFAR10 test set) and unknown classes (the CIFAR100 test set). The former is measured against the right-hand vertical axis and the latter against the left-hand vertical axis.
• Figure 5: Results when learning with imbalanced data-sets and testing on clean and corrupt images. (a) and (b) directly compare the performance produced by HEM and cross-entropy (CE) losses when used to train networks with long-tailed (LT) CIFAR10 and CIFAR100 each with imbalance factors 100 and 10. ResNet32 was used for the CIFAR10 data and ResNet18 for the CIFAR100 data. (c) and (d) show the same comparisons for HEM+ and LA losses. (e) Shows results averaged over all the data-sets and network architectures (and five trials in each condition) for all relevant losses: cross-entropy (CE), LogitNorm (LN), logit-adjusted (LA), DICE, multi-class margin (MM), HEM, and HEM+ losses. The format of this figure is otherwise the same as, and described in the caption of, \ref{['fig-standard_training_IND']}.