End-to-end Feature Selection Approach for Learning Skinny Trees

We propose a new optimization-based approach for feature selection in tree ensembles, an important problem in statistics and machine learning. Popular tree ensemble toolkits e.g., Gradient Boosted Trees and Random Forests support feature selection post-training based on feature importance scores, while very popular, they are known to havedrawbacks. We propose Skinny Trees: an end-to-end toolkit for feature selection in tree ensembles where we train a tree ensemble while controlling the number of selected features. Our optimization-based approach learns an ensemble of differentiable trees, andsimultaneously performs feature selection using a grouped $\ell_0$-regularizer. We use first-order methods for optimization and presentconvergence guarantees for our approach. We use a dense-to-sparse regularization scheduling scheme that can lead to more expressiveand sparser tree ensembles. On 15 synthetic and real-world datasets, Skinny Trees can achieve $1.5{\times}$--$620{\times}$ feature compression rates, leading up to $10{\times}$ faster inference over dense trees, without any loss in performance. Skinny Trees lead to superior feature selection than many existing toolkits e.g., in terms of AUC performance for 25\% feature budget, Skinny Trees outperforms LightGBM by 10.2\% (up to 37.7\%), and Random Forests by 3\% (up to 12.5\%).