MediTab: Scaling Medical Tabular Data Predictors via Data Consolidation, Enrichment, and Refinement

TL;DR

MediTab introduces a data engineering approach using large language models to consolidate and align heterogeneous medical tabular data, enabling scalable, zero-shot prediction across diverse datasets without fine-tuning.

Contribution

This work presents MediTab, a novel data engine leveraging LLMs for schema consolidation and data alignment, significantly improving zero-shot medical tabular prediction performance.

Findings

Achieves top rankings on multiple patient and trial outcome datasets.

Outperforms supervised models in zero-shot prediction accuracy.

Enables arbitrary tabular input inference without fine-tuning.

Abstract

Tabular data prediction has been employed in medical applications such as patient health risk prediction. However, existing methods usually revolve around the algorithm design while overlooking the significance of data engineering. Medical tabular datasets frequently exhibit significant heterogeneity across different sources, with limited sample sizes per source. As such, previous predictors are often trained on manually curated small datasets that struggle to generalize across different tabular datasets during inference. This paper proposes to scale medical tabular data predictors (MediTab) to various tabular inputs with varying features. The method uses a data engine that leverages large language models (LLMs) to consolidate tabular samples to overcome the barrier across tables with distinct schema. It also aligns out-domain data with the target task using a "learn, annotate, and refinement" pipeline. The expanded training data then enables the pre-trained MediTab to infer for arbitrary tabular input in the domain without fine-tuning, resulting in significant improvements over supervised baselines: it reaches an average ranking of 1.57 and 1.00 on 7 patient outcome prediction datasets and 3 trial outcome prediction datasets, respectively. In addition, MediTab exhibits impressive zero-shot performances: it outperforms supervised XGBoost models by 8.9% and 17.2% on average in two prediction tasks, respectively.