Mind the Missing Split: Resolving Feature Heterogeneity in Swarm Learning with Random Forests
Authors: Mohammad Tajabadi, Dominik Heider
Organizations: Institute of Medical Informatics, University of Münster, Münster, Germany
Abstract
Swarm Learning is a decentralized collaborative learning mechanism that allows multiple organizations to train a shared model without central coordination or direct data sharing. In typical horizontal Swarm Learning, datasets across sites are usually assumed to share the same feature set. However, in real-world applications, sites often have partially overlapping features because measurements, protocols, and available covariates differ across sites. This feature heterogeneity creates a practical issue for machine learning algorithms such as Random Forests. Specifically, when decision trees are pooled into a global Random Forest, inference at a given site can become ill-defined if a traversal encounters a split on a feature that is not available locally, often forcing organizations to discard site-specific variables upfront. In this paper, we address feature heterogeneity in Swarm Learning with Random Forests under partially overlapping feature spaces. We propose several deterministic and probabilistic inference-time strategies that resolve such missing splits without restricting training to the intersection of features. We evaluate the methods on nine datasets and demonstrate that they outperform both the intersection baseline and locally trained models across a broad range of scenarios.
Federated Learning (FL) enables collaborative training across decentralized clients, but most methods assume aligned feature schemas, an assumption that rarely holds in tabular settings where clients observe only partially overlapping feature subsets. In these heterogeneous feature spaces, parameter-averaging methods (e.g., FedAvg) transfer little information across weakly overlapping or disjoint feature groups, limiting their effectiveness for federated imputation. To overcome this, we propose \textbf{FedHF-Impute}, a federated imputation framework that separates structural feature unavailability from conventional missingness and uses a shared global feature graph to propagate information across statistically related features through message passing. This enables indirect cross-client knowledge transfer, even when features are never jointly observed locally, while preserving standard federated communication. Under simulated partial schema overlap on the SECOM and AirQuality datasets, FedHF-Impute improves imputation accuracy (RMSE) over FL baselines by 26.9%, and 8.4% respectively, while achieving comparable performance on PhysioNET, with only a 0.3% difference relative to the best baseline.
Multi-center survival prediction can improve robustness and generalizability, yet privacy regulations and institutional governance often prevent pooling patient-level clinical and genomic data across institutions. In practice, deployment is further complicated by feature-space heterogeneity, in which sites collect different covariates or use different sequencing panels, resulting in only partially overlapping feature sets. We present FederatedRSF, a Python package that implements federated random survival forests, aggregating locally trained survival trees and redistributing only feature-compatible trees to each site, enabling inference with partial overlap without sharing raw data. We evaluate FederatedRSF on the GBSG2 breast cancer cohort distributed with the scikit-survival package, simulating feature heterogeneity across clients by withholding subsets of features, and assessing discrimination using Harrell's concordance index (C-Index) under repeated cross-validation and site-splits. The results demonstrated that the federated model can achieve performance comparable to that of the centralized training setting.
Maryam Moradpour, Jonas Harriehausen, Amirreza Aleyasin +3
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.
Maya Ramchandran, Rajarshi Mukherjee, Giovanni Parmigiani