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

Gowthami Somepalli, Micah Goldblum+3 more|arXiv (Cornell University)|Jun 2, 2021

Imbalanced Data Classification TechniquesComputer Science43 references113 citations

TL;DR

SAINT uses self-attention over features and intersample attention across rows, with contrastive pre-training, to outperform traditional boosting methods on tabular data. It often surpasses XGBoost, CatBoost, and LightGBM across diverse benchmarks.

ABSTRACT

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.

Motivation & Objective

• Motivate a neural approach for tabular data that handles heterogeneous feature types and lacks inherent column order information.
• Propose SAINT, a transformer-based architecture that applies self-attention over features and intersample attention across rows.
• Introduce a contrastive self-supervised pre-training regime to improve performance in semi-supervised settings.
• Demonstrate empirical improvements of SAINT over tree-based methods and prior deep tabular models across a broad benchmark suite.

Proposed method

• Project continuous and categorical features into a common dense embedding space.
• Process embeddings with a transformer encoder that combines self-attention and a novel intersample attention (across batch rows).
• Embed continuous features via per-feature learnable projections before transformer processing.
• Pre-train with a hybrid objective: contrastive loss (InfoNCE) plus a denoising loss from augmented views (CutMix in input space and mixup in embedding space).
• Fine-tune by predicting targets from the [CLS] embedding via an MLP.
• Provide ablations and visualizations of attention to interpret model behavior.

Experimental results

Research questions

• RQ1 Can SAINT's combination of self-attention and intersample attention improve tabular data modeling beyond traditional boosting methods?
• RQ2 Does contrastive pre-training provide gains in semi-supervised settings for tabular data?
• RQ3 How does embedding continuous features affect performance compared to prior tabular transformers?
• RQ4 When is intersample attention most beneficial (e.g., many features, few labels)?

Key findings

• SAINT variants generally outperform baseline models across 14 binary classification datasets on AUROC, with SAINT often achieving the best results.
• On average, SAINT improves over traditional boosting methods (XGBoost, LightGBM, CatBoost) and other deep tabular models.
• In semi-supervised settings, pre-training SAINT (with both self and intersample attention) yields the best results, especially with limited labeled data.
• Embedding continuous features significantly boosts performance, as demonstrated by a TabTransformer comparison.
• Intersample attention provides robustness to noise and helps when the feature count is large or data is scarce.