Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Breaking the Prototype Bias Loop: Confidence-Aware Federated Contrastive Learning for Highly Imbalanced Clients

Tian-Shuang Wu, Shen-Huan Lyu, Ning Chen, Yi-Xiao He, Bing Tang, Baoliu Ye, Qingfu Zhang

TL;DR

Confidence-Aware Federated Contrastive Learning (CAFedCL) is proposed, a novel framework that improves the prototype aggregation mechanism and strengthens the contrastive alignment guided by prototypes, and which consistently outperforms representative federated baselines in both accuracy and client fairness.

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.

Breaking the Prototype Bias Loop: Confidence-Aware Federated Contrastive Learning for Highly Imbalanced Clients

TL;DR

Confidence-Aware Federated Contrastive Learning (CAFedCL) is proposed, a novel framework that improves the prototype aggregation mechanism and strengthens the contrastive alignment guided by prototypes, and which consistently outperforms representative federated baselines in both accuracy and client fairness.

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.
Paper Structure (19 sections, 2 theorems, 15 equations, 4 figures, 2 tables, 1 algorithm)

This paper contains 19 sections, 2 theorems, 15 equations, 4 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Federated learning
  4. Contrastive learning in FL
  5. Preliminaries
  6. Federated Prototype Learning
  7. Estimation Bias under Imbalance
  8. Method
  9. Prototype-guided Contrastive Learning for Federated Alignment
  10. Prototype Bias Loop in Federated Contrastive Learning
  11. Class-wise Confidence-weighted Aggregation
  12. Additional Stabilizers: Geometry Regularization and Tail Augmentation
  13. Experiment
  14. Experimental Setup
  15. Performance
  16. ...and 4 more sections

Key Result

Lemma 1

Under Assumption ass:variance, the estimation error of the local prototype satisfies where $n^{\mathrm{eff}}_{k,c}=n_{k,c}+\gamma m_{k,c}$ is the effective sample size, $m_{k,c}$ is the number of synthetic samples (if any), and $\gamma\in[0,1]$ discounts their reliability.

Figures (4)

  • 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.
  • Figure 3: The accuracy of CAFedCL and its ablated variants under the default settings.
  • Figure 4: The accuracy of CAFedCL with different hyperparameters.

Theorems & Definitions (3)

  • Lemma 1
  • Proposition 2
  • Remark 3