Paper: MediTab: Scaling Medical Tabular Data Predictors via Data Consolidation, Enrichment, and Refinement (ICLR-2024)

Reviewer comment (Experiments):
How are the “External Patient Databases” used (MIMIC-IV and PMC-Patients)? Are they used as supplementary databases during the psuedolabeling step?

Author response:
How are the “External Patient Databases” used (MIMIC-IV and PMC-Patients)? Are they used as supplementary databases during the pseudo-labeling step?

Exactly! In Step 2, MediTab retrieves samples from external databases and tries to annotate and audit the samples to build a clean, augmented dataset. That data will then be involved in enhancing the training of MediTab to obtain better performances.

Reviewer comment (Experiments):
How is the MIMIC-IV data processed? The information in Table 1 that shows that there are only 2 categorical, 1 binary, and 1 numerical feature in MIMIC-IV. As MIMIC-IV is much richer (potentially thousands of features) it is unclear which components of the database are actually used and no details are provided.

Author response:
How is the MIMIC-IV data processed? 

To simplify the setup, we use the simple baseline features of the patients, including their age and gender, and their clinical discharge notes as the external datasets to obtain shapley values and pseudolabels.

Reviewer comment (Theory):
The description of the psuedolabeling step is not entirely clear to me. Is the idea to take the initial model for the target task, make predictions for the target task on data collected for other tasks, and then use those predictions as pseudo-labels for further training? This seems peculiar because it is not clear that this should fundamentally improve performance for the target task given that the pseudo-labels are essentially just predictions of the target label derived from information in the target task database(s).

Author response:
Is the idea to take the initial model for the target task, make predictions for the target task on data collected for other tasks, and then use those predictions as pseudo-labels for further training? 

Yes you are correct! The model is trained on the same outcome prediction (predicting binary patient survival), and we refer to this as different tasks as they come from very different clinical trial datasets. We have clarified this in the text in section 2.2 step 1. We primarily use the pseudolabels to facilitate training the zero-shot (no access to the original data) and are not meant to improve the performance of the original supervised model. Thank you for pointing out this point, and we have clarified this in section 2.2 step 3.

Reviewer comment (Novelty):
The framework is quite straightforward, and there is not much technical contribution. It is mostly a combination of multiple existing models. The idea of transferring tabular data into text is not novel, as there are a number of existing works [1][2][3], including one of their baselines TabLLM [4]. The further incorporation of text information from samples from other datasets is just one trivial step forward. Furthermore, [4] actually proved that a template for transferring the tabular data works better than an LLM. Yet, in this paper, there is no comparison for such serialization methods.

Author response:
Thanks for pointing out the potential confusion about the difference between our method and previous papers. In TabLLM [1], the authors explore a suite of serialization strategies on nine public tabular prediction datasets and find the text template method works the best, better than the table-to-text variant. However, we would like to highlight that we propose two important aspects for the serialization:

We account for the potential *errors* (or *hallucinations*) that may occur during the table-to-text translation process (called *paraphrasing* in this paper), which we found to undermine the performance of the models because wrong information or noises are added to the translated data. We propose a new *auditing* process to control this type of error.
We also propose to distill the knowledge from large language models by augmenting the tabular data in this serialization process via multiple paraphrasing. 

Empirically, we found that both the above two aspects are crucial to the superior performance of MediTab, as shown in the table added to Appendix G Table 10, page 27:

|     Trial Name    | Metric | Simple Text | Praphrase | Audited Paraphrase |
|:-----------------:|:------:|:-----------:|:---------:|:------------------:|
|  Breast Cancer 1  | ROCAUC |    0.607    |   0.620   |        0.617       |
|                   |  PRAUC |    0.098    |   0.107   |        0.105       |
|  Breast Cancer 2  | ROCAUC |    0.753    |   0.753   |        0.876       |
|                   |  PRAUC |    0.083    |   0.083   |        0.135       |
|  Breast Cancer 3  | ROCAUC |    0.760    |   0.758   |        0.764       |
|                   |  PRAUC |    0.452    |   0.481   |        0.476       |
| Colorectal Cancer | ROCAUC |    0.695    |   0.691   |        0.705       |
|                   |  PRAUC |    0.259    |   0.264   |        0.256       |
|   Lung Cancer 1   | ROCAUC |    0.699    |   0.737   |        0.717       |
|                   |  PRAUC |    0.975    |   0.979   |        0.972       |
|   Lung Cancer 2   | ROCAUC |    0.699    |   0.697   |        0.716       |
|                   |  PRAUC |    0.679    |   0.680   |        0.715       |
|   Lung Cancer 3   | ROCAUC |    0.607    |   0.893   |        0.929       |
|                   |  PRAUC |    0.697    |   0.957   |        0.968       |


This work builds upon TabLLM by including data auditing of the paraphrased results as well as using external datasets for medical tabular prediction. 

Our results indicated that the basic approach of using "column name: column value," similar to what's demonstrated in TabLLM, was effective. However, we also observed that enhancing this format with paraphrased examples led to better performance. Furthermore, we find that audited examples improve performance the most. We believe that this performance benefit is useful and serves to justify our usage of more advanced paraphrasing and auditing techniques to address our first 2 points of addressing model hallucinations and data augmentation.

[1] Hegselmann S, Buendia A, Lang H, et al. Tabllm: Few-shot classification of tabular data with large language models[C]//International Conference on Artificial Intelligence and Statistics. PMLR, 2023: 5549-5581.

Reviewer comment (Experiments):
The author didn’t specify what exact features are included in these experimental datasets. It is also unclear how many columns are overlapped between different datasets. If there is a large portion of feature overlapping, simple concatenation and removal or recoding of the missing columns might work just as well. There is no discussion regarding this aspect.

Author response:
We have added more information regarding the dataset columns in Table 5, Appendix C.2, Page 17 and a discussion. 

|     Trial Name    | # Vars Shared | # Vars Total |
|:-----------------:|:-------------:|:------------:|
|  Breast Cancer 1  |       4       |      15      |
|  Breast Cancer 2  |       15      |      48      |
|  Breast Cancer 3  |       7       |      18      |
| Colorectal Cancer |       5       |      16      |
|   Lung Cancer 1   |       1       |      17      |
|   Lung Cancer 2   |       6       |      41      |
|   Lung Cancer 3   |       7       |      26      |

The columns are all quite diverse in terms of semantic meaning, and we believe that traditional methods like data imputation or renaming/removing would not apply here, as we have fundamentally different features.

Reviewer comment (Writing):
The Step 2 in section 2.2 is confusing, as the authors claimed that they used active learning. It is unclear whether the “active learning pipeline” method is the same as traditional active learning that selects informative samples to label. If not, the description can mislead readers.

Author response:
Thank you for pointing this out! Step 2 (Page 3) in the pipeline is conceptually similar to the active learning process. First, the model retrieves samples from other datasets that have different target variables to predict. Second, the model annotates the retrieved samples to build a noisy-label dataset. Third, the model runs the data audit process to select the samples that are most likely to be annotated correctly. Nonetheless, we agree that this was not the best choice of wording and have rephrased the sentence.

Reviewer comment (Experiments):
In Step 2, the authors claimed that they cleaned the supplementary dataset T_{1, sup} with a data audit module based on data Shapley scores. More experiments are expected to demonstrate the effectiveness of the audit module. Moreover, it would be better if the authors conducted more ablation studies to show whether the supplementary dataset improves the prediction performance.

Author response:
The external datasets were annotated with the pre-trained model and audits with our pipeline. Then, these clean pseudo-label data are used to fine-tune a model for the zero-shot and few-shot settings. 

Data importance ablation (Please see the Appendix G Table 11 on page 29! ): after each percentile was calculated, a single epoch of training was performed to obtain the ROC-AUC and PR-AUC. We see that empirically, using the 50 percentile cutoff performs the best. A small cutoff allows in too many irrelevant examples, and a high cutoff may remove too much diversity from the data.

| Trial Name        | Metric | 10\%  | 25\%  | 50\%  | 90\%  |
|-------------------|--------|-------|-------|-------|-------|
| Breast Cancer 1   | ROCAUC | 0.485 | 0.505 | 0.500 | 0.445 |
|                   | PRAUC  | 0.072 | 0.082 | 0.092 | 0.061 |
| Breast Cancer 2   | ROCAUC | 0.591 | 0.426 | 0.728 | 0.656 |
|                   | PRAUC  | 0.275 | 0.025 | 0.082 | 0.168 |
| Breast Cancer 3   | ROCAUC | 0.476 | 0.518 | 0.525 | 0.517 |
|                   | PRAUC  | 0.186 | 0.207 | 0.234 | 0.210 |
| Colorectal Cancer | ROCAUC | 0.517 | 0.497 | 0.489 | 0.466 |
|                   | PRAUC  | 0.139 | 0.132 | 0.135 | 0.118 |
| Lung Cancer 1     | ROCAUC | 0.429 | 0.458 | 0.535 | 0.391 |
|                   | PRAUC  | 0.931 | 0.929 | 0.898 | 0.917 |
| Lung Cancer 2     | ROCAUC | 0.514 | 0.536 | 0.573 | 0.512 |
|                   | PRAUC  | 0.472 | 0.452 | 0.439 | 0.438 |
| Lung Cancer 3     | ROCAUC | 0.464 | 0.500 | 0.671 | 0.607 |
|                   | PRAUC  | 0.766 | 0.745 | 0.812 | 0.765 |

Reviewer comment (Experiments):
The datasets in Table 1 contain fewer than 3000 patients, making it very easy for the LLMs (e.g., BioBERT) to overfit the training set. It is unclear how the authors prevent overfitting during the fine-tuning phase.

Author response:
The sample size is small due to the nature of clinical trials, where even a phase III trial may have at most 1K patients, and phase I and II trials usually have less than 100. That is the reason why we propose MediTab to augment the model with external datasets, e.g., from EHRs. Technically, we made train/val/test splits to make sure there was no data leakage. In experiments, we choose the model with the best performance on the validation set. The performance on the test set hence shows the generalization capability. Since the datasets are not large, we decided to only train the model for a small number of epochs to not overfit: with 3 epochs for pre-training and 1 epoch for fine-tuning. We have added the discussion to section 2.5 page 5.

Reviewer comment (Experiments):
In Table 3, the proposed MediTab exhibits the capability to access multiple datasets during its training, in contrast to the other baseline models, which are constrained to employing a single dataset. This discrepancy in data utilization introduces an element of unfairness in the comparison. It would be more appropriate to conduct a comparison against models that have undergone training on multiple datasets. For instance, TabLLM, being a large language model, can readily undertake multi-dataset training with minor adjustments to its data preprocessing procedures. Therefore, a more equitable comparison would involve evaluating MediTab and TabLLM under identical conditions, both in the context of training on a single dataset and across multiple datasets.

Author response:
We agree that TabLLM is capable of pretraining across multiple datasets like our method, but we also believe in preserving the integrity of its original design, as it was not created for the purpose of training across multiple datasets. However, we have performed an ablation regarding single dataset training, and we see that TabLLM struggles to perform on our domain-specific clinical datasets.

Reviewer comment (Experiments):
Most medical data, like MIMIC-IV, includes timestamp information of patients’ multiple visits or collections. This framework completely ignores this part of the medical data, which limits its application to real-world clinical environments.

Author response:
We agree that extending to longitudinal data is crucial. Since the focus of this work is tabular data prediction, we will explore it in future work. However, while we do not explicitly consider sequential patient data in this work, only traditional tabular data, we believe that this is a logical next step. In essence, we can structure sequential visits into longer textual descriptions, which can also be handled by BERT models. We believe that with appropriate adaptation, the ideas of paraphrasing, data auditing, and enrichment can also shine for sequential data with MediTab model.

Reviewer comment (Experiments):
Are there any overlaps of columns between the tabular data for the same tasks? Is it hard to do a simple concatenation? What’s the traditional method for dealing with missing columns? Are they applicable to this situation?

Author response:
We have added more information regarding the dataset columns in Table 5, Appendix C.2, Page 17 and a discussion. 

|     Trial Name    | # Vars Shared | # Vars Total |
|:-----------------:|:-------------:|:------------:|
|  Breast Cancer 1  |       4       |      15      |
|  Breast Cancer 2  |       15      |      48      |
|  Breast Cancer 3  |       7       |      18      |
| Colorectal Cancer |       5       |      16      |
|   Lung Cancer 1   |       1       |      17      |
|   Lung Cancer 2   |       6       |      41      |
|   Lung Cancer 3   |       7       |      26      |

The columns are all quite diverse in terms of semantic meaning, and we believe that traditional methods like data imputation or renaming/removing would not apply here, as we have fundamentally different features.

Reviewer comment (Experiments):
For the choice of BioBERT and the QA model for salinity check, the author did not provide a reason for choosing these models.

Author response:
We have added ablations regarding the choice of downstream model (BioBERT vs. ClinicalBERT vs. BERT), as well as further comparisons with TabLLM. Please see the Appendix G Table 9 on page 28! 

In detail, we compared NEW ablations of different base models (BERT, BioBERT, ClinicalBERT, and TabLLM) in terms of downstream performance, trained from scratch, respectively, as shown in the table below, about their average performance across seven datasets.

| Trial Name        | Metric | BERT  | ClinicalBERT | BioBERT | TabLLM |
|-------------------|--------|-------|--------------|---------|--------|
| Breast Cancer 1   | ROCAUC | 0.588 | 0.581        | 0.591   | -      |
|                   | PRAUC  | 0.097 | 0.082        | 0.094   | -      |
| Breast Cancer 2   | ROCAUC | 0.485 | 0.724        | 0.803   | -      |
|                   | PRAUC  | 0.023 | 0.026        | 0.060   | -      |
| Breast Cancer 3   | ROCAUC | 0.696 | 0.734        | 0.721   | 0.616  |
|                   | PRAUC  | 0.392 | 0.366        | 0.437   | 0.302  |
| Colorectal Cancer | ROCAUC | 0.613 | 0.700        | 0.705   | -      |
|                   | PRAUC  | 0.233 | 0.186        | 0.267   | -      |
| Lung Cancer 1     | ROCAUC | 0.555 | 0.479        | 0.699   | -      |
|                   | PRAUC  | 0.962 | 0.949        | 0.971   | -      |
| Lung Cancer 2     | ROCAUC | 0.544 | 0.616        | 0.711   | 0.619  |
|                   | PRAUC  | 0.483 | 0.616        | 0.691   | 0.562  |
| Lung Cancer 3     | ROCAUC | 0.357 | 0.893        | 0.893   | 0.804  |
|                   | PRAUC  | 0.695 | 0.957        | 0.957   | 0.826  |


We observed that the model selection choice is similar for BioBERT and ClinicalBERT. However, TabLLM does not converge for 4 out of the 7 datasets. We believe this may be due to the small amount of training data, but further research should be done to investigate this behavior fully. 

In the second experiment, we performed a NEW ablation on the effect of different serialization strategies. We’ve added the results to Appendix G Table 10, page 27, in the new version.

|     Trial Name    | Metric | Simple Text | Praphrase | Audited Paraphrase |
|:-----------------:|:------:|:-----------:|:---------:|:------------------:|
|  Breast Cancer 1  | ROCAUC |    0.607    |   0.620   |        0.617       |
|                   |  PRAUC |    0.098    |   0.107   |        0.105       |
|  Breast Cancer 2  | ROCAUC |    0.753    |   0.753   |        0.876       |
|                   |  PRAUC |    0.083    |   0.083   |        0.135       |
|  Breast Cancer 3  | ROCAUC |    0.760    |   0.758   |        0.764       |
|                   |  PRAUC |    0.452    |   0.481   |        0.476       |
| Colorectal Cancer | ROCAUC |    0.695    |   0.691   |        0.705       |
|                   |  PRAUC |    0.259    |   0.264   |        0.256       |
|   Lung Cancer 1   | ROCAUC |    0.699    |   0.737   |        0.717       |
|                   |  PRAUC |    0.975    |   0.979   |        0.972       |
|   Lung Cancer 2   | ROCAUC |    0.699    |   0.697   |        0.716       |
|                   |  PRAUC |    0.679    |   0.680   |        0.715       |
|   Lung Cancer 3   | ROCAUC |    0.607    |   0.893   |        0.929       |
|                   |  PRAUC |    0.697    |   0.957   |        0.968       |


Our results indicated that the basic approach of using "column name: column value," similar to what's demonstrated in TabLLM, was effective. However, we also observed that enhancing this format with paraphrased examples led to better performance. Furthermore, we find that audited examples improve performance the most. We believe that this performance benefit is useful and serves to justify our usage of more advanced paraphrasing and auditing techniques to address our points of addressing model hallucinations and data augmentation. For the QA model, we chose to use the most powerful version of UnifiedQA, a popular and effective model for general QA tasks, that we could reasonably run on all the paraphrased datasets. Although we explored using larger models such as LlaMA2, we found that running these models would take too much time for it to be practical.

We hope that these revisions help address some of your concerns regarding the paper, and look forward to further discussion!

Reviewer comment (Experiments):
It is not clear, how was splitting into train/val/test organised?

Author response:
Thank you for pointing this out, as this was accidentally left out during paper editing. We have added this info to the dataset description in Tables 1 and 2.

Reviewer comment (Experiments):
The statement “After the quality check step, we obtain the original task dataset $T$ and the supplementary dataset $T_{sup}$ 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.” is a bit unclear. Why it is suboptimal, could authors elaborate on this?

Author response:
Of course, we found that it was suboptimal due to the empirical lower performance obtained via simple pretraining on both the supplemental and the original task dataset, as opposed to the approach we used in the paper, which is the pretrain on the supplemental + model’s noisy labels and then finetune on the original task dataset with the real labels. We hypothesize that this is due to the much larger amount of supplemental data vs the real data.

Reviewer comment (Experiments):
Are these datasets prone to missing values, which is a critical concern in the medical domain? If so, what would be the recommended strategy for handling these missing values?

Author response:
In tabular prediction, it is a common practice to deal with missing values in the inputs—the same is true in clinical patient outcome prediction. The proposed method can conveniently deal with missing values: it serializes a row in tabular data into sentences and models on the textual description of the sample.

Reviewer comment (Experiments):
Results on Ablation studies on Different Learning strategies are provided. Could authors provide Ablation studies on the different model components?

Author response:
This is a great point! We have added ablations regarding the choice of downstream model (BioBERT vs. ClinicalBERT vs. BERT), as well as further comparisons with TabLLM. Please see Appendix G Table 9 on page 27! 

In detail, we compared NEW ablations of different base models (BERT, BioBERT, ClinicalBERT, and TabLLM) in terms of downstream performance, trained from scratch, respectively, as shown in the table below, about their average performance across seven datasets.

| Trial Name        | Metric | BERT  | ClinicalBERT | BioBERT | TabLLM |
|-------------------|--------|-------|--------------|---------|--------|
| Breast Cancer 1   | ROCAUC | 0.588 | 0.581        | 0.591   | -      |
|                   | PRAUC  | 0.097 | 0.082        | 0.094   | -      |
| Breast Cancer 2   | ROCAUC | 0.485 | 0.724        | 0.803   | -      |
|                   | PRAUC  | 0.023 | 0.026        | 0.060   | -      |
| Breast Cancer 3   | ROCAUC | 0.696 | 0.734        | 0.721   | 0.616  |
|                   | PRAUC  | 0.392 | 0.366        | 0.437   | 0.302  |
| Colorectal Cancer | ROCAUC | 0.613 | 0.700        | 0.705   | -      |
|                   | PRAUC  | 0.233 | 0.186        | 0.267   | -      |
| Lung Cancer 1     | ROCAUC | 0.555 | 0.479        | 0.699   | -      |
|                   | PRAUC  | 0.962 | 0.949        | 0.971   | -      |
| Lung Cancer 2     | ROCAUC | 0.544 | 0.616        | 0.711   | 0.619  |
|                   | PRAUC  | 0.483 | 0.616        | 0.691   | 0.562  |
| Lung Cancer 3     | ROCAUC | 0.357 | 0.893        | 0.893   | 0.804  |
|                   | PRAUC  | 0.695 | 0.957        | 0.957   | 0.826  |


We observed that the model selection choice is similar for BioBERT and ClinicalBERT. However, TabLLM does not converge for 4 out of the 7 datasets. We believe this may be due to the small amount of training data, but further research should be done to investigate this behavior fully.