UniTabE: A Universal Pretraining Protocol for Tabular Foundation Model in Data Science

Thank you sincerely for your reviewing and insightful suggestions.

Diversity of downstream datasets:

The selection of downstream evaluation datasets was a deliberate and considered process. Despite the expansive nature of our training data, covering 303 domains with 13 billion samples, ensuring a representative evaluation across all domains is challenging. We intentionally opted for specific downstream datasets in domains(e.g. investing, finance, healthcare, real estate, etc). These domains not only align with our pretraining data but also represent a diverse range of practical applications. The diversity of the downstream datasets is discussed in detail in Appendix A.2 (DATASETS). We acknowledge the importance of a comprehensive evaluation and would like to explore additional diverse tabular datasets to further enhance the breadth of our evaluation scope.

The pretraining methodology:

Our method is clearly different from TransTab and FT-Transformer.

TransTab employs contrastive learning on a limited set of tables from the same clinical domain, lacking the large-scale pretraining on diverse domains and the objective of multi-cell-masking undertaken in our work. Furthermore, TabUnit's cell processing strategy differs significantly from that of TransTab. Our model processes each cell in the model individually, while TransTab divides tabular data into three types and adopts textualization to convert each type of data into text before feeding it into an embedding module.

FT-Transformer, despite sharing a similar cell-wise processing concept, differs in representation and linking layer mechanisms from our approach. Our method involves utilizing a sequence of vectors for each cell, encompassing information on both the column name and cell value, with a specialized linking layer maintaining their relationship. We convert numerical values into a sequence of digits to preserve the original recording structure and numerical significance, while non-numerical values are treated as text and passed into the embedding module. In contrast, FT-Transformer considers only the cell value and represents each cell with a single vector. For numerical values, FT-Transformer acquires cell representations by multiplying them with a learning vector, while non-numerical values are mapped into a singular vector through the embedding module. Importantly, FT-Transformer is not designed for pretraining purposes.

The pretraining objective:

We think that applying the general idea of mask-then-predict is not a weakness of this work. The mask-then-predict idea is a simple but effective approach of providing self-supervised optimization, and is widely used in pretraining, especially for the pretraining in NLP. The application of the mask-then-predict paradigm in this work is not perceived as a weakness but rather as a deliberate choice. Our multi-cell-masking objective is the crafted variance of mask-then-predict. We manipulate the masking at the cell level with the consideration that filling in the entire content of a cell is more useful than merely filling in part of the cell content in practical applications. In addition, this objective can be regarded as learning the overall similarity of the table, while the contrastive learning focuses on learning the inner differences within the same table. This would bring the benefit of obtaining hierarchical understanding of the tabular data, distinguishing it from existing methods.

To Q1:

"UniTabE scratch" is a model trained on the downstream dataset without utilizing pretrained parameters. "UniTabE finetune" represents the model finetuned from pretrained parameters. "UniTabE + XGBoost" denotes the XGBoost model incorporating both the original input and the tabular representation generated by UniTabE, as discussed in the "Learned Feature + XGBoost" paragraph of our paper. In Tables 2 and 3, "UniTabE finetune" exhibits promising improvement compared to "UniTabE scratch." The results in Table 4 underscore the significant performance enhancement achieved through pretraining in the context of zero-shot prediction. Notably, compared to XGBoost, "UniTabE + XGBoost" demonstrates improved performance across the majority of evaluated datasets, indicating that the learned tabular representation from UniTabE contributes positively to XGBoost, even if the improvement on certain datasets is modest.

To Q2:

In the training process for Llama2's tabular predictions, we followed a specific template for both finetuning and predicting. Each example was structured as follows:

"I will provide you with a table. Your task is to predict the classification result based on this table, with options being 0 or 1. <table-description> The subsequent text contains the content of the table in Markdown: <table-content-in-markdown>."

We use the Hugging Face's transformers module as the framework, in which "LlamaForSequenceClassification" is used, for both finetuning and inferencing Llama2.

Thank you for dedicating your time to provide valuable feedback. As the deadline for author-reviewer discussions ends today, we kindly invite you to continue the dialogue if you have any further comments.

Thank you very much for your thorough review, insightful suggestions, and for pointing out the typos!

To Q1:

Data licensing, copyright, and ethics:

Our utilization of Kaggle's dataset for UniTabE pretraining adheres meticulously to Kaggle's terms of use and licensing agreements. The publicly available dataset complies with Kaggle's licensing terms, and any specific requirements imposed by dataset providers have been respected. The dataset was collected through the official Kaggle API (https://github.com/Kaggle/kaggle-api), specifying the license option as "gpl" and "odb". We exclusively collected data under licenses such as GNU General Public License (GPL) and Open Data Commons series (e.g., ODbL, PDDL, ODC-By), ensuring compatibility for open-sourcing our model. For example, the "sales-from-different-stores" dataset is obtained under the ODbL license (https://opendatacommons.org/licenses/odbl/1-0/), allowing a worldwide, royalty-free, non-exclusive license for data use. Our approach aligns with copyright laws, and we uphold ethical considerations, prioritizing responsible and fair data use within permitted scope and respecting Kaggle's data privacy policies. We intend to include information on data licensing, copyright, and ethical considerations in the final version of our paper if it is accepted.

Use of our model as a foundamental model:

1. Understanding tabular data: UniTabE serves as a powerful foundational model for understanding tabular data. It facilitates intuitive obtaintion of tabular representation through natural language prompts, enhancing usability in downstream applications.

2. Tabular data synthesis: UniTabE, trained to predict missing values in tables, holds the potential for synthesizing tabular data. By masking entire rows as missing values, the model can generate synthesized data, contributing to privacy protection by excluding real, sensitive information. Synthesized data further benefits applications such as data augmentation, enhancing diversity, and simulating scenarios.

3. Combining with large language models (LLMs): Inspired by the success of diffusion models in text-to-image tasks, future research may explore utilizing our pretrained model as a foundational component to generate semantic representations for tabular data combining with LLMs within the diffusion framework.

To Q2:

1. Baselines selection for Table 2 and Table 3: The 12 datasets(PID~HPA) in Table 2 (in our paper) come from Kaggle. Despite their exclusion from our pretraining data, we just conducted a preliminary comparison of our approach with TransTab-LSTM (a representative Transformer-based model), XGBoost, and "UniTabE scratch" on these datasets, considering potential domain similarities to pretraining data domains. Additionally, we employed other public datasets (CG~IO) for comprehensive evaluation in Table 3, ensuring a fair comparison with prior methodologies that may not have been trained on Kaggle's data domains.

2. Baselines selection for Table 4 and Table 5: In Table 4, we concentrated on evaluating zero-shot prediction performance. Given the limitations of most existing models that cannot be applied to table structures they were not trained with, hindering their capability for zero-shot prediction, we exclusively compared our model against a finetuned and a randomly initialized model. In Table 5, our objective was to assess the scalability of models in handling incremental columns added to tables(e.g. in clinical trials, incremental columns are collected across different phases). Similar to Table 4, we excluded methods that cannot support table structures different from those encountered during training; for example, TabPFN prohibits a different table structure at the test stage compared to the trained one. We chose TransTab-LSTM as a baseline for its adaptability to incremental columns and its representation as a Transformer-based model. Additionally, we incorporated "UniTabE scratch" as a comparative model to evaluate the benefits derived from pretraining.

To Q3:

We employ TabUnit to process each cell, transforming the column name and value into a sequence of hidden states that serve as the cell's representation. For tables with multi-tiered headers, we adopt the format "top-level-header-name :: second-level-header-name :: ..." to convey the nested structure. The following table presents multi-tiered headers, with actual data omitted for clarity of demonstration:

The column name for cell c11 is "Models", while the column name for cell c12 and c22 is set as "Task 1 :: Precision". Utilizing our feature processing module, TabUnit, our trained model exhibits flexibility in encoding tables with complex structures.

Thank you for dedicating your time to provide valuable feedback. As the deadline for author-reviewer discussions ends today, we kindly invite you to continue the dialogue if you have any further comments.

Thank you very much for your reviewing and insightful suggestions!

Diversity of Downstream Datasets:

Our choice of downstream datasets was deliberate, encompassing diverse domains, including investing, finance, medicine, healthcare, real estate, dress sales, and so on. The selection aimed at providing a comprehensive evaluation across practical domains, as detailed in Table 7.

Performance gap between the competitors:

While competitors are strong, our method consistently outperforms them on most datasets. We acknowledge the modest improvements in certain tasks (e.g., RSP, AD). However, the average score on CG~IO datasets demonstrates a clear improvement against prior approaches.

Ability of dealing with missing values:

In practical applications, modeling tabular data with existing missing values is common. Our proposed model is inherently equipped to effectively address missing values, as it is trained with a mask-the-predict objective. In Table 10, we conducted an evaluation of our model and competitors (including GPT) for filling in both numerical and textual missing values. It's noteworthy that a considerable number of prior methodologies do not accommodate this scenario, especially in predicting textual values, unless explicitly trained for such cases. For example, TabPFN can be additionally trained to predict numerical values, similar to regression tasks, but lacks support for filling in textual values. In contrast, our method can perform this task without requiring further training.

To Q1:

we plan to release our trained model and our code.

To Q2:

We further evaluated the performance of XGBoost in regression tasks and TUTA in classification tasks. Results are as follows.

The 12 datasets(PID~HPA) in Table 2 (in our paper) come from Kaggle. Despite their exclusion from our pretraining data, we just conducted a preliminary comparison of our approach with TransTab-LSTM (a representative Transformer-based model), XGBoost, and "UniTabE scratch" on these datasets, considering potential domain similarities to pretraining data domains. Additionally, we employed other public datasets (CG~IO) for comprehensive evaluation, ensuring a fair comparison with prior methodologies that may not have been trained on Kaggle's data domains.

To Q3:

Most methods faced considerable difficulties in handling datasets with pronounced class imbalances, exemplified by the highly skewed class distribution in the IO dataset (around 1:16 for positive label:negative label). Predicting minority classes accurately is challenging in such imbalanced settings. Our method exhibits a distinct relative improvement of about 7% compared to Transformer-based methods.

To Q4:

Except for the predictive performance, our method has other advantages:

1. Adaptability: Our model exhibits adaptability to diverse table schemata, leveraging TabUnit for versatile processing. The representation of each cell, akin to tokens in natural language processing, allows the application of Transformer-based architecture talented in sequence modeling. This inherent flexibility enables our model to effectively encode representations of cells, making it adaptable to various tables, analogous to pretrained language models handling different texts.

2. Flexibility: The proposed approach is designed for versatility, adapting seamlessly to different tasks within a unified framework. Equipped with a decoder, our model operates in a table-to-sequence format, encompassing classification, regression, and text generation. This flexibility enhances its applicability across a spectrum of scenarios.

3. Robustness: Evaluation on CG~IO datasets outside Kaggle's domains demonstrates our method's continued superior performance. Zero-shot classification experiments (Table 4) further affirm its robustness in out-of-domain scenarios. Notably, our model exhibits robustness in the presence of missing values in tabular data, as demonstrated in Table 10. Moreover, the AUC scores in Table 3 are computed as the average scores across 5 training runs. The mean standard deviation derived from five iterations on all datasets is presented bellow. UniTabE exhibits a lower variance compared to other competitive methods, indicating its consistency and reliability across different runs.

Thank you for dedicating your time to provide valuable feedback. As the deadline for author-reviewer discussions ends today, we kindly invite you to continue the dialogue if you have any further comments.