Introduction to ControlBurn: How to Train Gradient Boosted Trees for Better Interpretability

Often, the role of a data scientist isn't to build cutting-edge deep learning models, but rather to construct simple, interpretable models that get the job done. Interpretability is crucial because it allows us to communicate key insights about customer behavior to management and other non-technical stakeholders.

In this article, we will discuss the common problems with feature interpretability in traditional linear models and decision trees, and introduce a new method designed to overcome these limitations.

Two Main Types of Interpretability

Feature Interpretability using a Linear Model

Consider a standard linear model:

$$\hat{f}(x) = \theta_0 + \theta_1 x_1 + \theta_2 x_2 + \cdots + \theta_d x_d$$

Interpreting this model is straightforward. The magnitudes of the coefficients $\theta_1, \ldots, \theta_d$ inform us about the importance of the corresponding features.

However, this method of interpretation has serious flaws, which we will discuss in the next section.

Feature Interpretability using a Decision Tree

As a reminder, a decision tree performs recursive partitioning. A node is split into two child nodes based on a specific feature and a threshold. This feature and threshold are chosen to ensure the resulting child nodes have the minimum possible "impurity." For classification, impurity can be measured by the Gini index, cross-entropy, or simply the misclassification error. For regression, it is typically the mean squared error.

After a decision tree is built, we can compute the importance of each feature based on how effectively it split the nodes. Specifically, we look at a quantity called Mean Decrease in Impurity (MDI). This is calculated as the impurity of the parent node minus the weighted sum of the impurities of its two child nodes:

$$I_{\text{parent}} - w_1 I_{\text{left-child}} - w_2 I_{\text{right-child}},$$

Here, $I$ represents the impurity measure, while $w_1$ and $w_2$ represent the fraction of samples that fall into the left and right child nodes, respectively.

The features are then ranked in importance based on this mean decrease in impurity.

The problem with the above methods: Correlation Bias

Let's return to the linear regression model:

$$\hat{f}(x) = \theta_0 + \theta_1 x_1 + \theta_2 x_2 + \cdots + \theta_d x_d$$

Suppose a group of features are highly correlated. A linear regression model with quadratic regularization (like Ridge regression) will tend to assign similar weights to all features in this group. Consequently, when you look at the feature importance, all these features appear together. This is called correlation bias, and is also observed in decision trees

Now, why is this a problem? If a group of correlated features is highly predictive, they will all appear together in the feature importance chart. This can cause another feature, which is also important but not correlated with the group, to be pushed down in the rankings. As a result, we might mistakenly conclude this independent feature isn't as important as it truly is.

Furthermore, we don't gain much new insight from seeing many similar features listed as important. This redundancy doesn't inform us as much as identifying distinct, important factors would.

Hence, ideally, we want a feature importance list that highlights features that are different from each other.

A Potential Solution: LASSO Regression

In LASSO regression, we penalize the sum of the absolute values of the coefficients. This encourages sparsity, causing the model to select only a single feature from a group of highly correlated features. The weights of the other features in that group are pushed to zero. This appears to solve our problem.

However, LASSO regression isn't without its flaws. If you train a LASSO model multiple times on the same set of features (e.g., on different subsets of the data), it can arbitrarily pick a different feature from that same correlated group each time. This may (slightly) mess up interpretability

There's a second, related problem. When the model picks one feature from a group and drops the others, the predictive effect of those dropped features is "picked up" by the selected feature. This means if you observe a large coefficient in a LASSO model, you don't know if The feature is genuinely important on its own, or it is simply acting as a proxy for a larger group of correlated features. Again, this ambiguity harms interpretability.

So, what about decision trees and random forests? In practice, we rarely use linear models. We usually use ensemble models like random forests or boosted trees, as they perform exceptionally well on tabular data.

In a coming section, we'll discuss a technique specifically designed to make these more complex models interpretable.

The Idea Behind Gradient Boosting

Gradient boosting is a powerful ensemble technique where models are built sequentially, with each new model improving upon the previous ones.

The model at the $m$-th step, $F_m(x)$, is built upon the previous model $F_{m-1}(x)$:

$$F_m(x) = F_{m-1}(x) + \gamma h_m(x),$$

where $\gamma$ is a constant called the learning rate, and $h_m$ is also a model. In particular, $h_m$ is a simple model which is often called a weak learner. Let's first discuss how we train $h_m$ and then we can discuss why we do it this way

How the Weak Learner is Trained

$h_m$ is trained to learn the negative gradient of the loss function, evaluated at the previous model's predictions.

At each step $m$, we create a new training set

$$\{ (x_1, r_1), (x_2, r_2), \ldots, (x_N, r_N) \},$$

where

$$\{ (x_1, y_1), (x_2, y_2), \ldots, (x_N, y_N) \}$$

is the original training set. The target $r_i$ for the data point $x_i$ is defined as:

$$r_i = - \nabla_{F_{m-1}} L \big( y_i, F_{m-1}(x_i) \big)$$

This expression is simply the gradient of the loss function $L$ with respect to the model's current prediction, $\hat{y}_i = F_{m-1}(x_i)$:

$$\nabla_{F_{m-1}} L \big( y_i, F_{m-1}(x_i) \big) = \frac{\partial L(y_i, \hat{y}_i)}{\partial \hat{y}_i}$$

Once we compute these $r_i$'s, we train a weak learner $h_m$ to predict $r_i$ from $x_i$. This learner is intentionally kept simple to prevent overfitting.

The Intuition: Why This Works

Why do we fit the learner to the negative gradient? Let's use the squared loss function as a simple example to build intuition.

Let the loss function be:

$$L(y, \hat{y}) = \frac{1}{2} (y - \hat{y})^2$$

First, we find the gradient of $L$ with respect to the prediction $\hat{y} = F_{m-1}(x)$:

$$\begin{align*} \nabla_{F_{m-1}} L \big( y, F_{m-1}(x) \big) &= \frac{\partial L(y, \hat{y})}{\partial \hat{y}} \\[4pt] &= - (y - \hat{y}) \end{align*}$$

Now, $r$ (our new target for $h_m$) is the negative of this gradient:

$$\begin{align*} r &= - \big( - \left( y - \hat{y} \right) \big) \\[4pt] &= y - \hat{y} \end{align*}$$

This is exactly the residual—the actual error of the previous model $F_{m-1}$ on the data point $x$

Now, remember we train $h_m$ to approximate $r$, so $h_m \approx y - \hat{y}$. Let's plug this into our update equation, assuming a learning rate $\gamma = 1$ for simplicity:

$$\begin{align*} F_m(x) &= F_{m-1}(x) + \gamma h_m(x) \\ &= F_{m-1}(x) + h_m(x) \\ &\approx F_{m-1}(x) + \bigl(y - \hat{y}\bigr) \\ &\approx \hat{y} + \bigl(y - \hat{y}\bigr) \\ &\approx y \end{align*}$$

So, in essence, when the weak learner learns the negative gradient, the model learns the true target

Note: In the general case, the $r_i$'s are sometimes referred to as the pseudo-residuals

We haven't discussed how the initial model, $F_0$ is chosen. It is typically a simple, constant model. For example, in regression problems, $F_0$ is often just the average of all the true target values ($y$) in the training set

Until now, we have focused on the training process. For inference, it's important to remember that the entire sequence of models is involved. The final model's output, $F_m(x)$, depends on $F_{m-1}(x)$, which in turn depends on $F_{m-2}(x)$, and so on, all the way back to $F_0(x)$. The final prediction is therefore an additive aggregation of all the learners in the ensemble.

In a gradient boosted tree, the weak learner $h_m$ is typically a shallow regression tree.

To further prevent overfitting, GBTs often employ feature subsampling. This means that when a tree is being built, it considers only a random subset of features when finding the best split for a node. This restriction is applied per-split, not per-tree. A single tree can still use a variety of features, but it is forced to make decisions without considering all features at once.

Finally, since each weak learner in a gradient boosted tree is itself a tree, GBTs also suffer from correlation bias

Feature Interpretability in Gradient Boosted Trees

Tackling correlation bias in tree-based models requires a different approach than in linear models. As we saw, LASSO regression is a potential solution for linear models because it removes certain features. We need a similar solution for trees, but it isn't as straightforward due to their non-linear nature.

A starting goal, similar to LASSO, is to select a subset of learners (trees) so that the final ensemble doesn't use all the features. One post-processing approach involves growing a forest of n trees and then solving the following objective:

$$\begin{align*} &\text{minimize} && \frac{1}{m}L(y,Aw) + \lambda\sum_{i=1}^n u_i w_i, \\ &\text{subject to} && w \ge 0. \end{align*}$$

Here:

Let's analyze this objective. The first term,

$$\frac{1}{m}L(y,Aw)$$

is the loss term across the entire dataset. This term ensures we don't compromise on the accuracy of the model

The below diagram might make things clear:

The second term,

$$\lambda\sum_{i=1}^n u_i w_i$$

is the regularization term. It penalizes trees that use many features.

The constraint,

$$w \geq 0$$

ensures the final model is purely additive (i.e., we can only add, not subtract, the predictions of a weak learner). This non-negative regularization acts similarly to LASSO, encouraging sparsity and helping to remove redundancy by driving the weights of unhelpful or overly complex trees to zero.

This post-processing technique helps to prune trees that have used many features, ensuring the final ensemble uses only a subset of features. However, there's a significant catch: what if every tree in the initial forest has already used all (or most) of the features? In that scenario, this technique doesn't work

We need a better technique that ensures we grow a diverse forest from the beginning—one where the correlation between trees is low—to ensure the final ensemble doesn't rely on all the features.

In the next section, we’ll discuss how to grow a diverse forest using incremental depth boosting.

Incremental Depth Boosting

As discussed in the previous section, the goal is to grow a diverse forest such that the final ensemble doesn't use all the features.

Below is the algorithm for incremental depth boosting.

Here's the summary of the algorithm:

  1. Initialize the boosted model $F$ and set the tree depth to $d=1$.
  2. Start the boosting process by adding trees of depth $d$ until the training error converges.
  3. Increment $d \leftarrow d+1$ and continue the boosting process by adding trees of the new depth. Repeat this until the out-of-bag (OOB) error converges, then return the final forest $F$.

Note that the out-of-bag (OOB) error is computed by evaluating each training point only on the trees that did not see that point during bootstrap sampling.

In standard boosting, we usually fix the depth of each tree and then train the ensemble. In incremental depth boosting, however, we are much more careful about both adding new trees and increasing their depth. We only add complexity when doing so helps reduce the error. This ensures the final forest contains the minimum necessary number of trees at the minimum necessary depth, helping to achieve our goal of a diverse forest.

Note that in standard boosting, all trees usually have the same fixed depth. In incremental depth boosting, however, the final ensemble can contain trees of different depths.

The ControlBurn paper also proposes a related algorithm for bagging, called incremental depth bagging. In this article, however, we focus only on the boosting setting. I encourage interested readers to consult the original paper for details on incremental depth bagging.

The Idea Behind ControlBurn

The idea behind ControlBurn is simple. First, we grow a diverse forest using incremental depth boosting or incremental depth bagging. Then, we make the ensemble feature-sparse by using the technique discussed in the section Feature Interpretability in Gradient Boosted Trees.

ControlBurn also helps prevent overfitting. It controls model complexity by limiting tree depth during training, and it further simplifies the ensemble by assigning sparse weights to redundant trees.

Finally, ControlBurn is especially useful for datasets with highly correlated features.

Conclusion

I hope you enjoyed this article. I encourage interested readers to explore the original ControlBurn paper for additional insights.