[论文解读] Contrastive Learning of Temporal Distinctiveness for Survival Analysis in Electronic Health Records

Mohsen Nayebi Kerdabadi, Arya Hadizadeh Moghaddam+3 more|arXiv (Cornell University)|Aug 24, 2023

Machine Learning in HealthcareComputer Science被引用 3

一句话总结

该论文提出 OTCSurv，一种用于电子健康记录（EHR）生存分析的新型对比学习框架，通过利用时间差异性（以生存时长差异衡量）构建感知难度的负样本对。通过整合基于本体的序列编码器与监督加权对比损失，OTCSurv 在风险排序与生存时间预测方面均表现优异，在 AKI 预测任务中达到最先进性能，并通过注意力分析增强了模型可解释性。

ABSTRACT

Survival analysis plays a crucial role in many healthcare decisions, where the risk prediction for the events of interest can support an informative outlook for a patient's medical journey. Given the existence of data censoring, an effective way of survival analysis is to enforce the pairwise temporal concordance between censored and observed data, aiming to utilize the time interval before censoring as partially observed time-to-event labels for supervised learning. Although existing studies mostly employed ranking methods to pursue an ordering objective, contrastive methods which learn a discriminative embedding by having data contrast against each other, have not been explored thoroughly for survival analysis. Therefore, in this paper, we propose a novel Ontology-aware Temporality-based Contrastive Survival (OTCSurv) analysis framework that utilizes survival durations from both censored and observed data to define temporal distinctiveness and construct negative sample pairs with adjustable hardness for contrastive learning. Specifically, we first use an ontological encoder and a sequential self-attention encoder to represent the longitudinal EHR data with rich contexts. Second, we design a temporal contrastive loss to capture varying survival durations in a supervised setting through a hardness-aware negative sampling mechanism. Last, we incorporate the contrastive task into the time-to-event predictive task with multiple loss components. We conduct extensive experiments using a large EHR dataset to forecast the risk of hospitalized patients who are in danger of developing acute kidney injury (AKI), a critical and urgent medical condition. The effectiveness and explainability of the proposed model are validated through comprehensive quantitative and qualitative studies.

研究动机与目标

• 解决现有生存模型在处理右删失数据及从纵向 EHR 中学习判别性患者表征方面的局限性。
• 探索生存分析中的对比学习，特别是以生存时长差异作为负样本难度的基础。
• 通过多任务学习方法，同时提升生存建模中的风险排序与精确的时间到事件预测。
• 通过注意力机制突出患者病史中的关键就诊记录与医疗编码，以增强模型可解释性。
• 将医学本体知识融入表征学习，以提升泛化能力，尤其针对罕见或稀疏的医疗编码。

提出的方法

• 模型采用双编码器架构：本体编码器用于将 ICD-9 编码映射到分层医学概念，序列自注意力编码器用于建模患者就诊记录中的时间依赖性。
• 设计了一种监督加权对比损失（SupWCon），根据生存时长差异对比患者嵌入，并通过随时间间隔调整的温度参数反映负样本采样的难度。
• 负样本对基于删失与观测患者的生存时长构建，确保时间到事件的差异信息被用于对比学习。
• 框架联合优化三项目标：SupWCon 用于对比表征学习，排序损失用于 concordance 指标，回归损失用于时间预测。
• 采用实例级与就诊级注意力池化机制，生成可解释的注意力权重，用于识别关键就诊记录与单个医疗编码。
• 本体编码器利用医学本体（如 ICD-9）中的分层关系，通过关注更广泛的父级概念，提升稀有或低频编码的表征质量。

Figure 1. Architecture of the proposed OTCSurv model. There are three main components: 1) Sequential attention-based ontological encoder, which mainly consists of an ontological encoder, two attention-pooling blocks, and a transformer encoder to learn the ultimate instance-level representations of p

实验结果

研究问题

• RQ1 基于时间差异性（以生存时长差异衡量）的对比学习是否能提升 EHR 生存分析中的表征学习？
• RQ2 将生存时长差异纳入负样本难度设计，相较于均匀或随机负采样，能否显著提升模型性能？
• RQ3 在 EHR 中，整合医学本体知识在多大程度上能提升表征质量，特别是对罕见或稀疏的医疗编码？
• RQ4 结合对比学习、排序损失与回归损失的多任务学习框架，能否同时提升风险排序与时间预测的准确性？
• RQ5 模型中的注意力权重在多大程度上能反映临床有意义的模式，如关键就诊或高影响诊断，在患者层面决策中的作用？

主要发现

• OTCSurv 在大规模真实世界 EHR 数据集上的 AKI 预测任务中达到最先进性能，在 concordance index（c-index）与时间预测准确性方面均优于现有方法。
• 所提出的 SupWCon 损失通过根据生存时长差异动态调整对比温度，显著提升模型性能，增强对困难负样本的学习能力。
• 注意力分析显示，第 5 次就诊（注意力权重最高）以及编码 572（肝病）、410（心肌梗死）和 287（出血性疾病）对 AKI 发作最具预测性，与临床文献一致。
• 本体编码器通过关注更广泛且更频繁的父级概念（如 ICD-9 编码 514 通过其祖先编码 460-519 获得注意力），有效提升了稀有编码的表征质量。
• 消融实验表明，移除对比损失或本体组件会导致性能显著下降，验证了其关键贡献。
• 模型展现出强大的可解释性，注意力权重能清晰识别出临床相关的就诊记录与诊断，支持临床信任与透明决策。

Figure 2. Comparison of the mean survival curves of proposed model and baselines with Kaplan-Meier curve (for all patients).