Adaptive Split Balancing for Optimal Random Forest

We focus on the Lipschitz class here, with the Hölder class explored in Section 3. In the following, we illustrate how the auxiliary randomness introduced by the widely-used random feature selection affects the approximation error of tree models.

For any leaf $L\subseteq[0,1]^{d}$, let $\mathrm{diam}(L):=\sup_{\mathbf{x},\mathbf{x}^{\prime}\in L}\|\mathbf{x}-\mathbf{x}^{\prime}\|$ be its diameter. Under the Lipschitz condition, the approximation error can be controlled by the leaves’ diameters: $R_{2}\leq L_{0}^{2}\mathbb{E}_{\mathbf{x}}[\mathbb{E}_{\xi}[\mathrm{diam}^{2}(L(\mathbf{x},\xi))]].$ Therefore, it suffices to obtain an upper bound for the diameters. Although the Breiman’s algorithm selects up to $\mbox{mtry}\leq d$ features at random, we find it worthwhile to study the special case with $\mbox{mtry}=1$, as seen in centered and median forests. Recall that in the centered forest, a splitting direction is randomly selected (with probability $1/d$) for each split, and the splitting location is chosen as the center point of the parent node. For the moment, let $B=\infty$, meaning the forest is the ensemble of infinitely many trees. Suppose that each terminal leaf has been split for $\mathcal{M}$ times, and consider a sequence of $\mathcal{M}$ consecutive leaves containing $\mathbf{x}$, with the smallest (terminal) leaf denoted as $L(\mathbf{x},\xi)$.

For each $m\leq\mathcal{M}$ and $j\leq d$, let $\delta_{j,m}(\mathbf{x},\xi)=1$ if the $m$-th split is performed along the $j$-th coordinate, and $\delta_{j,m}(\mathbf{x},\xi)=0$ otherwise. Here, $\mathbb{P}_{\xi}$ and $\mathbb{E}_{\xi}$ represent the corresponding probability measure and the expectation taken with respect to $\xi$, respectively. For any given $j\leq d$, the sequence $(\delta_{j,m}(\mathbf{x},\xi))_{m=1}^{\mathcal{M}}$ is i.i.d., with $\mathbb{P}_{\xi}(\delta_{j,m}(\mathbf{x},\xi)=1)=1/d$ for centered forest. Let $\mathrm{diam}_{j}(L(\mathbf{x},\xi))$ be the length of the longest segment parallel to the $j$-th axis that is a subset of $L(\mathbf{x},\xi)$. Then,

The discrepancy introduced by the strict inequality (i),“Jensen gap”, stems from the variation in the quantity $c_{j}(\mathbf{x},\xi):=\sum_{m=1}^{\mathcal{M}}\delta_{j,m}(\mathbf{x},\xi)$, representing the count of splits along the $j$-th direction induced by the auxiliary randomness $\xi$. This discrepancy results in a relatively large approximation error $R_{2}\asymp(1-1/(2d))^{2\mathcal{M}}$ (up to logarithmic terms). By selecting an optimal $\mathcal{M}$ (or $k$) that strikes a balance between stochastic and approximation errors, the centered forest yields an overall IMSE of the size $N^{\frac{-2\log(1-1/(2d))}{2\log(1-1/(2d))-\log(2)}}$ – which is not minimax optimal for Lipschitz functions as long as $d>1$, as depicted in Figure 1.

The sub-optimality stems from the forests’ excessive reliance on auxiliary randomness, leading to a significant number of redundant and inefficient splits. When splitting directions are chosen randomly, there is a non-negligible probability that some directions are overly selected while others are scarcely chosen. Consequently, terminal leaves tend to be excessively wide in some directions and overly narrow in others. For centered and median forests, this long and narrow leaf structure is solely due to auxiliary randomness and is unrelated to the data. This prevalence of long and narrow leaves increases the expected leaf diameter, resulting in a significant approximation error that deviates from the optimal one. Averaging across the trees stabilizes the forest, but to attain minimax optimal rates, controlling the number of splits at the tree level is essential to mitigate the Jensen gap, as shown in (2.2).

Instead of selecting splitting directions randomly, we adopt a permutation-based and more balanced approach. By ensuring a sufficiently large $c_{j}(\mathbf{x},\xi)$ for each $j\leq d$ and reducing dependence on auxiliary randomness $\xi$, we achieve an approximation error of $O((N/k)^{-2/d})$, as shown in Lemma 2.1. This upper bound mirrors the right-hand side of (2.2) (whereas centered or median forests cannot improve the left-hand side) when $\mathcal{M}\approx\log_{2}(N/k)$, resulting in a minimax optimal rate for the overall IMSE when $\mathcal{M}$ (or $k$) is chosen to balance the approximation and the stochastic error.

In order to reduce the large approximation error caused by auxiliary randomness, we propose a permutation-based splitting method. Each time a leaf is split, we randomly select a direction from one of the sides that has been split the least number of times. Moreover, the splitting directions are chosen in a balanced fashion – we have to split once in each direction on any given tree path before proceeding to the next round.

To further enhance the practical performance, we employ data-dependent splitting rules of the “honest forests” approach that are dependent on both $\mathbf{X}_{i}$ and $Y_{i}$. This flexibility is particularly valuable when the local smoothness level varies in different directions and locations. For each tree, we partition the samples into two sets, denoted by $\mathcal{I}$ and $\mathcal{J}$. Additionally, we impose constraints on the child node $\alpha$-fraction and terminal leaf size $k$. With $\alpha\in(0,0.5]$ and $k\in\mathbb{N}$, we require the following conditions to hold: (a) each child node contains at least an $\alpha$-fraction of observations within the parent node, and (b) the number of observations within terminal leaves is between $k$ and $2k-1$. The splitting locations are then determined to minimize the empirical mean squared error (or maximize the impurity gain) within each parent node, selected from the set of points satisfying the above conditions. For more details, refer to Algorithm 1 and an illustration in Figure 2.

Any $\alpha$ controls the balance in leaf lengths, and any $\alpha>0$ prevents greedy splits near the endpoints of the parent node. For any $\alpha<0.5$, while the terminal leaves may still vary in width, they will be more balanced than before, hence the term “adaptive split-balancing.” Unlike centered and median forests, this time, the long and narrow leaf structure depends on the data. This distinction sets it apart from methods with data-independent splitting rules (e.g., [38, 40, 21]), allowing us to improve empirical performance, particularly when covariates have distinct local effects on the outcome.

While Algorithm 1 selects splitting points data-dependently and can adaptively learn local smoothness levels for different directions, it still encounters challenges in the presence of certain sparse structures. Notably, even if some covariates are entirely independent of the outcome, Algorithm 1 may still make splits along such directions. These splits are redundant, as they do not contribute to reducing the approximation error but increase the stochastic error, due to the smaller sample sizes in the resulting child nodes.

To address sparse situations, we introduce a tuning parameter, $\text{mtry}\in\{1,\dots,d\}$, representing the number of candidate directions for each split. Unlike existing methods, we select candidate directions in a balanced fashion instead of randomly. In each round, we initiate a collection of $d$ index sets, each containing mtry distinct directions, ensuring each direction appears in exactly mtry sets. At each split, we randomly select one unused index set for the current round. Alternatively, we can achieve the same result by shuffling the directions before each round and setting the index sets as $\{1,2,\dots,\text{mtry}\}$, $\{2,3,\dots,\text{mtry}+1\}$, etc., in a random order; see Figure 3. The selected index set provides the candidate splitting directions for the current node. Details are provided in Algorithm 2.

When the true conditional mean function is sparse, setting a relatively large mtry avoids splits on redundant directions. Conversely, when all directions contribute similarly, a smaller mtry prevents inefficient splits due to sample randomness. However, existing methods select candidate sets randomly, causing some directions to be frequently considered as candidate directions. This randomness can reduce approximation power by neglecting important features. In contrast, our balanced approach ensures all features are equally considered for selection at the tree level, making all the trees perform similarly and leading to a smaller estimation error for the forest. Due to the Jensen gap, the error of the averaged forest mainly stems from the poorer-performing trees, as the approximation power non-linearly depends on the number of splits in important features. The tuning parameters $\alpha$ (location) and mtry (direction) control the “greediness” of each split during the tree-growing procedure. Greedy splits, characterized by a small $\alpha$ and a large mtry, lead to better adaptability but are sensitive to noise, making them sub-optimal under the Lipschitz or Hölder class. Conversely, non-greedy splits, which use a larger $\alpha$ and a smaller mtry, achieve minimax optimal results under smooth functional classes but are non-adaptive to complex (e.g., sparse) scenarios.

For simplicity, we consider uniformly distributed $\mathbf{X}$ with support $[0,1]^{d}$, as seen in [1, 23, 5, 31, 10, 17, 42, 19, 11, 36, 35, 46]. We first highlight the advantages of the balanced approach.

The right most term of the above expression dominates the rest and determines the expected size of the diameter $(\lfloor wN\rfloor/(2k-1))^{-\frac{r\log(1-\alpha)}{d\log(\alpha)}}\asymp(w/2)^{-\frac{r\log(1-\alpha)}{d\log(\alpha)}}\cdot(N/k)^{-\frac{r\log(1-\alpha)}{d\log(\alpha)}}\asymp(k/N)^{\frac{r\log(1-\alpha)}{d\log(\alpha)}}$. Recall that the sample-splitting fraction $\omega\in(0,1]$ remains constant and that the minimum node size, $k$ is such that $k\leq\lfloor\omega N\rfloor$. According to Lemma 2.1, the proposed forests’ approximation error is at the order of $(k/N)^{\frac{2\log(1-\alpha)}{d\log(\alpha)}}$. Our balanced median forest, with $\alpha=0.5$ yields a rate of $(k/N)^{\frac{2}{d}}$, whereas, a standard median forest with random feature selection has a strictly larger rate of $(k/N)^{2\log_{2}\left(\frac{2d}{2d-1}\right)}$ [31], whenever $d>1$; for $d=1$, the rates are the same (no need to choose a splitting direction). Centered forests achieve comparable upper bounds, which suggests the sub-optimality of auxiliary randomness.

We further assume the following standard condition for the noise variable.

The following theorem characterizes the IMSE of the proposed forest in Algorithm 1.

The issue of excess auxiliary randomness is inherent at the tree level, meaning that merely collecting more samples or further averaging does not address the optimality of the approximation error. Our approach to resolving this involves a new splitting criterion that has yielded minimax optimal rates.

For instance, split-balancing median forest is the first minimax optimal median forest. When setting $\alpha=0.5$, split-balancing random forest becomes new, balanced median forest, with the optimal minimax rate of $N^{-\frac{2}{d+2}}$. In contrast, existing median and center forests [31, 17, 5, 1] do not achieve the minimax optimal rate; see Table 1.1 and Figure 1. The conclusions of Theorem 2.2 remain valid for any constant $\alpha$ in the range $(0,0.5]$. Greedy splits (small $\alpha$) enhance adaptability by leveraging information from both $\mathbf{X}_{i}$ and $Y_{i}$, whereas non-greedy splits attain minimax optimality, necessitating a balance in practice. The rates are represented in Figure 4.

To capture the higher-order smoothness of the conditional mean function $m(\cdot)$, we propose to fit a $q$-th local polynomial regression within the leaves of the split-adaptive balanced forest of Section 2.2. We first introduce the polynomial basis of order $q\in\mathbb{N}$. For any $\mathbf{x}\in[0,1]^{d}$ and $j\in\{0,1,\dots,q\}$, let $\mathbf{g}_{j}(\mathbf{x}):=(\mathbf{x}^{\bm{\gamma}})_{\bm{\gamma}\in\mathcal{A}_{j}}\in\mathbb{R}^{d^{j}}$, where $\mathcal{A}_{j}:=\{\bm{\gamma}=(\bm{\gamma}_{1},\dots,\bm{\gamma}_{d})\in\mathbb{N}^{d}:|\bm{\gamma}|=j\}$. For instance, $\mathbf{g}_{0}(\mathbf{x})=1$, $\mathbf{g}_{1}(\mathbf{x})=(\mathbf{x}_{1},\mathbf{x}_{2},\dots,\mathbf{x}_{d})^{\top}$, and $\mathbf{g}_{2}(\mathbf{x})=(\mathbf{x}_{1}^{2},\mathbf{x}_{1}\mathbf{x}_{2},\dots,\mathbf{x}_{d}^{2})^{\top}$. Denote $\mathbf{G}(\mathbf{x}):=(\mathbf{g}_{0}(\mathbf{x}),\mathbf{g}_{1}(\mathbf{x})^{\top},\dots,\mathbf{g}_{q}(\mathbf{x})^{\top})^{\top}\in\mathbb{R}^{\bar{d}}$ as the $q$-th order polynomial basis, where $\bar{d}:=\sum_{j=0}^{q}d^{j}$. For any $\xi\in\Xi$, define the weights $\omega_{i}(\mathbf{x},\xi)$ as in (2.1). Using the training samples indexed by $\mathcal{I}$, we consider the weighted polynomial regression:

The $q$-th order local adaptive split balancing forest is proposed as

Unlike Algorithm 1, where local averages are used as tree predictions, our approach here involves conducting polynomial regressions within the terminal leaves. To optimize the behavior of the final polynomial regressions, we minimize the following during the leaf construction:

where $\widehat{Y}_{i}:=Y_{i}-G(\mathbf{X}_{i})^{\top}\widehat{\bm{\beta}}$ with $\widehat{\bm{\beta}}$ denoting the least squares estimate within the parent node $L$, and $\widetilde{Y}_{j}$ is the average of $\widehat{Y}_{i}$ within the child node $L_{j}$ for each $j\in\{1,2\}$. As in Section 2.2, we also require that both child nodes contain at least an $\alpha$-fraction of samples from the parent node. Additional specifics are outlined in Algorithm 3. Note that Algorithm 1 is a special case of Algorithm 3 with $q=0$. To further improve the forests’ empirical behavior in the presence of certain sparse structures, we also introduce a generalized sparse version of the local adaptive split balancing forest following; see Algorithm 4 in the Supplement for more details.

The following theorem characterizes the convergence rate of the local adaptive split balancing forest proposed in Algorithm 3.

When $\alpha=0.5$, the rate given by (3.3) is $N^{-\frac{2(q+\beta)}{d+2(q+\beta)}}$, which is minimax optimal for the Hölder class $\mathcal{H}^{q,\beta}$; see, e.g., [44]. This accomplishment represents the first random forest with IMSE reaching minimax optimality when $q>1$; refer to Table 1.1 and Figure 1.

Theorem 3.1 is built upon several key technical components. To mitigate the issue of multicollinearity during local polynomial regression within terminal leaves, we first demonstrate that $\sup_{\mathbf{x}\in[0,1]^{d},\xi\in\Xi}\kappa_{N}(\mathbf{x})=O_{p}(1)$, where $\kappa_{N}(\mathbf{x}):=[1-\mathbf{d}(\mathbf{x})^{\top}\mathbf{S}^{-1}(\mathbf{x})\mathbf{d}(\mathbf{x})]^{-1}$ assesses the degree of multicollinearity. Here, $\mathbf{d}(\mathbf{x}):=\sum_{i\in\mathcal{I}}\omega_{i}(\mathbf{x},\xi)\mathbf{G}_{-1}(\mathbf{X}_{i}-\mathbf{x})$, $\mathbf{S}(\mathbf{x}):=\sum_{i\in\mathcal{I}}\omega_{i}(\mathbf{x},\xi)\mathbf{G}_{-1}(\mathbf{X}_{i}-\mathbf{x})\mathbf{G}_{-1}(\mathbf{X}_{i}-\mathbf{x})^{\top}$, and $\mathbf{G}(\mathbf{x})=(1,\mathbf{G}_{-1}(\mathbf{x})^{\top})^{\top}$. Unlike [19], which imposes an upper bound on the random quantity $\kappa_{N}$ by assumption, we rigorously prove that this quantity is uniformly bounded above with high probability; see the equivalent form in Lemma S.4. Additionally, by employing localized polynomial regression, we reduce the approximation error to $O(\mathbb{E}_{\mathbf{x}}[\mathrm{diam}^{2(q+\beta)}(L(\mathbf{x},\xi))])$ under the Hölder class $\mathcal{H}^{q,\beta}$, as detailed in (LABEL:consistency:R1) of the Supplement. Together with our proposed permutation-based splitting rule, previous results on leaf diameter (Lemma 2.1) ensure a sufficiently small rate.

[38] achieved minimax optimal IMSE for $q=1$ and $\beta\in(0,1/2]$ by leveraging forest structure and concentration across trees. However, for $\beta\in(1/2,1]$, their optimality is limited to interior points where $\mathbb{E}_{\xi}[\mathbb{E}_{\mathbf{x}^{\prime}}[\mathbf{x}^{\prime}\mid\mathbf{x}^{\prime}\in L(\mathbf{x},\xi)]]\approx\mathbf{x}$. Our approach broadens the range of Hölder smooth functions addressed, using polynomial approximation and split-balancing to eliminate the approximation error at the tree level.

Tessellation forest of [40] addressed this boundary issue, achieving minimax optimal IMSE for $q=1$ and any $\beta\in(0,1]$. Theorem 3.1 further highlights that the local split-balancing median forest (with $\alpha=0.5$) is minimax optimal not only for $q=1$ but also for any arbitrary $q>1$. This signifies a major breakthrough in attaining minimax optimality for higher-order smoothness levels.

Debiased Mondrian forests by [12] achieve minimax optimality for any $q\in\mathbb{N}$ and $\beta\in(0,1]$, in terms of point-wise MSE at interior points. However, their results do not yield an optimal rate for the IMSE, as their debiasing technique does not extend for boundary points. Additionally, [10] achieved a nearly optimal rates for any $q\in\mathbb{N}$ and $\beta\in(0,1]$ but their findings are limited to in-sample excess risk.

It’s important to note that while all previous works grow trees completely independently of the samples, our approach employs supervised splitting rules with $\alpha<0.5$ after tuning, significantly enhancing forest performance. Only [7, 19] considered data-dependent splitting rules and studied local linear forests under the special case $q=1$. However, [7] only demonstrated the consistency of their method, without providing any explicit rate of convergence. [19] provided asymptotic normal results at a given $\mathbf{x}\in[0,1]^{d}$. However, their established upper bound for the asymptotic variance is no faster than $N^{-(1+\frac{d\log(0.2)}{1.3\log(0.8)})^{-1}}\approx N^{-(1+5.55d)^{-1}}$, which is slow as the splitting directions are chosen randomly.

In the following, we extend our analysis to include uniform-type results for the estimation error of the forests. These are useful for our particular applications to causal inference and more broadly are of independent interest. As Algorithm 1 is a specific instance of Algorithm 3, we focus on presenting results for the latter.

To begin, we establish a uniform bound on the diameter of the leaves as follows.

The result in Lemma 3.2 holds for a single tree and also for an ensemble forest. In fact, obtaining uniform results for a single tree is relatively simple – it suffices to control the diameter of each terminal leaf with high probability and take the uniform bound over all the leaves, as the number of terminal leaves is at most $N/k$. However, such an approach is invalid for forests – the number of all terminal leaves in a forest grows with the number of trees $B$. Hence, such a method can be used only when $B$ is relatively small; however, in practice, we would like to set $B$ as large as possible unless constrained by computational limits. To obtain a uniform result for forests, we approximate all the possible leaves through a collection of rectangles.

Subsequently, we present a uniform upper bound for the estimation error of the forests.