One Training Fits All: Addressing Model-Heterogeneity Federated Learning via Architecture Probing

Feijie Wu, Yaqing Wang, Tianci Liu, Xingchen Wang, Lu Su, Jing Gao

Keywords: Federated Learning, Model Heterogeneity, Model Pruning

Abstract: Model-heterogeneity federated learning (FL) is a flexible setting where a client trains a model subject to its local computation capacity. Towards the scenario, partial averaging extracts the clients' models from a global model so that the aggregation of each model parameter is identified. While existing models can only generate submodels with predefined settings established during training, our approach utilizes a trainable probabilistic masking strategy named FedMAP, enabling the dynamic creation of customized model sizes aligned with the client's budget. In detail, the clients find the best model architectures based on their local datasets and computation resources, and the FL server merges these local optimal architectures into a probabilistic mask. In the end, we attain a stable probabilistic mask, with which we can generate arbitrary models for evaluation or update the counterpart of the model parameters while training with the clients' data. Our experiments validate the effectiveness of the proposed FedMAP from two aspects: (i) It can improve the state-of-the-art approaches to heterogeneous model updates, especially for those small-size models; and (ii) We can extract the submodels whose sizes never appear in training with exceptional performance.

Primary Area: general machine learning (i.e., none of the above)

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 6686