Robust Federated Learning Through Representation Matching and Adaptive Hyper-parameters

1 INTRODUCTION . The size of the data used to train machine learning models is steadily increasing , and the privacy concerns associated with storing and managing this data are becoming more pressing . Offloading model training to the data owners is an attractive solution to address both scalability and privacy concerns . This gives rise to the Federated Learning ( FL ) setting ( McMahan et al. , 2016 ) where several clients collaboratively train a model without disclosing their data . In synchronous FL , training proceeds in rounds where at the beginning of each round , a central party sends the latest version of the model to the clients . The clients ( or a subset of them ) train the received model on their local datasets and then communicate the resulting models to the central party at the end of the round . The central party aggregates the client models ( typically by averaging them ) to obtain the new version of the model which it then communicates to the clients in the next round . The FL setting poses a unique set of challenges compared to standard stochastic gradient descent ( SGD ) learning on a monolithic dataset . In real-world settings , data from each client may be drawn from different distributions , and this heterogeneity can lead the local learning processes to diverge from each other , harming the convergence rate and final performance of the aggregate model . We address this challenge in two ways : ( 1 ) adaptive on-line tuning of hyper-parameters ; and ( 2 ) adding a regularization term that punishes divergent representation learning across clients . Traditional hyper-parameter tuning schemes , such as random search ( Bergstra & Bengio , 2012 ) or Bayesian methods ( Bergstra et al. , 2011 ; Snoek et al. , 2012 ) , require several training runs to evaluate the fitness of different hyper-parameters . This is impractical in the FL setting , which strives to minimize unnecessary communication . To address this issue , we formulate the hyper-parameter 1Code used to run all experiments is available under the anonymous repository https : //gitlab.com/ anon.iclr2020/robust_federated_learning selection problem as an online reinforcement learning ( RL ) problem : in each round , the learners perform an action by selecting particular hyper-parameter values , and at the end of the round get a reward which is the relative reduction in training loss . We then update the hyper-parameter selection policy online to maximize the rewards . Unlike centralized SGD , where there is a single trajectory of parameter updates , the FL setting has many local parameter trajectories , one per each active client . Even though they share the same initial point , these trajectories could significantly diverge . Averaging the endpoints of these divergent trajectories at the end of a training round would then result in a poor model ( Zhao et al. , 2018 ) . Clients with heterogeneous data distributions exacerbate this effect as each client could quickly learn representations specific to its local dataset . We introduce a scheme that mitigates the problem of divergent representations : each client uses the representations it learns to reconstruct the representations in the initial model it receives . A client is thus discouraged from learning representations that are too specific and that discard information about the globally learned representations . We show that this representation matching scheme significantly improves robustness and accuracy in the presence of heterogeneous client-side data distributions . We evaluate the performance of our two-part scheme , representation matching and online hyperparameter adjustments , on image classification tasks : MNIST , CIFAR10 ; and on a keyword spotting ( KWS ) task . We show that for homogeneous client-side data distributions , our scheme consistently improves accuracy . For heterogeneous clients , in addition to improving accuracy , our scheme improves training robustness and stops catastrophic training failures without having to manually tune hyper-parameters for each task . 2 RELATED WORK . The simplest automated hyper-parameter tuning methods execute random searches in hyperparameter space to find the best-performing hyper-parameters ( Bergstra & Bengio , 2012 ) . Random search is outperformed by Bayesian methods ( Bergstra et al. , 2011 ; Snoek et al. , 2012 ) that use the performance of previously selected hyper-parameters to inform the choice of new hyper-parameter to try . In these methods , the fitness of a hyper-parameter choice is the post-training validation accuracy . Running the training process to completion in order to evaluate the hyper-parameter fitness is computationally intensive . This motivates running partial training runs instead ( Klein et al. , 2016 ; Li et al. , 2016 ) . However , in FL settings with tight communication budgets , running several ( partial ) training runs to select the right hyper-parameters could still be impractical . Hyper-parameters can be optimized using gradient descent to minimize the final validation loss . To avoid storing the entire training trajectory in order to evaluate the hyper-gradients , Maclaurin et al . ( 2015 ) uses reversible learning dynamics to allow the entire training trajectory to be reconstructed from the final model . Alternatively , the hyper-gradients can be calculated in a forward manner , albeit at an increased memory overhead ( Franceschi et al. , 2017 ) . Using the hyper-gradient to optimize hyper-parameters is an iterative process that requires several standard training runs ( hyper-iterations ) , again making this approach impractical in a communication-constrained FL setting . Hyper-parameters can be optimized using SGD in tandem with model parameters through the use of hyper-networks ( Ha et al. , 2016 ; Lorraine & Duvenaud , 2018 ) which map hyper-parameters to optimal model parameters . It is unclear how hyper-networks can be learned in a FL setting as we are interested in choosing hyper-parameters that result in the best post-aggregation model and not the hyper-parameters chosen by any particular client to optimize its own model on its own dataset . The hyper-parameter tuning approaches most related to ours are those based on RL methods . Daniel et al . ( 2016 ) ; Jomaa et al . ( 2019 ) learn policies and/or state-action values that are then used to select good hyper-parameters . However , several training runs are needed to learn the policies and the action values . Perhaps the closest approach to ours is the method in Xu et al . ( 2017 ) which trains an actor-critic network in tandem with the main model and uses it to select actions ( hyper-parameter choices ) that maximize the expected reduction in training loss . It is unclear how this approach can be applied in a FL setting as the training process on each client has no access to the loss of interest which is the loss of the post-aggregation model . FL typically struggles when the client-side data distributions are significantly different ( McMahan et al. , 2016 ; Zhao et al. , 2018 ) as this causes the parameter trajectories to significantly diverge across different clients . This severely degrades the performance of the averaged model ( Zhao et al. , 2018 ) . One straightforward solution is to reduce the SGD steps each client takes between the synchronization ( model averaging ) points . More frequent model synchronization , however , increases communication volume . Various compression methods ( Sattler et al. , 2019 ; Konečnỳ et al. , 2016 ) are able to compress the model updates , which makes it possible to synchronize more frequently and mitigate the effect of heterogeneous data distributions . Sattler et al . ( 2019 ) takes the extreme case of synchronizing after every SGD step , which reduces the FL setting to standard SGD . An alternative solution is to mix the datasets of the different clients to obtain more homogeneous client-side data distributions ( Zhao et al. , 2018 ) . This , however , compromises the privacy of the clients ’ data . 3 METHODS . We build upon the federated averaging ( FedAvg ) algorithm ( McMahan et al. , 2016 ) and introduce two novel aspects to it : global adaptive hyper-parameters and local ( per-client ) representation matching . We first give an informal description of the complete algorithm . We consider a synchronous FedAvg setting with K clients , where Dk is the dataset local to client k. Training proceeds in rounds . Figure 1 illustrates the steps involved in training round t. The server maintains a distribution P ( H|ψ ) over the space of hyper-parameters H. At the beginning of round t ( Fig . 1a ) , the selected clients receive the most recent model parameters from the server , wt , together with training hyper-parameters ht sampled from P ( H|ψt ) . Client i uses two copies of wt to initialize two models : a fixed modelMF and a trainable modelMit . The parameters wit ofMit are trained using SGD ( Fig . 1b ) . Client i maintains a set of local parameters θit that it uses to map the activations in Mit to the activations inMF . θit and wit are simultaneously trained to minimize a two-component loss : the first component is the standard training loss ofMit on Di ( for example , the cross-entropy loss ) ; the second component is the mean squared difference between the activations inMF and the activations reconstructed fromMit using θit . We denote this second component as the representation matching loss . In the final step in the round ( Fig . 1c ) , each participating client sends its final model parameters , wit , to the central server . The server aggregates the model parameters , and uses the loss of the resulting model to update the parameters ψ of the hyper-parameter distribution , P ( H|ψ ) . 3.1 REPRESENTATION MATCHING . We now describe the representation matching scheme . For an example model , Fig . 2 illustrates how the local parameters θi in client i are used to map the activations of the model being trained , Mi , to the activations of the fixed model MF . Mi is parameterized by wi while MF is parameterized by w which was received from the server at the beginning of training . To simplify notation , we dropped the index of the round , t. Given a data point x ∼ Di , x is fed toMi andMF to obtain the layer activations [ a1 ( x ; wi ) , . . .aM ( x ; wi ) ] and [ a1 ( x ; w ) , . . . , aM ( x ; w ) ] , respectively . These are not all the model activations , but only the M activations we are interested in matching . The activations in the two models are matched by a set of matching layers parameterized by θi ( θi ≡ [ θi1 , θ i 2 , θ i 3 ] in Fig . 2 ) . Instead of deriving the activations of one layer inMF from the activations of the corresponding layer in Mi , we found performance significantly improves if the activations of one layer are instead derived from the activations of the next layer of interest above it . As shown in the representative example in Fig . 2 , the activations/layers of interest are the input layer , the point-wise non-linearity layers and the top layer . We use convolutional matching layers if the input has spatial information ( such as within a convolutional stack ) , and fully-connected matching layers otherwise . To reverse a pooling operation , we follow the scheme in Zhao et al . ( 2015 ) and unpool pooled activations using the pooling positions used during the pooling operation as shown in Fig . 2 . Given a training point and label ( x , y ) ∼ Di , we define the matching loss at client i as : LiM ( x ; wi , w , θi ) = j=M−1∑ j=1 ||fj ( aj+1 ( x ; wi ) ; θij ) − aj ( x ; w ) ||22 ( 1 ) where fj denotes the jth matching layer . The per-point training loss at client i is : Li ( x , y ; wi , w , θi ) = LC ( aM ( x ; wi ) ; y ) + LiM ( x ; wi , w , θi ) ( 2 ) where LC is the cross-entropy loss and aM ( x ; wi ) is the top layer activation in the trainable model . Our representation matching architecture is similar in some respects to auto-encoding architectures , particularly ladder networks ( Rasmus et al. , 2015 ) . Unlike ladder networks , however , we do not reconstruct the activations of a non-noisy model from a noisy model but rather , reconstruct the activations of the most recent aggregate model from the activations of the model while it is being trained on a local dataset . This has a regularizing influence as it stops the clients from learning representations that are too specific to their local datasets . This should also ameliorate the effect of model divergence in the clients as all client models have to learn representations that are mappable to a common set of representations defined by the common aggregate model . Additional regularization comes about because a layer ’ s activations are used to reconstruct the activations of a layer below it . A similar technique was used in what-where autoencoders ( Zhao et al. , 2015 ) , and was shown to significantly improve accuracy . Our approach is fundamentally distinct , however , as we use the activations in one one model , Mi , to reconstruct the activations in a different model , MF . Our representation matching scheme is similar in spirit to the techniques used to learn invariant feature representations for domain adaptation where the goal is to minimize the discrepancy between the distributions of features extracted from different domains . However , current techniques for learning domain-invariant features ( Ganin et al. , 2016 ; Shen et al. , 2017 ) are incompatible with the FL setting as they require the learner to have access to data from different domains ( clients in our case ) in order to match the feature distributions across these domains .

Summary: This manuscript proposes two strategies to improve both the robustness and accuracy of local agents under the setting of federated learning. Specifically, online reinforcement learning is used to perform adaptive hyperparameter search in order to maximize the utility of local models. This is quite an interesting idea since traditional hyperparameter tuning techniques, including random search or Bayesian optimization, usually requires access to monolithic dataset, which is clearly impractical under federated learning. The second contribution is an idea on using local distribution matching in order to synchronize the learning trajectories of different local models. This again is a novel and interesting idea. Overall the paper is well-written and clear to follow, which is a plus. Detailed comments and questions follow: