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: This paper introduces a new method to make ensembles of decision trees differentiable, and trainable with (stochastic) gradient descent. The proposed technique relies on the concept of "oblivious decision trees", which are a kind of decision trees that use the same classifier (i.e. a feature and threshold) for all the nodes that have the same depth. This means that for an oblivious decision tree of depth d, only d classifiers are learned. Said otherwise, an oblivious decision tree is a classifier that split the data using d splitting features, giving a decision table of size 2^d. To make oblivious decision trees differentiable, the authors propose to learn linear classifiers using all the features, but add a sparsity inducing operator on the weights of the classifiers (the entmax transformation). Similarly, the step function used to split the data is replaced by a continuous version (here a binary entmax transformation). Finally, the decision function is obtained by taking the outer product of all the scores of the classifiers: [c_1(x), 1-c_1(x)] o [c_2(x), 1-c_2(x)] ... This "choice" operator transforms the d dimensional vectors of the classifier scores to a 2^d dimensional vector. Another interpretation of the proposed "differentiable oblivious decision trees" is a two layer neural network, with sparsity on the weights of the first layer,