Multimodal Medical Code Tokenizer

🕒 Published (v1): 2025-02-06 06:58 UTC · Source: Arxiv · Venue: ICML 25 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

MedTok is a multimodal tokenizer for medical codes that jointly encodes textual descriptions and graph-relational context from biomedical ontologies into a unified vector-quantized token space. It replaces standard text-only tokenizers in EHR foundation models and improves AUPRC by 4.10%–11.32% across MIMIC-III, MIMIC-IV, and EHRShot datasets. It also enhances few-shot medical QA when used as prefix tokens with general-purpose LLMs.

Problem

Standard EHR tokenizers treat medical codes (ICD, SNOMED, ATC, CPT, etc.) as isolated text tokens, discarding six critical properties: (1) hierarchical ontological structure, (2) relational knowledge (co-morbidities, drug-treatment links), (3) cross-system semantic redundancy, (4) vocabulary scalability beyond 600,000 codes, (5) poor embeddings for rare/low-frequency codes, and (6) absence of any multimodal representation. Existing LLM-enriched methods (GraphCare, RAM-EHR) depend on prompted knowledge triplets and lack a unified tokenization framework.

Method

MedTok operates in three stages over 617,490 codes from eight systems (ICD-9, ICD-10-CM, ICD-10-PCS, SNOMED CT, ATC, NDC, CPT, RxNORM):

1. Multimodal encoding: A language model encoder produces text embedding \(x_t\) from the code's natural-language description; a graph encoder produces \(x_g\) from a local subgraph of PrimeKG. Linear projectors yield modality-specific embeddings (\(e^s_t\), \(e^s_g\)); a cross-attention module yields cross-modality embeddings (\(e^c_t\), \(e^c_g\)).

2. Vector quantization into partitioned codebook: A unified codebook \(C \in \mathbb{R}^{N \times d}\) is divided into text-specific, graph-specific, and shared regions. Top-\(K\) nearest codebook vectors are selected per embedding via Euclidean distance; quantized vectors are computed as distance-weighted sums with straight-through gradient estimation.

3. Token packing loss: Cross-modality tokens are aligned via KL divergence on distance distributions and InfoNCE mutual information maximization (with redundancy penalty \(\beta\)). Modality-specific tokens are optimized with InfoNCE plus an orthogonality loss (weight \(\lambda\)). The full loss is \(L = L_{vq} + L_{KL} + L^c_{token} + L^s_{token}\).

After pretraining, MedTok is a drop-in tokenizer for any transformer-based EHR model. For medical QA, MedTok tokens are prepended as prefix to LLM input after disease-code extraction and prefix fine-tuning on MedMCQA.

Key Contributions

• Unified multimodal tokenizer that jointly encodes text and ontological graph structure for 617,490 medical codes across 8 coding systems.
• Partitioned codebook design with separate text-specific, graph-specific, and shared regions preserving modality-distinct and cross-modality information simultaneously.
• Token packing objective combining KL alignment, InfoNCE, and orthogonality losses to maximize inter-modality information while suppressing redundancy.
• Demonstrated drop-in compatibility with five EHR architectures (ETHOS, GT-BEHRT, MulT-EHR, TransformEHR, BEHRT) and three LLMs for QA (LLaMA3.1-8B, Qwen2.5-7B, MMedLM).
• Interpretability analysis mapping frequent token IDs to clinically coherent codes (e.g., Rosuvastatin, T2DM, Hyperlipidemia) for high-risk patient subsets.

Results

• MIMIC-III (5 tasks, 5 models): average AUPRC +4.10% over standard tokenizers; largest gain on drug recommendation (+4.10% average, e.g., ETHOS DrugRec 0.104→0.170).
• MIMIC-IV (5 tasks): average AUPRC +4.78%; drug recommendation ETHOS 0.131→0.240.
• EHRShot (outpatient, 7 subtasks): average AUPRC +11.32%, driven by mortality prediction; Hyperlipidemia new-diagnosis detection +6.00%.
• MedTok outperforms both BERT tokenizer and VQGraph tokenizer (graph-only baseline) on all tasks and datasets under TransformEHR backbone.
• Ablation: full loss (\(L_{vq} + L_{KL} + L^c + L^s\)) improves over \(L_{vq}\) alone by +3.9%, +5.7%, +9.1% on MIMIC-III, MIMIC-IV, EHRShot respectively.
• Medical QA: MedTok prefix tokens improve accuracy on MMLU, MedDDx, AfrimedQA for all three LLMs tested (exact numbers from Figure 6, e.g., LLaMA3.1-8B MMLU 0.437→0.494; MMedLM AfrimedQA 0.491→0.618).

Limitations

• Drug recommendation evaluation restricted to five specific drugs (Vancomycin, Levofloxacin, Heparin Sodium, Metoprolol, Atorvastatin), limiting generalization claims.
• Optimal codebook size is dataset-dependent (\(N=12{,}000\) for MIMIC, \(N=24{,}000\) for EHRShot); requires hyperparameter tuning per deployment.
• Relational context sourced exclusively from PrimeKG; codes absent or poorly represented in this KG may receive impoverished graph embeddings.
• Medical QA integration requires a disease-code extraction step before tokenization, introducing a preprocessing dependency not evaluated for noise sensitivity.
• No evaluation on clinical notes or unstructured EHR data; scope limited to structured medical codes.

Relevance to Foundation Models in Medicine

MedTok addresses a foundational infrastructure problem for EHR-based medical AI: how to tokenize structured clinical codes in a way that preserves the semantic and relational richness embedded in medical ontologies. As EHR foundation models (e.g., ETHOS, BEHRT) scale toward broader pretraining, the quality of their tokenization layer directly constrains what clinical knowledge can be learned from structured data. By demonstrating consistent plug-in improvements across five disparate EHR architectures and extending to LLM-based medical QA, MedTok offers a transferable tokenization component that could standardize how any transformer-based medical foundation model handles coded clinical data. The VQ-based multimodal design also connects to the broader trend of unified token spaces enabling cross-modal reasoning in medical AI.