GraphFL: A Federated Learning Framework for Semi-Supervised Node Classification on Graphs

Suppose we are given a set of $I$ clients $\mathbb{C}=\{C^{(1)},C^{(2)},\cdots,C^{(I)}\}$, where each client $C^{(i)}$ owns a graph ¹¹1We interchangeably use client and client graph in this paper. $G^{(i)}=(\mathbb{V}^{(i)},\mathbb{E}^{(i)})$ with $\mathbb{V}^{(i)}$ the node set and $\mathbb{E}^{(i)}$ the edge set. Each node $v^{(i)}\in\mathbb{V}^{(i)}$ is associated with a feature vector $\mathbf{x}_{v^{(i)}}$ and a label $y_{v^{(i)}}$ among the label set $\mathbb{K}=\{1,2,\cdots,K\}$. Moreover, each client graph $G^{(i)}$ has a set of a few labeled nodes $\mathbb{L}^{(i)}\subset\mathbb{V}^{(i)}$. We assume each client $C^{(i)}$ can learn a local semi-supervised node classification model $f_{\theta^{(i)}}$, parameterized by $\theta^{(i)}$, based on its labeled set $\mathbb{L}^{(i)}$ and graph $G^{(i)}$. We also consider a server $\mathbb{S}$ that can learn global model parameters $\theta$ by aggregating the clients’ local model parameters $\{\theta^{(i)}\}_{i=1}^{I}$, while not accessing the client graphs. Now, suppose we have a set of testing nodes $\mathbb{T}$ which may or may not have the same label domain as the training nodes. Then, our problem of federated semi-supervised node classification on graphs is defined as follows:

Given a set of training tasks $\{{T}_{i}\}$ drawn from an underlying task distribution $\mathcal{T}$, instead of learning a model that performs well on all tasks, MAML [8] learns a task-independent initialization $\theta$ that performs well on all tasks after a few steps of gradient updates. Specifically, each task $T_{i}$ splits its labeled training set $\mathbb{L}^{(i)}$ into a support set $\mathbb{L}_{S}^{(i)}$ and a disjoint query set $\mathbb{L}_{Q}^{(i)}$. Then, MAML has a two-level optimization: inner-optimization and meta-optimization. In the inner-optimization, for each task ${T}_{i}$, MAML trains a model $f_{\theta}$ on the support set $\mathbb{L}_{S}^{(i)}$ with the initialization $\theta$ and outputs a task-specific model parameters $\theta^{(i)}$. MAML then takes each $\theta^{(i)}$ as the initialization and evaluates the model $f_{\theta^{(i)}}$ on the corresponding query set $\mathbb{L}_{Q}^{(i)}$ with a task-specific loss. In the meta-optimization, MAML minimizes the total loss on the query sets of all tasks simultaneously to learn the task-independent initialization. Formally, the objective function of MAML is as follows:

where we use one step of gradient descent in the inner-optimization for simplicity; $\alpha$ is the learning rate; and the task-specific loss over the support set and the query set are respectively defined as

where $\ell(\cdot)$ is the loss function defined for the specific task.

The objective function of MAML in Equation 1 is solved via gradient descent with a meta-learning rate $\beta$, i.e.,

MAML proceeds in an episodic manner, where in each episode a batch of tasks are sampled from the task distribution $\mathcal{T}$ for training. When a new task arrives, MAML uses the learnt task-independent initialization as the initial model and updates the model via a few steps of gradient descent with respect to the loss defined in the new task. Then, the updated model is used to make predictions.

One key challenge to solving federated GraphSSC is that graph data are essentially and potentially highly non-IID across clients, while existing FL methods perform poorly with non-IID data [22, 41, 19]. This is because the goal of existing FL is to collaboratively learn a global model for all clients, but the learnt global model cannot generalize well on clients’ data when they are non-IID. We design a novel FL method GraphFL for federated GraphSSC that can handle the non-IID graph data across clients. Our GraphFL method is inspired by MAML, as it fits into the FL framework and trains a set of tasks simultaneously but does not require the data to be IID across the tasks. Specifically, our GraphFL method incorporates MAML into FL and learns a global model on the server that shows good generalization on testing nodes. GraphFL consists of two stages: Stage I learns a global model on the server by following the training scheme of MAML and thus can mitigate the issue caused by non-IID graph data. Stage II leverages the existing FL method to further update the global model such that it can achieve a good generalization ability.

We first define some notations. In each client $C^{(i)}$, we split the training set $\mathbb{L}^{(i)}$ into the support nodes $\mathbb{L}_{S}^{(i)}$ and the query nodes $\mathbb{L}_{Q}^{(i)}$. Suppose at episode $t$, the server $\mathbb{S}$ has global model parameters $\theta_{t}$, and each client $C^{(i)}$ has local model parameters $\theta_{t}^{(i)}$. We define the loss on the support nodes $\mathbb{L}_{S}^{(i)}$ in client $C^{(i)}$ as ${\mathcal{L}_{\mathbb{L}_{S}^{(i)}}(\theta_{t})=\frac{1}{|\mathbb{L}_{S}^{(i)}|}\sum_{v^{(i)}\in\mathbb{L}_{S}^{(i)}}\ell(f_{\theta_{t}}(\mathbf{x}_{v}^{(i)},G^{(i)}),y_{v}^{(i)})}$, and the loss on the query nodes $\mathbb{L}_{Q}^{(i)}$ as ${\mathcal{L}_{\mathbb{L}_{Q}^{(i)}}(\theta_{t}^{(i)})=\frac{1}{|\mathbb{L}_{Q}^{(i)}|}\sum_{v^{(i)}\in\mathbb{L}_{Q}^{(i)}}\ell(f_{\theta_{t}^{(i)}}(\mathbf{x}_{v}^{(i)},G^{(i)}),y_{v}^{(i)})}$, where $f_{\theta_{t}}$ and $f_{\theta_{t}^{(i)}}$ are the GraphSSC models learnt on the support nodes and query nodes in client graph $G^{(i)}$, respectively. Moreover, suppose at episode $t$, the server has learnt the global model parameters $\theta_{t}$. Then, our GraphFL method has the following steps.

Existing FL targets the problems where all data in clients share the same label domain. However, real-world graphs are dynamically changing and new types of nodes can emerge at any time. In this section, we study the graph-based problems where the training nodes and testing nodes have different label domains. One possible solution is leveraging transfer learning. Specifically, we first learn a global model on the server based on the existing FL methods (e.g., FedAvg [22]). Next, we adopt the global model as the initial model and finetune the model using a few labeled nodes with new label domains. Then, we use the finetuned model to predict the labels of testing nodes with new label domains. However, as shown in our experimental results (See Table 2), such transfer learning-based solution achieves unsatisfactory performance.

We design a novel GraphFL method for GraphSSC that can generalize to testing nodes with new label domains. Specifically, we propose to reformulate MAML in the FL framework and aim to learn a shared global model on the server for all clients, such that each client performs well after a few steps of gradient updates with respect to its loss defined by a specific GraphSSC method. Formally, we define our objective function as follows:

where $\theta$ is the shared initialization we aim to learn; $\mathcal{L}_{i}(\theta)$ is the loss defined on client $C^{(i)}$. Note that our loss function is completely different from that in the existing FL methods [22, 41, 19].

We now solve Equation (9) in the FL framework. Specifically, we first update the local model based on our defined client loss, and then update the global model by aggregating local models. Assume at episode $t$, server has global model parameters $\theta_{t}$. Then, GraphFL has the following steps:

By solving Equation (9) with several episodes, we learn a global semi-supervised node classification model on the server that can fast adapt to testing nodes with new label domains. Specifically, we use the global model as the initial model, and update the model with a few steps of gradient descent using a few labeled nodes from new label domains. Then, we adopt the updated model to predict labels of testing nodes from new label domains. Algorithm 2 details our GraphFL for federated GraphSSC with new label domains.

Existing FL methods are mainly for supervised learning, i.e., they only use labeled data. However, real-world graph-based problems often involve limited labeled nodes, while having a substantial number of unlabeled nodes. A natural question is how can we leverage unlabeled nodes to further enhance the performance of our GraphFL?

We propose a self-training method to leverage unlabeled nodes in client graphs. Specifically, given a graph-based semi-supervised node classification method, we first use the method to train a local model in each client using the client’s few labeled nodes. Next, in each client, we use its local model to predict unlabeled nodes and select a set of unlabeled nodes with the most confident predictions. Then, we treat the predicted labels of the selected nodes as their pseudo labels, and add each client’s selected nodes (as well as their pseudo labels) to the client’s training set. Finally, we train our GraphFL methods on the augmented training sets for federated semi-supervised node classification. We observe in our experimental results (See Figure 5) that such a method is simple yet effective.

In this section, we evaluate our method and compare it with individual learning (IL), and FL for federated GraphSSC with non-IID graph data across clients.

In this section, we evaluate our GraphFL and compare it with FL for federated semi-supervised node classification with testing nodes having new label domains. As existing FL cannot handle nodes with new label domains, we adopt the transfer learning (TL) solution as described above. We only show the results on using different number of labeled nodes per class. Note that we also find similar observations as in aforementioned results on the impact of the fraction of participating clients and the fraction of overlapped nodes between client graphs.

Table 2 shows the node classification accuracy of GCN and SGC using FL and our GraphFL with new label domains vs. number of labeled nodes per class on the four graph datasets. We observe that our GraphFL method consistently outperforms FL, with at least a 10% higher accuracy in almost all cases. This indicates that transfer learning is not effective enough in handling testing nodes with new label domains for federated GraphSSC. In addition, when the number of labeled nodes per class is bigger, the performance gain of our GraphSSL over FL + TL is larger.

In this experiment, we study the effectiveness of adopting self-training to further enhance GraphFL’s performance.

Figure 5 shows the node classification accuracy of GCN and SGC using our GraphFL vs. number of pseudo labeled nodes per class generated via self-training (Due to space limitation, we do not show results with new label domains. However, we have similar observations. First, our method obtains a higher accuracy on all datasets when it uses the pseudo labeled nodes for training. This demonstrates that the pseudo labeled nodes are indeed beneficial to the federated training and thus can further enhance the performance of our method for semi-supervised node classification. One reason is that most of the predicted labels of pseudo labeled nodes match these nodes’ true labels. Second, as the number of pseudo labeled nodes per class increases, node classification accuracy first increases and then decreases. This is because, when the number of pseudo labeled nodes is relatively small, most of these nodes’ labels are correctly predicted via self-training. However, as the number becomes larger, wrong predicted labels of the pseudo labeled nodes also becomes larger. Table 3 also shows the fraction of the pseudo labeled nodes whose predicted labels are correct via GCN self-training on the four datasets. We note that these numbers justify our above claims.

• •

  Our method significantly outperforms conventional FL for federated GraphSSC in terms of handling non-IID graph data and generalizing to testing nodes with new label domains.

• •

  GraphFL can be further enhanced when unlabeled nodes are leveraged via self-training.

• •

  Increasing the number of training nodes, the fraction of participating clients, the graph size in clients, or leveraging unlabeled nodes all can all enhance the performance of federated node classification on graphs.