EHR基础模型中ICD代码的分层建模

# 电子健康记录基础模型中ICD代码的层次建模
来源：https://arxiv.org/html/2606.15447
\\theorembodyfont\\theoremheaderfont\\theorempostheader

:\\theoremsep

\\NameDong Gyun Kang\\nametag1\\Emaildkang335@gatech\.edu\\NameRudra Pratap Singh\\nametag1\\Emailrudra\.singh@gatech\.edu\\NameShruthi Kashinath Hiremath\\nametag2\\Emailshruthi\.hiremath@optum\.com\\NameKatrin Hänsel\\nametag2\\Emailkatrin\_haensel@optum\.com\\NameThomas Plötz\\nametag1\\Emailthomas\.ploetz@gatech\.edu\\addr1School of Interactive ComputingGeorgia Institute of TechnologyAtlantaUnited States 2Optum AIUnited States

###### 摘要

电子健康记录基础模型通常将ICD诊断代码视为扁平化标记，忽略了具有临床意义的层次结构，该结构捕捉了疾病家族、子类别以及细粒度的诊断细节。因此，现有的电子健康记录表示学习方法并未明确利用编码系统中已有的层次结构。在这项工作中，我们将ICD-10-CM层次结构作为临床表示学习的*通用归纳偏置*进行研究。我们研究了两种互补的机制来引入层次结构：首先，通过在BERT风格的Transformer诊断序列中增加对应于ICD层次结构不同级别的标记；其次，通过将层次结构注入到基于图的代码表示中，利用层次感知边与诊断共现结构相结合。在这些设置下，我们评估了显式层次结构是否改善了下游预测、层次的哪些级别最为有用、层次编码是否改善了跨数据集的迁移，以及层次结构如何重塑嵌入相似性结构。我们在两个大规模真实世界临床数据集上进行了实验：MIMIC-IV（用于预训练和域内评估）和eICU（通过冻结编码器探测来评估跨数据集迁移）。我们的发现表明，在域内和跨数据集设置中，显式编码ICD层次结构均优于扁平化代码表示，同时揭示了最有用的层次级别取决于任务和建模方法。更广泛地说，我们关注层次感知的电子健康记录表示学习，并表明编码层次结构的益处可泛化到不同的建模设置和层次级别。我们的代码可在此处获取（https://meghathukral.github.io/HICD_EHR_FMs/）。

## 1 引言

电子健康记录通过诊断、程序、药物和就诊提供了患者临床病史的纵向视图。基于Transformer的电子健康记录模型（如BEHRT和Med-BERT）的进展表明，大规模预训练可以学习到对下游临床预测任务有用的患者表示。然而，这些进展大多将结构化临床代码视为扁平化符号，即使这些代码来自明确设计用于编码临床层次结构的本体。这在数据结构和呈现给模型的结构之间造成了不匹配。例如，ICD代码并非任意标识符：它们将疾病组织成具有临床意义的家族、子组以及越来越具体的类别（参见图1）。因此，一个代码既携带了细粒度的诊断，也携带了通过更广泛临床分类学的路径。然而，在实践中，大多数电子健康记录模型将ICD代码表示为独立标记，忽略了这种层次结构。结果，本应共享信息的相关诊断可能被当作孤立符号来学习。这在临床数据中尤其成问题，因为许多诊断是稀疏的、长尾的且特定于机构的。

我们假设显式建模层次结构可能以几种方式有益于电子健康记录表示学习。首先，层次结构为疾病之间的关系提供了自然的归纳偏置，因此罕见的诊断代码可以从与临床相似代码相关的数据中受益。其次，更宽泛和中间的组别可能捕获比最具体代码更具预测性的护理路径或解剖系统。第三，层次结构可以提高在*领域偏移*下的鲁棒性，因为更高级别的疾病家族通常在不同机构间更稳定。更广泛地说，编码层次结构可以帮助学习到的嵌入更好地反映临床语义，而不仅仅是经验共现。

先前的工作已经探索了将ICD代码层次结构融入医疗保健表示学习，例如GRAM，主要通过将本体结构嵌入代码表示的基于图的方法。然而，这些方法先于当前一代基于Transformer的电子健康记录基础模型，并且如何将ICD层次结构系统地融入此类模型，以及层次的哪个粒度最为有益，仍然探索不足。

为了解决这一差距，我们提出并评估了两种互补策略，用于将ICD层次结构注入当代电子健康记录表示学习中。第一种，HICD-BERT（层次ICD-BERT），直接将层次结构注入BERT风格编码器的标记表示中，将ICD层次的每一级作为附加标记嵌入与诊断代码一起嵌入。第二种，HICD-Graph（层次ICD-图），通过用本体派生边增强的诊断共现图来对层次结构进行关系编码，学习层次感知的代码嵌入，用于初始化患者级别的Transformer。结合起来，这些方法使我们能够比较层次信息是作为标记级信号还是作为代码词汇表上的关系结构被消耗得更好。

我们的结果表明，层次结构是电子健康记录表示学习的*稳健有效的归纳偏置*，因为在两种架构和两种预测任务中，绝大多数层次增强配置都优于其无层次基线。两种模型都受益于层次结构，其中基于图的编码利用了所有层次级别，而基于标记的编码从最细粒度级别中获益最多。至关重要的是，层次结构也改善了跨数据集迁移，特别是对于基于图的层次编码，它从MIMIC-IV稳健地迁移到eICU。

- •我们研究了ICD代码层次结构作为跨两种架构、两种临床预测任务以及三个粒度级别的系统消融的电子健康记录表示学习的归纳偏置，显示在28次比较中的26次中，层次结构显著提升了性能。
- •我们通过评估在冻结编码器探测协议下从MIMIC-IV到eICU的跨数据集迁移，研究了层次结构在分布偏移下的益处，显示基于图的层次编码稳健地迁移到新的数据集和任务设置。
- •通过分析学习到的代码表示的嵌入，我们表明层次结构产生了更连贯的代码簇，并且具有最紧密簇的配置也实现了最强的下游性能。

### 在医疗保健背景下关于机器学习的可泛化见解

我们的发现为健康领域的机器学习提供了两个更广泛的教训。首先，结构化临床本体（如ICD）为电子健康记录基础模型提供了有用的归纳偏置，即使是轻量级的编码方式也能在扁平化代码表示上产生一致的增益。其次，融入具有临床意义的结构也有助于应对分布偏移，并且融入的机制可能产生影响：基于图的层次编码比标记级层次注入在数据集间迁移得更稳健。这些见解不仅适用于ICD-10-CM，还可以扩展到其他结构化临床词汇表，如ATC和程序本体，以及其他表示学习方法。

## 2 相关工作

### 2.1 序列化和基于Transformer的电子健康记录表示学习

纵向电子健康记录建模的早期工作确立了从按时间排序的患者病史中学习的重要性。Doctor AI使用循环神经网络从就诊序列中预测未来临床事件，而RETAIN引入了逆向时间注意力机制以提高医疗保健预测的可解释性。Dipole进一步证明了基于注意力的序列模型可以通过捕获跨就诊的依赖性来改善诊断预测。这些方法表明电子健康记录中的序列结构信息量很大，但它们通常将诊断代码视为扁平化符号。

最近，基于Transformer的模型已成为结构化电子健康记录表示学习的核心方法。BEHRT是一个早期且有影响力的框架，它将患者病史表示为医疗代码序列，并应用自注意力来捕获跨就诊的依赖性。Med-BERT进一步表明，在大型电子健康记录语料库上的掩码预训练可以产生用于下游预测任务的可复用表示。后续模型如CEHR-BERT、ExBEHRT以及更大规模的基础模型如Foresight和ETHOS，通过更丰富的输入表示和更大的训练语料库进一步扩展了这一范式。尽管取得了这些进展，基于Transformer的电子健康记录模型通常将诊断代码视为原子标识符。

### 2.2 本体感知和层次感知的医疗表示学习

并行的研究工作探索了医疗本体如何改善表示学习。其中许多方法是在基础模型时代之前开发的，并且针对单个下游任务进行端到端训练，而不是作为可复用的预训练编码器。GRAM通过在一个循环患者编码器中利用注意力加权聚合来融入本体祖先，并引入了我们研究核心的直觉：临床相关的诊断应通过其在层次结构中的位置共享信息。KAME通过基于知识的注意力机制将其扩展到诊断预测。MiME通过围绕临床有意义的组别结构化医疗代码探索了相关思想。在利用预训练的本体感知方法中，G-BERT将图增强的BERT风格预训练应用于药物推荐，尽管它没有研究不同本体级别的贡献。

互补的研究工作探索了纯粹数据驱动的代码关系方法，从共现中学习结构，而不是从外部分类法。Med2Vec证明，跳跃式共现嵌入可以恢复临床有意义的代码和就诊表示，而无需任何本体监督。后续基于图的方法对诊断之间的经验依赖性进行建模以进行预测：构建患者级异构图，共同编码诊断、程序和人口统计学信息，通过消息传递在临床基础关系上学习表示；同时从患者记录中构建疾病共现图，并应用图神经网络来捕获超越成对嵌入的高阶关系信号。这些方法直接从数据中恢复关系结构，但与本体感知方法不同，它们不利用现有的基于领域知识的层次结构。

在这些工作路线中，先前的方法通常依赖于单一来源的层次信号。即使经过预训练，它们通常也不会跨层次级别进行消融，或明确检查层次结构本身如何塑造学习到的表示。这留下了一个问题：层次结构是否在任务和架构上一致地有帮助，ICD层次的哪些级别贡献最大，以及层次感知表示在跨数据集迁移下如何表现。相比之下，我们在现代基于Transformer的电子健康记录基础模型中直接比较两种范式：HICD-BERT通过字符串前缀截断在标记嵌入级别注入层次结构，在训练时无需访问本体；而HICD-Graph将本体派生的层次边与经验PMI加权共现边结合在一个混合诊断图中，通过图卷积网络学习层次感知的代码嵌入。在两种范式中，我们系统地消融了三个层次的粒度级别，评估了两个域内任务和一个跨数据集迁移任务，并分析了层次结构如何重塑学习到的嵌入几何形状，不仅隔离了层次结构*是否*有帮助，还隔离了*哪种*编码范式使其最有用、*哪些*级别有贡献，以及本体信息是否*严格必要*。

## 3 方法

我们将每位患者表示为按时间排序的就诊序列V=(v1,...,vT)，其中每个就诊vt包含一个或多个ICD诊断代码。我们的目标是学习既保留时间上下文又保留诊断本体层次结构的患者表示，并评估这种结构是否改善了下游预测和泛化。我们研究了两种互补机制，用于将ICD层次结构注入电子健康记录模型，我们统称为层次ICD模型（HICD）。第一种是*标记级*方法，HICD-BERT，它通过为每个诊断代码添加层次特定嵌入标记来扩展BEHRT（图2(a)）。第二种是*图级*方法，HICD-Graph，它构建一个层次增强的诊断共现图，使用图卷积网络学习代码嵌入，并将这些嵌入用于初始化基于Transformer的患者编码器（图2(b)）。完整的实现细节和超参数见附录B。

### 3.1 ICD层次级别

ICD-10-CM代码按层次组织。每个代码属于一个*章节*（广泛的身体系统或病因分组，例如第十九章“损伤、中毒及外因”），分为官方定义的*块*中的相关类别（例如S70--S79“髋部和股部损伤”），其中包含三个字符的*类别*（例如S72“股骨骨折”）。小数点后的字符可编码病因、解剖部位、严重程度和就诊细节（图1）。每个叶子代码嵌套在逐渐粗略的临床分组中，这种嵌套就是我们注入的归纳偏置。

参见图注图1：ICD-10-CM诊断代码的结构。对于每个诊断代码，我们定义三个嵌套的层次级别G0、G1、G2，从最粗到最细排列。这些级别是相对粒度索引，每种架构通过其自身机制（详情如下）对其进行操作。每个消融设置启用或禁用这些级别的子集，产生所有八种组合(G0,G1,G2)∈{0,1}^3，包括非层次基线(0,0,0)，对于HICD-BERT等同于BEHRT，对于HICD-Graph等同于没有层次边的标准GCN-Transformer。此设置比较了两种范式：数据驱动的标记编码，其中层次结构派生自