Federated Learning with Heterogeneous Architectures using Graph HyperNetworks

1 INTRODUCTION . Federated learning ( FL ) ( McMahan et al. , 2017a ; Yang et al. , 2019 ; Konečný et al. , 2015 ; 2017 ) allows multiple clients to collaboratively train a strong model that benefits from their individual data , without having to share that data . In particular , aggregating the parameters of locally trained models alleviates the need to share raw data , thereby preserving privacy to some extent and reducing the volume of data transferred . FL allows safe and efficient learning from edge nodes like smartphones , self-driving cars , and medical systems . While successful , FL approaches have one key limitation : all clients need to share the same network architecture . As a result , they are not applicable to many important cases that require learning with heterogeneous architectures . For example , clients often run different networks on their edge devices , due to computational limitations or OS versions . Also , in some cases , clients want to keep their network architectures private , for instance , when different organizations wish to benefit from each other ’ s access to data without having to share their proprietary architectures . In all these cases , one is interested in federated learning with heterogeneous architectures ( FLHA ) . Unfortunately , current FL approaches do not support mixtures of different architectures , but require that clients share the same model architecture . As an illustrative example , consider Federated Averaging ( FedAvg ) , perhaps the most widely used FL technique , where model parameters of different clients are averaged on a shared server McMahan et al . ( 2017a ) . Unless all clients share the same number of parameters and layer structure , it is not clear how to average their weights . The problem of aggregating weights across different architecture is not specific to FedAvg , but also exists with other FL approaches . This limitation raises the fundamental research question : Can federated learning handle heterogeneous architectures ? And can it be achieved with clients keeping their architectures undisclosed ? A possible alternative to FLHA could be achieved through knowledge distillation . For example , Lin et al . ( 2020 ) suggested distilling knowledge from ensembles of client models on a server using unlabeled or synthetic datasets . Although this solution could be beneficial in certain setups , it has two major limitations : ( 1 ) Client architectures must be disclosed and ( 2 ) external data must be provided . We could also design a distillation-based approach to FLHA that does not require clients to share their architectures . First , train a shared model using a standard FL approach . Once trained , distill knowledge from the shared model to train private client models using each client local data and architecture . While this solution can address the two limitations discussed above , it also has two drawbacks : ( 1 ) Training on a large shared architecture may not be feasible for clients with limited computing power ; ( 2 ) Local distillation using only local data may overfit and lose the benefits of FL . Finally , and perhaps most importantly , this approach bypasses , but does not addresses , the fundamental question : How can we learn to use knowledge about architectures when transferring information about model parameters across architectures . In this work , we propose a general approach to FLHA based on a hypernetwork ( HN ) ( Ha et al. , 2016 ) that acts as the knowledge aggregator . The hypernetwork takes as input a description of a client architecture and predicts weights to populate that architecture . Unlike FedAvg , which aggregates weights in a predefined way , the hypernetwork implicitly learns an aggregation operator from data and can therefore be significantly more powerful . For providing an architecture descriptor to the HN , we propose to model client architectures as graphs where nodes represent both parametric and non-parametric layers ( e.g. , a convolutional layer or a summation operation ) and directional edges represent network connectivity such as skip connections . To allow the hypernetwork to process any graph , regardless of its size or topology , we use a graph hypernetwork ( GHN ) similar to Zhang et al . ( 2020 ) . We therefore name our approach FLHA-GHN . The GHN operates on a graph representation of the architecture and predicts layer parameters at each node that represents a parametric layer . At each communication round , clients train local copies of the hypernetwork weights using their own architectures and data ( Fig . 1 ) . The updated weights are sent by the clients to the server where they are aggregated . Since different architectures have different layer compositions , representing layers as nodes allows meaningful knowledge aggregation across architectures . This forms an improved hypernetwork model that uses knowledge gleaned from different network types and datasets , to populate them with improved parameters . Critically , client architectures are not communicated . The ability to generalize across different client architectures is crucial for successful FLHA . Recent results in theory of Graph Neural Networks ( GNNs ) Yehudai et al . ( 2021 ) suggest that GNNs can indeed generalize across input distributions . We discuss this issue and demonstrate experimentally that our approach generalizes to unseen architectures in certain cases – allowing clients to modify their architectures after federation has occurred , without the need for client-wide FL retraining . In an extensive experimental study on three image datasets , we show that applying FLHA-GHN to FL with heterogeneous-architecture clients outperforms a distillation-based baseline , and a non-GNN based HN architecture ( Shamsian et al. , 2021 ) , by a large margin , and that margin generally grows when the amount of local data decreases . We further show that FLHA-GHN provides a large benefit in generalization to unseen architectures , improving the accuracy of converged models and drastically shortening the time needed for convergence . Such generalization could allow new clients , with new architectures , to benefit from models trained on different data and different architectures , lowering the bar for deploying new , personalized , architectures . 2 PREVIOUS WORK . Federated Learning . Federated learning McMahan et al . ( 2017a ) ; Kairouz et al . ( 2019 ) ; Yang et al . ( 2019 ) ; Mothukuri et al . ( 2021 ) is a learning setup in which multiple clients collaboratively train individual models while trying to benefit from the data of all the clients without sharing their data . The most well known FL technique is federated averaging , where all clients use the same architecture , which is trained privately and then sent to the server and averaged with other locally trained models . Many recent works have focused on improving privacy McMahan et al . ( 2017b ) ; Agarwal et al . ( 2018 ) ; Li et al . ( 2019 ) ; Zhu et al . ( 2020 ) and communication efficiency Chen et al . ( 2021 ) ; Agarwal et al . ( 2018 ) ; Chen et al . ( 2021 ) ; Dai et al . ( 2019 ) ; Stich ( 2018 ) . Another widely studied setup is the heterogeneous client data setup Hanzely & Richtárik ( 2020 ) ; Zhao et al . ( 2018 ) ; Sahu et al . ( 2018 ) ; Karimireddy et al . ( 2020 ) ; Zhang et al . ( 2021 ) . To solve this problem , personalized FL ( pFL ) methods were proposed that adapt global models to specific clients Kulkarni et al . ( 2020 ) . Most related to the current work is the recent work of Shamsian et al . ( 2021 ) that also used hyper-networks . Their method considers a simple case of clients with varying architectures , where three pre-defined simple architectures are “ hard '' encoded into the structure of the hypernetwork . In contrast , our framework allows clients to use a variety of layers and computational graphs , and facilitates better weight sharing as the same layers are used in different architectures . In the experimental section 5 we show that our approach outperforms Shamsian et al . ( 2021 ) by large margins . Hypernetworks . HyperNetworks ( HN ) Klein et al . ( 2015 ) ; Ha et al . ( 2016 ) are neural networks that predict input conditioned weights for another neural network that performs the task of interest . HNs are widely used in many learning tasks such as generation of 3D content Littwin & Wolf ( 2019 ) ; Sitzmann et al . ( 2019 ) , neural architecture search Brock et al . ( 2017 ) and language modelling Suarez ( 2017 ) . More relevant to our work are Graph Hypernetworks ( GHNs ) - hypernetworks that take graphs as an input . Nachmani & Wolf ( 2020 ) used GHNs for molecule property prediction . Even more relevant is the use of GHNs for Neural Architecture Search ( NAS ) Zhang et al . ( 2020 ) . In our work , we adapt GHNs to FLHA with unique challenges arising from the problem setup . 3 PROBLEM DEFINITION . Traditional FL addresses the setup of C clients working together to improve their local models , while each client c ∈ { 1 , . . . , C } has access only to its local data samples { ( xcj , ycj ) } ncj=1 , sampled from client specific data distributions Pc , c = 1 , . . . , C. In this paper , we generalize FL to FL with Heterogeneous Architectures ( FLHA ) . In this setup , clients can use different network architectures f1 ( · ; θ1 ) , f2 ( · ; θ2 ) , . . . , fC ( · ; θC ) . Here , fc ( · ; θc ) is a neural architecture from some predefined family of models with parameters θc ∈ Rmc ( see further discussion on architecture families in Section 4 ) . Moreover , we assume that all clients are connected to a server , and that the clients can share information among themselves only through the server . Importantly , in FLHA we are interested in the setup where both the architectures and the data can not be transferred due to privacy concerns . Our goal is to solve the following minimization problem : ( θ∗1 , ... , θ ∗ C ) = argmin ( θ1 , ... , θC ) C∑ c=1 E ( x , y ) ∼Pc ` ( fc ( x ; θc ) , y ) , ( 1 ) for a suitable loss function ` ( y , y′ ) . 4 APPROACH . 4.1 OVERVIEW . Standard FL methods rely on aggregating model parameters and are not directly applicable when clients use different architectures . In our FLHA setup , the parameter vectors θc of different clients have different shapes and sizes , and as a result , a direct aggregation can be meaningless or not well defined . To address this issue , we re-parameterize the weights θc of a model c as the output of a hypernetwork . The hypernetwork serves as a trainable knowledge aggregator . The parameters of this operator are learned from data using updates from all clients . Since the output of the hypernetwork must fit a given architecture fc , the hypernetwork should take as input a representation of fc . Here , we advocate graphs as a natural representation for neural architectures ( in agreement with Zhang et al . ( 2020 ) ) and apply graph hypernetworks to process them . The workflow of training our FLHA-GHN is illustrated in Figure 1 . At each communication round , several steps are followed . ( 1 ) First , a server shares the current weights of the GHN , W , with all clients ; ( 2 ) Each client c uses the GHN to predict weights θc for its own specific architecture fc , and updates the GHN weights Wc locally using its own architecture and data ; ( 3 ) Each client sends the locally updated GHN weights to the server ; ( 4 ) Weight averaging is performed on the server .

Summary: This paper proposes a method to tackle the problem of architecture-agnostic federated learning. The primary motivation of this paper is that in some cases of federated learning, different participants may need to use different neural architectures and may not want the exact architecture to be known. To this end, the authors propose a method based on graph hypernetworks. The main insight of this method is that neural architectures can be represented as graphs and their weights can be outputs of a larger GNN. Empirically, the proposed method compares favorably against state-of-the-arts.