MABSplit: Faster Forest Training Using Multi-Armed Bandits
Abstract
Ensemble learning methods such as random forest are some of the most widelyused machine learning models, especially in domains that necessitate interpretability.
We present an algorithm that accelerates the training of random forest andother popular tree-based learning methods.
At the core of our algorithm is a noveland fast node-splitting subroutine, dubbed MABSplit, used to efficiently find splitpoints when constructing decision trees.
Our algorithm borrows techniques fromthe multi-armed bandit literature to judiciously determine how to allocate computational power across potential split points.
We provide theoretical guarantees that
MABSplit improves the computational complexity from linear to logarithmic inthe number of data points.
Even on small datasets such as MNIST, our algorithmleads to 7x faster training (an 85% reduction in training time) without any decreasein test accuracy.
We demonstrate similar speedups when the MABSplit subroutineis used across a variety of forest-based variants, such as Extremely Random Forestsand Random Patches.
We also show our algorithm can be used in both classification and regression tasks.
Finally, MABSplit outperforms existing methods intest performance and feature importance calculations under a fixed computationalbudget.
Introduction
The Random Forest (RF) algorithm is a supervised learning technique that is widely used forclassification and regression tasks [21, 11].
In RF, an ensemble of decision trees (DTs) is trainedfor the same prediction task.
Typically, bootstrapping [10] is used to generate multiple sample setsfrom a dataset, each of which is used to train a separate DT.
Moreover, in the process of trainingeach DT, a random subset of the features is considered at each candidate split.
The prediction ofthe different DTs is then aggregated to provide a single prediction.
Common aggregation schemesinclude majority voting in classification tasks and averaging in regression tasks.
By aggregatingthe prediction of the different DTs in the ensemble, RFs tend to be more robust against noise andoverfitting [41] and are capable of capturing multiple patterns in the data, as opposed to a single DT[44].
RF has gained tremendous popularity due to its flexibility, usefulness in multi-class classificationand regression tasks, high performance across a broad range of data types, and relatively lowcomputational complexity [46, 42, 8, 27].
Furthermore, RF inherits the interpretability of decisiontrees, as the prediction of each individual 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[24], solar radiation analysis [8], and the Higgs Boson classification problem [2].
As we move intothe era of big data, a simple and flexible machine learning technique such as RF is expected to play akey 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-constrainedhardware (such as smartphones and Internet-of-Things devices) requires the development of newalgorithms that can deliver faster performance without sacrificing results [46].
For this reason, recentwork has proposed ways to accelerate the training of RFs, both at the algorithmic level and at thehardware level.
At the algorithmic level, the focus is on the fast construction of each individualdecision tree.
A decision tree is built by identifying a feature f and a threshold t that best split thedata points according to the prediction targets.
The data points are split based on f and t , and theprocess is repeated for each resulting subset.
The key computational bottleneck is the subroutinerequired for identifying the pair ( f, t ) that provides the best split for a set of N data points, whichtypically requires O ( N ) computations per split.
Recent proposals include computing (or estimating)f and t from a subsample of the data points and a subset of the features, or quantizing the featurevalues.
The latter technique creates a histogram of values of each feature across the datapoints andrestricts t to be at the edges of histogram bins.
In Subsection 1.1, we discuss these approaches inmore detail.
While existing approaches provide significant speedups in the training of RFs, they are often basedon a pre-defined strategy in terms of the number of data points and number of features to consider ateach split.
As such, they are not able to adapt to the data distribution.
Moreover, when comparingdifferent candidate features f for a split, all features are treated equally and the quality of their splitis computed based on the same number of data points.
Intuitively, this is wasteful as features that arenot informative for the prediction task can be identified based on a very small number of data points.
Alternatively, an adaptive scheme could better allocate computational resources towards a promisingset of candidate features and achieve a better tradeoff between computational cost and accuracy.
In this work, we propose MABSplit, a fast subroutine for the node splitting problem, which adaptivelyrefines the estimate of the “quality” of each feature-threshold pair ( f, t ) as a candidate split.
In thisway, bad splitting candidates can be discarded early, which can lead to significant computationalsavings.
The main idea behind our algorithm is to reduce the node-splitting task to a multi-armedbandit problem [26, 1, 4, 43], where each pair ( f, t ) is a distinct arm.
The unknown parameter ofeach arm, µ ft , corresponds to the quality of the split of the data points X
= { x 1 , . .
, x N } basedon feature f and threshold t .
The split quality is measured in terms of how much the split wouldreduce label impurity, for a chosen label impurity measure such as Gini impurity.
An arm ( f, t )can be “pulled” by computing the reduction of label impurity on a small subset of data points from X .
This allows us to compute an estimate ˆ µ ft and a corresponding confidence interval, which canbe used in a batched variant of the Upper Confidence Bound (UCB) and successive eliminationalgorithms [26, 48] to identify the best arm ( f, t ) .
We demonstrate the benefits of MABSplit on a variety of datasets, for both classification and regressiontasks.
Even on small datasets such as MNIST, our algorithm leads to 7x faster training (an 85%reduction in training time), without any decrease in test accuracy, over an exact implementation of RFthat searches for the optimal ( f, t ) pair via brute-force computation.
Additionally, we demonstratesimilar 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 [21].
In its basic formulation, a RF involvesconstructing n tree DTs, each based on a bootstrap sample of N data points and a random subset of thefeatures (known as random subspacing [22]).
More recently, the need for training RFs on large-scaledatasets has prompted the development of several techniques to accelerate the training process viaboth software and hardware techniques.
Software acceleration of RF: Most of the software/algorithmic acceleration techniques
are basedon accelerating the training of each individual DT.
FastForest [46] accelerates the node-splittingtask using three ideas: subsampling a pre-specified number of datapoints (subbagging) withoutreplacement, subsampling a pre-specified number of features dependent on the current number of datapoints (Dynamic Restricted Subspacing), and dividing values of a given feature into T bins where Tdepends on the number of data points at the node (Logarithmic Split-Point Sampling, inspired by thesingle-tree SPAARC algorithm
MABSplit is inspired by the ideas in FastForest, but does notrequire the number of datapoints or features to be pre-specified as these are determined adaptively.
Recent work has discussed using adaptivity to identify the best split.
For example, Very Fast Decision Tree (VFDT) [14] and Extremely Fast Decision Tree (EFDT) [31] are incremental decision treelearning algorithms in which trees can be updated in streaming settings.
Acceleration at the nodesplitting task is achieved by adaptively selecting a subset of data points sufficient to distinguishthe best and second best features.
This is similar to our approach, but the sampled data points areevaluated for all features, unlike in MABSplit, which adaptively discards unpromising feature splits.
The F-forest [17] algorithm also applies adaptivity and utilizes an upper bound on the impurityreduction of a split in order to discard candidate splits early and save on computation.
This is similarin spirit to the goal of MABSplit, but is based on a deterministic upper bound on the impurityreduction (as opposed to our statistical estimate) and is still linear in N per split computation, evenwhen considering a fixed number of splitting thresholds per feature.
Other variations of RF have been proposed to improve training time.
Random Patches [30] buildstrees based on a subset of data points and a subset of features (which is fixed for each entire DT,unlike in RF).
Extra-Trees (ETs) [18] draws a random subset of K feature at each node and, foreach one, chooses a number R of random splits, and then selects the split that yields the largestimpurity 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 [34], a form ofhyperparameter tuning.
The MABSplit subroutine can also be incorporated into these methods.
Hardware acceleration of RF: The training of RFs can also be significantly accelerated throughthe use of specialized hardware.
For instance, the implementation of RF available in Weka [15]allows trees to be trained on different cores and reduces forest build time.
A GPU-based parallelimplementation of RF has also been proposed in [19].
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 also desirable to develop techniques toimprove prediction performance and processing speed on standard platforms.
Preliminaries
We consider N data points { ( x i , y i ) } Ni =1 where each x i is an M -dimensional feature vector and y iis its target.
Following standard RF literature, we consider flexible feature types such as numericalor categorical; RFs also deal naturally with missing values [23].
We also consider both categoricaltargets 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.
A random forest contains n tree decision trees, each constructed using a set of bootstrapped N samples(sampled with replacement) and a random set of features (typically √ M [9]) at each node for a giventree; it outputs the majority vote across trees for the classification task and average prediction acrossall trees for the regression task [21, 11].
We focus on the top-down approach of constructing DTs bychoosing a variable at each step that best splits the set of labels at a given node [41].
Specifically,consider a node with n data points X and m features F .
Let X L ,ft and X R ,ft be the left and rightchild subsets of X when X is split according to feature f at threshold t .
The approach finds the splitthat best reduces the label impurity; i.e., it finds
[Equation]
where I ( X ) measures the impurity of targets { y i } i ∈X and T f is the set of feasible thresholds forfeature f .
Common impurity measures include Gini impurity or entropy for classification, andmean-squared-error (MSE) for regression [12]:
[Equation]
where K is the total number of classes, p k = 1 n (cid:80) i ∈X I ( y i = k ) is the proportion of class k in X inclassification and ¯ y = 1 n (cid:80) i ∈X y i is the average target value in regression.
We note that our proposedalgorithm, MABSplit, does not assume a particular structure of I ( · ) .
While the conventional RF considers all ( n − 1) possible splits at a feature f , in this work we focuson the histogram-based variation which chooses the threshold from a set of predefined values T f , e.g., |T f | equally-spaced histogram bin edges; this variation is substantially more computationallyefficient, offers comparable accuracy, and has been used in most state-of-the-art implementations[40, 13, 25, 47].
MABSplit
Without loss of generality, consider a classification task with K classes.
RF constructs n tree DTs, eachof which requires recursively partitioning data points at a given node across a set of √ M features.
Consider a node with n data points X , m features F , and T possible thresholds for each feature T f .
We note that n and m at the given node may or may not be the same as N and M of the entire dataset.
A naïve algorithm finds the best feature-threshold pairs in Eq.
(1) by evaluating the label impurityreduction for each feature-threshold over all n data points, which incurs computational cost linear inn .
We present MABSplit, which can substantially reduce this complexity to sublinear in n .
Note that Eq.(1) is equivalent to
[Equation]
and so we focus on solving Eq.
Let µ ft = |X L ,ft | n
I ( X L ,ft ) + |X R ,ft | n
I ( X R ,ft ) be the optimizationobjective for feature-threshold pair ( f, t ) .
Computing µ ft exactly is at least O ( n ) 1 , but we caninstead estimate µ ft with fewer computations by drawing n ′ <
n independent samples uniformlywith 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 δ determineestimation accuracy.
The width of CIs generally scales with (cid:113) log1 /δ n ′ .
To estimate the solution to
Problem (3) with high confidence, we can then choose to sample data points for different featuresto different degrees of accuracy.
Intuitively, promising features with small values of µ ft should beestimated with high accuracy, while less promising ones can be discarded without being evaluated ontoo many data points.
The specific adaptive estimation procedure is described in Algorithm 1.
It can be viewed as a batchedversion of the conventional UCB algorithm [26, 48] combined with successive elimination [16], andis straightforward to implement.
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 ) } .
Then, for eachpotential solution ( f, t ) ∈ S solution , Algorithm 1 maintains its mean objective estimate ˆ µ ft and (1 − δ )
CI C ft .
In each iteration, a new batch of data points X batch is evaluated for all potential solutions in S solutionand allows us to update the histogram for each feature f and make each estimate ˆ µ ft more accurate.
Based on the current estimate, if a candidate’s lower confidence bound ˆ µ ft − C ft is greater thanthe 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 sampledmore than n data points.
In the latter case, we know that the difference between the remainingcandidates in S solution is so subtle that an exact computation is warranted.
We then compute thosecandidates’ objectives exactly and return the best candidate in the set.
3.1 Point estimate and confidence interval for impurity metrics
We consider two widely used impurity metrics, Gini impurity and entropy, although mean estimatesand confidence intervals for other metrics can be derived similarly (more details are provided in
Appendix 3).
We consider 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.
Let p L ,k := 1 n (cid:80) ni
=
I ( x if < t, y i = n[-(cid:80) ni =
 I ( x if ≥ t, y i = k ) denote theproportion of the full n data points in class k and each of the two subsets (“left” and “right”),-]
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
▷ Update running mean and CI in 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: return ( f ∗ , t ∗ ) =
arg min ( f,t ) ∈S solution µ ft respectively).
Note that(cid:88)
[Equation] and (cid:88)
[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 (since for each random data point ( X i , Y i ) ,exactly one of these variables is incremented) satisfying
[Equation]
Using Eqs.
(2) and (4), and the definition of µ ft after Eq.
(3), we can 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 in the equations above to produce the point estimate ˆ µ ft .
In Algorithm 1, aftereach batch of B data points is obtained, we update ˆ p L ,k and ˆ p R ,k , which can then be used to updateour point estimate ˆ µ ft and the corresponding CIs. The CIs of ˆ µ ft are based on standard error derivedusing the delta method [45].
As in standard applications of the delta method, these estimates ˆ µ ft areasymptotically unbiased and the confidence intervals are asymptotically valid.
Appendix 3 providesfurther details including derivation of CIs and discussion of convergence properties.
3.2 Algorithmic and implementation details
We considered sampling with replacement Algorithm 1 primarily for the ease of theoretical analysis.
In practice, we found sampling without replacement was more computationally efficient and did notsignificantly 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.
This is in contrast with hard votes, in which each DT is only permitted to produce its best label andthe forest’s prediction is determined by majority voting.
For the fixed-budget experiments describedin Section 5, we terminate tree construction and do not split the next node if doing so would violateour budget constraints.
Finally, we note that a head-to-head comparison with common implementations of forest-basedalgorithms, such as Weka’s [15] or scikit-learn ’s
[37], would be unfair due to their extensivehardware- and language-level optimizations in those implementations.
As such, we reimplementedthese baselines in Python and focused on algorithmic improvements.
We verified the quality ofour own implementations agreed with that of scikit-learn in Appendix ?? .
The only differencebetween our model and the baselines is the call to the node splitting subroutine (MABSplit inour model, and the exact, brute-force solver for the baseline models); thus, any improvementsin runtime are due to improvements in the node-splitting algorithm.
This is verified by profilingthe implementations and measuring the relative time spent in the node-splitting subroutine versustotal runtime (see Appendix 5).
This allows us to focus on algorithmic improvements as opposedimplementation-specific optimizations.
Our implementation of these baselines may also be ofindependent interest.
Analysis of the Algorithm
In this section, we prove that MABSplit returns the optimal feature-threshold pair for a node splitwith high probability.
Furthermore, we provide bounds on computational complexity of MABSplit,leading to the logarithmic dependency on the number of data points n .
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 feature-threshold pairat 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 is the case for the confidence intervals of Gini impurity and entropy in Subsec.
Withthis assumption, we state the following theorem:
Theorem 1.
Assume ∃ c 0 > 0 s.t. ∀
δ > 0 , n ′ > 0 , C ft ( n ′ , δ ) <
c 0 (cid:113) log1 /δ n ′ .
For δ = 1 n 2 mT , withprobability at least 1 − 1 n , Algorithm 1 returns the correct solution to (3) .
Furthermore, Algorithm 1uses 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 instancewise bound Eq.
(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 a
O (log( n 2 mT )) computation that is sublinear in n .
In general, MABSplit takes only O (log( n 2 mT )) computations per feature-threshold pair if thereis reasonable heterogeneity among them.
As proven in Appendix 2 in [4], this is the case under awide range of distributional assumptions on the µ ft ’s, e.g., when the µ ft ’s follow a sub-Gaussiandistribution across the pairs ( f, t ) .
These assumptions ensure that MABSplit has an overall complexityof O ( mT log( n 2 mT )) , which is sublinear in the number of data points n .
However, we note thatin the worst case, MABSplit may take O ( n ) computations per feature-threshold pair when mostfeature-threshold pairs are equally good, for example in highly symmetric datasets where all splitsreduce the impurity equally.
Other recent work provides further discussion on the conversion betweenan instance-wise bound like Equation (7) and an instance-independent bound such as O ( mT log n )[4, 48, 6, 43, 3, 5].
Finally, we note that δ is a hyperparameter governing the error rate.
It is possible to prove resultsanalogous to Theorem 1 for arbitrary δ .
Experimental Results
We demonstrate the advantages of MABSplit in two settings.
In the first setting, the baselinemodels with and without MABSplit are trained to completion, and wall-clock training time and testperformances are measured.
In the second setting, we consider training each forest with a fixedcomputational budget and study effect of MABSplit on generalization performance.
Datasets: We run experiments on public datasets to validate the behavior of MABSplit, in bothclassification and regression.
For classification, we use the MNIST dataset [29] consisting of60,000 training and 10,000 test images of handwritten digits, where each black-and-white imageis represented as a 784-dimensional vector.
For regression, we generate 200,000 datapoints with50 features, 6 of which are correlated with the targets and 44 of which are pure noise, usingscikit-learn ’s make_regression function [36]; 160,000 datapoints are used for training and theremaining 40,000 for test.
We call this dataset the Random Linear Model dataset.
The two datasetsare available under MIT and BSD licenses, respectively.
Baselines: We compare the histogrammed versions of three baselines with and without MABSplit:
Random Forest (RF), Extremely Randomized Forests (ExtraTrees) [18], and Random Patches (RP)[30].
In Random Forests, each tree fits a bootstrap sample of N datapoints and only considersrandom subset of √ M features at every node split.
Extremely Randomized Forests are identical to
Random Forests except for two differences.
First, in regression problems, all features are consideredat every node split (in classification problems, we still use only √ M features).
Second, the histogramedges of a feature are randomly chosen from a uniform distribution over that feature’s minimum andmaximum value.
In classification problems, each histogram has √ M bins and in regression problems,each histogram has M bins.
These conventions follow standard implementations [36].
Note thatin Extremely Randomized Forests, the bins in a feature’s histogram need not be equally spaced.
In
Random Patches, the dataset is reduced to α n of its original datapoints and α f of its original features,where α n and α f are constants.
The reduced dataset is fixed for the training of the entire forest.
Fullsettings 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-clocktraining time.
Tables 1 and 2 show that MABSplit provides similar test accuracy but leads to fastertraining than the usual brute-force algorithm for node-splitting for almost all baselines in bothclassification and regression tasks.
Across various tasks, MABSplit leads to 2x-7x faster training, areduction of training time of 50-85%.
These benefits are wholly attributable to MABSplit as the onlydifference between successive rows is the node-splitting subroutine (see Appendix 5).
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’sthresholds are evenly spaced, in which the correct bin in which to insert a value can be indexeddirectly.
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 anevenly-spaced grid across the range of the given feature’s value.
Intuitively, the MABSplit algorithm allows for splitting a given node with less data point queries andhistogram insertions than the brute-force 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 thanforests trained with the brute-force solver.
Prior work suggests that increasing the number of trees ina forest improves generalization performance by reducing variance at the cost of slightly increasedbias [20].
Tables 3 and 4 demonstrates the generalization performance of different models as the computationalbudget is held constant for classification and regression, respectively. When using MABSplit, the [Table] [Table]
MABSplit accelerates these models by approximately 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 is not fair. trained forests consist of more trees and demonstrate better generalization performance across allmodels.
We also apply MABSplit to compute feature importances under a fixed budget.
We follow thecommon approach of computing the feature importances of multiple forests and then measuring thestability of feature selection across forests [32, 38]; see Table 5.
[Table] [Table] [Table]
MABSplit pares down the computation required to train a tree for each type of model.
This results in better performance for the Random Forest stability baseline.
scikit-learn datasets for both classification and regression were used.
Discussions and Conclusions
In this work, we presented a novel algorithm, MABSplit, for determining the optimal feature andcorresponding feature threshold at which to split a node in tree-based learning models.
Unlike priormodels such as Random Patches, in which the hyperparameters α n and α f must be chosen manually,
MABSplit requires no tuning and queries only as much data as is needed.
Indeed, robustness tochoice of hyperparameters is one of primary appeals of algorithms like RF [39].
MABSplit can avoid 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 usedin conjunction with existing software- and hardware-specific optimizations and with other methodssuch as Logarithmic Split-Point Sampling and boosting [46].
In boosting, MABSplit has the potentialadvantage of only needing to update target variables on-the-fly as needed, as opposed to currentboosting approaches which update target variables for the entire dataset at each iteration.
Additionally, MABSplit may permit easier parallelization due to lower memory requirements than existingalgorithms, which may enable greater use in edge computing, and may be adaptable to streamingsettings.
Checklist
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 statementof Theorem [-
-]1 contains a full description of assumptions, which is further expoundedupon before and after the theorem statement.
(b) Did you include complete proofs of all theoretical results?
[Yes] The proof of Theorem1 is provided in Appendix 13.
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, fullspecifications 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 Section5 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, and 4contains confidence intervals
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] Full specifications of computing
resources are provided in
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 existingassets 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 inthe 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.5.
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 ) asan arm and refer to its optimization object µ ft as the arm parameter.
Pulling an arm corresponds toevaluating the change in impurity induced by one data point i 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, whichtranslates directly to sample complexity.
First, we show that, with probability at least 1 − 1 n , all confidence intervals computed throughoutthe algorithm are valid, in the sense that they contain the true parameter µ ft .
For a fixed ( f, t ) and agiven iteration of the algorithm, the (1 − δ ) confidence interval satisfies
[Equation]
Note that there are at most nB rounds and hence at most nmTB ≤ nmT confidence intervals computedacross all arms and all steps of the algorithm, where B is the batch size.
With δ = 1 n 2 mT , we see that 514 µ ft ∈ [ˆ µ ft − C ft , ˆ µ ft + C ft ] for every arm ( f, t ) and for every step of the algorithm with probabilityat 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 desiredoutput of the algorithm.
First, observe that the main while loop in the algorithm can only runnB times, so the algorithm must terminate.
Furthermore, if all confidence intervals throughout thealgorithm 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 probabilityat 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 pullscomputed for each of the arms 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, wemust have C ft ≤ c 0 (cid:113) log1 /δ
n used by assumption.
With δ = 1 n 2 mT as above and ∆ ft = µ ft − µ f ∗ t ∗ , ifn 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 concludethat the total number of arm pulls M satisfies
[Equation]
all confidence intervals are correct
[Equation]
where we used the fact that the maximum number of computations for any arm is 2 n .
As arguedbefore, since each arm pull involves an O (
O (log n ) scaling of MABSplit
In Theorem 1, we demonstrated that MABSplit scales logarithmically in data set size.
In this section,we empirically validate this claim.
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 eachsample size, a sample is drawn with replacement from the original MNIST dataset.
The model istrained using Gini impurity in the typical classification setting.
Figure 1 (b) demonstrates the sameplot for various subset sizes of the Random Linear Model dataset.
Figure 1 also shows the best linear and logarithmic fits to each dataset.
The relatively high R 2 valuesof the logarithmic fits ( R 2 = 0 .
97 and R 2 = 0 .
compared to that of the linear fits ( R 2 = 0 .
66and R 2 = 0 . 43 ) suggests that the scaling of MABSplit is logarithmic (and therefore sublinear) withdataset size.
[Figure]
Appendix
Mean Estimation and Confidence Interval Constructions
In this section, we discuss the estimation of the means µ ft and construction of their confidenceintervals 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 amultinomial 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, write µ ft in terms of θ for the impurity metrics as
Gini
[Equation]
[Equation]
For a given impurity metric, let ∇ µ ft ( θ ) be the derivative of µ ft with respect to θ .
From the deltamethod,
[Equation]
where the CIs can be constructed accordingly.
These CIs are asymptotically valid as n ′ , n → ∞ .
Forother impurity metrics such as MSE, the CIs can be similarly derived by writing the correspondingµ ft in terms of θ .
[Table]
Comparison of baseline implementations and scikit-learn
In this section, we compare our re-implementation of common baselines to those in popular packagesto 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 as their correctness is implied by that of the
Random Forest model, as the Random Patches model consists of applying the Random Forest modelto subsampled data and features.
For classification, we compare our implementations on a the 20 newsgroups dataset filtered to twonewsgroups, alt.atheism and sci.space .
The dataset is embedded via TF-IDF and projectedonto their top 100 principal components, following standard practice [36].
The train-test split is thestandard one provided by scikit-learn .
For all classification problems, we average the predicted probabilities of each tree in the forest ("softvoting") as opposed to only allowing each tree to vote for a single class ("hard voting"), following theimplementation in scikit-learn [36].
For regression, we compare our implementations on the California Housing dataset, subsampledto 1,000 points as performing the regression on the full dataset of approximately 20,000 points iscomputationally prohibitive.
The train-test split is the standard one provided by scikit-learn [36].
Table 6 presents our results.
In all cases, our re-implemented baselines do not present a statisticallysignificant difference in performance than the models present in scikit-learn , which suggeststhat our re-implementations are correct.
Performance is measured over 20 random seeds to computeconfidence intervals.
Profiles
In this work, we focused on the reducing the runtime at the algorithmic level, i.e., reducing thecomplexity of computing the best feature-threshold split.
In this section, we justify this choice bydemonstrating that most of the time spent in our re-implementation of the baseline algorithms is spentin computing the best feature-threshold split.
Figure 2 demonstrates the wall-clock time spent inside various functions when fitting a Random
Forest classifier on two subsets of the MNIST dataset of sizes 5,000 and 10,000.
Most of the timeis spent inside the computation of the best feature-threshold split, which scales approximately asdataset size and motivates our work’s focus on improving the performance of the split-identificationsubroutine.
When using MABSplit, the time spent to identify the best feature-threshold split isreduced significantly (Figure 3).
Figure 4 also contains an example callgraph demonstrating callers and callees for the fitting procedureof a Random Forest, for easier interpretation of Figures 2 and 3.
[Figure]
All times are in nanoseconds.
When increasing the dataset size, the overhead spent outside of 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 aboutseconds to about 61 seconds 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 featurethreshold split.
This profile was generated with cProfile and visualized
[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.3.1 (Monterey) with a 2.4 GHz 8-Core Intel Core i9 processor,64 GB 2667 MHz DDR4 RAM, and an Intel UHD Graphics 630 1536 MB Graphics Card.
In allexperiments, trees of depth between 3 and 5 were trained.
6.1 Runtime Experiments
For the classification runtime experiments presented in Table 1, all models used the full MNIST traindataset of 60,000 images and test performance was measured on the 10,000 standard test images.
Performance was averaged over 5 random seeds.
For the regression runtime experiments presented in Table 2, all models used data generatedby scikit-learn.datasets.make_regression , with n_samples=200000, n_features=50 ,and n_informative=3 .
The first 160,000 datapoints were allocated to the training set and theremaining 40,000 to the test set; no validation set was constructed as no hyperparameter tuning wasperformed.
The performance of the Random Forest model was averaged over 5 random seeds and theperformances of the ExtraTrees and RandomPatches models were averaged over 20 random seeds.
6.2 Budget Experiments
For the classification budget experiments presented in Table 3, all models used the MNIST datasetsubsampled to the first 10,000 points for training data, whereas performance was still measured onthe full test dataset of 10,000 points.
The budget (number of histogram insertions permitted) for eachmodel was 784 ∗ 10000 ∗ 1 . 3 = 10192000 .
The performance of each model was averaged over 5random seeds.
For the regression budget experiments presented in Table 4, all models used data generatedby scikit-learn.datasets.make_regression , with n_samples=200000, n_features=50 ,and n_informative=3 .
The first 160,000 datapoints were allocated to the training set and theremaining 40,000 to the test set; no validation set was constructed as no hyperparameter tuningwas performed.
The budgets of Random Forest, ExtraTrees, and Random Patches models were setto 24 , 000 , 000 , 1 , 200 , 000 , and 24 , 000 , 000 respectively.
The performance of the Random Forestmodel was averaged over 5 random seeds and the performances of the ExtraTrees and RandomPatchesmodels were averaged over 20 random seeds.
6.3 Stability Experiments
Two measures for calculating feature importance are used in Table 5.
The first is the out-of-bag
Permutation Importance (OOB PI).
For a feature f , the OOB PI is calculated by measuring thedifference between the trained model’s out-of-bag error on the original data with its out-of-bag errorof all the data with all out-of-bag datapoints’ f values shuffled.
The second measure of featureimportance is the Mean Decrease in Impurity (MDI).
The MDI for a feature f is the average decreasein 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 areselected and the stability of these k features is measured via standard stability formula [33].
The stability comparison experiments are conducted on various models with the settingsdescribed in Table 5.
The Random Classification dataset is generated via scikit-learn ’sdatasets.make_classification function with n_samples=10000, n_features=60 ,and n_informative=5 .
The Random Regression dataset is generated by scikit-learn ’sdatasets.make_regression with n_samples=10000, n_features=100 , andn_informative=5 .
Limitations7.1 Theoretical Limitations
Crucial to the success of MABSplit are the assumptions described before and after Theorem 1.
Inparticular, we assume that their is reasonable heterogeneity amongst the true impurities of differentfeature-value splits.
Such assumptions are common in the literature and have been validated on manyreal-world datasets [4, 48, 6, 43, 3, 5].
We also note that the assumptions that each CI scales as (cid:113) log1 /δ n ′ may be violated when using certainimpurity metrics.
For example, the derivative of the entropy impurity criterion with respect to somep k approaches ∞ when p k → 0 .
In this case, we cannot apply the delta method from Appendix 3 tocompute finite CI intervals that scale in the way we require.
In such settings, it may be necessary tocompute the confidence intervals in other ways, e.g., following [35] or [7].
We note that in the worst case, even when all assumptions are violated, MABSplit is never worse thanthe brute-force algorithm in terms of sample complexity.
In the worst case, it is a batched version ofthe brute-force algorithm.
7.2 Practical Limitations
We note that MABSplit may perform worse than brute-force node-splitting on very small datasets,where the overhead of sampling the data in batches outweighs any potential benefits in samplecomplexity.
In this work, we avoided a direct runtime comparison with scikit-learn because scikit-learnutilizes a number of low-level implementation optimizations that would make the comparison unfair.
In order for practitioners to take full advantage of MABSplit, however, it may be necessary toimplement MABSplit within the scikit-learn library.
In doing so, it may be possible that MABSplit makes it difficult or impossible to use existingoptimizations in the scikit-learn library.
An example of this is vectorization: because thebrute-force node-splitting algorithm queries the data in a predictable way, each datapoint can bequeried more quickly than in MABSplit.
Despite MABSplit’s advantages in sample complexity, thedisadvantages of being unable to use implementation optimizations like vectorization may outweight
MABSplit’s benefits.
Many of these risks may be ameliorated by int.
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.
Weleave 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.
Figure 5 demonstratesthe performance of MABSplit, both in wall-clock training time and sample complexity, for varioussubset sizes of MNIST.
Our results that RF+MABSplit outperforms the standard RF algorithm, inboth sample complexity and wall-clock time, when the dataset size exceeds approximately 1100datapoints.
However, we also note that the main use case for MABSplit is when the data size is largeand it is computationally challenging to run standard forest-based algorithms.
[Figure]