ConPro: Learning Severity Representation for Medical Images using Contrastive Learning and Preference Optimization

TL;DR

This work tackles the problem of representing severity in medical images, where severity is inherently ordered rather than purely categorical. It introduces ConPrO, a two-stage framework that first uses contrastive learning to create a normal-vs-abnormal separation and then applies a preference-optimization stage guided by a normality anchor and cosine-distance rewards to embed severity ordering in the latent space. Key contributions include the ConPrO architecture, a new MAEE metric for severity prediction, and empirical evidence showing 6% relative improvement over SupCon and 20% over SimCLR on macro F1 across two medical imaging datasets (Papilledema and VinDr-Mammo). The results demonstrate improved severity-aware representations with potential to enhance clinical assessment, monitoring, and cross-modality reasoning, while the discussion outlines avenues for explainable AI and extended applications in multi-pathology contexts. $d_{cos}$ and the normality anchor play central roles in the learned ordering, enabling more interpretable and clinically aligned severity embeddings.$

Abstract

Understanding the severity of conditions shown in images in medical diagnosis is crucial, serving as a key guide for clinical assessment, treatment, as well as evaluating longitudinal progression. This paper proposes Con- PrO: a novel representation learning method for severity assessment in medical images using Contrastive learningintegrated Preference Optimization. Different from conventional contrastive learning methods that maximize the distance between classes, ConPrO injects into the latent vector the distance preference knowledge between various severity classes and the normal class. We systematically examine the key components of our framework to illuminate how contrastive prediction tasks acquire valuable representations. We show that our representation learning framework offers valuable severity ordering in the feature space while outperforming previous state-of-the-art methods on classification tasks. We achieve a 6% and 20% relative improvement compared to a supervised and a self-supervised baseline, respectively. In addition, we derived discussions on severity indicators and related applications of preference comparison in the medical domain.

Figures (5)

• Figure 1: Conventional (a) and target (b) representation for severity modeling in latent space. A darker color represents a higher severity level, and '0' represents normality. Our proposed method targets to embed distance relation to severity classes in representation space
• Figure 2: ConPro learning framework (a) includes contrastive learning and preference optimization to get the desired latent space (b)
• Figure 3: T-SNE visualization of representation vectors of (a) training and (b) test set after supervised contrastive learning (c) training (d) test set after preference optimization. The plots are samples from the Papilledema dataset. All figures use cosine distance. Label '0' denotes normality, and "1-5" denotes increasing level of severity.
• Figure 4: T-SNE visualization of representation vectors of (a) training and (b) test set after supervised contrastive learning (c) training (d) test set after preference optimization. The plots are 2000 random samples from the VinDr-Mammo dataset. All figures use cosine distance. Label '0' denotes normality, and "1-4" denotes increasing level of severity.
• Figure 5: Confusion matrices of same setting on two independent run. Two matrix have the same F1 score but different in MAE and MAEE