[Paper Review] TabTransformer: Tabular Data Modeling Using Contextual Embeddings

Xin Huang, Ashish Khetan+2 more|arXiv (Cornell University)|Dec 11, 2020

Domain Adaptation and Few-Shot LearningComputer Science48 references174 citations

TL;DR

TabTransformer uses contextual embeddings from Transformer layers to model tabular data, achieving superior accuracy vs. ML baselines and matching GBDT, with robust handling of missing/noisy data and a two-phase semi-supervised pre-training approach.

ABSTRACT

We propose TabTransformer, a novel deep tabular data modeling architecture for supervised and semi-supervised learning. The TabTransformer is built upon self-attention based Transformers. The Transformer layers transform the embeddings of categorical features into robust contextual embeddings to achieve higher prediction accuracy. Through extensive experiments on fifteen publicly available datasets, we show that the TabTransformer outperforms the state-of-the-art deep learning methods for tabular data by at least 1.0% on mean AUC, and matches the performance of tree-based ensemble models. Furthermore, we demonstrate that the contextual embeddings learned from TabTransformer are highly robust against both missing and noisy data features, and provide better interpretability. Lastly, for the semi-supervised setting we develop an unsupervised pre-training procedure to learn data-driven contextual embeddings, resulting in an average 2.1% AUC lift over the state-of-the-art methods.

Motivation & Objective

• Bridge the performance gap between MLP and gradient-boosted decision trees (GBDT) on tabular data by learning contextual embeddings for categorical features.
• Leverage Transformer-based self-attention to transform column embeddings into contextual representations for improved predictive accuracy.
• Demonstrate robustness to missing and noisy categorical features and provide interpretability of learned embeddings.
• Propose a two-phase semi-supervised learning pipeline (pre-training on unlabeled data, then fine-tuning) to improve performance when labeled data are scarce.

Proposed method

• Embed each categorical feature with a dedicated column embedding table that includes a missing-value embedding.
• Process the sequence of embeddings through N Transformer layers (multi-head self-attention followed by feed-forward blocks).
• Concatenate contextual embeddings from the top Transformer layer with continuous features and feed them to an MLP for final prediction.
• Optionally pre-train Transformer layers on unlabeled data using MLM (masked language modeling) or RTD (replaced token detection) tasks, then fine-tune with labeled data.
• Optimize end-to-end via gradient-based learning to minimize standard supervised loss (cross-entropy for classification, MSE for regression).
• In semi-supervised settings, perform a two-phase workflow: (i) pre-train on unlabeled data, (ii) fine-tune on labeled data.

Experimental results

Research questions

• RQ1 Can Transformer-based contextual embeddings for categorical features outperform traditional MLPs on tabular data?
• RQ2 Do contextual embeddings provide robustness to missing and noisy categorical features compared to baseline neural models?
• RQ3 How does TabTransformer perform relative to tree-based models (GBDT) and other deep tabular models across diverse datasets?
• RQ4 Does a two-phase semi-supervised pre-training/fine-tuning pipeline yield measurable gains in AUC with limited labeled data?

Key findings

• TabTransformer improves over the baseline MLP on 14 of 15 datasets with an average 1.0% AUC gain.
• TabTransformer matches or closely rivals GBDT in supervised learning, outperforming several deep tabular baselines (e.g., TabNet, VIB).
• Contextual embeddings become more predictive across Transformer layers, enabling linear models on embeddings to approach end-to-end performance.
• The model shows robustness to noisy and missing categorical features, outperforming MLP as noise or missingness increases.
• In semi-supervised settings, TabTransformer-RTD/MLM with pre-training yields meaningful AUC gains (up to 2.1% average lift) over competitors when unlabeled data are abundant.