FedCC: Towards Addressing Label Distribution Skews in Distillation-Based Federated Learning
Abstract
Federated Learning (FL) enables distributed clients to collaboratively train models without sharing raw data, making it promising for leveraging massive devices in communication networks. In distillation-based FL, each client applies its local model on an unlabeled public dataset, and shares only prediction results with the server. While heterogeneous local data introduces label distribution skew, thus biasing client models toward majority classes and leading to potentially inaccurate predictions. The lack of ground-truth labels in the public dataset hampers the server’s ability to calibrate predictions, which ultimately degrades overall performance. To address this, we propose FedCC, a simple and effective algorithm for mitigating client misclassification. Instead of being forced to classify and risking error propagation, clients are allowed to tag ambiguous samples as ‘unknown’. This additional class, together with calibrated pseudo-labels on the public data, balances confidence in majority classes against uncertainty in under-represented ones. Extensive experiments demonstrate that FedCC significantly outperforms existing methods, especially under severe label skew. In the extreme scenario where each client holds samples from only one of ten classes, FedCC achieves 67.3% accuracy, while baselines collapse to near-random results.
Index Terms:
Federated learning, Distillation-based FL, Label distribution skew, Pseudo-labelingI Introduction
Federated Learning (FL) [1] enables collaborative model training across decentralized clients without exposing their raw data. In the classical parameter-based approach, each client refines the global model with its private data, uploads the updated parameters to the server, and the server aggregates these updates to produce the next-round global model. While effective, this approach incurs heavy communication overhead on parameter transmission, and assumes a homogeneous model architecture, overlooking the diverse bandwidth, computation, and data distributions across real-world devices [2, 3, 4, 5].
Drawing on Knowledge Distillation (KD) [6], Distillation-based FL offers an alternative for client collaboration [7]. Each client runs its local model on a public dataset and uploads only the resulting prediction vectors (i.e., class-score logits); the server then aggregates these vectors into a global prediction vector and redistributes it, reducing communication overhead by orders of magnitude. In addition, because logits are architecture-agnostic, the approach natively supports heterogeneous models, enabling flexible deployment across varied networked devices. Regarding the public dataset, various strategies have been employed, including the use of a labeled dataset [8, 9], or the generation of synthetic data derived from local datasets or local model parameters [10]. Recent studies [11] have relaxed the requirement by adopting unlabeled ones, avoiding costly labeling processes and alleviating privacy concerns. Accordingly, this paper adopts an unlabeled public dataset as an integral component of the scenario setting.
Despite its promise, distillation-based FL is vulnerable to label distribution skew: heterogeneous client data cause local models to overfit the majority classes, and the resulting biases are reflected in the prediction vectors [2]. The challenge is further compounded by the absence of ground-truth labels in the public dataset, which prevents the server from calibrating logits and results in distorted supervision that degrades global performance. Existing approaches in parameter-based FL (e.g., weighted model aggregation [12], model regularization [13]) presume access to client model weights. This assumption breaks down in distillation-based FL, where only logit vectors are exchanged, providing far sparser signals. Recent efforts focus on weighting client updates by estimated reliability during server-side aggregation [14], moving beyond naive averaging as in FedDF[11]. Others aim at amplifying minority-class information and aligning local models with the global objective [15]. Nevertheless, label-skew bias still creeps in, and existing schemes cannot consistently filter out misleading updates during aggregation. Empirical evidence [16] demonstrates that, under label-distribution skew, local models overwhelmingly predict their majority classes, thereby injecting errors into the aggregated prediction (as in Fig. 1a).
In response, we propose FedCC, a simple and effective algorithm designed to mitigate Client misClassification. Treating the mismatch between a client’s skewed data and the ideal balanced distribution as an open-set problem [17], we introduce an ‘unknown’ class to absorb minority and missing classes. Instead of being forced to issue a class prediction and risking error propagation, clients are allowed to acknowledge their limitations and mark ambiguous inputs as ‘unknown’. Illustrated in Fig. 1b, this ‘unknown’ class (shown in blue shadow) enables clients to express uncertainty, yielding more informative aggregated predictions. To implement this, we incorporate the unlabeled dataset into the local objective function and employ a calibrated pseudo-label generation method, to weigh confidence in majority classes with the uncertainty hidden in unknown classes. This weight is adaptively adjusted per client and per sample, optimizing the model’s generalization under diverse data conditions.
To evaluate performance, we conduct experiments on three widely used datasets: CIFAR-10, CIFAR-100, and TinyImageNet, across various label skew scenarios. FedCC consistently outperforms state-of-the-art FL methods by at least 3% points in most scenarios, and the accuracy gap grows as label skew becomes more severe. In the extreme scenario where each client holds samples from only one of ten classes, FedCC achieves an accuracy of 67.3%, while other baselines fail to generate meaningful predictions. Beyond delivering strong global generalization, FedCC enhances local model performance on its minority classes by leveraging uncertainty-aware predictions.
To summarize, our key contributions are as follows:
- •
We introduce FedCC, a simple and effective method for mitigating client misclassification. We frame label skew as an open-set recognition problem, and extend the existing class framework with an ‘unknown’ class to absorb minority and missing classes.
- •
FedCC incorporates the unlabeled dataset into the local objective function, and proposes a novel pseudo-label generation method, enabling the model to recognize and account for unknownness. This method can be interpreted as adding a regularizer, mitigating overfitting.
- •
Extensive experiments demonstrate that FedCC significantly outperforms current FL methods, and the performance gap increases as label skew becomes more severe.
II Related Works
Parameter-based FL with label distribution skews: While FL safeguards privacy by keeping data local, the skewed label distributions across clients undermine model generalization. To stabilize heterogeneous training, methods like FedProx [13] and FedSAM [18] employ regularization, variance reduction, or gradient filtering. Other approaches target class imbalance, e.g., via rescaling logits (FedRS [19]), or calibrating predictions (FedLC [20], FedMR [15], FedVLS [21]). Alternatively, FedConcat [22] departs from model averaging by concatenating local encoders to collaboratively train a global classifier. While FedOV [16] uses open-set recognition to alleviate label skew, its reliance on synthesizing outlier samples with implicit labels limits its applicability to unlabeled data.
Distillation-based FL with label distribution skews: To curb label skew under distillation-based FL, FedNed [23] emphasizes distilling minority-class knowledge, LightTS [24] re-weights clients by strength, and Co-Boosting [10] alternates synthetic data generation with ensemble updates. However, their reliance on labeled public datasets or large-scale synthesis raises annotation and privacy concerns. When public data is unlabeled, the lack of ground-truth labels complicates bias correction. Existing solutions address this by retaining high-confidence predictions [8], clustering clients [25], or communicating only hard labels [26]. While emphasizing reliable local knowledge, these methods struggle to filter misinformation stemming from inherent local biases (As shown in Fig. 1a). Thus, effective solutions for fair client-side classification remain an open challenge.
III Scenario Setting
Consider an FL scenario with clients and classes (label set ). Each client holds a private local dataset . Additionally, a global unlabeled public dataset is accessible to all clients and the server. Due to label skew, client label priors vary, though we assume the conditional distribution is consistent across clients [2]. We partition into three mutually exclusive subsets () for each client : majority classes (sufficiently represented), minority classes (sparse, noisy estimates), and missing classes (absent).
Collaborative distillation protocol: Client deploys a heterogeneous classifier . In each communication round, client computes logits for public samples . These are converted to probabilities via a temperature-scaled softmax (with temperature omitted for brevity).
The server aggregates these probabilities into an ensemble teacher:
| (1) |
Clients then synchronize their models by minimizing the Kullback-Leibler divergence with the teacher distribution:
| (2) |
Post-synchronization, clients refine their models on local private data before sharing predictions in the next round.
IV Approach
FedCC targets client-side misclassification, and by lowering local errors, it ultimately improves global generalization. Sections IV-A-IV-B detail the proposed client-level algorithm, and Section IV-C outlines the overall federated training pipeline. In Sections IV-A–IV-B we focus on a single client, and therefore omit the client index for notational brevity (e.g., , and instead of , and ).
IV-A Local learning objective
Semi-supervised learning: We begin by reviewing Semi-Supervised Learning (SSL), a paradigm that leverages both labeled and unlabeled data to build effective models. A classical SSL strategy is self-training [27], which bootstraps a model with its own high-confidence predictions. Firstly, a classifier is fitted on labeled data . The trained model then assigns pseudo-labels to the unlabeled data by selecting the class of maximum posterior probability: . Merging the original and pseudo-labeled examples yields an enlarged dataset on which the model is re-optimized. The overall training objective is to minimize the following loss function:
| (3) | ||||
where the cross-entropy loss quantifies the discrepancy between the label and prediction logits , and is given as the probability after the softmax operation; balances the labeled and unlabeled terms.
Our method: As described in the scenario setting, client maintains a labeled dataset and unlabeled dataset , mirroring the typical SSL structure. To optimally utilize both datasets, we design the objective function based on Eq. 3 as:
| (4) | ||||
where and denote the losses on the labeled and unlabeled datasets, respectively.
As introduced before, client first fits a model on its labeled data, detailed as . Ideally, would master classes it has encountered, i.e., . However, in practice, the severe under-representation of minority classes leads to poorly calibrated posterior probabilities, which often drift toward majority classes [20]. Furthermore, lacks the ability to recognize samples from missing classes as it has never learned them during training.
To mitigate misclassification of minority and missing classes, we enlarge the label space with an auxiliary ‘unknown’ class , which acts as a catch-all for samples with unreliable predictions. With the label space augmented to , the model then satisfies .
Building on this, we ideally partition the unlabeled dataset into samples from majority classes , and all others. If the sample’s true class belongs to client ’s majority set , the pseudo-label is assigned to: , as abundant local examples enable the client to identify these classes reliably [28]; if , the sample is intended to be classified as . Then, the ideal pseudo-label assignment is
| (5) |
The unlabeled loss is therefore reformulated as:
| (6) | ||||
where the indicator equals 1 if the event is true and 0 otherwise. Since the truth labels are unavailable, we approximate this indicator with a soft weight :
| (7) | ||||
The comparison between Eq. 4 and Eq. 7 yields the pseudo-label in probabilistic-vector form:
| (8) |
where is the one-hot vector for the predicted class , and represents that of class . This pseudo-label distributes probability mass over only two classes: the predicted class and the unknown class . We detail the concrete design of in the following subsection.
IV-B Weight design
The weight re-weighs the local loss on unlabeled data so that learning is guided by global rather than local class prevalence. Instead of relying on class probabilities, which risk privacy leakage, we derive the weights solely from entropies, since they are inexpensive to compute and reflect uncertainty.
Weight towards majority classes (): Following [29], we regard a solid-color image as featureless, and probe the classifier’s intrinsic bias by feeding it into the pre-trained model, recording . To keep our entropy-based weighting scheme responsive even for near-certain predictions, we apply label smoothing [30]:
| (9) |
The weight is then , where denotes entropy. The smoothed entropy quantifies how strongly the model already favors its local label set and stays strictly positive even under a one-class bias.
Weight over unknown classes (): Classes under-represented by client are collected in with number being . Since the client lacks prior knowledge for those classes, we assign them the maximum uncertainty, defined as the entropy of a uniform distribution. This yields . Here, we add one virtual category to account for residual ambiguity introduced by the majority classes, reflecting that the client remains uncertain even with only one under-represented class.
Since quantifies the client’s concentration of probability over majority classes and captures its uncertainty about unknown classes, a natural way to trade off these two effects is through their ratio: . Note that and are client-specific and vary across clients.
To further account for sample-specific prediction uncertainty, we quantify the pretrained model’s confidence for each unlabeled instance by . As high entropy indicates low confidence in assigning the sample to a class in , we penalize with , leading to the final expression:
| (10) |
This weighting factor limits the extent to which a local model’s bias affects pseudo-labeling. reserves a non-zero probability for the ‘unknown’ class, allowing the model to consider an alternative label rather than forcing a majority-class assignment. By deferring these uncertain cases, curbs early mislabelling of minority samples and prevents such errors from cascading through subsequent training rounds.
After generating pseudo-labels, each client updates its local model via Eq. 4 to obtain , which is then used to generate predictions shared with the server.
IV-C Federated training pipeline
The training process runs over communication rounds. For notational clarity, we use the superscript and to denote the client index and the server-aggregated information, respectively, and to indicate the communication round.
At the beginning of round , each client downloads the aggregated soft targets from the previous round, denoted as , and then use Eq. 2 to synchronize its local model. Then each client adapts its model to its private labeled dataset , resulting in the intermediate model . Unlike prior distillation-based FL schemes that share predictions immediately, FedCC inserts an extra refinement step. Each client generates the pseudo-label via Eq. 8 and obtains the final round-specific parameters by optimizing the objective in Eq. 4. Each client then evaluates all samples in and produces temperature-scaled prediction vectors, , where the last entry corresponds to the unknown class. These prediction vectors are uploaded to the server. Clients’ contributions are weighted by known-class confidence: we take as a preliminary weight, then normalize it across clients to produce . The server then aggregates the received predictions as
| (11) |
It then zeroes out the unknown-class component and renormalizes the remaining entries to yield a valid probability vector . The resulting soft targets are then redistributed to all clients to initiate round .
During inference, we adopt the ensemble strategy used in FedMD [31] and FedOV [16]. After completing all rounds, clients transmit their model with parameters to the server. The server then aggregates the predictions from all client models using Eq. 11 and selecting the most probable label among the original classes . Formally, the predicted label is given by .
Algorithm 1 details the complete FedCC procedure.
V Experiments
V-A Experimental details
Method TinyImageNet CIFAR-100 CIFAR-10 FedAvg FedProx FedNova FedRS FedSAM FedLC FedMR FedConcat FedVLS FedMD FedDF LSR FedET FedHKT FedCT FedCC
Test data Local model FedAvg Local FedLC FedMR SSL FedCC FedAvg Global FedLC FedMR SSL FedCC
Dataset settings: Following the previous work [20, 21], we conduct classification tasks using CIFAR-10, CIFAR-100 and TinyImageNet datasets. We randomly divide each dataset into two parts: a public dataset and a private dataset. The public dataset, which comprises of the total training samples, has all labels removed and is shared among all participants. To simulate different label skews in the private dataset, we use two partition methods: (1) Dirichlet distribution : for each class , we sample , then assign client a fraction of the class- samples to its private set. A lower value indicates greater heterogeneity. (2) Pathological partition : each client has data from a fixed number of classes , with an equal number of samples for each class distributed among the assigned clients.
Baseline: We benchmark FedCC against nine parameter-based methods (FedAvg [1], FedProx [13], FedNova [12], FedRS [19], FedSAM [18], FedLC [20], FedMR [15], FedConcat [22], FedVLS [21]) and six distillation-based methods (FedMD [31], FedDF [11], LSR [30], FedET [32], FedHKT [25], FedCT [26]). For fairness, we exclude methods requiring a labeled public dataset or those based on generating labeled data. For details, please refer to the corresponding paper.
Experiment settings: Inspired by [25], parameter-based FL uses ResNet20 on all nodes; in distillation-based FL, clients split evenly between ResNet14 and ResNet20. The setup includes clients, global epochs, and local epochs per step in client model training. We adopt Adam optimizer with batch size , learning rate . Temperature is set to , tradeoff in Eq. 4 to , in (Eq. 9) to . We categorize classes by client ’s bias (Eq. 9): a class is a majority if , minority if , and missing if . We report results over three random seeds, with model accuracy shown without percent signs.
V-B An overall comparison
Experiment results, as summarized in Table I, show that FedCC outperforms current FL methods across almost all tested scenarios, with its advantage widening as label skew intensifies. On heavy skew ( or pathological partition), it is ahead by at least 3% points and often by more than 10%. In the harshest scenario where each client has access to only one class of CIFAR-10, baselines collapse to near-random guessing ( 12.7%), yet FedCC still delivers a robust 67.3%. This robustness stems from permitting clients to only predict familiar data and to label uncertain samples as ‘unknown’, thereby abstaining from unreliable predictions and reducing misleading information.
Comparison of local model performance: We evaluate the performance of local models on both local and global test data, as shown in Table II. The models are generated using FedAvg [1], FedLC [20], FedMR [15], standard SSL [27], and our FedCC. For the global test data, if the true label does not belong to the majority classes, classifying it as the ‘unknown’ class is considered valid. As expected, FedCC shows significant improvement on the global test data. On the local test set, FedCC leverages uncertainty to supply richer signals for minority classes and increases their likelihood of being correctly learned, consistently outperforming all baselines.
V-C Performance visualization
Decision margin: The decision margin measures the gap between the model’s score for the correct class and its highest incorrect guess. Evaluating under two challenging data partitions, Figure 2 shows margin histograms with a logarithmic y-axis. FedCC markedly shifts the distribution rightward to achieve a positive mean, whereas the best competing method remains negative. This rightward shift confirms that FedCC effectively ensures clear class separation and suppresses client-side misclassification, even under severe label skew.
T-SNE analysis of learned representations: Figure 3 visualizes the learned feature representations where each client holds data from a single class () on CIFAR-10, where each color denotes a distinct class. The embeddings are obtained by applying PCA whitening followed by t-SNE in a unified pipeline. Baseline: each client trains a local model using standard cross-entropy on its data, with the embedding of a representative client shown in Fig.3(a), and that of the FedAvg global model in Fig.3(b). Due to the lack of inter-class variation, local representations collapse along class-specific axes, and averaging these misaligned feature spaces yields weak global performance. FedCC: we next train local models with FedCC’s objective and aggregate them by ensembling their soft predictions. The corresponding client-side embedding (Fig.3(c)) cleanly isolates the client’s own class from the “unknown” region, and the aggregated embedding (Fig.3(d)) recovers distinct clusters for all classes. FedCC thus curbs error propagation by explicitly separating learned versus unseen classes at the client level, allowing the server to combine information without blurring minority categories.
VI Conclusion
We present FedCC, a novel distillation-based FL algorithm that mitigates client misclassification by incorporating an auxiliary ‘unknown’ class. By enabling clients to defer ambiguous predictions to this class, FedCC acts as an implicit regularizer, preventing overconfidence on majority classes and suppressing noisy updates. Extensive empirical evaluation confirms that it consistently outperforms state-of-the-art baselines in both local and global performance. Crucially, FedCC’s lightweight reliance on logit exchanges and robust performance make it well suited for deployment in real-world communication networks, which involves large numbers of devices with heterogeneous compute power, bandwidth, and data distributions.
References
- [1] (2017) Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics, Cited by: §I, §V-A, §V-B.
- [2] (2021) Advances and open problems in federated learning. Foundations and trends® in machine learning. Cited by: §I, §I, §III.
- [3] (2023) Advancing federated learning in 6g: a trusted architecture with graph-based analysis. In GLOBECOM 2023-2023 IEEE Global Communications Conference, pp. 56–61. Cited by: §I.
- [4] (2026) Foundation models for generalizable semantic and goal-oriented communication. In ICC 2026 - IEEE International Conference on Communications, Vol. , pp. 1–6. External Links: Document Cited by: §I.
- [5] (2025) LLaVA steering: visual instruction tuning with 500x fewer parameters through modality linear representation-steering. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Cited by: §I.
- [6] (2015) Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531. Cited by: §I.
- [7] (2026) Select to think: unlocking SLM potential with local sufficiency. In Forty-third International Conference on Machine Learning, Cited by: §I.
- [8] (2022) Not all knowledge is created equal: mutual distillation of confident knowledge. In Workshop on trustworthy and socially responsible machine learning, NeurIPS 2022, Cited by: §I, §II.
- [9] (2025) PRISM: self-pruning intrinsic selection method for training-free multimodal data selection. ArXiv abs/2502.12119. Cited by: §I.
- [10] (2024) Enhancing one-shot federated learning through data and ensemble co-boosting. arXiv preprint arXiv:2402.15070. Cited by: §I, §II.
- [11] (2020) Ensemble distillation for robust model fusion in federated learning. Advances in neural information processing systems. Cited by: Fig. 1, Fig. 1, §I, §I, §V-A.
- [12] (2020) Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems. Cited by: §I, §V-A.
- [13] (2020) Federated optimization in heterogeneous networks. Proceedings of machine learning and systems. Cited by: §I, §II, §V-A.
- [14] (2021) Ensemble attention distillation for privacy-preserving federated learning. In Proceedings of the IEEE/CVF international conference on computer vision, Cited by: §I.
- [15] (2024) Federated learning under partially class-disjoint data via manifold reshaping. arXiv preprint arXiv:2405.18983. Cited by: §I, §II, §V-A, §V-B.
- [16] (2023) Towards addressing label skews in one-shot federated learning. In International conference on learning representations, Cited by: §I, §II, §IV-C.
- [17] (2012) Toward open set recognition. IEEE transactions on pattern analysis and machine intelligence. Cited by: §I.
- [18] (2022) Generalized federated learning via sharpness aware minimization. In International conference on machine learning, Cited by: §II, §V-A.
- [19] (2021) Fedrs: federated learning with restricted softmax for label distribution non-iid data. In Proceedings of the ACM SIGKDD conference on knowledge discovery & data mining, Cited by: §II, §V-A.
- [20] (2022) Federated learning with label distribution skew via logits calibration. In International Conference on Machine Learning, Cited by: §II, §IV-A, §V-A, §V-A, §V-B.
- [21] (2025) Exploring vacant classes in label-skewed federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, Cited by: §II, §V-A, §V-A.
- [22] (2024) Exploiting label skews in federated learning with model concatenation. In Proceedings of the AAAI Conference on Artificial Intelligence, Cited by: §II, §V-A.
- [23] (2023) Federated learning with label-masking distillation. In Proceedings of the 31st ACM International Conference on Multimedia, Cited by: §II.
- [24] (2023) LightTS: lightweight time series classification with adaptive ensemble distillation. Proceedings of the ACM on management of data. Cited by: §II.
- [25] (2023) A hierarchical knowledge transfer framework for heterogeneous federated learning. In IEEE conference on computer communications, Cited by: §II, §V-A, §V-A.
- [26] (2025) Little is enough: boosting privacy by sharing only hard labels in federated semi-supervised learning. In Proceedings of the AAAI Conference on Artificial Intelligence, Cited by: §II, §V-A.
- [27] (1995) Unsupervised word sense disambiguation rivaling supervised methods. In Annual meeting of the association for computational linguistics, Cited by: §IV-A, §V-B.
- [28] (2018) Unsupervised domain adaptation for semantic segmentation via class-balanced self-training. In Proceedings of the European conference on computer vision, Cited by: §IV-A.
- [29] (2024) CDMAD: class-distribution-mismatch-aware debiasing for class-imbalanced semi-supervised learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, Cited by: §IV-B.
- [30] (2020) Revisiting knowledge distillation via label smoothing regularization. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, Cited by: §IV-B, §V-A.
- [31] (2019) Fedmd: heterogenous federated learning via model distillation. arXiv preprint arXiv:1910.03581. Cited by: §IV-C, §V-A.
- [32] (2022) Heterogeneous ensemble knowledge transfer for training large models in federated learning. International joint conferences on artificial intelligence. Cited by: §V-A.
Comments
· 0