Multimodal Medical Code Tokenizer

2025-02-06

AI Summary

To address insufficient multimodal semantic modeling of medical codes in electronic health records (EHRs), this paper proposes MedTok—the first multimodal medical code tokenizer integrating textual descriptions and ontology-based relational graph structures. MedTok innovatively unifies language models (text modality) and graph neural networks (relational modality) via joint encoding, with a shared vector quantization module mapping both modalities into a unified, semantically enriched discrete token space—explicitly capturing clinical logical relationships and departing from conventional isolated tokenization paradigms. Evaluated on MIMIC-III, MIMIC-IV, and EHRShot, MedTok achieves average AUPRC improvements of 4.10–11.30%, with the largest gains observed in drug recommendation tasks. It has been integrated into five EHR foundation models and a medical question-answering system, demonstrating strong generalizability and practical utility.

Technology Category

• Machine Learning: Multimodal Learning
• Computer Vision: Multi-modal Vision
• Natural Language Processing: Language Grounding & Multi-modal NLP

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG

Abstract

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.30% 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.

Problem

• Enhances medical code tokenization
• Integrates text and relational data
• Improves EHR model performance

Innovation

• Multimodal tokenizer for EHRs
• Graph encoder for relational structure
• Unified token space integration