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 aims to train a GHN model that can predict weights of each layer. The architecture is represented in adjacent matrix format, where a node represents one layer and an edge indicates whether two nodes are connected in given architecture. Each type of layer is represented with a one-hot vector as their node attribute. The outputs are predicted weights of the corresponding layer. Finally, a task-specific object function is minimized by updating the GHN.