Building a Soft Decision Tree in TensorFlow | by Benoit Descamps | Towards Data Science

Deep learning products have grown in interests within the industry. While being computationally expensive, they have proven efficient in tackling classification of unstructured data such as in natural language or image processing.
The goal of a high-value data science product lies not only within its high-value delivery but also in its acceptance by the business. Hence, interpretability remains an important factor.
Google Brain recently released a paper proposing the implementation of soft decision trees.
In this blog, I take some key points from their paper and illustrate how a Neural Network can mimic a tree structure.
What is a decision tree in Data Science?How can a neural network have a tree-like structure?Implementing a Soft Decision Tree in TensorFlow
What is a decision tree in Data Science?
How can a neural network have a tree-like structure?
Implementing a Soft Decision Tree in TensorFlow
For the sake of informality, in data science, a tree consists of three main ingredients:
Split conditions
Pruning conditions
Leaves
The split conditions partition the data over the leaf. The pruning conditions decide when the partitioning stops. Finally, and most importantly, the leaves contain the predicted target value which has been fitted to the data within the partition.
As illustrated in Figure 1, each split condition is a function of the features.
The splits in each inner node and the target output of the leaves are determined by the optimization of a loss function. Each data point is assigned a cost for the prediction.
where
is the target of the data-point x and
is the probability of the target class of the leaf. We have also denoted the leaf selector:
This notation will be clear in the next section.
A neural network consists of a set of smooth operations on its input features. By smooth, I mean that each operation is continuously differentiable in their variables.
The differentiability is a crucial ingredient, as the optimal variables are calculated using the gradients of the loss and the operations. This seems like a problem at first, if we want to design trees as each inner node of the tree, assign each data point of one of its children based on the split condition.
Hence, our first steps towards building a tree with a neural network is to smear out the split condition, by allowing the data to go both left and right at each split but with a probability.
Consequently, each data point does not have a unique path through the tree. They now belong to every leaf of the tree, with a certain probability, i.e. the path probability.
Let us denote the path probability from the root of the tree to the n’th-(inner)-node,
We can look back at the loss function, we have written previously, for each data point x:
This is nothing other than a weighting of the categorical cross-entropy,
with the probability of x belonging to a leaf, which is the path probability.
We have now rewritten the split condition. The leaf output can remain a probability, such as a softmax.
A last questions remains however, will the optimization converge to a proper local minimum?
The updates of the weights of the inner-nodes is dependent on the gradient of the loss and therefore the path probability
Let us have a closer look at its variation at some inner-node j:
We immediately see an issue here. If the probability of some leaf
is small for a lot of data, the weight is then barely be updated in the next step. Even more so, after this update the probability remains very small. Our model might therefore converge and get stuck in these useless local minima.
A performant tree does not only have maximally pure leafs, but distributes the data as evenly as possible over its leafs. If our model gets stuck in these local minima, we end in only a few highly populated leafs and a lot of sparsely populated leafs. In this case, we lose the usefulness of the tree, and we end up with a small ensemble of logistic regression models.
Google Brain’s team proposes a remedy to this issue by introducing the regularization,
This regularization favors evenly binary split, and penalizes early purification on the level of the inner-nodes. The decay-rate gamma adds emphasis in the early inner nodes of the tree.
Now that everything is set-up, we can discuss the TensorFlow implementation.
TensorFlow is a popular open source software library by Google Brain mainly used for the development of Deep Learning models. The user defines operations on a static computational graph and runs different sections of the graph depending on the need.
Let us summarize the different ingredients with some basic formula:
Split conditions
Pruning conditions: let’s stick to a simple maximum depth
Leaf output
Loss for each data point x
We split our work into three classes:
TreeProperties: contains all the hyper-parameter related properties of the model.
Node: defines all operation on node level, such as forwarding the probability shown below in the case of an inner node or the softmax in case of a leaf:
SoftDecisionTree: builds all the inner-nodes, leafs and outputs the predicted target for each data point.
Let’s have a closer look!
As seen in section 3.2.2,the hyperparameters self.decay_penalty and self.regularisation_penalty, corresponds to the decay-rate gamma and weight lambda respectively.
The next class is more interesting.
At instantiation, we define the weight and bias of the split probability of the Node. When building the node, we first forward the output of the node, which is either the softmax prediction probability or the probability of the split of the inner node.
The loss is also updated on the node level. Inner nodes only contribute to the loss through the regularization described earlier.
Finally, all the nodes are combined together into a tree. The output prediction of the target and the probability distribution of the nodes are combined into self.output and self.leafs_distribution.
Let us conclude with some testing on MNIST. Constructing a tree with the following hyperparameters:
max_depth = 7regularisation_penality=10.decay_penality=0.9
and choosing a training batch size of 32 and validation batch size of 256. We arrive at the following result:
Not bad. Certainly not as competitive as CNN but we see a lot of promising ideas!
Make sure to check out the full implementation.
@ Tuning Hyperparameters (part I): SuccessiveHalving
@ Custom Optimizer in TensorFlow
@ Regression Prediction Intervals with XGBOOST