Is it possible to ensure that all classes are represented in the output of a scikit-learn decision tree?

Question

I am working with an ordinal classification problem with six ordered classes and I want to compare a neural network classifier with a baseline classifier that is as simple and parameter-free as possible.

In my case, I want the baseline to use only the most important single feature $X$ that the neural network uses. For this particular problem, it makes sense to look for a set of class thresholds $\{t_i\}$ that lets me classify as:

      X <= t_0 -> class 0
t_0 < X <= t_1 -> class 1
t_1 < X <= t_2 -> class 2
t_2 < X <= t_3 -> class 3
t_3 < X <= t_4 -> class 4
t_4 < X        -> class 5

I have looked at a few options, and decision trees seem to fit the bill. They are very simple, and do not require choosing any free parameters (unlike e.g. $k$-nearest neighbours which requires choosing a value for $k$). If I create the tree with the argument max_leaf_nodes=6 it is quite easy to extract the thresholds $\{t_i\}$ from the resulting decision tree after fitting. (If there is another method I could use that would fit the bill as well or better to achieve this, please let me know in the comments!)

I have divided my data into six folds while ensuring that the class distribution is very similar in all folds. For this baseline I balance the classes by oversampling, as the classes are originally somewhat unbalanced. (The ratio between samples in the most common class and the rarest class is around 5:1).

For five out of my six folds, the decision tree method works very well when I use that fold for testing and the rest for training. However, for the last combination of folds used for training, I get a decision tree where class 2 is not represented in the output. Instead, two of the leaf nodes represent class 1:

Is there any way to force the decision tree to build itself in such a way that the six leaf nodes of the tree represent the six different classes?

Answer 1

It's expected behavior that the tree will not define rules for all classes if the models doesn't see a need. If you just want to compare a neural net's class predictions, you can use tree.predict_proba(...) to get all classification probabilities. Use the outputs of the neural net after softmax and before binarize.

For the fold that had little knowledge of class 2, the probability would be 0 or close to 0 given how the tree was built. If you want to force the tree to represent class 2 and assuming that it is present in all folds of your cross validation, you can change the learning params to default.

Another, but not great way, it to upsample your data to a large degree. This is bad because it would invalidate the purpose of CV, but you could at least see a small representation of all classes and your mean classification scores will remain similar.

If there are plenty of each class, you can stratify the classes across folds.

Also, if the problem is that you have insufficient samples of a minority class, you should consider up sampling the minority during training. https://towardsdatascience.com/dealing-with-imbalanced-classes-in-machine-learning-d43d6fa19d2