Federated XGBoost on Sample-Wise Non-IID Data

Originally, Federated Learning (FL) [1] emerged as a paradigm for collaboratively training a single global machine learning model without having to share or expose the local party’s raw data distributions among the federation, thus providing a layer of security for each party that participates in the training process. Typically these algorithms are often implemented as linear-based models or Neural Network-based approaches. In recent years, the adoption of Decision Tree-based models for Federated Learning [3] has begun to emerge as a new type of model specifically for Federated Learning models, compared to traditional models such as linear models and neural network-based approaches. The majority of tree-based approaches proposed in Federated Learning are based on either a Random Forest (RF) approach, or a variant of the Gradient Boosted Decision Tree-based algorithm, such as XGBoost [6] or LightGBM [7].

One of the largest challenges when training Federated Learning models involves utilizing data across heterogeneous sources, which often are non-IID. Many prior works have explored these areas of training Federated Learning algorithms under various non-IID scenarios. Zhao et. al [20] explored the effects of non-IID data for Neural Network-based Federated Learning algorithms and demonstrated significant performance degradation of the model. Sattler et. al [21] proposes a compression framework that allows for robustness against non-IID data when training models in Federated Learning. However, these approaches mainly evaluate Federated Learning models over Neural Network-based approaches. Fan et. al [22] on the other hand, propose a method for Gradient Boosted Decision Tree-based models for Federated Learning and evaluates the effects of skewed label data. However, their work does not perform a systematic evaluation of the effects of training a Federated Learning model over different types of sample-based non-IID data. In our work, we design very specific experiments that test the Federated XGBoost algorithm proposed by Ong et. al [5] under different non-IID scenarios to evaluate its performance over different types of data distribution skews in terms of sample size distributions. By studying how different allocation of sample sizes impacts the overall model performance, we can better evaluate the efficacy of Federated Gradient Boosted Decision Trees to see what the potential weakness or possible performance issues are that arise from these types of data distributions.

In this section, we introduce a slightly modified version of the XGBoost-based Federated Learning training method, also known as the Party-Adaptive XGBoost (PAX), proposed by Ong et. al [5]. One of the key differences we introduce in this paper is around the methodology of selecting the bin sizes for the individual party’s data, and how we simplify the implementation by utilizing a fixed bin size across all parties when generating the underlying histogram representation of the local party’s data. As noted previously, one of the major challenges in training a gradient-boosted decision tree is to find the optimal feature and value to split using the computed gain score. Similar to the method of decomposing the routines of the algorithm as described in the previous section, the same principles of the aggregator querying each party for its distribution statistics, fusing those statistics, and finding the optimal partition based on those fused statistics holds for the Gradient Boosted Decision Tree. In this case, instead of simple data count statistics, we deal with histograms of data distributions instead.

Optimized GBDT methods such as XGBoost [6] and LightGBM [7] utilize a quantile-based approximation to efficiently reduce the overall search space of the split finding process by approximating the raw data distribution as a surrogate histogram representation. Empirically, it has been shown that quantile approximations work just as well as the exact greedy solutions [6, 7, 24]. This method for data approximations can serve as a semi-secure approach for training Gradient Boosted Decision Trees under a Trusted Federations security policy. By quantizing or reducing the resolution of the raw distribution, we can effectively generate a surrogate representation of the raw data containing relatively lower fidelity of information than the original data distribution. Therefore, this does not directly reveal the raw data distribution of the original data source.

Many methods for building distributed quantile sketches exist, including GKMethod [25] and its extended variants [26]. However, each comes with its trade-offs in performance, speed, and reconstruction accuracy. For our proposed method, this algorithm implements the Distributed Distribution Sketch or DDSketch [27], an efficient and robust process that constructs highly accurate quantile sketch approximations of data distributions with the ability to merge multiple quantile sketches. Furthermore, another advantage of using DDSketch enables training of XGBoost where parties can join during intermediate steps of the boosting process when a new party joins the federation. Due to its ability to merge quantile sketch histograms efficiently and accurately, this method enables dynamic adaptations to new data distributions in the data as new parties join the federation.

To train a gradient-boosted decision tree within an FL-setting, the aggregator first initializes a global null model, $f_{\emptyset}^{(\mathcal{A})}$ (line 1-2). The null model in this scenario is based on computing the average values of the target label from each of the local party’s data distributions. For each party $p$, we obtain the sum of the target values as $Y_{p}$ and the corresponding sample size count, $n_{p}$ (Line 1). These values are collected from each local party and transmitted to the aggregator $\mathcal{A}$, where the aggregator computes the federation average of the target label values (Line 2).

Next, we then compute a federation-wise surrogate data representation that will be used in the boosting process when updating the tree model. To compute this surrogate data representation, we utilize a histogram-based method to compute approximations of the data. For each party $p$, we construct a data sketch of the local party’s input features for a given dataset $\mathcal{D}_{X}^{(p_{i})}$ (Line 5). For our implementation, we utilize DDSketch [27] to construct the histogram representation for each of the party’s local data distribution. During this process, we also provide a parameter $\delta$, which represents the relative error parameter used by the DDSketch algorithm. Controlling this parameter $\delta$ can dictate the amount of fidelity of the histogram used to approximate the data. Hence, the lower the error parameter used, the higher the risk of exposing the local party’s data distribution. After computing the local party’s histogram sketch representation, each party then transmits its histogram representation to the aggregator $\mathcal{A}$ (Line 6). When each of the parties finishes transmitting their local data distribution histograms to the aggregator, the aggregator then merges the histogram into a single histogram representation $\tilde{\mathcal{D}}^{\mathcal{A}}_{X}$.

After computing the surrogate histogram representation of the data, we then initiate the iterative federated learning process. First, the aggregator $\mathcal{A}$ transmits their global model, $f^{(\mathcal{A})}_{t}$ to each party, $p_{i}$, which is assigned to each party’s respective local model $f^{(p_{i})}_{t}$ (Line 13). We then evaluate $f_{t}^{(p_{i})}$ on $\tilde{\mathcal{D}}^{(p_{i})}_{X}$ to obtain the model’s predictions, $\hat{y}_{t}^{(p_{i})}$ (Line 14). Afterward, given the predictions, we compute the loss function which is used to compute the gradient, $g^{(p_{i})}$, and Hessian’s, $h^{(p_{i})}$ for each of the corresponding surrogate input feature value split candidates (Line 15). Gradient and Hessian statistics that fall under a certain bin interval are grouped within their respective value buckets [6]. The gradient and Hessian for each party are then sent back as replies to the aggregator and then collected until some quorum criterion has been met (Line 16). Given the collected results from each party, we perform a fusion operation to merge the final histogram representation used towards boosting the decision tree model, as formulated in the method of DDSketch [27] (Line 20).

Given the aggregated gradients $G^{(\mathcal{A})}_{m}$, Hessians, $H^{(\mathcal{A})}_{m}$, and the aggregated surrogate data representation histogram, $\tilde{\mathcal{D}}^{\mathcal{A}}_{X}$, we perform the boosting process to grow a new tree to derive a new global model, $f_{t}^{(\mathcal{A})}$ (Line 21). With a new $f_{t}^{(\mathcal{A})}$ generated, we repeat our training process for $T$ rounds, or until some other stopping criteria depending on whether early-stopping or other heuristics are considered (line 22).

We evaluate our algorithm with 7 different datasets ranging from Kaggle [28] and the University of California, Irvine Machine Learning Repository (UCI ML Repository) [29]. All datasets are classification task-based and are tabular data. They are a mix of binary and multiclass classification tasks and have a varying number of instances. We consider datasets of skewed labels and differing numbers of features, spanning from 9 to 756. The datasets span multiple subjects including science, finance, healthcare, and miscellaneous areas, which allows us to evaluate various non-IID scenarios. The entire list of the dataset utilized in this experiment is presented in Table II.

The Bank Marketing Dataset (Bank) was taken from the UCI ML Repository [30]. The focus of this dataset is classifying whether a customer subscribed to a marketing campaign done by a bank based on the characteristics of the customer and the specifics of the campaign. These 20 characteristics make up the feature values and included how many times and day of the week the customer was contacted. There are 45,211 instances in this dataset and is formulated as a binary classification task.

Bitcoin Heist Ransomware Address Dataset (Bitcoin) is sourced also from the UCI ML Repository [31]. This dataset has features on executed bitcoin transactions, such as information on the merging pattern and if the coins were split. It classifies if the transaction was hit with ransomware, and if so, classifies the type of ransomware. This dataset is significantly larger, originally with 2,916,697 instances, but was preprocessed to only include one transaction per originating address. There were also originally many types of ransomware, but only the top 6 are utilized. This was performed since the aggregator needs all the parties to have the same number of classes and many ransomware only have single instances.

Credit Card Fraud Dataset (Credit Card) is taken from Kaggle [32]. The dataset was preprocessed through PCA from the original data from transactions to preserve the privacy of the original dataset. The dataset contains 284,807 instances with only 492 instances being classified as a fraudulent transactions, making this a binary classification task. This makes the dataset highly imbalanced, with less than 1% of the whole dataset being in a positive class.

We use an already encoded image classification dataset, which is the Dry Bean Dataset from the UCI ML Repository [33]. The encoding done makes up the features of the dataset as characteristics of the beans, such as area. It has 13,611 instances and is based on a multiclass classification task with 7 labels.

The HTRU 2 Dataset is also sourced from the UCI ML Repository and is a binary dataset with 9 features [34]. It determines whether a pulsar candidate is a pulsar and is imbalanced with 1,639 of the 17,898 instances classified as positive. The features all pertain to statistics of the integrated pulse profile or the DM-SNR curve of the pulsar.

The Internet Firewall Dataset (Firewall) is a multiclass dataset of whether an action was allowed or not by a firewall system [35]. It is taken from the UCI ML Repository and has 65,532 instances spread across 4 labels. With features pertaining to the transaction like source and destination port. This dataset only has experiments run on 3 of the partitions since the fourth would create a party with less than 4 labels, and thus could not be run on the Federated Learning framework. To run, all parties need to have data with the same number of labels.

The final dataset is the Parkinson’s Dataset which is sourced from UCI ML Repository and is binary of whether a patient has Parkinson’s [36]. This is based on 754 features, which makes an interesting problem for the algorithm since as the number of features increases, so does the complexity of the decision trees thus increasing the time spent on each iteration. This dataset is 756 instances in size.

The datasets have limited preprocessing done other than encoding the data into numeric values or deleting non-unique values.

Models are evaluated using a 5-party Federated Learning setup. Each party has access to one partition of data as well as an evaluation dataset which is then shared with the aggregator to evaluate the model created by the partitions of data. Each partition contributes to the training dataset used by the aggregator and has access to the test dataset used.

We run experiments with $\epsilon=255$ max bins for the response histograms and with 100 max rounds. For the DDSketch relative error parameter we utilize a value of $\delta=0.01$. We tune hyperparameters using an l2-regularization of 0.1 and a learning rate also of 0.1. The experiments are run using the IBM Federated Learning framework [37].

We split datasets to simulate non-IID cases by reallocating some samples that each party has, thus becoming increasingly imbalanced. This heterogeneity in data was done to evaluate how the model would adapt to distribution changes in the parties. Real-world applications of the framework will use parties that have data distributed differently across them, in most cases. Thus the framework will need to have consistent results despite the heterogeneity of the parties’ data.

These partitions are all done using 80% of the data and the rest is preserved for evaluation. First, starting with each party having an equal number of samples as Sample Even experiments. Then move on to creating Sample Split A, which is done by reallocating 50% the number of instances from Party 2 and moving it to Party 1. Then 75% of the instances from Party 3 and moving it to Party 2. This was similarly done for Party 4 and 5, where 62.5% was moved to Party 3 and where 56.25% was moved to Party 4 respectively. This creates an unequal distribution of the number of instances for each party. This process was repeated 3 more times for all datasets, except for the firewall dataset which was only repeated twice, creating an increasing imbalanced distribution of instances for the parties. Thus, in the end, Party 1 has 68% of the number of instances in the dataset, and Party 5 has less than 1%. The ratio of each of the parties, relative to the whole training set, for each of the partitions, is shown in Table II. The actual size of each party used for each reallocation for each dataset is noted in Table I.

This method of reallocating the data should keep the label distribution of each of the parties consistent with the original dataset as well as each other party thus label non-IID is not an underlying factor for influencing the results. We perform such data reallocation, as the data is randomly sampled, the sample taken should continue to follow the same distribution as the original population. Thus the instances that are reallocated to the previous party should hold the same label and feature distribution as the originating party and the party it is being added to. Additionally, once the reallocated instances are added to the party, the label and feature distribution should not change for the originating or new party based on the loss or addition of instances.

The partitions of data are evaluated against both a centralized case as well as against the local parties. The centralized case is done by using a global holdout set of 20% of the data to evaluate the performance of the model that uses the remainder of the data, or 80%. The local parties are each party’s data for every partition, so the five from Sample Even and five from Sample Split A, and so on, are evaluated against the same holdout set as before. The local party results are averaged for each of the datasets.

We evaluate the partitions using 5 parties and the same global holdout set as previous evaluations to test model performance with the non-IID data. Thus, the aggregator receives responses from five parties to build the model, instead of one which the previous evaluations use. We use the F1 score for the evaluation metric of all the models and compare the centralized model as well as each of the partitions.