[论文解读] Learning the Graphical Structure of Electronic Health Records with Graph Convolutional Transformer

Edward Choi, Zhen Xu+5 more|arXiv (Cornell University)|Jun 11, 2019

Machine Learning in Healthcare参考文献 28被引用 4

一句话总结

该论文提出图卷积注意力变换器（GCT），一种新颖的模型，能够联合学习电子健康记录（EHR）的隐藏图结构，并在缺乏结构信息的情况下执行监督预测任务。通过利用条件概率先验和注意力掩码引导自注意力机制，GCT 在合成和真实世界EHR数据上的图结构重建与再入院预测任务中均优于基线模型。

ABSTRACT

Effective modeling of electronic health records (EHR) is rapidly becoming an important topic in both academia and industry. A recent study showed that using the graphical structure underlying EHR data (e.g. relationship between diagnoses and treatments) improves the performance of prediction tasks such as heart failure prediction. However, EHR data do not always contain complete structure information. Moreover, when it comes to claims data, structure information is completely unavailable to begin with. Under such circumstances, can we still do better than just treating EHR data as a flat-structured bag-of-features? In this paper, we study the possibility of jointly learning the hidden structure of EHR while performing supervised prediction tasks on EHR data. Specifically, we discuss that Transformer is a suitable basis model to learn the hidden EHR structure, and propose Graph Convolutional Transformer, which uses data statistics to guide the structure learning process. The proposed model consistently outperformed previous approaches empirically, on both synthetic data and publicly available EHR data, for various prediction tasks such as graph reconstruction and readmission prediction, indicating that it can serve as an effective general-purpose representation learning algorithm for EHR data.

研究动机与目标

• 解决在缺乏完整结构信息（如诊断-治疗关系）的情况下，从EHR数据中学习有效表征的挑战。
• 探究在无显式监督的情况下，Transformer中的自注意力机制是否能够被引导以发现有意义的EHR图结构。
• 开发一种通用的EHR表示学习框架，联合学习图结构并执行下游预测任务。
• 通过实证验证，利用注意力掩码和正则化引入先验知识可提升结构学习与预测性能。

提出的方法

• 提出图卷积注意力变换器（GCT），一种改进的Transformer架构，通过注意力掩码和基于条件概率的正则化将先验知识融入自注意力机制。
• 使用条件概率矩阵P作为先验，引导注意力学习，确保注意力头聚焦于EHR特征之间的合理关系。
• 对注意力机制应用归一化和可学习参数，以平衡先验知识与数据驱动的注意力模式。
• 采用多头自注意力机制，其中每个头均受源自先验概率矩阵P的掩码约束。
• 在EHR数据上端到端训练模型，使用图重建和诊断-治疗分类等监督任务。
• 利用KL散度和注意力可视化评估模型注意力对真实邻接结构的逼近程度。

实验结果

研究问题

• RQ1 当未提供显式结构时，Transformer模型中的自注意力机制能否学习到有意义的EHR图结构？
• RQ2 通过注意力掩码和条件概率正则化引入先验知识，是否能提升结构学习与预测性能？
• RQ3 在合成和真实EHR数据上，GCT与标准Transformer及GCN模型在学习隐藏EHR结构方面表现如何？
• RQ4 注意力引导在下游任务（如再入院预测与图结构重建）中的性能提升程度如何？

主要发现

• GCT在合成和真实EHR数据上的图结构重建与再入院预测任务中，均优于基线模型，包括标准Transformer和使用固定先验的GCN模型。
• GCT在预测邻接矩阵与真实邻接矩阵之间的KL散度平均为0.48，低于标准Transformer的0.62，表明其结构恢复能力更强。
• GCT学习到的注意力分布与真实结构高度吻合，尤其在早期自注意力块中表现明显，且在深层网络中逐渐逼近真实邻接结构。
• 与使用固定先验的GCN（GCN P）相比，GCT在图结构重建与诊断-治疗分类任务中表现更优，证明了端到端结构学习的优势。
• 注意力可视化结果表明，GCT最初依赖先验矩阵P，但逐渐优化注意力以匹配真实结构，而未引导的Transformer则不具备此特性。
• 即使注意力图未与真实结构完全对齐，GCT在图结构重建任务中仍表现出色，凸显了注意力机制在结构发现中的鲁棒性。