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

1 Introduction

Tabular data are structured as tables or spreadsheets in a relational database. Each row in the table represents a data sample, while columns represent various feature variables of different types, including categorical, numerical, binary, and textual features. Most previous papers focused on the model design of tabular predictors, mainly by (1) augmenting feature interactions via neural networks [Arik and Pfister, 2021], (2) improving tabular data representation learning by selfsupervised pre-training [Yin et al., 2020; Yoon et al., 2020; Bahri et al., 2022], and (3) performing cross-tabular pretraining for transfer learning [Wang and Sun, 2022b; Zhu et al., 2023]. Tabular data predictor was also employed in medicine, such as patient health risk prediction [Wang and Sun, 2022b; Ma et al., 2023], clinical trial outcome prediction [Fu et al., 2022], modeling Electronic Health Record (EHR) data [Ma et al., 2020; Gao et al., 2024], and unifying heterogeneous EHRs via text embeddings [Hur et al., 2022]. Additionally, LLMs can sample synthetic and yet highly realistic tabular data as well [Borisov et al., 2022; Theodorou et al., 2023].

Despite these significant advances, it is worth noting that the data-centric approaches have received comparatively less attention in prior research. Some prominent examples lie in the detection and mitigation of label noise [Wang et al., 2020; Northcutt et al., 2021], but they only address a fraction of the challenges in medical tabular data prediction. As illustrated in Figure 1, there is typically substantial heterogeneity among different data sources in medical data, and within each data source, the available sample sizes are small. Harnessing multi-source data requires extensive manual effort in terms of data cleaning and formatting. As such, current medical tabular prediction methods are often built on small handcrafted datasets and, hence, do not generalize across tabular datasets.

In this paper, we embrace a data-centric perspective to enhance the scalability of predictive models tailored for medical tabular data. Our core aim revolves around training a single tabular data predictor to accommodate inputs with diverse feature sets. Technically, our framework, namely MediTab, encompasses three key components: data consolidation, enrichment, and refinement modules:

• Data consolidation and enrichment involves consolidating tabular samples with varying features and schemas using natural language descriptions. We also expand the training data by distilling knowledge from large language models and incorporating external tabular datasets.
• Data refinement rectifies errors and hallucinations introduced during the consolidation and enrichment stages. It also aligns a diverse set of tabular samples with the target task through a distantly supervised pipeline.

As illustrated in Figure 1, MediTab offers the advantages:

• Multi-task learning and prediction: the model can learn from and make predictions for multiple medical tabular datasets without requiring modifications or retraining.
• Few-shot and zero-shot learning: the model can quickly adapt to new prediction tasks using only a small amount of training data or even make predictions for any new tabular input when no training data is available.

In Section 2, we provide a detailed description of our approach. We present the experimental results in Section 3, where we demonstrate the effectiveness of our method on 7 patient outcome prediction datasets and 3 trial outcome prediction datasets, achieving an average performance ranking of 1.57 and 1.00, respectively, across tabular prediction baselines. Furthermore, our method shows impressive few-shot and zero-shot performances that are competitive with supervised baselines: the zero-shot MediTab outperforms supervised XGBoost by 8.9% and 17.2% on average in two prediction tasks, respectively. We discuss related work in Section 4 and conclude our findings in Section 5.

2 Method

2.1 ProblemFormulation

We characterize tabular prediction tasks by dataset D and task T, where a task T = {D1, D2, . . . } consists of multiple in-domain datasets with varying features and schema but the same target label. For example, the patient mortality prediction task contains samples from many clinical trials (where input features differ between trials). For T1, the datasets from other tasks T2, T3, . . . are considered out-domain since they differ in prediction objectives. As illustrated by Figure 1, existing methods for tabular prediction fall short in transfer learning across datasets, as each model learns from a single dataset D and needs to learn from scratch when encountering new datasets. On the contrary, MediTab extends the training data to all available tasks T = {T1,T2, . . . }, demonstrating its flexibility to encode and predict for arbitrary tabular samples. After training, it serves all D ∈ T1 without further fine-tuning. Our method eliminates the need to keep as many models as datasets, paving the way for the efficient and streamlined deployment of tabular prediction models. Depending on the use case, the problems that our method can handle can be classified into the following categories.

Problem 1 (Multi-task learning (MTL)). MTL is a machine learning technique where a single model is trained to perform multiple tasks simultaneously. Define f : X ↦→ Y as a model that takes a consolidated tabular sample x as input and predicts the target label y. The training dataset is formed by combining all the tabular inputs in D∗ ∈ T. Once trained, the model f is fixed and can be used to make predictions on any new samples x ∼ D, ∀D ∈ T.

Problem 2 (Zero-shot/Few-shot learning). The model f is trained on T = {D1, . . . ,DN}. Then, it makes predictions for a new dataset DN+1 that has not been included in the training data. Model f performs zero-shot learning if no label is available for all samples in DN+1; Model f performs few-shot learning to predict for DN+1 if a few labeled samples are available.

2.2 The MediTab Framework

As illustrated in Figure 2, our method consists of:

Step 1: Tabular Data Consolidation. The tabular datasets D differ in their features, schema, and particularly in their target objectives if they are from distinct tasks T. The consolidation is accomplished by converting each row of the table into natural language descriptions that consider the data schema. This conversion process transforms all tabular data into text data that share the same semantic space, enabling them to be utilized in language modeling. Additionally, we can produce diverse consolidated samples by describing one sample in multiple different ways, which allows for data augmentation. To prevent hallucinations that may occur during this transformation, an audit module that utilizes LLMs is employed to perform self-check and self-correction. Our goal of patient survival classification is the same for each dataset; however, we use a diverse number of datasets, so the task is indeed different.

Step 2: Learn, Annotate, & Audit. Our method can benefit from out-of-domain datasets T∗ ∈ T through our annotation and data importance pipeline. Once it is trained on T1, it is used to produce pseudo labels for samples from all other tasks, which yields a big but noisy supplementary dataset ˜T1,sup. This dataset is further cleaned by a data audit module based on data Shapley scores, leading to a smaller but cleaner dataset T1,sup.

Step 3: Learning & Deployment. The final prediction model learns from the combination of the original task 1 data T1 and the supplementary data T1,sup. The resulting multitask model fMTL can be used for all datasets D∗ ∈ T1 without any modifications. Furthermore, the model can predict on new datasets D ∈ T in a zero-shot manner and perform few-shot learning for any D ∈ T or D /∈ T. (Note that the primary purpose of the pseudolabels is to facilitate training the zero-shot and few-shot models, and are not meant to improve the performance of the original model.)

We will elaborate on the technical details of these steps in the following sections.

2.3 Tabular Data Consolidation & Sanity Check

Data scarcity is a challenge in scaling medical tabular data predictors. Existing tabular prediction models often struggle on these datasets due to the vague semantic meanings of the varying columns and values. We transform each row into natural language sentences that describe the sample using generative language models like GPT-3.5. Specifically, we combine the linearization, prompting, and santity check steps to obtain input data that the language model can use to generate coherent and meaningful natural language descriptions.

Linearization. A function linearize(c, v) takes the column names c and the corresponding cell values v from a row as the input, then linearizes the row to a concatenation of paired columns and cells {c : v}. Notably, we identify sparse tabular datasets that have many binary columns. In linearization, we exclude binary columns that have a cell value of False and only include those with positive values. This approach has two key benefits. First, it ensures that the linearized output does not exceed the input limit of language models. Second, it helps in reducing hallucinations arising from failed negation detection during the generation process of the LLMs. An ablation on the different types of tabular-totext serialization, shown in Table 10 in the Appendix, shows that audited examples improve performance via augmentation. We believe that this performance benefit is useful, and serves to justify our usage of more advanced paraphrasing and auditing techniques.

Prompting. We combine the linearization with prefix p and suffix s to form the LLM prompt as

(p, linearize(c, v), s). The schema definition is added to p to provide the context for LLM when describing the sample. For each column, we provide the type and explanation as {column}({type}): {explanation} (e.g., "demo1(numerical): the age of the patient in years."). The suffix s represents the instruction that steers LLM to describe the target sample or generate paraphrases for data augmentation. We describe the specific prompt templates we used in Appendix C.1 and display some consolidated examples in Appendix C.2.

The text descriptions x are hence sampled from LLMs by

The paired inputs and target {x, y} will be the training data for the tabular prediction model f. We can adjust the suffix s in Eq. 1 to generate multiple paraphrases of the same sample as a way for instance-level data augmentation. Some instance-level augmentation examples are available in Appendix C.3.

Sanity Check via LLM's Reflection. To reduce low-quality generated samples, a sanity check function evaluates the fidelity of the generated text x to address potential hallucinations or loss of information that occurs during the translation process {c, v} → x in Eq. 1, particularly for numerical features. Specifically, we query LLM with the input template "What is the {column}? {x}" to check if the answer matches the original values in {c, v}. The descriptions are corrected by re-prompting the LLM if the answers do not match. We provide some examples of sanity checks in Appendix C.4, and the quantitative analysis of this correction is available in Appendix C.5.

2.4 Data Enrichment &Refinement

Through the consolidation and sanity check process, we are able to aggregate all tabular samples {x, y} from the target task T1 and train a prediction model. We can use the dataset T1 to train a multi-task learning model, denoted as fMTL, which can be applied to all datasets within the task. Nevertheless, there still lacks a route to leverage data from out-domain tasks T∗ ∼ T , ∀T∗ ∈ T \ {T1} for data enrichment. It is particularly valuable for low-data applications such as healthcare, where there may only be a few dozen data points for each dataset. Specifically, we propose to align out-domain task datasets via a learn, annotate, and audit pipeline for data enrichment.

Learn and Annotate. We train an initial model fMTL on all available training data fromT1 (we will omit the subscript 1 from now on to avoid clutter). The model fMTL then makes pseudo labels for a set of external samples that are retrieved from all other tasks T \ {T}, formulating the noisy supplementary data ˜Tsup = {(xi, ˜ yi)}: xi are consolidated textual description samples and ˜ yi are noisy labels that are aligned with the objective format of the target task.

Quality Audit. It is vital to audit the quality of noisy training data to ensure optimal prediction performance. To this end, we clean ˜T by estimating the data Shapley values for each instance [Ghorbani and Zou, 2019]. We denote the value function by V (S), which indicates the performance score evaluated on the target task T of the predictor trained on training data S. Correspondingly, we have the data Shapley value ϕi for any sample (xi, ˜ yi) ∈ ˜T defined by

where the summation is over all subsets of ˜T not with sample i; n is the number of samples in ˜T; | · | implies to the size of the set; C is an arbitrary constant. Intuitively, ϕi measures the approximated expected contribution of a data point to the trained model's performance. Therefore, a sample with low ϕ is usually of low quality itself.

Computing the exact data Shapley value with Eq. 2 requires an exponentially large number of computations with respect to the number of train data sources. Instead, we follow [Jia et al., 2021; Jia et al., 2019] to use K-Nearest Neighbors Shapley (KNN-Shapley), which offers an avenue for efficient data Shapley computation. Moreover, we are able to achieve a 10× speedup by parallelizing the algorithm, which completes computing scores for 100K+ samples in minutes. Upon acquiring Φ = {ϕi}, we execute a representative stratified sampling corresponding with the distribution of the sample classes to establish the cleaned supplementary dataset Tsup. Appendix G explores the Shapley value and pseudolabel distributions of different supplemental datasets.

2.5 Learning &Deployment

After the quality check step, we obtain the original task dataset T and the supplementary dataset Tsup and have two potential options for model training. The first is to combine both datasets for training, but we have found that this approach results in suboptimal performance. Instead, we employ a two-step training approach: (1) pre-train the model on Tsup, and (2) finetune the model using T. The resulting model will be deployed to provide predictions for any tabular samples belonging to the target task. Because the model trained on the supplementary model has not seen any examples from the original dataset, it can make zero-shot predictions for the test samples from a new dataset D′ /∈ T or adapt to a new task T′ via few-shot learning when a few labeled data are available. Due to the small number of samples in some datasets, we thought it would be best to use all possible training samples in the learning phase, without leaking information in the testing phase. As the label distribution is highly skewed, it may also bias our model if validation samples were chosen randomly, and we did not have the expertise to choose a representative sample. In our case, we chose to simply train the model for 3 epochs on all of the datasets, and then perform a single pass of fine-tuning, without significant hyperparameter optimization, due to the small amount of data and the good performance that it gives.

3 Experiments

We conduct an extensive evaluation of MediTab's performance in supervised learning (Q1), few-shot learning (Q2), and zero-shot prediction (Q3). We also compare the different training strategies for the final deployment of our method (Q4).

3.1 Experimental Setting

Datasets: In our experiments, we introduce the following types of tabular prediction tasks. Patient Outcome Datasets. This dataset includes the patient records collected separately from seven oncology clinical trials 1 . These datasets each have their unique schema and contain distinct groups of patients in different conditions. A CTGAN model [Xu et al., 2019] was trained on the raw data to generate the synthetic patient data for the experiments. We train the model to predict the patient's morbidity, which is a binary classification task. The statistics of the datasets are available in Table 1. Clinical Trial Outcome Datasets. We use clinical trial data from the HINT benchmark [Fu et al., 2022] and ClinicalTrials.gov 2 . The HINT benchmark contains drug, disease, and eligibility information for 17K clinical trials. The trial database contains 220K clinical trials with information about the trial setup (such as title, phase, enrollment, conditions, etc.). Both datasets cover phases I, II, and III trials, but only the HINT benchmark includes the trial outcome labels in {success, failure}. We have also included the MIMIC-IV dataset and PMC-Patients dataset as the external patient database and clinical trial documents as the external trial outcome prediction dataset. Please refer to Appendix D for details.

Implementations: For the patient outcome prediction task, we choose a tree ensemble method (XGBoost) [Chen and Guestrin, 2016a], Multilayer Perceptron, FT-Transformer [Gorishniy et al., 2021], TransTab [Wang and Sun, 2022b], and TabLLM [Hegselmann et al., 2022] as the baselines. For the trial outcome prediction task, we choose XGBoost, feedforward neural network (FFNN) [Tranchevent et al., 2019], DeepEnroll [Zhang et al., 2020], COMPOSE [Gao et al., 2020], HINT [Fu et al., 2022], and SPOT [Wang et al., 2023b] as the baselines. We use PyTrial [Wang et al., 2023a] to implement most baselines and provide the parameter tuning details of the selected baselines in Appendix E.

1 https://data.projectdatasphere.org/projectdatasphere/html/access

2 https://clinicaltrials.gov/

We use a pre-trained bidirectional transformer model named BioBERT [Lee et al., 2020] as the classifier for MediTab. We utilize GPT-3.5 [Brown et al., 2020] via OpenAI's API 3 for the data consolidation and enhancement. We use UnifiedQA-v2-T5 3B [Khashabi et al., 2020] 4 for the sanity check. The evaluation metrics selected are ROC-AUC and PR-AUC, with the details in Appendix F. Further ablations on base model choice is shown in Appendix H. All experiments were run with 2 RTX-3090 GPUs and AMD Ryzen 3970X 32-Core CPU.

3.2 Results on Patient Outcome Prediction and Trial Outcome Prediction

We report the supervised results for patient outcome prediction: the AUROC and PRAUC on the test sets of all clinical trials, in Table 3. Note that we train a single classifier for MediTab and predict on all datasets, while the baselines need to be trained on each dataset separately. Our findings demonstrate that a single MediTab model achieves the highest ranking in 5 out of 7 datasets, with an overall ranking of 1.57 across all datasets. Conversely, MLP and FTTransformer fail to converge in certain cases due to imbalanced target labels (e.g., Lung Cancer 1) or limited availability of data (e.g., Lung Cancer 3). This highlights the data-hungry nature of deep learning algorithms and emphasizes the importance of augmenting training data through data consolidation and enrichment. Additionally, we observe that TabLLM fails in both the single-dataset and multi-dataset settings. We see thatonly the text-template serialization performs poorly in this setting, with multiple datasets not converging. The small amount of data and the clinical-specific terms may be too niche for the general-purpose TabLLM. Furthermore, it is not able to generalize across datasets, as the column names are quite diverse (Table 5).

3 Engine

gpt-3.5-turbo-0301:

https://platform.openai.com/docs/models/gpt-3-5

4 Huggingface: allenai/unifiedqa-v2-t5-large-1363200

MediTab also leads to substantial improvements in trial outcome prediction tasks, as illustrated in Table 4. Notably, our approach outperforms all other methods in every phase of the trials. We observe remarkable improvements of 5.9%, 9.5%, and 3.2% over the previous state-of-the-art baselines in the three phases, respectively. This provides insight into the benefits of increased data availability and the utilization of transfer learning in deep learning-based tabular prediction algorithms.

3.3 Results on Zero-Shot and Few-Shot Learning

We assess the zero-shot prediction capability of MediTab on two tasks. For the evaluation of the dataset D, we deliberately exclude D from the training data during step 2, where pseudo labels are generated for the external database. When computing the data Shapley values for out-domain samples during the quality check process, D is also excluded. Subsequently, we train a model solely on the cleaned supplementary data Tsup and evaluate its performance on the target dataset D. The results of this evaluation are illustrated in Figure 3. MediTab exhibits impressive zero-shot performances: it wins over supervised XGBoost models in 5 out of 7 datasets in patient outcome prediction and all three datasets in trial outcome prediction by a significant margin. On average, MediTab achieves gains of 8.9% and 17.2% improvements in the two tasks, respectively.

The encouraging zero-shot learning result sheds light on the development of task-specific tabular prediction models that can offer predictions for new datasets even before the label collection stage. This becomes particularly invaluable in scenarios where acquiring training labels is costly. For instance, it enables us to predict the treatment effect of a drug on a group of patients before conducting clinical trials or collecting any trial records. Consequently, it allows us to make informed decisions regarding treatment adjustments or trial discontinuation.

We further visualize the few-shot learning results in Figure 4. We witness consistent performance improvement with more labeled training samples for both methods. Additionally, for all tested cases, XGBoost is unable to surpass the zero-shot score of MediTab.

3.4 Results on Ablations on Different Learning Strategies

Section 2.5 discusses a two-stage training strategy for the final learning & deployment stage. Here, we investigate the different training regimens of our method: single-stage training (augment), two-stage training (finetune), training on the original datasets from scratch (scratch), and zero-shot (zeroshot). We list their rankings in Figure 5 and detailed performances across datasets in the Appendix (Tables 7 and 8). Results show that finetune generally performs the best. We conjecture that jointly training on the target task and supplementary data improves the model's overall utility, but it may affect the performance of specific samples in the target task T. Furthermore, we also identify that zeroshot reaches comparable performances with scratch.

Ablations

Figure 5: ROC-AUC Ranking (lower is better) of the variations of MediTab. Left: patient dataset; Right: Trial dataset.

4 RelatedWork

Tabular Prediction has traditionally relied on tree ensemble methods [Chen and Guestrin, 2016b; Ke et al., 2017]. In recent years, the powerful representation learning abilities of neural networks have motivated the new design of deep learning algorithms for tabular prediction [Arik and Pfister, 2021; Kadra et al., 2021; Chen et al., 2023; Bertsimas et al., 2022]. They involve using transformer-based architectures [Huang et al., 2020; Gorishniy et al., 2021; Wang and Sun, 2022a] to enhance automatic feature interactions for better prediction performances. In addition, self-supervised learning (SSL) has been extended to tabular prediction tasks. This includes approaches such as generative pretraining objective by masked cell modeling [Yoon et al., 2020; Arik and Pfister, 2021; Nam et al., 2023], and discriminative pretraining objective by self-supervised [Ucar et al., 2021; Somepalli et al., 2022; Bahri et al., 2022] or supervised contrastive learning [Wang and Sun, 2022b]. Moreover, transfer learning was also adapted to tabular prediction, employing prompt learning based on generative language models [Hegselmann et al., 2022] and multi-task learning [Levin et al., 2023]. Multitask learning and transfer learning were also performed in the medical domain for EHR-based predictive modeling [Hur et al., 2023; Hur et al., 2022]. Nonetheless, these approaches primarily focus on algorithm design, including model architecture and objective functions, often overlooking the engineering of the underlying data.

Data-Centric AI underscores the importance of data for building advanced machine learning prediction systems [Zha et al., 2023]. Notable progress in the domain of tabular data includes efforts to detect [Wang et al., 2020] and debug the noises in labels [Kong et al., 2021]; automate feature selection [Liu et al., 2023]; and streamline feature generation [Su et al., 2021]. Additionally, LM for finetuning on tabular data has been proposed by [Dinh et al., 2022], but it uses strict templates to create the sentence, which limits its expressivity. These methods were proposed for general tabular data while not covering the challenges of heterogeneity and limited samples in medical tabular data. Though there were efforts in enhancing medical codes in EHRs with text descriptions [Hur et al., 2022], there is no further exploration on augmenting medical tabular data that include more diverse features. In contrast, we present a data engineering framework designed to consolidate diverse tabular datasets, by distilling the knowledge from large language models with hallucination detection and distilling from out-domain datasets with data auditing. MediTab is hence able to build a versatile prediction model for the target task.

5 Conclusion

In conclusion, we proposed a novel approach to train universal tabular data predictors for medical data. While there were many efforts in developing new algorithms for tabular prediction, the significance of data engineering has raised much less attention. Specifically, medicine faces challenges in limited data availability, inconsistent dataset structures, and varying prediction targets across domains. To address these challenges, MediTab generates large-scale training data for tabular prediction models by utilizing both in-domain tabular datasets and a set of out-domain datasets. The key component of this approach is a data engine that utilizes large language models to consolidate tabular samples by expressing them in natural language, thereby overcoming schema differences across tables. Additionally, the out-domain tabular data is aligned with the target task using a learn, annotate, and refine pipeline. By leveraging the expanded training data, MediTab can effectively work on any tabular dataset within the domain without requiring further finetuning, achieving significant improvements compared to supervised baselines. Moreover, MediTab demonstrates impressive performance even with limited examples (few-shot) or no examples (zero-shot), remaining competitive with supervised approaches across various tabular datasets. A discussion on the ethical implications and limitations can be found in App. A and App. B. Additionally, further ablation on Data Shapley and Auditing can be found in App. H.

Acknowledgments

This work was supported by NSF award SCH-2205289, SCH-2014438, and IIS-2034479.

Contribution Statement

Zifeng Wang and Chufan Gao contributed equally to this work.