Adaptive Split Balancing for Optimal Random Forest
Yuqian Zhang, Weijie Ji, Jelena Bradic
TL;DR
The paper introduces Adaptive Split Balancing Forests (ASBF) and Local ASBFs (L-ASBF) that replace conventional random feature-based splits with permutation-based, balanced splitting rules to reduce the Jensen gap and leaf diameters. By enforcing alpha-fraction constraints and balanced splitting across directions, and by integrating honest, data-dependent impurity measures, the methods achieve minimax-optimal IMSE rates over Lipschitz and Hölder smoothness classes, with L-ASBF extending to higher-order smoothness (q>1). Theoretical results establish IMSE and uniform convergence rates, and practical benefits are demonstrated through causal inference applications (ATE via AIPW and double ML) and extensive simulations/real-data experiments showing improved accuracy over baseline random forests and local linear forests. The findings highlight that controlled reduction of auxiliary randomness, together with local polynomial fitting, yields faster convergence and more reliable inference in moderate to high dimensions, with clear implications for causal estimation and structured prediction tasks.
Abstract
In this paper, we propose a new random forest algorithm that constructs the trees using a novel adaptive split-balancing method. Rather than relying on the widely-used random feature selection, we propose a permutation-based balanced splitting criterion. The adaptive split balancing forest (ASBF), achieves minimax optimality under the Lipschitz class. Its localized version, which fits local regressions at the leaf level, attains the minimax rate under the broad Hölder class $\mathcal{H}^{q,β}$ of problems for any $q\in\mathbb{N}$ and $β\in(0,1]$. We identify that over-reliance on auxiliary randomness in tree construction may compromise the approximation power of trees, leading to suboptimal results. Conversely, the proposed less random, permutation-based approach demonstrates optimality over a wide range of models. Although random forests are known to perform well empirically, their theoretical convergence rates are slow. Simplified versions that construct trees without data dependence offer faster rates but lack adaptability during tree growth. Our proposed method achieves optimality in simple, smooth scenarios while adaptively learning the tree structure from the data. Additionally, we establish uniform upper bounds and demonstrate that ASBF improves dimensionality dependence in average treatment effect estimation problems. Simulation studies and real-world applications demonstrate our methods' superior performance over existing random forests.