[论文解读] 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 Healthcare被引用 4

一句话总结

该论文提出了一种密度感知的个性化训练框架，用于处理类别不平衡的医疗数据中的风险预测问题。通过将表征学习与分类解耦，并引入可学习的成本矩阵及密度感知损失函数，该方法在TOPCAT和MIMIC-III数据集上实现了优于基线模型的AUC-ROC、AUC-PRC和Brier Skill Score，同时在模型校准方面表现更优。

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.

研究动机与目标

• 为解决医疗风险预测中不良事件罕见导致的类别不平衡问题，避免模型性能偏差。
• 通过利用数据内在的密度差异而非消除这些差异，提升模型的校准性与判别能力。
• 开发一种个性化训练方案，以考虑高风险患者内部的异质性。
• 减少对需要手动调参的启发式重采样或重加权技术的依赖。
• 在真实世界医疗数据集（如TOPCAT和MIMIC-III）中，同时提升预测性能与概率校准能力。

提出的方法

• 将特征提取与分类过程解耦，以减轻类别密度差异带来的偏差。
• 引入一种密度感知损失函数，用于建模特征空间中低风险患者（密集分布）与高风险患者（稀疏分布）的几何分布特性。
• 使用可学习的成本矩阵，根据类别特定的风险特征个性化设置误分类惩罚。
• 采用交叉熵损失结合密度感知损失与可学习成本矩阵进行模型训练。
• 在随机临床试验（TOPCAT）与观察性电子健康记录数据集（MIMIC-III）上应用该框架，并与强基线方法进行对比。
• 对可学习成本矩阵参数θ进行消融研究与超参数调优，以优化性能。

实验结果

研究问题

• RQ1 将表征学习与分类解耦是否能缓解类别不平衡医疗数据中的偏差？
• RQ2 建模类内密度差异（低风险患者密集 vs. 高风险患者稀疏）如何提升风险预测性能？
• RQ3 可学习成本矩阵是否能在无需外部校准数据的情况下提升模型校准性与性能？
• RQ4 所提出的框架是否在AUC-ROC与Brier Skill Score上均优于标准的重采样与重加权技术？
• RQ5 该框架在高风险概率区间（过自信现象常见）中，对校准性的改善程度如何？

主要发现

• 完整框架在TOPCAT与MIMIC-III数据集上均取得了最高的AUC-ROC与AUC-PRC，优于所有基线方法。
• 仅解耦表征学习显著提升了AUC-PRC，但导致校准性能差，表现为Brier Skill Score（BSS）值较低。
• 仅使用密度感知损失即可改善校准性，实现正的BSS，表明概率预测更具可靠性。
• 将解耦、密度感知损失与可学习成本矩阵结合后，BSS与校准性达到最优，校准图中预测概率更接近对角线。
• 超参数调优显示，θ=5时AUC-ROC最大，而θ=10时AUC-PRC最优，两者之间统计差异极小。
• 模型在无需额外校准数据的情况下实现了优越的校准性能，尤其在高风险概率区域表现出强鲁棒性。