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

TL;DR

MediTab addresses the challenge of scaling medical tabular predictors across heterogeneous sources with limited samples by building a data engine that converts tabular rows into unified natural language descriptions using LLMs. It then enriches and refines this data through a Learn-Annotate-Audit pipeline that leverages distantly supervised pseudo-labels from multiple tasks and data Shapley auditing to filter high-quality samples. A final multi-task predictor is trained in a two-stage process (pre-training on augmented data $\mathbf{T}_{\text{sup}}$ then fine-tuning on the target data $\mathbf{T}$), enabling zero-shot and few-shot inference on unseen datasets $\mathbf{D}$ not in the original task set. Across seven patient-outcome datasets and three trial-outcome datasets, MediTab achieves superior average rankings (e.g., $1.57$ for patient outcomes) and notable zero-shot gains over supervised baselines, demonstrating a practical, scalable approach to data-centric medical tabular prediction that reduces per-dataset engineering. This framework advances cross-task generalization in clinical settings and highlights the potential of combining LLM-driven data consolidation with principled data auditing for robust, deployable tabular predictors.

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.

Figures (7)

• Figure 1: MediTab vs. existing tabular prediction methods. Existing methods learn and predict on a per-dataset basis, while MediTab can use data from the target task and all other tasks to improve performance.
• Figure 2: The demonstration of scaling medical tabular data predictors models (MediTab). It encompasses three steps: Step 1 consolidates tabular datasets using LLM; Step 2 aligns out-domain datasets with the target task; Step 3 facilitates the predictor with cleaned supplementary data. More details are presented in Section \ref{['sec:method_framework']}.
• Figure 3: Zero-shot MediTab is better than a fully supervised baseline (XGBoost). The evaluation is across 7 patient outcome prediction datasets (left) and 3 trial outcome prediction datasets (right). The compared baseline XGBoost model is fitted on each dataset, respectively.
• Figure 4: Few-shot MediTab compared with XGBoost with varying training data sizes. The compared baseline XGBoost model is fitted on each dataset, respectively.
• Figure 5: ROC-AUC Ranking (lower is better) of the variations of MediTab.