TabTransformer: Tabular Data Modeling Using Contextual Embeddings

In plain words

Spreadsheets full of categories, such as job title, country or product type, are still most often handled by tree ensembles like gradient-boosted trees. Neural networks usually turn each category into a fixed vector and treat the vectors separately, so the word "manager" means the same thing in every row.

TabTransformer lets the vectors of one row talk to each other. It runs the category vectors through the same attention layers that power language models, so each value's vector ends up shaped by the other values in that row. Those context-aware vectors, plus the plain numbers in the row, go to a small network that makes the prediction.

The authors also show how to use unlabeled rows. The attention layers are first trained on a fill-in-the-blank or spot-the-swapped-value task that needs no labels, and then fine-tuned on the few labeled rows. This matters when labels are scarce and unlabeled data is plentiful.

TL;DR

• Architecture: per-column category embeddings, a stack of self-attention Transformer layers, then an MLP over the concatenated contextual embeddings and continuous features §2 eq. 1.
• Supervised: on 15 binary datasets it reports 82.8 ± 0.4 mean AUC, 1.0 above an MLP and level with GBDT (82.9 ± 0.4) tab. 2.
• Semi-supervised: pretraining with replaced-token detection or masking beats the paper's semi-supervised baselines on its 8 largest datasets when labels are 50 to 500 rows tab. 3.
• Other claims (robustness to noise and missing values, interpretability) are shown by plots and are not quantified in the text fig. 2 fig. 3 fig. 4 fig. 5.
• Caveat: the paper's own counts of datasets are inconsistent in places; see Pitfalls.

Problem

Tree ensembles such as GBDT are the strongest standard choice for tabular data, but they do not suit streaming training, end-to-end training with image or text encoders, or the semi-supervised methods that work for neural networks §1. Plain MLPs train end-to-end and handle some of these, but they use context-free, per-feature embeddings, are not interpretable, are not robust to missing or noisy values, and lag GBDT on most datasets §1. The experiments cover binary classification only, with 5 cross-validation splits per dataset §3.

Method

The model has three parts: a column embedding layer, N Transformer layers and a top MLP §2 fig. 1.

1. Column embedding. Each categorical column i with d_i classes gets its own lookup table with d_i + 1 rows; the extra row stands for a missing value §2. The embedding of class j in column i is the concatenation of a column-identifier vector c_i (length ℓ) and a class-specific vector w_ij (length d − ℓ) §2. There are no positional encodings, because the columns have no order §2. Continuous features are not embedded §2.

2. Contextual layers. Each Transformer layer is multi-head self-attention followed by a position-wise feed-forward block that first expands the embedding to four times its size and then projects it back §2. Residual addition and layer normalization are applied after each layer §2. The attention weights are softmax(QKᵀ/√k) and the output is that weight matrix times V §2 eq. 2. The output of the top layer is one contextual embedding h_i per column, each of dimension d §2.

3. Prediction head. The contextual embeddings h_1 … h_m are concatenated with the continuous features x_cont into a vector of size d·m + c, which an MLP g_ψ maps to the prediction §2. The loss is H(g_ψ(f_θ(E_φ(x_cat)), x_cont), y), where H is cross-entropy for classification §2 eq. 1. All parameters (column embeddings φ, Transformer θ and head ψ) are trained end-to-end with first-order gradient methods §2.

4. Semi-supervised pretraining. The Transformer layers are first trained on unlabeled rows, then fine-tuned with the supervised loss on the labeled rows §2. The column embeddings are trained together with the Transformer during pretraining §2. Two pretraining objectives are used:

   • MLM (TabTransformer-MLM). Choose k% of the categorical features in each row and set them to the missing-value embedding. A multi-class classifier predicts the original class of each masked feature from its contextual embedding, trained with cross-entropy §2. Whether that classifier is shared across columns is not stated; the per-column classifier is described for RTD only §2.
   • RTD (TabTransformer-RTD). Replace the chosen values with a random value of the same feature instead of masking them. A binary classifier, one per column rather than shared, predicts whether each feature was replaced §2 app. A. Unlike ELECTRA's generator-based replacement, the text describes a plain random value §2.
   • The replacement rate is k = 30 §2. The paper's ablation tries 15, 30 and 50 app. A.

# TabTransformer forward pass  [§2] [eq. 1]
# Input:  cat_ids (B, m) integer class ids; cont (B, c) numeric features
# Output: logits (B, classes)

def forward(cat_ids, cont):
    # Look up each column's class vector and prepend its column identifier
    # (concatenation, not addition; no positional encoding)  [§2]
    embeddings = column_embedding(cat_ids)              # (B, m, d)
    # Stacked self-attention layers make each column's vector depend on
    # the other columns of the row  [§2] [eq. 2]
    contextual = transformer_layers(embeddings)         # (B, m, d)
    # Join the flattened contextual vectors to the numeric features  [§2]
    joined = concat(flatten(contextual), cont)          # (B, d*m + c)
    # Top MLP makes the prediction  [§2] [eq. 1]
    return top_mlp(joined)

# Pretraining on unlabeled rows, MLM variant  [§2]
# Input:  cat_ids (B, m) from unlabeled data; rate = 0.30

def mlm_loss(cat_ids):
    # Choose k% of the features in each row  [§2]
    masked = sample_mask(cat_ids.shape, rate=0.30)      # bool (B, m)
    # Set the chosen features to the missing-value embedding  [§2]
    corrupted = where(masked, MISSING_ID, cat_ids)
    contextual = transformer_layers(column_embedding(corrupted))
    # Predict the original class of each masked feature  [§2]
    logits = class_head(contextual)
    return cross_entropy(logits[masked], cat_ids[masked])

# Pretraining on unlabeled rows, RTD variant  [§2] [app. A]
# Input:  cat_ids (B, m) from unlabeled data; rate = 0.30

def rtd_loss(cat_ids):
    # Choose k% of the features in each row  [§2]
    replaced = sample_mask(cat_ids.shape, rate=0.30)    # bool (B, m)
    # Replace each chosen value with a random value of its own column  [§2]
    random_ids = random_class_per_column(cat_ids)       # (B, m)
    corrupted = where(replaced, random_ids, cat_ids)
    contextual = transformer_layers(column_embedding(corrupted))
    # One binary classifier per column says "replaced or not"  [§2] [app. A]
    logits = per_column_heads(contextual)               # (B, m)
    return binary_cross_entropy(logits, replaced)

Fine-tuning then fine-tunes the pretrained Transformer layers together with the top MLP, using the supervised loss above on the labeled rows §2. The text does not say how the top MLP is initialized.

Hyperparameters

Results

Author-reported. Setup: 15 binary datasets from the UCI repository, the AutoML Challenge and Kaggle, 5 cross-validation splits, mean AUC §3. Hardware: 8 CPUs, 60 GB RAM and one NVIDIA V100; GBDT via LightGBM, other models in PyTorch; the baselines were re-implemented by the authors app. B.1.

Reading the semi-supervised rows: the paper's margin statement (at least 1.2, 2.0 and 2.1 points over the others) is computed from the better of RTD and MLM at each label count tab. 3 §3. Below 30K rows, MLP (ER) is best at 50 labels, and RTD is best at 200 and 500 labels tab. 4 §3.

Paper vs code

No official code.

Checked: the arXiv abstract page lists no code link. The TeX source has no repository link. It refers to a "Code and Data Appendix" as supplementary material app. B.1, which was not found in the fetched source. GitHub search for "TabTransformer", restricted to the awslabs and amazon-science organisations, returned nothing. A title-word search for tabular contextual embeddings returned only unrelated owners. Third-party TabTransformer implementations were not used.

Without code, every implementation detail below is the paper's own statement. The unspecified items are listed under Pitfalls.

Pitfalls

• Dataset count is inconsistent. Abstract and §3 say 15 datasets, but the semi-supervised tables cover 20 datasets. The five extra are jannis, volkert, htru2, fabert and sylvine app. C tab. 3 tab. 4. The text also says the second semi-supervised group has the "remaining 9 datasets" §3. The table captions say 8 and 12, and the 12 is consistent with the 20-dataset list. Implementers should treat the 15-dataset supervised set and the 20-dataset semi-supervised set as separate benchmarks.
• Volkert positive-class rate differs between appendix tables. The dataset table gives 12.7 % app. C. The 50-label per-dataset table gives 1.0 %, and the 200-label table gives 12.7 % app. C. Which value the experiments used is not stated.
• Margin claim uses the better of two variants. The "at least 1.2 points" margin at 50 labels is MLM at 66.8 over MLP (ER) at 65.6 tab. 3. RTD alone is 1.0 points above it. Read the claim as best-of-two, not per variant.
• HPO budget is described two ways. §3 says each model gets 20 HPO rounds per split and the TabTransformer size is fixed §3. Appendix B.1 says the fixed size was chosen from 50 HPO rounds on 5 datasets app. B.1. The search spaces for width, depth and heads are given there app. B.1.
• Pretraining hyperparameters are missing. Epochs, batch size and learning rate for the pretraining phase are not stated. The fine-tuning learning rate is only given as a search range app. B.1.
• Replacement distribution is unspecified. RTD replaces a value with "a random value of that feature" §2. Whether the draw is uniform over observed classes, includes the missing class, or can equal the original value is not stated. The ablation resamples the replacement each epoch ("dynamic"), which the authors found better than a fixed draw app. A.
• MLM head sharing is unspecified. The text describes a per-column head only for RTD §2 app. A. For MLM it says "a multi-class classifier" without saying whether one head serves all columns.
• Missing values at test time. The model has a learned missing-value embedding per column. The benchmarks did not have enough missing values to train it well, so the missing-value experiments impute with the column's average embedding instead §3.2. An implementer must choose which to use.
• Continuous features are not transformed by the paper. Continuous features bypass the Transformer and are concatenated to the contextual embeddings §2. The appendix lists rescaling options (quantiles, normalization, log) and says "the best strategy is likely to use all ... in parallel", but says this was not fully explored app. B.1. The scaling used in the reported runs is not stated.
• Layer normalization placement is unspecified. The text says addition and layer normalization happen after each layer §2. It does not say whether this is post-norm or pre-norm, or where dropout is applied.
• Attention key size is unspecified. The formula uses a key dimension k with a 1/√k scale §2 eq. 2. The per-head key and value sizes are not stated.
• Semi-supervised fine-tuning limit. The authors say the fine-tuned linear head is learned only from labeled rows, so the method loses to ER and PL with 50 labels §3.4. They offer this as the reason and as a direction for future work §3.4.
• Robustness and interpretability are qualitative. The text describes noise and missingness results and the t-SNE and probe plots §3.2 fig. 2 fig. 3 fig. 4 fig. 5. No numeric values are given in the text.
• Pretraining is semi-supervised only. The paper does not pretrain in the fully labeled setting, because no gain was observed there §3 §3.3.
• Regression is not evaluated. The method section mentions mean squared error for regression §2. All experiments are binary classification §3.

Lineage

• Builds on: the Transformer §2; masked language modeling as used in BERT §2; replaced-token detection from ELECTRA §2. All three are cited in the paper.
• Compared against: LightGBM GBDT and logistic regression §3; MLP and a sparse MLP with pruning §3; TabNet §3; the variational information bottleneck model §3; pseudo-labeling and entropy regularization on MLP, TabTransformer and GBDT, and a swap-noise denoising autoencoder §3 app. B.1.
• Superseded by: not stated in this card.