May 17, 2021
Building trustworthy machine learning algorithms for biological applications requires adapting to data heterogeneity from different sources, batches, distributions, or studies. We propose the Cross-Cluster Weighted Forest (CCWF), an ensembling approach that explicitly leverages heterogeneity in the feature distribution to produce more accurate and more generalizable predictors than the standard Random Forest in cases when data can be naturally clustered. CCWF generalizes the RF architecture to an outer unsupervised layer, supervised subtasks, and ensembling. Specifically it involves unsupervised clustering of the training data, fitting a Random Forest on each cluster, and combining the forests via stacked regression weights that reward cross-cluster generalizability. We provide a theoretical analysis of an analytically tractable forest model showing that cluster-based ensembling is asymptotically more accurate than training a single forest on the full data, with the gain driven by bias reduction. In simulations, we find that CCWF is robust across data-generating regimes and outcome models; furthermore, we explore the influence of data partitioning and ensemble weighting strategies on the benefits of our method. Finally, we apply our approach to cancer molecular profiling and gene expression datasets that are naturally divisible into clusters; in both simulations and real data examples, we illustrate that our approach outperforms classic Random Forest by margins of 30-40%, aligning with our theoretical results. Overall, we show that CCWF provides a statistically grounded prediction algorithm for data spanning multiple domains or sub-populations, a structure common in biological applications.
Natural clusters and batch effects are common in biological datasets pooled across studies, sites, instruments, or biological sub-populations. This heterogeneity poses a central challenge for trustworthy and generalizable machine learning: a model that is accurate on the full training distribution can fail unpredictably on a single sub-population or new domain [1]–[5]. A growing body of work addresses this phenomenon through mixed-effects models, mixture-of-experts architectures, partition-then-ensemble strategies, and clustered federated learning, all of which reflect the principle that learning across domains is most reliable when algorithms explicitly account for heterogeneity rather than averaging it away [6]–[18].
Random Forest (RF) represent a highly popular tabular data algorithm and are extensively used in genomic and biomedical applications [15], [19]–[24]. Theoretical work has characterized their consistency, convergence rates, and high-dimensional behavior [25]–[27]. Empirically, [15] showed that ensembles of RFs trained on heterogeneous datasets can outperform a single forest trained on all available data, aligning with the multi-study ensembling approach of [19], in which learners trained on separate studies are stacked to improve cross-study performance. These findings suggest that, for datasets with cluster-like heterogeneity, partitioning data before ensembling forests may improve over the classic RF algorithm.
We evaluate this hypothesis when the outcome model is shared across clusters and the primary heterogeneity lies in the covariate distribution. This setting extends the classic covariate shift paradigm by considering internal heterogeneity within the training distribution itself [28]–[30]. It is also closely related to domain generalization and distributionally robust learning, which seek predictors that maintain performance across heterogeneous training and test populations [31]–[33]. Throughout the paper, we distinguish between ‘estimated clusters’ obtained algorithmically and ‘true clusters’ arising from the data-generating process.
We introduce the Cross-Cluster Weighted Forest (CCWF), a learning approach designed to improve prediction performance by leveraging covariate heterogeneity. CCWF first clusters the training data using algorithms such as k-means, then trains separate RFs within each estimated cluster, and finally uses stacked regression to assign weights to these forests based on their stability and predictive accuracy across clusters [34]–[36]. This strategy reduces within-cluster heterogeneity while preserving cross-cluster learner diversity, following the ensemble learning principle that combining diverse learners can improve generalization [8], [37]–[39]. We study both the theoretical and empirical properties of CCWF, demonstrating its ability to outperform standard RF in simulated settings and real biological datasets with natural cluster structure.
Our approach builds on prior work using clustering to partition data before training ensemble learners, including methods based on simultaneous co-clustering and model fitting or clustering as a preprocessing step [13], [14], [16]. CCWF differs by using generic clustering algorithms to identify substructure and cross-cluster weighting to favor learners that generalize across partitions. It also differs from prior work linking RFs to clustering, which typically uses RF itself as a clustering method rather than using external cluster structure to form weighted forest ensembles [40]–[42]. Finally, it differs from [43], who develop an ensemble of RF’s method for data with independent groups that each have a known dependence structure, such as repeated measures/longitudinal data. The method’s primary goal is to achieve better precision by accounting for this known correlation structure within groups.
The remainder of the paper is organized as follows. We first present a theoretical analysis of CCWF using an analytically tractable formulation to study how bias–variance tradeoffs determine when cluster-based ensembling improves prediction on heterogeneous data. We then validate CCWF empirically through experiments on simulated datasets with cluster structure, examining the impact of each algorithmic step and assessing robustness across data-generating scenarios. Finally, we apply CCWF to real gene expression and clinical datasets with feature-distribution heterogeneity, showing improved performance relative to classic RF.
CCWF is a structural generalization of the RF architecture using an outer unsupervised layer, supervised subtasks, and ensembling. While in theory it can be compared to any tree-based algorithm such as gradient-boosted trees or deep tabular models, we benchmark against baselines where the supervised subtasks are the same in our algorithm and the comparators. In this way, we are able to provide direct evidence on the value of the architecture. Our baselines therefore hold the learner class fixed and vary only the partitioning and ensembling procedure, isolating whether covariate-aware clustering and weighting improve over standard RF training.
The CCWF approach is outlined in Algorithm 1. Standard RF averages equally weighted trees trained on bootstrap samples. Our approach introduces a non-random data partitioning step in which we break the data into \(k\) subsections, potentially improving performance when the data can be naturally clustered. We use k-means for scalability, though the framework permits other clustering methods [44]–[46]. We train RFs on each cluster and combine in-sample training predictions using nonnegative ridge stacking; weights are learned only from training data, and generalization is evaluated on held-out test data [15], [19], [34], [47].
In this section, we examine the asymptotic risk of CCWF compared to the classic RF algorithm using an analytical setup that approximates the full algorithm as closely as is possible given current theoretical results on Random Forest [48]. We show that bias reduction is the primary mechanism through which the general cluster-based ensembling framework produces improvements for forest-based learners.
We denote the training data set as \(\mathcal{D}_n=(\mathbb{X},\mathbb{Y})\in\mathbb{R}^{n\times S}\times\mathbb{R}^n\), where \(S\) is the number of covariates and \(n\) is the training sample size. The training dataset can be further expressed as \((\mathbb{X},\mathbb{Y})=\left(\left[\mathbb{X}_1 \ldots \mathbb{X}_k\right]^T,\left[\mathbb{Y}_1 \ldots \mathbb{Y}_k\right]\right)\), where \(k\) is the number of clusters, \(\mathbb{X}_b\in\mathbb{R}^{n_b\times S}\), and \(\mathbb{Y}_b\in\mathbb{R}^{n_b}\) denote the covariates and outcomes for cluster \(b\). We specify that \(n_b=n/k\) and that the training covariate set consists of \(k\) equally sized, non-overlapping, uniform clusters spanning the \([0,1]^S\) hypercube: \(\mathbb{X}_1\stackrel{\rm i.i.d.}{\sim}\left[\mathrm{U}\left(0,\frac{1}{k}\right)\right]^S\), \(\mathbb{X}_2\stackrel{\rm i.i.d.}{\sim}\left[\mathrm{U}\left(\frac{1}{k},\frac{2}{k}\right)\right]^S\), \(\ldots\), \(\mathbb{X}_k\stackrel{\rm i.i.d.}{\sim}\left[\mathrm{U}\left(\frac{k-1}{k},1\right)\right]^S\). Equivalently, each training sample is \((\mathbf{x}_i,Y_i)\), where \(A_i\sim\mathrm{Categorical}\left(\frac{1}{k},\ldots,\frac{1}{k}\right)\) denotes the cluster assignment, \(\mathbf{x}_i\mid A_i=b\) is uniform on \(\left[\frac{b-1}{k},\frac{b}{k}\right]^S\), and \(Y_i=f(\mathbf{x}_i)\). We make no parametric assumptions on \(f(\cdot)\). Test point \((\mathbf{x}_\star,Y_\star)\) follows the same distribution as the training set.
We analyze the Centered Random Forest (CRF) of [49], which shares the recursive axis-aligned-split structure of standard RFs but uses coordinate-uniform splits independent of the response. CRFs are the standard analytical surrogate in the RF theory literature [24], [50] because conclusions about bias and variance scaling carry over qualitatively to standard RFs while remaining tractable. We later evaluate whether the bias-dominance prediction holds for standard RFs (Section 4). Full definitions are in Appendix 7.1.
The Ensemble method representing an analytically tractable formulation of CCWF functions by training a CRF on \((\mathbb{X}_b, \mathbb{Y}_b)\) for each cluster \(b\), yielding predictions \(\hat{Y}_b(\mathbf{x}_\star; \theta, \mathcal{D}_n)\) which are then summed to form the final Ensemble prediction: \(\hat{Y}_E(\mathbf{x}_\star) = \sum_{b=1}^k \hat{Y}_b(\mathbf{x}_\star)\). Because CRFs do not extrapolate beyond their training data, \(\hat{Y}_b(\mathbf{x}_\star) = 0\) if the cluster assignment of \(\mathbf{x}_\star \neq b\) and thus \(\hat{Y}_E(\mathbf{x}_\star)\) is in fact the prediction of the single forest whose training data range contains \(\mathbf{x}_\star\). This construction isolates the mechanism that CCWF is designed to exploit: when clusters are well separated and learners do not extrapolate outside their training ranges, cluster-specific forests specialize to the relevant region of covariate space. As a baseline, we train a single CRF on the full dataset termed the Merged, with prediction \(\hat{Y}_M(\mathbf{x}_\star)\). The aim of this section is to derive the limit of \(\frac{\textit{RMSE}(\hat{Y}_E)}{\textit{RMSE}(\hat{Y}_M)} = \frac{\sqrt{\mathbb{B}_E^2 + \mathbb{V}_E}}{\sqrt{\mathbb{B}_M^2 + \mathbb{V}_M}}\), where \(\mathbb{B}, \mathbb{V}\) are the bias and variance of the Ensemble and Merged.
In Lemma 1 presented in the Appendix, we show that the Ensemble and Merged predictions can be represented as weighted averages of the training outcome values, with the choice of weights as the differentiating factor between the methods. To develop theoretical results regarding the performance of these two approaches, we build on [50] in characterizing the asymptotic behavior of the bias and variance components of the MSE for CRFs. We extend upon these results to allow for clusters in the training and test data and to remove any requirement on the functional form of the outcome model \(f(\mathbf{x})\). We note that the latter represents a significant step forward upon previous analytical work on Random Forests, which typically assume a linear or piecewise linear outcome model.
We begin by stating supplementary results that support the primary performance comparison between the Ensemble and the Merged in Theorem 1 below. First, we can show that the variance terms \(\mathbb{V}_E\) and \(\mathbb{V}_M \to 0\) at a rate of \(O\left(\frac{1}{n}\right)\), the formal statement and proof of which we provide in Lemma 2 in the Appendix. This follows directly from the results in [[50], Proposition 2] and its corresponding proof. We further show in Remark 3 in the Appendix that the bias terms \(\mathbb{B}_E\) and \(\mathbb{B}_M\) converge more slowly than \(O(1/n)\) under mild assumptions on \(S\) and \(c_n\). Therefore, we can characterize the limit of \(\frac{\textit{RMSE}(\hat{Y}_E(\mathbf{x}_{\star}))}{\textit{RMSE}(\hat{Y}_M(\mathbf{x}_{\star}))}\) by the limit of \(\frac{\mathbb{B}_E}{\mathbb{B_M}}\); that is, the bias terms dominate the RMSE as \(n \to \infty\). We can then derive the limit of the performance ratio between the Ensemble and the Merged as presented in Theorem 1 below.
Theorem 1. Given a training set with \(n\) samples, \(S\) covariates, and \(k \geq 2\) equally-sized uniform clusters \(\{(\mathbf{x}_i,Y_i): i=1,\ldots,n\}\), where \(A_i\sim\mathrm{Categorical}\left(\frac{1}{k},\ldots,\frac{1}{k}\right)\) denotes the cluster assignment, \(\mathbf{x}_i \mid A_i=b\) is drawn uniformly from the hypercube \(\left[\frac{b-1}{k},\frac{b}{k}\right]^S\), and \(Y_i=f(\mathbf{x}_i)\), we can represent the predictions of the Ensemble and the Merged on test point \((\mathbf{x}_\star,Y_\star)\) drawn from the same distribution, as \(\hat{Y}_E\) and \(\hat{Y}_M\) respectively. Then, \[\begin{align} \lim_{n \to \infty} \frac{\textit{RMSE}(\hat{Y}_E(\mathbf{x}_{\star}))}{\textit{RMSE}(\hat{Y}_M(\mathbf{x}_{\star}))} &= \frac{1}{\sqrt{2}} . \end{align}\]
Theorem 1 demonstrates that the asymptotic performance improvement of the Ensemble over the Merged is independent of the number of true clusters \(k\) for \(k \geq 2\). To understand this phenomenon, consider that ensembling involves a trade-off: it reduces the number of samples used to train each learner (decreasing efficiency compared to the Merged), but it also allows each learner to focus on regions of the covariate space that better reflect the distribution of the test point, thereby improving accuracy. The proof of Theorem 1 decomposes the ratio \(\frac{\textit{RMSE}(\hat{Y}_E(\mathbf{x}{\star}))}{\textit{RMSE}(\hat{Y}_M(\mathbf{x}{\star}))}\) into two components that capture this trade-off. It shows that when these components are multiplied, terms involving \(k\) cancel out, resulting in a final ratio of \(\frac{1}{\sqrt{2}}\). Although each learner in the Ensemble is trained on \(\frac{1}{k}\) times the number of samples as the Merged, the benefit of the Ensemble approach lies in its ability to more effectively parse out the heterogeneity introduced by clusters from the true covariate-outcome relationship. This advantage outweighs the loss in efficiency due to the reduced sample size as \(n \to \infty\). This result provides a theoretical characterization of how covariate heterogeneity can affect forest ensembling without imposing a parametric form on \(f(\mathbf{x})\). We verify that the ratio in Theorem 1 holds exactly in a simulation study presented in Appendix section 8.1.
This section evaluates ensembling approaches on simulated data designed to mimic real genomic datasets. We first examine the importance of data partitioning on overall results, highlighting the benefits of clustering algorithms. Next, we analyze the impact of ensemble weighting strategies, comparing stacked regression weights to simple averaging. All results are averaged over 250 simulations per scenario (run on 4 CPU cores), with 95% confidence bands computed as mean \(\pm 1.96\) standard errors. We interchangeably refer to CCWF as the Cluster method for clarity, as this naming aligns more naturally with the other methods we compare it to. We then evaluate it alongside three variations of Algorithm 1 that ablate various steps, defined as follows. The Random method uses equally sized random partitions instead of clusters, keeping the number of total learners equal to \(k\). The Multi method trains RFs on true clusters (known in simulation but unavailable in practice) and combines them using stacking weights, thus named for its equivalence to multi-study stacking [19]. The Merged method serves as the baseline, training a single RF on the entire dataset. To ensure fair comparison, all methods are trained using the same number of total trees in order to isolate each algorithmic step’s contribution to the performance. We use 100 trees per component forest for ensemble methods and \(100k\) trees for the Merged.
We employ two methods to generate clustered datasets for our simulations: multivariate Gaussian mixture models using the R package clusterGeneration [51], and plasmodes more closely mimicking authentic biological clusters using the ‘monte’ function in the R package fungible [52]. For each simulation per setting, we generate training data with ntrain clusters (baseline: 5) and ntest test datasets with 2 clusters each (baseline: 5), all containing ncoef covariates
(baseline: 20). We create response models by randomly selecting 10 covariates, drawing coefficients uniformly from \([-5,-0.5] \cup [0.5, 5]\), adding Gaussian noise, and considering non-linear relationships. Each cluster
contains 500 samples at baseline, with between-cluster coefficient perturbation drawn uniformly from [0, 0.25]. Cluster separation is set at median values, allowing some overlap while remaining distinguishable. This setup enables comprehensive evaluation
of feature distribution heterogeneity effects on ensemble performance across various scenarios. All modifications to this baseline strategy are outlined in the relevant results sections.
We train Random Forests using the R package randomForest [53] with the number of trees as described above. We
keep ‘mtry’ and ‘nodesize’ hyperparameters, representing the number of predictors randomly sampled as candidates at each split and the minimum size of terminal nodes respectively, at their default values of \(\sqrt{S}\) and
\(1\). For stacked regression weights, we tune the ridge regression regularization parameter through 5-fold cross-validation within the training set. In experiments where we choose one optimal value of \(k\) for k-means rather than display results across a variety of \(k\) values, we choose this optimal value through maximizing the silhouette score over candidate \(k\)’s ranging from 2 to \((ntrain * 500) // 30\).
For fixed total tree count \(T\), CCWF with k-means can be more efficient than the Merged baseline. Stacking and k-means add minimal overhead in comparison to forest training. CCWF trains \(k\) forests on \(n/k\) samples, reducing total training complexity to \(O(\frac{TSn}{k} \log \frac{n}{k}) < O(TSn \log n)\).
We investigate how data partitioning methods affect ensemble accuracy by comparing the approaches outlined in the previous section. Both the Cluster and Random methods train component forests on the same number of partitions, but differ in how those partitions are determined. To keep notation consistent, we refer to \(k\) as the parameter used to train k-means, reflecting the number of partitions of the data also used for the Cluster and Random approaches. We explore the impact of varying \(k\) and the data-generating mechanism across six scenarios in Figure 2.
We see that the Cluster and Random outperform the Multi and Merged when \(k\) exceeds the true number of clusters, with improvements over 20% at optimal values. The Multi, which uses true clusters, shows inconsistent performance relative to Merged. The Cluster and Random are robust across outcome models and data distributions; however, the Cluster generally outperforms the Random, suggesting k-means clustering is preferable to random partitioning. Both methods improve as \(k\) increases, plateauing when each forest trains on 35–40 out of 2500 observations. Correspondingly, average tree depth decreases, indicating that high \(k\) values effectively ensemble locally weak learners akin to local smoothing approaches. This supposition is explored further in Appendix section 9.3. Note that the theoretical result equates to a \(29\%\) improvement of ensembling over merging in a more idealized setting, which is squarely within the range of improvements we see empirically. For further comparison of how the theoretical results translate in practice, we compare the Merged and Multi approaches and confirm that bias reduction remains the primary mechanism for performance improvements of the Multi; a full analysis can be found in Appendix section 8.2.
The differences between the Cluster and the Random reveal key advantages of clustering as a partitioning strategy. Increasing the number of partitions generally improves ensemble performance, with the Multi often outperforming other methods when \(k\) is smaller than the true number of clusters. Clustering allows learners to be less correlated and more dissimilar, a property effective in other ensemble approaches. The improvement of the Cluster over the Merged suggests that, in these heterogeneous settings, clustering creates useful partitions for training component trees. While clustering may reduce accuracy in high-density regions, this loss does not outweigh the benefits.
A central and initially counterintuitive finding: even when \(k\) matches the true cluster count, k-means partitions do not recover the true clusters, and yet they yield better ensembles than partitioning by ground truth. We hypothesize the mechanism: k-means minimizes within-cluster covariate range, tightening the local function approximation each forest must learn. True clusters, defined by data-generating provenance, may not minimize covariate range and often span the same covariate region as other clusters. Empirically, k-means partitioning achieves a lower average covariate range (3.17 CI:(3.08, 3.26)) than random partitioning (4.22 CI:(4.13, 4.31)), while true-cluster partitioning gives a much larger range (7.33 CI:(7.21, 7.45)). By restricting each forest to a smaller subset of the covariate space, learners can more precisely discern localized patterns, especially in distribution tails. The same phenomenon persists in multi-study settings (Appendix 9.2). These results suggest that for downstream RF ensembles, covariate-range minimization may be more useful than provenance recovery.
Silhouette analysis offers a practical diagnostic: Cluster outperforms Merged only when silhouette analysis selects an optimal \(k > 1\), regardless of whether true clusters exist or how separated they are. Similar behavior holds for other clustering methods, so the unsupervised clustering step itself is a reasonable indicator of downstream gains.
We next evaluate method robustness across varying dataset characteristics, keeping \(k\) at the optimal value determined by Silhouette analysis. Using a quadratic outcome model and non-Gaussian cluster generation, we examine how prediction performance is affected by covariate-outcome signal strength (by varying coefficient norms), the number of true clusters in the training set, and overall sample size.
Figure 3a shows that as covariate signal strength increases (achieved by increasing the norm of the coefficient norm in the simulated outcome model), all ensembling methods improve over Merged, with performance gains plateauing at 20-60%. The Cluster consistently outperforms others, especially at higher coefficient norms. The Multi’s limited improvement suggests that reducing within-cluster variation is more effective than using true clusters for discerning covariate-outcome relationships. Figure 3b demonstrates that as true cluster count increases (keeping total sample size at 2500), the Cluster remains superior, while the Multi approaches the Random’s performance. This occurs because Multi’s forest count increases with the number of true clusters, aligning with previous observations that more ensemble members improve overall accuracy even as individual model complexity decreases. The Cluster’s consistent advantage highlights its benefits regardless of the true data composition. Figure 3c shows that as sample size increases from 500 to 5000, the performance gap between the Cluster and others widens. While the Multi becomes less distinguishable from the Merged (as the Merged accuracy increases proportionally with sample size), the Cluster and Random continue to improve upon the Merged since the optimal number of partitions \(k\) also increases proportionally with sample size.
| No.of Clusters \(K\) | 2 | 5 | 10 | 20 | 30 | 50 | 70 | 80 |
|---|---|---|---|---|---|---|---|---|
| Simple Averaging | 11.93 | 29.97 | 43.60 | 59.63 | 68.17 | 79.09 | 86.65 | 89.48 |
| Stacked Regression | -9.88 | -13.35 | -19.73 | -25.21 | -27.73 | -29.99 | -31.96 | -32.27 |
Our next simulation explores the impact of ensemble weighting on prediction performance, focusing on the advantages of cross-cluster stacked regression weights over other approaches. Table 1 compares the percent change in average RMSE relative to the Merged for ensembles using simple averaging versus stacked regression weights. As \(k\) increases to its optimal value around 80, stacking-based ensembles continuously improve while equal weighting-based ensembles show declining performance. The best performance for equal weighting occurs at \(k=2\), at which it is still significantly worse than the Merged. This indicates that not all clusters from k-means produce accurate forest predictors, but certain combinations greatly outperform the Merged. Multi-cluster stacking up-weights the cluster-specific learners that display the best cross-cluster prediction ability. Strong cross-cluster performance suggests better learning of the true covariate-outcome relationship and greater generalizability, as clusters are designed for maximal distributional separation.
To explore classifier performance on real biological data, we apply our methods to data on tumor samples from patients with various types of brain cancer from the Cancer Genome Atlas project (TCGA) [54]. We consider two outcomes: a binary variable indicating tumor grade and type (0 for low-grade gliomas, 1 for high-grade glioblastoma) and a continuous variable measuring the mutational burden, the total
number of DNA mutations found in the tumor. When reporting performance, we use cross-entropy loss for the binary outcome and RMSE for the continuous outcome. Two covariate sets are used in separate analyses: clinical data (513 patients, 50 variables, 12 of
which are continuous) with missing data imputed using Random Forest via the ‘mice’ package, and gene expression data for the same patients previously shown by authors to be separable into clusters (513 patients, 100 genes). For the clinical covariates
analysis, we employ the vscc package in R, which uses mclust‘s Gaussian mixture approach for simultaneous variable selection and clustering. Because mclust handles a single data type, we apply it to the 12 continuous
covariates, demonstrating that CCWF can use clustering algorithms beyond k-means. Gene expression data clustering is based on the original authors’ cluster analysis. We implement our ensembling strategies as follows: Merged remains a single
500-tree forest trained on all 50 covariates. Subset Merged refers to a single 500-tree forest trained only on the 12 continuous variables. Sample Weighted refers to first clustering training data using vscc, training
100-tree forests on each cluster, and ensembling the learners with weights proportional to each cluster’s sample size. Stack Ridge follows the same clustering and forest-training paradigm as Sample Weighted, but instead ensembles learners
using stacked regression weights with a ridge constraint. For each iteration, we randomly select 100 samples for testing and use the rest for training, repeating this process 500 times to obtain distributional and median prediction accuracy measures.
Figure 4 displays the results: we observe substantial improvement of clustering methods over merging approaches across all covariate types and outcomes considered. For molecular profiling data (Fig. 4A), ensembles show 20-30% median improvement over merged-based learners in predicting mutation numbers, indicating vscc’s effectiveness for ensemble cluster construction. Gene expression data (Fig. 4b) shows larger improvements, with Sample Weighted and Stack Ridge achieving around 75% median improvement over the Merged. The gene expression clusters used to build the ensemble were previously found to
align with biologically relevant characteristics [54]. While Stack Ridge consistently outperforms Sample Weighted across prediction tasks, the
difference is modest, suggesting the weighting scheme’s influence on prediction ability varies by context. These results demonstrate CCWF’s robustness across different outcomes, covariates, and cluster constructions in a real biological setting. Note that
in this experiment, we have displayed the top performing approaches; in this case, the Unweighted and Random approaches are not competitive.
We introduce Cross-Cluster Weighted Forests (CCWF): ensembles of Random Forests (RFs) trained on estimated clusters with weights rewarding cross-cluster performance. Across theoretical analyses, simulations, and real biological data with feature-distribution heterogeneity, CCWFs yield more generalizable and accurate predictors than standard RFs. Theoretically, we analyze Centered Random Forest learners in a simplified CCWF setup and derive an asymptotic limit of the ensemble-to-merged RMSE ratio, showing that ensembling strictly improves on merging through bias reduction; this limit holds well in our experiments. Our findings highlight the role of data partitioning in RF algorithms. Dividing data by true clusters or studies is not always optimal; reducing within-cluster covariate heterogeneity enhances learning efficiency even at smaller sample sizes. Cluster-based ensembling further outperforms merging when the variables driving cluster formation have low overlap with those most predictive of the outcome. RFs are among the most widely applied machine learning tools, with thousands of published applications across diverse data types. We have identified a simple method that yields substantial improvements, and we believe that the architecture we outlined, as a general concept, may provide further important insight into combining simple and successful machine learning tools via multi-layer approaches. Several limitations remain. First, the theoretical result currently requires uniform clusters, CRFs, and partitions based on known structure; extensions to flexible data distributions, overlapping clusters, and algorithmically-determined cluster partitions remain future work. Second, the scope of our empirical analysis is restricted to biological tabular datasets with moderate dimension and defined heterogeneity. Future work can investigate whether the CCWF approach holds utility for non-biological datasets, which could greatly increase the scope of its utility. Third, our baselines are restricted to RF variants; future work should compare CCWF against other tabular methods such as gradient-boosted trees.
All code to replicate analyses from this paper can be found at https://github.com/m-ramchandran/cross-cluster. The CCWF method is implemented within the rCCWF R
package at https://github.com/m-ramchandran/rCCWF and the PyCCWF Python package available through PyPi.
Maya Ramchandran was supported by NIH-NCI Training Grant T32CA009337 at Harvard T.H. Chan School of Public Health. Rajarshi Mukherjee was partially supported by NSF Grant EAGER-1941419. Giovanni Parmigiani was supported by grant NSF-DMS 1810829.
We define a Random Forest as an ensemble of base regression trees \(\{f_n(\mathbf{x}_{\star};\theta_m, \mathcal{D}n), m \geq 1\}\), where \(\theta_m\) are i.i.d. copies of a randomizing variable \(\theta\) determining node splits. The forest-level prediction is \(\bar{f}_n(\mathbf{x}_{\star};\theta, \mathcal{D}_n) = \mathbb{E}_{\theta}\left[f_n(\mathbf{x}_{\star};\theta, \mathcal{D}_n)\right]\), typically estimated by averaging predictions from multiple trees trained on bootstrap samples [23]. In this section, we analyze the Centered Random Forest method proposed in [49], which simplifies the original algorithm by assuming \(\theta\) is independent of the training sample \(D_n\). This excludes bootstrapping and data-dependent tree-building strategies but permits \(\theta\) to be based on a second sample \(\mathcal{D}_n^{'} \stackrel{\rm i.i.d.}{\sim} \mathcal{D}_n\), enabling the splits to be optimized in a manner closely resembling how they would have been with the actual training set. We build upon the analytical work to characterize this model in [50] and [55], keeping our notation consistent whenever possible.
Centered Random Forest trees are constructed over \(\log_2c_n\) iterations, where \(c_n \geq 2\) is a user-defined parameter that possibly depends on \(n\). Each tree is built as follows: at each node, a coordinate of \(\mathbf{x}\) is selected, with the \(j^{th}\) feature having a probability \(p_{nj} \in (0,1)\) of being selected, and splits occur at the midpoint of the chosen coordinate. Following from Klusowski (2020), \(p_{nj} \to 1/S\) as \(n \to
\infty\). Each tree outputs the average \(Y_i\) for \(\mathbf{x}_i\) in the same leaf node partition as \(\mathbf{x}_{\star}\), denoted \(A_n(\mathbf{x}_{\star}, \theta)\). The individual tree predictor can therefore be represented as \(f_n(\mathbf{x}_{\star};\theta, \mathcal{D}_n) = \frac{\sum_{i = 1}^n Y_i
\mathbb{1}\{\mathbf{x}_{\mathbf{i}} \in A_n(\mathbf{x}_{\star}, \theta)\}}{\sum_{i = 1}^n \mathbb{1}\{ \mathbf{x}_{\mathbf{i}} \in A_n(\mathbf{x}_{\star}, \theta )\}} \mathbb{1}\{\epsilon_n(\mathbf{x}_{\star}, \theta)\}\), where \(\epsilon_n(\mathbf{x}_{\star}, \theta )\) is the event that
\(\sum_{i = 1} \mathbb{1}\{\mathbf{x}_{\mathbf{i}} \in A_n(\mathbf{x}_{\star}, \theta )\} > 0\); that is, that there is at least one training point that falls within the same leaf node as \(\mathbf{x}_{\star}\). We can then obtain the expected prediction made by the Centered Random Forest by taking the expectation of the individual tree predictors with respect to the randomizing variable \(\theta\): \(\bar{f}_n(\mathbf{x}_{\star};\theta, \mathcal{D}_n) = \sum_{i = 1}^n Y_i \mathbb{E}_{\theta}\left[\frac{\mathbb{1}\{\mathbf{x}_{\mathbf{i}} \in A_n(\mathbf{x}_{\star}, \theta)\}}{\sum_{i =
1}^n \mathbb{1}\{\mathbf{x}_{\mathbf{i}} \in A_n(\mathbf{x}_{\star}, \theta)\}} \mathbb{1}\{\epsilon_n(\mathbf{x}_{\star}, \theta)\} \right]\)
Lemma 1.
For data divided into \(k\) clusters, the predictions of the \(k\) cluster-level forests within the Ensemble learner on new point \(\mathbf{x}_{\star}\) can be expressed as \[\begin{align} \hat{Y}_b(\mathbf{x}_{\star}; \theta, \mathcal{D}_n) &= \sum_{i = 1}^n Y_i\mathbb{E}_{\theta}\left[ W_{ib}(\mathbf{x}_{\star}, \theta)\right] \end{align}\] where \[\begin{align} W_{ib}(\mathbf{x}_{\star}, \theta) &= \frac{\mathbb{1}_{\{\mathbf{x}_{\mathbf{i}} \in A_n(\mathbf{x}_{\star}, \theta)\}} \mathbb{1}_{\{\mathbf{x}_i \in [\frac{b-1}{k}, \frac{b}{k}]^S\}} \mathbb{1}_{\{\mathbf{x}_{\star} \in [\frac{b-1}{k}, \frac{b}{k}]^S\}}}{N_b(\mathbf{x}_{\star}, \theta)} \\ & \times \mathbb{1}_{\{\epsilon_{nb}(\mathbf{x}_{\star}, \theta)\}}\\ N_b(\mathbf{x}_{\star}, \theta) &= \sum_{i = 1}^n \mathbb{1}_{\{\mathbf{x}_{\mathbf{i}} \in A_n(\mathbf{x}_{\star}, \theta)\}} \mathbb{1}_{\mathbf{x}_i \in [\frac{b-1}{k}, \frac{b}{k}]^S\}} \mathbb{1}_{\mathbf{x}_{\star} \in [\frac{b-1}{k}, \frac{b}{k}]^S\}} \end{align}\] and \(\epsilon_{nb}(\mathbf{x}_{\star}, \theta)\) is the event that \(N_b(\mathbf{x}_{\star}, \theta) > 0\) for \(b = 1, \ldots, k\); \(N_b(\mathbf{x}_{\star}, \theta)\) represents the number of total training points from \(\mathbb{X}_b\) that fall into the same partition of the test point \(\mathbf{x}_{\star}\) given that \(\mathbf{x}_{\star} \in [\frac{b-1}{k}, \frac{b}{k}]^S\)
Specifically, the predictions of the overall Ensemble as \[\begin{align} \hat{Y}_E(\mathbf{x}_{\star} ;\theta, \mathcal{D}_n) &= \sum_{i = 1}^n \sum_{b = 1}^k Y_i\mathbb{E}_{\theta}\left[ W_{ib}(\mathbf{x}_{\star}, \theta)\right] \\ &= \sum_{i = 1}^n Y_i \mathbb{E}_{\theta}\left[ W_{i}(\mathbf{x}_{\star},\theta) \right] \end{align}\]
where \(W_{i}(\mathbf{x}_{\star},\theta) = \sum_{b = 1}^k W_{ib}(\mathbf{x}_{\star}, \theta)\).
The predictions of the Merged learner on new point \(\mathbf{x}_{\star}\) can be expressed as \[\begin{align} \hat{Y}_M(\mathbf{x}_{\star}; \theta, \mathcal{D}_n) &= \sum_{i =
1}^n Y_i\mathbb{E}_{\theta}\left[ H_{i}(\mathbf{x}_{\star}, \theta)\right]
\end{align}\] where \(H_{i}(X, \theta) = \frac{\mathbb{1}_{\{\mathbf{x}_{\mathbf{i}} \in A_n(\mathbf{x}_{\star}, \theta) \}}}{N_n(\mathbf{x}_{\star}, \theta)} \mathbb{1}_{\left\{\epsilon_n(\mathbf{x}_{\star}, \theta)
\right\}}\) and
\(N_n(\mathbf{x}_{\star}, \theta) = \sum_{i = 1}^n \mathbb{1}_{\{\mathbf{x}_{\mathbf{i}} \in A_n(\mathbf{x}_{\star}, \theta) \}}\), the number of total training samples falling into the same box as \(\mathbf{x}_{\star}\). Finally, \(\epsilon_n(\mathbf{x}_{\star}, \theta)\) is the event that \(N_n(\mathbf{x}_{\star}, \theta) > 0\).
In this section, we provide a a derivation for the limit of the ratio between the bias of the Ensemble and the bias of the Merged when the number of clusters in the training and testing data is equal to \(k\) for \(k \geq 2\).
Throughout, we use the following result regarding the leaf node containing test point \(\mathbf{x}_{\star}\).
Remark 2. Let \(a_{nj}(\mathbf{x}_{\star}, \theta)\) and \(b_{nj}(\mathbf{x}_{\star}, \theta)\) be the left and right endpoints of \(A_{nj}(\mathbf{x}_{\star}, \theta)\), the j\(^{th}\) side of the box containing \(\mathbf{x}_{\star}\), for \(j = 1, \ldots, S\). For ease of notation, we will henceforth refer to \(a_{nj}(\mathbf{x}_{\star}, \theta)\) as \(a_{nj}\) and \(b_{nj}(\mathbf{x}_{\star}, \theta)\) as \(b_{nj}\). \(c_{nj}(\mathbf{x}_{\star}, \theta)\) represents the number of times that the j\(^{th}\) coordinate is split, with the total number of splits across all coordinates set to equal \(\log_2{c_n}\) for some constant \(c_n > 2\). For ease of notation, we will suppress the dependencies of \(a_{nj}, b_{nj}\), and \(c_{nj}\) on \((\mathbf{x}_{\star}, \theta)\). We then observe that each endpoint of \(A_{nj}(\mathbf{x}_{\star}, \theta)\) is a randomly stopped binary expansion of \(\mathbf{x}_{\star}^{(j)}\): \[\begin{align} a_{nj} &\overset{D}{=} \frac{1}{k} \sum_{l = 1}^{c_{nj}} B_{kj} 2^{-l} + T \\ b_{nj} &\overset{D}{=} \frac{1}{k}\sum_{l = 1}^{c_{nj}} B_{kj} 2^{-l} + 2^{-{ c_{nj}}} + T \end{align}\] No matter the value of T, the length of the j\(^{th}\) side of the box is given by: \[\begin{align} {\lambda(A_{nj}) = b_{nj} - a_{nj} = \frac{1}{k} 2^{-c_{nj}}} \end{align}\] The division by the number of clusters \(k\) reflects the fact that the test point falls within a cluster of width \(\frac{1}{k}^{S}\), so we divide the typical uniform [0,1] binary expansion by the same factor. The measure of the box \(A_n\) is therefore equal to \[\begin{align} \lambda(A_n) &= \prod_{j =1}^{S} \lambda(A_{nj}) \\ &= \left(\frac{1}{k}\right)^{S} 2^{-\lceil \log_2c_n \rceil} \end{align}\] since by construction, \(\sum_j c_{nj} = \lceil \log_2c_n \rceil\).
The rest of the proof is structured as follows:
We derive the expression for the ratio of the Ensemble and Merged biases
We derive expressions for the terms \(W_1 W_2^{'}\) and \(H_1 H_2^{'}\)
We use (ii) to factor the expressions for the Ensemble and Merged biases into two terms
We provide upper and lower bounds for the first term of the factorization
We use these bounds to bound the overall limiting ratio
We simplify the second term of the factorization
We combine all previous steps and take limits to complete the proof
Proof.
The leading term of the squared bias for the Ensemble can be represented as \[\begin{align} & n(n-1)\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\left[ \mathbb{E}_{\theta}\left[ W_{1}\right] (f(\mathbf{x}_1) - f(\mathbf{x}_{\star}))\mathbb{E}_{\theta}\left[ W_{2} \right] (f(\mathbf{x}_2) - f(\mathbf{x}_{\star}))\right] \\ =\text{ } &n(n-1) \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\left[ \mathbb{E}_{\theta}\left[ W_{1} \right] \mathbb{E}_{\theta}\left[ W_{2} \right] (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) (f(\mathbf{x}_2) - f(\mathbf{x}_{\star}))\right] \\ = \text{ } &n(n-1) \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n, \theta, \theta^{'}}\left[ W_{1} W_{2}^{'} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) (f(\mathbf{x}_2) - f(\mathbf{x}_{\star})) \right] \end{align}\] where \(W_{i}(\mathbf{x}_{\star},\theta) = \sum_{b = 1}^k W_{ib}(\mathbf{x}_{\star}, \theta)\), as defined in Lemma 1 (i) and \(\theta^{'}\) is an independent copy of \(\theta\); thus, for any quantity indexed with the \('\) notation, we replace \(A_n(\mathbf{x}_{\star}, \theta)\) with \(A_n(\mathbf{x}_{\star}, \theta^{'})\).
The leading term of the squared bias for the Merged is equal to \[\begin{align} & n(n-1)\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\left[ \mathbb{E}_{\theta}\left[ H_{1} \right] (f(\mathbf{x}_1) - f(\mathbf{x}_{\star}))\mathbb{E}_{\theta}\left[ H_{2} \right] (f(\mathbf{x}_2) - f(\mathbf{x}_{\star}))\right] \\ = \text{ } &n(n-1) \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n, \theta, \theta^{'}}\left[ H_{1} H_{2}^{'} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) (f(\mathbf{x}_2) - f(\mathbf{x}_{\star})) \right] \end{align}\] where \(H_1\) and \(H_2^{'}\) are defined in Lemma 1 (ii).
Therefore, the limiting ratio of the Ensemble bias vs. the Merged bias is
\[\begin{align} \lim_{n \to \infty} \sqrt{\frac{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n, \theta, \theta^{'}}\left[ W_{1} W_{2}^{'} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) (f(\mathbf{x}_2) - f(\mathbf{x}_{\star}))\right]}{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n, \theta, \theta^{'}}\left[H_{1} H_{2}^{'} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star}))(f(\mathbf{x}_2) - f(\mathbf{x}_{\star}))\right]}} \end{align}\]
For the rest of the proof, we focus on evaluating the ratio within the square root, and applying the square root function as the last step.
To simplify this further, we expand the quantities \(W_{1} W_{2}^{'}\) and \(H_{1} H_{2}^{'}\).
\[\begin{align} \;& W_{1} W_{2}^{'}\\ &= \mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_{\mathbf{2}} \in A_n^{'}\}} \\ &\times \left[ \sum_{b = 1}^k \frac{\mathbb{1}_{\{\mathbf{x}_1 \in [\frac{b-1}{k}, \frac{b}{k}]^S\}} \mathbb{1}_{\{\mathbf{x}_{\star} \in [\frac{b-1}{k}, \frac{b}{k}]^S\}}}{N_b} \mathbb{1}_{\{\epsilon_{nb}\}} \right] \times \left[ \sum_{b = 1}^k \frac{\mathbb{1}_{\{\mathbf{x}_1 \in [\frac{b-1}{k}, \frac{b}{k}]^S\}} \mathbb{1}_{\{\mathbf{x}_{\star} \in [\frac{b-1}{k}, \frac{b}{k}]^S\}}}{{N_b}^{'}} \mathbb{1}_{\{\epsilon_{nb}^{'}\}} \right] \\ &= \mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_{\mathbf{2}} \in A_n^{'}\}} \\ &\times\left[\sum_{b = 1}^k \frac{\mathbb{1}_{\{\mathbf{x}_1 \in [\frac{b-1}{k}, \frac{b}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_2 \in [\frac{b-1}{k}, \frac{b}{k}]^S\}}[\mathbb{1}_{\{\mathbf{x}_{\star} \in [\frac{b-1}{k}, \frac{b}{k}]^S\}}]^2 \mathbb{1}_{\{\epsilon_{nb}\}} \mathbb{1}_{\{\epsilon_{nb}^{'}\}}}{N_b N_b^{'}} \right] \end{align}\] The cross terms of the product are equal to 0, since \(\mathbf{x}_{\mathbf{1}}, \mathbf{x}_{\mathbf{1}}, \text{ and } \mathbf{x}_{\mathbf{\star}}\) have to fall in the same cluster for any term to be non-zero.
Similarly, \[\begin{align} H_{1} H_{2}^{'} &= \mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_{\mathbf{2}} \in A_n^{'}\}} \times \frac{\mathbb{1}_{\left\{\epsilon_n \right\}}}{N_n} \times \frac{\mathbb{1}_{\left\{\epsilon_n^{'} \right\}}}{N_n^{'}} \end{align}\]
We now simplify these quantities using the following definitions and assumptions: Define \[\begin{align} T_b &= \sum_{i \geq 3} \mathbb{1}_{\{\mathbf{x}_i \in A_n\}}\mathbb{1}_{\{\mathbf{x}_i \in [\frac{b-1}{k}, \frac{b}{k}]^S\}} \mathbb{1}_{\{\mathbf{x}_{\star} \in [\frac{b-1}{k}, \frac{b}{k}]^S\}} \end{align}\] for \(b = 1, \ldots, k\) and let \(T_b^{'}\) be the equivalent expressions based on \(\theta^{'}\). We make the assumption that the quantities \(\mathbb{1}_{\{\epsilon_{nb}\}}\), \(b = 1, \ldots, k\), as well as \(\mathbb{1}_{\{\epsilon_{n}\}}\), are all non-zero; that is, that there is always at least one training point that falls within \(A_n(\mathbf{x}_{\star}, \theta)\), within the box (leaf node) containing the test point, regardless of the test point’s true cluster assignment. Without loss of generality, we can therefore order the dataset so that one of \(\mathbf{x}_1\), \(\mathbf{x}_2\) falls within \(A_n(\mathbf{x}_{\star}, \theta)\) and the other falls in another cluster. For any test point, we can identify at least \(k\) training data points that belong to each distinct cluster out of which one is guaranteed to fall within the same leaf node as the test point. There are therefore a minimum of \(k/2\) different arrangements of the training set that meet this criteria for all \(b = 1, \ldots, k\).
For test points belonging to cluster \(b\), we denote the event that this assumption is met by
\(d_b = \mathbb{1}_{\left\{\epsilon_{nb} >0, \sum_{i = 1}^2 \mathbb{1}_{\{\mathbf{x}_i \in A_n\}} = 1, \sum_{i = 1}^2 \mathbb{1}\left\{\mathbf{x}_i \in \left[\frac{b-1}{k}, \frac{b}{k}\right]^S\right\} = 1\right\}}\).
For any test point belonging to an arbitrary cluster, we denote the event that this assumption is met by \(d_n = \mathbb{1}_{\{\epsilon_{n} >0, \sum_{i = 1}^2 \mathbb{1}\{\mathbf{x}_i \in A_n \} = 1\}}\). This event is
in fact exactly equivalent to our overall assumption, so by definition, \(P(\mathcal{D}_n) = 1\). Therefore, \[\begin{align}
\mathbb{E}\left[\frac{\mathbb{1}_{\{\epsilon_{nb}\}}\mathbb{1}_{\{\epsilon_{nb}^{'}\}}}{N_b N_b^{'}} \bigg| \mathbf{x}_{\star}\right] &= \mathbb{E}\left[\frac{1}{(1 + T_b)(1 + T_b^{'})} \bigg| \mathbf{x}_{\star}, d_b \right] \\
&= \frac{\mathbb{E}\left[\frac{1}{(1 + T_b)(1 + T_b^{'})} \bigg| \mathbf{x}_{\star} \right]}{P(d_b)} \\
&= \frac{k}{2}\mathbb{E}\left[\frac{1}{(1 + T_b)(1 + T_b^{'})} \bigg| \mathbf{x}_{\star} \right]
\end{align}\] for \(b = 1, \ldots, k\). For the Merged learner, we can similarly define \[\begin{align} U &= \sum_{i \geq 3} \mathbb{1}_{\{\mathbf{x}_i \in A_n\}}
\end{align}\] and \[\begin{align}
\mathbb{E}\left[\frac{\mathbb{1}_{\{\epsilon_{n}\}}\mathbb{1}_{\{\epsilon_{n}^{'}\}}}{N_n N_n^{'}} \bigg| \mathbf{x}_{\star}\right] &= \mathbb{E}\left[\frac{1}{(1 + U)(1 + U^{'})} \bigg| \mathbf{x}_{\star}, d_n \right] \\
&= \mathbb{E}\left[\frac{1}{(1 + U)(1 + U^{'})} \bigg| \mathbf{x}_{\star} \right] \\
\end{align}\]
Given these definitions, we can further simplify the ratio of the squared biases: \[\begin{align} &\frac{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n, \theta, \theta^{'}}\left[ W_{1} W_{2}^{'} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) (f(\mathbf{x}_2) - f(\mathbf{x}_{\star}))\right]}{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n, \theta, \theta^{'}}\left[H_{1} H_{2}^{'} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star}))(f(\mathbf{x}_2) - f(\mathbf{x}_{\star}))\right]} \\ \end{align}\]
The numerator (corresponding to the leading term of the squared bias of the Ensemble) simplifies to: \[\begin{align} &\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n, \theta, \theta^{'}}\left[ W_{1} W_{2}^{'} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) (f(\mathbf{x}_2) - f(\mathbf{x}_{\star}))\right] \\ = \text{ } & \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n, \theta, \theta^{'}}\Bigg[ \mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_{\mathbf{2}} \in A_n^{'}\}} \times \left[\sum_{b = 1}^k \frac{\mathbb{1}_{\{\mathbf{x}_1 \in [\frac{b-1}{k}, \frac{b}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_2 \in [\frac{b-1}{k}, \frac{b}{k}]^S\}}[\mathbb{1}_{\{\mathbf{x}_{\star} \in [\frac{b-1}{k}, \frac{b}{k}]^S\}}]^2 \mathbb{1}_{\{\epsilon_{nb}\}} \mathbb{1}_{\{\epsilon_{nb}^{'}\}}}{N_b N_b^{'}}\right] \\ = \text{ } & \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n, \theta, \theta^{'}}\Bigg[ \mathbb{E}\left[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_{\mathbf{2}} \in A_n^{'}\}} \times k\left[\frac{\mathbb{1}_{\{\mathbf{x}_1 \in [0, \frac{1}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_2 \in [0, \frac{1}{k}]^S\}}[\mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}}]^2 \mathbb{1}_{\{\epsilon_{n1}\}} \mathbb{1}_{\{\epsilon_{n1}^{'}\}}}{N_1 N_1^{'}}\right] \bigg| \mathbf{x}_{\star}, \theta, \theta^{'} \right] \Bigg] \\ = \text{ } & k \times \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \mathbb{E}\left[\frac{1}{(1 + T_1)(1 + T_1^{'})} \bigg| \mathbf{x}_{\star}, \theta, \theta^{'}, d_b \right] \\ & \times \mathbb{E}_{\mathbf{x}_1, \mathbf{x}_2}\bigg[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_{\mathbf{2}} \in A_n^{'}\}}\mathbb{1}_{\{\mathbf{x}_1 \in [0, \frac{1}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_2 \in [0, \frac{1}{k}]^S\}}[\mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}}]^2 \\ & \times (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) (f(\mathbf{x}_2) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}, \theta, \theta^{'} \bigg] \Bigg]\\ = \text{ } & \frac{k^2}{2} \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \mathbb{E}\left[\frac{1}{(1 + T_1)(1 + T_1^{'})} \bigg| \mathbf{x}_{\star}, \theta, \theta^{'} \right] \\ & \times \left(\mathbb{E}_{\mathbf{x}_1}\bigg[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_1 \in [0, \frac{1}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\Bigg]\right)^2 \Bigg]\\ \end{align}\]
Various independence relations are used to simplify the expression: \(\mathbf{x}_1 \perp \mathbf{x}_2\), \(\theta \perp \theta^{'}\), and \(T_1 \perp \mathbf{x}_1, \mathbf{x}_2\).
The denominator (corresponding to the leading term of the squared bias of the Merged) simplifies to:
\[\begin{align} &\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n, \theta, \theta^{'}}\left[H_{1} H_{2}^{'} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star}))(f(\mathbf{x}_2) - f(\mathbf{x}_{\star}))\right] \\ \text{ } & \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n, \theta, \theta^{'}}\left[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_{\mathbf{2}} \in A_n^{'}\}} \left(\frac{1}{1 + U}\right)\left(\frac{1}{1 + U^{'}}\right) (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) (f(\mathbf{x}_2) - f(\mathbf{x}_{\star}))\right] \\ \text{ } & \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n, \theta}\left[\left(\frac{1}{1 + U}\right)\left(\frac{1}{1 + U^{'}}\right) \left(\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \right)^2\right] \\ = \text{ } & \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \mathbb{E}\left[\frac{1}{(1 + U)(1 + U^{'})} | \mathbf{x}_{\star}, \theta, \theta^{'} \right] \times \left(\mathbb{E}_{\mathbf{x}_1}\left[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) | \mathbf{x}_{\star}\right]\right)^2 \Bigg]\\ \end{align}\]
We now focus on calculating \(\mathbb{E}\left[\frac{1}{(1 + T_1)(1 + T_1^{'})} | \mathbf{x}_{\star}, \theta, \theta^{'} \right]\) and \(\mathbb{E}\left[\frac{1}{(1 + U)(1 + U^{'})} | \mathbf{x}_{\star}, \theta, \theta^{'} \right]\). To upper bound these quantities, we use the fact that for \(Z \sim \text{Binomial}(m, p)\), \(\mathbb{E}\left[\left(\frac{1}{1+Z} \right)^2\right] \\\leq \frac{1}{(m + 1)(m + 2) p^2}\) along with the Cauchy-Schwarz inequality. To lower bound these quantities, we use the fact that for \(Z \sim \text{Binomial}(m, p)\), one has that \(\mathbb{E}\left[\frac{1}{1+Z}\right] = \frac{1}{(m + 1) p}\left[1-(1-p)^{m+1}\right]\) along with the definition of covariance, which tells us that Cov\((X, Y) = E[XY] - E[X]E[Y] \Rightarrow E[X]E[Y] \leq E[XY]\) for positive Cov\((X, Y)\). To move forward, we need to calculate the binomial probabilities associated with \(T_1\) and \(U\).
\(T_1 \sim \text{Binomial}(n-2, p_e)\), where \[\begin{align} p_e &=\mathbb{E}\left[\mathbb{1}_{\{\mathbf{x}_i \in A_n\}}\mathbb{1}_{\{\mathbf{x}_i \in [0, \frac{1}{k}]^S\}} \mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}} \right] \\ &= \mathbb{E}\left[\mathbb{1}_{\{\mathbf{x}_i \in A_n\}} \bigg| \mathbb{1}_{\{\mathbf{x}_i \in [0, \frac{1}{k}]^S\}} \mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}} \right] \times \mathbb{E}\left[\mathbb{1}_{\{\mathbf{x}_i \in [0, \frac{1}{k}]^S\}} \right] \times \mathbb{E}\left[\mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}} \right] \\ &= \frac{1}{k^2} \prod_{j = 1}^S \frac{b_{nj} - a_{nj}}{\frac{1}{k}} \\ &= \frac{1}{k^2} 2^{-\lceil \log_2c_n \rceil} \end{align}\]
\(U \sim \text{Binomial}(n-2, p_m)\) where \[\begin{align} p_m &= E_{\mathbf{x}_1} \left[ \mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n \}} \right] \\ &= \prod_{j = 1}^S k \times P\left(\mathbb{1}_{\{\mathbf{x}_1^{(j)} \in A_{nj}\}} \bigg|\mathbf{x}_1^{(j)} \in \left[0, \frac{1}{k}\right] \right)\text{P}\left(\mathbf{x}_1^{(j)} \in \left[0, \frac{1}{k}\right]\right) \\ &= \prod_{j = 1}^S P\left(\mathbb{1}_{\{\mathbf{x}_1^{(j)} \in A_{nj}\}} \bigg|\mathbf{x}_1^{(j)} \in \left[0, \frac{1}{k} \right], \mathbf{x}_{\star}^{(j)} \in \left[0, \frac{1}{k} \right] \right) \text{P}\left(\mathbf{x}_{\star}^{(j)} \in \left[0, \frac{1}{k}\right]\right) \\ &= \prod_{j = 1}^S \frac{b_{nj} - a_{nj}}{\frac{1}{k}} \times \frac{1}{k} \\ & = \frac{1}{k} 2^{-\lceil \log_2c_n \rceil} \end{align}\]
Now, we can state that by the definition of covariance,
\[\begin{align} \mathbb{E}\left[\frac{1}{(1 + T_1)(1 + T_1^{'})} \bigg| \mathbf{x}_{\star}, \theta, \theta^{'} \right] &\geq \mathbb{E}\left[\frac{1}{(1 + T_1)} \bigg| \mathbf{x}_{\star}, \theta, \theta^{'} \right] \mathbb{E}\left[\frac{1}{(1 + T_1^{'})} \bigg| \mathbf{x}_{\star}, \theta, \theta^{'} \right]\\ &= \left(\mathbb{E}\left[\frac{1}{(1 + T_1)} \bigg| \mathbf{x}_{\star}, \theta, \theta^{'} \right] \right)^2 \\ &= \frac{k^4 4^{\lceil \log_2c_n \rceil}}{(n-1)^2} \left[1 - \left(\frac{1}{k^2} 2^{-\lceil \log_2c_n \rceil}\right)^{n-1}\right]^2 \end{align}\]
By the Cauchy-Schwarz inequality, \[\begin{align} \;& \mathbb{E}\left[\frac{1}{(1 + T_1)(1 + T_1^{'})} \bigg|\mathbf{x}_{\star}, \theta, \theta^{'} \right]\\ &\leq \sqrt{\mathbb{E}\left[\left(\frac{1}{1 + T_1}\right)^2 \bigg|\mathbf{x}_{\star}, \theta, \theta^{'}\right]} \sqrt{\mathbb{E}\left[\left(\frac{1}{1 + T_1^{'}}\right)^2 \bigg|\mathbf{x}_{\star}, \theta, \theta^{'}\right]} \\ &=\left(\frac{k^2 2^{\lceil \log_2c_n \rceil}}{\sqrt{n(n-1)}}\right)^2 \\ &= \frac{k^44^{\lceil \log_2c_n \rceil}}{n(n-1)} \end{align}\]
Putting it all together, \[\begin{align} \frac{k^4 4^{\lceil \log_2c_n \rceil}}{(n-1)^2} \left[1 - \left(\frac{1}{k^2} 2^{-\lceil \log_2c_n \rceil}\right)^{n-1}\right]^2 \leq \mathbb{E}\left[\frac{1}{(1 + T_1)(1 + T_1^{'})} \bigg|\mathbf{x}_{\star}, \theta, \theta^{'} \right] &\leq \frac{k^44^{\lceil \log_2c_n \rceil}}{n(n-1)} \end{align}\]
Now, we can do a similar calculation for \(\mathbb{E}\left[\frac{1}{(1 + U)(1 + U^{'})} | \mathbf{x}_{\star}, \theta, \theta^{'} \right]\), yielding:
\[\begin{align} \frac{k^2 4^{\lceil \log_2c_n \rceil}}{(n-1)^2} \left[1 - \left(\frac{1}{k} 2^{-\lceil \log_2c_n \rceil}\right)^{n-1}\right]^2 \leq \mathbb{E}\left[\frac{1}{(1 + U)(1 + U^{'})} \bigg|\mathbf{x}_{\star}, \theta, \theta^{'} \right] &\leq \frac{k^2 4^{\lceil \log_2c_n \rceil}}{n(n-1)} \end{align}\]
To get an upper bound on our ratio simplified in part (iii), we can take the max of the numerator and the min of the denominator:
\[\begin{align} &\max \frac{\frac{k^2}{2} \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \begin{array}{c}\mathbb{E}\left[\frac{1}{(1 + T_1)(1 + T_1^{'})} | \mathbf{x}_{\star}, \theta, \theta^{'} \right] \\ \times\left(\mathbb{E}_{\mathbf{x}_1}\bigg[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_1 \in [0, \frac{1}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\Bigg]\right)^2 \end{array}\Bigg]}{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \mathbb{E}\left[\frac{1}{(1 + U)(1 + U^{'})} | \mathbf{x}_{\star}, \theta, \theta^{'} \right] \left(\mathbb{E}_{\mathbf{x}_1}\left[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\right]\right)^2 \Bigg]} \\ &=\frac{\frac{k^2}{2} \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \frac{k^44^{\lceil \log_2c_n \rceil}}{n(n-1)} \left(\mathbb{E}_{\mathbf{x}_1}\bigg[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_1 \in [0, \frac{1}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\Bigg]\right)^2 \Bigg]}{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \frac{k^24^{\lceil \log_2c_n \rceil}}{(n-1)^2} \left[1 - \left(\frac{1}{k} 2^{-\lceil \log_2c_n \rceil}\right)^{n-1}\right]^2 \left(\mathbb{E}_{\mathbf{x}_1}\left[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\right]\right)^2 \Bigg]} \\ &= \left(\frac{k^4}{2}\right)\left(\frac{n-1}{n}\right)\left(\frac{1}{\left[1 - \left(\frac{1}{k} 2^{-\lceil \log_2c_n \rceil}\right)^{n-1}\right]^2}\right) \\ &\times \frac{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \left(\mathbb{E}_{\mathbf{x}_1}\bigg[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_1 \in [0, \frac{1}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\Bigg]\right)^2 \Bigg]}{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \left(\mathbb{E}_{\mathbf{x}_1}\left[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\right]\right)^2 \Bigg]} \end{align}\]
We find the lower bound of the ratio by taking the min of the numerator and max of the denominator, yielding: \[\begin{align} &\min \frac{\frac{k^2}{2} \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[\begin{array}{c} \mathbb{E}\left[\frac{1}{(1 + T_1)(1 + T_1^{'})} | \mathbf{x}_{\star}, \theta, \theta^{'} \right] \\ \times \left(\mathbb{E}_{\mathbf{x}_1}\bigg[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_1 \in [0, \frac{1}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\Bigg]\right)^2 \end{array}\Bigg]}{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \mathbb{E}\left[\frac{1}{(1 + U)(1 + U^{'})} | \mathbf{x}_{\star}, \theta, \theta^{'} \right] \left(\mathbb{E}_{\mathbf{x}_1}\left[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\right]\right)^2 \Bigg]} \\ &= \left(\frac{k^4}{2}\right)\left(\frac{n}{n-1}\right) \left[1 - \left(\frac{1}{k^2} 2^{-\lceil \log_2c_n \rceil}\right)^{n-1}\right]^2 \\ &\times \frac{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \left(\mathbb{E}_{\mathbf{x}_1}\Bigg[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_1 \in [0, \frac{1}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\Bigg]\right)^2 \Bigg]}{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \left(\mathbb{E}_{\mathbf{x}_1}\left[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\right]\right)^2 \Bigg]} \end{align}\]
We now simplify the quantity \[\begin{align} \frac{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \left(\mathbb{E}_{\mathbf{x}_1}\bigg[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_1 \in [0, \frac{1}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\Bigg]\right)^2 \Bigg]}{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \left(\mathbb{E}_{\mathbf{x}_1}\left[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\right]\right)^2 \Bigg]} \end{align}\]
We begin with the numerator, evaluating the expression within the inner expectation: \[\begin{align} & \mathbb{E}_{\mathbf{x}_1}\Bigg[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_1 \in [0, \frac{1}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}} f(\mathbf{x}_1) - f(\mathbf{x}_{\star}) \bigg| \mathbf{x}_{\star}\Bigg] \\ = \text{ } &\mathbb{E}_{\mathbf{x}_1}\Bigg[ f(\mathbf{x}_1) - f(\mathbf{x}_{\star})\bigg| \mathbf{x}_{\mathbf{1}} \in A_n, \mathbf{x}_1 \in \left[0, \frac{1}{k}\right]^S, \mathbf{x}_{\star} \in \left[0, \frac{1}{k}\right]^S, \mathbf{x}_{\star}\Bigg] \\ \times &\mathbb{E}\left[ \mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n \}} | \mathbf{x}_1 \in \left[0, \frac{1}{k}\right]^S, \mathbf{x}_{\star} \in \left[0, \frac{1}{k}\right]^S, \mathbf{x}_{\star} \right] \times \mathbb{E}\left[\mathbb{1}_{\{\mathbf{x}_1 \in \left[0, \frac{1}{k}\right]^S} \right] \times \mathbb{E}\left[\mathbb{1}_{\{\mathbf{x}_{\star} \in \left[0, \frac{1}{k}\right]^S\}} \right] \\ = \text{ } & \left(\frac{\lambda(A_n)}{k} \right) \mathbb{E}_{\mathbf{x}_1}\Bigg[ f(\mathbf{x}_1) - f(\mathbf{x}_{\star})\bigg| \mathbf{x}_{\mathbf{1}} \in A_n, \mathbf{x}_{\star}\Bigg] \end{align}\]
Moving onto the denominator, we again evaluate the expression within the inner expectation: \[\begin{align} &\mathbb{E}_{\mathbf{x}_1}\left[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\right] \\ = \text{ } &\mathbb{E}_{\mathbf{x}_1}\left[ f(\mathbf{x}_1) - f(\mathbf{x}_{\star}) \bigg| \mathbf{x}_{\mathbf{1}} \in A_n, \mathbf{x}_{\star}\right] \times \mathbb{E}\left[ \mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n \}} \bigg| \mathbf{x}_{\star} \right] \\ = \text{ } &\mathbb{E}_{\mathbf{x}_1}\left[ f(\mathbf{x}_1) - f(\mathbf{x}_{\star}) \bigg| \mathbf{x}_{\mathbf{1}} \in A_n, \mathbf{x}_{\star}\right] \frac{\lambda(A_n)}{\frac{1}{k}}\\ \end{align}\]
Therefore, \[\begin{align} &\frac{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \left(\mathbb{E}_{\mathbf{x}_1}\bigg[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_1 \in [0, \frac{1}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\Bigg]\right)^2 \Bigg]}{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \left(\mathbb{E}_{\mathbf{x}_1}\left[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) | \mathbf{x}_{\star}\right]\right)^2 \Bigg]} \\ = \text{ }& \frac{\left(\frac{\lambda(A_n)}{k} \right)^2 \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \left( \mathbb{E}_{\mathbf{x}_1}\Bigg[ f(\mathbf{x}_1) - f(\mathbf{x}_{\star})\bigg| \mathbf{x}_{\mathbf{1}} \in A_n, \mathbf{x}_{\star}\Bigg] \right)^2 \Bigg]}{\left(k\lambda(A_n) \right)^2 \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \left(\mathbb{E}_{\mathbf{x}_1}\left[ f(\mathbf{x}_1) - f(\mathbf{x}_{\star}) \bigg| \mathbf{x}_{\mathbf{1}} \in A_n, \mathbf{x}_{\star}\right]\right)^2 \Bigg]} \\ = \text{ }& \frac{1}{k^4} \end{align}\]
Note that we do not need to make any assumptions on the form that \(f(x)\) takes, since the expressions involving the outcome model cancel between numerator and denominator.
Putting everything together from parts (v) and (vi), \[\begin{align} &\frac{\frac{k^2}{2} \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \begin{array}{c}\mathbb{E}\left[\frac{1}{(1 + T_1)(1 + T_1^{'})} | \mathbf{x}_{\star}, \theta, \theta^{'} \right]\\ \times \left(\mathbb{E}_{\mathbf{x}_1}\bigg[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_1 \in [0, \frac{1}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\Bigg]\right)^2 \end{array}\Bigg]}{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \mathbb{E}\left[\frac{1}{(1 + U)(1 + U^{'})} | \mathbf{x}_{\star}, \theta, \theta^{'} \right] \left(\mathbb{E}_{\mathbf{x}_1}\left[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\right]\right)^2 \Bigg]} \\ \geq \text{ }& \left(\frac{k^4}{2}\right)\left(\frac{n}{n-1}\right)\left[1 - \left(\frac{1}{k^2} 2^{-\lceil \log_2c_n \rceil}\right)^{n-1}\right]^2 \times \frac{1}{k^4} \\ = \text{ }& \left(\frac{n}{n-1}\right) \frac{1}{2} \left[1 - \left(\frac{1}{k^2} 2^{-\lceil \log_2c_n \rceil}\right)^{n-1}\right]^2 \end{align}\]
and
\[\begin{align} &\frac{\frac{k^2}{2} \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \begin{array}{c}\mathbb{E}\left[\frac{1}{(1 + T_1)(1 + T_1^{'})} | \mathbf{x}_{\star}, \theta, \theta^{'} \right]\\ \times \left(\mathbb{E}_{\mathbf{x}_1}\bigg[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_1 \in [0, \frac{1}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\Bigg]\right)^2 \end{array} \Bigg]}{\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \mathbb{E}\left[\frac{1}{(1 + U)(1 + U^{'})} | \mathbf{x}_{\star}, \theta, \theta^{'} \right] \left(\mathbb{E}_{\mathbf{x}_1}\left[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) \bigg| \mathbf{x}_{\star}\right]\right)^2 \Bigg]} \\ \leq \text{ }& \left(\frac{k^4}{2}\right)\left(\frac{n-1}{n}\right) \times \frac{1}{k^4} \left(\frac{1}{\left[1 - \left(\frac{1}{k} 2^{-\lceil \log_2c_n \rceil}\right)^{n-1}\right]^2}\right) \\ = \text{ }& \left(\frac{n}{n-1}\right) \frac{1}{2} \left(\frac{1}{\left[1 - \left(\frac{1}{k} 2^{-\lceil \log_2c_n \rceil}\right)^{n-1}\right]^2}\right) \end{align}\]
We now take square roots of the minimum and maximum values calculated above since we are interested in the ratio of the Ensemble bias vs. the Merged bias (not the squared ratio), and take limits as \(n \to \infty\).
\[\begin{align} \lim_{n \to \infty} \frac{\textit{Ensemble bias}}{\textit{Merged bias}} \in & \left[\begin{array}{c} \lim_{n \to \infty}\sqrt{\left(\frac{n}{n-1}\right) \frac{1}{2} \left[1 - \left(\frac{1}{k^2} 2^{-\lceil \log_2c_n \rceil}\right)^{n-1}\right]^2},\\ \lim_{n \to \infty}\sqrt{\left(\frac{n}{n-1}\right) \frac{1}{2} \left(\frac{1}{\left[1 - \left(\frac{1}{k} 2^{-\lceil \log_2c_n \rceil}\right)^{n-1}\right]^2}\right)}\end{array}\right] \\ &\to \left[\sqrt{\frac{1}{2}}, \sqrt{\frac{1}{2}}\right] \\ \Rightarrow &\lim_{n \to \infty} \frac{\textit{Ensemble bias}}{\textit{Merged bias}} \to \frac{1}{\sqrt{2}} \end{align}\] for any \(k \geq 2\)
◻
Remark 3. Convergence of bias terms
From the proof of Theorem 1, we can express the squared bias of the Ensemble* as: \[\begin{align} & n(n-1)\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n, \theta, \theta^{'}}\left[ W_{1} W_{2}^{'} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) (f(\mathbf{x}_2) - f(\mathbf{x}_{\star}))\right] \\ \to \text{ } &\frac{k^{4-2S}}{2}\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \left( \mathbb{E}_{\mathbf{x}_1}\Bigg[ f(\mathbf{x}_1) - f(\mathbf{x}_{\star})\bigg| \mathbf{x}_{\mathbf{1}} \in A_n, \mathbf{x}_{\star}\Bigg] \right)^2\\ \geq \text{ } & O\left(\frac{1}{n} \right) \end{align}\]*
The convergence of this term depends on the function \(f(x)\), distributions of \(\mathbf{x}_{1} \text{ and } \mathbf{x}_{\star}\), number of non-sparse variables \(S\), the number of clusters \(k\), and the total number of splits \(c_n\) (which influences the size of the box containing \(\mathbf{x}_{\star}\)). Assuming that \(\frac{k}{n}\), \(\frac{c_n}{n}\), and \(\frac{S}{n} \to 0\) as \(n \to \infty\), the squared bias converges at a rate greater than \(O\left(\frac{1}{n} \right)\)
We can express the squared bias of the Merged* as: \[\begin{align} & n(n-1)\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n, \theta, \theta^{'}}\left[ H_{1} H_{2}^{'} (f(\mathbf{x}_1) - f(\mathbf{x}_{\star})) (f(\mathbf{x}_2) - f(\mathbf{x}_{\star}))\right] \\ \to \text{ } &k^{4-2S}\mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \left( \mathbb{E}_{\mathbf{x}_1}\Bigg[ f(\mathbf{x}_1) - f(\mathbf{x}_{\star})\bigg| \mathbf{x}_{\mathbf{1}} \in A_n, \mathbf{x}_{\star}\Bigg] \right)^2\\ \geq \text{ } & O\left(\frac{1}{n} \right) \end{align}\] using similar arguments as above.*
We now show that the variance of the Ensemble and Merged asymptotically converge to zero. We assume that \(\text{Var}[f(\mathbf{x}) | \mathbf{x}] \leq \sigma^2\) for some \(\sigma > 0\).
Lemma 2.
Using [Biau 2012, Proposition 2], we can express the variance of the Ensemble as \[\begin{align} \mathbb{E}[\hat{Y}_E(\mathbf{x}_{\star}) - \mathbb{E}[\hat{Y}_E(\mathbf{x}_{\star})|\mathbf{x}_{\star}]]^2 = \text{ }& n\sigma^2\mathbb{E}\left[\mathbb{E}_{\theta, \theta^{'}}\left[W_1 W_2^{'} \right] \right] \\ = \text{ }& n\sigma^2\frac{k^2}{2} \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \mathbb{E}\left[\frac{1}{(1 + T_1)(1 + T_1^{'})} \bigg| \mathbf{x}_{\star}, \theta, \theta^{'} \right] \\ &\times \left(\mathbb{E}_{\mathbf{x}_1}\bigg[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \mathbb{1}_{\{\mathbf{x}_1 \in [0, \frac{1}{k}]^S\}}\mathbb{1}_{\{\mathbf{x}_{\star} \in [0, \frac{1}{k}]^S\}} \bigg| \mathbf{x}_{\star}\Bigg]\right)^2 \Bigg]\\ = \text{ }& n\sigma^2\frac{k^2}{2}\frac{k^44^{\lceil \log_2c_n \rceil}}{n(n-1)} \left(\frac{\lambda(A_n)}{\frac{1}{k}} \left(\frac{1}{k} \right)^2 \right)^2 \\ = \text{ }& \frac{\sigma^2 k^{(4-2S)}}{2(n-1)} \\ = \text{ }& O\left(\frac{1}{n} \right) \end{align}\]
using calculations from the proof of Theorem 1, step (iiia), (iv), and (vi) and finite assumptions on \(S\) and \(k\) outlined in Remark 3. Note that we do NOT need a finite assumption on the number of total splits \(c_n\) for the variance to be \(O(1/n)\).
Similarly, we can express the variance of the Merged as \[\begin{align} &\mathbb{E}[\hat{Y}_M(\mathbf{x}_{\star}) - \mathbb{E}[\hat{Y}_M(\mathbf{x}_{\star})|\mathbf{x}_{\star}]]^2\\ = \text{ }& n\sigma^2\mathbb{E}\left[\mathbb{E}_{\theta, \theta^{'}}\left[H_1 H_2^{'} \right] \right] \\ = \text{ }& n\sigma^2 \mathbb{E}_{\mathbf{x}_{\star}, \mathcal{D}_n}\Bigg[ \mathbb{E}\left[\frac{1}{(1 + U)(1 + U^{'})} \bigg| \mathbf{x}_{\star}, \theta, \theta^{'} \right] \times \left(\mathbb{E}_{\mathbf{x}_1}\bigg[\mathbb{1}_{\{\mathbf{x}_{\mathbf{1}} \in A_n\}} \bigg| \mathbf{x}_{\star}\Bigg]\right)^2 \Bigg]\\ = \text{ }& n\sigma^2\frac{k^24^{\lceil \log_2c_n \rceil}}{n(n-1)} \left(\frac{\lambda(A_n)}{\frac{1}{k}}\right)^2 \\ = \text{ }& \frac{\sigma^2 k^{(3-2S)}}{n-1} \\ = \text{ }& O\left(\frac{1}{n} \right) \end{align}\]
using calculations from the proof of Theorem 1, step (iiib), (iv), and (vi).
Therefore, the variance terms converge to 0 as \(n \to \infty\) and the squared bias asymptotically dominates the MSE.
Figure 5 presents simulation results illustrating Theorem 1’s relevance to clustered data generated from three distributional paradigms: uniform, Gaussian, and Laplace. The uniform distribution simulation matches our theoretical setup, with non-overlapping clusters of width \(\frac{1}{k}\). In contrast, the Gaussian and Laplace simulations feature overlapping clusters characterized by different location parameters and identity scale parameters. Across all distributions and the number of simulated clusters, the ratio of Ensemble to Merged RMSEs is consistently close to \(\frac{1}{\sqrt{2}}\); as expected, the uniform simulation exhibits a high consistency with the theory. When assumptions are violated, in the Gaussian and Laplace cases, ratios are close to \(\frac{1}{\sqrt{2}}\) as \(k\) gets large, but can deviate for smaller values. This illustrates that our analytical framework effectively captures the performance improvements of the Ensemble across diverse and realistic scenarios at least for large \(k\).
This simulations decomposes the MSE of Merged and Multi into their bias and variance terms to individually characterize each using the Gaussian cluster framework outlined in the Simulation setup. This bridges the more simplified setup we explore in the theory with the setup we analyze analytically through simulations and real data examples; instead of non-overlapping uniform clusters, CRFs as the base learners, and equal ensemble weights, we allow for overlapping Gaussian clusters, full RFs for the base learners, and stacking weights for the ensemble. As the training sample size increases, the MSE of both ensembling approaches decreases, and across all sample sizes considered, the squared bias term accounts for almost the entirety of the MSE, as expected. There is a far smaller decrease from 2500 to 5000 than from 500 to 2500, suggesting an asymptotic limit to how well the ensembles can perform, and supporting the analytical bounds on the risk from Theorems 1. All empirical improvements in performance from Multi over Merged arise through a smaller bias, supporting the validity of analytical results. Furthermore, the average ratio of MSE’s between the Multi and Merged is very similar to the theoretical limiting ratio. We see that in fact the variance is slightly higher for the Multi over the Merged, but is proportionally so small as to have almost no effect; this directly confirms our theoretical results showing that the variance term converges to 0 faster than the bias term. We limit our comparison in this simulation to the Multi rather than full CCWF, since this ensemble construction method can be analytically characterized much more easily than the Cluster or Random. These results directly show the applicability of our theoretical results to the case in which we use the standard RF algorithm and stacked regression weights.
We next evaluate method robustness across varying dataset characteristics, keeping \(k\) at the optimal value determined by Silhouette analysis. Using a quadratic outcome model and non-Gaussian cluster generation, we examine how prediction performance is affected by covariate-outcome signal strength (by varying coefficient norms), the number of true clusters in the training set, and overall sample size.
Figure 3a shows that as covariate signal strength increases (achieved by increasing the norm of the coefficient norm in the simulated outcome model), all ensembling methods improve over Merged, with performance gains plateauing at 20-60%. The Cluster consistently outperforms others, especially at higher coefficient norms. The Multi’s limited improvement suggests that reducing within-cluster variation is more effective than using true clusters for discerning covariate-outcome relationships. Figure 3b demonstrates that as true cluster count increases (keeping total sample size at 2500), the Cluster remains superior, while the Multi approaches the Random’s performance. This occurs because Multi’s forest count increases with the number of true clusters, aligning with previous observations that more ensemble members improve overall accuracy even as individual model complexity decreases. The Cluster’s consistent advantage highlights its benefits regardless of the true data composition. Figure 3c shows that as sample size increases from 500 to 5000, the performance gap between the Cluster and others widens. While the Multi becomes less distinguishable from the Merged (as the Merged accuracy increases proportionally with sample size), the Cluster and Random continue to improve upon the Merged since the optimal number of partitions \(k\) also increases proportionally with sample size.
We examine whether we can more optimally partition the total amount of data when multiple studies are available for training that measure the same covariates and outcome variable. The traditional multi-study ensembling paradigm is to train a single learner on each study and combine learners using some weighting strategy such as stacking [19]. This is analogous to the Multi method when we are able to separate data into its true clusters. We now explore whether training k-means on the the merged data (comprising covariate data from all of the training studies) produces improvements comparatively to the single dataset setting. We furthermore evaluate the performance of these approaches on real gene expression data from the CuratedOvarianData repository from Bioconductor in R [56]. Finally, we explore whether the general strategy of ensembling learners built on clusters also works for Neural Network base learners, and compare the results to ensembles using Random Forest.
The general framework for the simulation setup used to create Figure 8 is drawn from Ramchandran et. al. (2020) [15]. We use all 15 studies in CuratedOvarianData that include survival information without any missing data in the features. For N = 250 iterations per value of \(k\), we randomly separate the 15 datasets from CuratedOvarianData into 10 training and 5 validation sets. We then generate the outcome using a non-linear model in order to test the ability of the candidate learning methods to detect more difficult covariate-outcome relationships. We simulate baseline levels of coefficient perturbation per study similarly to the setup discussed in Section 4.1. Using either Random Forest or Neural Nets as the base learner, we then construct ensembles using the four main approaches compared throughout this paper. The Multi in this case trained a learner on each study to form the final ensemble. All ensembles are built using stacked regression weights with a ridge constraint.
For Random Forest-based ensembles, clustering outperforms study-based partitioning for \(k\) greater than the number of training studies. The Cluster approach remains most effective, followed by Random, Multi, and the Merged. The true cluster or study-membership does not represent the most effective data partitioning for Random Forest; again, it is more effective to partition the data based on minimizing within-cluster heterogeneity and maximizing across. Overall, these results suggest that for either single or multiple training datasets, the Cluster approach should be used for Random Forest learners.
For Neural Network-based models, the Cluster method still produces improvements over training on true studies or the merged data, though less so than with Random Forest. Notably, the Cluster approach yields similar accuracy with both base learning algorithms. While Neural Nets outperform Random Forest using the Merged and Multi approaches, the Cluster method elevates Random Forest-based ensembles to match Neural Net performance. This can be intuitively understood by comparing the Cluster strategy to convolutional neural networks, in which the estimated clusters represent the convolutional neighborhoods and the stacking weights delineate the relationship between the tree-learning and the stacking layers. While this analogy does not numerically explain the similar performance of the two algorithms, it provides insight into why we may be seeing these results.
In this section, we consider datasets with a single covariate, so that we can visualize the differences between the true outcome model and the estimated outcome models from the Merged and Ensemble as a function of the input covariate. We consider a cubic outcome model and the gaussian cluster-generating framework with two clusters.
The true model is visualized by the black line in all of the panels and the true outcome values corresponding to the generated covariate data is displayed by the black points. The predictions by the Merged and Ensemble are depicted in green and red respectively. In comparing the prediction accuracy of the two learning approaches, we find that in this instance, the Ensemble approach outperforms the Merged by 15%. Figure 9 allows us to visualize where these improvements occur; while the Merged is good at learning the true covariate-outcome at the densest regions of the covariate space (in this case, corresponding to the centers of the two clusters), the Ensemble is able to additionally more closely learn the true model at the tails of the covariate distribution, akin to local smoothing. By focusing each forest on a single cluster, the Ensemble method is able to focus more on all areas of the covariate space without getting pulled overly to the densest regions, while the Merged concentrates on those with most representation within the training set.