Density-Aware Personalized Training for Risk Prediction in Imbalanced Medical Data
This addresses the challenge of accurate risk prediction for medical events like mortality in imbalanced datasets, which is incremental as it builds on existing methods like resampling and reweighting.
The paper tackled the problem of training models for risk prediction in imbalanced medical data, where low event rates lead to suboptimal performance, and demonstrated improved AUC-ROC, AUC-PRC, and Brier Skill Score on real-world datasets like TOPCAT and MIMIC-III compared to baselines.
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.