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: This paper proposes SAINT, a neural network model for handling tabular data with both continuous and discrete values. SAINT uses both self-attention among variables and inter-sample attention among different samples. Using both InfoNCE and denoising objectives for pre-training strategies, SAINT was able to outrank all baseline methods in experiments with 30 different tabular datasets, consisting of binary-classification, multi-class classification, and regression.