TabNet: Attentive Interpretable Tabular Learning

1 INTRODUCTION . Deep neural networks have been demonstrated to be very powerful in understanding images ( He et al. , 2015 ; Simonyan & Zisserman , 2014 ; Zagoruyko & Komodakis , 2016 ) , text ( Conneau et al. , 2016 ; Devlin et al. , 2018 ; Lai et al. , 2015 ) and audio ( van den Oord et al. , 2016 ; Amodei et al. , 2015 ; Chiu et al. , 2018 ) , yielding many important artificial intelligence use cases . For these data types , a major enabler of the rapid research and development progress is the availability of canonical neural network architectures to efficiently encode the raw data into meaningful representations . Integrated with simple decision-making layers , these canonical architectures yield high performance on new datasets and related tasks with small extra tuning effort . For example , consider the image understanding task – variants of convolutional layers with residual connections , e.g . the notable ResNet ( He et al. , 2015 ) architecture , can yield reasonably good performance on a new image dataset ( e.g . in medical imaging ) or a slightly different visual recognition problem ( e.g . segmentation ) . Our focus in this paper is tabular ( structured ) data . Tabular data is indeed the most common data type in the entire addressable artificial intelligent market ( Chui et al. , 2018 ) . Yet , canonical neural network architectures for tabular data understanding have been under-explored . Instead , variants of ensemble decision trees still dominate data science competitions with tabular data ( Kaggle , 2019b ) . A primary reason for the popularity of tree-based approaches is their representation power for decision manifolds with approximately hyperplane boundaries that are commonly observed for tabular data . In addition , decision tree-based approaches are easy to develop and fast to train . They are highly-interpretable in their basic form ( e.g . by tracking decision nodes and edges ) and various interpretability techniques have been shown to be effective for their ensemble form , e.g . ( Lundberg et al. , 2018 ) . On the other hand , conventional neural network architectures based on stacked convolutional or multilayer perceptrons , may not be the best fit for tabular data decision manifolds due to being vastly overparametrized – the lack of appropriate inductive bias often causes them to fail to find robust solutions for tabular decision manifolds ( Goodfellow et al. , 2016 ) . We argue here that neural network architectures for tabular data should be redesigned to account for a ‘ decision-tree-like ’ mapping . Given the aforementioned benefits and reasonable performances of tree-based methods , why is deep learning worth exploring for tabular data ? One obvious motivation is pushing the performance albeit an increased computational cost , especially with more training data . In addition , introduction of a high-performance deep neural network architecture unlocks the benefits of gradient descent-based end-to-end deep learning for tabular data . For example , decision tree learning ( even with gradient boosting ) does not utilize back-propagation into their inputs to use an error signal to guide efficient learning of complex data types . On the other hand , with a deep neural network architecture , complex data types like images can be integrated into tabular data efficiently . Another well-known challenge for tree-based methods is learning from streaming data . Most algorithms for tree learning need global statistical information to select split points and straightforward modifications such as ( Ben-Haim & Tom-Tov , 2010 ) typically yield lower accuracy compared to learning from full data at once , yet deep neural networks show great potential for continual learning ( Parisi et al. , 2018 ) . Lastly , deep learning models learn meaningful representations which enable new capabilities such as data-efficient domain adaptation ( Goodfellow et al. , 2016 ) , generative modeling ( e.g . using variational autoencoders or generative adversarial networks ( Radford et al. , 2015 ) or semi-supervised learning ( Dai et al. , 2017 ) . As one example of these potential new capabilities , we demonstrate the potential of semi-supervised learning in the Appendix , showing the potential benefits of information extraction from unlabeled data which non-deep learning models are much weaker at . In this paper we propose TabNet , a deep neural network architecture to make a significant leap forward towards the optimal model design for tabular data learning . Motivated by the key problems for tabular data discussed above , the design of TabNet has two goals that are often not considered jointly : state-of-the-art performance and interpretability . With decision tree motivations , TabNet brings sparsity-controlled soft feature selection , but individually for each instance . Unlike other instance-wise feature selection methods like ( Chen et al. , 2018 ) or ( Yoon et al. , 2019 ) , TabNet employs a single deep learning architecture with end-to-end learning , to map the raw data to the final decision with soft feature selection . The key aspects and contributions of TabNet are : 1 . In order to learn flexible representations and enable flexible integration into end-to-end learning , unlike most tabular data methods , TabNet inputs raw tabular data without any feature preprocessing and is trained using conventional gradient descent-based optimization . 2 . To improve performance and interpretability , TabNet utilizes a sequential attention mechanism to choose which features to reason from at each decision step , as shown in Fig . 1 . We design this feature selection to be instance-wise such that the model can decide which features to focus on separately for each input – e.g. , for income classification capital gain may be a more important feature to focus on for a middle-aged individual . Explicit selection of sparse features enables interpretability as well as more efficient learning as the model parameters are fully utilized for the most salient features at the corresponding decision step . 3 . Overall , our careful architecture design leads to two valuable properties for real world tabular learning problems : ( 1 ) TabNet outperforms other tabular learning models on various datasets for classification and regression problems from different domains , particularly those which are not saturated in performance ; and ( 2 ) TabNet enables two kinds of interpretability : local interpretability that visualizes the importance of input features and how they are combined , and global interpretability that quantifies the amount of contribution of each input feature to the trained model . 2 RELATED WORK . Feature selection : Feature selection in machine learning broadly refers to the techniques for judicious selection of a subset of features that are useful to build a good predictor for a specified response variable . Commonly-used feature selection techniques such as forward feature selection and LASSO regularization ( Guyon & Elisseeff , 2003 ) attribute importance to the features based on the entire training data set , and are referred as global methods . On the other hand , instance-wise feature selection refers to selection of the most important features , individually and differently for each input . Instance-wise feature selection was studied in ( Chen et al. , 2018 ) by training an explainer model to maximize the mutual information between the selected features and the response variable , and in ( Yoon et al. , 2019 ) by using an actor-critic framework to mimic a baseline model while optimizing the feature selection . ( Ma et al. , 2018 ) uses partial variational autoencoder to dynamically decide which piece of information to acquire next sequentially , that can be adapted for instance-wise feature selection . Unlike these approaches , our proposed method employs soft feature selection with controllable sparsity in end-to-end learning – a single model jointly performs feature selection and output mapping , enabled by the specific design of the architecture . Thus , we demonstrate superior performance with very compact representations . Tree-based learning : Tree-based models are the most common machine learning approaches for tabular data learning . The prominent strength of tree-based models is their efficacy in picking global features with the most statistical information gain ( Grabczewski & Jankowski , 2005 ) . To improve the performance of standard tree-based models by reducing the model variance , one common approach is ensembling . Among ensembling methods , random forests ( Ho , 1998 ) use random subsets of data with randomly selected features to grow many trees . XGBoost ( Chen & Guestrin , 2016 ) and LightGBM ( Ke et al. , 2017 ) are the two recent ensemble decision tree approaches that dominate most of the recent data science competitions . They are based on learning the structures of trees at first , and then updating the leaves with the streaming data . Integration of neural networks into decision trees : One direction to address the limitations of decision trees is integration of neural networks . Representing decision trees with canonical neural network building blocks , as in ( Humbird et al. , 2018 ) , yields redundancy in representation and inefficient learning . Soft ( neural ) decision trees ( Wang et al. , 2017 ; Kontschieder et al. , 2015 ) are proposed with differentiable decision functions , instead of non-differentiable axis aligned splits to construct trees . Yet , abandoning axis-aligned splits loses the automatic feature selection ability , which is important for learning from tabular data . In ( Yang et al. , 2018 ) , a soft binning function is proposed to simulate decision trees in neural networks , which needs to enumerate all possible decisions and is inefficient . In ( Ke et al. , 2019 ) , a novel neural network architecture is proposed , with the motivations of explicitly leveraging expressive feature combinations and reducing model complexity . However , learning is based on transferring knowledge from a gradient boosted decision tree . Thus , it yields very limited performance improvement compared to it , and interpretability was not considered . In ( Tanno et al. , 2018 ) , a deep learning framework is proposed based on adaptively growing the architecture from primitive blocks while representation learning into edges , routing functions and leaf nodes of a decision tree . Our proposed model TabNet differs from these methods as it embeds the soft feature selection ability into a sequential attention-based network architecture , with controllable sparsity . Attentive table-to-text models : Table-to-text models extract textual information from tabular data . Recent works ( Liu et al. , 2017 ) ( Bao et al. , 2019 ) propose an architecture based on sequential mechanism for field-level attention . Despite the high-level similarities in the architecture , TabNet aims to perform the ultimate classification or regression task considering the entire input features , rather than mapping them to a different data type . 3 TABNET MODEL . 3.1 PRINCIPLES . We initially consider the implementation of a decision tree-like output manifold using conventional neural network building blocks ( Fig . 2 ) . Individual feature selection is the key idea to obtain decision boundaries in hyperplane form , which can be generalized for linear combination of features where constituent coefficients determine the proportion of each feature in the decision boundary . We aim to generalize this type of tree-like functionality by : • Utilizing sparse instance-wise feature selection , learned based on the training dataset . • Constructing a sequential multi-step architecture , where each decision step can contribute to a portion of the decision that is based on the selected features . • Improving the model capacity by non-linear processing of the selected features . • Ensembling via higher feature dimension and more decision steps . 3.2 OVERALL ARCHITECTURE . Fig . 3 depicts the TabNet architecture . Tabular data inputs are comprised of numerical and categorical features . We use the raw numerical features and we consider mapping of categorical features with FC BN Sparsem ax + Attentive transformer FC BN Sparsem ax + Prior scales + Features Feature transformer Feature transformer MaskAttentive transformer BN Split Split … trainable embeddings1 . We do not consider any global normalization of features , but merely apply batch normalization . We pass the same D-dimensional features f ∈ < B×D to each decision step , where B is the batch size . TabNet is based on sequential multi-step processing , with Nsteps decision steps . The ith step inputs the processed information from the ( i− 1 ) th step to decide which features to use and outputs the processed feature representation to be aggregated into the overall decision . The idea of top-down attention in sequential form is inspired from its applications in processing visual and language data such as for visual question answering ( Hudson & Manning , 2018 ) or in reinforcement learning ( Mott et al. , 2019 ) while searching for a small subset of relevant information in high dimensional input . Ablation studies in Appendix focus on the impact of various TabNet design choices , explained next . Overall , the performance is not too sensetitive to most hyperparameters , and guidelines on selection of the important hyperparameters are also provided in Appendix . Feature selection : We employ a learnable sparse mask M [ i ] ∈ < B×D for soft selection of the salient features . Through sparse selection of the most salient features , the learning capacity of a decision step is not wasted on irrelevant features , and thus the model becomes more parameter efficient . The masking is in multiplicative form , M [ i ] · f . We use an attentive transformer ( see Fig . 3 ) to obtain the masks using the processed features from the preceding step , a [ i− 1 ] : M [ i ] = sparsemax ( P [ i− 1 ] · hi ( a [ i− 1 ] ) ) . ( 1 ) Sparsemax normalization ( Martins & Astudillo , 2016 ) encourages sparsity by mapping the Euclidean projection onto the probabilistic simplex , which is observed to be superior in performance and aligned with the goal of sparse feature selection for most real-world datasets . Note that Eq . 1 has the normalization property , ∑D j=1 M [ i ] b , j = 1. hi is a trainable function , shown in Fig . 3 using a fully-connected layer , followed by batch normalization . P [ i ] is the prior scale term , denoting how much a particular feature has been used previously : P [ i ] = i∏ j=1 ( γ −M [ j ] ) , ( 2 ) where γ is a relaxation parameter – when γ = 1 , a feature is enforced to be used only at one decision step and as γ increases , more flexibility is provided to use a feature at multiple decision steps . P [ 0 ] is initialized as all ones . To further control the sparsity of the selected features , we propose sparsity regularization in the form of entropy ( Grandvalet & Bengio , 2004 ) : Lsparse = − 1 Nsteps ·B ∑Nsteps i=1 ∑B b=1 ∑D j=1 Mb , j [ i ] log ( Mb , j [ i ] + ) , ( 3 ) where is a small number for numerical stability . We add the sparsity regularization to the overall loss , with a coefficient λsparse . Sparsity may provide a favorable inductive bias for convergence to higher accuracy for some datasets where most of the input features are redundant . Feature processing : We process the filtered features using a feature transformer ( see Fig . 3 ) to obtain the features that are split for the decision step output and information for the subsequent step : [ d [ i ] , a [ i ] ] = fi ( M [ i ] · f ) , ( 4 ) where d [ i ] ∈ < B×Nd and a [ i ] ∈ < B×Na . For parameter-efficient learning and efficient convergence with the high model capacity , a feature transformer should comprise layers that are shared across all decision steps ( as the same features are input across different decision steps ) , as well as decision step-dependent layers . In Fig . 3 , we show the implementation of a block as concatenation of two shared layers and two decision step-dependent layers . Each fully-connected layer is followed by batch normalization and gated linear unit ( GLU ) nonlinearity ( Dauphin et al. , 2016 ) 2 , eventually connected to a normalized residual connection with normalization . Normalization with √ 0.5 helps to stabilize learning by ensuring that the variance throughout the network does not change dramatically , as discussed in ( Gehring et al. , 2017 ) . For faster training , we aim for very large batch sizes in practice . To improve performance with large batch sizes , all batch normalization operations , except the one applied to the input features , are implemented in ghost batch normalization ( Hoffer et al. , 2017 ) form , with a virtual batch size BV and momentum mB . For the input features , we observe the benefit of 1For example , if there are three possible categories A , B and C for a particular features , they can be learned to be mapped to scalars 0.4 , 0.1 , and -0.2 respectively , along with the training of the model . 2In GLU , first a linear mapping is applied to the intermediate representation and the dimensionality is doubled , and then second half of the output is used to determine nonlinear processing on the first half . low-variance averaging and hence avoid ghost batch normalization . After the last layer , we split the processed representation into d [ i ] and a [ i ] . Inspired by decision-tree like aggregation as in Fig . 2 , we construct the overall decision embedding as : dout = Nsteps∑ i=1 ReLU ( d [ i ] ) . ( 5 ) Finally , we apply a linear mapping Wfinaldout , for the final decision . When discrete outputs are required , we additionally employ a softmax function during training ( and argmax during inference ) .

Summary: In this paper, the authors propose a novel architecture with sequential attention modules for tabular learning. An attention module is trained to select some elements from the (normalized) input feature, and a feature transformer takes the selected features for overall feature embedding. The model is evaluated on multiple tasks, and the proposed method outperforms prior approaches in the paper.