[Paper Review] Unbiased Measurement of Feature Importance in Tree-Based Methods

Zhengze Zhou, Giles Hooker|arXiv (Cornell University)|Mar 12, 2019

Statistical Methods and InferenceMathematics42 references17 citations

TL;DR

This paper proposes an unbiased feature importance measure for tree-based models by using out-of-sample split-improvement scores to correct the inherent bias in traditional split-improvement (Gini importance) that favors features with more potential splits. The method reduces spurious importance for continuous or high-cardinality categorical features and ensures unbiased estimates under the null, with minimal computational overhead and ready integration into existing libraries like scikit-learn and randomForest.

ABSTRACT

We propose a modification that corrects for split-improvement variable importance measures in Random Forests and other tree-based methods. These methods have been shown to be biased towards increasing the importance of features with more potential splits. We show that by appropriately incorporating split-improvement as measured on out of sample data, this bias can be corrected yielding better summaries and screening tools.

Motivation & Objective

• To address the well-documented bias in split-improvement feature importance measures that favor features with more potential splits, especially continuous or high-cardinality categorical variables.
• To develop a simple, computationally efficient fix that uses out-of-sample data to estimate split-improvement, avoiding overfitting to training data.
• To provide a theoretically grounded, unbiased alternative to the default feature importance in popular libraries like scikit-learn and randomForest, which currently use biased in-sample measures.
• To validate the method empirically across simulated and real-world datasets, showing improved reliability in feature screening and ranking.
• To enable more trustworthy model interpretation and downstream analysis by ensuring features with no predictive power receive zero importance in expectation.

Proposed method

• The method replaces in-sample split-improvement with out-of-sample split-improvement, computed on out-of-bag samples or held-out data, to avoid overfitting.
• For each tree in the forest, the improvement in impurity reduction at each split is measured on out-of-sample data, rather than the training set used to build the tree.
• The final feature importance is the average of these out-of-sample split-improvement values across all trees and splits, ensuring unbiased estimation.
• The approach is compatible with existing tree-building frameworks and requires minimal code changes, as it only modifies how split-improvement is aggregated.
• The method can be extended to honest trees or any setting with a natural test set, such as cross-validation or held-out partitions.
• Theoretical justification is provided, showing that under the null hypothesis (no predictive power), the expected importance is zero.

Experimental results

Research questions

• RQ1 Does the standard split-improvement measure in Random Forests exhibit systematic bias toward features with more potential splits?
• RQ2 Can out-of-sample split-improvement estimation correct this bias and yield more reliable feature importance rankings?
• RQ3 How does the proposed unbiased measure compare to existing methods like permutation importance and cforest in real-world and simulated data?
• RQ4 Is the proposed method robust to varying feature types, such as binary, categorical, and continuous variables?
• RQ5 Can the unbiased measure support formal statistical inference, such as hypothesis testing or confidence intervals, in tree-based models?

Key findings

• The standard split-improvement measure in Random Forests is significantly biased toward features with more potential splits, such as continuous or high-cardinality categorical variables.
• In simulations with 10 noisy features, the proposed unbiased feature importance (UFI) correctly identified the single true predictor, while standard split-improvement (SI) ranked a noisy feature higher.
• On the Adult and Boston Housing datasets, UFI produced more plausible and consistent feature rankings compared to SI, with the random feature receiving near-zero importance.
• UFI's importance scores were consistent with those from two state-of-the-art methods, cforest and ranger, suggesting strong empirical validity.
• The method requires no additional computational cost beyond standard training, as out-of-bag samples are already available in Random Forests.
• Theoretical analysis confirms that features with no predictive power receive an expected importance of zero under the proposed method, satisfying unbiasedness.