Xiaorui Su, Shvat Messica, Yepeng Huang, Ruth Johnson, Lukas Fesser, Shanghua Gao, Faryad Sahneh, Marinka Zitnik
4 min
Standard electronic health record (EHR) models typically treat medical codes as isolated textual tokens. This approach fails to capture the rich, structured information inherent in medical terminologies, such as hierarchical relationships, disease co-occurrences, and drug-treatment associations. The authors ask whether a multimodal tokenizer that explicitly incorporates both textual descriptions and graph-based relational context can improve the performance of downstream clinical and operational tasks.
The authors introduce MedTok, a multimodal medical code tokenizer that operates in three primary stages. First, it uses a language model encoder to process textual descriptions and a graph encoder to process subgraphs extracted from biomedical knowledge graphs. Second, it employs vector quantization to map these modality-specific and cross-modality embeddings into a shared, unified token space. Finally, it uses a token-packing objective to ensure that the resulting tokens preserve both modality-specific nuances and cross-modality interactions. MedTok is designed as a general-purpose, plug-and-play component that can be integrated into any transformer-based EHR model or medical QA system.
MedTok was evaluated by replacing standard tokenizers in five different EHR models across three datasets: MIMIC-III, MIMIC-IV, and EHRShot. The results show consistent improvements in AUPRC across all models and tasks, including diagnosis classification, patient risk stratification, and drug recommendation. Specifically, MedTok achieved performance gains of 4.10% on MIMIC-III, 4.78% on MIMIC-IV, and 11.32% on EHRShot. The most substantial improvements were observed in drug recommendation tasks, suggesting that the graph-based relational context is particularly beneficial for capturing complex medication dependencies.
Medical foundation models are limited by how they represent the vast, complex vocabulary of clinical codes. By moving beyond simple textual tokenization, MedTok provides a more robust, semantically rich representation of medical data. This approach not only improves predictive accuracy in clinical settings but also offers a scalable solution for integrating diverse coding systems, bridging semantic gaps, and enhancing the reasoning capabilities of medical AI systems.
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.
Alex: Does that cleaner representation actually translate to better performance?
Sam: The results suggest it does. Across several medical AI tasks, swapping in MEDTOK's approach produced a meaningful boost in performance. The most notable gains were in drug recommendations — which makes sense, because that's exactly the kind of task that depends on understanding relationships between conditions and treatments.
Alex: So the key wasn't giving the AI more data — it was giving it a better way to understand the data it already had.
Sam: That's the central idea. There's also a mechanism the researchers call a packing loss — a kind of internal pressure that forces the model to be selective about what it encodes. Think of it like a strict luggage limit before a flight. You can't bring everything, so you have to decide what's truly essential. That constraint pushes the system to capture only the most clinically meaningful information in each token, which means the AI is learning the structure of medical knowledge rather than just memorising raw data.
Alex: So it's not just about adding more information — it's about being disciplined about which information actually matters.
Sam: Exactly. And that discipline is what makes the tokens genuinely useful across different models and tasks, rather than just useful in one narrow setting.
Alex: Does adopting this require hospitals or researchers to completely rebuild their existing AI tools?
Sam: That's one of the practical strengths of the approach. MEDTOK is designed as a tokenizer — essentially a pre-processing step that runs before the main model even sees the data. You can plug it into existing medical AI systems without changing their underlying architecture. Think of it as replacing a blurry lens with a sharper one on a camera that already works — the camera itself stays the same.
Alex: So it's a more thoughtful way of feeding medical knowledge into models we already have, rather than a new model entirely.
Sam: That's a fair summary. The researchers' argument is that the bottleneck in medical AI isn't always the model itself — it's the quality of the representation going in. If the input is impoverished, even a sophisticated model is working with one hand tied behind its back. MEDTOK is an attempt to fix that foundation. The underlying medical knowledge was always there, encoded in those graphs and descriptions. This just builds a representation that lets the model actually use it.
Alex: A smarter front door. Thanks for walking us through it — and thanks to everyone listening to ResearchPod.