Breaking the Prototype Bias Loop: Confidence-Aware Federated Contrastive Learning for Highly Imbalanced Clients
Authors: Tian-Shuang Wu, Shen-Huan Lyu, Ning Chen, Yi-Xiao He, Bing Tang, Baoliu Ye, Qingfu Zhang
Organizations: Key Laboratory of Water Big Data Technology of Ministry of Water Resources, College of Computer Science and Software Engineering, Hohai University, Nanjing, China · Department of Computer Science, City University of Hong Kong, Hong Kong, China · State Key Laboratory for Novel Software Technology, Nanjing University, Nanjing, China · School of Artificial Intelligence and Information Technology, Nanjing University of Chinese Medicine, Nanjing, China
Local class imbalance and data heterogeneity across clients often trap prototype-based federated contrastive learning in a prototype bias loop: biased local prototypes induced by imbalanced data are aggregated into biased global prototypes, which are repeatedly reused as contrastive anchors, accumulating errors across communication rounds. To break this loop, we propose Confidence-Aware Federated Contrastive Learning (CAFedCL), a novel framework that improves the prototype aggregation mechanism and strengthens the contrastive alignment guided by prototypes. CAFedCL employs a confidence-aware aggregation mechanism that leverages predictive uncertainty to downweight high-variance local prototypes. In addition, generative augmentation for minority classes and geometric consistency regularization are integrated to stabilize the structure between classes. From a theoretical perspective, we provide an expectation-based analysis showing that our aggregation reduces estimation variance, thereby bounding global prototype drift and ensuring convergence. Extensive experiments under varying levels of class imbalance and data heterogeneity demonstrate that CAFedCL consistently outperforms representative federated baselines in both accuracy and client fairness.
Figures & tables
Figure 1: The bias loop in federated contrastive learning.
Figure 2: The data distributions of CIFAR-10, EMNIST, and CIFAR-100 in the default settings.
Practical heterogeneous ( α=0.1,K=20 )
Pathological heterogeneous ( K=20,IR=10 )
Method
CIFAR-10
CIFAR-100
EMNIST
CIFAR-10
CIFAR-100
EMNIST
Acc. ↑
Std. ↓
Acc. ↑
Std. ↓
Acc. ↑
Std. ↓
Acc. ↑
Std. ↓
Acc. ↑
Std. ↓
Acc. ↑
Std. ↓
FedAvg
60.32
12.03
29.74
4.08
82.95
11.59
48.73
19.15
24.47
7.58
72.38
14.96
FedProx
61.82
10.07
31.42
4.25
84.27
8.13
60.29
9.11
30.60
5.49
72.45
7.21
MOON
82.06
11.19
46.23
5.15
92.77
6.53
84.94
13.44
48.70
6.95
93.24
5.40
FedProto
88.33
8.13
52.81
6.58
95.50
1.69
89.45
7.78
52.14
5.90
94.88
2.63
Table 1: Test accuracy (%) on three datasets under different heterogeneous settings. Results are reported as accuracy and standard deviation over 5 runs.
Method
Dirichlet parameter
Imbalance ratio
Class per client
Client amount
α=0.05
α=1.0
IR=50
IR=100
C/Client=10
C/Client=50
50 Clients
100 Clients
Acc. ↑
Std. ↓
Acc. ↑
Std. ↓
Acc. ↑
Std. ↓
Acc. ↑
Std. ↓
Acc. ↑
Std. ↓
Acc. ↑
Std. ↓
Acc. ↑
Std. ↓
Acc. ↑
Std. ↓
FedAvg
52.54
12.24
86.67
6.19
43.28
22.81
35.29
17.50
20.87
7.12
29.66
2.59
56.73
12.88
54.14
13.65
FedProx
55.48
14.05
87.93
6.44
48.30
11.32
37.33
13.72
28.56
6.28
32.70
2.85
58.50
10.12
55.94
10.47
MOON
65.79
10.40
90.95
5.52
54.23
15.73
42.61
18.25
56.33
4.95
25.49
5.22
79.51
11.58
78.19
10.66
FedProto
66.86
8.35
91.27
5.24
55.48
9.56
44.46
10.29
62.12
4.66
35.06
4.41
83.07
9.92
80.44
10.25
Table 2: Performance (%) under heterogeneous settings. Results are reported as accuracy and standard deviation over 5 runs.
Figure 3: The accuracy of CAFedCL and its ablated variants under the default settings.
Figure 4: The accuracy of CAFedCL with different hyperparameters.
Federated learning (FL) enables collaborative learning of computer vision models, where privacy and regulatory constraints prevent centralizing data across devices or organizations. However, practical FL deployments often exhibit severe class imbalance and label skew, causing standard aggregation protocols to overfit dominant clients and degrade minority-class performance. We propose a data-free, class-wise contribution estimation and aggregation framework based on logit maximization (CELM) that does not require sharing raw data, client metadata, or auxiliary public datasets. The FL server probes client updates to obtain class-wise evidence scores and assembles a cross-client evidence matrix, which quantifies both per-class competence and class coverage. Using this matrix, we compute contribution weights that upweight clients providing strong, discriminative evidence for underrepresented classes. The resulting aggregation is stable due to simplex constraints and momentum smoothing, and it remains compatible with standard FL training pipelines. We evaluate the approach on representative vision benchmarks under controlled non-IID and pathological label splits, demonstrating that CELM-based aggregation improves robustness to imbalance and statistical heterogeneity, while yielding better performance without requiring any additional data exchange.
Asim Ukaye, Nurbek Tastan, Mubarak Abdu-Aguye +1
Mohamed bin Zayed University of Artificial Intelligence, UAE · Michigan State University, USA
In federated learning, mitigating class imbalance is essential to improve minority-class performance. A common approach to address this problem is to augment minority-class samples to achieve local class balance. Existing approaches treat augmentation as a heuristic and do not establish how the amount of augmentation influences the convergence of federated learning, leading to excessive augmentation and increased training time. To address this limitation, we first establish the relationship between augmentation and the convergence behavior of federated learning. Leveraging this insight, we propose DAFL, a distributed augmentation framework that determines the minimum augmentation required for each client-class pair by jointly minimizing augmentation and training time while constraining global class imbalance, thereby improving minority-class F1-score. Experimental results demonstrate that DAFL consistently improves minority-class F1-score while substantially reducing training time, particularly under severe global class imbalance and high label proportion imbalance.
Arathi Nair M, J. Harshan, Anwitaman Datta
Amar Nath and Shashi Khosla School of Information Technology, Indian Institute of Technology Delhi, India · Department of Electrical Engineering, Indian Institute of Technology Delhi, India · College of Computing and Data Science, Nanyang Technological University, Singapore
Federated Learning (FL) enables collaborative training of distributed clients while protecting privacy. To enhance generalization capability in FL, prototype-based FL is in the spotlight, since shared global prototypes offer semantic anchors for aligning client-specific local prototypes. However, existing methods update global prototypes at the prototype-level via averaging local prototypes or refining global anchors, which often leads to semantic drift across clients and subsequently yields a misaligned global signal. To alleviate this issue, we introduce hyper-prototypes, defined by a set of learnable global class-wise prototypes to preserve underlying semantic knowledge across clients. The hyper-prototypes are optimized via gradient matching to align with class-relevant characteristics distilled directly from clients' real samples, rather than prototype-level descriptors. We further propose FedHPro, a Federated Hyper-Prototype Learning framework, to leverage hyper-prototypes to promote inter-class separability via mutual-contrastive learning with client-specific margin, while encouraging intra-class uniformity through a consistency penalty. Comprehensive experiments under diverse heterogeneous scenarios confirm that 1) hyper-prototypes produce a more semantically consistent global signal, and 2) FedHPro achieves state-of-the-art performance on several benchmark datasets. Code is available at https://github.com/mala-lab/FedHPro.
Huan Wang, Jun Shen, Haoran Li +6
School of Computing and Information Technology, University of Wollongong, Wollongong, Australia · School of Computing and Information Systems, Singapore Management University, Singapore, Singapore · Monash University, Australia +3