Neural Oblivious Decision Ensembles for Deep Learning on Tabular Data

1 INTRODUCTION . The recent rise of deep neural networks ( DNN ) resulted in a substantial breakthrough for a large number of machine learning tasks in computer vision , natural language processing , speech recognition , reinforcement learning ( Goodfellow et al. , 2016 ) . Both gradient-based optimization via backpropagation ( Rumelhart et al. , 1985 ) and hierarchical representation learning appear to be crucial in increasing the performance of machine learning for these problems by a large margin . While the superiority of deep architectures in these domains is undoubtful , machine learning for tabular data still did not fully benefit from the DNN power . Namely , the state-of-the-art performance in problems with tabular heterogeneous data is often achieved by “ shallow ” models , such as gradient boosted decision trees ( GBDT ) ( Friedman , 2001 ; Chen & Guestrin , 2016 ; Ke et al. , 2017 ; Prokhorenkova et al. , 2018 ) . While the importance of deep learning on tabular data is recognized by the ML community , and many works address this problem ( Zhou & Feng , 2017 ; Yang et al. , 2018 ; Miller et al. , 2017 ; Lay et al. , 2018 ; Feng et al. , 2018 ; Ke et al. , 2018 ) , the proposed DNN approaches do not consistently outperform the state-of-the-art shallow models by a notable margin . In particular , to the best of our knowledge , there is still no universal DNN approach that was shown to systematically outperform the leading GBDT packages ( e.g. , XGBoost ( Chen & Guestrin , 2016 ) ) . As additional evidence , a large number of Kaggle ML competitions with tabular data are still won by the shallow GBDT methods ( Harasymiv , 2015 ) . Overall , at the moment , there is no dominant deep learning solution for tabular data problems , and we aim to reduce this gap by our paper . We introduce Neural Oblivious Decision Ensembles ( NODE ) , a new DNN architecture , designed to work with tabular problems . The NODE architecture is partially inspired by the recent CatBoost package ( Prokhorenkova et al. , 2018 ) , which was shown to provide state-of-the-art performance on a large number of tabular datasets . In a nutshell , CatBoost performs gradient boosting on oblivious decision trees ( decision tables ) ( Kohavi , 1994 ; Lou & Obukhov , 2017 ) , which makes inference very efficient , and the method is quite resistant to overfitting . In its essence , the proposed NODE architecture generalizes CatBoost , making the splitting feature choice and decision tree routing differentiable . As a result , the NODE architecture is fully differentiable and could be incorporated in any computational graph of existing DL packages , such as TensorFlow or PyTorch . Furthermore , NODE allows constructing multi-layer architectures , which resembles ” deep ” GBDT that is trained end-to-end , which was never proposed before . Besides the usage of oblivious decision tables , another important design choice is the recent entmax transformation ( Peters et al. , 2019 ) , which effectively performs a ” soft ” splitting feature choice in decision trees inside the NODE architecture . As discussed in the following sections , these design choices are critical to obtain state-of-the-art performance . In a large number of experiments , we compare the proposed approach with the leading GBDT implementations with tuned hyperparameters and demonstrate that NODE outperforms competitors consistently on most of the datasets . Overall , the main contributions of our paper can be summarized as follows : 1 . We introduce a new DNN architecture for machine learning on tabular data . To the best of our knowledge , our method is the first successful example of deep architectures that substantially outperforms leading GBDT packages on tabular data . 2 . Via an extensive experimental evaluation on a large number of datasets , we show that the proposed NODE architecture outperforms existing GBDT implementations . 3 . The PyTorch implementation of NODE is available online1 . The rest of the paper is organized as follows . In Section 2 we review prior work relevant to our method . The proposed Neural Oblivious Decision Ensembles architecture is described in Section 3 and experimentally evaluated in Section 4 . Section 5 concludes the paper . 2 RELATED WORK . In this section , we briefly review the main ideas from prior work that are relevant to our method . The state-of-the-art for tabular data . Ensembles of decision trees , such as GBDT ( Friedman , 2001 ) or random forests ( Barandiaran , 1998 ) , are currently the top choice for tabular data problems . Currently , there are several leading GBDT packages , such as XGBoost ( Chen & Guestrin , 2016 ) , LightGBM ( Ke et al. , 2017 ) , CatBoost ( Prokhorenkova et al. , 2018 ) , which are widely used by both academicians and ML practitioners . While these implementations vary in details , on most of the tasks their performances do not differ much ( Prokhorenkova et al. , 2018 ; Anghel et al. ) . The most important distinction of CatBoost is that it uses oblivious decision trees ( ODTs ) as weak learners . As ODTs are also an important ingredient of our NODE architecture , we discuss them below . Oblivious Decision Trees . An oblivious decision tree is a regular tree of depth d that is constrained to use the same splitting feature and splitting threshold in all internal nodes of the same depth . This constraint essentially allows representing an ODT as a table with 2d entries , corresponding to all possible combinations of d splits ( Lou & Obukhov , 2017 ) . Of course , due to the constraints above , ODTs are significantly weaker learners compared to unconstrained decision trees . However , when used in an ensemble , such trees are less prone to overfitting , which was shown to synergize well with gradient boosting ( Prokhorenkova et al. , 2018 ) . Furthermore , the inference in ODTs is very efficient : one can compute d independent binary splits in parallel and return the appropriate table entry . In contrast , non-oblivious decision trees require evaluating d splits sequentially . Differentiable trees . The significant drawback of tree-based approaches is that they usually do not allow end-to-end optimization and employ greedy , local optimization procedures for tree construction . Thus , they can not be used as a component for pipelines , trained in an end-to-end fashion . To address this issue , several works ( Kontschieder et al. , 2015 ; Yang et al. , 2018 ; Lay et al. , 2018 ) 1https : //github.com/Qwicen/node propose to ” soften ” decision functions in the internal tree nodes to make the overall tree function and tree routing differentiable . In our work , we advocate the usage of the recent entmax transformation ( Peters et al. , 2019 ) to ” soften ” decision trees . We confirm its advantages over the previously proposed approaches in the experimental section . Entmax . The key building block of our model is the entmax transformation ( Peters et al. , 2019 ) , which maps a vector of real-valued scores to a discrete probability distribution . This transformation generalizes the traditional softmax and its sparsity-enforcing alternative sparsemax ( Martins & Astudillo , 2016 ) , which has already received significant attention in a wide range of applications : probabilistic inference , topic modeling , neural attention ( Niculae & Blondel , 2017 ; Niculae et al. , 2018 ; Lin et al. , 2019 ) . The entmax is capable to produce sparse probability distributions , where the majority of probabilities are exactly equal to 0 . In this work , we argue that entmax is also an appropriate inductive bias in our model , which allows differentiable split decision construction in the internal tree nodes . Intuitively , entmax can learn splitting decisions based on a small subset of data features ( up to one , as in classical decision trees ) , avoiding undesired influence from others . As an additional advantage , using entmax for feature selection allows for computationally efficient inference using the sparse pre-computed choice vectors as described below in Section 3 . Multi-layer non-differentiable architectures . Another line of work ( Miller et al. , 2017 ; Zhou & Feng , 2017 ; Feng et al. , 2018 ) promotes the construction of multi-layer architectures from nondifferentiable blocks , such as random forests or GBDT ensembles . For instance , ( Zhou & Feng , 2017 ; Miller et al. , 2017 ) propose to use stacking of several random forests , which are trained separately . In recent work , ( Feng et al. , 2018 ) introduces the multi-layer GBDTs and proposes a training procedure that does not require each layer component to be differentiable . While these works report marginal improvements over shallow counterparts , they lack the capability for end-to-end training , which could result in inferior performance . In contrast , we argue that end-to-end training is crucial and confirm this claim in the experimental section . Specific DNN for tabular data . While a number of prior works propose architectures designed for tabular data ( Ke et al. , 2018 ; Shavitt & Segal , 2018 ) , they mostly do not compare with the properly tuned GBDT implementations , which are the most appropriate baselines . The recent preprint ( Ke et al. , 2018 ) reports the marginal improvement over GBDT with default parameters , but in our experiments , the baseline performance is much higher . To the best of our knowledge , our approach is the first to consistently outperform the tuned GBDTs over a large number of datasets . 3 NEURAL OBLIVIOUS DECISION ENSEMBLES . We introduce the Neural Oblivious Decision Ensemble ( NODE ) architecture with a layer-wise structure similar to existing deep learning models . In a nutshell , our architecture consists of differentiable oblivious decision trees ( ODT ) that are trained end-to-end by backpropagation . We describe our implementation of the differentiable NODE layer in Section 3.1 , the full model architecture in Section 3.2 , and the training and inference procedures in section 3.3 . 3.1 DIFFERENTIABLE OBLIVIOUS DECISION TREES . The core building block of our model is a Neural Oblivious Decision Ensemble ( NODE ) layer . The layer is composed ofm differentiable oblivious decision trees ( ODTs ) of equal depth d. As an input , allm trees get a common vector x ∈ Rn , containing n numeric features . Below we describe a design of a single differentiable ODT . In its essence , an ODT is a decision table that splits the data along d splitting features and compares each feature to a learned threshold . Then , the tree returns one of the 2d possible responses , corresponding to the comparisons result . Therefore , each ODT is completely determined by its splitting features f ∈ Rd , splitting thresholds b ∈ Rd and a d-dimensional tensor of responses R ∈ R 2× 2× 2︸ ︷︷ ︸ d . In this notation , the tree output is defined as : h ( x ) = R [ 1 ( f1 ( x ) − b1 ) , . . . ,1 ( fd ( x ) − bd ) ] , ( 1 ) where 1 ( · ) denotes the Heaviside function . To make the tree output ( 1 ) differentiable , we replace the splitting feature choice fi and the comparison operator 1 ( fi ( x ) − bi ) by their continuous counterparts . There are several existing approaches that can be used for modelling differentiable choice functions in decision trees ( Yang et al. , 2018 ) , for instance , REINFORCE ( Williams , 1992 ) or Gumbel-softmax ( Jang et al. , 2016 ) . However , these approaches typically require long training time , which can be crucial in practice . Instead , we propose to use the α-entmax function ( Peters et al. , 2019 ) as it is able to learn sparse choices , depending only on a few features , via standard gradient descent . This function is a generalization of softmax in its variational form : softmax ( x ) = argmaxp∈∆ [ 〈p , x〉+H ( p ) ] where H ( p ) is Shannon entropy . We can define α-entmax by replacing H ( p ) with Tsallis α-entropy2 . The choice function is hence replaced by a weighted sum of features , with weights computed as α-entmax ( α=1.5 ) over the learnable feature selection matrix F ∈ Rd×n : f̂i ( x ) = n∑ j=1 xj · entmaxα ( Fij ) ( 2 ) Similarly , we relax the Heaviside function 1 ( fi ( x ) − bi ) as a two-class entmax , which we denote as σα ( x ) =entmaxα ( [ x , 0 ] ) . As different features can have different characteristic scales , we use the scaled version ci ( x ) = σα ( fi ( x ) −bi τi ) , where bi and τi are learnable parameters for thresholds and scales respectively . Based on the ci ( x ) values , we define a ” choice ” tensorC ∈ R 2× 2× 2︸ ︷︷ ︸ d of the same size as the response tensor R by computing the outer product of all ci : C ( x ) = [ c1 ( x ) 1− c1 ( x ) ] ⊗ [ c2 ( x ) 1− c2 ( x ) ] ⊗ · · · ⊗ [ cd ( x ) 1− cd ( x ) ] ( 3 ) The final prediction is then computed as a weighted linear combination of response tensor entries R with weights from the entries of choice tensor C : ĥ ( x ) = ∑ i1 , ... id∈ { 0,1 } d Ri1 , ... , id · Ci1 , ... , id ( x ) ( 4 ) Note , that this relaxation equals to the classic non-differentiable ODT h ( x ) ( 1 ) iff both feature selection and threshold functions reach one-hot state , i.e . entmax always returns non-zero weights for a single feature and ci always return exactly zeros or ones . 2If one is unfamiliar with this definition , we highly recommend reading Peters et al . ( 2019 ) Finally , the output of the NODE layer is composed as a concatenation of the outputs ofm individual trees [ ĥ1 ( x ) , . . . , ĥm ( x ) ] . Multidimensional tree outputs . In the description above , we assumed that tree outputs are onedimensional ĥ ( x ) ∈ R. For classification problems , where NODE predicts probabilities of each class , we use multidimensional tree outputs ĥ ( x ) ∈ R|C| , where |C| is a number of classes .

Summary: The paper tries to ask if there is a good neural net architecture that works as effectively as gradient boosting decision trees on tabular data. The authors propose an architecture (NODE) that satisfies this conditions. NODE is an architecture consisting of differentiable oblivious decision trees that can be trained end to end via back propagation. The paper is readable and the experiments are well presented. They make use of an alpha-entmax transformation to obtain a differentiable architecture. The approach seems well motivated in the literature. It is unclear how novel the contribution is. It is unclear if in the experimental section the datasets used are standard for this classes of tasks. Would be good to mention if it is the case.