Multimodal Medical Code Tokenizer

📄 arXiv: 2502.04397v3 📥 PDF

作者: Xiaorui Su, Shvat Messica, Yepeng Huang, Ruth Johnson, Lukas Fesser, Shanghua Gao, Faryad Sahneh, Marinka Zitnik

分类: cs.CL, cs.AI, cs.LG

发布日期: 2025-02-06 (更新: 2025-06-29)

备注: ICML'25

💡 一句话要点

提出MedTok：一种融合文本描述和关系信息的医疗代码多模态Tokenizer

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 医疗代码 Tokenizer 多模态学习 电子病历 图神经网络

📋 核心要点

  1. 现有医疗代码Tokenizer将EHR中的医疗代码视为孤立的文本Token，忽略了代码的文本描述和关系信息。
  2. MedTok通过语言模型编码文本描述，图编码器编码关系结构，并将两种模态信息量化到统一的Token空间。
  3. 实验表明，MedTok在多种EHR建模任务中显著提升了性能，尤其是在药物推荐任务上，并可应用于医疗问答系统。

📝 摘要（中文）

本文提出了一种多模态医疗代码Tokenizer，名为MedTok，旨在解决现有Tokenizer将电子病历（EHR）中的医疗代码视为孤立文本Token的问题。MedTok利用代码的文本描述和关系上下文，通过语言模型编码文本信息，并使用图编码器编码关系结构。然后，它将两种模态量化到统一的Token空间，从而保留模态特定和跨模态信息。MedTok被集成到五个EHR模型中，并在住院和门诊数据集上进行了评估，任务包括结果预测、诊断分类、药物推荐和风险分层。实验结果表明，用MedTok替换标准EHR Tokenizer后，所有EHR模型的AUPRC均得到提高，在MIMIC-III上提高4.10%，在MIMIC-IV上提高4.78%，在EHRShot上提高11.32%，药物推荐任务的提升最为显著。此外，MedTok还可用于医疗问答系统。实验结果表明，MedTok作为医疗代码的统一Tokenizer具有巨大潜力，能够改进医疗基础模型的Tokenization。

🔬 方法详解

问题定义：现有方法将电子病历中的医疗代码视为孤立的文本Token，忽略了医疗代码的文本描述、在本体层次结构中的位置以及与其他代码的关系（如疾病共现和药物治疗关联）。这种处理方式丢失了大量对临床推理至关重要的信息，限制了EHR模型的性能。

核心思路：MedTok的核心思路是将医疗代码的文本描述和关系上下文信息融入到Tokenization过程中。通过同时考虑文本和关系信息，MedTok能够生成更具信息量的Token表示，从而提升下游任务的性能。这样设计的目的是为了让模型能够更好地理解医疗代码的语义和上下文关系。

技术框架：MedTok包含两个主要模块：文本编码器和图编码器。文本编码器使用预训练的语言模型（如BERT）来编码医疗代码的文本描述。图编码器使用图神经网络（GNN）来编码医疗代码之间的关系结构。然后，MedTok使用量化模块将文本和图编码器的输出映射到统一的Token空间。整体流程是：输入医疗代码，分别通过文本编码器和图编码器提取特征，然后通过量化模块得到最终的Token表示。

关键创新：MedTok的关键创新在于它是一种多模态的医疗代码Tokenizer，能够同时利用医疗代码的文本描述和关系上下文信息。与现有方法相比，MedTok能够生成更具信息量的Token表示，从而提升下游任务的性能。这种多模态融合的方式是现有方法所不具备的。

关键设计：文本编码器可以使用各种预训练的语言模型，例如BERT、RoBERTa等。图编码器可以使用各种图神经网络，例如GCN、GAT等。量化模块可以使用各种量化方法，例如矢量量化（VQ）。损失函数的设计需要考虑文本和图编码器的输出，以及量化模块的输出。具体参数设置需要根据具体任务进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，用MedTok替换标准EHR Tokenizer后，所有EHR模型的AUPRC均得到提高，在MIMIC-III上提高4.10%，在MIMIC-IV上提高4.78%，在EHRShot上提高11.32%，药物推荐任务的提升最为显著。这些结果表明，MedTok能够有效提升EHR模型的性能，尤其是在药物推荐等关键任务上。

🎯 应用场景

MedTok可广泛应用于电子病历分析、医疗问答系统、药物研发等领域。通过提升医疗代码的Tokenization质量，MedTok能够提高EHR模型的性能，从而改善临床决策支持、疾病预测和个性化治疗方案的制定。未来，MedTok有望成为医疗领域基础模型的重要组成部分，推动医疗人工智能的发展。

📄 摘要（原文）

Foundation models trained on patient electronic health records (EHRs) require tokenizing medical data into sequences of discrete vocabulary items. Existing tokenizers treat medical codes from EHRs as isolated textual tokens. However, each medical code is defined by its textual description, its position in ontological hierarchies, and its relationships to other codes, such as disease co-occurrences and drug-treatment associations. Medical vocabularies contain more than 600,000 codes with critical information for clinical reasoning. We introduce MedTok, a multimodal medical code tokenizer that uses the text descriptions and relational context of codes. MedTok processes text using a language model encoder and encodes the relational structure with a graph encoder. It then quantizes both modalities into a unified token space, preserving modality-specific and cross-modality information. We integrate MedTok into five EHR models and evaluate it on operational and clinical tasks across in-patient and out-patient datasets, including outcome prediction, diagnosis classification, drug recommendation, and risk stratification. Swapping standard EHR tokenizers with MedTok improves AUPRC across all EHR models, by 4.10% on MIMIC-III, 4.78% on MIMIC-IV, and 11.32% on EHRShot, with the largest gains in drug recommendation. Beyond EHR modeling, we demonstrate using MedTok tokenizer with medical QA systems. Our results demonstrate the potential of MedTok as a unified tokenizer for medical codes, improving tokenization for medical foundation models.