[论文解读] Bootstrapping Your Own Positive Sample: Contrastive Learning With Electronic Health Record Data

Tingyi Wanyan, Jing Zhang+4 more|arXiv (Cornell University)|Apr 7, 2021

Machine Learning in Healthcare参考文献 37被引用 8

一句话总结

本文提出一种对比学习框架，采用新颖的正样本采样策略——基于特征的和基于属性的——用于电子健康记录（EHR）数据，以提升新冠肺炎患者死亡率预测性能。通过利用特征空间中的患者相似性及预计算的相似性度量，该方法增强了类间分离性并保留了类内异质性，AUROC达到0.959，优于交叉熵损失（0.873）和焦点损失（0.931）。

ABSTRACT

Electronic Health Record (EHR) data has been of tremendous utility in Artificial Intelligence (AI) for healthcare such as predicting future clinical events. These tasks, however, often come with many challenges when using classical machine learning models due to a myriad of factors including class imbalance and data heterogeneity (i.e., the complex intra-class variances). To address some of these research gaps, this paper leverages the exciting contrastive learning framework and proposes a novel contrastive regularized clinical classification model. The contrastive loss is found to substantially augment EHR-based prediction: it effectively characterizes the similar/dissimilar patterns (by its "push-and-pull" form), meanwhile mitigating the highly skewed class distribution by learning more balanced feature spaces (as also echoed by recent findings). In particular, when naively exporting the contrastive learning to the EHR data, one hurdle is in generating positive samples, since EHR data is not as amendable to data augmentation as image data. To this end, we have introduced two unique positive sampling strategies specifically tailored for EHR data: a feature-based positive sampling that exploits the feature space neighborhood structure to reinforce the feature learning; and an attribute-based positive sampling that incorporates pre-generated patient similarity metrics to define the sample proximity. Both sampling approaches are designed with an awareness of unique high intra-class variance in EHR data. Our overall framework yields highly competitive experimental results in predicting the mortality risk on real-world COVID-19 EHR data with a total of 5,712 patients admitted to a large, urban health system. Specifically, our method reaches a high AUROC prediction score of 0.959, which outperforms other baselines and alternatives: cross-entropy(0.873) and focal loss(0.931).

研究动机与目标

• 解决EHR数据在临床预测任务中面临的类别不平衡和高类内方差挑战。
• 克服传统对比学习在EHR中应用的局限性，即与视觉任务相比，数据增强的可行性较低。
• 开发针对EHR数据的正样本采样策略，以捕捉语义相似性，同时保留同一结果类别内患者的多样性。
• 通过最小监督的对比正则化，提升在类别不平衡EHR数据上的临床模型预测性能。
• 证明基于有信息的正样本采样策略的对比学习可带来更好的泛化能力和鲁棒性，适用于真实世界EHR应用场景。

提出的方法

• 提出一种结合焦点损失与对比损失的对比正则化分类模型，以应对类别不平衡并提升特征学习能力。
• 提出一种基于特征的正样本采样策略，通过在学习到的特征空间中选择邻居来构建正样本对，增强局部邻域学习。
• 开发一种基于属性的正样本采样策略，利用预计算的患者相似性度量（例如基于临床特征）来定义语义上有意义的正样本对。
• 应用具有“拉近-推开”机制的对比损失函数：将正样本（相似患者）在嵌入空间中拉近，将负样本（不相似患者）推开。
• 使用UMAP可视化和嵌入分离评分（ESS）评估学习表示空间中的类间与类内聚类质量。
• 在对比嵌入上训练线性分类器，以评估下游预测性能，并验证所学表示的质量。

实验结果

研究问题

• RQ1 与标准监督损失相比，采用自定义正样本采样策略的对比学习是否能提升在类别不平衡EHR数据上的预测性能？
• RQ2 基于特征和基于属性的正样本采样策略在捕捉EHR数据类内异质性方面，与随机采样相比表现如何？
• RQ3 所提出的采样策略在多大程度上增强了学习嵌入空间中的类间分离性？
• RQ4 在EHR中存在高数据异质性和类别不平衡的情况下，对比正则化是否能带来更鲁棒和泛化能力更强的表示？
• RQ5 所提出的方法是否能在真实世界临床预测任务（如新冠肺炎患者院内死亡率预测）中实现更优性能？

主要发现

• 所提出的对比正则化模型在真实世界新冠肺炎EHR数据上实现了0.959的AUROC，显著优于交叉熵损失（0.873）和焦点损失（0.931）。
• 基于属性的正样本采样策略表现最佳，其线性分类器准确率比随机采样高出4.5%。
• 基于特征和基于属性的采样策略获得的嵌入分离评分（ESS）高于随机采样，表明类间分离性更优。
• 两种新型采样策略保留的类内方差（到聚类中心距离的标准差）高于随机采样，反映出对患者异质性的建模更充分。
• UMAP可视化结果表明，与随机采样相比，所提方法为死亡和非死亡患者创建了更清晰、更分离的聚类。
• 对比正则化器通过在每一类中保持多样化的子群体表示，有效缓解了随机正样本采样常导致的嵌入空间坍塌问题。