Federated Graph Learning with Adaptive Importance-based Sampling

Existing work on efficient intra-graph FedGCN training on large graphs can be divided into three branches. The first branch uses missing neighbor generation to obtain accurate node embeddings (Zhang et al. 2021b, a). However, it only focuses on improving prediction accuracy without considering the computation and communication overhead caused by additional training of the generative model. The second branch focuses on graph sampling approaches (Zhang et al. 2021b), e.g., FedGraph (Chen et al. 2021) uses deep reinforcement learning (DRL) to select neighbor nodes for embedding aggregation. However, since it ignores the graph topology and heterogeneity of clients, it results in inaccurate node embeddings and large variance in the gradients. Besides, it would incur large computation overhead as each local client needs to separately train two additional DRL networks. The third branch periodically transfer cross-client neighbor node embeddings to reduce communication costs (Chen et al. 2021; Du and Wu 2022; Zhang et al. 2022; Deng et al. 2023). However, they fail to capture the dynamics of model training to determine the optimal transfer period, resulting in inferior model performance.

One approach to scaling up GCN training is graph sampling, which can be categorized into: 1) node-wise sampling, 2) layer-wise sampling, and 3) subgraph sampling. Node-wise sampling methods (Hamilton et al. 2017; Cong et al. 2020) iteratively sample a number of neighbors for each node based on specific probabilities (e.g., calculated based on node centrality). Layer-wise sampling methods (Chen et al. 2018; Zou et al. 2019) independently sample a number of nodes for each GCN layer. Since multiple nodes are jointly sampled in each layer, the time cost of the sampling process is reduced by avoiding the exponential extension of neighbors. However, since nodes of different layers are sampled independently, some sampled nodes may have no connections with the ones in the previous layer, which would deteriorate the training performance. Subgraph sampling methods (Chiang et al. 2019; Zeng et al. 2019) sample a number of subgraphs for each batch in GCN training. However, graph partitioning of large graphs is time-consuming and the model performance is highly sensitive to the cluster size. Those three categories of methods, however, require direct access to data features, which would risk privacy leakage of local clients in FL settings. Besides, those methods neglect communication costs and would incur substantial communication costs when the volume of the graph data is large. Thus, we propose FedAIS to improve trade-offs between accuracy and efficiency (Fig. 1). Here, FedLocal is the federated GraghSage (Hamilton et al. 2017), which conducts random selection of within-client neighbor nodes, where the cross-client neighbor information is ignored. FedPNS performs periodic selection of both within-client and cross-client neighbor nodes.

To construct a global FedGCN model with fast convergence speed and low prediction error, we need to first analyze the sources of variances and biases when applying graph sampling strategies. Existing graph sampling approaches suffer from high variances and biases introduced to the stochastic gradients due to the approximation of node embeddings at different layers (Cong et al. 2020; Du and Wu 2022). We denote $\tilde{h}_{v}^{(l)}$ as the embedding approximation of node $v\in V_{k}$ in the $l$-th layer. Specifically, the variance of stochastic gradient estimator $\tilde{g}=\sum_{k=1}^{K}\frac{1}{|B_{k}|}\sum_{v\in B_{k}}\nabla f(\tilde{h}_{v}^{(L)},\theta_{t},y_{v})$, can be decomposed as:

$\mathbb{E}[||\tilde{g}-g||]$ denotes the variance of estimated gradients to their exact values $g=\sum_{k=1}^{K}\frac{1}{n_{k}}\sum_{v\in V_{k}}\nabla f(h_{v}^{(L)},\theta_{t},y_{v})$, resulting from the inner layer embedding approximation in forward propagation. The term $\mathbb{E}[||g-\nabla{F}(h^{(L)},\theta)||]$ denotes the variance of mini-batch gradients to the full gradients due to the mini-batch sampling.

Theorem 1 lets us immediately derive an upper error bound for the estimated gradients, i.e.,

From the decomposition of variance in Eq. (3) and the upper error bound for the estimated gradients in Eq. (5), we conclude that any graph sampling method introduces two sources of variance (i.e., the embedding approximation variance and the stochastic gradient variance). Therefore, to accelerate model convergence and reduce computation and communication overhead, both kinds of variance needs to be accounted in designing a graph sampling strategy.

FedAIS consists of two modules (as shown in Fig. 2).

1. 1.

   Historical Embedding-based Graph Sampling. In training round $t$, client $k$ updates the importance scores of its local training samples based on historical embeddings and training losses of individual samples. Then, client $k$ selects its most influential samples to be used for FedGCN model training. Using historical embeddings and training losses helps reduce both embedding approximation variance and stochastic gradient variance, thereby enabling accurate FedGCN model training.

2. 2.

   Adaptive Embedding Synchronization and Model Updating. With the influence estimation and sampling results, each client $k$ updates its historical embeddings and performs node aggregation. Then, it updates its local model and estimates the next optimal synchronization interval via the joint analysis of the overhead and error-convergence. Finally, client $k$ sends the updated local model to the server, which then aggregates the local models to produce the global model.

While evaluating embedding $h_{v}^{(l)}$, it is prohibitively costly to calculate all $h_{v}^{(l)}$ terms since they need to be computed and transmitted recursively (i.e., we again need the embeddings $h_{w}^{(l-1)}$ of all $v$’s neighbor nodes $w$). To reduce computation and communication costs, we introduce the historical embedding for FedGCN as an affordable approximation.

Here, we separate the within-client neighbors into two parts: 1) within-client in-batch nodes $w\in N(v)\cap B$, which are part of the current batch $B_{k}\in V_{k}$; and 2) within-client out-of-batch nodes $w\in N(v)\setminus B_{k}\cap V_{k}$, which are part of the client $k$ but not included in the current batch $B_{k}$. For both neighbor nodes, we approximate their embeddings $h_{w}^{(l)}$ via historical embeddings $\bar{h}_{w}^{(l)}$ acquired in previous iterations. Compared to the previous approach which incurs exponentially computation and communication costs, that incurred by historical embedding estimator increases linear with $L$, i.e., $O(\sum_{k=1}^{K}V_{k}|\cup_{v\in V_{k}}N(v)\cup\{v\}|\cdot LTJ)$ computation operations and $O(\sum_{k=1}^{K}\sum_{v\in V_{k}}\sum_{w\in N_{v}\setminus V_{k}}\sum_{l=1}^{L}d^{l}\cdot TJ)$ communication cost, where $T$ and $J$ are numbers of global training rounds and local training epochs.

Based on the historical embedding estimator and the variance analysis, we select training nodes that contribute most to the objective function and accelerate model convergence. It can be casted into the following optimization problem,

The most straightforward solution is to use the $L_{2}$ norm of gradient as the probability. However, it requires the calculation of $n_{k}$ derivatives for each client $k$ at each local epoch $j\in[\tau_{t}]$, which is computationally prohibitive (Li et al. 2021). To solve this issue, we instead use the difference $\Delta_{j}=f(\tilde{h}_{v}^{(L)},\theta_{j+1},y_{v})-f(\tilde{h}_{v}^{(L)},\theta_{j},y_{v})$ of training losses between two consecutive local model updates to approximate the gradient $\Delta_{j}\approx\nabla f(\tilde{h}_{v}^{(L)},\theta_{j},y_{v})$. Then, client $k$ calculates the probability $p_{v}$ by normalizing the differences across its all training samples,

Thus, the computational complexity is $O(n_{k})$ since client $k$ only requires one forward propagation.

To analyze the effect of $\tau$ on the expected runtime, we consider the following delay model. In round $t$, the time taken by client $k$ to conduct a local model update at the $j$-th epoch is modeled as a random variable $c_{k,j}^{t}$ and the total communication delay is a random variable $o_{\tau}^{t}$. Then, the communication cost is $o_{k,\tau}^{t}b_{t}$, where $b_{t}$ is the average network bandwidth during round $t$. For the full synchronization, the total time to complete each round is $c_{\rm syn}=\max\{c_{1,1}^{t},\cdots,c_{K,1}^{t}\}+o_{\tau}^{t}$, while for the periodic synchronization, the average time to complete each round is $c_{\rm avg}=\max\{\bar{c}_{1}^{t},\cdots,\bar{c}_{K}^{t}\}+o_{\tau}^{t}/\tau_{t}$ where $\bar{c}_{k}^{t}=\frac{1}{\tau_{t}}\sum_{j=1}^{\tau_{t}}c_{k,j}^{t}$. Consider the simplest case where $c_{k,\tau}^{t}=c$ and $o_{\tau}^{t}=o$ are constants, and $c/o$ is the ratio of communication delay to computation cost, which depends on the size of FedGCN model, network bandwidth and client computing capacity, etc.

From the optimization error bound in Eq. (9), the error-runtime trade-off for different synchronization communication intervals can be derived. While a larger $\tau$ reduces the runtime per iteration and makes the first term in Eq. (9) smaller, it also adds noise and increases the last term.

Then, we determine the optimal embedding synchronization interval $\tau$ to minimize the optimization error for each training batch. We start with infrequent cross-client embedding synchronization for improved convergence speed, and gradually transiting to higher embedding synchronization frequencies to reduce the prediction error of the learned global model. Specifically, at each training round $t$, the server selects the optimal embedding transmission interval that achieves the fast test loss decay of the global model $\theta_{t}$ for the next interval. Theorem 2 illustrates that there is an optimal value $\tau_{t}$ that minimizes the optimization error bound at round $t$ between the server and all selected clients,

It can be observed from Eq. (10) that the generated synchronization period sequence decreases along with the objective value on the test set when the learning rate is fixed. It is consistent with the intuition that the trade-off between error-convergence and overhead varies over time. Compared to the initial training phase, the benefit of using a large synchronization interval diminishes as the model converge since a lower error is preferred in the latter training phase. In some scenarios where the Lipschitz constant $\lambda$ and the gradient variance bound $\zeta^{2}$ are unknown and estimating these constants are difficult due to the highly non-convex and high-dimensional loss surface. As an alternative, we propose a simpler rule where we approximate $F_{inf}$ by 0, and divide Eq. (10) by $\tau_{0}$ to obtain the basic synchronization interval,

where $\lceil r\rceil$ is the ceil function to round $r$ to the nearest integer. In practical implementations, we take the test loss as the objective function value and the average batch number $\sum_{k=1}^{K}n_{k}B/K$ as the initial synchronization period $\tau_{0}$, both of which can be easily obtained during training.

The proposed FedAIS is illustrated in Algorithm 1. Specifically, in the $t$-th global round, the server randomly selects a set $M_{t}$ of $m$ clients, and distributes the current model $\theta_{t}$ to them (Lines 4-6). Each selected client $k$ calculates the loss for each sample $i\in V_{k}$ and updates its selection probability $p_{k,i}^{t}$ (Lines 11-12). Then, during each local epoch $j$, client $k$ selects a batch $B_{k}$ of samples with $x_{i}\propto p_{k,i}^{t},i\in V_{k}$. For each layer of FedGCN, when the number of local batch training epoch $j$ satisfies $j\%\tau_{t}==0$, client $k$ firstly calculates $\tilde{h}_{i}^{(l+1)}$ with Eq. (6) by aggregating embeddings of both within-client neighbors and cross-client neighbors. Then, it performs embedding synchronization by asking neighbor client $q\in Q$ to update the selected cross-client neighbor embeddings and transmit them back (Lines 13-18). Then, client $k$ updates its local model $\theta_{j}^{k}$ and sends $\theta_{J}^{k}$ to the server. The server aggregates updates by conducting model aggregation and updating interval $\tau_{t+1}$ (Lines 7-8). In this way, the server and clients collaboratively train a FedGCN model $\theta^{*}$ with high prediction accuracy and low overhead.

Without loss of generality, we analyze an arbitrary interval sequence $\{\tau_{1},\cdots,\tau_{R}\}$ with $R$ synchronization iterations.

The key idea of proof is as follows. To understand the condition (12), we consider the case when $\tau_{0}=\cdots=\tau_{R}$ is a constant. Then, the converge condition is identical to that for mini-batch SGD: $\sum_{r=0}^{R}\eta_{r}\rightarrow\infty,\sum_{r=0}^{R}\eta_{r}^{2}<\infty$. Provided that the sequence of communication periods is bounded, the learning rate in mini-batch SGD can be easily adjusted to satisfy condition (12). Specifically, when the communication period sequence decreases, the last two terms in (12) become easier to satisfy, and the differences between the objective values of two consecutive rounds are bounded. The full proof of FedAIS convergence is presented in Appendix.

Implementation. We implemented FedAIS and deployed it in an FL system consisting of one server and 100 clients. To further investigate the performance of FedAIS in large-scale FL systems, we also tested it in an environment with up to 1,000 clients. Our implementation is based on Python 3.11 and Pytorch Geometric 2.0.1 (Fey and Lenssen 2019). All the experiments are performed on Ubuntu 20.04 operating system equipped with a 32-core AMD Ryzen Threadripper PRO 5965WX @ 3.800GHz CPU, 192G of RAM and a NVIDIA RTX A5000 GPU with 24GB memory.

Datasets and Models. We use five real-world graph datasets of different scales, i.e., Coauthor (Shchur et al. 2018), Pubmed (Sen et al. 2008), Yelp (Zeng et al. 2019), Reddit (Hamilton et al. 2017), Amazon2M (Hu et al. 2020). We partition training/ validation sets over 100 clients in both independent and identically distributed (iid) setting and non-iid setting. We use a non-iid partition by $p_{i}\sim Dir_{k}(\alpha)$, $\alpha=0.5$ with a Dirichlet distribution and allocate a $p_{i,k}$ ratio of instances of class $i$ to client $k$ (Li et al. 2022; Yurochkin et al. 2019). Since the original graph is extremely dense, we downsample the edges in local subgraphs by 50% (Hamilton et al. 2017). The test dataset is located at the server and the statistics of the datasets are presented in Table 1. We use the widely adopted GraphSage model (Hamilton et al. 2017) with FedAvg to construct FedGCN models: FedAuthor, FedPubmed, FedYelp, FedReddit, FedAmazon (Hu et al. 2020). Each model has two hidden layers with 256, 128 neurons. We set the ratio of sample selection to 0.7, the number of neighbors sampled to 10, and the minimum embedding synchronization interval $\tau_{0}$ to 2 batch training epochs. We use Adam as the optimizer with the weight decay 0.001 and ReLU as the activation function. We set the learning rate $\eta=0.001$, the fixed batch number is 10, and the global warm-up round is 1. We conduct training until a pre-specified test accuracy is reached, or a maximum number of iterations has elapsed (e.g., 100 rounds). We perform 5-fold cross validation and report the average results.

Comparison Baselines. 1) FedAll: It conducts training using all local samples and performs random neighbor node selection of both local subgraph neighbors and cross-client neighbors. 2) FedRandom: It performs random selection for both local samples and neighbor nodes in each batch training. 3) FedSage+ (Zhang et al. 2021b): It proposes a GNN-based neighbor generative model for each client to predict the features of each node’s cross-client neighbors. 4) FedPNS (Du and Wu 2022): It conducts training using all local samples and conducts periodic neighbor node selection for cross-client neighbor nodes. We set the periodic interval to 2 local batch training epochs. 5) FedGraph (Chen et al. 2021): It performs FL training using all local samples and selects local subgraph neighbors and cross-client neighbors by adjusting sampling policies based on DRL.

FedAIS achieves comparable or higher accuracy. We adopt three metrics (Falessi et al. 2021; Herbold et al. 2018), i.e., test accuracy, F1-score, and area under the curve (AUC), to evaluate the accuracy of FedGCN models. We compare FedAIS with other baseline methods by training different FedGCN models in both iid and non-iid settings. We present the results of those accuracy metric scores and the standard deviations of the final global models in Table 2. It shows that FedAIS achieves test accuracy, F1 score, and AUC scores that are comparable to or better than other methods. For example, for model FedYelp trained on Yelp dataset in the iid setting, the test accuracy of FedAIS is 4.55%, 3.87%, 5.20%, 5.14 higher than other methods, respectively.

FedAIS significantly saves computation and communication costs. We present the test accuracy with the size of communication cost for training FedReddit and FedAmazon in iid settings in Fig. 4. The test accuracy, F1-score and AUC scores with the size of communication cost for training FedAuthor, FedPubmed and FedYelp in both iid and non-iid settings are much similar to that in Fig. 4. The results in Fig. 4 show that FedAIS requires much less amount of communication volume than the other baselines to achieve the target test accuracy, which leads to less computation time for transmitting those node embedding bytes and updating model parameters. Besides, we present the total computation and communication overhead for training FedGCN models in Fig. 4. It shows that FedAIS achieves significantly savings of both computation and communication costs than other baselines.

We perform ablation studies to show the effectiveness of each component of FedAIS. We compare FedAIS against the following ablation baselines: 1) FedAll; 2) FedAIS1: it only conducts the proposed dynamic importance sampling method of local samples without adaptive embedding synchronization; 3) FedAIS2: it conducts training using all local samples with the proposed adaptive embedding synchronization. We present the test accuracy with the size of communication cost and the total communication costs in Fig. 6. The results show that FedAIS, FedAIS1 and FedAIS2 achieve higher performance in saving much communication costs to reach the target accuracy scores and FedAIS performs the best among them. Thus, both the dynamic importance sampling module and the adaptive embedding synchronization module are effective to construct FedAIS.

Impact of the number of clients. We conduct experiments with different numbers of clients engagement, i.e., $K=100,300,500,700,1,000$. We present the test accuracy and communication cost of the model FedReddit trained with different number of clients in Fig. 6. It shows that the test accuracy of FedAIS is consistently high, i.e., above 75.0%, as the number of client increases to 1,000, and achieves test accuracy that is comparable to or higher than others. Besides, it shows that communication costs increase as the number of clients increases and FedAIS achieves substantial cost savings than other baselines in all settings.

Impact of the non-iid degree. We present the test accuracy of the model FedReddit with different non-iid degrees, i.e., $\alpha=0.05,0.1,0.5,1.0,10,100$, in Fig. 7(a). It shows that FedAIS achieves accuracy scores that are comparable to or higher than other baselines. Besides, these accuracy scores increase as $\alpha$ increases, and when $\alpha$ is greater than 0.5, the accuracy scores are relative high. The accuracy scores with the sizes of communication cost for training the model FedReddit with different non-iid degrees are much similar to that in Fig. 3(a), which shows that FedAIS consistently requires less communication cost than others.

Impact of the ratio of samples selected. Fig. 7 presents the test accuracy and communication costs of the model FedReddit when the selection ratio is $r=$0.1, 0.3, 0.5, 0.7, 0.9. Here, we adjust FedPNS and FedGraph so that they can select corresponding ratios of nodes. It shows that both the test accuracy and communication cost increase as the sampling ratio increases and FedAIS performs much better. Besides, due to the large size of the Reddit dataset, FedAIS can construct a FedGCN model with high test accuracy and less communication cost by sampling only 0.1 proportion of local samples, which achieves significantly more advantageous trade-offs between accuracy and efficiency.