[Paper Review] Density-Aware Personalized Training for Risk Prediction in Imbalanced Medical Data

Zepeng Huo, Xiaoning Qian+3 more|arXiv (Cornell University)|Jul 23, 2022

Artificial Intelligence in Healthcare4 citations

TL;DR

This paper proposes a density-aware personalized training framework for risk prediction in imbalanced medical data by decoupling representation learning from classification and introducing a learnable cost matrix with a density-aware loss. The method improves AUC-ROC, AUC-PRC, and Brier Skill Score while achieving superior calibration over baselines on TOPCAT and MIMIC-III datasets.

ABSTRACT

Medical events of interest, such as mortality, often happen at a low rate in electronic medical records, as most admitted patients survive. Training models with this imbalance rate (class density discrepancy) may lead to suboptimal prediction. Traditionally this problem is addressed through ad-hoc methods such as resampling or reweighting but performance in many cases is still limited. We propose a framework for training models for this imbalance issue: 1) we first decouple the feature extraction and classification process, adjusting training batches separately for each component to mitigate bias caused by class density discrepancy; 2) we train the network with both a density-aware loss and a learnable cost matrix for misclassifications. We demonstrate our model's improved performance in real-world medical datasets (TOPCAT and MIMIC-III) to show improved AUC-ROC, AUC-PRC, Brier Skill Score compared with the baselines in the domain.

Motivation & Objective

• To address class imbalance in medical risk prediction where adverse events are rare, leading to biased model performance.
• To improve model calibration and discrimination by leveraging inherent data density discrepancies rather than eliminating them.
• To develop a personalized training regime that accounts for intra-class heterogeneity in high-risk patients.
• To reduce reliance on heuristic resampling or reweighting techniques that require manual tuning.
• To enhance both predictive performance and probabilistic calibration in real-world medical datasets like TOPCAT and MIMIC-III.

Proposed method

• Decouple feature extraction and classification to reduce bias from class density discrepancy.
• Introduce a density-aware loss function that models the geometric distribution of low-risk (dense) and high-risk (scattered) patients in feature space.
• Use a learnable cost matrix to personalize misclassification penalties based on class-specific risk profiles.
• Train the model using cross-entropy loss combined with the density-aware loss and learnable cost matrix.
• Apply the framework to both a randomized clinical trial (TOPCAT) and an observational EHR dataset (MIMIC-III) with strong baselines.
• Conduct ablation studies and hyperparameter tuning on the learnable cost matrix parameter θ to optimize performance.

Experimental results

Research questions

• RQ1 Can decoupling representation learning from classification mitigate bias in imbalanced medical data?
• RQ2 How does modeling intra-class density differences (dense low-risk vs. scattered high-risk) improve risk prediction?
• RQ3 Can a learnable cost matrix enhance model calibration and performance without requiring external calibration data?
• RQ4 Does the proposed framework outperform standard resampling and reweighting techniques in both AUC-ROC and Brier Skill Score?
• RQ5 To what extent does the framework improve calibration in high-risk probability ranges where overconfidence is common?

Key findings

• The full framework achieved the highest AUC-ROC and AUC-PRC on both TOPCAT and MIMIC-III datasets, outperforming all baseline methods.
• Decoupling alone improved AUC-PRC significantly but led to poor calibration, as shown by low Brier Skill Score (BSS) values.
• The density-aware loss alone improved calibration, achieving positive BSS, indicating better probabilistic reliability.
• Combining decoupling with the density-aware loss and learnable cost matrix yielded the best BSS and calibration, with predicted probabilities closer to the diagonal in calibration plots.
• Hyperparameter tuning showed θ=5 maximized AUC-ROC, while θ=10 was optimal for AUC-PRC, with minimal statistical difference between them.
• The model achieved superior calibration without requiring additional calibration data, demonstrating robustness in high-risk probability regions.