PaperlayerView exact revision on arXiv

Source-linked AI summary

SAINT: Improved Neural Networks for Tabular Data via Row Attention and Contrastive Pre-Training

Gowthami Somepalli, Micah Goldblum, Avi Schwarzschild, C. Bayan Bruss, Tom Goldstein

arXiv:2106.01342v1cs.LGcs.AIstat.ML

TL;DR

Tabular data’s heterogeneous and unordered structure poses challenges for deep learning, despite its importance and the dominance of tree-based methods. SAINT combines feature and row attention, enhanced embeddings, and contrastive pre-training, and it outperforms evaluated deep and tree-based baselines on average across benchmark tasks.

  • Problem

    Tabular data mixes heterogeneous feature types without inherent positional order, while deep learning methods have lagged behind widely used tree-based approaches.

  • Method

    SAINT projects categorical and continuous features into a shared embedding space, applies self-attention and intersample row attention, and uses contrastive pre-training with CutMix and mixup views.

  • Results

    SAINT variants outperform all baselines on average across supervised and semi-supervised benchmark tasks, often surpassing XGBoost, CatBoost, and LightGBM.

  • Takeaways & Limitations

    SAINT shows that neural tabular models can often improve upon boosting methods across datasets with varying characteristics.

  • Takeaways & Limitations

    The benchmarks may not represent noisy or imbalanced real-world datasets, and SAINT was tuned for the tested settings.

Abstract

from arXiv · show

Tabular data underpins numerous high-impact applications of machine learning from fraud detection to genomics and healthcare. Classical approaches to solving tabular problems, such as gradient boosting and random forests, are widely used by practitioners. However, recent deep learning methods have achieved a degree of performance competitive with popular techniques. We devise a hybrid deep learning approach to solving tabular data problems. Our method, SAINT, performs attention over both rows and columns, and it includes an enhanced embedding method. We also study a new contrastive self-supervised pre-training method for use when labels are scarce. SAINT consistently improves performance over previous deep learning methods, and it even outperforms gradient boosting methods, including XGBoost, CatBoost, and LightGBM, on average over a variety of benchmark tasks.

1 Introduction

Tabular data’s heterogeneous, unordered structure has limited deep learning’s transfer from vision and language. SAINT addresses these challenges with hybrid feature- and row-level attention, improved embeddings, and contrastive pre-training.

  • Tabular datasets combine continuous, categorical, and ordinal features whose relationships may be independent or correlated.
  • Because tabular columns have no inherent positional order, methods designed for sequential data do not transfer directly.
  • SAINT projects categorical and continuous features into a shared dense space and applies self-attention within samples plus intersample attention across rows.
  • Intersample attention relates each row to other rows using an end-to-end learned distance, resembling nearest-neighbor classification.
  • Across diverse datasets, SAINT averages better than other evaluated methods on supervised and semi-supervised tasks and often surpasses XGBoost, CatBoost, and LightGBM.

2 Related Work

Related work spans classical tree-based models, emerging deep tabular architectures, row-and-column attention, and self-supervised objectives. SAINT combines these directions while extending attention to continuous features and contrastive learning for tabular data.

  • Classical Models: Classical tabular practice favors non-parametric tree-based models such as XGBoost, CatBoost, and LightGBM in complex settings.
  • Deep Tabular Models: Recent deep tabular models include TabNet, VIME, and TABERT, but they use specialized feature selection, denoising, or language-oriented settings.
  • Deep Tabular Models: TabTransformer applies self-attention only to categorical features, leaving correlations between categorical and continuous features outside its attention block.
  • Deep Tabular Models: SAINT projects both categorical and continuous features into the transformer embedding space and lets data points attend to one another.
  • Axial Attention: Prior axial-attention work applies separate row and column attention in localized or structured inputs, including protein sequences and tables.
  • Self-Supervised Learning: For tabular self-supervision, prior objectives include masking, denoising, and replaced token detection; SAINT adds contrastive learning and couples it with denoising.

3 Self-Attention and Intersample Attention Transformer (SAINT)

SAINT embeds heterogeneous tabular features as tokens and stacks self-attention with intersample attention. Its row-level block communicates across batch samples, producing contextual representations for downstream prediction or self-supervision.

  • Encoding the Data: Each data sample receives a learned [CLS] token alongside its categorical or continuous feature tokens.
  • Encoding the Data: SAINT uses an embedding layer that maps each feature into a d-dimensional space, potentially with different embedding functions for different features.
  • Architecture: The architecture stacks L stages, each containing one standard self-attention transformer block and one intersample attention transformer block.
  • Architecture: The resulting contextual representation can support self-supervision or classification downstream.
  • Intersample attention: Intersample attention computes attention across rows after concatenating a sample’s feature embeddings, allowing rows in the batch to communicate.
  • Architecture: The pipeline applies multi-head self-attention to feature tokens and reshapes inputs so intersample attention operates across samples before returning contextual representations.

4 Pre-training & Finetuning

SAINT pre-trains on augmented tabular views using contrastive and denoising objectives, then fine-tunes on labeled prediction data. Its augmentations combine CutMix in input space with mixup in embedding space, while projection heads support contrastive learning.

  • Pre-training: CutMix augments samples in input space, while mixup creates additional views in embedding space.CutMix first mixes raw feature values using a binary mask; mixup is then performed after embedding.
  • Pre-training: SAINT passes clean and mixed embeddings through the network and two projection-head MLPs before computing contrastive loss.The projection heads reduce dimensionality before contrastive-loss computation.
  • Pre-training: Contrastive learning makes representations of two views from the same data point close and different points far apart.
  • Pre-training: The pre-training objective combines contrastive and denoising losses, with feature-specific cross-entropy or mean squared error used for reconstruction.The reconstruction loss depends on whether each feature is categorical or continuous.
  • Finetuning: After pre-training on unlabeled data, SAINT fine-tunes on labeled samples and predicts from the contextual embedding of the [CLS] token.

5 Experimental Evaluation

The evaluation tests SAINT across diverse tabular datasets, model variants, baselines, supervision settings, robustness conditions, and attention visualizations. SAINT performs strongly across supervised and semi-supervised tasks, while its components show distinct advantages under specific data conditions.

  • Experimental setup: SAINT is evaluated on 16 diverse datasets spanning binary and multiclass classification, with 200–495,141 samples and 8–784 features.The datasets include categorical and continuous features, missing and complete data, and balanced and skewed class distributions.
  • Supervised setting: 13 out of 16 datasets have a SAINT variant outperforming all baselines, while performance is close to the best model on the remaining three.Results are means from five trials with different random seeds.
  • Semi-supervised setting: In semi-supervised experiments with 50, 200, or 500 labeled points, the pretrained SAINT model performs best; with all data labeled, pretraining contributes little.The semi-supervised comparisons use average AUROC across the evaluated datasets.
  • Model variants: SAINT-i consistently outperforms other variants when feature counts are large, especially when few training points accompany many features.The paper highlights Arcene and Arrhythmia as examples and reports faster execution than SAINT-s despite more parameters.
  • Robustness: SAINT and SAINT-i remain comparatively robust under feature corruption, whereas SAINT and SAINT-s are more robust when many training features are missing.Performance declines sharply after 70% CutMix corruption, while SAINT-s and SAINT remain non-dramatically degraded with 90% missing data.
  • Interpreting attention: Attention visualizations show self-attention focusing on important foreground pixels and intersample attention concentrating on a few potentially difficult or pivotal samples and classes.The attended classes vary between SAINT and SAINT-i, while using a few classes as pivots is prevalent; SAINT-i performance varies little across batch sizes 32–256.

6 Discussion, Limitations, and Impact

SAINT introduces intersample attention, contrastive pre-training, and improved embeddings for tabular data, showing that neural models can often improve on boosting methods. The authors caution that benchmark performance may not transfer to noisy or imbalanced real-world datasets.

  • SAINT combines intersample attention, contrastive pre-training, and an improved embedding strategy for tabular data.
  • SAINT often improves on boosting methods across numerous datasets with varying characteristics.
  • Limitations: Benchmark results should not be assumed to transfer directly to practitioners’ own settings.The authors note that real-world datasets may be highly noisy or imbalanced, and that SAINT was tuned for the tested settings.

B Datasets

The study uses 16 tabular datasets, spanning binary and multiclass classification, with dataset-specific preprocessing for feature types, missing values, and degenerate features. MNIST is represented as rows of 784 pixel features with positional encodings.

  • Dataset details and download links are provided in Tables 4 and 5.
  • The benchmark contains 16 datasets: 14 binary-classification datasets and 2 multiclass datasets with 10 classes.
  • Data preprocessing: Categorical features are label encoded, continuous features are z-normalized, and each column receives a distinct missing-value token.
  • Data preprocessing: Zero-standard-deviation features are removed from the Arcene, Arrhythmia, and KDD99 datasets.
  • Data preprocessing: For MNIST, each image becomes one row with 784 features, using a shared embedding space and positional encodings.

C Complete training details

Training uses a single GPU with dataset-dependent embedding sizes, layer counts, and attention heads constrained partly by memory. The implementation also applies specified dropout settings and positional encodings where needed.

  • Experiments use one Nvidia GeForce RTX 2080Ti GPU, with individual runs lasting 5 minutes to 6 hours and about 4 GPU days in total.
  • All attention layers use dropout 0.1; feed-forward dropout is 0.1 for SAINT-s and 0.8 for SAINT-i and SAINT.
  • The models use h = 8 attention heads generally and h = 4 for several lower-embedding-size datasets.
  • Positional Encoding: Positional encodings are used for MNIST because all features share the same embedding function and require feature distinction.

D Additional results

Additional analyses report uncertainty across runs, results for seven datasets omitted from the main paper, and robustness under corrupted training data. SAINT variants show consistency comparable to boosting methods and remain robust until substantial corruption.

  • Standard errors: Arrhythmia and Arcene have high AUROC standard errors across models, attributed to their small sizes of 400 and 200 datapoints.
  • Standard errors: SAINT variants exhibit consistency comparable to boosting methods, despite previous deep learning approaches being less consistent.
  • Remaining datasets: Table 7 reports average AUROC scores over 5 runs for seven binary-classification datasets not shown in the main paper.

E Additional analyses

Additional analyses examine robustness to corrupted training data and sensitivity to batch size. SAINT remains robust across substantial corruption, while AUROC variance is minimal across batch sizes from 32 to 256.

  • Robustness to data corruptions: SAINT’s performance drop remains minimal until 70% of training data is corrupted across missing-data and noisy-data settings.SAINT-s is more robust to missing data, while SAINT-i is more robust to noisy data.
  • Effect of batch size: AUROC variance is minimal when training batch size varies from 32 to 256 across five datasets and different SAINT variants.The datasets vary in numbers of features and samples.
  • Benchmark reporting: Table 7 reports average AUROC scores for SAINT variants and competitors across seven remaining binary-classification datasets.Scores are computed over five runs.
  • Benchmark reporting: Table 8 reports standard errors for the AUROC scores in Table 7 across seven remaining binary-classification datasets.Standard errors are computed over five runs.

E.1 Pre-training Ablations

The pre-training ablations evaluate projection heads, losses, and augmentations, with additional analyses of robustness and batch-size effects. Distinct projection heads, combined contrastive and denoising losses, and joint CutMix and mixup produce the best reported configurations.

  • Pre-training ablations: The ablation studies vary projection heads, pre-training losses, and data augmentation methods in SAINT’s pre-training pipeline.Table 9 reports average AUC over 14 datasets when only 50 points are labeled.
  • Additional analyses: Figure 6 evaluates how mean AUROC changes as the percentage of corrupted training data increases.It separates missing-data corruption from noisy-data corruption.
  • Projection heads: Distinct projection heads outperform weight-shared heads and configurations without projection heads.The study compares distinct heads, shared-weight heads, and no projection heads.
  • Pre-training loss: The combination of contrastive and denoising losses consistently yields the best results across all SAINT variants.The comparison also includes denoising alone, contrastive losses, and cosine similarity loss on positive pairs.
  • Pre-training augmentations: Using CutMix and mixup together produces the best performance across all SAINT variants.The study evaluates various configurations of these pre-training augmentations.
  • Projection-head and temperature analysis: Figure 8 studies the relationship between temperature τ and projection-head type in the pre-training pipeline.The broader ablation varies projection heads, losses, and augmentation methods.

F Additional interpretability plots

Interpretability plots visualize self-attention and intersample attention in SAINT variants. On the harder Volkert dataset, intersample attention is less sparse and all classes are attended to equally, unlike MNIST.

  • Self-attention: A one-stage SAINT-s model on MNIST has self-attention behavior similar to a one-stage SAINT model.With more stages, attention in the last stage is not quite as interpretable.
  • Intersample attention: Volkert’s intersample attention is less sparse than the corresponding behavior observed on MNIST.The authors hypothesize that attention sparsity depends on class separability, with Volkert harder than MNIST.
  • Value-vector visualization: On Volkert, all classes are attended to equally in the value-vector visualizations for SAINT and SAINT-i.Figure 10 plots 3000 points per figure with classes uniformly represented.
  • Intersample attention: Figure 9 compares intersample attention between different-class points for SAINT and SAINT-i on Volkert.The plots show attention between a batch of points from different classes.
Loading 2106.01342v1…