Class Distance Weighted Cross Entropy Loss for Classification of Disease Severity

TL;DR

This paper introduces the Class Distance Weighted Cross Entropy (CDW-CE) loss function, designed to improve ordinal disease severity classification by penalizing misclassifications based on class distance, leading to better performance and interpretability.

Contribution

The paper proposes a novel loss function, CDW-CE, tailored for ordinal classification, and demonstrates its superiority over existing loss functions through extensive evaluations and visualizations.

Findings

CDW-CE achieves higher Silhouette Scores indicating better class separation.

Models trained with CDW-CE focus more on clinically relevant regions in CAM visualizations.

CDW-CE outperforms other loss functions in ROC-AUC metrics for disease severity classification.

Abstract

Assessing disease severity with ordinal classes, where each class reflects increasing severity levels, benefits from loss functions designed for this ordinal structure. Traditional categorical loss functions, like Cross-Entropy (CE), often perform suboptimally in these scenarios. To address this, we propose a novel loss function, Class Distance Weighted Cross-Entropy (CDW-CE), which penalizes misclassifications more severely when the predicted and actual classes are farther apart. We evaluated CDW-CE using various deep architectures, comparing its performance against several categorical and ordinal loss functions. To assess the quality of latent representations, we used t-distributed stochastic neighbor embedding (t-SNE) and uniform manifold approximation and projection (UMAP) visualizations, quantified the clustering quality using the Silhouette Score, and compared Class Activation Maps (CAM) generated by models trained with CDW-CE and CE loss. Feedback from domain experts was incorporated to evaluate how well model attention aligns with expert opinion. Our results show that CDW-CE consistently improves performance in ordinal image classification tasks. It achieves higher Silhouette Scores, indicating better class discrimination capability, and its CAM visualizations show a stronger focus on clinically significant regions, as validated by domain experts. Receiver operator characteristics (ROC) curves and the area under the curve (AUC) scores highlight that CDW-CE outperforms other loss functions, including prominent ordinal loss functions from the literature.