cs.LGMar 3, 2026

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

Abstract

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

Explore similar work

CardsList
  1. Class-wise Contribution Estimation via Logit Maximization for Federated Learning

    May 17, 2026Asim Ukaye, Nurbek Tastan, Mubarak Abdu-Aguye +1Class-Imbalanced LearningFederated Learning Aggregation

  2. Fast Convergence through Distributed Augmentation for Class-Imbalanced Federated Learning

    Oct 4, 2026Arathi Nair M, J. Harshan, Anwitaman DattaClass-Imbalanced LearningData Augmentation

  3. FedHPro: Federated Hyper-Prototype Learning via Gradient Matching

    May 13, 2026Huan Wang, Jun Shen, Haoran Li +6Federated LearningHeterogeneous Federated Learning