SALT : Sharing Attention between Linear layer and Transformer for tabular dataset

Handling tabular data with deep learning models is a challenging problem despite their remarkable success in vision and language processing applications . Therefore , many practitioners still rely on classical models such as gradient boosting decision trees ( GBDTs ) rather than deep networks due to their superior performance with tabular data . In this paper , we propose a novel hybrid deep network architecture for tabular data , dubbed SALT ( Sharing Attention between Linear layer and Transformer ) . The proposed SALT consists of two blocks : Transformers and linear layer blocks that take advantage of shared attention matrices . The shared attention matrices enable transformers and linear layers to closely cooperate with each other , and it leads to improved performance and robustness . Our algorithm outperforms tree-based ensemble models and previous deep learning methods in multiple benchmark datasets . We further demonstrate the robustness of the proposed SALT with semi-supervised learning and pre-training with small dataset scenarios . 1 INTRODUCTION . In the fields of vision and natural language processing , deep networks such as CNN , RNN , LSTM and Transformer have gained great popularity with its impressive performance . In particular , Transformer ( Vaswani et al. , 2017 ) designed as language model , improves the performance of lots of deep learning models in various domains , so that there are many powerful models based on Transformer ( e.g . Devlin et al . ( 2019 ) , Brown et al . ( 2020 ) and Dosovitskiy et al . ( 2021 ) ) . Although the deep networks are powerful and used in natural language processing and vision , they are sub-optimal for other types of real-world problems that require to use tabular data , such as fraud detection ( Luo et al. , 2019 ) , product recommendation ( Guo et al. , 2017 ) and disease prediction ( Koppu et al. , 2020 ) . Tabular data has different characteristics from other data . Unlike text data that include vocabularies and words , and image data that include RGB values , the tabular data are usually mixed with different types of complex variables . For example , the tabular data contain continuous variables such as such as age , height , and weight have different ranges of values , and categorical variables which are independent from one another like gender and nationality . For these reasons , deep learning models were not quite successful with tabular data despite its strength in natural language processing and vision fields . Instead of deep learning models , classical models such as tree-based ensemble models are mainly used for tabular data . However , these classical approaches have some limitations . Continuous learning with real-time data is quite difficult using these classical methods . When the tabular data is high-dimensional sparse , the performance of tree-based methods is degraded . Also , the tree-based ensemble methods do not perform well for multi-modality learning and end-to-end systems . Therefore , it is an interesting to study a deep learning models for tabular data . Many studies have been attempted not only to overcome the shortcomings of the classical model with a deep learning model but also to overcome its performance . ( Arik & Pfister , 2020 ) , ( Huang et al. , 2020 ) and ( Somepalli et al. , 2021 ) . In this paper , we propose a new hybrid deep learning model architecture , named as SALT ( Sharing Attention between Linear layer and Transformer ) . We summarize the contributions of our paper as follows : • SALT shares the attention matrices between two blocks . Sharing the attention matrix allows to learn two blocks strongly and effectively . We demonstrate that sharing attention matrices improves the performance better . • SALT introduce the improved embedding method for continuous variables . We demonstrate that this method performs better than others . • SALT has four types of variants depending on which block is used and in what direction the attention matrix is shared . Each variant shows strength in different data . • SALT outperforms the other models on average over a variety of benchmark datasets . It also performs well even in small data environments via semi-supervised and self-supervised learning 2 RELATED WORK . 2.1 TREE-BASED MODELS . Decision trees ( Quinlan , 1986 ) are well-known for their high predictive performance compared to computational complexity . In addition , it has the strength of having explanatory power in units of variables with most statistical information gain . However , decision trees are likely to work well only on specific data because their decision boundaries are perpendicular to the data axis . To improve this drawback and performance , there are many ensemble models of decision trees such as Random forest ( Breiman , 2001 ) and GBDTs ( Gradient Boosting Decision Trees ) . Especially , GBDTs methods such as XGBoost ( Chen & Guestrin , 2016 ) , LightGBM ( Ke et al. , 2017 ) and CatBoost ( Dorogush et al. , 2018 ) are commonly used with powerful performance in lots of machine learning competitions and industrial sites . 2.2 DEEP LEARNING MODELS FOR TABULAR DATA . There are several studies on deep networks to overcome the limitations of tree-based models and outperform the performance of GBDTs . Especially there are deep learning models based on attention mechanism . TabNet ( Arik & Pfister , 2020 ) , for example , is designed to learn similarly to a decision trees , and it has interpretability with the attentive layer . It shows better performance than GBDTs in some dataset . TabTransformer ( Huang et al. , 2020 ) is designed based on Transformer and has contextual embedding values . However , it embeds only categorical variables , so that there is a limitation for continuous variables . There are some models that improve the limitation of TabTransformer ( Song et al. , 2018 ) , ( Somepalli et al. , 2021 ) . They have the contextual embedding values of not only categorical variables but also continuous variables . Especially a model called SAINT ( Somepalli et al. , 2021 ) introduces inter-sample self-attention method and shows powerful performance . But it has a fatal disadvantage that training costs are high . 2.3 TRANSFORMER . Self-attention is the core module of Transformer ( Vaswani et al. , 2017 ) . The self-attention consists of three parameter matrices : K ( keys ) , Q ( queries ) , and V ( values ) . Formally , input embedding values X ∈ Rn×d of n features of dimensions d , are projected using WQ ∈ Rd×dq , WQ ∈ Rd×dq , and WQ ∈ Rd×dq to extract feature representations Q , K , and V . With Q , K , and V , self-attention can be written as , Q = XWQ , K = XWK , V = XWV ( 1 ) Attention ( Q , K , V ) = Softmax ( QKT√ dk ) V ( 2 ) MHA ( Multi Head Attention ) is having multiple attention heads . Multi head allows the attention matrix to have the abundant representations ( Michel et al. , 2019 ) , ( Voita et al. , 2019 ) . Each head attention has different Q , K , and V weight matrices and calculates the attention values with the equation 2 . All heads are then concatenated and multiplied by the weight matrix to generate the final output of the layer . MHA ( Q , K , V ) = Concat ( head1 , ... , headh ) WO ( 3 ) where headi = Attention ( Q , K , V ) and WO is weighted matrix for final output . The dimension dh of each head is typically given as d/Nh . 2.4 GATING MECHANISM . The gating mechanism is an effective method of learning by controlling the information flow path of the LSTM ( Hochreiter & Schmidhuber , 1997 ) . However , the more layers of these mechanisms are stacked , the more likely the gradient is vanishing . As an improved mechanism , gate linear units ( GLU ) is introduced by ( Dauphin et al. , 2017 ) . GLU has been used in many deep learning models and shows the better performance ( Arik & Pfister , 2020 ) , ( Shazeer , 2020 ) . To briefly explain the GLUs , they divide the input in half , take an activation function on one side , and multiply by element with the other . Therefore , the output dimension is half the input dimension . h ( x ) = ( W1X + b1 ) ⊗ σ ( W2X + b2 ) ( 4 ) The gating mechanism is used to control the information flow slightly similar to the self-attention mechanism . The main difference between gating mechanism and self-attention is that gating mechanism controls only the bandwidth of a single element , while self-attention considers information of two different elements . 3 SALT : SHARING ATTENTION BETWEEN LINEAR LAYER AND TRANSFORMER . 3.1 SALT ARCHITECTURE . In this section , we introduce our model , SALT ( Sharing Attention between Linear layer and Transformer ) shown in Figure 1 . SALT uses the input embedding values obtained from the embedding layer . The embedding values have the shape of ’ features ( n ) × embedding dimensions ( d ) ’ . The main body of SALT has two blocks : Transformers block inspired by ( Vaswani et al. , 2017 ) and Linear layer block inspired by ( Liu et al. , 2021 ) . Each block has subblocks by feature-wise and dimensionwise . These two subblocks allow communications between different features and different embedding elements and makes the model robust ( Tolstikhin et al. , 2021 ) . The output values from two blocks become the contextual embedding values . SALT performs fin-tuning and pre-training with contextual embedding values . 3.2 SALT FOR LEARNING . The learning process of SALT is as follows . Let D = { xi , yi } mi=1 be tabular dataset of m samples . The feature ( n ) variables xfeatures ∈ Rn include categorical xcat and continuous variables xcont . SALT adds the special token [ cls ] to the feature variables and takes them as an input values like BERT ( Devlin et al. , 2019 ) . So , xi = [ [ cls ] , xcat , xcont ] is the input values consisting of the feature variables and a special token [ cls ] . The embedding layer E ( · ) converts the input values xi ∈ R ( n+1 ) into d-dimensional values E ( xi ) ∈ R ( n+1 ) ×d . Let Transformer be Transformer ( · ) and Linear layer block be Linearlayer ( · ) . SALT consists of an L stack of these two blocks . Each block returns the output value z and the sharing attention matrices , sf and sd by feature-wise and dimension-wise as the following equations : z ( 1 ) t , s ( 1 ) tf , s ( 1 ) td = Transformer1 ( E ( xi ) ) ( 5 ) z ( 1 ) l , s ( 1 ) lf , s ( 1 ) ld = Linearlayer1 ( E ( xi ) ) ( 6 ) z ( 1 ) t and z ( 1 ) l are the output values of Transformer block and Linear layer block from the first stack , respectively . s ( 1 ) ft and s ( 1 ) dt are the feature-wise attention matrix and dimension-wise attention matrix from Transformer block . s ( 1 ) fl and s ( 1 ) dl are from Linear layer block . The matrices from blocks are calculated in the head direction by function S ( · ) and the calculated matrix becomes the sharing attention matrix for the next blocks . Sf ( sft , sfl ) = concat ( sft , sfl ) Wf ( 7 ) Sd ( sdt , sdl ) = concat ( sdt , sdl ) Wd ( 8 ) Wf , Wd ∈ Rh×2h are weight for projection in the head direction . The two attention matrices obtained from the above equation are sent to the blocks of the next stack . This procedure is repeated until the last stack as the following equations . z ( i ) t , s ( i ) ft , s ( i ) dt = Transformeri ( z ( i−1 ) t , s̃ ( i−1 ) f , s̃ ( i−1 ) d ) ( 9 ) z ( i ) l , s ( i ) fl , s ( i ) dl = Linearlayeri ( z ( i−1 ) l , s̃ ( i−1 ) f , s̃ ( i−1 ) d ) ( 10 ) s̃ ( i ) f = Sf ( s ( i ) ft , s ( i ) fl ) s̃ ( i ) d = Sd ( s ( i ) dt , s ( i ) dl ) ( 11 ) Finally , the output values of the last stack , zLl and z L t are added as the output of SALT . This value also becomes the contextual embedding valuesẼ ( xi ) = [ ẽ ( [ cls ] ) , ẽ ( x1 ) , · · · , ẽ ( xn ) ] . The cls token , ẽ ( [ cls ] ) is used for fine-tuning , and the other values , ẽ ( xi ) are used for pre-training using MLM ( Masked Language Model ) ( Devlin et al. , 2019 ) .

Summary: This paper proposes a hybrid deep network, named SALT (Sharing Attention between Linear layer and Transformer). The SALT consists of two blocks: Transformer and linear layer blocks that take advantage of shared attention matrices. They compare SALT with tree-based ensemble models and previous deep learning models on multiple benchmark datasets. It furher shows robustness of the proposed SALT with semi-supervised learning and pre-training with small dataset scenarios.   The main body of SALT has two blocks: Transformers block inspired by (Vaswani et al., 2017) and Linear layer block inspired by (Liu et al., 2021). Each block has subblocks by feature-wise and dimension- wise. These two subblocks allow communications between different features and different embed- ding elements and makes the model robust (Tolstikhin et al., 2021). The output values from two blocks become the contextual embedding values. SALT performs fin-tuning and pre-training with contextual embedding values.  For supervised learning setting, the mean AUROC score of the proposed SALT improves upon the treebased LightGBM by mere 0.09%, and it improves upon the best deep learning based model SAINT by 0.29%