FedSelect

❗Check out our ICML 2023 Recap: https://aiuiuc.substack.com/p/icml-2023-recap

Recent advancements in federated learning (FL) seek to increase client-level performance by finetuning client parameters on local data or personalizing architectures for the local task. Existing methods for such personalization either prune a global model or fine-tune a global model on a local client distribution. However, these existing methods either personalize at the expense of retaining important global knowledge, or predetermine network layers for fine-tuning, resulting in suboptimal storage of global knowledge within client models. Enlightened by the lottery ticket hypothesis, we first introduce a hypothesis for finding optimal client subnetworks to locally fine-tune while leaving the rest of the parameters frozen. We then propose a novel FL framework, FedSelect, using this procedure that directly personalizes both client subnetwork structure and parameters, via the simultaneous discovery of optimal parameters for personalization and the rest of parameters for global aggregation during training. We show that this method achieves promising results on CIFAR-10.

In FedSelect, the input parameters $C, \theta_0, G, K, R, L$, and $p$ represent clients, the first global initialization, participation rate, GradLTN iterations, and personalization rate, respectively. The key step in FedSelect is performing LocalAlt on the shared and local parameter partition identified by GradLTN.

By the end of GradLTN, $v_k$ is identified as the set of appropriate parameters for dedicated local fine-tuning via LocalAlt; $u$ is also updated in LocalAlt and then averaged for global knowledge acquisition and retention. LocalAlt was introduced to update a defined set of shared and local parameters, $u$ and $v_k$, by alternating full passes of stochastic gradient descent between the two sets of parameters (Pillutla et al., 2022).

To the best of our knowledge, this is the first method to choose parameters for alternating updates in federated learning during training time.

A visualization on how parameters are aggregated and redistributed between communication rounds is shown below:

We consider a cross-silo setting in which the number of clients is low, but participation is high (Liu et al., 2022). We performed all experiments using a ResNet18 (He et al., 2015) backbone pretrained on ImageNet (Deng et al., 2009). We show results for our experimental setting on non-iid samples from CIFAR-10. Each client was allocated 20 training samples and 100 testing samples per class.

We would like to thank the anonymous reviewers for their valuable comments. We are thankful for the help of Chulin Xie and Wenxuan Bao for their valuable advising and support on this project. This research is part of the Delta research computing project, which is supported by the National Science Foundation (award OCI 2005572), and the State of Illinois. Delta is a joint effort of the University of Illinois at Urbana-Champaign and its National Center for Supercomputing Applications. We would also like to thank Amazon, Microsoft, and NCSA for providing conference travel funding, as well as ICML for providing registration funding.

This work appeared at FL-ICML 2023. It can be cited with the following BibTex:

Mini Presentation

An AI@UIUC Project © 2023