Estimating before Debiasing: A Bayesian Approach to Detaching Prior Bias in Federated Semi-Supervised Learning

1 Introduction

Federated Learning (FL) [McMahan et al., 2017] is a distributed learning paradigm that can facilitate collaborative model training among multiple clients while preserving data privacy. Presently, most FL methods are confined to supervised learning (SL) settings, wherein it is presumed that each client maintains a fully labeled dataset. Nevertheless, in realworld applications, data labeling is notably laborious and time-consuming. Therefore, a more realistic case involves Number of Samples each client possessing a mix of unlabeled and labeled data. This specific scenario, known as Federated Semi-Supervised Learning (FSSL), has been explored in various studies [Jeong et al., 2021; Lin et al., 2021; Diao et al., 2022] and is garnering increasing interest within the FL research community.

∗ Jianwei Niu is the corresponding author.

In this study, we focus on an FSSL setting where the data on each client are class-imbalanced. Moreover, it is assumed that both intra-client and inter-client data heterogeneity exist. Specifically, intra-client data heterogeneity implies that both the labeled data and unlabeled data on an individual client originate from diverse distributions. Inter-client data heterogeneity means that the overall distributions across clients are non-independent and identically distributed (Non-IID).

In the described scenario, the model's prediction can skew to some certain classes during the training, i.e., prediction bias. Figure 1 presents the experimental results conducted in the above scenario, where the overall distributions of labeled and unlabeled data are balanced. It can be observed that due to class imbalance in the local client, the local model's predictions gradually skew towards the major classes in the local data. More importantly, this prediction bias cannot be alleviated after model aggregation, even if the overall distributions are balanced. Instead, it evolves into a different form of bias due to the influence from other clients. This bias can disrupt the pseudo-labeling process, further creating a 'vicious cycle' between pseudo-labeling and local model training.

Existing FSSL methods attribute the above issue to the divergence across clients caused by heterogeneous data and primarily address it by promoting consistency between model parameters or outputs [Zhang et al., 2021; Jiang et al., 2022; Liang et al., 2022]. However, as both the local and global models are biased, merely constraining their consistency cannot fundamentally mitigate the model prediction bias.

In this paper, we delve into the essential reason for the prediction bias in FSSL from a Bayesian perspective. Based on Bayes' rule, the model prediction is as follows:

where p(y|x) is the model's prediction, p(x|y) is the class conditional likelihood, p(y) is the label prior. As shown in Figure 2, both the label prior of local labeled data (i.e., pl(y)) and unlabeled data (i.e., ˆ pu(y)) are biased. Consequently, the model can gradually absorb these biases during training. These biases are eventually injected into the global model through model aggregation, causing its output priors ps(y) to skew towards certain classes. When conducting inference on a balanced test dataset (i.e., pt(y)), the model may suffer severe performance degradation, as ps(y) ̸= pt(y).

Nevertheless, ps(y) is commonly challenging to estimate. On the one hand, in local clients, the ambiguity of pseudolabels for unlabeled data makes the label prior bias during local training intractable. On the other hand, in the server, model aggregation combines influences from participating clients, further complicating the estimation of prior bias.

Taking the class-wise accuracy on a balanced test dataset as the ground truth for prior bias, we find that the Average Prediction Probability of Unlabeled Data (APP-U) serves as a robust metric to approximate this bias. Figure 3 illustrates the Jensen-Shannon (JS) divergence [Lin, 1991] between the ground truth bias and either the labeled data distribution or APP-U, where the solid line and shaded area represent the mean and range across clients, respectively. Interestingly, it reveals that for both the global and local models, prior bias does not consistently align with the labeled data distribution. Rather, it shows a stronger correlation with APP-U, indicating that APP-U can effectively quantify the prior bias.

Building upon the above insights, we introduce a hierarchical debiasing method for FSSL termed FedDB, to mitigate the prior bias at both the local training and global aggregation stages. During the local training, FedDB implements debiased pseudo-labeling (DPL) based on Bayes' theorem, with APP-U serving as the approximation of bias prior. This approach promotes a more balanced pseudo-labeling process for unlabeled data, substantially reducing the label prior bias during local training. At the global aggregation stage, FedDB utilizes APP-U from the participating clients to determine optimal aggregation weights. The above process, termed debiased model aggregation (DMA), effectively mitigates bias within the global model. It should be noted that DPL can be seamlessly integrated with FSSL methods that utilize pseudolabeling with minimal cost. This demonstrates its substantial potential for practical application of FSSL.

The main contributions of this paper are as follows:

• We analyze the prediction bias in class-imbalanced FSSL from a Bayesian perspective.
• We propose FedDB, a Bayesian debiasing method for FSSL that uses APP-U as an approximation of prior bias.
• We conduct extensive experiments on multiple datasets to demonstrate the effectiveness of FedDB.

2 RelatedWork

2.1 Federated Learning

Data heterogeneity is a substantial challenge in FL, which can lead to considerable divergence across clients, thereby degrading the model performance [Zhao et al., 2018; Li et al., 2020a]. To address this issue, various strategies are explored, including reducing the divergence across local models [Li et al., 2020b; Acar et al., 2020; Karimireddy et al., 2020], enhancing aggregation schemes [Wang et al., 2020; Acar et al., 2020; Reddi et al., 2021], promoting representation consistency across clients [Tan et al., 2022; Zhu et al., 2023; Liao et al., 2023], developing personalized models for individual clients [Collins et al., 2021; Liu et al., 2023]. However, these methods primarily focus on SL settings, which is impractical as data labeling is laborious and time-consuming.

2.2 Semi-Supervised Learning

SSL aims to mitigate the reliance on labeled data, which prompts various mechanisms to leverage the latent information within unlabeled data. Pseudo-labeling [Lee and others, 2013; Wang et al., 2023], also known as self-training, involves assigning pseudo-labels to unlabeled samples with high confidence, enabling their incorporation into the training process. Consistency regularization [Miyato et al., 2018] introduces arbitrary perturbations to unlabeled samples and promotes the consistent predictions between different views of unlabeled data. Additionally, hybrid methods that amalgamate these approaches are also developed, such as MixMatch [Berthelot et al., 2019], FixMatch [Sohn et al., 2020]. Recently, SSL has focused on class imbalance, leading to various studies such as class-rebalancing sampling [Wei et al., 2021], and pseudo label sampling [Guo and Li, 2022]. However, simply combining these methods with FL is challenging, as they ignore the collaboration across clients.

2.3 Federated Semi-Supervised Learning

FSSL can be divided into three distinct scenarios [Bai et al., 2023]: (1) Labels-at-Partial-Clients, where only a few clients have full labels, while the rest possess only unlabeled data [Liang et al., 2022; Li et al., 2023]; (2) Labels-atServer, where labeled data are only available at the server, with local clients merely having unlabeled data [Zhang et al., 2021; Jeong et al., 2021; Diao et al., 2022]; (3) Labels-atClients, where each client has mostly unlabeled data and a few labeled samples [Jeong et al., 2021; Bai et al., 2023].

This paper focuses on the Labels-at-Clients scenario. Currently, several works have been proposed for this scenario. For instance, SemiFed [Lin et al., 2021] assigns pseudolabels to unlabeled data only when multiple models provide consistent predictions. FedMatch [Jeong et al., 2021] enforces prediction consistency across multiple models. However, these methods primarily concentrate on encouraging consistency across clients, overlooking the inherent prior biases within the model - a critical factor leading to performance degradation in FSSL with class imbalance.

3 Preliminary and Background

In this section, we present the notations used in this paper, followed by a detailed discussion of the framework of FSSL.

3.1 Problem Setting and Notation of FSSL

We focus on a FSSL setting for K-class classification task with totally M clients participating in the training. Each client m maintains a labeled dataset Dm l = {(xn , yn )}N m l n=1 and an unlabeled dataset Dm u = {(xn )}N m u n=1 , where N m l and N m u are the counts of labeled and unlabeled samples, respectively (typically, N m u ≫ N m l ), xn ∈ X ⊆ Rd is the input sampled from a d-dimensional space, yn ∈ Y ⊆ {0, 1}K is the one-hot label. For clarity, we sometimes omit the superscript denoting the client index in the following contents.

With a slight abuse of notation, we denote N k l and N k u as the numbers of samples in class k under Dl and Du for an arbitrary client, i.e., ∑Kk=1 N k l = Nl and ∑Kk=1 N k u = Nu. In this paper, we assume that both Dl and Du exhibit class imbalance, that is, ∃i, j ∈ {1, 2, . . . ,K} for which the ratio N i l N j l is significantly greater than 1. In other words, the label prior distribution {p1l , . . . , pKl } shifts from a uniform distribution { 1K }K . This assumption is similarly applicable for Du.

Furthermore, we consider the setting that both intra-client and inter-client data heterogeneity exist in the FL system. Intra-client heterogeneity refers to the varied distributions of labeled and unlabeled data within a single client, that is,

∀m ∈ {1, 2, . . . ,M}, Dm u ̸= Dm l . Inter-client heterogeneity, on the other hand, pertains to the dissimilar mixture distributions of both labeled and unlabeled data across clients, i.e., ∀i, j ∈ {1, 2, · · ·M}, i ̸= j, it holds that Di l +Diu ̸= Dj l +Dju .

The final objective of FSSL is to learn a global model f(x;w) : X → Y parameterized by w that can generalize well to a balanced test dataset whose label prior distribution is { 1K }K . Given the input xn , we denote its corresponding output logits as z(xn ) := f(xn ;w), and the normalized prediction probability after softmax layer as p(y|xn ) := σ(f(y|xn ;w)), where σ(·) is the softmax function. The detailed framework of FSSL is explained as follows.

3.2 Framework of FSSL

During each global round, the server first selects a random subset of clients S based on the activation rate C and broadcasts the global model w to these clients. Subsequently, these clients perform local training for E epochs using w as initial weights, resulting in the updated local model wm. Finally, the selected clients upload their local models wm to the server for model aggregation. The training paradigms of labeled and unlabeled data in local clients are as follows.

For labeled data, the standard cross-entropy loss is applied to the weakly augmented version of samples to promote the discriminative objective, as shown below:

where α(·) is the weak augmentation function, p(y|α(xn )) is the prediction probability for α(xn ), and H(p1, p2) is entropy between probability distributions p1 and p2.

For unlabeled data, the samples are pseudo-labeled using the trained model, after which they are incorporated into the training process. Specifically, for a given unlabeled sample xn , the model first generates the probability on its weakly augmented version. Then the pseudo-label is calculated by:

where argmax(·) is the function that converts a probability distribution into a one-hot label based on its maximum value.

To enhance the model generalization, the consistency loss is applied to unlabeled data by minimizing the entropy between the pseudo-label and the prediction of its strong augmented version.

During the training, only those unlabeled samples that exhibit high confidence are selected to participate in further training. Consequently, the overall optimization objective for the unlabeled data can be expressed as follows:

where τ is the threshold, 1(·) is the indicator function, A(·) is the strong augmentation function.

The overall optimization objective of local training on clients is expressed as:L = Ls + λLu, (6)

where λ is used to balance these two loss terms.

After local training, the selected clients send the latest local models to the server for model aggregation, as shown below:

where |St| is the number of selected clients in round t, wtm is the local model on client m in round t, βm is the aggregate weight for wtm, wt+1 is the global model in round t + 1.

4 FedDB: Detaching Prior Bias in FSSL

This section details the framework of FedDB and its two key techniques: debiased pseudo-labeling (DPL) and debiased model aggregation (DMA).

4.1 Framework Overview of FedDB

Figure 4 illustrates the framework of FedDB. During the training, each global round consists of the following steps:

1. The server selects a subset of clients for training and broadcasts the global model to these clients;
2. The clients perform inference on unlabeled data and calculate APP-U to estimate the prior bias;
3. The clients performDPL on unlabeled data usingAPP-U;
4. The clients train the model utilizing both labeled data and pseudo-labeled data;
5. The clients upload local models and APP-U to the server. The server performs DMA using APP-U from clients;
6. Repeating steps 1-5 until the global model converges.

4.2 Prior Bias Estimation

In this paper, we consider an FSSL setting where both the labeled data and unlabeled data are class imbalanced. In such a case, the model's predictions can skew towards certain classes, owing to the biased label prior in the training data. This skew contradicts the training objective of FSSL, which is to achieve uniform performance across all classes.

To investigate the impact of class imbalance on model training in FSSL, we conduct preliminary experiments using the CIFAR10 dataset. We establish a scenario with 10 clients, each participating in model training in every round. The number of labeled and unlabeled samples is set to 4000 and 46000, respectively. The class imbalance is created by the Dirichlet distribution, as declared in Section 5.

As shown in Figure 1, both the local and global models exhibit a biased prediction towards certain classes. However, estimating the above bias in FSSL is challenging due to the data heterogeneity and imprecision in pseudo-labeling. By extensive experiments, we discover that the prior bias can be effectively approximated by the Average Prediction Probability on Unlabeled Data (APP-U). Specifically, for client m, the APP-U, denoted by pm, can be calculated by:

where N m u denotes the total number of unlabeled samples on client m, p(y|α(xnu )) is the prediction probability of the weak augmentation of sample xnu .

We adopt JS divergence as a metric to quantify the disparity between two distributions. A larger JS divergence indicates a greater disparity between the distributions. Taking the class-wise accuracy on a balanced test dataset as the ground truth bias, we calculate the JS divergence between it and either APP-U or the labeled data distribution. As shown in Figure 3, for both local and global models, the JS divergence between APP-U and the ground truth is significantly smaller than that between the labeled data distribution and the ground truth. This demonstrates the effectiveness of APP-U as a metric for quantifying prior bias in FSSL.

4.3 Debiased Pseudo-Labeling

In this subsection, we detail the procedure of DPL. Given an FL model parameterized by w, we first obtain the prediction probability ps(y|x) by applying a softmax function to unnormalized logits, as illustrated below:

where z(x)[y] is the y-th unnormalized logit.

By applying the Bayes' theorem to ps(y|x), we obtain:

Due to the class imbalance in our FSSL settings, the prior distribution ps(k), as outputted by the model, is biased towards certain majority classes. This leads to a biased prediction probability ps(y|x), causing the model to be overconfident in these majority classes. The objective of DPL is to

Algorithm 1: DPL: Debiased Pseudo-labeling

Input: Confidence threshold τ Output: Debiased pseudo-labels ˆ Y , APP-U p 1 p = ∑Nu n=1 p(y|α(xnu )) Nu , ˆ Y := {} 2 for n = 1, 2, ...,Nu do 3 ˆ pn := p(y|α(xnu ))/p ∑Kk=1 p(k|α(xnu ))/pk 4 if max( ˆ pn ) ≥ τ then 5 ˆ Y := ˆ Y ⊕ arg max( ˆ pn ) 6 else 7 ˆ Y := ˆ Y ⊕ {0}K 8 Return ˆ Y , p

seek a conditional probability pt(y|x) that is robust across all classes, given the estimation of the model's biased prior p, as defined in Eq. (8).

Following previous studies [Tian et al., 2020; Kairouz et al., 2021; Hong et al., 2021], we assume that the class conditional likelihoods are the same in both the biased and debiased predictions, i.e., pt(x|y) = ps(x|y). By rearranging Eq. (9) and Eq. (10), we have:

Recalling that:

We derive the following debiased posterior probability:

where pt(k) is a uniform distribution that is robust for all classes. By applying the estimated bias p as the approximation of the prior bias ps, we can obtain the debiased prediction probability of unlabeled data as follows:

Intuitively, Eq. (14) serves as a regularization term that smooths the prediction probabilities of the majority classes and sharpens these of the minority classes, which can alleviate the prior bias introduced by the heterogeneous data. The detailed procedures of DPL are shown in Algorithm 1.

4.4 DebiasedModel Aggregation

The objective of DMA is to computing aggregation weights that enable the model to perform uniformly across all classes. During each local updating round, the activated clients send

Algorithm 2: DMA: Debiased Model Aggregation

their accumulated APP-U pm and their latest models wm to the server. Then we can get the aggregated APP-U as follows:

where βm denotes the aggregation weight for client m. To achieve a more balanced model, we expect paggr to be more uniform, leading to the following optimization objective:

where pt = { 1K }K is the uniform distribution over K classes, identical to the test dataset. In FedDB, we utilize the gradient descent algorithm to solve the above optimization problem.

After obtaining the aggregation weights β, we aggregate client models and update the global model as follows:

where wtm is the local model of client m at last round, wt+1 is the global model. wt+1 is then broadcast to the activated client for further updates. The processes of DMA and FedDB are presented in Algorithms 2 and 3, respectively.

5 Experiments

This section details the experimental results in various settings to demonstrate the effectiveness of FedDB.

5.1 Experimental Setup

Datasets. We evaluate FedDBon three benchmark datasets, including CIFAR10, SVHN, and CIFAR100. Initially, a balanced labeled dataset is separated from the original training dataset, with the residual data designated as the unlabeled dataset. When distributing these training data to clients, we sample data from a Dirichlet distribution q ∼ Dir(δp), where p is the class-wise prior distribution and δ is a parameter that modulates the heterogeneity among clients. A higher value of δ correlates with reduced data heterogeneity. To enrich the unlabeled dataset, we add the samples from the labeled dataset to the unlabeled dataset after discarding their labels. We conduct experiments in IID setting and Non-IID settings with δ = {0.1, 0.3}. In the IID setting, the total number of labeled samples is set to 4000, 1000, 10000 for CIFAR10,

Algorithm 3: FedDB: Detaching Prior Bias in FSSL

Input: Client number M, client activate rate C, global rounds T , update epochs E and Eaggr, learning rate η and ηaggr, threshold τ , unlabeled loss weight λ, momentum accumulation coefficient γ Output: Global model wT 1 Server executes: 2 Initialize w0 3 for t = 1, 2, ..., T do 4 St ← randomly select M · C clients 5 for each client in m ∈ St in parallel do 6 wtm, ptm ← ClientUpdate(wt-1 ) 7 wt ←DMA({wtm}m∈St , {ptm}m∈St , Eaggr, ηaggr) 8 Return wT 9 ClientUpdate(wt ) 10 ˆ Y , p ← DPL(τ ) 11 for e = 1, 2, ...,E do 12 pe = ∑Nu n=1 p(y|α(xnu )) Nu 13 Ls = 1 Nl ∑Nl n=1 H(yn , p(y|α(xnl ))) 14 Lu = 1 Nu ∑Nu n=1 1(max( ˆ Y n )) ≥ τ )· H( ˆ Y n , p(y|A(xnu))) 15 L←Ls +λLu 16 we ←we-1 -η∇L; p← γp+ (1 -γ)pe 17 Return wE , p

SVHN and CIFAR100, respectively. For Non-IID setting, the total number of labeled data is set to 4000 for CIFAR10 and SVHN, and 10000 for CIFAR100. The test dataset from the original dataset is used for model evaluation.

BenchmarkMethods. We compare FedDB against the following benchmark methods:

• FedAvg [McMahan et al., 2017]: The FedAvg method is applied in a constrained scenario where each client utilizes only the small labeled dataset for training.
• FixMatch [Sohn et al., 2020]: This method is a basic adaptation of FixMatch within FedAvg framework.
• FedMatch [Jeong et al., 2021]: FedMatch introduces the inter-client consistency loss to maximize the agreement between local models.
• FedRGD [Zhang et al., 2021]: It mitigates the model bias by reducing gradient divergence among clients.
• SemiFL [Diao et al., 2022]: SemiFL adopts alternate training between server and clients. Here, we adopts its client-side training due to the lack of training samples on the server in our scenario.
• Methods combining DPL. We also conduct experiments that integrate DPL with benchmark methods. These hybrid methods are denoted as Method-FedDPL.

Implementation Details. We primarily follow the experimental settings adopted in prior works of FSSL [Jeong et al., 2021]. There are a total of 100 clients participating in the training, with 10 active clients (C = 0.1) engaged in each global round. The local training epoch is set to E = 5 and the epoch for updating the model aggregation weights is set to Eaggr = 100. All experiments are executed for 800 global rounds. We employ Wide ResNet28x2 in our experiments. The SGD optimizer is adopted for model training, operating at learning rates η = 0.03 for local updating and ηaggr = 1.0 for aggregation, complemented by a momentum of 0.9. Due to the limited number of samples on clients, we feed all training data simultaneously to the model during local training. The confidence threshold for pseudo-labeling is set to τ = 0.95. The data augmentation operation is consistent with those described in FixMatch [Sohn et al., 2020]. All experiments are repeated for 4 times and we report the mean and standard deviation of the best accuracy during training.

5.2 Results on Benchmark Datasets

The experimental results are presented in Tables 1 - 3, where values inside the parentheses represent the mean, and values outside the parentheses represent the standard deviation of multiple experiments. It can be observed that with the same number of labeled samples, the accuracy of all methods decreases as δ decreases, demonstrating that data heterogeneity is a key factor harming model performance. FedAvg, despite its simplicity, serves as a reliable benchmark method, particularly as the dataset difficulty increases (e.g., CIFAR100). This issue is also noted by [Diao et al., 2022]. This demonstrates that improperly incorporating unlabeled data into training can negatively impact the model's training. Compared with other FSSL methods, FedDB enhances test accuracy, demonstrating the effectiveness of FedDB in the FSSL scenario. The same conclusion can also be drawn from Figure 5.

5.3 Effectiveness of DPL

As illustrated in Table 4, employing DPL results in substantial gains for FedDB. Furthermore, DPL can be regarded as a convenient plug-in that can be easily integrated into existing FSSL methods utilizing pseudo-labeling. As shown in Tables 1 - 3, introducing DPL to existing FSSL methods effectively enhances their performance. Figure 6 displays the accuracy of pseudo-labels during training. It indicates that DPL effectively enhances the accuracy of these pseudo-labels, which in turn benefits FSSL training. Figure 7 presents the ratio of pseudo-labeled samples in the unlabeled data. However, introducing DPL does not consistently improve the ratio of pseudo-labeled samples, as the model in FSSL is challenging to train, making it difficult for samples to be pseudo-labeled.

5.4 Effectiveness of DMA

As shown in Table 4, DMA generally contributes positively to FedDB in most scenarios. However, its impact differs among various datasets. More specifically, DMA consistently results in improved outcomes on the CIFAR10 and CIFAR100 datasets. Conversely, on the SVHN dataset, DMA can lead to performance decline in certain scenarios. Upon detailed analysis, we ascribe this issue to the imbalanced distribution of the SVHN dataset, which contravenes the objective of FSSL that seeks for a balanced model.

6 Conclusion

In this paper, we propose FedDB to detach prior bias in FSSL with class imbalance. At the local training level, FedDB debiases the pseudo-labeling using APP-U based on Bayes' theorem, encouraging a more balanced training data during the training. At the global aggregation level, FedDB leverages APP-U across different clients to derive optimal aggregation weights, aiming to debias the global model. Extensive experiments have shown the effectiveness of FedDB.

Acknowledgements

This work was supported by the National Natural Science Foundation of China under Grants 62372028 and 62372027.