Semi-Supervised Federated Learning with non-IID Data: Algorithm and System Design

Semi-supervised federated learning attempts to use semi-supervised learning techniques [19, 20, 21, 22, 23] to further improve the performance of the FL model in scenarios where there is unlabeled data on the client side [11]. For example, Long et al. in [10] proposed a semi-supervised federated learning (SSFL) system, FedSemi, which unifies the consistency-based semi-supervised learning model [24], dual model [15], and average teacher model [25] to achieve SSFL. The DS-FL system [26] was proposed to solve the communication overhead problem in SSFL. Reference [27] proposes a method to study the distribution of non-IID data, which introduces a probability distance metric to evaluate the difference in client data distribution in SSFL. Different from the literature [9, 10, 26], in this paper, we focus on labels-at-server scenario and also solve the problem of data availability and data heterogeneity in SSFL.

If the local data set distribution of each client is inconsistent (i.e., non-IID problem) [28, 29, 30, 31, 18, 32, 33], the local objective loss function of the client will be inconsistent with the global objective[34]. In particular, when the model of the local client is updated larger, such a difference will be more obvious. Therefore, we need to design some robust FL systems to solve the above problems. Some studies try to design a robust federated learning algorithm to solve the non-IID problem. For example, FedProx [6] limits the distance between the local model and the global model by introducing an additional $\mathcal{L}_{2}$ regularization term in the local target function to limiting the size of the local model update. However, this method has a disadvantage that each client needs to individually adjust the local regularization term to obtain good model performance. FedNova [35] improved FedAvg in the aggregation phase, which normalized and scaled the model update according to the local training batch of the client. Although previous studies have alleviated the problem of non-IID to some extent, they only evaluated the data distribution at specific non-IID levels and lacked extensive experimental verification for different non-IID scenarios. Therefore, we propose a more comprehensive data distribution and data partition strategy, i.e., we introduce the Dirichlet distribution function to simulate different non-IID levels of client data.

Federated learning solves the problem of data island on the premise of privacy protection. In particular, the FL is a distributed machine learning framework, which requires clients to hold data locally, where these clients coordinate to train a shared global model $\omega^{\ast}$. In FL, there is a server $\mathcal{S}$ and $k$ clients, each of which holds an IID or non-IID datasets $\mathcal{D}_{k}$. Specifically, for a training sample $x$ on the client side, let $\ell(\omega;x)$ be the loss function at the client, where $\omega\in{\mathbb{R}^{d}}$ denotes the model’s trainable parameters. Therefore, we let $\mathcal{L}(\omega)=\mathbb{E}_{x\sim\mathcal{D}}[\ell(\omega;x)]$ be the loss function at the server. Thus, FL needs to optimize the following objective function at the server:

where ${p_{k}\geq 0},\sum\limits_{k}{{p_{k}}}=1$ indicates the relative influence of $k$-th client on the global model. In FL, to minimize the above objective function, the server and clients execute the following steps:

• •

  Step 1, Initialization: The server sends the initialized global model $\omega_{0}$ to the selected clients.

• •

  Step 2, Local training: The client uses the local optimizer (e.g., SGD, Adam) on the local dataset $\mathcal{D}_{k}$ to train the received initialization model. Then, each client uploads the local model $\omega_{t}^{k}$ to the server.

• •

  Step 3, Aggregation: The server collects and uses a certain algorithm (e.g., FedAvg [1]) to aggregate the model updates uploaded by these clients to obtain a new global model, i.e., $\omega_{t+1}$ = $\omega_{t}$ + $\sum\limits_{k=1}^{K}\dfrac{{D}_{k}}{D}\omega_{t}^{k}$. Then, the server sends the updated global model $\omega_{t+1}$ to all selected clients.

Note that FL repeats the above steps until the global model converges.

In the real world, (e.g., financial and medical fields), unlabeled data is easy to gain, while labeled data is often difficult to obtain. Meanwhile, annotating data requires a lot of manpower and material resources. To this end, the researchers proposed a machine learning paradigm namely Semi-supervised Learning [36, 37] to learn a high-precision model on a mixed dataset (part of the data is labeled, and some of the data is unlabeled). Thus, in recent years, semi-supervised learning has become a hot research direction in the field of deep learning. In this section, we introduce a basic assumption and two methods of semi-supervised learning.

According to the above assumption, we adopt two common semi-supervised learning methods as follows:

Consistency Regularization: The main idea of this method is that the model prediction results should be the same whether noise is added or not on an unlabeled training sample [15, 36]. We generally use data augmentation (such as image flipping and shifting) methods to add noise to increase the diversity of the dataset. Specifically, for an unlabeled sample $u_{i}$ in unlabeled dataset $\mathbf{u}=\{u_{i}\}_{i=1}^{m}$ and its perturbation form $\hat{u_{i}}$, our goal is to minimize the distance $d(f_{\theta}(u_{i}),f_{\theta}(\hat{u_{i}}))$, where $f_{\theta}(u_{i})$ is the output of sample $u_{i}$ on model $\theta$. The common distance measurement method is Kullback-Leiber (KL) divergence. Thus, we can calculate the consistency loss as follows:

where $m$ is the total number of unlabeled samples and $f_{\theta}(u_{i})$ indicates the model output of unlabeled sample $u_{i}$.

Pseudo-label: The pseudo-label method [24] is to utilize some labeled samples to train a model that to set the pseudo labels for unlabeled samples. Previous work generally used sharpening [39] and argmax [24] methods to set pseudo-label, where the former make the distribution of model output extreme of unlabeled samples and the latter change the model output to one-hot of unlabeled samples.

In our system setting, the server $\mathcal{S}$ holds a labeled dataset $\mathcal{D}_{s}=\{x_{i},y_{i}\}_{i=1}^{n}$, where $n$ indicates the number of labeled samples. For $K$ clients, we assume that $k$-th client holds a local unlabeled dataset $\mathcal{D}_{k}=\{u_{i}\}_{i=1}^{m}$, where $m$ denotes the number of unlabeled samples in the local client. Similar to the traditional FL system, the server and clients in the SSFL are cooperative to train a high-performance global model $\omega^{*}$. The goal of previous work is to optimize the objective function mentioned above, i.e., Equation (5). However, they ignore the implicit contribution between iterations of the global model, which results in the learned global model not being optimal. Inspired by the above facts, we propose an SSFL algorithm called $\mathrm{FedMix}$ that focuses on the implicit contributions between iterations of the global model in a fine-grained manner. We define the supervised model trained on the labeled dataset as $\sigma$, the unsupervised model trained on the unlabeled dataset as $\psi$, and the aggregated global model as $\omega$. Specifically, we design a strategy that assigns three weights $\alpha,\beta,$ and $\gamma$ to the unsupervised model $\psi$, supervised model $\sigma$, and the previous round of global model, respectively. The designed algorithm can capture the implicit relationship between each iteration of the global model in a fine-grained manner. Thus, the steps of our proposed $\mathrm{FedMix}$ algorithm are as follows:

• •

  Step 1, Initialization: The server randomly selects a certain proportion of $F$ ($0<F<1$) clients from all local clients to send the initialized global model $\omega_{0}$. Note that the global model $\omega_{0}$ also remains on the server-side.

• •

  Step 2, Server Training: Unlike FL, in our SSFL system, the server not only aggregates the model uploaded by the clients, but also trains the supervised model $\sigma$ (i.e., $\sigma_{t}\leftarrow{\omega_{t}}$) on the labeled dataset $\mathcal{D}_{s}$. Thus, the server uses the local optimizer on the labeled dataset $\mathcal{D}_{s}$ to train the supervised model $\sigma$. The minimization of the objective function is defined as follows:

  $$\mathop{\min}\limits_{\sigma\in{\mathbb{R}^{d}}}{{\cal L}_{s}}({\sigma_{t}}),\text{where}\quad{{\cal L}_{s}}({\sigma_{t}})\buildrel\textstyle.\over{=}{\lambda_{s}}CE(y,{f_{{\sigma_{t}}}}(x)),$$

  (6)

  where $\lambda_{s}$ is the hyperparameter, ${x}$ and ${y}$ are from labeled dataset $\mathcal{D}_{s}$, and $f_{\sigma_{t}}({x})$ means the output of labeled samples on supervised model $\sigma$ at $t$-th training round.

• •

  Step 3, Local Training: The $k$-th client utilizes the local unlabeled data to train the received global model $\omega_{t}$ (i.e., $\psi_{t}^{k}\leftarrow{\omega_{t}}$) and then obtains the unsupervised model $\psi_{t+1}^{k}$. Thus, we define the following objective function:

  	$\displaystyle\mathop{\min}\limits_{\psi\in{\mathbb{R}^{d}}}{{\cal L}_{u}}(\psi_{t}^{k}),{{\cal L}_{u}}(\psi_{t}^{k})\buildrel\textstyle.\over{=}{\lambda_{2}}||{f_{\psi_{t}^{k}}}({\pi_{1}}(u))-{f_{\psi_{t}^{k}}}({\pi_{2}}(u))|{|^{2}}$		(7)
  	$\displaystyle+{\lambda_{1}}CE(\widehat{y},{f_{\psi_{t}^{k}}}(u))+{\lambda_{L1}}||{\sigma_{t}}-\psi_{t}^{k}|{|^{2}},$

  where $\lambda_{1},\lambda_{2},$ and $\lambda_{L1}$ are hyperparameters to control the ratio between the loss terms, $\psi_{t}^{k}$ is the unsupervised model of the $k$-th client at $t$-th training round, ${u}$ is from unlabeled dataset $\mathcal{D}_{k}$, $\pi(\cdot)$ is the form of perturbation, i.e., $\pi_{1}$ is the shift augmentation, $\pi_{2}$ is the flip augmentation, $||\sigma_{t}-\psi_{t}^{k}||^{2}$ is a penalty term that aims to let the $k$-th client unsupervised model $\psi_{t}^{k}$ learn the knowledge of the supervised model $\sigma_{t}$ (note that $\sigma_{t}$ is inferred from Equation (9)), and $\hat{y}$ is pseudo label obtained by using our proposed argmax method. The argmax method is defined as follows:

  $$\hat{y}=\mathsf{\textbf{1}}(\mathrm{Max}(\sum\limits_{i=1}^{A}f_{\psi_{t}^{k}}(\pi_{i}(u)))),$$

  (8)

  where $\mathrm{Max}(\cdot)$ is a function that can output the maximum probability that unlabeled data belongs to a certain class, $\mathsf{\textbf{1}}(\cdot)$ is the one-hot function that can change the numerical value to 1, $A$ represents the number of unlabeled data after data augmentation, and ${u}$ is from unlabeled dataset $\mathcal{D}_{k}$. Specifically, we discard low-confident predictions below confidence threshold $\tau=0.80$ when generating pseudo-labels.

  

• •

  Step 4, Aggregation: The server uses the proposed $\mathrm{FedFreq}$ (see Section IV-B) aggregation algorithm to aggregate the unsupervised models uploaded by the clients to obatin the global unsupervised model, i.e.,$\psi_{t+1}=\sum\limits_{k=1}^{K}w_{t+1}^{k}\psi_{t+1}^{k},$ where $\psi_{t+1}^{k}$ is the unsupervised model of the $k$-th client at $t+1$-th training round and $w_{t+1}^{k}$ is the weight of the $k$-th client. The server then aggregates the global unsupervised model $\psi_{t+1}$, the supervised model $\sigma_{t+1}$, and the global model $\omega_{t}$ from the previous round $t$ to obtain a new global model $\omega_{t+1}$:

  $$\omega_{t+1}=\alpha\psi_{t+1}+\beta\sigma_{t+1}+\gamma\omega_{t},$$

  (9)

  where $\alpha$, $\beta$, and $\gamma$ are the corresponding weights of the three models and $(\alpha,\beta,\gamma)\in\{\alpha+\beta+\gamma=1\wedge\alpha,\beta,\gamma\geqslant 0\}$.

Repeat all the above steps until the global model converges. The proposed $\mathrm{FedMix}$ algorithm is shown in Algorithm 1.

In this section, we present the designed $\mathrm{FedFreq}$ aggregation algorithm, which can dynamically adjust the weight of the corresponding local model according to the training frequency of the client to alleviate the non-IID problem. We observe that the parameter distribution of the global model will be biased towards clients that often participate in federated training, which is obviously not friendly to the robustness of the global model. Therefore, our insight is to reduce the influence of clients with high training frequency on the global model to improve the robustness of the model. Thus, the formal expression of the $\mathrm{FedFreq}$ aggregation algorithm is as follows:

where $F$ is the sample proportion of the server, $K$ is the total number of clients, $p_{t+1}^{k}=\dfrac{q_{t+1}^{k}}{\sum\nolimits_{k\in S_{t+1}}q_{t+1}^{k}}$, $q_{t+1}^{k}$ is the number of times that the $k$-th client has been trained up to the $t+1$-th round, and $S_{t+1}$ denotes the set of clients selected by the server in round $t+1$. Then, for the client, the $\mathrm{FedFreq}$ aggregation rule is expressed as follows: $\psi_{t+1}=\sum\limits_{k=1}^{K}w_{t+1}^{k}\psi_{t+1}^{k}.$

To better evaluate the robustness of the designed system to non-IID data, in this paper, we introduce the Dirchlet distribution function [40, 41], which is a popular non-IID function, to adjust the non-IID level of the local client data. Specifically, we generate data distributions of different non-IID levels by adjusting the parameter (i.e., $\mu$) of the Dirchlet distribution function. We assume that the local dataset $\mathcal{D}_{k}$ of $k$-th client has $c$ classes, and thus, the definition of Dirichlet distribution function is as follows:

where $\Theta$ is a set of $c$ samples randomly selected from the Dirichlet function, i.e., $\Theta=\{{\varphi_{1}},\ldots,{\varphi_{c}}\}$ and $\Theta\sim Dir(\mu_{1},\cdots,\mu_{c})$, $\mu,\mu_{1},..,\mu_{c}$ are the parameters of the Dirichlet distribution function (where $\mu=\mu_{1}=\mu_{2}=...=\mu_{c}$), and $p_{k}(\varphi_{c})$ denotes the proportion of the $c$-th class data in all data of the client. In particular, the smaller the $\mu$, the higher the non-IID level of the data distribution of each client; otherwise, the data distribution of the client tends to the IID setting. Therefore, we adjust the parameters of the Dirichlet distribution function to simulate the different non-IID levels of the client’s local dataset. For example, as shown in Fig. 3, we demonstrate the data distribution when $\mu=\{0.1,1,10\}$.

Dataset 1) CIFAR-10 dataset under IID setting: We use the CIFAR-10 dataset including 56,000 training samples and 2,000 test samples as the validation dataset in our experiment. The training set includes 55,000 unlabeled samples and 1,000 labeled samples, where the former is used to train the unsupervised model at the local and the latter is used to train the supervised model at the server. The unlabeled samples are equally distributed to 100 clients at a ratio of $1:100$, of which there are 550 samples for each client (i.e., 55 for each class, and 10 classes in total). Similarly, the labeled samples have a total of 1,000 and contain 10 classes on the server, of which there are 100 samples in each class. Meanwhile, we set a participation rate $F=0.05$ of clients, i.e., 5 clients are randomly selected for training in each round.

Dataset 2) CIFAR-10 dataset under non-IID setting: Our setting is similar to the above IID setting, except that the Dirchlet distribution function is introduced to adjust the non-IID level of the local client data. Specifically, we generate data distributions of different non-IID levels by adjusting the parameter (i.e., $\mu$) of the Dirchlet distribution function. Meanwhile, we simulate quantity imbalance and class imbalance the local client data. In particular, we make each client hold a different number of training samples. For example, some customers have 580 samples, while some customers have less than 50 samples. Second, we make the client hold a different sample size for each category in the data. For example, some clients have ten types of data, and some clients have less than two types of data.

Baseline and training details: Our baseline is $\mathrm{FedMatch}$ [12] naively using unsupervised model and supervised model parameter decomposition strategy, i.e., $\omega=\psi+\sigma$. In the training process, both our model and baseline use Stochastic Gradient Descent (SGD) to optimize the ResNet-9 neural network with initial learning rate $\eta=1e-3$. We set training round $t=150$, the number of labeled samples on sever is $N_{s}=1000$, local client training epoch $E_{client}=1$ and mini-batch size $B_{client}=64$, the server training epoch $E_{server}=1$ and mini-batch size $B_{server}=64$. Second, we set the data augmentation number in the argmax method $A=5$.

As shown in Fig. 4, under IID and non-IID settings, our method $\mathrm{FedMix}$ is better than baseline under each different aggregation method settings. For example, under a non-IID setting, the convergence accuracy of our method is 47.5% about 3% higher than that of the baseline. In particular, the accuracy of our method increases faster and more stable in the early stage of model training. The reason is that: (1) The $\mathrm{FedMix}$ focuses on the implicit contributions between iterations of the global model in a fine-grained manner, while the $\mathrm{FedMatch}$ only naively uses model parameter decomposition. (2) Frequency-based aggregation method $\mathrm{FedFreq}$ is more suitable for non-IID settings. Notably, the $\mathrm{FedFreq}$ only requires the server to give appropriate weights in the aggregation process according to the training frequency of each client, which does not bring additional computational overhead to the server and local clients.

Fig. 5 shows the performance comparison of the proposed method under different hyperparameter settings. To be specific, the three hyperparameters are the weights of the global unsupervised model, the supervised model, and the previous round of global model. From Fig. 5, we can find that under the non-IID setting, as $\alpha$ decreases, the accuracy curve of the proposed method becomes unstable. The reason for this phenomenon is that: with the decrease of global unsupervised model weight, $\mathrm{FedFreq}$ is losing its effect. In particular, when $\alpha=0.5$, $\beta=0.3$, and $\gamma=0.2$, our global model is the aggregation of the optimal weights of the three models. In addition, we find that the proposed method is easy to achieve better performance when these three parameters are relatively close under the IID setting, as shown in Fig. 5.

Fig. 6 shows the performance comparison of the proposed method on different non-IID levels of client data. In this experiment, we let $\mu=0.1$ denote the highest non-IID level of the client data. In this case, as the value of $\mu$ increases, the local client data distribution is closer to the IID setting. It can be seen from Fig. 6 that for different non-IID levels, our method all can achieve stable accuracy. Meanwhile, the model convergence accuracy under $\mu=\{0.1,1,10,100\}$ settings does not differ by more than 1%. Therefore, our method is not sensitive to the different levels of client data distribution, i.e., it is robust to different types of data distribution settings.

Fig. 6 shows the performance comparison of the proposed method in the case of different numbers of labeled samples at the server. Obviously, the converged accuracy of our method is 47% with 800 labeled samples, which is 2% higher than $\mathrm{FedMatch}$. However, when the number of labeled samples is reduced to 700, the accuracy of our model decreases greatly. Therefore, we regard $N_{s}=800$ as the best setting for our method.

In this section, we further analyze the advantages of $\mathrm{FedMix}$ compared to $\mathrm{FedMatch}$ in labels-at-server scenario.

1) The performance of $\mathrm{FedMix}$ training model under the CIFAR-10 dataset is better than $\mathrm{FedMatch}$. This is due to $\mathrm{FedMatch}$ simply uses the strategy of parameter decomposition of unsupervised model and supervised model in the training process, i.e., $\omega_{t}=\psi_{t}+\sigma_{t}$. In this way, the learned global model will be biased towards unlabeled data (unsupervised model) or labeled data (supervised model) instead of the overall data. Thus, in order to avoid the drift problem of the global model, $\mathrm{FedMix}$ adds the global model from the previous round to the model parameter aggregation, i.e., $\omega_{t}=\alpha\psi_{t}+\beta\sigma_{t}+\gamma\omega_{t-1}$. Meanwhile, we conducted a sensitivity experiment of model performance to different hyperparameter weights to find the optimal weight combination.

2) $\mathrm{FedMix}$ is robust to different levels of non-IID data. In our experiment, we introduced the Dirichlet distribution function to simulate the local client non-IID data in FL. In details, we generate data distributions of different non-IID levels by adjusting the parameters of the Dirchlet distribution function, i.e., $\mu=\{0.1,1,10,100\}$ respectively correspond to different levels of non-IID. The results show that the performance difference of our model does not exceed 1% under different levels of non-IID settings. $\mathrm{FedMatch}$ uses a pseudo-random method to generate the non-IID data distribution of each client. However, in reality there is no such distribution, which will cause the model to lose its robustness.