A Deep Dive into Information Gain for Decision Trees
Decision trees are a cornerstone of machine learning, cherished for their simplicity, interpretability, and surprising predictive power. While conceptually straightforward, the devil is in the details when it comes to actually constructing these models. Perhaps the single most important design choice is the criterion used to select the best feature and split point at each node. Among the various metrics proposed over the years, information gain remains one of the most popular and theoretically justified. In this article, we‘ll take an in-depth look at what information gain is, where it comes from, how it‘s used, and why it works.
Entropy: The Foundation of Information Gain
At the heart of information gain is the concept of entropy, which quantifies the amount of uncertainty or "disorder" in a random variable. Specifically, we‘ll be dealing with Shannon entropy, named after the father of information theory, Claude Shannon. For a discrete random variable X with possible values {x_1, …, x_n} and probability mass function P(X), the Shannon entropy H(X) is defined as:
$H(X) = -\sum_{i=1}^{n} P(x_i) \log_2 P(x_i)$
In plain English, entropy is the expected amount of information conveyed by each event, measured in bits (hence the base 2 logarithm). If we have a skewed distribution where some outcomes are much more likely than others, entropy will be low – there is less uncertainty about what will happen. For a uniform distribution where all outcomes are equally probable, entropy is maximized.

So how does this relate to decision trees? Imagine we‘re trying to classify samples into different classes based on their features. At each node, we want to split the data in a way that reduces the entropy or uncertainty about the class labels. In other words, we want our splits to produce child nodes that are as "pure" or homogeneous as possible with respect to the target variable.
Let‘s consider a toy binary classification example with features X_1 and X_2 and class labels Y:

Intuitively, X_1 looks more useful for separating the classes than X_2. If we split on X_1, the resulting child nodes will be relatively pure – mostly one class or the other. Splitting on X_2 would yield child nodes that are still fairly mixed.
We can quantify this notion of purity by calculating the entropy of the class labels within each node. For node t with N_t total samples, if it has k different classes with proportions p_1, …, p_k, then its entropy is:
$H(t) = -\sum_{i=1}^{k} p_i \log_2 p_i$
where $pi = N{ti} / Nt$ and $N{ti}$ is the number of samples of class i in node t.
Information Gain: Measuring Reduction in Entropy
With entropy as our measure of node impurity, we can now define information gain as the expected reduction in entropy by splitting on a particular feature. The information gain IG(t, f) for splitting node t on feature f is:
$IG(t, f) = H(t) – \sum{v \in Values(f)} \frac{N{tv}}{N_t} H(t_v)$
where Values(f) is the set of possible values for feature f, $N_{tv}$ is the number of samples in node t with f=v, and $H(t_v)$ is the entropy of the child node with f=v.
In other words, information gain is the difference between the entropy of the parent node and the weighted average entropy of the child nodes, where the weights are proportional to how many samples go into each child.
Let‘s revisit our earlier example and calculate the information gain for splitting on X_1 vs. X_2. We‘ll discretize X_1 into {X_1 <= 0.5, X_1 > 0.5} and X_2 into {X_2 <= 0.5, X_2 > 0.5}.
-
Parent entropy:
H(t) = -0.5 log2(0.5) – 0.5 log2(0.5) = 1.0 -
X_1 split:
- Left child (X_1 <= 0.5): 4/5 class 0, 1/5 class 1. Entropy = 0.72
- Right child (X_1 > 0.5): 1/5 class 0, 4/5 class 1. Entropy = 0.72
- Information gain = 1.0 – 0.5 0.72 – 0.5 0.72 = 0.28
-
X_2 split:
- Left child (X_2 <= 0.5): 1/2 class 0, 1/2 class 1. Entropy = 1.0
- Right child (X_2 > 0.5): 1/2 class 0, 1/2 class 1. Entropy = 1.0
- Information gain = 1.0 – 0.5 1.0 – 0.5 1.0 = 0.0
As expected, X_1 provides a positive information gain while X_2 provides none at all. This confirms our intuition that X_1 is the better split.
Applying Information Gain in Decision Trees
Now that we understand what information gain represents, let‘s see how it‘s typically used in the decision tree learning algorithm:
- Start with all the training samples at the root node
- Calculate the information gain for each feature
- Choose the feature with the highest information gain as the split feature for the current node
- Create child nodes based on the selected feature and recursively repeat steps 2-4 on each child
- Stop recursing when a node meets some stopping criteria (e.g. pure class, maximum depth, minimum samples)
Here‘s some pseudocode to make it concrete:
function BuildTree(samples, features):
if stopping_criteria_met(samples):
return LeafNode(samples)
best_feature, best_gain = None, 0
for feature in features:
gain = information_gain(samples, feature)
if gain > best_gain:
best_feature, best_gain = feature, gain
tree = InternalNode(best_feature)
for value in feature_values(best_feature):
child_samples = [sample for sample in samples if sample[best_feature] == value]
child = BuildTree(child_samples, features - {best_feature})
tree.add_child(child, value)
return tree
Some important things to note:
- Choosing the split feature greedily based on information gain doesn‘t guarantee finding the globally optimal tree, but it works well in practice
- There are other ways to incorporate information gain, such as splitting on the best feature at each level rather than recursively
- The stopping criteria control the complexity of the tree and can significantly impact generalization
- The computational complexity is O(m n log n) for m features and n samples, assuming balanced splits
- Continuous features must be discretized, which can be done greedily or in a preprocessing step
- There are techniques like reduced error pruning to prevent overfitting after growing the full tree
Why Information Gain Works
On a conceptual level, information gain is a principled way to quantify what we intuitively want in a decision tree – splits that reduce uncertainty and help distinguish between classes. By greedily optimizing the objective of reducing entropy, it leads to shorter trees that focus on the most relevant features.
Mathematically, there are close connections between information gain and other important quantities like mutual information and Kullback-Leibler divergence. In fact, maximizing information gain is equivalent to maximizing the mutual information between the features and class labels. Mutual information I(X;Y) measures how much knowing one variable reduces uncertainty about the other:
$I(X;Y) = H(X) – H(X|Y) = H(Y) – H(Y|X)$
where H(X|Y) is the conditional entropy of X given Y. Maximizing I(X;Y) means we‘re looking for features X that have high individual entropy H(X) but low conditional entropy H(X|Y) once we know the class labels Y. In other words, good features are ones that have a strong correlation with the target variable.
Another perspective is that by minimizing the conditional entropy H(Y|X), information gain is approximating the Bayes optimal classifier. The Bayes classifier predicts the most probable class label given the features:
$y^* = \arg\max_y P(Y=y|X)$
Minimizing H(Y|X) corresponds to maximizing the likelihood of the training data under a model that predicts Y based on X. With enough data, this converges to the true conditional distribution P(Y|X) and the Bayes classifier.
Limitations and Alternatives
Despite its popularity and strong theoretical grounding, information gain is not without flaws. One well-known issue is its bias towards features with many possible split points. All else equal, features with higher cardinality will tend to have higher information gain because they can more finely partition the data. In the extreme case, a unique identifier feature would produce pure leaf nodes and maximal information gain, but would fail to generalize to new data.
To address this, there are entropy-based splitting criteria with normalization terms to penalize high cardinality, such as gain ratio and symmetric uncertainty. These divide information gain by the entropy of the feature values to account for how much "choice" there is in splitting:
$GainRatio(t, f) = \frac{IG(t, f)}{IV(f)}$
$SymmU(t, f) = 2 \frac{IG(t, f)}{H(t) + IV(f)}$
where the intrinsic value IV(f) is just the entropy of the feature values:
$IV(f) = -\sum{v \in Values(f)} \frac{N{tv}}{N_t} \log2 \frac{N{tv}}{N_t}$
In practice, these modifications tend to perform similarly to information gain and can depend on the dataset and other hyperparameters.
Another fundamental limitation of information gain (and other greedy splitting criteria) is that it‘s a univariate approach – each feature is evaluated independently at each split. This can miss interactions between features that only appear when considered jointly. One way to capture feature interactions is to use a multivariate splitting criterion like chi-square or F-test. These statistical tests compare the observed distribution of class labels within child nodes to what would be expected under a null hypothesis of independence between the features and class.
Alternatively, there are tree-based methods that don‘t rely on explicit splitting criteria at all, such as Bayesian CART and conditional inference trees. These use statistical models and hypothesis tests to determine splits in a principled way without resorting to greedy search. Of course, the tradeoff is increased complexity and computational cost.
Beyond variations of decision trees, an entirely different paradigm is to learn feature combinations end-to-end using gradient descent, as in neural networks. With enough data and careful regularization, deep learning can learn highly expressive feature interactions that identify patterns a simple decision tree might miss. The downside is loss of interpretability and potential instability.
Current Research and Future Directions
Even after decades of use, information gain remains an active area of research with ongoing efforts to improve and extend it. Some recent developments:
- Efficient approximation algorithms for finding splits on streaming or distributed data
- Optimizations for sparse, high-dimensional feature spaces common in text and genomics
- Generalizations of information gain to handle missing data, multi-label classification, and regression
- Combinations of information gain with other criteria like Gini impurity and chi-square
- Applications of information gain for feature selection and dimensionality reduction outside of decision trees
- Information theoretic regularization terms to penalize complex trees and encourage generalization
- Integration of information gain into probabilistic graphical models and causal inference frameworks
- Adaptation of information gain for non-tabular data types like images, time series, and graphs
As the field progresses, information gain will continue to evolve and find new use cases. Nonetheless, it remains a fundamental concept that all data scientists should understand. By deeply engaging with the underlying math and intuition behind it, we can all become better equipped to make principled decisions when constructing and interpreting machine learning models.
Conclusion
We‘ve covered a lot of ground in this deep dive into information gain – from the basics of entropy to its mathematical properties, algorithmic implementation, theoretical justification, and current frontiers. While it can seem daunting at first, the core idea is quite intuitive: choosing splits that reduce uncertainty and yield purer child nodes. By greedily maximizing information gain, we can build concise, interpretable trees that hone in on the most discriminative features.
Of course, information gain is not a panacea and it‘s important to understand its limitations and alternatives. It‘s biased towards high cardinality features, can miss interactions, and is prone to overfitting without careful regularization. There will always be a tradeoff between simplicity and expressivity that depends on the specific application.
Ultimately, the best splitting criterion will depend on the nature of your data, computational constraints, and prediction goals. However, information gain is a solid default choice that has stood the test of time. By deeply understanding this foundational concept, you‘ll be well poised to make informed, data-driven decisions in all your machine learning endeavors. So get out there and start experimenting – and always keep an eye out for ways to gain new insights!