Federated Learning with Partial Model Personalization

Krishna Pillutla, Kshitiz Malik, Abdelrahman Mohamed, Michael Rabbat, Maziar Sanjabi, Lin Xiao

Keywords: personalization, federated learning, partial personalization, adapter modules, nonconvex minimization

Abstract: We propose and analyze a general framework of federated learning with partial model personalization. Compared with full model personalization, partial model personalization relies on domain knowledge to select a small portion of the model to personalize, thus imposing a much smaller on-device memory footprint. We propose two federated optimization algorithms for training partially personalized models, where the shared and personal parameters are updated either simultaneously or alternately on each device, but only the shared parameters are communicated and aggregated at the server. We give convergence analyses of both algorithms for minimizing smooth nonconvex functions, providing theoretical support of them for training deep learning models. Our experiments on real-world image and text datasets demonstrate that (a) partial model personalization can obtain most of the benefit of full model personalization with a small fraction of personalized parameters, and, (b) the alternating update algorithm often outperforms the simultaneous update algorithm.

One-sentence Summary: We personalize a part of the model for federated learning and find that it gives most of the benefits of full personalization and that an alternating optimization algorithm works the best.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2204.03809/code)