PaperlayerView exact revision on arXiv

Source-linked AI summary

FedProc: Prototypical Contrastive Federated Learning on Non-IID data

Xutong Mu, Yulong Shen, Ke Cheng, Xueli Geng, Jiaxuan Fu, Tao Zhang, Zhiwei Zhang

arXiv:2109.12273v1cs.LGcs.DC

TL;DR

Non-IID client data can make local federated training diverge from the global optimum, leaving image-classification performance unsatisfactory. FedProc uses global class prototypes with a tailored local architecture and prototypical contrastive loss to align local training. It reports higher accuracy than existing federated methods, including 70.6% top-1 accuracy versus 61.8% on CIFAR-100 with 100 clients.

  • Problem

    Non-IID client data causes local training to drift from the global optimum, and effective federated image classification remains an open problem.

  • Method

    FedProc broadcasts aggregated global class prototypes and combines a tailored local network with a global prototypical contrastive loss to regulate local training.

  • Results

    70.6% top-1 accuracy is achieved on CIFAR-100 with 100 clients versus 61.8% for the best existing approach, with acceptable computation cost.

  • Takeaways & Limitations

    FedProc provides a simple federated learning framework that demonstrates an advantage on non-IID image classification across multiple datasets.

  • Takeaways & Limitations

    The problem formulation assumes a multi-classification task with locally held datasets and labeled samples represented by feature vectors.

Abstract

from arXiv · show

Federated learning allows multiple clients to collaborate to train high-performance deep learning models while keeping the training data locally. However, when the local data of all clients are not independent and identically distributed (i.e., non-IID), it is challenging to implement this form of efficient collaborative learning. Although significant efforts have been dedicated to addressing this challenge, the effect on the image classification task is still not satisfactory. In this paper, we propose FedProc: prototypical contrastive federated learning, which is a simple and effective federated learning framework. The key idea is to utilize the prototypes as global knowledge to correct the local training of each client. We design a local network architecture and a global prototypical contrastive loss to regulate the training of local models, which makes local objectives consistent with the global optima. Eventually, the converged global model obtains a good performance on non-IID data. Experimental results show that, compared to state-of-the-art federated learning methods, FedProc improves the accuracy by $1.6\%\sim7.9\%$ with acceptable computation cost.

1 Introduction

Non-IID client data can cause local training to drift from the global optimum, limiting federated image classification. FedProc addresses this with global class prototypes, tailored local training, and prototypical contrastive learning, achieving stronger accuracy at acceptable computation cost.

  • Non-IID client distributions cause local model training to drift from the global optimum, making effective federated learning an open problem.
  • FedProc introduces global class prototypes as shared knowledge to correct each client’s local training.The server aggregates client class prototypes and broadcasts the resulting global prototypes.
  • FedProc is evaluated on CIFAR-10, CIFAR-100, and Tiny-ImageNet, where it significantly outperforms state-of-the-art federated learning algorithms.
  • 70.6% top-1 accuracy is achieved on CIFAR-100 with 100 clients, compared with 61.8% for the best existing approach.The paper reports accuracy improvements of 1.6% on CIFAR-10 and more than 7% on CIFAR-100 and Tiny-ImageNet, with acceptable computation cost.
  • FedProc combines a hybrid local network architecture with a global prototypical contrastive loss to align local objectives with the global optimum.The loss pulls samples toward their class prototype and pushes them away from prototypes of other classes.

2 Background and Related Work

Federated learning trains a global model through local client updates and server aggregation, while related work addresses non-IID data through aggregation or local-training stabilization. FedProc instead uses supervised prototypical contrastive learning and prototypes to reduce embedding-space inconsistency across clients.

  • 2.1 Federated Learning: Federated learning follows a four-step protocol: server initialization and distribution, local SGD updates, client uploads, and server averaging.
  • 2.1 Federated Learning: Prior non-IID methods mainly improve model aggregation or stabilize local training by limiting local-model deviation from a global model.
  • 2.2 Contrastive Learning: Contrastive learning learns embeddings by pulling positive pairs together and pushing negative pairs apart; supervised variants use labels to form those pairs.
  • 2.3 Contrastive Learning in Federated Learning: Federated contrastive-learning methods include unsupervised approaches for distributed unlabeled data and MOON’s model-based comparison strategy.
  • 2.3 Contrastive Learning in Federated Learning: FedProc targets supervised learning and introduces prototypes to address inconsistency in clients’ embedding spaces.

3 Prototypical Contrastive Federated Learning

FedProc uses global class prototypes to correct local training under non-IID data. Its local architecture and hybrid objective jointly learn representations and classifiers through federated optimization.

  • 3 Prototypical Contrastive Federated Learning: FedProc uses global class-prototypes as knowledge to correct each client’s local training and align local objectives with global optima.Samples are pulled toward their class prototype and away from prototypes of other classes.
  • Local Network Architecture: The local network contains a base encoder, projection head, and output layer that produce representations z and class-prediction logits s.The representation z supports prototypical contrastive learning, while logits s support classifier training.
  • Local Objective: The local loss combines global prototypical contrastive loss ℓgpc with cross-entropy loss ℓce, weighted by a communication-round coefficient α.The weighting progressively transitions local learning from feature learning toward classifier learning as rounds increase.
  • Global Prototypical Contrastive Loss: The global prototypical contrastive loss makes client samples close to their class prototype and far from prototypes of other classes.Each client’s class prototype is the mean representation of samples belonging to that class.
  • FedProc Framework: FedProc follows a federated training loop in which clients receive the global model and prototypes, update locally with SGD, and return model and prototype information.The framework operates over communication rounds, local epochs, classes, clients, and learning rate η.

4 Experiment

The experiments compare FedProc with several federated learning baselines, including MOON, FedAvg, FedProx, SCAFFOLD, and SOLO.

  • 4 Experiment: FedProc is compared with MOON, FedAvg, FedProx, SCAFFOLD, and SOLO as state-of-the-art or baseline federated learning approaches.The implementation uses PyTorch and runs on machines equipped with two NVIDIA GeForce RTX 3090 GPUs.

4.1 Experimental Setup

The experimental setup evaluates FedProc on three standard image-classification datasets under client distributions controlled by a Dirichlet concentration parameter.

  • 4.1 Experimental Setup: Experiments use CIFAR-10, CIFAR-100, and Tiny-ImageNet, containing 10, 100, and 200 classes, respectively.CIFAR-10 and CIFAR-100 each contain 60,000 images, while Tiny-ImageNet contains 100,000 images.
  • 4.1 Experimental Setup: A Dirichlet distribution allocates each class’s instance proportion to clients, with concentration parameter β controlling client identicalness.The approaches use the same local-network modules for fair comparison.

4.2 Accuracy Results

FedProc achieves the strongest reported top-1 accuracy across the evaluated non-IID image-classification settings, including varying heterogeneity and training parameters. Its gains come with slower convergence and acceptable computation cost.

  • Overall comparison: FedProc consistently outperforms other federated learning methods across the evaluated datasets, exceeding MOON by 1.6%∼7.9% accuracy.FedAvg, FedProx, and SCAFFOLD generally perform worse, while MOON improves over FedAvg by 1.3%∼3%.
  • Communication rounds: FedProc achieves the best final performance but converges more slowly as communication rounds increase.The paper attributes this pattern to stronger early feature learning followed by classifier learning in later stages.
  • Local epochs: Accuracy is highest for most methods at E = 10 local epochs; smaller E undertrains local networks, whereas E > 10 causes overfitting on skewed data.The resulting overfitting reduces the accuracy of the global model.
  • Data heterogeneity: FedProc remains the most accurate method across all tested CIFAR-100 heterogeneity levels, outperforming MOON by 7.6%, 7.1%, and 4.9%.The comparisons correspond to β = 5, 0.5, and 0.1, respectively; smaller β indicates more skewed data.
  • Loss and architecture: FedProc's hybrid local network and global prototypical contrastive loss improve representations while maintaining compatibility between feature and classifier learning.Two-stage training performs worse because it breaks this compatibility.

4.3 Computation Cost

FedProc incurs additional computation because it adds loss terms to the FedAvg training process, while FedAvg remains the fastest baseline.

  • FedAvg has the lowest average training time per round among the compared methods.The comparison measures all methods on the same machines.
  • FedProx, MOON, and FedProc take additional time because they introduce loss items beyond FedAvg.
  • SCAFFOLD adds computation through additional control variables for the server and clients.

4.4 Scalability

FedProc is evaluated with different client counts and sampling rates on CIFAR-100, showing strong scalability and higher accuracy than competing methods.

  • Experimental setting: FedProc is evaluated on CIFAR-100 with 50 clients at full participation and 100 clients with a 20% sampling rate.With 100 clients, 20 participate in each training round.
  • Results: FedProc achieves accuracy far higher than the other methods as the number of clients varies.
  • Results: 9.3% accuracy is FedProc’s advantage over MOON at T = 200 rounds and m = 50 clients.
  • Interpretation: The authors attribute FedProc’s scalability to prototypical contrastive learning, which aligns local objectives with global optima.They state that this keeps performance from being affected as the client count increases.

5 Conclusion

FedProc addresses non-IID federated learning by using class prototypes to correct local training. Its architecture and loss align local objectives with global optima, supporting good global-model classification performance.

  • FedProc is presented as a simple and effective framework for addressing non-IID data in federated learning.
  • Class prototypes provide global knowledge that corrects local training across clients.
  • A local network architecture and global prototypical contrastive loss make local objectives consistent with global optima.
  • The resulting global model achieves good classification performance according to extensive experiments.
Loading 2109.12273v1…