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

1 INTRODUCTION . While machine learning for image and language processing has seen major advances over the last decade , many critical industries , including financial services , health care , and logistics , rely heavily on data in structured format . Tabular data is unique in several ways that have prevented it from benefiting 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 classifiers 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 difficulties 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 classification 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 classification , where the distance metric is learned end-to-end rather than fixed . 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 . 2 RELATED WORK . 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 benefits 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 & Pfister , 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 . ( 2020 ) 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-specific 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 ; first , 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 ) ( Veličković 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 classification 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 finetuning 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 & Pfister , 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 first 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 . 3 SELF-ATTENTION AND INTERSAMPLE ATTENTION TRANSFORMER . ( SAINT ) In this section , we introduce our model , Self-Attention and INtersample attention Transformer ( SAINT ) and explain its components . Suppose D = { xi , yi } mi=1 is a tabular dataset with m points , xi is an n-dimensional feature vector , and yi 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 xi = [ [ 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 Rd . Note that E may use different embedding functions for different features . For a given xi ∈ R ( n+1 ) , we get E ( xi ) ∈ 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 significantly improve the performance of the TabTransformer model as discussed in Section 5.1 . An additional discussion concerning positional encodings can be found in Appendix C . 3.1 ARCHITECTURE . 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 . ( 2017 ) . 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 : z ( 1 ) i = LN ( MSA ( E ( xi ) ) ) +E ( xi ) z ( 2 ) i = LN ( FF1 ( z ( 1 ) i ) ) + z ( 1 ) i ( 1 ) z ( 3 ) i = LN ( MISA ( { z ( 2 ) i } b i=1 ) ) + z ( 2 ) i ri = LN ( FF2 ( z ( 3 ) i ) ) + z ( 3 ) i ( 2 ) where ri is SAINT ’ s contextual representation output corresponding to data point xi . This contextual embedding can be used in downstream tasks such as self-supervision or classification .

Summary: The paper presents a new deep-learning network architecture for tabular data classification and regression problems, called SAINT. While deep-learning provided significant improvements in many domains, tabular data problems are still dominated by classical algorithms like XGBoost and Catboost. The new architecture is based on a new attention mechanism, applying attention both between samples (rows) and features (columns). It also embeds continuous features rather than embedding just categorical ones. The new model is studied on 30 diverse datasets (10 for binary classification, 10 for multiclass classification and 10 for regression), and compared to 10 previous types of models, including both classical models and recently proposed deep-learning models. The authors show that the average rank of SAINT across the various datasets is the best compared to these 10 algorithms (average rank of 2.7, vs. average rank of 4.0 for Catboost). Ablation studies are performed to evaluate the impact of each direction of attention and the impact of the embedding, and the explainability is demonstrated on MNIST. Additionally, the paper presents a contrastive learning approach for pre-training on unlabeled data and fine-tuning on a small number of labels, and it show that pre-training improves the rank for these small numbers of labels (on average by 1.4 for 50 labels and by 0.5 for 200 labels, not improving the rank when all labels are available).