ControlBurn: Feature Selection by Sparse Forests

Correlation bias reduces the interpretability of feature rankings. When features importance scores are split evenly amongst groups of correlated features, influential features are harder to detect: feature ranks can become cluttered with many features that represent similar relationships. ControlBurn allows a data scientist to quickly identify a sparse subset of influential features for further analysis.

In many experimental design problems, features are costly to observe. For example, in medical settings, certain tests are expensive and risky to the patient. In these situations it is crucial to avoid redundant tests. In this setting, ControlBurn can be used to identify a subset of $k$ features (for any integer $k>0$) with (nearly the) best model performance, eliminating redundancy.

form the base learners used in tree ensembles. We give a quick overview of classification trees below. A classification tree of depth $d$ separates a dataset into at most $2^{d}$ distinct partitions. Within each partition, all data points are predicted as the majority class. Classification trees are grown greedily, and splits are selected to minimize the Gini impurity, misclassification error, or deviance in the directly resulting partitions. The textbook (friedman2001elements, , Chapter 9) provides a comprehensive reference on how to construct decision trees.

reduces the variance of tree ensembles through averaging the predictions of many trees. We sample the training data uniformly with replacement to create bags (bootstrapped datasets with the same size as the original) and fit decision trees on each bag. The final prediction is the average (or majority vote, for classification) of the predictions from all the trees. Reducing correlation between trees improves the effectiveness of bagging. To do so, we can restrict the number of features considered each time a tree is built or a split is selected (friedman2001elements, , Chapter 15).

fits an additive sequence of trees to reduce the bias of a tree ensemble. Each decision tree is fit on the residuals of its predecessor; weak learners are sequentially combined to improve performance. Popular boosting algorithms include AdaBoost (freund1999short, ), gradient boosting (friedman2001greedy, ), and XGBoost (chen2016xgboost, ). The textbook (friedman2001elements, , Chapter 10) provides a comprehensive reference on boosting algorithms.

can be computed from decision trees by calculating the decrease in node impurity for each tree split. Gini index and entropy are two commonly used metrics to measure this impurity. Averaging this metric over all the splits for each feature yields that feature’s importance score. This score is known as Mean Decrease Impurity (MDI); features with larger MDI scores are considered more important. MDI scores can be calculated for tree ensembles by averaging over all the trees in the ensemble (2001.04295, ).

feature selection algorithms remove features suspected to be irrelevant based on metrics such as correlation or chi-square similarity (ambusaidi2016building, ; lee2011filter, ). Features that are correlated with the response are preserved while uncorrelated features are removed. Groups of intercorrelated features correlated with the response are entirely preserved.

Filter-based algorithms are model-agnostic, their performance depends on the relevancy of the similarity metric used. For example, using Pearson’s correlation to select features from highly non-linear data will yield poor results.

feature selection algorithms iteratively select features through retraining the model across different feature subsets. (mafarja2018whale, ; maldonado2009wrapper, ).

For example, in recursive feature elimination (RFE) the model is initially trained on the full feature set. The features are ranked by feature importance scores and the least important feature is pruned. The model is retrained on the remaining features and the procedure is repeated until the desired number of features is reached. Other wrapper-based feature selection algorithms include sequential and stepwise feature selection. All of these algorithms are computationally expensive since the model must be retrained several times.

feature selection algorithms select important features as the model is trained. (langley1994selection, ; blum1997selection, ). For example in LASSO regression, the regularization parameter $\lambda$ controls the number of coefficients forced to zero, i.e. the number of features excluded from the model. For tree ensembles, MDI feature importance scores are computed during the training process. These scores can be used to select important features, however this algorithm is highly susceptible to correlation bias since MDI feature importance is distributed evenly amongst correlated features. ControlBurn extends LASSO to the nonlinear setting and is an embedded feature selection algorithm robust to correlation bias.

(Algorithm 1) increases forest diversity by incrementally increasing the depth $d$ of trees grown. A tree of depth $d$ can use no more than $d^{2}-1$ features. As we increase $d$, the number of features used per tree increases, allowing the model to represent more complex feature interactions.

The algorithm has a single hyperparameter, $d^{\text{max}}$, the maximum depth of trees to grow. This hyperparameter is easy to choose: bagging prevents overfitting, so bigger is better, and $d^{\text{max}}=p$ is allowed.

We represent a tree ensemble as a set of trees $F=\{t_{1}\ldots t_{n}$}. With some abuse of notation, we let $F(x)$ be the output of the tree ensemble given input $x\in\mathbb{R}^{p}$.

Incremental depth bagging increments the depth of trees grown when the performance of the model saturates at the lower depths. At each depth level, the training loss is recorded each time a new tree is added to the ensemble. When the training loss converges, the algorithm increments the depth $d$ and begins adding deeper trees. We say the loss converges if the tail $N$ elements of the sequence fall within an $\epsilon$ tube. Experimentally, choosing $N=5$ and $\epsilon=10^{-3}$ yields good results. This methodology differs from typical early stopping algorithms used in machine learning, since it stops based on training loss rather than test loss. We use training loss because it is fast to evaluate and overfitting is not a concern, since bagged tree ensembles tend not to overfit with respect to the number of trees used.

(Algorithm 2) uses out-of-bag (OOB) error to automatically determine the maximum depth of trees to grow, and requires no hyperparameters. The algorithm uses bagging ensembles as the base learner in gradient boosting to increase diversity without sacrificing accuracy.

The number of trees grown per boosting iteration is determined by how fast the training error of the ensemble converges, as in §3.2.1. Each boosting iteration increments the depth of trees grown and updates the pseudo-residuals. This is done by setting $e_{i}=-\frac{\partial L(y_{i},z)}{\partial z}|_{z=F(x_{i})}$ for $i\in\{1\ldots m\}$.

For each row in $x\in X$, the OOB prediction is obtained by only using trees in $F^{{}^{\prime}}$ fit on bags $\{X^{\prime}\mid x\not\in X^{\prime}\}$. Let $\delta$ be the improvement in OOB error between subsequent boosting iterations. The algorithm terminates when $\delta<0$.

In Figure 6, incremental depth bag-boosting is conducted on the abalone dataset from UCI MLR (Dua:2019, ). The scatterplot shows how the improvement in OOB error, $\delta$, changes as the number of boosting iterations increase. The procedure will terminate at the sixth boosting iteration, when $\delta<0$. The lineplot shows how out-of-sample (OOS) test error varies with the number of boosting iterations. OOS error is minimized at the seventh boosting iteration. This is remarkably close to the OOB error stop, which is determined on-the-fly.

The following lemma shows that the average performance improvement calculated using OOB samples is a good estimator of the true performance improvement for the data distribution. The lemma applies to any pair of ensembles $F$ and $\hat{F}$ where every element in ensemble $F$ appears in ensemble $\hat{F}$.

We found empirically that incremental depth bag-boosting performs better than incremental depth bagging when used in ControlBurn. In addition, the OOB early stopping rule eliminates all hyperparameters and limits the size of the forest. This reduces computation time in the following optimization step. Forests grown by incremental depth bag-boosting are more feature diverse than those grown via incremental depth bagging. Consider the case where a single highly influential binary feature dominates a dataset. While each iteration of incremental depth bagging will include this feature, the feature will only appear in the first boosting iteration in incremental depth bag-boosting. Boosting algorithms fit on the errors of previous iterations, so once the relationship between the feature and the outcome is learned, the feature may be ignored in subsequent iterations.

include feature-sparse decision trees and explainable boosting machines (EBM). These algorithms do not perform as well as incremental depth bag-boosting, but may be more interpretable

The procedure to build feature-sparse decision trees is a simple extension of the tradition tree building algorithm. For classification, decision trees typically choose splits to minimize the Gini impurity of the resulting nodes. Sparse decision trees add an additional cost $c$ to this impurity function when the tree splits on a feature that has not been used before. As a result, when groups of correlated features are present, the algorithm consistently splits on a single feature from each group. The cost parameter $c$ controls the sparsity of each tree. Ensembles of feature-sparse trees are not necessarily sparse, so ControlBurn must be used to obtain a final sparse model.

Explainable boosting machines are generalized additive models composed of boosted trees trained on a single feature (lou2013accurate, ). Each tree in an EBM is feature-sparse with sparsity $1$. The additive nature of EBMs improve interpretability as the contribution of each feature in the model can be easily obtained. However, this comes at a cost of predictive performance since EBMs cannot capture higher order interactions. When ControlBurn is used on EBMs, the matrix $G\in\mathbb{R}^{p\times p}$ in the optimization step is the identity matrix.

In many applications, it may be necessary to restrict feature selection to a set of features rather than individual ones. For example, in medical diagnostics, a single blood test provides multiple features for a fixed cost. We can easily incorporate this feature grouping requirement into our framework by redefining $u_{i}$. Suppose features $1,\dots,p$ are partitioned into various groups. Instead of defining $u_{i}$ to be the number of features tree $i$ uses, we now define

where $c_{k}$ is the cost of the group of features $k$ and $T_{k}$ is the set of feature groups tree $i$ uses. Our original formulation corresponds to the case where all groups are singleton.

In some problems the cost of obtaining one feature may be very different from the cost of obtaining another. We can easily accommodate this non-homogeneous feature costs requirement into our framework as well. Instead of defining $u_{i}$ to be the number of features used by tree $t_{i}$, we define

where $c_{j}$ is the cost of feature $j$ and $T_{i}$ is the set of features tree $i$ uses. Our original formulation corresponds to the case where all $c_{j}=1$.

The optimization step of ControlBurn scales linearly with the number of rows in our data. We can sketch $X$ to improve runtime. Let $S\in\mathbb{R}^{s\times m}$ be the sampled (e.g., Gaussian) sketching matrix. Our optimization problem becomes:

A relatively small $s$, around $O(\text{log}(m))$, can provide a good solution (pilanci2015randomized, ).

To evaluate how well ControlBurn handles correlation bias, we synthetically add correlated features to existing data by duplicating columns and adding Gaussian white noise to each replicated column. This procedure experimentally controls the degree of feature correlation in the data.

We consider a wide range of benchmark datasets selected from the Penn Machine Learning Benchmark Suite (olson2017pmlb, ). This curated suite aggregates datasets from well-known sources such as the UCI Machine Learning Repository (Dua:2019, ). We select all datasets in this suite with a binary target, 43 datasets in total.

Finally, we select several real-world datasets known to contain correlated features, based on machine learning threads on Stack Overflow and Reddit. We select the Audit dataset, where risk scores are co-linear, and the Adult Income dataset, where age, gender, and marital status form a group of correlated features (Dua:2019, ). We also evaluate ControlBurn on the DNA microarray data used in West (2001) (West:2001ft, ), since gene expression levels are highly correlated.

The regularization parameter $\lambda$ in Problem 2 controls the number of features selected by ControlBurn, $k$. We use bisection to find the value of $\lambda$ that yields the desired number of selected features $k$. Our goal in feature selection is to produce an interpretable model, so we mandate that $k\leq k^{\text{max}}:=10$ in our experiments. Tree ensembles with more features can be difficult to interpret.

In fact, we can use bisection with ControlBurn to generate feature sets of any sparsity between $0$ and $k^{\text{max}}$. Start with a value of $\lambda$ large enough to guarantee $k=0$. Let $k^{{}^{\prime}}$ denote the number of features we want ControlBurn to select in the following iteration and initialize $k^{{}^{\prime}}=1$. Bisect $\lambda=\frac{\lambda}{2}$ and solve Problem 2 to obtain $k$. If $k<k^{{}^{\prime}}$, update $\lambda=\frac{\lambda}{2}$; otherwise if $k>k^{{}^{\prime}}$, update $\lambda=\lambda+\frac{\lambda}{2}$. If $k=k^{{}^{\prime}}$, update $k^{{}^{\prime}}=k^{{}^{\prime}}+1$. Repeat until $k^{{}^{\prime}}>k^{\text{max}}$.

We use ControlBurn to build a classification model to predict whether an adult’s annual income is above $50k. The Adult Income plot in figure 10 shows how well ControlBurn performs compared to our random forest baseline. We see that ControlBurn performs substantially better than our baseline when $k=3$ features are selected. ControlBurn selects “EducationNum”, “CapitalGain”, and “MaritalStatus_Married-civ-spouse” while our random forest baseline selects “fnlwgt”, “Age” and “CapitalGain”.

Figure 11 shows a correlation heat map of the top 20 features in the Adult Income dataset ranked by MDI importance. The top 3 features selected by ControlBurn are in bold. From this plot, it is apparent that the MDI feature importance of “MaritalStatus_Married-civ-spouse” is diluted because the feature is correlated with many other influential features. ControlBurn is able to select this feature while our random forest baseline ignores it.

It is also interesting to note that the feature “fnlwgt” is uninformative: it represents the population size of a row and is unrelated to income. This feature is given high importance by our random forest baseline but ignored by ControlBurn. We discuss this phenomenon further in Section 6.1.

The Audit dataset demonstrates that highly correlated and even duplicated features sometimes appear in real-world data. Many of these risk scores are highly correlated and several scores are linear transformations of existing scores. We use ControlBurn to build a classification model to predict whether a firm is fradulent based on risk scores. Figure 10 shows that ControlBurn outperforms our baseline for any sparsity level $k$.

In this experiment, we use ControlBurn to determine which gene expressions are the most influential towards predicting the hormone receptor status of breast cancer tumors. Our data contains 49 tumor samples and over 7000 distinct gene expressions. We formulate this task as a binary classification problem to predict receptor status using gene expressions as features. The Gene Expression plot in Figure 10 shows the difference in accuracy between a model selected using ControlBurn and our random forest baseline. ControlBurn outperforms our baseline for every sparsity $k$. These preliminary findings suggest that ControlBurn is effective at finding a small number of genes with important effects on tumor growth. We discuss implications and extensions appropriate for gene expression data further in Section 6.2.