SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Navigating Data Heterogeneity in Federated Learning: A Semi-Supervised Federated Object Detection
Taehyeon Kim $^{1}$  Eric Lin $^{2}$  Junu Lee $^{3}$  Christian Lau $^{2}$  Vaikkunth Mugunthan $^{2}$
$^{1}$ KAIST  $^{2}$ DynamoFL  $^{3}$ The Wharton School
potter32@kaist.ac.kr
Federated Learning (FL) has emerged as a potent framework for training models across distributed data sources while maintaining data privacy. Nevertheless, it faces challenges with limited high-quality labels and non-IID client data, particularly in applications like autonomous driving. To address these hurdles, we navigate the uncharted waters of Semi-Supervised Federated Object Detection (SSFOD). We present a pioneering SSFOD framework, designed for scenarios where labeled data reside only at the server while clients possess unlabeled data. Notably, our method represents the inaugural implementation of SSFOD for clients with  $0\%$  labeled non-IID data, a stark contrast to previous studies that maintain some subset of labels at each client. We propose FedSTO, a two-stage strategy encompassing Selective Training followed by Orthogonally enhanced full-parameter training, to effectively address data shift (e.g. weather conditions) between server and clients. Our contributions include selectively refining the backbone of the detector to avert overfitting, orthogonality regularization to boost representation divergence, and local EMA-driven pseudo label assignment to yield high-quality pseudo labels. Extensive validation on prominent autonomous driving datasets (BDD100K, Cityscapes, and SODA10M) attests to the efficacy of our approach, demonstrating state-of-the-art results. Remarkably, FedSTO, using just  $20 - 30\%$  of labels, performs nearly as well as fully-supervised centralized training methods.
Federated Learning (FL) enables decentralized training across distributed data sources, preserving data privacy [27]. It has emerged in response to the need for privacy, security, and regulatory compliance such as GDPR [36] and CCPA [31]. FL trains models on local devices and shares only model updates, thereby improving privacy and efficiency. In a typical FL cycle, each client updates a shared model with local data, sends the updates to a server for parameter aggregation, and then updates its local model with the newly aggregated global model sent back by the server.
Despite the potential of FL, the assumption of fully labeled data restricts its practicality [12, 11]. In order to acquire high-quality labels, data is often transmitted from edge clients to a central server, thereby compromising the privacy assurances provided by FL. Limited labels at the edge necessitate the adoption of transfer learning, self-supervised learning, and semi-supervised learning (SSL) techniques. However, the separation of labeled and unlabeled data complicates the application of these techniques to FL, which can undermine the system's effectiveness. This issue is amplified in labels-at-server scenarios where only the server possesses labeled data, and clients hold only unlabeled data [5, 10, 43, 2, 18, 42]. In autonomous driving, a novel approach is required to bridge the knowledge gap between labeled and unlabeled data without the need for direct data exchange.
While Semi-Supervised Federated Learning (SSFL) has been explored for image classification tasks,[5, 10, 43, 2, 18, 42], these studies have faced the following challenges:
Upload updated Backbone
1. Representation Learning with Selective Training
Orthogonal Enhancement
Figure 1: An overview of our FedSTO method within the SSFOD framework with key components: selective training, orthogonal enhancement, and local Exponential Moving Average (EMA)-driven pseudo label assignment, organized into two stages. Algorithm steps are numbered accordingly.
2. Full Parameter Training with Orthogonal Enhancement
1. Limited scale and complexity of tasks with datasets such as CIFAR and ImageNet, while semi-supervised federated object detection (SSFOD) presents sizable greater difficulties.
2. Non-IID data shift from labeled to unlabeled data. Our investigation stands apart in tackling the most challenging FL situations where clients hold exclusively unlabeled data from a different distribution from labeled server data. This acknowledges inherent heterogeneity of real-world FL settings, such as diverse weather conditions across clients. For instance, one client's dataset may predominantly consist of images captured under cloudy conditions, while others may include images from overcast, rainy, snowy, etc. conditions.
To surmount these inadequately addressed challenges of SSFOD, we introduce FedSTO (Federated Selective Training followed by Orthogonally enhanced training), a two-stage training strategy tailored specifically for our SSFOD framework (Figure 1). Our key contributions include:
- Selective Training and Orthogonal Enhancement: FedSTO begins with selective training of the model's backbone while other components remain frozen, fostering more consistent representations and establishing a robust backbone. This promotes generalization across non-IID clients, even in the absence of local labels. The subsequent stage involves fine-tuning all parameters with orthogonal regularizations applied to the non-backbone part of the model. This enhancement step is designed to imbue the predictors with resilience against skewed representations induced by local data heterogeneity, thereby promoting representation divergence and robustness.
- SSFL with a Personalized EMA-Driven Semi-Efficient Teacher: To prevent deterioration of teacher pseudo labeling models for non-IID unlabeled clients, we showcase for the first time an SSFOD framework that applies an alternate training methodology [5], integrated with a Semi-Efficient Teacher [38], driven by a local Exponential Moving Average (EMA). Our empirical observations suggest that this personalized EMA-driven model provides superior quality pseudo labels for detection, contrary to the commonly used global model
Figure 2: Performance comparison on BDD100K dataset [41]. "Partially Supervised Training" shows lower-bound performance using partial labels in a centralized setting. "Vanilla Semi-Supervised Federated Learning" and "Our FedSTO" demonstrate improved performance with non-IID federated data. FedSTO approaches the "Fully Supervised Training" upper-bound performance under full label use in a centralized setting. The x-axis shows the number of labeled examples, and the y-axis displays the mean average precision (mAP@0.5) on the test set.
for pseudo labeling in related studies [5]. This approach further enhances the quality of the learning process, mitigating potential pitfalls of noisy pseudo labeling.
- Performance Improvements: FedSTO achieves 0.082 and 0.035 higher mAP@0.5 when compared to partially supervised and SSFL baselines respectively, nearly matching the fully supervised model's performance (0.012 gap) on BDD100K [41] with a mere  $25\%$  of labeled data. We demonstrate similar considerable improvements in model generalization (Figure 2) on rigorous benchmark and ablation experiments with 20k-120k datapoints from Cityscapes [4] and SODA10M [9], utilizing the YOLOv5 object detector [13].
Our above contributions present a pioneering approach for utilizing unlabeled data in FL to enhance non-IID detection performance, especially for dynamic objects—an aspect not yet considered in previous research. Despite the paucity of research on SSFOD, we contend that our methods and experiments offer a valuable benchmark for future investigations across diverse domains.
2.1 Federated Learning (FL): Challenges and Advances
FL has gained significant attention in recent years as a privacy-preserving approach to harness the potential of distributed data [27, 20, 21, 15, 28, 33]. Despite the progress made in FL, most research has focused primarily on classification tasks, which may limit its applicability and generalizability to a broader range of real-world problems. Advanced FL techniques are essential to revolutionize various fields, including autonomous driving, healthcare, and finance, by enabling collaborative learning from distributed data sources [6, 30].
Addressing data heterogeneity is of paramount importance in FL, as clients frequently hold data with diverse distributions, which may impact the performance of the global model. To tackle this challenge, researchers have proposed various techniques to handle non-IID data, including adaptive aggregation algorithms and local fine-tuning of models [20, 33]. Personalization constitutes another vital aspect of FL, since clients may exhibit unique requirements or preferences not entirely captured by the global model [3, 19, 39]. Methods such as model distillation [23] and meta-learning [7] have been investigated to facilitate client-specific model adaptation and personalization. Finally, communication efficiency is crucial in FL, as exchanging model updates can be resource-intensive. To alleviate this issue, researchers have introduced strategies like quantization [34], sparsification [29], and the utilization of a supernet containing multiple subnetworks, with only the selected subnetwork transmitted to the server to reduce communication overhead while preserving model performance [17].
2.2 Semi-Supervised Object Detection (SSOD) with YOLO Object Detector
SSOD has been an active research area, focusing on improving the quality of pseudo labels to enhance overall detection performance [35, 40]. The evolution of SSL techniques in object detection primarily revolves around using pretrained architectures and applying strong data augmentation strategies to generate consistent and reliable pseudo labels.
Traditional single-stage detectors, such as the family of YOLO detectors, have faced notable challenges in leveraging SSL techniques, often underperforming compared to their two-stage counterparts (e.g., Faster RCNN). The limited efficacy of existing SSL methods for single-stage detectors has inspired researchers to develop innovative solutions to overcome these limitations [44]. Recently, a novel pipeline incorporating EMA of model weights has exhibited remarkable enhancements in the performance of single-stage detectors like YOLO detectors [38]. By utilizing the EMA model for pseudo labeling, researchers have adeptly addressed the inherent weaknesses in single-stage detectors, substantially elevating their performance in SSL contexts for object detection tasks.
2.3 Semi-Supervised Federated Learning (SSFL)
SSFL has emerged as a promising approach to address the challenge of limited labeled data in FL scenarios [5, 10, 43, 2, 18, 42]. SSFL aims to jointly use both labeled and unlabeled data owned by participants to improve FL. Two primary settings have been explored: Labels-at-Client and Labels-at-Server [10, 11]. In the Labels-at-Client scenario, clients possess labeled data, while the server only has access to unlabeled data. Conversely, in the Labels-at-Server scenario, the server holds labeled data, and clients have only unlabeled data. Despite the progress in SSFL, there remain limitations in the current research landscape. The majority of existing SSFL research predominantly focuses on image classification tasks, leaving other applications relatively unaddressed. In this study, we address these limitations by tackling the more realistic and challenging scenarios with edge clients having (1) no labels and (2) non-IID data (domain shift from the server labeled data), specifically in the context of object detection tasks.
SSFOD We tackle a semi-supervised object detection task involving a labeled dataset  $S = \{x_i^s, y_i^s\}_{i=1}^{N_S}$  and an unlabeled dataset  $\mathcal{U} = \{x_i^u\}_{i=1}^{N_U}$ , focusing on scenarios where  $N_S \ll N_U$ . In our SSFOD setup, as illustrated in Figure 1, we assume  $M$  clients each possessing an unsupervised dataset  $x^{u,m}$ . The server retains the labeled dataset  $\{x^s, y^s\}$  and a model parameterized by  $W^s$ . Each client model, parameterized by  $W^{u,m}$ , shares the object detection architecture denoted by  $f:(\boldsymbol{x}, W) \mapsto f(\boldsymbol{x}, W)$ . We assume that all models share the same object detection architecture, denoted by  $f:(\boldsymbol{x}, W) \mapsto f(\boldsymbol{x}, W)$ , which maps an input  $\boldsymbol{x}$  and parameters  $W$  to a set of bounding boxes and their corresponding class probabilities on the  $K$ -dimensional simplex (e.g., using the sigmoid function applied to model outputs unit-wise).
Data Heterogeneity Our study addresses non-IID data resulting from varying weather conditions such as cloudy, overcast, rainy, and snow, inspired by feature distribution skew or covariate shift [14]. We utilize three datasets, BDD100K [41], CityScapes [4], and SODA10M [9], each displaying class distribution heterogeneity and label density heterogeneity. Our aim is an SSFOD framework that can manage this heterogeneity, maintaining performance across diverse conditions and distributions. Data is considered IID when each client exhibits a balanced weather condition distribution.
Evaluation In our framework, we assess the performance of all detection tasks based on mean average precision (mAP@0.5), a standard metric in object detection literature that provides a comprehensive view of model performance across various object classes and sizes. Importantly, we evaluate the post-training performance of our method by assessing the personalized models of the server and client on their respective datasets. This approach ensures a fair and context-specific evaluation, reflecting the true performance of the personalized models in their intended environments.
Baseline Training Our work explores two principal baselines: "Centralized Training" and "Federated Learning". Depending on the degree of labeled data utilization, we categorize the training into "Partially Supervised" and "Fully Supervised". An ideal situation is one where a fully supervised model is trained in a centralized fashion, utilizing all labeled data. In contrast, a more challenging scenario involves a partially supervised model trained solely on the server's limited labeled data. Under our problem setup, we initially establish a baseline by performing partial supervision on the
Table 1: Performance under different weather conditions for non-IID and IID data splits with 1 server and 3 clients. The results are presented for centralized (Fully Supervised and Partially Supervised) and federated approaches with a pseudo label assigner (Global Model and Local EMA Model).
Type	Method	Non-IID	IID
Cloudy	Overcast	Rainy	Snowy	Cloudy	Overcast	Rainy	Snowy
Centralized	Fully Supervised	0.600	0.604	0.617	0.597	0.600	0.604	0.617	0.597
Partially Supervised	0.540	0.545	0.484	0.474	0.528	0.545	0.533	0.510
Federated	Global Model [5]	0.555	0.560	0.497	0.488	0.540	0.551	0.576	0.542
Local EMA Model [38]	0.560	0.566	0.553	0.553	0.572	0.588	0.593	0.610
Figure 3: Performance of various methods on the BDD100K dataset [41], with the server containing labeled data for the "Cloudy" category and 3 clients having unlabeled data for "Rainy", "Snowy", and "Overcast" categories. Baseline SSFL (red hatched boxes) struggles in comparison to centralized learning (bars in green dotted boxes). "Fully Supervised" and "Partially Supervised" refer to training a centralized model with the complete labeled dataset and only the "Cloudy" labeled data, respectively.
server's limited labeled data, which serves as a pretraining step. Following this, each client conducts "Unsupervised learning with unlabeled data". Upon completion, clients transmit their model weights to the server. The server then aggregates these weights and fine-tunes the amalgamated model using its labeled data. The updated model is subsequently disseminated back to the clients, culminating one learning round. This cyclical process, known as alternate training in Diao et al. [5], continues effectively. It merges the strength of supervised and unsupervised learning to capitalize on unlabeled data while preventing model deterioration, thereby optimizing model performance.
Personalized Pseudo Labeling for Unlabeled Clients A crucial obstacle in SSFOD lies in precise pseudo label assignment, as improper allotments can result in label inconsistencies, thus negatively impacting mutual learning performance. Building upon the foundation by Xu et al. [38] in centralized settings, we present the first extension of this approach to federated settings, leveraging a personalized Pseudo Label Assigner (PLA) equipped with local EMA models. This technique bifurcates pseudo labels into reliable and unreliable categories using high and low thresholds, thus ensuring a robust and precise learning mechanism in federated environments. In FL, the PLA can be applied to both global and local models. However, the global model may fall short in capturing unique features of local data, compromising pseudo label quality. As demonstrated in our evaluation (Table 1), locally updated EMA models outperform global models. While it is feasible to federate the local EMA model, it introduces certain trade-offs, such as increased communication costs and minor performance degradation compared to the local EMA model. Our SSFOD framework, therefore, incorporates a local PLA with a local EMA model, optimally balancing communication efficiency and model stability, ensuring an effective learning process for SSOD tasks in distributed environments.
SSFOD with YOLO We utilize the YOLOv5 model, a single-stage object detector, in our evaluation. Existing literature shows a scarcity of research on SSFOD within FL like FedAvg [27], particularly for single-stage detectors like YOLO. Figure 3 compares various learning approaches in centralized and federated settings, denoted by green dotted and blue hatched boxes, respectively. We highlight non-IID scenarios with labeled (cloudy) and unlabeled data (overcast, rainy, snowy). In the CL scenario, fully supervised methods noticeably surpass partially supervised ones, and SSL approaches almost match the performance of fully supervised methods. However, baseline training for FL falls substantially short of these high standards, particularly with unlabeled data.
Table 2: Performance on the BDD dataset with 1 labeled server and 3 unlabeled clients as each element of our FedSTO approach within the SSFOD framework is added. It highlights how each added method contributes to the overall performance under both Non-IID and IID conditions.
Method	Non-IID	IID
Cloudy	Overcast	Rainy	Snowy	Total	Cloudy	Overcast	Rainy	Snowy	Total
Partially Supervised	0.540	0.545	0.484	0.474	0.511	0.528	0.545	0.533	0.510	0.529
+ SSFL [5] with Local EMA Model	0.560	0.566	0.553	0.553	0.558	0.572	0.588	0.593	0.610	0.591
+ Selective Training	0.571	0.583	0.557	0.556	0.567	0.576	0.578	0.594	0.599	0.587
+ FPT with Orthogonal Enhancement [16]	0.596	0.607	0.590	0.580	0.593	0.591	0.634	0.614	0.595	0.609
4 Main Method: FedSTO
To mitigate these inherent hurdles presented by FL, we introduce FedSTO, a method that unfolds in two stages, preceded by a warmup stage. The process commences with an emphasis on robust representation learning for pretraining (Subsection 4.1), followed by full parameter training (Subsection 4.2). The initial stage of pretraining integrates a warm-up period utilizing labeled data at the server, transitioning into selective training. This groundwork is fortified by the orthogonal enhancement implemented in the subsequent full parameter training phase.
4.1 Selective Training (ST)
Selective Training (ST) is designed to address the primary challenge of establishing a robust backbone for object detector in FL. The approach unfolds as follows:
1. Labeled dataset training: All model parameters are updated using a labeled dataset. This step ensures training commences on high quality labeled data, mitigating the potential destabilizing effect of noisy, unlabeled data, and heterogeneity from diverse weather conditions.
2. Client-side training with unlabeled dataset: The model, updated in the previous step, is dispatched to the clients. Each client trains the model on their local unlabeled dataset. However, during this phase, only the backbone part of the model is updated, leaving other components frozen. This selective updating procedure fosters more consistent representations by sharing the same non-backbone part (e.g., neck and head), and thus enhances its potential generalization capabilities by concentrating on feature extraction.
3. Server-side aggregation: The server aggregates the updated backbone parameters from clients, effectively synthesizing the learned information from diverse unlabeled datasets. The aggregated backbone is then utilized in the first step for further training, repeating the process until performance convergence.
By adhering to this procedure, ST effectively navigates the challenges inherent in the progression of FL while simultaneously accruing substantial benefits. Ensuring stability in semi-supervised object detection tasks is paramount. The exposure to heterogeneous unlabeled data, potentially characterized by noise or variable quality, can induce biases into the neck and head components of the model, thereby risking performance degradation by inadvertently generating low-quality or imprecise pseudo annotations. To mitigate this, ST employs a selective update strategy specifically targeting the backbone of the model, which is predominantly entrusted with the task of extracting salient features from the input data. By concentrating on the backbone during training, ST aids in the preservation of model stability and the enhancement of its generalization capabilities. Furthermore, in this stage, the communication cost between the server and clients is significantly reduced by uploading only the backbone part from clients to a server. Consequently, it significantly minimizes the deleterious impacts of heterogeneous unlabeled data on overall model performance (Table 2). While ST brings marginal improvements in IID conditions, it presents potent effects under Non-IID circumstances, emphasizing its efficacy in handling heterogeneous data distributions.
4.2 Full Parameter Training (FPT) with Orthogonal Enhancement
Inspired by the critical need for personalized models to exhibit robustness against feature distribution skewness—predominantly due to diverse weather conditions—we integrate orthogonality regularization presented by Kim et al. [16] which penalizes the symmetric version of the spectral restricted isometry property regularization  $\sum_{\theta}\sigma (\theta^T\theta -I) + \sigma (\theta \theta^T -I)$  within the SSFOD framework where  $\sigma (\cdot)$  calculates the spectral norm of the input matrix and  $\theta$  is a weight matrix from non-backbone parts. This regularization is applied during both server and client training stages and targets non-backbone
Algorithm 1: FedSTO Algorithm within the SSFOD Framework
INPUT: server model parameterized by  $W_{s}$ , the number of rounds for each phase  $T_{1}, T_{2}$ , client models parameterized by  $\{W_{u,1}, \dots, W_{u,M}\}$ , client backbone part parameterized by  $\{B_{u,1}, \dots, B_{u,M}\}$   
1:  $W_{s} \gets \text{WARMUP}(x_{s}, y_{s}, W_{s})$  // Supervised training at server /* Phase 1: Selective Training for Pretraining */  
2: for  $t \gets 0, \dots, T_{1} - 1$  do  
3:  $S^{t} \gets \text{SAMPLECLIENTS}$   
4: for each client  $k \in S^{t}$  in parallel do  
5:  $W_{u,k} \gets \text{CLIENT-BACKBONEUPDATE}(x_{u,k}, B_{u,k})$  // Client-Update  
6: end for  
7:  $W_{s} \gets \sum_{k \in S^{t}} p_{k} W_{u,k}$  // Aggregation  
8:  $W_{s} \gets \text{SERVER-UPDATE}(x_{s}, y_{s}, W_{s})$  // Server-Update  
9: end for  
/* Phase 2: Full Parameter Training with Orthogonal Enhancement */  
10: for  $t \gets 0, \dots, T_{2} - 1$  do  
11:  $S^{t} \gets \text{SAMPLECLIENTS}$   
12: for each client  $k \in S^{t}$  in parallel do  
13:  $W_{u,k} \gets \text{CLIENT-ORTHOGONALUPDATE}(x_{u,k}, W_{u,k})$  // Client-OrthogonalUpdate  
14: end for  
15:  $W_{s} \gets \sum_{k \in S^{t}} p_{k} W_{u,k}$  // Aggregation  
16:  $W_{s} \gets \text{SERVER-ORTHOGONALUPDATE}(x_{s}, y_{s}, W_{s})$  // Server-OrthogonalUpdate  
17: end for