Model-Contrastive Federated Learning

Note: If you use this baseline in your work, please remember to cite the original authors of the paper as well as the Flower paper.

Paper: arxiv.org/abs/2103.16257

Authors: Qinbin Li, Bingsheng He, Dawn Song

Abstract: 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: modelcontrastive 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.

About this baseline

What’s implemented: The code in this directory replicates the experiments in Model-Contrastive Federated Learning (Li et al., 2021), which proposed the MOON algorithm. Concretely ,it replicates the results of MOON for CIFAR-10 and CIFAR-100 in Table 1.

Datasets: CIFAR-10 and CIFAR-100

Hardware Setup: The experiments are run on a server with 4x Intel Xeon Gold 6226R and 8x Nvidia GeForce RTX 3090. A machine with at least 1x 16GB GPU should be able to run the experiments in a reasonable time.

Contributors: Qinbin Li

Description: MOON requires to compute the model-contrastive loss in local training, which requires access to the local model of the previous round (Lines 14-17 of Algorithm 1 of the paper). Since currently FlowerClient does not preserve the states when starting a new round, we store the local models into the specified model.dir in local training indexed by the client ID, which will be loaded to the corresponding client in the next round.

Experimental Setup

Task: Image classification.

Model: This directory implements two models as same as the paper:

• A simple-CNN with a projection head for CIFAR-10

• A ResNet-50 with a projection head for CIFAR-100.

Dataset: This directory includes CIFAR-10 and CIFAR-100. They are partitioned in the same way as the paper. The settings are as follow:

Dataset	partitioning method
CIFAR-10	Dirichlet with beta 0.5
CIFAR-100	Dirichlet with beta 0.5

Training Hyperparameters:

Description	Default Value
number of clients	10
number of local epochs	10
fraction fit	1.0
batch size	64
learning rate	0.01
mu	1
temperature	0.5
alg	moon
seed	0
service_device	cpu
number of rounds	100
client resources	{‘num_cpus’: 2.0, ‘num_gpus’: 0.0 }

Environment Setup

To construct the Python environment follow these steps:

# Set local python version via pyenv
pyenv local 3.10.6
# Then fix that for poetry
poetry env use 3.10.6
# Then install poetry env
poetry install

# Activate the environment
poetry shell

Running the Experiments

First ensure you have activated your Poetry environment (execute poetry shell from this directory). To run MOON on CIFAR-10 (Table 1 of the paper), you should run:

python -m moon.main --config-name cifar10 

To run MOON on CIFAR-100 (Table 1 of the paper), you should run:

python -m moon.main --config-name cifar100

You can also run FedProx on CIFAR-10:

python -m moon.main --config-name cifar10_fedprox

To run FedProx on CIFAR-100:

python -m moon.main --config-name cifar100_fedprox

Expected Results

You can find the output logs of a single run in this link. After running the above commands, you can see the accuracy list at the end of the output, which is the test accuracy of the global model. For example, in one running, for CIFAR-10 with MOON, the accuracy after running 100 rounds is 0.7071.

For CIFAR-10 with FedProx, the accuracy after running 100 rounds is 0.6852. For CIFAR100 with MOON, the accuracy after running 100 rounds is 0.6636. For CIFAR100 with FedProx, the accuracy after running 100 rounds is 0.6494. The results are summarized below:

	CIFAR-10	CIFAR-100
MOON	0.7071	0.6636
FedProx	0.6852	0.6494

Figure 6

You can find the curve comparing MOON and FedProx on CIFAR-10 and CIFAR-100 below.

You can tune the hyperparameter mu for both MOON and FedProx by changing the configuration file in conf.

Figure 8(a)

You can run the experiments in Figure 8 of the paper. To run MOON (mu=10) on CIFAR-100 with 50 clients (Figure 8(a) of the paper):

python -m moon.main --config-name cifar100_50clients

To run FedProx on CIFAR-100 with 50 clients (Figure 8(a) of the paper):

python -m moon.main --config-name cifar100_50clients_fedprox

You can find the curve presenting MOON and FedProx below.

You may also run MOON on CIFAR-100 with 100 clients (Figure 8(b) of the paper):

python -m moon.main --config-name cifar100_100clients