Publish AI, ML & data-science insights to a global community of data professionals.

Understand Decision Tree Classifiers

Understand how Decision Tree Classifier works in plain language and minimum math equations. Figure out how Gini Impurity and Information…

By my son, Charles Zhu
By my son, Charles Zhu

Compare with machine learning models like Neural Network, I thought Decision Tree Classifier should be the most simple one. But I was wrong, this model is a bit complex than I thought. And the model also lands the foundation for other advanced models like LightBGM and Random Forest Decision Tree. So, I spent some time learning it and try to figure out how Decision Tree Classifier works.

How decision tree works

The model works very much like how a human mind classifying objects in the real world.

Apple and Lemon, image from Unsplash
Apple and Lemon, image from Unsplash

When you see the above fruits, you instantly identify the lemonade from the apple. But, how? By following the decision tree model, your mind will ask 2 questions in a blink.

  1. What color is this looks like?
  2. What size is this fruit?
A simplified fruit identifier decision tree
A simplified fruit identifier decision tree

The decision tree already well maintained in your memory which will lead you to the right answer.

But, when you are trying to teach a computer model to do a similar thing, how to build a decision tree like this? or let’s ask the same question in other two more specified questions.

  1. Which question to ask first? (which feature column to target)
  2. What question I should ask? (what is the data partition rule)

The key of the Decision Tree Classifier is to build a decision tree like the one showing above.

A plate with 3 fruits – the sample data

Let’s say we have a training data set in hand. This time, we have three fruits on the plate: Apple, Grape, and Lemon. As showing below:

3 fruits in a Pandas Dataframe, Image by Andrew
3 fruits in a Pandas Dataframe, Image by Andrew

To solve the questions: How do I know which question to ask first? or ask the color question or diameter question first? and set the right data dividing point?

We might want our first question to divide the dataset so that in much better shape. If you have experience cleaning the kid’s toy, the question we ask should divide lego and woodblock the best. or in another narrative, we want the question to lead to the next level of data with less diversity.

Wikipedia will tell you the metric to measure diversity is called Gini Impurity.

What is Gini impurity and why it works like this

For example, you have two fruit arrays in front of you. Which fruit array is more diverse?

Fruit Array A, images from Unsplash
Fruit Array A, images from Unsplash
Fruit Array B, images from Unsplash
Fruit Array B, images from Unsplash

Obviously, our human brain will reach a quick answer to Array B. But how could we measure the diversity in a quantity way? I would suggest you stop here try to come up with your solution before moving on.

One concise solution is to calculate the probability of each fruit, then sum up the squared probability. the sum-up number will have one property, the more diversity of the array is, the lower the number will be.

For example, in fruit array A.

Then sum up the squared probabilities.

Now, in fruit array B. with the same logic, you will get the sum up squared probabilities as 0.5.

See? the more diverse array will generate a lower result. And Gini Impurity is simply:

Now, you may ask, why square the probability? without a square, the array A and B will generate the same result. (you can give it a try).

Use Information Gain to determine the right question to ask

Let us back to the fruit plate, we can produce the Gini impurity is 0.64 in its initial status.

Next assume we have a question to ask, the question will partition the dataset into two subsets.

In a real program, we can use a nest 2-for loop to try all cells and produce all Gini Impurity #.

And what is Information Gain? Information Gain is a measurement that used to measure how much uncertainty is reduced by a question.

In the above sample. The information gained by the question "Fruit diameter ≥3?" is:

Provided that we have the best question to ask, the question must reduce the Gini Impurity the most, or in other words, gain the most information.

Build the decision tree recursively

Now that you have the Gini and Info Gain solution, you can recursively apply the process to your dataset until the process can’t generate any more Information Gain. In the above sample, the right child contains only grapes. you can no longer reduce impurity(or gain info) because there is no need to ask any questions.

In the end, you will have a functioning decision tree model that will accept new data and predict the fruit type for you.

Links and Code


Towards Data Science is a community publication. Submit your insights to reach our global audience and earn through the TDS Author Payment Program.

Write for TDS

Related Articles