This is the official implementation of paper [FedTP: Federated Learning by Transformer Personalization].

Federated learning is an emerging learning paradigm where multiple clients collaboratively train a machine learning model in a privacy-preserving manner. Personalized federated learning extends this paradigm to overcome heterogeneity across clients by learning personalized models. Recent works have shown that the self-attention mechanism in Transformer models is robust to distribution shifts. As such, there have been some initial attempts to apply Transformers to federated learning. However, the impacts of federated learning algorithms on self-attention have not yet been studied.

We propose FedTP, a novel Transformer-based federated learning framework with personalized self-attention to better handle data heterogeneity among clients. FedTP learns a personalized self-attention layer for each client while the parameters of the other layers are shared among the clients. Additionally, the server learns a hypernetwork that generates projection matrices in the selfattention layers to produce client-wise queries, keys, and values. This hypernetwork is an effective way of sharing parameters across clients while maintaining the flexibility of a personalized Transformer. Notably, FedTP offers a convenient environment for performing a range of image and language tasks using the same federated network architecture –all of which benefits from Transformer personalization. Extensive experiments on standard personalized federated learning benchmarks with different non-IID data settings show that FedTP yields state-of-the-art performance.

To install requirements:

pip install -r requirements.txt

Additionally, To run FedTP+KNN, FAISS should be installed. Instructions for the installation of FAISS can be found here

We provide three federated benchmark datasets spanning image classification task (CIFAR10 and CIFAR100) and language modeling(Shakespeare).

For CIFAR10 and CIFAR100 dataset, download and unzip data under 'data' file catalog. Or simply run any algorithms with CIFAR10/CIFAR100 dataset, the program will download data automatically.

Shakespeare dataset was naturally partitioned by assigning all lines from the same characters to the same client. See the README.md files in data/shakespeare for instructions on generating data before running experiments.

The following table summarizes the datasets and models

Here is one example to run our FedTP:

python main.py --model=vit \
    --dataset=cifar10 \
    --alg=FedTP \
    --lr=0.01 \
    --batch-size=64 \
    --epochs=10 \
    --n_parties=10 \
    --mu=0.01 \
    --rho=0.9 \
    --comm_round=50 \
    --partition=noniid-labeldir \
    --beta=0.5\
    --device='cuda:0'\
    --datadir='./data/' \
    --logdir='./logs/' \
    --noise=0 \
    --sample=0.1 \
    --init_seed=0

Three demos for each dataset are provided in scripts folder.

To simulate non-IID scenarios for CIFAR-10/CIFAR-100, we follow two common split designs. You can call function get_partition_dict() in main.py to access net_dataidx_map. net_dataidx_map is a dictionary. Its keys are party ID, and the value of each key is a list containing index of data assigned to this party. For our experiments, we usually set init_seed=0. The default value of noise is 0 unless stated. We list the way to get our data partition as follow.

• Dirichlet Partition: partition=noniid-labeldir/noniid-labeldir100. The former is for CIFAR-10 dataset and the later is for CIFAR-100 dataset. beta controls degree of data heterogeneity.
• Pathological Partition: partition=noniid-labeluni. For CIFAR-10 and CIFAR-100 dataset.

Here is explanation of parameter for function get_partition_dict().