MABSplit: Faster Forest Training Using Multi-Armed Bandits
Abstract
Random forests are some of the most widely used machine learning models today, especially in domains that necessitate interpretability.
We present an algorithm that accelerates the training of random forests and other popular tree-based learning methods.
At the core of our algorithm is a novel node-splitting subroutine, dubbed MABSplit, used to efficiently find split points when constructing decision trees.
Our algorithm borrows techniques from the multi-armed bandit literature to judiciously determine how to allocate samples and computational power across candidate split points.
We provide theoretical guarantees that MABSplit improves the sample complexity of each node split from linear to logarithmic in the number of data points.
In some settings, MABSplit leads to 100x faster training (an 99% reduction in training time) without any decrease in generalization performance.
We demonstrate similar speedups when MABSplit is used across a variety of forest-based variants, such as Extremely Random Forests and Random Patches.
We also show our algorithm can be used in both classification and regression tasks.
Finally, we show that MABSplit outperforms existing methods in generalization performance and feature importance calculations under a fixed computational budget.
All of our experimental results are reproducible via a one-line script at https://github.com/ThrunGroup/FastForest .
Introduction
Random Forest (RF) is a supervised learning technique that is widely used for classification and regression tasks [28, 15].
In RF, an ensemble of decision trees (DTs) is trained for the same prediction task.
Each DT consists of a series of nodes that represent if/then/else comparisons on the feature values of a given datapoint and are used to produce an output label. In RF, each DT is typically trained on a bootstrap sample of the original dataset and considers a random sample of available features at each node split [14].
The prediction of the each DT is aggregated to provide an output label for the whole RF.
By aggregating the prediction of each DT in the ensemble, RFs tend to be more robust to noise and overfitting [52] and are capable of capturing more complex patterns in the data than a single DT [59].
RF has gained tremendous popularity due to its flexibility, usefulness in multi-class classification and regression tasks, high performance across a broad range of data types, natural support for missing features, and relatively low computational complexity [61, 53, 10, 35, 30].
Furthermore, RF inherits the interpretability of decision trees because the prediction of each constituent DT can be explained through a sequence of binary decisions.
RFs have been successfully applied in contexts as varied as the prediction of legal court decisions [32], solar radiation analysis [10], and the Higgs Boson classification problem [3].
In the era of big data, a simple and flexible machine learning technique such as RF is expected to play a key role in processing large datasets and providing accurate and interpretable predictions.
The need for training prediction models on massive datasets and doing so on compute-constrained hardware, such as smartphones and Internet-of-Things devices, requires the development of new algorithms that can deliver faster results without sacrificing generalization performance [61].
For this reason, recent work has proposed ways to accelerate the training of RFs, both at the algorithmic level and at the hardware level.
At the algorithmic level, most work focuses on fast construction of each individual DT.
Each DT is built by identifying the feature f and threshold t that best split the data points according to the prediction targets.
The data points are split into subsets based on whether their feature f has a value less than t or feature f has a value greater than t . The process is then recursed for each resulting subset.
Most of the complexity in this process is in identifying the pair ( f, t ) that provides the best split for a set of N data points, which typically requires O ( N ) computation per split.
Recent proposals include computing (or estimating) f and t from a subsample of the data points and features, or quantizing the feature values.
The latter technique creates a histogram of values of each feature across the data points and restricts t to be at the edges of histogram bins.
While existing approaches provide significant speed up in the training of RFs, they often require prespecification of fixed hyperparameters, such as the proportion of data points or features to subsample, and are not adaptive to the underlying data distribution.
Moreover, when comparing different candidate features for a split, all features are treated on equal footing and the quality of their split is computed based on the same number of data points.
Intuitively, this is wasteful because features that are not informative for the prediction task can be identified based on a smaller number of data points.
Alternatively, an adaptive scheme could better allocate computational resources towards a promising set of candidate features and achieve a better tradeoff between computational cost and generalization performance.
In this work, we propose MABSplit, a fast subroutine for the node-splitting problem, which adaptively refines the estimate of the “quality” of each feature-threshold pair ( f, t ) as a candidate split.
Bad split candidates can be discarded early, which can lead to significant computational savings.
The core idea behind our algorithm is to formulate the node-splitting task as a multi-armed bandit problem [34, 2, 5, 58], where each pair ( f, t ) is a distinct arm.
The unknown parameter of each arm, µ ft , corresponds to the quality of the split based on feature f and threshold t , where the split quality is measured in terms of how much the split would reduce label impurity.
An arm ( f, t ) can be “pulled” by computing the reduction of label impurity induced by a new data point sampled from the dataset.
This allows us to compute an estimate ˆ µ ft and a corresponding confidence interval, which can be used in a batched variant of the Upper Confidence Bound (UCB) and successive elimination algorithms [34, 63] to identify the best arm ( f, t ) .
Crucially, MABSplit uses the adaptive sampling tools of multi-armed bandits to avoid computing the split qualities over the entire dataset.
We demonstrate the benefits of MABSplit on a variety of datasets, for both classification and regression tasks.
In some settings, MABSplit algorithm leads to 100x faster training (a 99% reduction in training time), without any decrease in test accuracy, over an exact implementation of RF that searches for the optimal ( f, t ) pair via brute-force computation.
Additionally, we demonstrate similar speedups when using MABSplit across a variety of forest-based variants, such as Extremely Random
Forests and Random Patches.
1.1 Related work
Random Forests were originally proposed by Ho [28].
In its original formulation, RF constructs n tree DTs, where each DT is trained on a bootstrap sample of all N data points and a random subset of the features at each node (a technique known as random subspacing [29]).
More recently, the need for training RFs on large datasets has prompted the development of several techniques to accelerate training at both the software and hardware levels.
Software acceleration of RF: Most of the software and algorithmic acceleration techniques focus on the training of each individual DT.
FastForest [61] accelerates the node-splitting task using three ideas: subsampling a pre-specified number of data points without replacement (subbagging), subsampling a pre-specified number of features dependent on the current number of data points (Dynamic Restricted Subspacing), and dividing values of a given feature into T bins, where T depends on the number of data points at the node (Logarithmic Split-Point Sampling, inspired by the single-tree SPAARC algorithm [62]).
MABSplit is inspired by the ideas in FastForest, but does not require the number of data points or features to be prespecified and, instead, determines them by adapting to the data distribution.
Other recent work has also used adaptivity to identify the best split.
For example, Very Fast Decision
Trees (VFDTs) [19] and Extremely Fast Decision Trees (EFDTs) [40] are incremental decision tree learning algorithms in which trees can be updated in streaming settings.
Acceleration of the node-splitting task is achieved by adaptively selecting a subset of data points sufficient to distinguish the best and second best splits.
These approaches are similar to ours, but the sampled data points are used to evaluate all possible splits. MABSplit, in contrast, adaptively discards unpromising splits early.
The F-forest [23] algorithm also applies adaptivity and uses an upper bound on the impurity reduction of each split in order to discard candidate splits.
This is similar in spirit to the goal of MABSplit, but is based on a deterministic, conservative upper bound on the impurity reduction (as opposed to MABSplit’s statistical estimate) and incurs computation linear in the node size for each split, even when considering a fixed number of possible split thresholds per feature.
Other variations of RF have been proposed to improve training time.
Random Patches [39] builds trees based on a subset of data points and features that is fixed for each entire forest.
ExtraTrees (ETs) [24] draw a random subset of K features at each node and, for each one, chooses a number R of random splits. It then selects the split that yields the largest impurity reduction from among these KR candidate splits.
Other recent work attempts to accelerate RF training by identifying the optimal number of decision trees needed in the forest [45], a form of hyperparameter tuning.
The MABSplit subroutine can also be incorporated into these methods.
Hardware acceleration of RF: The training of RFs can also be significantly accelerated through the use of specialized hardware.
For instance, the implementation of RF available in Weka [21] allows trees to be trained on different cores and reduces forest training time.
A GPU-based parallel implementation of RF has also been proposed in [26].
These solutions require specialized hardware (e.g., GPU-based PC video cards) and are inappropriate for everyday users locally executing datamining tasks on standard PCs or smartphones.
As such, it is still desirable to develop techniques to improve prediction performance and processing speed at an algorithmic, platform-independent level.
Preliminaries
We now formally describe the RF algorithm and other tree-based models, all of which rely on a nodesplitting subroutine.
We consider N data points { ( x i , y i ) } Ni =1 where each x i is an M -dimensional feature vector and y i is its target.
Following standard literature, we consider flexible feature types such as numerical or categorical.
We consider both categorical targets for classification and numerical targets for regression.
With a slight abuse of notation, we use X to mean either the set of indices { i } or the values { ( x i , y i ) } , with the meaning clear from context.
An RF contains n tree decision trees, each trained on a set of N bootstrapped data points (sampled with replacement) and a random subset of features at each node. The whole RF, an ensemble model, outputs the trees’ majority vote in classification and the trees’ average prediction in regression [28, 15].
We focus on the top-down approach of constructing DTs by choosing the feature-threshold pair at each step that best splits the set of targets at a given node [52]. We now discuss the method by which the best split is chosen in each node.
Consider a node with n data points X and m features F .
Note that n and m at the given node may or may not be the same as N and M of the entire dataset (for example, RF considers a random subset of m = √ M features at each node [12]).
Let X L ,ft and X R ,ft be the left and right child subsets of X when X is split according to feature f at threshold t .
The approach finds the split that best reduces the label impurity; i.e., finds
[Equation]
where I ( S ) measures the impurity of targets { y i } i ∈S and T f is the set of allowed thresholds for feature f .
Common impurity measures include Gini impurity or entropy for classification, and mean-squared-error (MSE) for regression [16]:
[Equation]
where K is the total number of classes and p k = 1 n (cid:80) i ∈X I ( y i = k ) is the proportion of class k in X in classification, and ¯ y = 1 n (cid:80) i ∈X y i is the average target value in regression.
We note that our proposed algorithm, MABSplit, does not assume any particular structure of I ( · ) .
While the conventional RF considers all ( n − 1) possible splits among the n generally different values in the dataset for a given feature f , in this work we focus on the histogram-based variant that chooses the threshold from a set of predefined values T f , e.g., |T f | equally-spaced histogram bin edges; this variant is substantially more efficient, offers comparable accuracy, and has been used in most state-of-the-art implementations such as XGBoost and LightGBM [51, 17, 33, 62].
A naïve algorithm finds the best feature-threshold pairs in Equationh (1) by evaluating the label impurity reduction for each feature-threshold over all n data points, which incurs computation linear in n .
MABSplit
We now discuss MABSplit and how it can reduce the complexity of the node-splitting problem to logarithmic in n .
With the same notation as in Section 2, we note that Equation (1) is equivalent to
[Equation]
and so we focus on solving Equation (3).
Let µ ft = |X L ,ft | n
I ( X L ,ft ) + |X R ,ft | n
I ( X R ,ft ) be the optimization objective for feature-threshold pair ( f, t ) .
Omitting the dependence on K , computing µ ft exactly is at least O ( n ) . MABSplit, however, estimates µ ft with less computation by drawing n ′ < n independent samples with replacement from X .
As shown in Subsec.
3.1, it is possible to construct a point estimate ˆ µ ft ( n ′ ) and a (1 − δ ) confidence interval (CI) C ft ( n ′ , δ ) for the parameter µ ft , where n ′ and δ determine estimation accuracy.
The width of these CIs generally scales with (cid:113) log1 /δ n ′ .
To estimate the solution to Problem (3) with high confidence, we can then choose to sample different amounts of data points for different features so as to estimate their impurity reductions to varying degrees of accuracy.
Intuitively, promising features-threshold splits with high impurity reductions (lower values of µ ft ) should be estimated with high accuracy with many data points, while less promising ones with low impurity reductions (higher values of µ ft ) can be discarded early.
The exact adaptive estimation procedure, MABSplit, is described in Algorithm 1.
It can be viewed as a batched version of the conventional UCB algorithm [34, 63] combined with successive elimination [22], is straightforward to implement, and has been used in other applications [58, 55, 56, 6, 7].
Algorithm 1 uses the set S solution to track all potential solutions to Problem (3); S solution is initialized as the set of all feature-threshold pairs { ( f, t ) } and Algorithm 1 maintains the mean objective estimate ˆ µ ft and (1 − δ
CI C ft for each potential solution ( f, t ) ∈ S solution .
In each iteration, a new batch of data points X batch is used to evaluate the split quality for all potential feature-threshold splits in S solution , which allows the estimate of each ˆ µ ft to be made more accurate.
Based on the current estimate, if a candidate’s lower confidence bound ˆ µ ft − C ft is greater than the upper confidence bound of the most promising candidate min f,t (ˆ µ ft + C ft ) , we remove it from
S solution .
This process continues until there is only one candidate in S solution or until we have sampled more than n data points.
In the latter case, we know that the difference between the remaining candidates in S solution is so subtle that an exact computation is warranted.
MABSplit then compute those candidates’ objectives exactly and returns the best candidate in the set.
3.1 Point estimates and confidence intervals for impurity metrics
We now discuss MABSplit’s construction of point estimates and confidence intervals of µ ft based on a set of n ′ points, { ( X i , Y i ) } n ′ i =1 , sampled independently and with replacement from X . We consider two widely used impurity metrics in classification, Gini impurity and entropy, although mean estimates and confidence intervals for other settings and metrics can be derived similarly (more details are provided in Appendix 3).
Let p L ,k := 1 n (cid:80) ni
=
I ( x if < t, y i = k ) and p R ,k := 1 n (cid:80) ni
{+=
+}
I ( x if ≥ t, y i = k ) denote the proportion of the full n data points in class k and each of the two subsets created by the split ( f, t ) (we call these subsets “left” and “right”, respectively).
Note that
[Equation]
Furthermore, let ˆ p L ,k := 1 n ′ (cid:80) n ′ i
=1 I ( X if < t, Y i = k ) and ˆ p R ,k := 1 n ′ (cid:80) n ′ i =
I
( X if ≥ t, Y i =
k ) denote the empirical estimates of p L ,k and p R ,
k based on the n ′ subsamples drawn thus far.
Then { ˆ p L ,k , ˆ p R ,k } Kk =1 jointly follow a multinomial distribution satisfying
[Equation]
since for each random data point ( X i , Y i ) , exactly one element of the set { ˆ p L ,k , ˆ p R ,k } Kk =1 is incremented.
Using Equations (2) and (4), and the definition of µ ft after Equation (3), we write
[Equation]
where we can use the empirical parameters { ˆ p L ,k , ˆ p R ,k } Kk =1 as plug-in estimators for the true parameters { p L ,k , p R ,k } Kk =1 to produce the point estimate ˆ µ ft .
In MABSplit (Algorithm 1), each batch of B data points is used to update each ˆ p L ,k and ˆ p R ,k , which are used in turn to update our point estimates ˆ µ ft and the corresponding CIs.
The CIs of ˆ µ ft are based on standard error derived using the delta method [60].
As in standard applications of the delta method, the estimates ˆ µ ft are asymptotically unbiased and their corresponding CIs are asymptotically valid.
Appendix 3 provides further details, including a derivation of the CIs and discussion of convergence properties.
3.2 Algorithmic and implementation details
We considered sampling with replacement in MABSplit (Algorithm 1) primarily for the ease of theoretical analysis.
In practice, we found sampling without replacement was more computationally efficient and did not significantly change the results, and was used in the actual implementation.
For DTs, in classification tasks, we allow individual DTs to provide soft votes and average their predicted class probabilities to determine the forest’s predicted class label, following existing approaches [47].
This is in contrast with hard votes, in which each DT is only permitted to produce its best label and the forest’s prediction is determined by majority voting.
For the fixed-budget experiments in
Section 5, we terminate tree construction and do not split nodes further if doing so would violate our budget constraints.
Analysis of the Algorithm
In this section, we prove that MABSplit returns the optimal feature-threshold pair for a node split with high probability.
Furthermore, we provide bounds on computational complexity of MABSplit, which can lead to a logarithmic dependence on the number of data points n under weak assumptions.
1: S solution ← { ( f, t ) , ∀ f ∈ F , ∀ t ∈ T f } ▷ Set of potential solutions to Problem (3) 2: n used ←▷ Number of data points sampled 3: For all ( f, t ) ∈ S solution , set ˆ µ ft ← ∞ , C ft ← ∞ ▷ Initialize mean and CI for each arm 4: for all f ∈ F do 5:
Create empty histogram, h f , with |T f |
=
T equally spaced bins 6: while n used < n and |S solution | > 1 do 7:
Draw a batch sample X batch of size B with replacement from X 8: for all unique f in S solution do 9: for all x in X batch do 10: Insert x f into histogram h f ▷ Each insertion is O (1) 11: for all ( f, t ) ∈ S solution do 12:
Update ˆ µ ft and C ft based on histogram h f ▷ For fixed f , this is O ( T ) 13:
S solution ← { ( f, t ) : ˆ µ ft − C ft ≤ min f,t (ˆ µ ft + C ft ) } ▷ Retain only promising splits 14: n used ← n used + B 15: if |S solution | = 1 then 16: return ( f ∗ , t ∗ ) ∈ S solution 17: else 18:
Compute µ ft exactly for all ( f, t ) ∈ S solution 19:
( f
) =
arg min ( µ ft
As above, consider a node with n data points X , m features F , and T possible thresholds for each feature ( |T f | =
T for all f ).
Suppose ( f ∗ , t ∗ ) =
arg min f ∈F ,t ∈T f µ ft is the optimal featurethreshold pair at which to split the node.
For any other feature-threshold pair ( f, t ) , define ∆ ft := µ ft − µ f ∗ t ∗ .
To state the following results, we will assume that, for a fixed feature-threshold pair ( f, t ) and n ′ randomly sampled datapoints, the (1 − δ ) confidence interval scales as C ft ( n ′ , δ ) =
O ( (cid:113) log1 /δ
n ′ ) .
(This assumption is justified for the confidence intervals of Gini impurity and entropy under weak assumptions on the µ f,t ’s [60]).
With this assumption, we state the following theorem:
Theorem 1.
Assume ∃ c 0 > 0 s.t. ∀
δ > 0 , n ′ > 0 , we have C ft ( n ′ , δ ) <
c 0 (cid:113) log1 /δ n ′ .
For δ = 1 n 2 mT , with probability at least 1 − 1 n , Algorithm 1 returns the correct solution to Equation (3) .
Furthermore, Algorithm 1 uses a total of M computations, where
[Equation]
Theorem 1 is proven in Appendix 1.
Intuitively, Theorem 1 states that with high probability, MABSplit returns the optimal feature-threshold pair at which to split the node.
The bound Equation (7) suggests the computational cost of a feature-threshold pair ( f, t ) , i.e., min (cid:104) 4 c 20 ∆ 2 ft log( n 2 mT ) + B, 2 n (cid:105) , depends on ∆ ft , which measures how close its optimization parameter µ ft is to µ f ∗ t ∗ .
Most reasonably different features f ̸ = f ∗ will have a large ∆ ft and incur computational cost O (log( n 2 mT )) that is sublinear in n .
In turn, this implies MABSplit takes only O ( mT log( n 2 mT )) computations per feature-threshold pair if there is reasonable heterogeneity among them.
As proven in Appendix 2 of [5], this is the case under a wide range of distributional assumptions on the µ ft ’s, e.g., when the µ ft ’s follow a subGaussian distribution across the pairs ( f, t ) .
Such assumptions ensure that MABSplit has an overall complexity of O ( mT log( n 2 mT )) , which is sublinear in the number of data points n .
We note that in the worst case, however, MABSplit may take O ( n ( m + T )) computations per feature-threshold pair when most splits are equally good, in which case MABSplit reduces to a batched version of the naïve algorithm. This may happen, for example, in highly symmetric datasets where all splits reduce the impurity equally.
Other recent work provides further discussion on the conversion between a bound like Equation (7), which depends on the ∆ i ’s, and a bound in terms of other problem parameters such as O ( mT log( n 2 mT )) under various assumptions on the µ ft ’s [5, 63, 7, 58, 4, 6].
Finally, we note that δ is a hyperparameter governing the error rate.
It is possible to prove results analogous to Theorem 1 for arbitrary δ .
Experimental Results
We demonstrate the advantages of MABSplit in two settings.
In the first setting, the baseline models with and without MABSplit are trained to completion and we report wall-clock training time and generalization performance.
In the second setting, we consider training each forest with a fixed computational budget and study effect of MABSplit on generalization performance.
We provide a description of each dataset in Appendix 6.
We note that head-to-head wall-clock time comparisons with common implementations of forestbased algorithms, such as Weka’s [21] or scikit-learn ’s [48], would be unfair due to their extensive hardware- and language-level optimizations. As such, we reimplement these baselines in Python and focus on algorithmic improvements. The only difference between our model and the baselines is the call to the node-splitting subroutine (MABSplit in our model, and the exact, brute-force solver for the baseline models); thus, any improvements in runtime are due to improvements in the nodesplitting algorithm. This is verified by profiling the implementations and measuring the relative time spent in the node-splitting subroutine versus total runtime (see Appendix 5).
Our approach allows us to focus on algorithmic improvements as opposed implementation-specific optimizations. Our implementation of these baselines may also be of independent interest and we verify the quality of our implementations, via agreement with scikit-learn , in Appendix 4. We also provide a brief discussion of how an optimized version of our reimplementations may outperform scikit-learn in
Appendix 7. On the MNIST dataset, our optimized implementation trains approximately 4x faster than scikit-learn ’s DecisionTreeClassifier .
Baseline Models: We compare the histogrammed versions of three baselines with and without
MABSplit: Random Forest (RF), ExtraTrees [24], and Random Patches (RP) [39]. In Random
Forests, each tree fits a bootstrap sample of N datapoints and considers random subset of √ M features at every node split.
Extra Trees (also known as Extremely Randomized Forests) are identical to Random Forests except for two differences.
First, in regression problems, all features are considered at every node split (in classification problems, we still use only √ M features).
Second, the histogram edges of a feature are randomly chosen from a uniform distribution over that feature’s minimum and maximum value.
In classification problems, each histogram has √ M bins and in regression problems, each histogram has M bins.
These conventions follow standard implementations [47].
Note that in ExtraTrees, the bins in a feature’s histogram need not be equally spaced.
Random Patches is identical to Random Forests but the training dataset is reduced to α n of its original datapoints and α f of its original features, where α n and α f are prespecified constants, and the subsampled dataset is fixed for the training of the entire forest.
Full settings for all experiments are given in Appendix 6.
5.1 Wall-clock time comparisons
In the first setting, we compare baseline models with and without MABSplit in terms of wall-clock training time.
Tables 1 and 2 show that MABSplit provides similar generalization performance but faster training than the usual naïve algorithm for node-splitting for almost all baselines in both classification and regression tasks.
Across various tasks, MABSplit leads to approximately 2x-100x faster training, a reduction of training time of 50-99%.
These benefits are wholly attributable to MABSplit as the only difference between successive minor rows in each table is the node-splitting subroutine (see Appendix 5 for further discussion).
5.2 Fixed budget comparisons
In the second setting, we consider training models under a fixed computational budget.
As before, insertion into a histogram is taken to be an O (1) operation.
This is justified if the histogram’s thresholds are evenly spaced, wherefore the correct bin in which to insert a value can be indexed into directly.
(If the bins are unevenly spaced, we may perform binary searches to locate the correct bin, which is O (log T ) and does not depend on n , or cache the results of these binary searches for an evenly-spaced grid across the range of the given feature’s value.) [Table] [Table]
MABSplit can accelerate these models by up to 2x (an 50% reduction in training time) while achieving comparable results.
We omit the number of histogram insertions in favor of wall-clock time for simplicity; unlike in classification, the different baseline regression models have widely varying histogram bin counts.
Since the histogram insertion complexity scales as O (log T ) , the comparison across models would not be fair.
[Table] to be trained and leads to better generalization performance. Intuitively, the MABSplit algorithm allows for splitting a given node with less data point queries and histogram insertions than the naïve solution.
As such, when the computational budget is fixed, forests trained with MABSplit should be able to split more nodes and therefore train more trees than forests trained with the naïve solver.
Prior work suggests that increasing the number of trees in a forest improves generalization performance by reducing variance at the cost of slightly increased bias [27].
Tables 3 and 4 demonstrate the generalization performance of different models as the computational budget is held constant for different classification and regression tasks.
When using MABSplit, the trained forests consist of more trees and demonstrate better generalization performance across all baseline models.
5.3 Feature stability comparisons
We also apply MABSplit to compute feature importances under a fixed budget.
We follow the common approach of computing the feature importances of multiple forests and then measuring the stability of feature selection across forests using Permutation Feature Importance and Mean
Decrease in Impurity (MDI) [42, 49] (see Appendix 6 for a further discussion of these metrics).
The forests trained with MABSplit demonstrate better feature stabilities than those trained with the naïve algorithm; see Table 5.
Note that the datasets used for these experiments are different from the real-world datasets used in the other experiments and are described in Appendix 6.
Discussions and Conclusions
In this work, we presented a novel algorithm, MABSplit, for determining the optimal feature and corresponding threshold at which to split a node in tree-based learning models.
Unlike prior models such as Random Patches, in which the subsampling hyperparameters α n and α f must be prespecified manually, MABSplit requires no tuning and queries only as much data as is needed by virtue of its [Table] [Table]
MABSplit allows more trees to be trained, which leads to greater feature stabilities across the forests. adaptivity to the data distribution.
Indeed, robustness to choice of hyperparameters is one of primary appeals of algorithms like RF [50].
MABSplit avoids the expensive O ( n log n ) sort used in many existing baselines and the O ( n ) computational complexity of their corresponding histogrammed versions.
MABSplit can be used in conjunction with existing software- and hardware-specific optimizations and with other methods such as Logarithmic Split-Point Sampling and boosting [61].
In boosting, MABSplit has the potential advantage of only needing to update data points’ targets on-the-fly, as needed by its sampling, as opposed to current approaches that update targets for the entire dataset at each iteration.
Additionally, MABSplit may permit easier parallelization due to lower memory requirements than existing algorithms, which may enable greater use in edge computing and may be adaptable to streaming settings.
We also note that log n 2 term in the complexity of MABSplit (Equation 7) comes from a union bound over all possible n 2 confidence intervals computed for each feature-threshold pair across all stages of the algorithm (see the proof of Theorem 1 in Appendix 1). This union bound may be weak. Instead, one can more precisely compute the number of confidence intervals in terms of the arm gaps (the ∆ i ’s). In this way, the complexity result (Equation 7) can be phrased in terms of the ∆ i ’s instead of n .
Intuitively, this would lead to a complexity bound in terms of the data-generating distribution that is independent of dataset size. Indeed, this is how complexity results are usually stated for multi-armed bandit algorithms [31]. We leave a more detailed investigation of this topic to future work. Acknowledgements: M. T. was funded by a J.P. Morgan AI Fellowship, a Stanford Indisciplinary
Graduate Fellowship, a Stanford Data Science Scholarship
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] We have all reviewed the claims in the abstract and introduction and confirmed that they reflect the paper’s contributions and scope.
(b) Did you describe the limitations of your work?
[Yes] We discuss the assumptions necessary for MABSplit to work in Section 4 and also discuss failure modes when these assumptions are violated.
We also describe further room to improve our work in Section(c)
Did you discuss any potential negative societal impacts of your work?
[Yes] We discuss potential negative societal impacts in the Future Work section (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] We have all reviewed the ethics review guidelines and ensured that our paper conforms to them.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] The statement of Theorem 1 contains a full description of assumptions, which is further expounded upon before and after the theorem statement.
(b) Did you include complete proofs of all theoretical results?
[Yes] The proof of Theorem 1 is provided in Appendix3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] Yes, full specifications of all experiments are provided in Section 5 and Appendix 6, and full, documented source code is provided to reproduce our results.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] Yes, full specifications of all experiments are provided in Section 5 and Appendix 6, and full, documented source code is provided to reproduce our results.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] Every numerical measurement in Tables 1, 2, 3, andcontains confidence intervals (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] Full specifications of computing resources are provided in4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] Yes, the only existing assets we used were publicly available datasets, which we cited in Section ??
(b) Did you mention the license of the assets?
[Yes] We explicitly mention the licenses in Section 5.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] Yes, the only new assets created are in the form of code, which is entirely submitted in the supplemental material (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[Yes] Yes, both datasets we use are publicly available and we mention the licenses in Section 5.
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A]
The two datasets are the MNIST dataset and a mathematical dataset; neither have any personally identifiable information or offensive content.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A]
We did not use crowdsourcing or conduct research with human subjects.
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A]
We did not use crowdsourcing or conduct research with human subjects.
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
We did not use crowdsourcing or conduct research with human subjects.
Appendix Proofs
In this section, we present the proof of Theorem 1.
Proof.
Following the multi-armed bandit literature, we refer to each feature-threshold pair ( f, t ) as an arm and refer to its optimization objective µ ft as the arm parameter.
Pulling an arm corresponds to evaluating the change in impurity induced by one data point at one feature-threshold pair ( f, t ) (i.e., arm) and incurs an O (1) computation.
This allows us to focus on the number of arm pulls, which translates directly to sample complexity.
First, we show that, with probability at least 1 − 1 n , all confidence intervals computed throughout the algorithm are valid, in that they contain the true parameter µ ft .
For a fixed ( f, t ) and a given iteration of the algorithm, the (1 − δ ) confidence interval satisfies
[Equation]
Let B denote the batch size chosen for MABSplit. Note that there are at most nB rounds in the main while loop (Line 6) of Algorithm 1 and hence at most nmT B ≤ nmT confidence intervals computed across all arms and all steps of the algorithm.
With δ = 1 n 2 mT , we see that µ ft ∈ [ˆ µ ft − C ft , ˆ µ ft + C ft ] for every arm ( f, t ) and for every step of the algorithm with probability at least 1 − 1 n , by a union bound over at most nmT confidence intervals.
Next, we prove the correctness of Algorithm 1.
Let ( f ∗ , t ∗ ) =
arg min f ∈F ,t ∈T f µ ft be the desired output of the algorithm.
Since the main while loop in the algorithm can only run nB times, the algorithm must terminate.
Furthermore, if all confidence intervals throughout the algorithm are correct, it is impossible for ( f ∗ , t ∗ ) to be removed from the set of candidate arms.
Hence, ( f ∗ , t ∗ ) (or some ( f, t ) with µ ft = µ f ∗ t ∗ ) must be returned upon termination with probability at least 1 − 1 n .
This proves the correctness of Algorithm 1.
Finally, we consider the complexity of Algorithm 1.
Let n used be the total number of arm pulls computed for each arm remaining in the set of candidate arms at a given point in the algorithm.
Notice that, for any suboptimal arm ( f, t ) ̸
= ( f ∗ , t ∗ ) that has not left the set of candidate arms, we must have C ft ≤ c 0 (cid:113) log1 /δ
n used by assumption.
With δ = 1 n 2 mT as above and ∆ ft = µ ft − µ f ∗ t ∗ , if n used > 4 c 20 ∆ 2 ft log( n 2 mT ) then
[Equation]
and
[Equation]
which means that ( f, t ) must be removed from the set of candidate arms at the end of that iteration.
Hence, the number of data point computations M ft required for any arm ( f, t ) ̸
= ( f ∗ , t ∗ ) is at most
[Equation]
Notice that this holds simultaneously for all arms ( f, t ) with probability at least 1 − 1 n .
We conclude that the total number of arm pulls M satisfies
E [ M ] ≤ E [ M | all confidence intervals are correct ]
[Equation]
[Equation] ) ≤ (cid:88) f [Equation] t ∈T
[Equation]
where we used the fact that the maximum number of computations for any arm is 2 n .
As argued before, since each arm pull involves an O (1) computation, M also corresponds the total number of computations.
O (log n ) scaling of MABSplit
In Theorem 1, we demonstrated that MABSplit scales logarithmically in dataset size.
In this section, we empirically validate this claim.
Appendix Figure 1 (a) demonstrates the number of data points queried by MABSplit for a single node split, i.e., a single call to MABSplit, as the dataset size increases, for various subset sizes of
MNIST. For each sample size, a sample is drawn with replacement from the original MNIST dataset.
The model is trained using Gini impurity in for the usual digitr classification task.
Appendix Figure 1 (b) demonstrates the same plot for various subset sizes of the Random Linear Model dataset.
The Random Linear Model dataset consists of 200,000 datapoints with 50 features, 6 of which are correlated with the targets and 44 of which are pure noise, using scikit-learn ’s make_regression function [47]; 160,000 datapoints are used for training and the remaining 40, for test.
Appendix Figure 1 also shows the best linear and logarithmic fits to each dataset.
The relatively high R 2 values of the logarithmic fits ( R 2 = 0 .
97 and R 2 = 0 .
compared to that of the linear fits ( R 2 = 0 .
66 and R 2 = 0 . 43 ) suggests that the scaling of MABSplit is logarithmic (and therefore sublinear) with dataset size.
[Figure]
Appendix
Mean Estimation and Confidence Interval Constructions
In this section, we discuss the estimation of the means µ ft and construction of their confidence intervals via plug-in estimators and the delta method.
Let p L ,k , p R ,k , ˆ p L ,k , and ˆ p R ,k be the same as defined in Subsection 3.1.
Furthermore, let p =
[ p L , 1 , · · · , p L ,K , p R , 1 , · · · , p R ,K ] T and ˆ p =
[ˆ p L , 1 , · · · , ˆ p L ,K , ˆ p R , 1 , · · · , ˆ p R ,K ] T .
Then, n ′ ˆ p follows a multinomial distribution with parameters ( n ′ , 2 K, p ) .
Let θ =
[ p L , 1 , · · · , p L ,K , p R , 1 , · · · , p R ,K − 1 ] T and ˆ θ =
[ˆ p L , 1 , · · · , ˆ p L ,K − 1 , ˆ p R , 1 , · · · , ˆ p R ,K − 1 ] T .
Then, by the Central Limit Theorem,
[Equation]
where Σ ii =
θ i (1 − θ i ) and Σ ij = − θ i θ j .
Next, we write µ ft in terms of θ for the impurity metrics as
Gini impurity
[Equation]
Entropy : µ ft ( θ ) =
[Equation]
For a given impurity metric, let ∇ µ ft ( θ ) be the derivative of µ ft with respect to θ .
From the delta method,
[Equation]
where the CIs can be constructed accordingly.
These CIs are asymptotically valid as n ′ , n → ∞ .
For other impurity metrics such as MSE, the CIs can be similarly derived by writing the corresponding µ ft in terms of θ and computing ∇ µ ft ( θ ) .
[Table]
Comparison of baseline implementations and scikit-learn
In this section, we compare our re-implementation of common baselines to those in popular packages to verify the accuracy of our re-implementation.
Specifically, we compare our implementations of Random Forest Classifiers, Random Forest Regressors, Extremely Random Forest Classifiers, and Extremely Random Forest Regressors to those of scikit-learn .
We omit comparisons of the Random Patches models because their correctness is implied by that of the Random Forest model, as the Random Patches model consists of applying the Random Forest model to subsampled data and features.
For classification, we compare our implementations on the 20 newsgroups dataset filtered to two newsgroups, alt.atheism and sci.space .
The dataset is embedded via TF-IDF and projected onto their top 100 principal components, following standard practice [47].
The train-test split is the standard one provided by scikit-learn .
For all classification problems, we average the predicted probabilities of each tree in the forest ("soft voting") as opposed to only allowing each tree to vote for a single class ("hard voting"), following the implementation in scikit-learn [47].
For regression, we compare our implementations on the California Housing dataset, subsampled to 1,000 points as performing the regression on the full dataset of approximately 20,000 points is computationally prohibitive.
The train-test split is the standard one provided by scikit-learn [47].
Table 1 presents our results.
In all cases, our re-implemented baselines do not present a statistically significant difference in performance from the models present in scikit-learn , which suggests that our re-implementations are correct.
Performance is measured over 20 random seeds to compute averages and standard deviations.
Profiles
In this work, we focused on the reducing the runtime at the algorithmic level, i.e., reducing the complexity of computing the best feature-threshold split.
In this section, we justify this choice by demonstrating that most of the time spent in our re-implementation of the baseline algorithms is spent in computing the best feature-threshold split.
Appendix Figure 2 demonstrates the wall-clock time spent inside various functions when fitting a
Random Forest classifier without MABSplit on two subsets of the MNIST dataset of sizes 5, and 10,000.
Most of the time is spent inside the computation of the best feature-threshold split, which scales approximately as dataset size and motivates our focus on improving the performance of the split-identification subroutine.
When using MABSplit, the time spent to identify the best feature-threshold split is reduced significantly (Appendix Figure 3).
Appendix Figure 4 also contains an example callgraph demonstrating callers and callees for the fitting procedure of a Random Forest, for easier interpretation of Appendix Figures 2 and 3.
[Figure]
Appendix for the node-splitting algorithm using the /naïve computation, the algorithm computing When increasing the dataset size the overhead spent the solve_exactly function grows negligibly from about 0.5 seconds to about 1 second.
However, the time spent in the solve_exactly function and any called subroutines grows from about 28 seconds to aboutseconds and constitutes approximately 98% of the increase in wall-clock time.
This observation motivates our focus on improving the subroutine used to identify the best feature-threshold split.
This profile was generated with cProfile and visualized with pyprof2calltree [37].
[Figure]
Appendix
Experiment Details
Here we provide full details for the experiments in Section 5.
All experiments were run on[-
-]
MacBook Pro running MacOS 12.5.1 (Monterey) with an Apple M1 Max processor, and 64 GB RAM. 6.1 Datasets Classification Datasets: We use the MNIST [38], APS Failure at Scania Trucks [25, 20], and Forest
Cover Type [13, 20] datasets.
The MNIST dataset consists of 60,000 training and 10,000 test images of handwritten digits, where each black-and-white image is represented as a 784-dimensional vector and the task is to predict the digit represented by the image. The APS Failure at Scania Trucks dataset consists of 60,000 datapoints with 171 features and the task is to predict component failure. The Forest Covertype dataset consists of 581,012 datapoints with 54 feature and the task is to predict the type the forest cover type from cartographic variables.
Regression Datasets: We use the Beijing Multi-Site Air-Quality [64, 20] and the SGEMM GPU
Kernel Performance [8, 44, 20] datasets.
The Beijing Multi-Site Air-Quality dataset consists of 420,768 datapoints with 18 features and the task is to predict the level of air pollution.
The SGEMM
GPU Kernel Performance dataset consists of 241,600 datapoints and the task is to predict the running time of a matrix multiplication.
For all datasets except MNIST (which has predefined training and test datasets), all datasets were randomized into 9:1 train-test splits.
All datasets are publicly available.
6.2 Runtime Experiments
For the runtime experiments presented in Tables 1, all performances were measured from 5 random seeds. For all datasets, the maximum depth was set to 1 except for the MNIST dataset, in which the maximum depth was set to 5.
The number of trees in each model was set to 5.
All experiments used the Gini impurity criterion and the minimum impurity decrease required from performing a split was set to 0.005. For the Random Patches (RP) model, α n was set to 0.7 and α f was set to 0.85. For the regression runtime experiments presented in Table 2, all performances were measured fromrandom seeds.
For the Beijing Multi-Site Air-Quality Dataset, the maximum depth was set to 1 and for the SGEMM GPU Kernel Performance Dataset, the maximum number of leaf nodes was set to 5.
The number of trees in each model was set to 5. All experiments used the MSE impurity criterion and the minimum impurity decrease required from performing a split was set to 0.005. For the Random Patches (RP) model, α n was set to 0.7 and α f was set to 0.85.
6.3 Budget Experiments
For the classification budget experiments presented in Table 3, all performances were measured from 5 random seeds. The budget for each model on the MNIST, APS Failure at Scania Trucks, and Forest
Covertype datasets were set to 10,192,000, 784,000, and 9,408,000, respectively. For the Random
Patches (RP) model, α n was set to 0.6 and α f was set to 0.8. The maximum number of trees in any model was set to 100 and the maximum depth of each tree was set to 5. For the regression budget experiments presented in Table 4, all performances were measured fromrandom seeds.
The budget for each model on the Beijing Multi-Site Air-Quality Dataset was set to 76,800,000 and the budget for each model on the SGEMM GPU Kernel Performance Dataset was set to 24,000,000.
For the Random Patches (RP) model, α n was set to 0.8 and α f was set to 0.5. The maximum number of trees in any model was set to 100 and the maximum depth of each tree was set to 5.
6.4 Stability Experiments
Two metrics for calculating feature importance are used in Table 5: out-of-bag Permutation Importance (OOB PI) and Mean Decrease in Impurity (MDI) [42, 49].
For a feature f , the OOB PI is calculated by measuring the difference between the trained model’s out-of-bag error on the original data with its out-of-bag error on all the data with all out-of-bag datapoints’ f values shuffled.
The
MDI for a feature f is the average decrease in impurity of all nodes where f is selected as the splitting criterion.
Once feature importances have been calculated, the top k most important features for the model are selected and the stability of these k features is measured via standard stability formulas [43].
The results of the stability experiments are shown in Table 5.
The Random Classification dataset is generated via scikit-learn ’s datasets.make_classification function with n_samples=10000, n_features=60 , and n_informative=5 .
The Random
Regression dataset is generated by scikit-learn ’s datasets.make_regression with n_samples=10000, n_features=100 , and n_informative=5 . [Table]
Limitations 7.1 Theoretical Limitations
Crucial to the success of MABSplit are the assumptions described before and after Theorem 1.
In particular we assume that their is reasonable heterogeneity amongst the true impurity reductions of different feature-value splits.
Such assumptions are common in the literature and have been validated on many real-world datasets [5, 63, 7, 58, 4, 6].
We also note that the assumptions that each CI scales as (cid:113) log1 /δ n ′ may be violated when using certain impurity metrics.
For example, the derivative of the entropy impurity criterion with respect to some p k approaches ∞ when p k → 0 .
In this case, we cannot apply the delta method from Appendix 3 to compute finite CIs that scale in the way we require.
In such settings, it may be necessary to compute the CIs in other ways, e.g., following [46] or [9].
We note that in the worst case, even when all assumptions are violated, MABSplit is never worse than the naïve algorithm in terms of sample complexity.
In the worst case, it is a batched version of the naïve algorithm.
7.2 Practical Limitations
We note that MABSplit may perform worse than naïve node-splitting on very small datasets, where the overhead of sampling the data in batches outweighs any potential benefits in sample complexity (see Appendix 8 for further discussion).
In this work, we avoided a direct runtime comparison with scikit-learn because scikit-learn utilizes a number of low-level implementation optimizations that would make the comparison unfair. To provide a brief comparison to the popular scikit-learn implementation, however, we attempted to optimize our implementation using numba [36], a package that translates Python code to optimized machine code. Our numba -optimized implementation is 4x faster than scikit-learn ’s DecisionTreeClassifier and achieves comparable performance on the MNIST dataset; see
Appendix Table 2. In order for practitioners to take full advantage of MABSplit, however, it may be necessary to implement MABSplit within the scikit-learn library.
In doing so, it may be possible that MABSplit makes it difficult or impossible to use existing optimizations in the scikit-learn library.
An example of this is vectorization: because the naïve node-splitting algorithm queries the data in a predictable way, each datapoint can be queried more quickly than in MABSplit.
Despite MABSplit’s advantages in sample complexity, the disadvantages of being unable to use implementation optimizations like vectorization may outweigh MABSplit’s benefits.
Many of these risks may be ameliorated by addressed MABSplit into existing RF implementations such as the one in scikit-learn .
We anticipate that many optimizations will still apply: for example pre-fetching data to have it in caches close to the CPU, manual loop unrolling, etc.
We leave an optimization implementation of MABSplit inside the scikit-learn library to future work.
Comparison on Small Datasets
In this section, we investigate the performance of MABSplit on small datasets.
Appendix Figuredemonstrates the performance of MABSplit, both in wall-clock training time and sample complexity, for various subset sizes of MNIST.
Our results that RF+MABSplit outperforms the standard RF algorithm, in both sample complexity and wall-clock time, when the dataset size exceeds approximately 1100 datapoints.
However, we also note that the main use case for MABSplit is when the data size is large and it is computationally challenging to run standard forest-based algorithms.
Indeed, the use of big data in many applications that necessitate sampling was the primary motivation for our work [11, 54, 18, 57, 41, 1].
[Figure]
Appendix
Description of Other Node-Splitting Algorithms
For completeness, we provide a brief description of various baseline models’ node-splitting algorithms here to enable easier comparison with MABSplit.
Consider a node with n datapoints each with m features, and T possible thresholds at which to split each feature.
We discuss the classification setting for simplicity, though the same arguments apply to regression.
A very naïve approach would be to iterate over all mT feature-value splits, and compute the probabilities p L,k and p R,k from all n datapoints.
This results in complexity O ( mTn ) , which is O ( mn 2 ) when T = n (for example, T = n in the un-histogrammed setting).
Instead, the usual RF algorithm sorts all n datapoints in O ( n log n ) time for each of the m features, resulting in total computational cost O ( mn log n ) .
Then the algorithm scans linearly from lowest value to highest value for each feature and update the parameters p L,k and p R,k via simple counting to find the best impurity reduction for each of the T potential splits.
The complexity of this step is O ( mT + mn ) , where the “ + mn ” comes from the allocations of each data point to the left or right node during the scan (each data point is re-allocated only once per feature).
Thus the total complexity of this approach is O ( mn log n + mT + mn )
=
O ( mn log n + mT ) .
This is O ( mn log n ) when T = n .
The binned (a.k.a. histogrammed) method does not require the per-feature sort and avoids the
O ( mn log n ) computation when T < n ,.
Instead, each of the n points must be inserted into the correct bin (which can be done in O (1) time for each datapoint if the bins are equally spaced) for each of the m features, incurring total computational cost O ( mn ) .
Then, the same linear scanning approach as in the “standard” algorithm is performed with complexity O ( mT + mn ) .
The total complexity of this approach is O ( mn + mT + mn )
=
O ( m ( n + T )) .
This is O ( mn )
when T = n .
In general, we do not assume T = n , i.e., that every feature value is a potential split point, unless otherwise specified.
In our paper, the “standard” approach refers to the un binned approach which requires an O ( mn log n ) sort and “linear” refers to the binned approach that is O ( m ( n + T )) , which is O ( mn ) when T =
O ( n ) .
Crucially, when T = o ( N ) (as is often the case in practice, e.g., for a constant number of bins) and the necessary gap assumptions are satisfied, MABSplit scales as O ( mT log n ) .
In many cases, this is much better than O ( m ( n + T )) , e.g., for large datasets, because the dependence on n is reduced from linear to logarithmic.
More concretely, treating T as a constant and ignoring the dependence on m , we reduce the complexity of the binned algorithm from O ( n ) , what we refer to as “linear,” to
O ( log n ) .