Federated Learning with Partial Model Personalization

1 INTRODUCTION . Federated Learning ( McMahan et al. , 2017 ) has emerged as a powerful paradigm for distributed and privacy-preserving machine learning over a large number of edge devices ( see Kairouz et al. , 2021 , and references therein ) . We consider a typical setting of Federated Learning ( FL ) with n devices ( also called clients ) , where each device i has a training dataset of Ni samples zi,1 , · · · , zi , Ni . Let w ∈ Rd represent the parameters of a ( supervised ) learning model and fi ( w , zi , j ) be the loss of the model on the training example zi , j . Then the loss function associated with device i is Fi ( w ) = ( 1/Ni ) ∑Ni j=1 fi ( w , zi , j ) . A common objective of FL is to find model parameters that minimize the weighted average loss across all devices ( without transferring the datasets ) : minimize w n∑ i=1 αiFi ( w ) , ( 1 ) where weights αi are nonnegative and satisfy ∑n i=1 αi = 1 . A common practice is to choose the weights as αi = Ni/N where N = ∑n k=1Nk , which corresponds to minimizing the unweighted average loss across all samples from the n devices : ( 1/N ) ∑n i=1 ∑Ni j=1 fi ( w , zi , j ) . The main motivation for minimizing the average loss over all devices is to leverage their collective statistical power for better generalization , because the amount of data on each device can be very limited . This is especially important for training modern deep learning models with large number of parameters . However , this argument assumes that the datasets from different devices are sampled from the same , or at least very similar , distributions . Given the diverse characteristics of the users and increasing trend of personalized on-device services , such an i.i.d . assumption may not hold in practice . Thus , the one-model-fits-all formulation in ( 1 ) can be less effective and even undesirable . Several approaches have been proposed for personalized FL , including ones based on multi-task learning ( Smith et al. , 2017 ) , meta learning ( Fallah et al. , 2020 ) , and proximal methods ( Dinh et al. , 2020 ; Li et al. , 2021 ) . A simple formulation that captures their main idea is minimize w0 , { wi } ni=1 n∑ i=1 αi ( Fi ( wi ) + λi 2 ‖wi − w0‖2 ) , ( 2 ) where wi for i = 1 , . . . , n are personalized model parameters at the devices , w0 is a reference model maintained by the server , and the λi ’ s are regularization weights that control the extent of personalization . A major disadvantage of the formulation ( 2 ) , which we call full model personalization , is that it requires twice the memory footprint of the model , wi and w0 at each device , which severely limits the size of trainable models . On the other hand , the flexibility of full model personalization can be unnecessary . Modern deep learning models are composed of many simple functional units and are typically organized into layers or a more general interconnected architecture . Personalizing the “ right ” components , selected with domain knowledge , may result in a substantial benefit with only a small increase in memory footprint . In addition , partial model personalization can be less susceptible to “ catastrophic forgetting ” ( McCloskey & Cohen , 1989 ) , where a large model finetuned on a small local dataset forgets the original ( non-personalized ) task , leading to a degradation of test performance . We propose a framework for FL with partial model personalization . Specifically , we partition the model parameters into two groups : the shared parameters u ∈ Rd0 and the personal parameters vi ∈ Rdi for i = 1 , . . . , n. The full model on device i is denoted as wi = ( u , vi ) , and the local loss function is Fi ( u , vi ) = ( 1/Ni ) ∑Ni i=1 fi ( ( u , vi ) , zi , j ) . Our goal is to solve the optimization problem minimize u , { vi } ni=1 n∑ i=1 αiFi ( u , vi ) . ( 3 ) Notice that the dimensions of vi can be different across the devices , allowing the personal components of the model to have different number of parameters or even different architecture . We investigate two FL algorithms for solving problem ( 3 ) : FedSim , a simultaneous update algorithm and FedAlt , an alternating update algorithm . Both algorithms follow the standard FL protocol . During each round , the server randomly selects a subset of the devices for update and broadcasts the current global version of the shared parameters to devices in the subset . Each selected device then performs one or more steps of ( stochastic ) gradient descent to update both the shared parameters and the personal parameters , and sends the updated shared parameters to the server for aggregation . The updated personal parameters are kept local at the device to serve as the initial states when the device is selected for another update . In FedSim , the shared and personal parameters are updated simultaneously during each local iteration . In FedAlt , the devices first update the personal parameters with the received shared parameters fixed and then update the shared parameters with the new personal parameters fixed . We provide convergence analysis and empirical evaluation of both methods . The main contributions of this paper are summarized as follows : • We propose a general framework of FL with partial model personalization , which relies on domain knowledge to select a small portion of the model to personalize , thus imposing a much smaller memory footprint on the devices than full model personalization . This framework unifies existing work on personalized FL and allows arbitrary partitioning of deep learning models . • We provide convergence guarantees for the FedSim and FedAlt methods in the general ( smooth ) nonconvex setting . While both methods have appeared in the literature previously , they are either used without convergence analysis or with results on limited settings ( assuming convexity or full participation ) Our analysis provides theoretical support for the general nonconvex setting with partial participation . The analysis of FedAlt with partial participation is especially challenging and we develop a novel technique of virtual full participation to overcome the difficulties . • We conduct extensive experiments on image classification and text prediction tasks , exploring different model personalization strategies for each task , and comparing with several strong baselines . Our results demonstrate that partial model personalization can obtain most of the benefit of full model personalization with a small fraction of personalized parameters , and FedAlt often outperforms FedSim . • Our experiments also reveal that personalization ( full or partial ) may lead to worse performance for some devices , despite improving the average . Typical forms of regularization such as weight decay and dropout do not mitigate this issue . This phenomenon has been overlooked in previous work and calls for future research to improve both performance and fairness . Related work . Specific forms of partial model personalization have been considered in previous works . Liang et al . ( 2019 ) propose to personalize the input layers to learn a personalized representation per-device ( Figure 1a ) , while Arivazhagan et al . ( 2019 ) and Collins et al . ( 2021 ) propose to personalize the output layer while learning a shared representation with the input layers ( Figure 1b ) . Both FedSim and FedAlt have appeared in the literature before , but the scope of their convergence analysis is limited . Specifically , Liang et al . ( 2019 ) , Arivazhagan et al . ( 2019 ) and Hanzely et al . ( 2021 ) use FedSim , while Collins et al . ( 2021 ) and Singhal et al . ( 2021 ) proposed variants of FedAlt . Notably , Hanzely et al . ( 2021 ) establish convergence of FedSim with full device participation in the convex and non-convex cases , while Collins et al . ( 2021 ) prove the linear convergence of FedAlt for a two-layer linear network where Fi ( · , vi ) and Fi ( u , · ) are both convex for fixed vi and u respectively . We analyze both FedSim and FedAlt in the general nonconvex case with partial device participation , hence addressing a more general and practical setting . While we primarily consider the problem ( 3 ) in the context of partial model personalization , it can serve as a general formulation that covers many other problems . Hanzely et al . ( 2021 ) demonstrate that various full model personalization formulations based on regularization ( Dinh et al. , 2020 ; Li et al. , 2021 ) , including ( 2 ) , as well as interpolation ( Deng et al. , 2020 ; Mansour et al. , 2020 ) are special cases of this problem . The rates of convergence we prove in §3 are competitive with or better than those in previous works for full model personalization methods in the non-convex case . 2 PARTIALLY PERSONALIZED MODELS . Modern deep learning models all have a multi-layer architecture . While a complete understanding of why they work so well is still out of reach , a general insight is that the lower layers ( close to the input ) are mostly responsible for feature extraction and the upper layers ( close to the output ) focus on complex pattern recognition . Depending on the application scenarios and domain knowledge , we may personalize either the input layer ( s ) or the output layer ( s ) of the model ; see Figure 1 . In Figure 1c , the input layers are split horizontally into two parts , one shared and the other personal . They process different chunks of the input vector and their outputs are concatenated before feeding Algorithm 1 Federated Learning with Partial Model Personalization ( FedSim / FedAlt ) Input : initial states u ( 0 ) , { v ( 0 ) i } ni=1 , number of rounds T , number of devices per round m 1 : for t = 0 , 1 , · · · , T − 1 do 2 : server randomly samples m devices as S ( t ) ⊂ { 1 , . . . , n } 3 : server broadcasts u ( t ) to each device in S ( t ) 4 : for each device i ∈ S ( t ) in parallel , do 5 : ( u ( t+1 ) i , v ( t+1 ) i ) = LocalSim / LocalAlt ( u ( t ) , v ( t ) i ) . v ( t+1 ) i = v ( t ) i if i /∈ S ( t ) 6 : send u ( t+1 ) i back to server 7 : server updates u ( t+1 ) = 1m ∑ i∈S ( t ) u ( t+1 ) i to the upper layers of the model . As demonstrated in ( Bui et al. , 2019 ) , this partitioning can help protect user-specific private features ( input 2 in Figure 1c ) as the corresponding feature embedding ( through vi ) are personalized and kept local at the device . Similar architectures have also been proposed in context-dependent language models ( e.g. , Mikolov & Zweig , 2012 ) . A more structured partitioning is illustrated in Figure 2a , where a typical transformer layer ( Vaswani et al. , 2017 ) is augmented with two adapters . This architecture is proposed by Houlsby et al . ( 2019 ) for finetuning large language models . Similar residual adapter modules are proposed by Rebuffi et al . ( 2017 ) for image classification models in the context of multi-task learning . In the context of FL , we treat the adapter parameters as personal and the rest of the model parameters as shared . Figure 2b shows a generalized additive model , where the outputs of two separate models , one shared and the other personalized , are fused to generate a prediction . Suppose the shared model is h ( u , · ) and the personal model is hi ( vi , · ) . For regression tasks with samples zi , j = ( xi , j , yi , j ) , where xi , j is the input and yi , j ∈ Rp is the output , we let Fi ( u , vi ) = ( 1/Ni ) ∑Ni j=1 fi ( ( u , vi ) , zi , j ) with fi ( ( u , vi ) , zi , j ) = ‖yi , j − h ( u , xi , j ) − hi ( vi , xi , j ) ‖2 . In this special case , the personal model fits the residual of the shared model and vice-versa ( Agarwal et al. , 2020 ) . For classification tasks , h ( u , · ) and hi ( vi , · ) produce probability distributions over multiple classes . We can use the cross-entropy loss between yi , j and a convex combination of the two model outputs : θh ( u , xi , j ) + ( 1− θ ) hi ( vi , xi , j ) , where θ ∈ ( 0 , 1 ) is a learnable parameter . Finally , we can cast the formulation ( 2 ) of full model personalization as a special case of ( 3 ) by letting u← w0 , vi ← wi , Fi ( u , vi ) ← Fi ( vi ) + ( λi/2 ) ‖vi − u‖2 . Many other formulations of full personalization can be reduced to ( 3 ) ; see Hanzely et al . ( 2021 ) .

Summary: To overcome statistical heterogeneity among data shards in collaborative federated learning, this paper proposes a novel personalization schema which only requires smaller memory footprint in clients and possibly is less susceptible to catastrophic forgetting. The main idea is to split the trainable parameters into  shared and personalized parameters where, unlike existing personalization schema, only the shared model is exchanged with the server in communication rounds to be aggregated- thus partial personalization. In regression, this corresponds to learning the residual error of shared model via personalized model and in the classification setting corresponds to output averaging (unlike interpolation based personalization methods that do parameter mixing). The authors propose two algorithms FedSim (simultaneous updating of shared and personalized models locally) and FedAlt (alternative updating of shared and personalized models locally), to learn in this setting, theoretically analyze the convergence rates in non-convex settings and conduct empirical studies on image classification and next-word prediction to evaluate the proposed methods.