Improving Federated Learning Personalization via Model Agnostic Meta Learning

1 INTRODUCTION . In recent years , the growth of machine learning applications was driven by aggregation of large amounts of data in a datacenter , where a model can be trained using large scale distributed system ( Dean et al. , 2012 ; LeCun et al. , 2015 ) . Both the research community and general public are becoming increasingly aware that there is a variety of scenarios where this kind of data collection comes with significant risks , mainly related to notions of privacy and trust . In the presence of user generated data , such as activity on mobile phones , Federated Learning ( FL ) ( McMahan & Ramage , 2017 ) proposes an alternative approach for training a high quality global model without ever sending raw data to the cloud . The FL system proposed by Google ( Bonawitz et al. , 2019 ) selects a sample of available devices and sends them a model to be trained . The devices compute an update to the model based on an optimization procedure with locally available data , and the central system aggregates the updates from different devices . Such iteration is repeated many times until the model has converged . The users ’ training data does not leave their devices . The basic FL algorithm , Federated Averaging ( FedAvg ) ( McMahan et al. , 2017 ) , has been used in production applications , for instance for next word prediction in mobile keyboard ( Hard et al. , 2018 ) , which shows that Federated Learning can outperform the best model trained in a datacenter . Successful algorithmic extensions to the central idea include training a differential private model ( McMahan et al. , 2018 ) , compression ( Konečný et al. , 2016b ; Caldas et al. , 2018a ) , secure aggregation ( Bonawitz et al. , 2017 ) , and a smaller number of always-participating nodes ( Yang et al. , 2019 ) . FL applications generally face non-i.i.d and unbalanced data available to devices , which makes it challenging to ensure good performance across different devices with a FL-trained global model . Theoretical guarantees are only available under restrictive assumptions and for convex objectives , cf . Li et al . ( 2019b ) . In this work , we are interested in personalization methods that adapt the model for data available on each device , individually . We refer to a trained global model as the initial model , and the locally adapted model as the personalized model . Existing FL personalization work directly takes a converged initial model and conducts personalization evaluation via gradient descent ( Beaufays et al. , 2019 ) . However , in this approach , the training and personalization procedures are completely disconnected , which results in potentially suboptimal personalized models . Meta Learning optimizes the performance after adaptation given few-shot adaptation examples on heterogeneous tasks , and has increasing applications in the context of Supervised Learning and Reinforcement Learning . Model Agnostic Meta Learning ( MAML ) introduced by Finn et al . ( 2017 ) is a solely gradient-based Meta Learning algorithm , which runs in two connected stages ; metatraining and meta-testing . Meta-training learns a sensitive initial model which can conduct fast adaptation on a range of tasks , and meta-testing adapts the initial model for a particular task . Both tasks for MAML , and clients for FL , are heterogeneous . For each task in MAML and client in FL , existing algorithms use a variant of gradient descent locally , and send an overall update to a coordinator to update the global model . If we present the FL training process as meta-training in the MAML language , and the FL personalization via gradient descent as meta-testing , we show in Section 2 that FedAvg ( McMahan et al. , 2017 ) and Reptile ( Nichol et al. , 2018 ) , two popular FL and MAML algorithms , are very similar to each other ; see also Khodak et al . ( 2019 ) . In order to make FL personalization useful in practice , we propose that the following objectives must all be addressed , simultaneously . ( 1 ) Improved Personalized Model – for a large majority of the clients . ( 2 ) Solid Initial Model – some clients have limited or even no data for personalization . ( 3 ) Fast Convergence – reach a high quality model in small number of training rounds . Typically , the MAML algorithms only focus on objective ( 1 ) ; that was the original motivation in Finn et al . ( 2017 ) . Existing FL works usually focus on objectives ( 2 ) and ( 3 ) , and take the personalized performance as secondary . This is largely due to the fact that it was not obvious that getting a solid initial model is feasible or practical if devices are available occasionally and with limited resources . In this work , we study these three objectives jointly , and our main contributions are : • We point out the connection between two widely used FL and MAML algorithms , and interpret existing FL algorithm in the light of existing MAML algorithms . • We propose a novel modification of FedAvg , with two stages of training and fine-tuning , for optimizing the three above objectives . • We empirically demonstrate that FedAvg is already a meta learning algorithm , optimizing for personalized performance , as opposed to quality of the global model . Furthermore , we show that the fine tuning stage enables better and more stable personalized performance . • We observe that different global models with the same accuracy , can exhibit very different capacity for personalization . • We highlight that these results challenge the existing objectives in the FL literature , and motivate new problems for the broader Machine Learning research community . 2 INTERPRETING FEDAVG AS A META LEARNING ALGORITHM . In this section , we highlight the similarities between the FL and MAML algorithms and interpret FedAvg as a linear combination of a naive baseline and a collection of existing MAML methods . Algorithm 1 presents a conceptual algorithm with nested structure ( left column ) , of which the MAML meta-training algorithm , Reptile ( middle column ) , and FL-training algorithm , FedAvg ( right column ) , are particular instances . We assume thatL is a loss function common to all of the following arguments . In each iteration , a MAML algorithm trains across a random batch of tasks { Ti } . For each task Ti , it conducts an inner-loop update , and aggregates gradients from each sampled task with an outer-loop update . In each training round , FL uses a random selection of clients { Ti } . For each client Ti and its weight wi , it runs an optimization procedure for a number of epochs over the local data , and sends the update to the server , which aggregates the updates to form a new global model . If we simplify the setting and assume all clients have the same amount of data , causing the weights wi to be identical , Reptile and FedAvg in fact become the same algorithms . Several other MAML algorithms ( Finn et al. , 2017 ; Antoniou et al. , 2018 ) , or other non-MAML/FL methods ( Zhang et al. , 2019 ) , can also be viewed as instances of the conceptual method in left column of Algorithm 1 . Algorithm 1 Connects FL and MAML ( left ) , Reptile Batch Version ( middle ) , and FedAvg ( right ) . OuterLoop/Server learning rate α InnerLoop/Client learning rate β Initial model parameters θ while not done do Sample batch of tasks/clients { Ti } for Sampled task/client Ti do if FL then gi , wi = ClientUpdate ( θ , Ti , β ) else if MAML then gi = InnerLoop ( θ , Ti , β ) end if end for if FL then θ = ServerUpdate ( θ , { gi , wi } , α ) else if MAML then θ = OuterLoop ( θ , { gi } , α ) end if end while Require : : Reptile StepK . function InnerLoop ( θ , Ti , β ) SampleK-shot dataDi , k from Ti . θi = θ for each local step i from 1 to K do θi = θi − β∇θL ( θi , Di , k ) end for Return gi = θi − θ end function Require : : Meta Batch SizeM . functionOuterLoop ( θ , { gi } , α ) θ = θ + α 1M ∑M i=1 gi Return θ end function Require : FedAvg Local EpochE . function ClientUpdate ( θ , Ti , β ) Split local dataset into batchesB θi = θ for each local epoch i from 1 to E do for batch b ∈ B do θi = θi − β∇θL ( θi , b ) end for end for Return gi = θi − θ end function Require : Clients per training roundM . functionServerUpdate ( θ , { gi , wi } , α ) θ = θ+α ∑M i=1 wigi/ ∑M i=1 wi Return θ end function In the following , we rearrange the summands comprising the update formula of FedAvg/Reptile algorithm to reveal the connection with other existing methods – a linear combination of the Federated SGD ( FedSGD ) ( McMahan et al. , 2017 ) and First Order MAML ( FOMAML ) algorithms ( Finn et al. , 2017 ) with different number of steps . For clarity , we assume identical weights wi in FedAvg . Consider T participating clients and let θ be parameters of the relevant model . For each client i , define its local loss function as Li ( θ ) , and let gij be the gradient computed in j th iteration during a local gradient-based optimization process . FedSGD was proposed as a naive baseline against which to compare FL algorithms . For each client , it simply takes a single gradient step based on the local data , which is sent back to the server . It is a sensible baseline because it is a variant of what a traditional optimization method would do if we were to collect all the data in a central location , albeit inefficient in the FL setting . That means that FedSGD optimizes the performance of the initial model , as is the usual objective in datacenter training . The local update produced by FedSGD , gFedSGD , can be written as gFedSGD = −β T T∑ i=1 ∂Li ( θ ) ∂θ = 1 T T∑ i=1 gi1 . ( 1 ) Next , we derive the update of FOMAML in similar terms . Assuming client learning rate β , the personalized model of client i , obtained after K-step gradient update is θiK = U i K ( θ ) = θ − β ∑K j=1 g i j = θ − β ∑K j=1 ∂Li ( θj ) ∂θ . Differentiating the client update formula , we get ∂U iK ( θ ) ∂θ = I − β ∂ ∑K j=1 g i j ∂θ = I − β K∑ j=1 ∂2Li ( θj ) ∂θ2 . ( 2 ) Directly optimizing the current model for the personalized performance after locally adapting K gradient steps , results in the general MAML update proposed by Finn et al . ( 2017 ) . gMAML = ∂LMAML ∂θ = 1 T T∑ i=1 ∂Li ( U i K ( θ ) ) ∂θ = 1 T T∑ i=1 L′i ( U i K ( θ ) ) ( I − β K∑ j=1 ∂2Li ( θj ) ∂θ2 ) . ( 3 ) MAML requires to compute 2nd-order gradients , which can be computationally expensive and creates potentially infeasible memory requirements . To avoid computing the 2nd-order term , FOMAML simply ignores it , resulting in a first-order approximation of the objective ( Finn et al. , 2017 ) . FOMAML ( K ) then uses the ( K + 1 ) th gradient as the local update , after K gradient steps . gFOMAML ( K ) = 1 T T∑ i=1 L′i ( U i K ( θ ) ) I = 1 T T∑ i=1 L′i ( θ i K ) = 1 T T∑ i=1 giK+1 . ( 4 ) Now we have derived the building blocks of the FedAvg . As presented in Algorithm 1 , the update of FedAvg , gFedAvg , is the average of client updates , which are the sums of local gradient updates . Rearranging the terms presents its interpretation as a linear combination of the above ideas . gFedAvg = 1 T T∑ i=1 K∑ j=1 gij = 1 T T∑ i=1 gi1 + K−1∑ j=1 1 T T∑ i=1 gij+1 = gFedSGD + K−1∑ j=1 gFOMAML ( j ) ( 5 ) Note that interpolating to special case , gFedSGD can be seen as gFOMAML ( 0 ) – optimizing the performance after 0 local updates , i.e. , the current model . This sheds light onto the existing Federated Averaging algorithm , as the linear combination of algorithms optimizing personalized performance after a range of local updates . Note , however , this does not mean that FedAvg optimizes for the linear combination of the objectives of the respective algorithms . Nevertheless , we show in the following section that using K = 1 results in a model hard to personalize , and increasing K significantly improves the personalization performance , up until a certain point where the performance of initial model becomes unstable .

Summary: This paper considers personalization federated learning problem in which the goal is to personalize the global model on a given client/device based on available data on that device/client.  The paper claims not only their proposed method can lead to fast convergence time but also provide a solid initial model per device/client and results in a better-personalized model. To evaluate the performance of their method, EMNIST-62 and Shakespeare data are used.