SAINT: I MPROVED N EURAL N ETWORKS FOR T ABULAR D ATA VIA R OW A TTENTION AND C ONTRASTIVE P RE -T RAINING
Anonymous authors
Paper under double-blind review
A BSTRACT
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 performs competitively with gradient boosting methods, including XGBoost, CatBoost, and LightGBM, on average over 30 benchmark datasets in regression, binary classiﬁcation and multi-class classiﬁcation tasks.
I NTRODUCTION
While machine learning for image and language processing has seen major advances over the last decade, many critical industries, including ﬁnancial services, health care, and logistics, rely heavily on data in structured format.
Tabular data is unique in several ways that have prevented it from beneﬁting from the success of deep learning in vision and language.
First, tabular data often contain heterogeneous features that represent a mixture of continuous, categorical, and ordinal values, and these values can be independent or correlated.
Second, there is no inherent positional information in tabular data, meaning that the order of columns is arbitrary.
This differs from text, where tokens are always discrete, and ordering impacts semantic meaning.
It also differs from images, where pixels are typically continuous, and nearby pixels are correlated.
Tabular models must handle features from multiple discrete and continuous distributions, and they must discover correlations without relying on the positional information.
Deep learning systems with specialized architectures that embrace these differences have the potential to improve performance beyond what is achieved by classical methods, like linear classiﬁers and random forests.
Furthermore, without performant deep learning models for tabular data, we lack the ability to exploit compositionality, end-to-end multitask models, fusion with multiple modalities (e.g. image and text), and representation learning.
We introduce SAINT, the Self-Attention and INtersample attention Transformer, a specialized architecture for tabular data.
SAINT leverages several mechanisms to overcome the difﬁculties of training on tabular data.
SAINT projects all features – categorical and continuous – independently into a dense vector space.
These projected values are passed as token embeddings into a transformer encoder which uses attention in the following two ways.
First, “self-attention” attends to individual features within each data sample.
Second, we propose a novel “intersample attention”, which enhances the classiﬁcation of a row (i.e., a data sample) by relating it to other rows in the table.
Intersample attention is akin to a nearest-neighbor classiﬁcation, where the distance metric is learned end-to-end rather than ﬁxed.
In addition to this hybrid attention mechanism, we also leverage selfsupervised contrastive pre-training to boost performance for semi-supervised problems.
We provide comparisons of SAINT to a wide variety of deep tabular architectures and commonly used tree-based methods using a diverse battery of tabular datasets.
We observe that SAINT, on av- erage, outperforms all other deep learning methods on supervised and semi-supervised tasks.
More importantly, SAINT often out-performs boosted trees (including XGBoost (Chen & Guestrin, 2016),
CatBoost (Dorogush et al., 2018), and LightGBM (Ke et al., 2017)), which have long been an industry favorite for complex tabular datasets.
Finally, we visualize the attention matrices produced by our models to gain insights into how they behave.
R ELATED W ORK
Classical Models.
The most widely adopted approaches for supervised and semi-supervised learning on tabular datasets eschew neural models due to their black-box nature and high compute requirements.
When one has reasonable expectations of linear relationships, many modeling approaches are available (Wright, 1995; Weisberg, 2005; Starkweather & Moske, 2011; McCulloch & Neuhaus, 2005).
In more complex settings, non-parametric tree-based models are used.
Commonly used tools such as XGBoost (Chen & Guestrin, 2016), CatBoost (Dorogush et al., 2018), and LightGBM (Ke et al., 2017) provide several beneﬁts such as interpretability, the ability to handle a variety of feature types including null values, as well as performance in both high and low data regimes.
Deep Tabular Models.
While classical methods are still the industry favorite, some recent work brings deep learning to the tabular domain.
For example, TabNet (Arik & Pﬁster, 2019) uses neural networks to mimic decision trees by placing importance on only a few features at each layer.
The attention layers in TabNet do not use the regular dot-product self-attention common in transformerbased models, but a sparse layer that allows only certain features to pass through.
Yoon et al.
propose VIME, which employs MLPs in a technique for pre-training based on denoising.
TABERT (Yin et al., 2020), a more elaborate neural approach inspired by the large language transformer model
BERT (Devlin et al., 2018), is trained on semi-structured test data to perform language-speciﬁc tasks.
Several other studies utilize tabular data, but their problem settings are outside of our scope (Pathak et al., 2016;
Chen et al., 2019;
Padhi et al., 2021; Shavitt & Segal, 2018; Katzir et al., 2020).
Transformer models for more general tabular data include TabTransformer (Huang et al., 2020), which uses a transformer encoder to learn contextual embeddings only on categorical features.
The continuous features are concatenated to the embedded features and fed to an MLP.
The main issue with this model is that continuous data do not go through the self-attention block.
That means any information about correlations between categorical and continuous features is lost.
In our model, we address the issue by projecting continuous and categorical features to the higher dimensional embedding space and passing them both through the transformer blocks.
In addition, we propose a new type of attention to explicitly allow data points to attend to each other to get better representations.
Axial Attention.
Ho et al.
(2019) propose row and column attention in the context of localized attention in 2-dimensional inputs (like images) in the Axial Transformer.
For a given pixel, the attention is computed only on the pixels that are on the same row and column, rather than using all pixels.
The MSA Transformer (Rao et al., 2021) extends this work to protein sequences and applies both column and row attention across similar rows (tied row attention).
TABBIE (Iida et al., 2021) is an adaptation that applies self-attention to rows and columns separately, then averages the representations and passes them as input to the next layer.
In all these works, different features from the same data point communicate with each other and with the same feature from a whole batch of data.
Our approach, intersample attention, is hierarchical in nature; ﬁrst, features of a data point interact with each other, then data points interact with each other using entire rows/samples.
In a similar vein, Graph Attention Networks (GAT) (Veliˇckovi´c et al., 2017) compute attention over neighbors on a graph, thereby learning which neighbor’s information is relevant to a node’s prediction.
One way to view our intersample attention is as a GAT on a complete graph where all rows are connected to all other rows.
Yang et al.
(2016) explore hierarchical attention for document classiﬁcation where attention is computed between words in a given sentence and then between the sentences, but they did not attempt to compute the attention between entire documents themselves.
Self-Supervised Learning.
Self-supervision via a ‘pretext task’ on unlabeled data coupled with ﬁnetuning on labeled data is widely used for improving language and computer vision models.
Some of the pretext tasks previously used include imputing masked values (Pathak et al., 2016; Arik & Pﬁster, 2019; Huang et al., 2020), denoising (Vincent et al., 2008; Yoon et al., 2020), and replaced token detection (Huang et al., 2020; Iida et al., 2021).
Inspired by the success of contrastive learning for images (Chen et al., 2020;
He et al., 2020; Grill et al., 2020), we propose a variant of contrastive pre-training for tabular data.
To the best of our knowledge, this is the ﬁrst work to adopt contrastive learning for tabular data.
We couple it with denoising to pre-train on a plethora of datasets with varied volumes of labeled data, and we show that our method outperforms boosting methods.
S ELF -A TTENTION AND I NTERSAMPLE A TTENTION T RANSFORMER (SAINT)
[Figure]
In this section, we introduce our model, Self-Attention and INtersample attention Transformer (SAINT) and explain its components.
Suppose D
= { x i , y i } mi =1 is a tabular dataset with m points, x i is an n -dimensional feature vector, and y i is a label or target.
Similar to BERT (Devlin et al., 2018), we append a [CLS] token with a learned embedding to each sample.
Let x i =
[ [CLS] , f { 1 } i , f { 2 } i , .., f { n } i ] be a single data-point with categorical or continuous features f { j } i , and let E be the embedding layer that embeds each feature into R d .
Note that E may use different embedding functions for different features.
For a given x i ∈ R ( n +1) , we get E ( x i ) ∈ R ( n +1) ×
d .
Encoding the Data.
In language models, all tokens are embedded using the same procedure.
However, in the tabular domain, different features can come from distinct distributions, necessitating a heterogeneous embedding approach.
Note that tabular data can contain multiple categorical features which may use the same set of tokens.
Unless it is known that common tokens possess identical relationships within multiple columns, it is important to embed these columns independently.
Unlike the embedding of TabTransformer (Huang et al., 2020), which uses attention to embed only categorical features, we propose also projecting continuous features into a d − dimensional space before passing their embedding through the transformer encoder.
To this end, we use a separate single fully-connected layer with a ReLU nonlinearity for each continuous feature, thus projecting the 1 − dimensional input into d − dimensional space.
With this simple trick alone, we signiﬁcantly improve the performance of the TabTransformer model as discussed in Section 5.1.
An additional discussion
[Figure]
3. A RCHITECTURE
SAINT is inspired by the transformer encoder of Vaswani et al.
(2017), designed for natural language, where the model takes in a sequence of feature embeddings and outputs contextual representations of the same dimension.
A graphical overview of SAINT is presented in Figure 1(a).
SAINT is composed of a stack of L identical stages.
Each stage consists of one self-attention transformer block and one intersample attention transformer block.
The self-attention transformer block is identical to the encoder from Vaswani et al.
It has a multi-head self-attention layer (MSA) (with h heads), followed by two fully-connected feed-forward (FF) layers with a GELU non-linearity (Hendrycks & Gimpel, 2016).
Each layer has a skip connection (He et al., 2016) and layer normalization (LN)
(Ba et al., 2016).
The intersample attention transformer block is similar to the self-attention transformer block, except the self-attention layer is replaced by an intersample attention layer (MISA).
The details of intersample attention layers are presented below.
The SAINT pipeline, with a single stage ( L = 1 ) and a batch of b inputs, is described by the following equations.
We denote multi-head self-attention by MSA, multi-head intersample attention by MISA, feed-forward layers by FF, and layer norm by LN:
[Equation]
where r i is SAINT’s contextual representation output corresponding to data point x i .
This contextual embedding can be used in downstream tasks such as self-supervision or classiﬁcation.
3. I NTERSAMPLE ATTENTION
We introduce intersample attention (a type of row attention) where attention is computed across different data points (rows of a tabular matrix) in a given batch rather than just features of a single data point.
Speciﬁcally, we concatenate the embeddings of each feature for a single point, then compute attention over samples (rather than features).
This enables us to improve the representation of a point by inspecting others.
When a feature is missing or noisy in one row, intersample attention enables SAINT to borrow the corresponding features from other similar data samples in the batch.
An illustration of how intersample attention is performed in a single head is shown in Figure 2 and the pseudo-code is presented in Algorithm 1.
Unlike the row attention used in Ho et al.
Child et al.
; Rao et al.
(2021); Iida et al.
(2021), intersample attention allows all features from different samples to communicate with each other.
In our experiments, we show that this ability boosts performance appreciably.
In the multi-head case, instead of projecting q, k, v to a given dimension d , we project them to d/h where h is the number of heads.
Then we concatenate all the updated value vectors, v i , to get back a vector of length d .
[Table]
P RE - TRAINING & F INETUNING
Contrastive learning, in which models are pre-trained to be invariant to reordering, cropping, or other label-preserving “views” of the data (Chen et al., 2020;
He et al., 2020; Pathak et al., 2016;
Grill et al., 2020; Vincent et al., 2008), is a powerful tool in vision domains that has never (to our knowledge) been applied to tabular data.
We present a contrastive pipeline for tabular data, a visual description of which is shown in Figure 1.
Existing self-supervised objectives for tabular data include denoising (Vincent et al., 2008), a variant of which is used by VIME (Yoon et al., 2020), masking, and replaced token detection as used by TabTransformer (Huang et al., 2020).
We ﬁnd that, while these methods are effective, superior performance can be achieved by contrastive learning.
Generating augmentations.
Contrastive methods in vision craft different “views” of images via crops and ﬂips.
For tabular data, crafting such transformations is a non-trivial problem.
VIME (Yoon et al., 2020) uses mixup in the non-embedded space as data augmentation, but this is limited to continuous data.
We instead propose to use CutMix (Yun et al., 2019) to augment samples in input space and mixup (Zhang et al., 2017) in embedding space.
These augmentations yield a challenging and effective self-supervision task.
Assume that l of m data points are labeled.
We denote the embedding layer by E , the SAINT network by S , and 2 projection heads as g 1 ( · ) and g 2 ( · ) .
The
CutMix augmentation probability is denoted p cutmix and the mixup parameter is α .
Given point x i , the original embedding is p i =
E ( x i ) , while the augmented representation is generated as follows:
[Equation]
CutMix in raw data space
[Equation]
mixup in embedding space
[Equation]
[Equation]
where x a , x b are random samples from the current batch, x (cid:48) b is the CutMix version of x b , m is the binary mask vector sampled from a Bernoulli distribution with probability p cutmix , and α is the mixup parameter.
We ﬁrst obtain a CutMix version of every data point in a batch by randomly selecting a partner to mix with.
We then embed the samples and choose new partners before performing mixup.
SAINT and projection heads.
Now that we have both the clean p i and mixed p (cid:48)
i embeddings, we pass them through SAINT, then through two projection heads, each consisting of an MLP with one hidden layer and a ReLU.
The use of a projection head before computing contrastive loss is common in vision (Chen et al., 2020;
He et al., 2020; Grill et al., 2020) and indeed also improves results on tabular data.
Ablation studies and further discussion are available in Appendix F.
Loss functions.
We consider two losses for the pre-training phase.
(i)
The ﬁrst is a contrastive loss that pushes the latent representations of two views of the same data point ( z i and z (cid:48) i ) close together and encourages different points ( z i and z j , i (cid:54) =
j ) to lie far apart.
For this, we borrow the InfoNCE loss from metric-learning works (Sohn, 2016; Oord et al., 2018; Chen et al., 2020; Wu et al., 2018); (ii) The second loss comes from a denoising task.
For denoising, we try to predict the original data sample from a noisy view.
Formally, we are given r (cid:48)
i and we reconstruct the inputs as x (cid:48)(cid:48)
i to minimize the difference between the original and the reconstruction.
The combined pre-training loss is:
[Equation]
where r i = S ( p i ) , r (cid:48)
i =
S ( p (cid:48) i ) , z i = g 1 ( r i ) , z (cid:48)
i = g 2 ( r (cid:48) i ) .
L j is cross-entropy loss or mean squared error depending on the j th feature being categorical or continuous.
Each MLP j is a single hidden layer perceptron with a ReLU non-linearity.
There are n in number, one for each input feature.
λ pt is a hyper-parameter and τ is temperature parameter and both of these are tuned using validation data.
Finetuning.
Once SAINT is pre-trained on all unlabeled data, we ﬁnetune the model on the target prediction task using the l labeled samples.
The pipeline of this step is shown in Figure 1(b).
For a given point x i , we learn the contextual embedding r i .
For the ﬁnal prediction step, we pass the embedding corresponding only to the [CLS] token through a simple MLP with a single hidden layer with ReLU activation to get the ﬁnal output.
We evaluate cross-entropy loss on the outputs for classiﬁcation tasks and mean squared error for regression tasks.
E XPERIMENTAL E VALUATION
We now discuss SAINT variants and evaluate them in both supervised and semi-supervised scenarios on 30 datasets.
We also analyze each component of SAINT and perform ablation studies to understand the importance of each component.
Finally, we probe the behavior of attention maps in SAINT by treating the MNIST dataset (LeCun et al., 1998) as tabular data and generating visualizations.
[Table]
Datasets.
We evaluate SAINT onbenchmark datasets, with 10 for each of binary classiﬁcation, multiclass classiﬁcation, and regression.
We chose datasets based on (i) usage in previous tabular work (Arik & Pﬁster, 2019;
Huang et al., 2020) and (ii) availability on the OpenML (Vanschoren et al., 2013) platform where the datasets are uniformly processed and accessible.
The datasets are diverse, including from 452 up to 581,012 samples, with 2 to 100 class labels, and containing from 6 to 1,777 features – both categorical and continuous features.
Some datasets are missing data while others are complete, and some are balanced, while others have highly skewed class distributions.
dataset IDs and additional details can be found in Appendix B. We pre-process each dataset by Z-normalizing all continuous features and by label-encoding all categorical features before data is passed to the embedding layer.
Model variants.
The SAINT architecture discussed in the previous section has one self-attention transformer encoder block stacked with one intersample attention transformer encoder block in each stage.
We also consider variants with only one of these blocks.
SAINT-s has only self-attention, while SAINT-i has only intersample attention.
See Table 1 for an architectural comparison.
Baselines.
We compare our model to traditional methods like Random Forests (Breiman, 2001),
Extra Trees (Geurts et al., 2006), and k-NN (Altman, 1992) as well as against powerful boosting libraries XGBoost (Chen & Guestrin, 2016), LightGBM (Ke et al., 2017), and CatBoost (Dorogush et al., 2018).
We also compare to neural networks, like 2-layer multi-layer perceptrons, TabNet (Arik & Pﬁster, 2019), and TabTransformer (Huang et al., 2020).
For methods that use unsupervised pretraining as preprocessing, we use Masked Language Modeling (MLM) (Devlin et al., 2018) for
TabNet and Replaced Token Detection (RTD) (Clark et al., 2020) for TabTransformer as in the respective papers.
We split the data into 65% , 15% , and 25% for training, validation, and testing, respectively.
Hyperparameter tuning is done on validation data and results are reported on test data.
Please refer to Appendix D for training details and hyperparameter search spaces for each method.
[Table]
Metrics.
We used different metrics for each of the three tasks we examine in this paper.
For binaryclassiﬁcation, we use area under ROC curve (AUROC), accuracy for multi-class classiﬁcation, and root-mean-squared error (RMSE) for regression.
Training.
We train (including pre-training runs) using AdamW (Loshchilov & Hutter, 2017) with β 1 = 0 .
9 , β 2 = 0 .
999 , decay = 0 .
01 , and with a learning rate of 0 .
with batches of size(except for datasets with more than 100 features where we vary the batch size depending on the number of features to ﬁt the batch into a single GPU memory).
We split the data into 65% , 15% , and 25% for training, validation, and testing, respectively.
We use only validation split for selecting all other hyper-parameters.
The ﬁnal conﬁgurations for each of the datasets is presented in Appendix
C. We use CutMix mask parameter p cutmix = 0 .
3 and mixup parameter α = 0 .
2 for all standard pretraining experiments.
We use pre-training loss hyper-parameters λ pt = 10 and temperature τ = 0 .for
all settings.
5. R ESULTS
Supervised setting.
In Table 2, we report average ranks of each model on 10 datasets from each
-wise on binary
, we respectively to
For regression, we measure RMSE.
SAINT ranks 2.9 across 10 datasets compared to the next best method, CatBoost, which ranks 3.9.
For complete results on all 30 datasets, see Appendix E.
Semi-supervised setting.
We perform 2 sets of experiments with 50, and 200 labeled data points (the rest are unlabeled).
See Table 3 for average ranks across all 30 datasets and standard errors in both settings.
We also examine the case when all data is labeled, to see if pre-training still helps.
In all 3 cases, the pre-trained SAINT model performs the best.
Pre-training had the largest positive impact on the SAINT-s variant with a rank improvement in all cases.
Interestingly, we note that when all the training data samples are labeled, pre-training does not contribute appreciably (except for SAINT-s variant), hence the results with and without pre-training are fairly close.
Effect of embedding continuous features.
To understand the effect of learning embeddings for continuous data, we perform a simple experiment with TabTransformer.
We modify TabTransformer by embedding continuous features into d dimensions using a single-layer ReLU MLP, just as they use on categorical features, and we pass the embedded features through the transformer block.
We keep the entire architecture and all training hyper-parameters the same for both TabTransformer and its modiﬁed version.
In binary classiﬁcation, the average AUROC across 10 datasets increases from 0.82 to 0.85 (an increment of 3.65%), and in multiclass classiﬁcation, the average accuracy across 10 datasets increases from 0.72 to 0.74 (an increment of 2.77%).
In regression, the average RMSE over 10 datasets decreases from 5,071.30 to 4,998.2 (a decrement of 1.44%).
This experiment shows that embedding the continuous data is important and can boost performance signiﬁcantly.
When to use intersample attention?
In our experiments
[Table]
we observe that SAINT-i outperforms other variants whenever the number of features is large ( > 500 ) as observed in the cases of MNIST and Bioresponse.
Another advantage of SAINT-i is that execution is fast compared to SAINT-s, despite the fact that the number of parameters in SAINT-i is much higher than that of SAINT-s (see
Table 1).
This can be attributed to quadratic complexity in terms of number of features to compute attention in SAINT-s while in SAINT-i, this depends on the size of the batch which we can control.
[Table]
How
robust is SAINT to data cor-
ruptions?
We evaluate the robustness of SAINT variants by corrupting the training data via noise and missing values.
To simulate noise, we apply CutMix, replacing 10% to 90% of features with values of other randomly selected samples.
We examine the resulting trends in binary classiﬁcation and regression.
In trials with noise, the drop in mean
AUROC or mean scaled-RMSE is minimal until 70% of data is corrupted, when the performance drops signiﬁcantly.
SAINT and SAINT-i are comparatively more robust than
SAINT-s.
This shows that using row attention improves the model’s robustness to noisy training data.
However, we ﬁnd the opposite trend when many features are missing in the training data.
In this scenario, SAINT-s and SAINT are quite robust, and the drop in AUROC is not drastic until 70% or more of the data is missing.
We conclude that SAINT is reliable for training on corrupted training data.
Trend line plots for these scenarios are available in Appendix F, Figure 6.
Effect of batch size on intersample attention performance.
As we discuss in Section 3.2, intersample attention is computed between batches of data points.
We consider multiple multi-class classiﬁcation datasets and examine the impact of batch size by varying it from 8 to 1,024.
In most cases, we ﬁnd that the variation in SAINT-i’s performance is small and is comparable to that of
SAINT-s, which has no intersample attention component and thus has no dependence on batch size.
We note that when the number of training samples is very small (i.e. less than 1,000), the variation in performance increases, but we observe the same trend in SAINT-s as well.
We present corresponding plots in Appendix F, Figure 7.
Further discussion and analysis is presented in Appendix F.
5. I NTERPRETING ATTENTION
One advantage of transformers is that attention lends itself to interpretability tools as opposed to
MLPs, which are hard to interpret.
In particular, when we use only one transformer stage, the attention maps reveal which features and which data points are being used by the model to make decisions.
We use MNIST (LeCun et al., 1998) data to examine how self-attention and intersample attention behave in our models.
While MNIST is not a typical tabular dataset, it has the advantage that its features can easily be visualized as an image.
Figure 3a depicts the attention on each of the pixels/features in a self-attention layer of SAINT.
Without any explicit supervision, the model learns to focus on the foreground pixels, and we clearly see from the attention map which features are most important to the model.
The self-attention plots of SAINT-s are similar (Appendix G, Figure 8a).
Figures 3b and 3c depict a similar visualization on a batch of 20 points, two from each class in
MNIST.
Figure 3b shows intersample attention in SAINT.
This plot shows which samples attend to which other samples in the batch.
Surprisingly, very few points in a batch receive attention.
We hypothesize that the model focuses on a few points that are critical because they are particularly difﬁcult to classify without making direct comparisons to exemplars in the batch.
In Figure 3c, we
[Figure] [Figure]
show the intersample attention plot from a SAINT-i model.
The same sparse attention behaviour persists here too, but the points being attended to are different in this model.
Interestingly, we ﬁnd this behavior to be signiﬁcantly different on the Volkert data (Dataset ID: 41166, multiclass withclasses), where a wide range of data becomes the focus of attention depending on the input.
The intersample attention layer gets dense with the difﬁculty (to classify) of the datasets.
See Appendix G for additional MNIST and Volkert attention maps.
Figure 4 shows the behavior of attention at the dataset (rather than batch) level.
We visualize a t-SNE (Van der Maaten & Hinton, 2008) embedding for value vectors generated in intersample attention layers, and we highlight the points that are most attended to in each batch.
In Figure 4 (left), the value vectors and attention are computed on the output representations of a self-attention layer.
In contrast, value vectors and attention in Figure 4 (right) are computed on the embedding layer output, since SAINT-i does not use self-attention.
In these plots, the classes to which the model attends vary dramatically.
Thus, the exact classes to which an attention head attends change with the architecture, but the trend of using a few classes as a ‘pivot’ seems prevalent in intersample attention heads.
C ONCLUSION
Even though tabular data is an extremely common data format used by institutions in various domains, deep learning methods are still lagging behind tree-based boosting methods in production.
In this paper, we highlight the limitations of existing deep learning approaches, and the reasons why they do not perform competitively on the tabular datasets.
We further introduce three novel improvements to counter these limitations - (1) intersample attention, (2) contrastive pre-training, (3) and a strategy for embedding continuous columns in the tabular datasets.
Our method, SAINT, performs competitively across a large number of tabular datasets with diverse characteristics on the tasks of binary classiﬁcation, multi-class classiﬁcation and regression.
We quantitatively and qualitatively analyse the importance of each of our improvements.
E THICS AND R EPRODUCIBILITY S TATEMENT
The tremendous success of deep learning in vision and language domains is yet to be realized for tabular datasets.
In our work, we discuss some of the potential limitations of existing deep learning models and how we can mitigate them by deploying not just features but other data points for the inference.
We believe that our proposed architecture, along with the self-supervised and data augmentation techniques, will greatly help with representation learning and machine learning applications in tabular data domain.
With that in mind, like most deep learning models, SAINT in its current form is prone to capturing and even amplifying the biases in the dataset on which it is trained on.
Our data augmentation strategy might help to mitigate this issue to some extent but by no means is it intended to resolve the data imbalance.
We caution the end users to carefully consider their use cases and use our method only after appropriate validation and testing.
Tabular datasets often contain demographics or personally identiﬁable information (PII).
Leaking of information from training data with large deep learning models is a topic of great concern and has been demonstrated to a certain extent in the language domain.
However, we are not aware of such works in the case of tabular data.
Reproducibility.
We are fully committed to open sourcing our code for reproducibility, as well as for the broader community to build upon it.
Please ﬁnd the code and instructions to use it attached in the supplementary material.
A A DDITIONAL ILLUSTRATIONS
[Figure]
B D ATASETS
Data sources.
As mentioned in main text, all 30 datasets are downloaded from OpenML (Vanschoren et al., 2013) repository.
We present the dataset IDs and some statistics of the datasets in Table 4.
All datasets can be downloaded via https://www.openml.org/d/<dataset_id>
Data preprocessing.
In each dataset, the categorical features are label encoded, and continuous features are z-normalized before passing them into the embedding layer.
Each feature (or column) has a different missing value token to account for missing data.
For MNIST, we unravel each image into a vector of 784 features and consider each image as a single row.
Since each feature is of same type in this dataset, we encode all the features into the same embedding space.
To distinguish the features, we also use positional encodings in the encoding layer.
C C OMPLETE TRAINING DETAILS
In each of our experiments, we use a single Nvidia GeForce RTX 2080Ti GPU.
Individual training runs take between 5 minutes and 6 hours.
In total, the experiments in this paper account for around 30 GPU days (including semi-supervised experiments and ablation studies).
[Table]
multiclasstopo
N/A    regressionsocmob
N/A    regressionhouse prices nominal
N/A    regression
Allstate Claims Severity
N/A    regression
Yolanda
N/A    regression
OnlineNewsPopularity
N/A    regressionabalone
N/A    regressioncolleges
N/A    regression
Airlines DepDelay 10M
N/A    regressionnyc-taxi-green-decN/A    regression
For most of the datasets, we use embedding size d = 32 .
For MNIST, we use d = 12 , for datasets with > 100 features we used d = 4 .
The variance in the embedding size is only due to the memory constraints of a single GPU.
We used L = 6 layers in the SAINT-s variant and L = 1 for SAINTi and SAINT variants.
We use dropout of 0.1 in all attention layers.
In feed-forward layers, use dropout of 0.1 in the SAINT-s variant, and we use 0.8 in SAINT-i and SAINT models.
We use attention heads h = 8 for SAINT-s model and h = 4 for SAINT-i and SAINT models.
Inside the self-attention layer, the q , k , and v vectors are of dimension 16, and in the intersample attention layer, they are of size 64.
Other minor details are shared in the code.
Feature Encoding.
Transformers for vision and language typically employ same embedding function for each feature, or a column since each word, or pixel, or a patch comes from a similar distribution.
This is not the case with all of the 30 datasets presented in this paper; each feature may be of a different type and drawn from different distribution.
Hence we use a unique embedding function for each column.
Positional Encoding.
In tabular datasets, different features, or columns do not have any order.
Different feature encodings used to embed different columns is sufﬁcient to provide the model the necessary information about the feature.
Hence we skip the positional encodings in all the datasets (except in the case of MNIST, where all columns share the embedding function, and we provide positional encodings instead).
Note that the attention operation itself is order invariant in both the self-attention and intersample attention cases.
D B ASELINES
We use Auto-Gluon 2 (Erickson et al., 2020) framework to run the train the traditional models as well as to perform the hyperparameter sweep.
For TabNet, we use the famous PyTorch implementation 3 , and for TabTransformer since the ofﬁcial implementation is not available, we used the widely used PyTorch re-implementation 4 .
In the following paragraphs we will give details on hyperparameter conﬁgurations for each baseline.
All the baseline models use early stopping on the validation dataset’s performance.
All the deep learning methods take embedding for categorical data as input and we tuned them with variable embedding sizes from [16 , 32 , 64] .
Random Forest & Extra trees Classiﬁers.
We tuned on number of estimators ( n estimators ) from the list [100 , 200] , on criterion from [‘gini’,‘entropy’].
Random Forest & Extra trees Regressors.
We tuned on number of estimators ( n estimators ) from the list [100 , 200] , on criterion from [‘squared error’,‘absolute error’].
KNeighbors models.
We tuned the weights parameter from [‘uniform’,‘distance’], and p value from [1,2].
For large datasets, k-NN model is computationally intractable and hence for datasets with > 50000 datapoints, we randomly choose 50k points and train the model on the smaller dataset.
LightGBM.
We tried with different values of min data in leaf in [10 , 20] , num leaves in [30 , 50] and learning rate in the interval [0 . 01 , 0 . 1]
XGBoost.
We used the values of booster as ’gbtree’, n estimators of 200 and tuned over
L2 regularization lambda in [1 , 10] and the learning rate in the interval [0 . 01 , 0 . 3] .
CatBoost.
We used the values of reg lambda as 100 and early stopping rounds asand tuned the learning rate in the interval [0 . 01 , 0 . 1] .
MLP.
We used a 2 hidden-layered neural network of sizes [200 , 100] , with 1d BatchNorm optimized on Adam (Kingma & Ba, 2014) with learning rate set to 0.001.
TabNet.
We have used the hyperparameters suggested by the PyTorch implementation of Tabnet.
We used Adam optimizer with learning rate set to 0.02 and momentum set to 0.3.
We used learning rate scheduler with gamma as 0.95 and step-size as 20.
The rest of the parameters we have tuned based on the ranges provided in the original paper.
TabTransformer.
We used AdamW optimizer with learning rate set to 0.001.
We also tried different levels of feed-forward drop-outs from [0 . 1 , 0 . 5 , 0 . 8] .
E E XTENDED RESULTS
In the main part of the paper, we have presented the average ranks of the models across the benchmark datasets.
The supervised results are shared across the following three tables.
In Table 5, we share the AuROC scores of all the baselines and SAINT variants for the binary classiﬁcation task.
In Table 6, we share accuracies for multiclass classiﬁcation.
In both these cases, higher the score, better the model.
While in case of regression, we share the RMSE scores in Table 10, smaller the value, better the model.
In case of semi-supervised/ pre-training experiments, see the exhaustive results for the case oflabeled samples in Table 11, for 200 samples in Table 12, and all samples labeled in Table 13.
Please note that in these 3 tables, we have scores from different tasks together in the same table.
F A DDITIONAL ANALYSES
How robust is SAINT to data corruptions?
(contd.)
As discussed in the main body, we evaluate the robustness of SAINT variants in cases of data corruption or missing data.
We examine the cases of Binary classiﬁcation and Regression and plot the average scores across the 10 datasets correponding to that task.
We present the results in Figure 6.
[Table] [Table]
[Table]
[Table]
Effect of batch size on intersample attention performance.
(cont.)
As discussed in the main body, we examine the affect of batch size on different SAINT variants in Figure 7.
We pickmulticlass classiﬁcation datasets with varying numbers of features and samples.
In all cases, we see that the variance in Accuracy is minimal when varying the batch size from 8 to 1024.
In all the plots, we see log 2 (batchsize) plotted on the X-axis with accuracy on Y-axis . [Table]
F. P RE - TRAINING A BLATIONS
In Table 9, we study various conﬁgurations of pre-training components.
We perform 3 primary studies: we vary (1) projection head, (2) pre-training loss, and (3) data augmentation method.
Note, the ﬁnal result in all 3 studies refers to the same experiment (hence the row is repeated), which is the ﬁnal chosen conﬁguration for our model.
Effect of projection heads.
As described in Section 4, we use two different projection heads, g 1 ( · ) and g 2 ( · ) , to project the contextual representations to lower dimensions and then compute contrastive losses.
We study three different options for the heads: (1) distinct projection heads (2) heads with weight sharing, and (3) no projection heads at all.
Table 9 shows that using distinct projection heads performs best.
Varying pre-training loss.
We train SAINT’s variants with different loss functions, as shown in
Study 2 of Table 9.
We try denoising and contrastive losses, in addition to a cosine similarity loss on positive pairs (inspired by Grill et al.
; Chen & He (2020)).
The combination of contrastive and denoising consistently yields the best results in all SAINT variants.
Varying the pre-training augmentations.
We also try to understand how important it is to use
CutMix and mixup to generate augmented embeddings in the pre-training pipeline.
We tinker with
[Table] various conﬁgurations in Study 3 of Table [Table] and we observe that using these two augmentations in unison results in the best performance across all SAINT variants
[Table]
G A DDITIONAL INTERPRETABILITY PLOTS
In Figure 8a
[Table]
we show a self-attention plot for the SAINT-s variant (with L = 1 ) on MNIST.
The self-attention in one stage SAINT-s model behaves similar to a one stage SAINT model.
However, when there are more stages, the attention in the last stage is not quite as interpretible.
In Figure 8b, we show the intersample attention between a batch of points from different classes in
SAINT model on the Volkert dataset (OpenML dataset ID: 41166).
Similarly in Figure 8c, we show intersample attention in the SAINT
As mentioned in the main body [Figure] the intersample behaviour is not quite as sparse as that of MNIST.
We hypothesize that the sparsity of the intersample attention layer depends on how separable the classes in the dataset are.
(Volkert is a harder dataset than MNIST).
In Figure [Figure] we show the t-SNE plots on value vectors for SAINT and SAINT-i variants on Volkert.
Unlike MNIST [Figure] all the classes are attended to equally.
[Figure] [Figure]
[Table]
T a b l e 11 :
S e m i - s up e r
v
i s e d l ea r n i ng w it
h
l a b e l e d s a m p l e s ca s e :
A c t u a l s c o r e s f o r eac
h
o f t
h e m od e l f o r a g
i v e n d a t a s e t .
F i r s t 10 d a t a s e t s c
o rr e s pond t
o m u lti c l a ss c l a ss i ﬁ ca ti on , h e n ce t
h e s c
o r e s a r e acc
u r
ac
i e s , t h e
n e x t b e l ong t o b i n a r y c l a ss
i ﬁ ca ti on , s o t
h e s c
o r e s a r e A u R O
C s , a nd t h e l a s t 10 c o l u m n s a r e r e g r e ss i on t yp e , h e n ce t
h e y ca rr y
R M
S E v a l u e s .
R e m ov e d t h e n e g a ti v e s i
gn i n
R M
S E s c
o r e s
T a s k /
M u l t i c
l a ss C l a ss i ﬁ c a t i o n
B i n a r y
C l a ss i ﬁ c a t i o n
R e g re ss i o n
M o d e l/ d a t a s e t
[Table]
R a ndo m F o r e s t 0 .0 .0 .0 .1 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .
[Table]
[Table]
E x t r a T r ee s 0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .
[Table]
KN e i ghbo r s D
i s t 0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .
[Table]
[Table]
KN e i ghbo r s U n i f 0 .0 .0 .
[Table]
a t
[Table]
oo s t
[Table]
S
[Table] [Table]
+
[Table]
r
[Table] [Table]
r
[Table]
n
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
S A I N T
+ p r e - t r a
i n i ng
[Table]
[Table] [Table]
[Table] [Table]
T a b l e 12 :
S e m i - s up e r
v
i s e d l ea r n i ng w it
h l a b e l e d s a m p l e s ca s e :
A c t u a l s c o r e s f o r eac
h
o f t
h e m od e l f o r a g
i v e n d a t a s e t .
F i r s t 10 d a t a s e t s c
o rr e s pond
[Table]
l
[Table]
c l a ss i ﬁ ca ti on , h e n ce t
h e s c
o r e s a r e acc
u r
ac
i e s , t h e
n e x t b e l ong t o b i n a r y c l a ss
i ﬁ ca ti on , s o t
h e s c
o r e s a r e A u R O
C s , a nd t h e l a s t 10 c o l u m n s a r e r e g r e ss i on t yp e , h e n ce t
h e y ca rr y
R M
S E v a l u e s .
R e m ov e d t h e n e g a ti v e s i
gn i n
R M
S E s c
o r e s
T a s k /
M u l t i c
l a ss C l a ss i ﬁ c a t i o n
B i n a r y
C l a ss i ﬁ c a t i o n
R
[Table]
i
[Table]
M o d e l/ d a t a s e t 18
R a ndo m F o r e s t 0 .0 .0 .0 .1 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .
E x t r a T r ee s 0 .0 .0 .0 .1 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .
[Table]
KN e i ghbo r s D
i s t 0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .
[Table]
[Table] [Table]
[Table] [Table]
[Table]
KN e i ghbo r s U n i f 0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .
[Table]
[Table] [Table]
L i gh t G B
M 0 .0 .0 .0 .1 .0 .0 .0 .0 .0 .0 .0 .0 .0 .
[Table]
[Table]
XG B oo s t 0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .
[Table]
C a t B oo s t
[Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
M L P 0 .
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
T a bn e t +
M L M
[Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
T a b T r a n s f o r m e r + R T D 0 .0 .0 .0 .1 .0 .0 .0 .
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
S A I N T - s
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
S A I N T - i
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
S A I N T 0 .
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
S A I N T - s + p r e - t r a i n i ng 0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .
S A I N T - i + p r e - t r a
i n i ng 0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .
S A I N T
+ p r e - t r a
i n i ng 0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .
T a b l e 13 :
S e m i - s up e r
v
i s e d l ea r n i ng w it h
A
LL s a m p l e s l a b e l e d ca s e :
A c t u a l s c o r e s f o r eac
h
o f t
h e m od e l f o r a g
i v e n d a t a s e t .
F i r s t
d a t a s e t s c o rr e s pond t o m u lti
c l a ss c l a ss i ﬁ ca ti on , h e n ce t
h e s c
o r e s a r e a cc u r
ac
i e s , t h e
n e x t b e l ong t o b i n a r y c l a ss
i ﬁ ca ti on , s o t
h e s c
o r e s a r e
A u R
O C s , a nd t h e l a s t c o l u m n s a r e r e g r e ss i on t yp e , h e n ce t
h e y ca rr y
R M
S E v a l u e s .
R e m ov e d t h e n e g a ti v e s i
gn i n
R M
S E s c
o r e s
T a s k / A ll
M u l t i c
l a ss C l a ss i ﬁ c a t i o n
B i n a r y
C l a ss i ﬁ c a t i o n
R e g re ss i o n
M o d e l/ d a t a s e t
[Table]
[Table]
R a ndo m F o r e s t
[Table]
[Table]
E x t r a T r ee s
[Table]
[Table]
r s D i s
[Table]
KN e
[Table]
r s U n i f 0 .
[Table] [Table]
S A I N T - i 0 .0 .0 .0 .1 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .
[Table]
[Table] [Table]
[Table] [Table]
[Table]
S A I N T 0 .0 .0 .1 .1 .0 .1 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .
[Table]
[Table] [Table]
S A I N T - s + p r e - t r a i n i ng 0 .0 .0 .0 .1 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .0 .
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
S A I N T - i + p r e - t r a
i n i ng 0 .0 .0 .1 .1 .0 .0 .0 .0 .0 .0 .
[Table]