[论文解读] Few-Shot Electronic Health Record Coding through Graph Contrastive Learning

Shanshan Wang, Pengjie Ren+6 more|arXiv (Cornell University)|Jun 29, 2021

Machine Learning in Healthcare参考文献 53被引用 6

一句话总结

本文提出 CoGraph，一种用于少样本电子健康记录（EHR）编码的图对比学习框架，显著提升了常见与罕见 ICD 编码的性能。通过构建异质词-实体（HEWE）图，并应用两种对比学习方案——图采样对比学习（GSCL）与图演化对比学习（GECL）——CoGraph 从结构模式中学习可迁移表征，在罕见 ICD 编码上实现显著性能提升（F1 最高提升 2.95%），同时保持对常见编码的强性能。

ABSTRACT

Electronic health record (EHR) coding is the task of assigning ICD codes to each EHR. Most previous studies either only focus on the frequent ICD codes or treat rare and frequent ICD codes in the same way. These methods perform well on frequent ICD codes but due to the extremely unbalanced distribution of ICD codes, the performance on rare ones is far from satisfactory. We seek to improve the performance for both frequent and rare ICD codes by using a contrastive graph-based EHR coding framework, CoGraph, which re-casts EHR coding as a few-shot learning task. First, we construct a heterogeneous EHR word-entity (HEWE) graph for each EHR, where the words and entities extracted from an EHR serve as nodes and the relations between them serve as edges. Then, CoGraph learns similarities and dissimilarities between HEWE graphs from different ICD codes so that information can be transferred among them. In a few-shot learning scenario, the model only has access to frequent ICD codes during training, which might force it to encode features that are useful for frequent ICD codes only. To mitigate this risk, CoGraph devises two graph contrastive learning schemes, GSCL and GECL, that exploit the HEWE graph structures so as to encode transferable features. GSCL utilizes the intra-correlation of different sub-graphs sampled from HEWE graphs while GECL exploits the inter-correlation among HEWE graphs at different clinical stages. Experiments on the MIMIC-III benchmark dataset show that CoGraph significantly outperforms state-of-the-art methods on EHR coding, not only on frequent ICD codes, but also on rare codes, in terms of several evaluation indicators. On frequent ICD codes, GSCL and GECL improve the classification accuracy and F1 by 1.31% and 0.61%, respectively, and on rare ICD codes CoGraph has more obvious improvements by 2.12% and 2.95%.

研究动机与目标

• 解决 EHR 编码中严重的类别不平衡问题，即由于训练样本不足，导致罕见 ICD 编码预测性能差。
• 不仅提升常见 ICD 编码的预测性能，尤其关注现有方法常被忽略的罕见编码。
• 开发一种少样本学习框架，通过自监督图表征学习，将常见 ICD 编码的知识迁移至罕见编码。
• 通过利用内在图结构进行对比预训练，降低模型仅学习对常见 ICD 编码有用的特征的风险。
• 设计利用 EHR 中图内与图间相关性的对比学习方案，以学习鲁棒且可迁移的表征，而无需依赖外部知识。

提出的方法

• 为每份 EHR 构建一个异质词-实体（HEWE）图，其中词语与医学实体为节点，关系（如共现、提及）为边。
• 将每个 ICD 编码预测视为少样本学习设置下的元任务，以实现从常见编码到罕见编码的知识迁移。
• 引入图采样对比学习（GSCL），通过对比同一 HEWE 图中采样的子图，学习不变性。
• 提出图演化对比学习（GECL），通过对比同一患者在不同临床阶段的历史与未来 HEWE 图，捕捉时间动态。
• 使用这些对比目标预训练图编码器，以在微调前学习可迁移的特征。
• 在训练中采用元学习策略，使用支持集与查询集模拟少样本场景，确保模型在有限监督下对罕见 ICD 编码具备泛化能力。

实验结果

研究问题

• RQ1 当仅有少量标注样本时，少样本学习框架是否能有效提升罕见 ICD 编码的预测性能？
• RQ2 如何利用 EHR 中的内在结构模式，在不依赖外部医学知识的前提下学习可迁移表征？
• RQ3 与标准监督微调相比，GSCL 和 GECL 等图对比学习方案在 EHR 编码表征学习方面的提升程度如何？
• RQ4 随着少样本支持集大小（K）增加，模型在罕见 ICD 编码上的性能是否提升？这对常见编码的性能有何影响？
• RQ5 不同的训练类别采样策略（如随机采样与顶置采样）如何影响模型对罕见 ICD 编码的泛化能力？

主要发现

• 与最先进方法相比，CoGraph 在常见 ICD 编码上的分类准确率提升 1.31%，F1 得分提升 0.61%。
• 在罕见 ICD 编码上，CoGraph 展现出更显著的增益，准确率提升 2.12%，F1 得分提升 2.95%，证明其在低资源 ICD 编码上的有效性。
• 当少样本支持集大小（K）超过 1 后，常见 ICD 编码的性能略有下降，表明常见与罕见编码学习之间存在权衡。
• 即使训练 ICD 类别数量增加，模型在罕见 ICD 编码上的性能仍保持稳定且较高，表明从常见编码到罕见编码的知识迁移有效。
• 罕见 ICD 编码的精确率显著高于召回率与准确率，表明模型倾向于保守预测罕见编码，对极罕见编码的召回率仍有提升空间。
• ‘顶置’采样策略因训练分布中类别不平衡，导致训练动态更不稳定，但对罕见编码的最终性能影响不大。