Source-linked AI summary
Model-Contrastive Federated Learning
Qinbin Li, Bingsheng He, Dawn Song
TL;DR
Federated learning must handle heterogeneous local data without exchanging private data, and existing approaches perform poorly on deep-learning image datasets. MOON applies model-level contrastive learning to correct local training, and experiments show significant improvements over state-of-the-art federated methods across image classification tasks.
Problem
Heterogeneous local data distributions make federated learning difficult, while existing methods have not performed well on image datasets with deep-learning models.
Method
MOON corrects local training by maximizing agreement between representations learned by local and global models, using a supervised network with encoder, projection head, and output layer.
Results
MOON significantly outperforms existing federated learning approaches across CIFAR-10, CIFAR-100, and Tiny-Imagenet, improving accuracy by at least 2% in most cases.
Takeaways & Limitations
MOON introduces model-level contrastive learning for non-IID federated learning and potentially extends beyond vision because it does not require image inputs.
Takeaways & Limitations
The study leaves comparison with other contemporary work for future studies.
Abstract
from arXiv · showhide
Federated learning enables multiple parties to collaboratively train a machine learning model without communicating their local data. A key challenge in federated learning is to handle the heterogeneity of local data distribution across parties. Although many studies have been proposed to address this challenge, we find that they fail to achieve high performance in image datasets with deep learning models. In this paper, we propose MOON: model-contrastive federated learning. MOON is a simple and effective federated learning framework. The key idea of MOON is to utilize the similarity between model representations to correct the local training of individual parties, i.e., conducting contrastive learning in model-level. Our extensive experiments show that MOON significantly outperforms the other state-of-the-art federated learning algorithms on various image classification tasks.
1. Introduction
Federated learning avoids exchanging private data but struggles with heterogeneous local distributions, especially for deep-learning image tasks. MOON addresses this by contrasting model representations during local training and reports substantial gains over existing methods.
- Data are dispersed among parties, while privacy concerns and regulations prevent sending private data to a centralized server.
- Non-IID local data can move each party’s objective away from the global objective, degrading the averaged model’s performance.
- FedProx and SCAFFOLD address local non-IID updates, but can perform as poorly as FedAvg on deep-learning image datasets.
- MOON corrects local updates by maximizing agreement between representations learned by the current local model and the global model.
- At least 2% accuracy improvement over existing approaches occurs in most cases across CIFAR-10, CIFAR-100, and Tiny-Imagenet.
- 61.8% top-1 accuracy on CIFAR-100 with 100 parties exceeds the best existing result of 55%.
2. Background and Related Work
Federated learning research has targeted local training, aggregation, and personalization, while contrastive learning traditionally compares representations of different image views. MOON instead applies model-level contrastive learning in supervised federated learning.
- Federated learning: FedAvg sends a global model to parties, aggregates their locally updated models, and produces the next global model by averaging weights.
- Federated learning: Methods addressing non-IID FedAvg training modify either local training or server aggregation; this study belongs to the local-training category.
- Federated learning: FedProx limits local updates with a global-to-local ℓ2-norm proximal term, while SCAFFOLD uses control variates to correct updates.
- Federated learning: Prior evidence for FedProx and SCAFFOLD on image datasets with deep-learning models was limited, and experiments found little or no advantage over FedAvg.
- Federated learning: Aggregation methods include layer-wise weight matching, server momentum, and normalized local updates, which are potentially compatible with MOON’s local-training focus.
- Federated learning: Personalized federated learning seeks separate local models, whereas this paper studies one global model for all parties.
- Contrastive learning: Contrastive learning brings representations of augmented views of the same image together and separates representations from different images.
- Contrastive learning: SimCLR uses a base encoder, projection head, and NT-Xent loss to maximize agreement between differently augmented views.
3. Model-Contrastive Federated Learning
MOON addresses non-IID federated learning by correcting local updates through contrastive comparisons among model representations. It adds a model-contrastive term to supervised local training while retaining FedAvg-style aggregation.
- MOON targets non-IID federated learning, where skewed local objectives can move the averaged global model away from the global optimum.
- MOON uses the global model's representation as a positive reference and the previous local model's representation as a contrasting reference during local updates.For each input, representations are extracted from the global, previous local, and currently updated local models.
- The local objective combines supervised loss ℓsup with model-contrastive loss ℓcon, whose weight is controlled by hyper-parameter µ.The model-contrastive loss aims to decrease distance to the global representation and increase distance from the previous local representation.
- MOON builds on FedAvg: the server sends the global model, parties train locally with stochastic gradient descent, and the server updates the global model by weighted averaging.The framework can also operate when only a sampled subset of parties participates in each round.
- When local and global representations already match, MOON reduces to the same result as FedAvg, supporting robustness to different amounts of model drift.
- MOON differs from SimCLR by comparing representations learned by different models for supervised federated learning rather than augmented views for unsupervised centralized learning.
4. Experiments
Experiments evaluate MOON against federated and local-training baselines across image datasets, architectures, communication settings, party counts, and data heterogeneity levels. MOON consistently achieves stronger accuracy and communication efficiency, while its model-contrastive loss mitigates degradation from excessive local updates.
- Experimental Setup: MOON is compared with FedAvg, FedProx, SCAFFOLD, and SOLO on CIFAR-10, CIFAR-100, and Tiny-Imagenet using CNN and ResNet-50 architectures.The default setup uses Dirichlet-partitioned non-IID data, typically 10 parties, and reports tuned hyperparameters for MOON and FedProx.
- Accuracy Comparison: MOON is consistently the best approach across tasks and outperforms FedAvg by 2.6% accuracy on average.Under non-IID settings, SOLO performs much worse, while FedProx remains close to FedAvg and SCAFFOLD performs worse on CIFAR-100 and Tiny-Imagenet.
- Communication Efficiency: MOON needs about half as many communication rounds as FedAvg on CIFAR-100 and Tiny-Imagenet, with a CIFAR-10 speedup close to 4.The comparison measures rounds needed to reach the accuracy obtained by FedAvg after 100 rounds on CIFAR-10/100 or 20 rounds on Tiny-Imagenet.
- Local Epochs: MOON maintains the best accuracy across local-epoch settings and mitigates the negative effects of drift from too many local updates.With one local epoch, all approaches are relatively low and close in accuracy; excessive local epochs reduce accuracy for all approaches.
- Heterogeneity: MOON remains best across three CIFAR-100 heterogeneity levels and exceeds FedAvg by more than 2% when β = 5.For β = 5, FedProx is worse than FedAvg, whereas MOON retains its advantage.
- Loss Function: Model-contrastive loss improves accuracy over both no additional term and ℓ2-norm regularization.The ℓ2 approach does not improve over FedAvg on CIFAR-10 and remains below MOON on CIFAR-100 and Tiny-Imagenet.
5. Conclusion
The paper concludes that MOON addresses non-IID data in federated deep learning through model-level contrastive learning. Experiments report significant improvements over state-of-the-art approaches on image classification tasks, with potential applicability beyond vision because the method does not require image inputs.
- Conclusion: MOON is proposed as a simple federated-learning approach that applies contrastive learning at the model level to improve deep learning on non-IID data.The approach targets federated learning settings involving data silos and heterogeneous party distributions.
- Conclusion: MOON achieves significant improvement over state-of-the-art approaches on various image classification tasks.The conclusion summarizes the reported experimental outcome across the evaluated tasks.
- Conclusion: Because MOON does not require image inputs, it potentially can be applied to non-vision problems.This is presented as a potential scope beyond the reported image-classification experiments.
A. More Details of the Datasets
The dataset appendix reports dataset statistics and explains how the experiments evaluate test accuracy under Dirichlet-generated non-IID partitions. Additional figures visualize partitions with β = 0.1 and β = 5.
- Dataset Details: All datasets provide training and test splits, and reported accuracies are computed on the test datasets.The dataset statistics are summarized in Table 6 for β = 0.5.
- Dataset Details: Figures 9 and 10 visualize party data distributions for β = 0.1 and β = 5, respectively.These partitions are used in the heterogeneity experiments.
B. Projection Head
The projection-head study evaluates MOON with and without a projection head on CIFAR-10 and CIFAR-100. MOON benefits substantially from retaining the projection head.
- Projection Head: Removing the projection head changes the network architecture for all approaches and is evaluated on CIFAR-10 and CIFAR-100.The comparison uses top-1 accuracy with and without the projection head.
- Projection Head: A projection head improves MOON accuracy by about 2% on average.The study reports that MOON benefits substantially from including the projection head.
C. IID Partition
With IID local data partitions, MOON’s model-contrastive loss has little effect, so its accuracy remains close to FedAvg.
- MOON’s model-contrastive loss has little influence when local datasets are IID.
- MOON remains applicable without data heterogeneity, maintaining accuracy very close to FedAvg even with large µ.
D.1. Effect of µ
The model-contrastive loss benefits MOON when its weight is sufficiently large, while µ = 1 provides a practical untuned choice.
- µ = 0.1 keeps MOON’s accuracy very close to FedAvg because the model-contrastive loss has little impact.
- The best µ values are 5 for CIFAR-10 and 1 for both CIFAR-100 and Tiny-Imagenet.
- µ = 1 is a reasonable default, yielding at least 2% higher accuracy than FedAvg without parameter tuning.
D.2. Effect of temperature and output dimension
MOON is generally stable across temperature and projection-head dimensions, while it can also improve FedAvgM and incurs additional training cost.
- Effect of temperature and output dimension: The best temperatures are 0.5 for CIFAR-10, 1.0 for CIFAR-100, and 0.5 for Tiny-Imagenet.
- Effect of temperature and output dimension: The best projection-head output dimensions are 128 for CIFAR-10, 256 for CIFAR-100, and 128 for Tiny-Imagenet.
- Effect of temperature and output dimension: MOON is generally stable when temperature and output dimension change, and users may tune them for higher accuracy.
- Combining MOON and FedAvgM: MOON further improves FedAvgM by 2-3%, although FedAvgM is already better than FedAvg.
- Computation Cost: MOON trains more slowly than FedAvg because it introduces an additional loss term, with acceptable computation overhead especially on CIFAR-10 and CIFAR-100.
G. Number of Negative Pairs
MOON can use local models from previous rounds as negative pairs, but increasing their number does not improve accuracy; one negative pair performs best.
- MOON’s negative pair compares the current local model with the local model from the previous round.
- Using multiple negative pairs would incorporate local models from multiple previous rounds, up to the available number of rounds.
- k = 1 reproduces the MOON objective presented in the main paper.
- Increasing the number of negative pairs does not increase MOON’s accuracy, which is best when k = 1.