Federated Learning with Heterogeneous Architectures using Graph HyperNetworks
Authors: Or Litany
Haggai Maron
David Acuna
Jan Kautz
Gal Chechik
Sanja Fidler
Venue: ICLR 2022 Submitted (2022)
Keywords: federated learning
graph neural networks
hypernetworks

Abstract: Standard Federated Learning (FL) techniques are limited to clients with identical network architectures. As a result, inter-organizational collaboration is severely restricted when both data privacy and architectural proprietary are required. In this work, we propose a new FL framework that removes this limitation by adopting a graph hypernetwork as a shared knowledge aggregator. A property of the graph hyper network is that it can adapt to various computational graphs, thereby allowing meaningful parameter sharing across models. Unlike existing solutions, our framework makes no use of external data and does not require clients to disclose their model architecture. Compared with distillation-based and non-graph hypernetwork baselines, our method performs notably better on standard benchmarks. We additionally show encouraging generalization performance to unseen architectures.

Under review as a conference paper at ICLR 2022
FEDERATED LEARNING WITH HETEROGENEOUS AR- CHITECTURES USING GRAPH HYPERNETWORKS
Anonymous authors Paper under double-bling review
ABSTRACT
Standard Federated Learning (FL) techniques are limited to clients with identical network architectures. As result, inter-organizational collaboration is severely restricted when both data privacy and architectural proprietary are required. In this work, we propose new FL framework that removes this limitation by adopting a graph hypernetwork as shared knowledge aggregator. A property of the graph hyper network that it can adapt to various computational graphs, thereby allow- ing meaningful parameter sharing across models. Unlike existing solutions, our framework makes no use of external data and does not require clients to disclose their model architecture. Compared with distillation-based and non-graph hypernet- work baselines, our method performs notably better on standard benchmarks. We additionally show encouraging generalization performance to unseen architectures.
INTRODUCTION
Federated learning (FI.) (McMahan et al. 2017a| Yang et al. 2019 Konečný al. 2015 2017 allows multiple clients to collaboratively train a strong model that benefits from their individual data, without having 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 share the same network architecture. As 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, 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?
not clear how average their weights.
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 distillation-basel 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 a large shared architecture may not be feasible for clients with

---

Under review as a conference paper at ICLR 2022
Figure 1: An overview of our approach. We tackle a feder- ated learning setup with clients that have different architectures by using shared Graph Hyper- network (GHN) that can adapt any network architecture. At each communication round, lo- cal copy of the GHN is trained by each client using their own data and network architecture (il- lustrated as graphs within each client; nodes represent layers and edges represent computa- tional flow), before returning to the server for aggregation.
FL with heterogeneous architettures using GHNs
Local
update
Client 1
GHN Parmeter Averaging
GHN Parmeters
W
W = 2
Local
W
update
Client
GHN
Local
update
Client C
Updated GHN Parmeters W
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 learn to knowledge about architectures when transferring information about model parameters across architectures.
fundamental question: How
In this work, propose general approach to FLHAA based a hypernetwork (HN) (Ha 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 hypernetwrk 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-parametri layers (e.g., a convolutional layer or 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. (0200). 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 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 sent by the clients the server where they aggregated. Since different architectures have different layer compositions, representing layers as nodes allows meaningful knowledge aggregation across architectures. This forms an improved hypernetoor 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 al. (02T 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 clients outperforms distillation-basel baseline, and non-GNN
FL with 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

---

Under review as a conference paper at ICLR 2022
2 PREVIOUS WORK
Federated Learning. Federated learning McMahan et al. 2017a Kairouz et al. (2019); Yang et al. (2019); Mothukuri et al. (202I) 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. (201bb; Agarwal al. (2018): ; Li et al. (2019); /hu et al. (2020) and communication efficiency Chen et al. (2021 Agarwal ct al. (0018) Chen et al. (2021) )) Dai et al. (011)); Stich 2018). Another widely studied setup is the
heterogeneous client data setup Hanzely & Richtárik (2000); Zhao al. (2018) Sahu et al. (2018);
Karimireddy et al. 2200); Zhang et al. (2021). To solve this problem. personalized FL (pFL) methods 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. ) that also used hyper-networks. Their method considers 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 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 approach outperforms Shamssin et (021) by large margins.
we show that our
Hypernetworks. HyperNetworks (HN) Klein et al. (2015) Ha et al. (2016) 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 generation of 3D content
content[Littiin 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 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 e E (1,...,C) C) has access only to local data samples {(Jrej Yes )es)) sampled from client specific data distributions Pc,
I,..., C.
In this paper, we generalize FL to FL with Heterogeneous Architectures (FLHA). In this setup, clients can use different network architectures fi(:0 f2(:; 02), ... fo(; 0c). Here, fel: 0.) is neural architecture from some predefined family of models with parameters 0g E RT (see further discussion architecture families in Section 4 Moreover, assume that all clients to a server, and that the clients can share information among themselves only through the server Importantly, in FLHAA we interested in the setup where both the architectures and the data cannot be transferred due to privacy concerns.
connected
Our goal is to solve the following minimization problem:
(0 ..., 0%) argmin ...... 8c) Σ ((
for a suitable loss function ely. v/).
4 APPROACH
4.1 OVERVIEW
Standard FL methods rely aggregating model parameters and are not directly applicable when clients use different architectures. In our FLHA setup, the parameter vectors 9, of different clients have different shapes and sizes, and result, direct aggregation defined. To address this issue, we re- parameterize the weights 8, of a model c as the output of a hypernetwork. The hypernetoork serves as trainable knowledge aggregator. The parameters of this
be meaningless not well
3

---

Under review as a conference paper at ICLR 2022
Graph representation of architecture
GNN output:
MUP output: Layer weights
Node (layer) features
SDEEP
MBP
MLP
Graph neural Network G
ML?
SUM
MLP
Nen-poramenti
FC
Figure 2: The GHN architecture. GHN is composed two sub-networks: a Graph neural network (GNN) G and set of MLPs LeL, for each layer type. We input graph representation of an architecture to the GHN, with the colored node features corresponding to different layer types. We then apply the GNN to produce node features. Finally, the node features for the parametric layers further processed by an MLP H1 to attain the final layer weights.
operator are learned from data using updates from all clients. Since the output of the hypernetwork must fit given architecture fer the hypemnetorr should take imput 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 process them.
The workflow of training our FLHA-GHN is illustrated in Figure [ several steps are followed. (1) First, a server shares the current weights of the GHN, W, with all clients: (2) Each client e uses the GHN to predict weights 0g for its own specific architecture fe, and updates the GHN weights We 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.
At each communication round,
4.2 METHOD
Representing neural architectures as graphs. A neural architecture can be represented in many different ways. Early struuustrectured architectures such as VGG Simonyyan & Zisserman (2014) and AlexNet Krizhevsky et al. (2012) can easily be represented as sequences of layers. Recently, several architectures that have richer connectivity structures have been proposed Ronneberger al. 2015); Targ et al. (2016): Huang et al. (2017). Sequential representation is not suitable for these architectures. In contrast, all neural architectures are computational graphs, which makes graph representation much more natural choice. We use graph representations to ensure generality in architecture space.
Given neural architecture, represent it graph AA (V, E') in the following way. The set of vertices V contains a vertex u for each parametric layer in the architecture. For example, a convolutional layer with weights of dimensions 3 3 x 64 128. similar way, use another type of nodes to represent non-parametric operation in the network, such as summation or concatenation of matrices, which are frequently used ResNet-like architectures. See example in Figure The set of edges E represents the computational flow of the architecture: there is an edge between the nodes and u (i.e., e = (u, u) E E) if the output of the layer represented by is the input of layer represented by u. X I x is a matrix that holds the input node features. Initially, each node is equipped with categorical (one-hot) features indicating the layer type they represent, denoted by = ( J.). Each categorical layer type specifies the following: linear/conv, stride, kernel size, activation and feature dimensions. As example, the convolution layer discussed above may be represented as /1. According to the desired architecture family, the nodes in the graph can represent different computational blocks of different granularity. Granularity can range from single layer to complex blocks, so nodes can even represent complex mechanisms like attention. Throughout this paper, we consider the ResNets architecture family and model layers as nodes.
4

---

Under review as a conference paper at ICLR 2022
Architecture. Originally introduced for neural architecture search, GHN Zhang et al. (2020) is generalization of standard hypernetorks that allows generating weights for heterogeneous network architectures. The GHN weight prediction be broken down into two stages: (1) Processing graph representation an architecture using a GNN and then (2) Predicting weights for each layer in that architecture. We now explain these stages in detail.
In the first stage, our hypernetwork takes as input a graph representation A of an architecture and processes using T-layer graph neural network G(;; Wa) with learnable parameters Wa (we use T 6 in our implementation). This process outputs latent representations h for each node {u E V/ We use maximally expressive message passing GNN layers, as introduced in Morris al. 20199) These layers have the following form:
A10)5-- B(t)
1
Σ (ul(u,))EE)
A + b(0) )
Here, & {1...., T represents the depth of the layer, A B(+), ((+) are layer-specific parameters. and a is a non-linear activation function such as ReLU. As mentioned above, we denote the concate- nation of the parameters A (()), ((0) 1 {1, ..., T} Wa. As shown in several recent works, the node features extracted by such GNNs are a representation of the local neighbourhood around each node (Xu et al., 2018: Morris et al., 2019 Yehudai et al.) 2021).
At the second stage, we use a set of MLPs { (Hi(: Wh) with learnable parameters W le L map latent node representations ! layer weights 0%:
8" = Hi( (h? wH
VE V.
where ((u) is categorical variable that determines the type of the layer represented by the node v. We note that a single MLP will not suffice since there are multiple layer types with different output sizes. We denote the concatenation of {W1 # hie as WH. For a particular client c, the weights {0") Jvvv are concatenated to form the client's weight vector (A WG, WH) mentioned above.
Objective and training procedure. Based on the GHN formulation, we can now state our training objective: we are looking for optimal GHN paramters (W&, WH) that simultancously minimize the empirical risk of all clients:
C
(Wa Wi) = ung. (Wo.WH)
(fe(tess) Oc(As; Wa, WH)), Mey).
(4)
The training procedure of our GHN is based on local updates of the GHN weights, performed by all clients, followed by a GHN weight averaging process on the shared server. More specifically, the local optimization at each client aims to solve the following client-specific minimization problem:
(WG* W = argmin (W),Wa)) WH)
((Se(res: Ac Wa, Wa)), #=1
Σ
(5)
and performs predefined number of SGD iterations (or similar gradient- based optimization method). The locally updated weights are then averaged at the server side and redistributed to the clients for further updates. The procedure is summarized in Algorithm I
We note that FedAvg can be seen as a specific instance of FLHA-GHN. To see that, consider the standard FL setting where all client architectures identity map. Consequently, all clients will have the same latent node representations h.. the 1-hot linear mapping (i.e.,
the same and the GNN implements the
encoding of the layer. If the MLPs (Hi(: wh))) tez are all implemented
0" Hw) (hw) = Wi( h.,) then averaging their parameters Wi()) is equivalent to directly averaging the client network parameters.
Hypernetwork weight initialization. Poor weight initialization of CNNs have detrimental effect to the training behavior causing either vanishing or exploding gradients. This has been carefully studied in Glorot & Bengio (010): He al. (2015b). When the network weights are outputs of
5

---

Under review as a conference paper at ICLR 2022
Algorithm 1: FLHA-GHN
Input: R: number communication rounds. C: number of clients, L: local updates.
Initialize shared GHN weights:
for r= R...,, do
Server shares current GHN weights (Wa, Wa) with all clients c E {I,...,") C+; for corc= I..,,ooo
Update GHN weights by local optimization for update steps client (see Eq. equation 5
Send updated GHN weights (W& Wa) to server;
end
Average GHN weights: We + J Σ w. Wn + WW::
end
another network, proper initialization of the hypernetwork parameters must be carried out so as to output the proper main network initialization. This has been observed by Chang et al. (2019) who proposed variance formula. We empirically found that initializing WH using zero bias and Xavier-normal initialization multiplied by /2cm/ daa translates to accurate Kaiming initialization of the client network. Here Gin and dia are the number channels in each layer and the hidden layer dimension of Hi See supplementary for more details.
Inference and local refinement.
inference time, the trained GHN applied to the graph weights (A; Wa. WH). In
representation of given architecture A, yielding architecture-specd contrast to the training stage, where these weights are used as-is, at the inference stage, they can be used initialization and then locally refined by using the client's 's local dataset. The amount of improvement depends on the amount local data. Furthermore, since this data has already been used for training the shared GHN model, minor improvements should be expected before overfitting might occur. We found that a single refinement epoch of just the linear prediction head worked well.
4.3 IMPLEMENTATION DETAILS
For full reproducibility of our results, will release code upon publication. We ran extensive hyperparameter search using Biewald (2020) for 4 architecture setup using a fixed number 500 epochs. We found the optimal values (which introduced in (Morris et al. 2019) with T = 6 layers, a latent dimension of 51, SGD with learning rate of 0.009 and cosine scheduling. Inspired by Zhang et al. (2020), we experimented with both synchronous and asynchronous message-passing that respects the directional graph structure, yet experiment did indicate an advantage the latter. For the client architectures, we define k: 10 different layer types. The different architectures and non parametric layers are described in detail in the supplementary material. For each layer type, a 2-layer MLP hypemneworr is used with latent dimension of 16, and leaky-ReLU activation function. We adopted the idea from Litwwin Wolf 2019) and separately predict the scale s and weight W for each node so that the layer parameters become sW. Implementation
used in all our experiments)
be: the GNN
PyTorchPaszkee done in Py Torch Paszke et al. (2017) using PyTorcl
Geometric (Fey Lenssen 2019) and training was done on a cluster with NVIDIA DGX V100 GPUs. The range of parameter weep and implementation details for basclines are provided in the supplementary.
5 EXPERIMENTS
Here we describe comparative experimental study of our approach. In Section (5.]) we review the methods that we compare with. Sections 5.2 provide results for two applications. Finally, we show ablations including generalization to unseen architectures (Sec. (S44. the effect of communication rate (Sec. (.5), and the importance of the GNN (Sec. 5.6). In the supplementary material, also demonstrate strong results with our method in unbalanced data distribution scenarios.
We demonstrate our FLHA performance two tasks. For natural image classification we use two standard benchmarks in FL: CIFAR-10 10 and CIFAR-100. Another task involves disease classification in chest x-ray scans. Each experiment consists of four different network architectures, named Arch 0-3, Arch is a standard ResNet18 He et al. (2015a) and the other are variants of it: Arch has 10 layers and no skip connections, Arch 2 has 12 layers and skip connections used only in the first 8

---

Under review as a conference paper at ICLR 2022
layers, and Arch 3 has 12 (different) layers with skip connection in the last layers. Detailed designs can be found in Sigure[Sjin in the supplementary. In all experiments, we split the data (either uniformly or following& Dirichlet distributions, indicated in the setup) between 4, 8, and 16 clients while maintaining an even distribution of the different architectures among clients (for instance, in the 8 clients case, we use 2 clients for each architecture).
5.1 METHODS IN COMPARISON
Since we are the first to address the FLHA setup (FL with hetcrogencous architectures without disclosing architectures or relying external data), compare performance with two variants of previous methods that we adapted to the FLHA setup. (1) A local-distillation baseline, which we name FLHA-distilaation. As variant of the distillation technique proposed by Lin et al. (2020), this baseline is based on locally available data and does not require the clients to disclose their architectures. In detail, for the first step we use standard federated averaging McMahan et al. (2017b) to train a shared architecture. The trained model is then sent to the client for distillation using only local data. (2) new variant of pFedHN (Shamsian et al. 202I). where the individual client heads where changed to account for different architectures (different number of outputs according to the number of parameters). In addition to these baselines, we also report results of Local training where each client performs standard (non-FL) training on its local train samples. This can be seen as a lower bound. For completion, we also include an Upper bound score. This is the result if all the clients used the most expressive Arch O, and trained in a standard FL.
5.2 RESULTS CIFAR
We experiment with the widely known IIFAR-IO/CIFARR 100 Krizhevsky et al. (2009) image classi- fication datasets that contain 60,000 32 x 32 natural color images in 10/100 classes with predefined train/test split. The results under FLHA are summarized in Table T Our method outperforms the local distillation and PFedHN baselines consistently expect on CIFAR-10 when local data is 25% where we perform similarly or slightly worse than local distillation. Importantly, as can be expected, dittilitiilaooon deteriorates rapidly. By contrast,
as the amount of locally available data decreases, our proposed method shows gradual degradation, resulting in an improvement of between 17 and 25 points in accuracy over the baseline when local data percentage is at 6.25%. Evidently, despite being tuned to its best performing parameters per setup (dataset+split) pFedHN performs poorly. This can be attributed to weight sharing only occurring the shared MLP responsible for the encoding the architecture, so the separate predictions heads do not benefit from federation.
5.3 RESULTS ON CHEST X-RAY
As discussed in the introduction, main motivation for FLHA is to allow cross-organizational collaborations. This is especially important when multiple entities have access to valuable data that cannot be shared, but use different neural architectures for processing it. Such is the case for medical clinics and hospitals. To showcase the application of our method to such a real-world problem, we tested it using medical data of X-ray images. The Chest X-ray Wang al. (2017) dataset comprises 1 112,120 120 frontal-view X-ray images of 30,805 unique patients with fourteen disease image labels (images can have multiple labels), extracted from medical reports using natural language processing. Table reports the average AUC score across 14 binary classification tasks for our method as well as the baselines. Our method achieves consistently superior performance.
5.4 GENERALIZATION TO UNSEEN ARCUITECTURES
Generalization to unseen architectures. In standard FL., a client that did not participate in the training procedure may still benefit from the pre- treitraindd network. We want FLHA-GHN to enjoy similar generalization, namely, that the GHN could generalize to clients whose architecture has not been observed during training. Since neural architectures often share local structures, with the various variants ResNets, when a new architecture is composed of local structures that have previously been seen, it may be feasible to generalize to new composition of these known components (compositional generalization). The reason is that T -llayer the T -hop local connectivity pattern around each node. When the same T-hop neighborhood appears in new architecture, possibly in a different position in the computational graph, the GNN will be able to predict reasonable weights since it was trained on such local structures (Yehudai et al. 2021).
paassage-psssnng GNNs essentially encode
7

---

Under review as a conference paper at ICLR 2022
CIFAR-100
CIFAR-10
data
method
Arch 0 Original
Arch
Arch
Arch Skip last
Arch Original
Arch No skip
Arch Skip first
Arch
No Skip
Skip first
Skip last
Upper-bound
73.2
93.6
Local training
49.5
50.2
52.2
48.2
86.3
84.6
84.7
85.4
25%
pledHN
38.0
38.7
46.5
34.7
85.8
84,
85.8
84.4
FLHAA Distillation
51.1
44.8
49.0
46.8
90.0
89.4
without graph
55.0
52.1
45.2
88.9
87.5
87.5
FLHA-GHN (ours)
56.0
52.5
50.8
54.4
90.4
89.0
87.3
88.5
Upper-bound
71.1
93.2
Local truining
36,4
32.5
38.1
33.5
79.2
75.6
74.5
80.0
12.50%
pFedHN
17.8
25.0
24.0
20.2
77.6
76.0
78.3
74.8
FLHA Distillation
44.2
33.2
41.3
35.2
85.7
84.6
84.2
83.3
without graph
55.4
52.9
49.5
54.3
89.0
88.2
86.2
89.1
FLHA-GHN (ours)
52.2
50.3
49.1
51.9
86.5
Upper-bound
71.3
92.9
Local training
22.7
19,4
22.0
22.0
76.9
75.6
75.6
74.5
6.25%
pFcdHN
13.6
14.4
16.4
12.6
58.2
56.6
57.3
50.7
FLHA Distillation
25.3
23.1
30.0
23.6
77.4
75.9
75.7
74.8
without graph
49.7
51.6
34.4
50.6
87.9
86.1
83.8
87.5
FLHA-GHN (ours)
50.3
50.4
46.9
49.2
86.8
85.6
83.9
85.4
Table 1: Accuracy CIFAR-10/10 datasets
Generalization architectures
clear benefit of such generalization is that given a new architecture, GHN immediately populate to give better initialization. To test this we ran a "leave one out" experiment using CIFAR- dataset, where we let 3 clients with 3 different architectures train in an FLHA fashion. We then introduce 4th ar- chitecture and refine it using only local data. We com- pare against training that architecture from scratch on that same data. The results shown in Figure and Table 6 in the supplementary are very encouraging. When refining local data, the unseen architecture performance ramps up very quickly and reaches sim- ilar performance that achieved by the architectures that participated the FLHA-GHN process. contrast, the training faiming-from scratch alternative takes much longer converge and reaches consistently lower final performance. We further stress tested the generalization capabilities, by introducing much smaller CNN with only 4 layers. A practical scenario of this sort is training model
Accuracy
60
From scratch
FLHA
2500
5000
7500
10000
15000 17500
# iterations
Generalization to unseen archi- tectures: our method (light blue) quickly ramps up to high performance and maintains considerable gap compared to training from scratch (dark blue).
Figure 3:
edge device together with stronger models on the server. Table 7 in the supplementary shows the performance of the architectures used in our main experiments, and how they are influenced when each is replaced by a smaller architecture. The average drop in performance of 2.2 = 1.4 keeps the performance well above the local-training alternative. Despite the different architecture, when initialized with model trained on the four main architectures, the fast convergence property persist (see Figure
5.5 COMMUNICATION RATE
Local epochs accuracy
In FL, the amount of data transfer could become bottleneck. A common remedy to reduce the fre- quency of clients-server communication To study the effect of communication rates on the FLHA-GHN performance, we trained 4 architectures on CIFAR- 10. At each experiment, clients train locally for L epochs before sending local GHN weights to server for averaging. Client-server communication occurs R times, keeping L x R at a constant value of shows the results for R between and 20. While more frequent communication leads better performance, we observe that even with times less frequent communication the overall performance decreases by less than 1%.
Accuracy
mean
std
5.0
10.0
12.5
17.5 20.0
local epochs
Figure 4: FLHA-GHN CIFAR-10 with different communication rates.
1000.
8

---

Under review as a conference paper at ICLR 2022
5.6 GNN IMPORTANCE
For our FLHA-GHN, all models created from the same of layer types. The more these layers are repeated in the computational graph, the more sharing could be done, allowing better utilization of the federation. Importantly, while same layers could be directly averaged, it is important that their location within the graph is encoded. This is done via message passing in our GNN. To verify its importance, we implemented variant of our method that does not use GNN and instead directly averages MLPs corresponding to the same layer types. Table shows this ablation, for which we used CNNs with 4,7.10 and 13 layers all of type 3 64 x 64 except for the first and last. With this amount of sharing, the results highlight that the GNN is instrumental for performance gains.
CIFARIO
CIFARIO0
layers
layers
10 layers
13 layers
layers
layers
10
layers
55.3
75.5
75.1
69.9
27.8
40.0
36.5
29.4
25%
/
63.6
77.4
77.0
74.4
26.4
41.1
40.4
33.5
60.2
75.5
73.8
70.1
32.4
41.5
41.6
12.50%
67.7
78.2
77.8
74.4
36.1
43.5
42.3
35.0
Table 2: The importance of the GNN. Under an FL.HA setup we use architectures sharing a repeated layer type. The result show significant gains due message passing informing the layer encoding with their placement within the architecture.
ChestX-ra Wang et al 2017
method
Arch 0 Original
Arch No Skip 2
Arch Skip lirst
Arch
Skip last
Upper-bound
78.0
Local training
60.5
61.8
58.7
60.3
25%
DFodHN
62.4
60.5
63.0
63.7
FLHA Distilation
67.5
67.8
65.4
69.8
FLHAA GIIN (ours)
72.2
73.2
70.5
Upper-bound
78.0
Local training
59.5
59.4
57.1
59.8
pFcdHN
62.5
39.9
58.4
61.0
Distillaionn
66.4
FLHI GHN (ours)
74.0
71.9
70.6
73.5
Upper-hound
77.0
Local training
59.2
58.5
56.9
58.9
6.25%
pledHN
60.4
FLIIA Distillation
65.7
62.8
61.72
66.5
(ours)
Table 3: Average AUC NIH Chest X-ray dataset.
CONCLUSION AND LIMITATIONS
In this work we have proposed a new setup for federated learning called FLHA, as well as first solution based GHN. The results of experiments promising: by modeling neural networks as graphs with layers as nodes, we have shown that GNNs can utilize recurring structures and facilitate efficient federation during learning and even generalization a few limitations to our solution. Firstly, it is limited to architectures built from predefined building blocks. In theory, this set could be arbitrarily large, but architectures that use disjoint sets of nodes may lose efficiency when sharing knowledge. Secondly, as indicated by the bound performance, our solution still is not yet comparable with came-architecture FL performance. Finally, a more robust aggregation technique would alleviate the undesirable degradation of performance when local data limited. Our hope that this research will lead to further study of this important setup.
new architectures. There are, however,
9

---

Under review as a conference paper at ICLR 2022
Ethics Statement. In this paper, study federated learning setup in which the clients use different neural architectures. We term this setup Federated Learning with Heterogeneous Architecture (FLHA) and propose graph hypernetwork solution. By relaxing the single architecture constraint, FLHA- GHN could facilitate useful cooperation beyond a single organization and create opportunities for smaller-scale entities train efficient models. We helieve that medical applications would be the primary beneficiary of our framework. While we include experiments on image classification, our contributions are not lask-specific and could be applied a variety of applications. We expect that our method will be used for positive ends, but similarly to all federated learning approaches, adversarial behavior possible.
Reproducibility Statement. To guarantee reproducibility we will release code upon publication. The main method as well as baselines described in the paper are hosted in a single repository and run from a docker container. Implementation details along with optimal parameters for the main method are described in section 4.3 A discussion on how to initialize the hypernetwork included
in 4.2 and demonstrated in C Network architectures are described in A
search for FLHA- is described in D Implementation details for the local distillation and pFedHN baselines
given in and respectively.
REFERENCES
Naman Agarwal, Ananda Theertha Suresh, Felix Yu, Sanjiv Kumar, and H Brendan Mema- han. cpsgd: Communication-fficientt and differentially-private distributed sgd. arXiv preprint arXiv: /805. 10559, 2018.
Lukas Biewald. Experiment tracking with weights and biases, 2020. URL https:: /tttps::/ww.... www wandb com/ Software available from wandb.com.
Andrew Brock, Theodore Lim, James M Ritchie, and Nick Weston. Smash: one-shot model architecture search through hypernetworks. arXiv preprint urXivv/708.000 05344, 2017.
Oscar Chang, Lampros Flokas, and Hod Lipson. Principled weight initialization for hypernetworks. In International Conference on Learning Representations, 2019.
Mingzhe Chen, Nir Shlezinger, H Vincent Poor, Yonina Eldar, and Shuguang Cui. Communication- efficient federated learning. Proceedings of the National Academy of Sciences, 118(17), 2021.
Xinyan Dai, Xiao Yan, Kaiwen Zhou, Han Yang. Kelvin KW Ng. James Cheng, and Yu Fan. Hyper-sphere quantization: Communication--fficieee sgd for federated learning, arXiv preprint arXiv:1911.04655, 2019.
Matthias Fey and Jan E. Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-256. JMLR Workshop and Conference Proceedings, 2010.
David Ha, Andrew Dai, and Quoe V Le. Hypernetworks. arXiv preprint urXiv:1609.09100 2016.
William L. Hamilton, Rex Ying, and Jure Leskovec. Inductive representation learning on large graphs, 2018.
Filip Hanzely and Peter Richtárik. Federated learning of mixture of global and local models. arXiv preprint arXiv:2002.05576. 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015a.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification, 2015b.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network, 2015.
10

---

Under review as a conference paper at ICLR 2022
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification, 2019.
Gao Huang. Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4700-4708, 2017.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. arXiv preprint 1rXiv:1912.04977, 2019.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning, pp. 5132-5143. PMLR, 2020.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.69823 2014.
Benjamin Klein, Lior Wolf, and Yehuda Afek. A dynamic convolutional layer for short range weather prediction. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. pp. 4840-4848, 2015.
Jakub Konečný, Brendan McMahan, and Daniel Ramage. Federated optimizatioon::istribueee opti- mization beyond the datacenter, 2015.
Jakub Konečný, H. Brendan McMahan, Felix X. Yu, Peter Richtárik, Ananda Theertha Suresh, and
Dave Bacon. Federated learning: Strategies for improving communication efficiency, 2017.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. CIFAR-10 (canadian institute for advanced research), 2009. URL http::tttt.ww.... cs toronto. www.
htmll
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. Imagenet classification with deep con- volutional neural networks. Advances in neural information processing systems, 25:1097-1105, 2012.
Viraj Kulkarni, Milind Kulkarni, and Aniruddha Pant. Survey of personalization techniques for federated learning. In 2020 Fourth World Conference on Smart Trends in Systems, Security and Sustainability (WorldS4), pp. 794-797 IEEE, 2020.
Wenqi Li, Fausto Milletari, Daguang Xu, Nicola Rieke, Jonny Hancox, Wentao Zhu, Maximilian Baust, Yan Cheng, Sébastien Ourselin, M Jorge Cardoso, et al. Privacy-preserving federated brain tumour segmentation. International Workshop on Machine Learning in Medical Imaging, pp. 133-141. Springer, 2019.
Yujia Li, Daniel Tarlow, Mare Brockschmidt, and Richard Zemel. Gated graph sequence neural networks, 2017.
Tao Lin, Lingjing Kong, Sebastian U. Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning, 2020.
Gidi Littwin and Lior Wolf. Deep meta functionals for shape representation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 1824-1833, 2019.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication-efficee learning of deep networks from decentralized data. Artificial Intelli- gence and Statistics, pp. 1273-1282. PMLR, 2017a
H Brendan McMahan, Daniel Ramage, Kunal Talwar, and Li Zhang. Learning differentially private recurrent language models. arXiv preprint arXiv:1710.06963, 2017b.
Christopher Morris, Martin Ritzert, Matthias Fey, William L Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe Weisfeiler and leman go neural: Higher-order graph neural networks, In Proceedings of the AAAI Conference Artificial Intelligence, volume 33, pp. 4602-6600 2019.

---

Under review as a conference paper at ICLR 2022
Viraaji Mothukuri, Reza M Parizi, Seyedamin Pouriyeh, Yan Huang, Ali Dehghantanha, and Gautam Srivastava. A survey on security and privacy of federated learning. Future Generation Computer Systems, 115:619-644 640, 2021.
Eliya Nachmani and Lior Wolf. Molecule property prediction and classification with graph hypernet- works. arXiv preprint 12Xi::2002.002-4:: 2020.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary De Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in PyTorch. In Advances in neural information processing systems, 2017.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer- assisted intervention, pp. 234-241. Springer, 2015.
Anit Kumar Sahu, Tian Li, Maziar Sanjabi, Manzil Zaheer, Ameet Talwalkar, and Virginia Smith. On the convergence of federated optimization in heterogeneous networks. arXiv preprint arXiv:1182..0112,3 3, 2018.
Aviv Shamsian, Aviv Navon, Ethan Fetaya, and Gal Chechik. Personalized federated learning using hypemeetwrrs. In International Conference on Machine Learning, 2021.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXivv1409.1555 2014.
Vincent Sitzmann, Michael Zollhöfer, and Gordon Wetzstein. Scene representation networks: Contin- uous 3d-structure-aware neural scene representations. arXiv preprint arXiv:1906.01611 2019.
Sebastian U Stich. Local sgd converges fast and communicates little. arXiv preprint arXiv:1805.0970: 2018.
Joseph Suarez. Character-level language modeling with recurrent highway hypernetworks. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pp. 3269-33 3278, 2017.
Sasha Targ, Diogo Almeida, and Kevin Lyman. Resnet in resnet: Generalizing residual architectures. arXiv preprint arXiv: 1603:00039 2016.
Xiaosong Wang, Yifan Peng, Le Lu, Zhiyong Lu, Mohammadhadi Bagheri, and Ronald Summers. Chestx-ray8: Hospital-scale chest x-ray database and benchmarks on weakly-supervised classifica- tion and localization of common thorax discases. In 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 3462-3471, 2017.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? arXiv preprint 11X1v:1810.00826, 2018
Qiang Yang, Yang Liu, Tianjian Chen, and Yongxin Tong. Federated machine learning: Concept and applications. ACM Transactions Intelligent Systems and Technology (TIST), 10(2):1-19, 2019.
Gilad Yehudai, Ethan Fetaya, Eli Meirom, Gal Chechik, and Haggai Maron. From local structures to size generalization in graph neural networks, 2021.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Trong Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks, 2019.
Chris Zhang, Mengye Ren, and Raquel Urtasun. Graph hypernetworks for neural architecture search, 2020.
Michael Zhang, Karan Sapra, Sanja Fidler, Serena Yeung, and Jose M. Alvarez. Personalized federated learning with first order model optimization. In International Conference on Learning Representations, 2021.
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv preprint /806 00582, 2018.
12

---

Under review as a conference paper at ICLR 2022
Wennan Zhu, Peter Kairouz, Brendan McMahan, Haicheng Sun, and Wei Li. Federated heavy hitters discovery with differential privacy. In International Conference on Artificial Intelligence and Statistics, pp. 3837-3847. PMLR, 2020.
13

---

Under review as a conference paper at ICLR 2022
A NETWORK ARCHITECTURES
Each experiment in this paper used four different types of architectures split among the different clients plus an additional small architecture for the stress test. There are ten different types of nodes (layers) in each architecture. Figure
the architectures used in the CLFAR-IO/10) experiments. The types aatiis_iin___hhhaa <channels, out> <kernel Bxernel_siz>>> <stride>". In the chest X-ray experiment, where the input images grayscale and of size 224 x 224 the first convolutional layer type is "cl_64_k7_s2" type.
of convolutional layers are denoted by e<channels in>_
CIFAR-10, CIFAR-100, 100, and Chest X-rays, respectively, use different types linear layers with 10, 100, and 14 output dimensions.
Arch 0
Arch 1
AURUUT
FOR 1
(2
0138 281
INUUUTT
(296 290 A2
BUAUU
AULUU
MUASLE
HURUU
128.1188.2.11
BURUU
Arch 2
Arch 3
GUBDUT
SUPU
PUB
DE
AUAUU
AUJUU
Figure 5: The four client architectures used in 1IFAR-10/100 experiments.
B NON PARAMETRIC LAYERS
Since our main architectures all use residual connections and same activation type, the graph con- nectivity suffices to express the architecture. However, in general, different non-parametric layers might be desired. For that, we implemented three versions of ResNet in which several different residual layers were replaced by concatenation, resulting in hybrid addition-concatenatioo To that end, we included two new non-parametric types of nodes in our layer type set: "add" and "concat". Participating message passing, these nodes produce latent embedding, depending on where they reside in the graph. No additional parameters are required for them. We trained these 3 architectures together with the vanilla ResNet with four and eight clients with the CIFAR and CIFAR (00 datasets. On CIFAR10 10, the average results by architecture type are: 88.96, 89, 89, 89.2 for 4 clients and 88.6, 88.1 1, 88.3, 86.8 for clients. On CIFARI00, the results are: 56, 56.2, 56.5, 48.200 for 4 clients and 50.6. 50.4, 47.7, 46.0 for 8 clients. The results are on par with those shown in Table In particular, the performance obtained by the vamilla ResNet architecture on CIFAR 10/100 with 4 and 8 clients respectively is: 88.96. 88.6. 56 and 50.6, whereas its performance under the FLHA-de -is--iiitlllation baseline 90, 85.7. 51.1, 44.2. In this example, FLHA-GHN shows success in training with two commonly used non-parametric layers.
14

---

Under review as a conference paper at ICLR 2022
C HYPERNETWORK INITIALIZATION
As described in Section 4.2 of the main paper, a proper initialization of the hypernetwork weights instrumental to successful training of the client networks. In @ is shown on an example convolutional layer with dimensions: 3 x 64 128. In both plots, the blue colored histogram shows the distribution of the desired Kaiming weight initialization He et al. (2015b). When the weights generated by a hypernetwork, a standard initialization of the hypernetwork would generate the orange histogram shown the left. We show what that histogram looks like after our initialization scheme on the right.
lat-dim: dim-in: 64 dim-out: 128 kemeel-siee:
lat-dim:
dim-in: dim-out: 128 kernel-size:
12000
30000
10000
8000
4000
2000
0
6.0
0.05
Standard Initialization
Our initialization
Figure 6: Hypernetwork weight initialization. We compare (blue:) a direct Kaiming weight initializa- tion He et al. (20156) of convolutional layer with (orange:) the resulting weight initialization by the hypernetwork, without (left) and with (right) our initialization scheme.
D HYPERARAMETER SEARCH FOR OUR FLHA-GHN
We ran an extensive hyperparameter search using Biewald 2020) for a architecture setup using fixed number of 500 epochs with different GNN types: GraphConv(Morr)
al.) (019). et al. 2018); number of GNN layers
GatedGraphConv et al. 2017), and GraphSAGE(Haminn from with latent dimensions between 16 and 128: hypernetwork Hi bottleneck dimension between 16 and 64: learning rates between lc-4 and 0.1; SGD (with and without cosine scheduler) and Adam(Kingma Ba 2014) optimizers with weight decay values between 5e-4 and 5e-6.
E IMPLEMENTATION DETAILS: LOCAL DISTILLATION
Baseline distillation from teacher model trained via standard (same- architerchit FL done using distillation loss Hinton et al. (2015): ) (1 a)CE((prrdd: y) + aKL(Ypred: Vteacker) x 2T where CE and KI. denote Entso-Entropy and Kullback Leibler, respectively. The softmax in the KI. loss is taken with respect to a temperature 7 20.0 and 0.7. We trained distillation as well as the main FL models for 200 epochs with SGD. The learning rate for training the teacher network (with same-arco eeme-archieecture FL) is set to 0.1 with cosine scheduling, When distilling from the teacher network to the student, we use a learning rate to 0.01.
F IMPLEMENTATION DETAILS: PFEDHN
We found this architecture to be quite sensitive to hyperparameters, hence unlike our method which uses the same of parameters in all experiments, here chose the best performing parameters per setting. Hyperparameters sweep on the following parameters: learning rate, number of hidden layers in shared mlp, latent dimension size, optimizer (adam and sgd), and weight decay.
15

---

Under review as a conference paper at ICLR 2022
Tha noet Ply Dest
Transet Jetaset P. Dlyy Dat
Parrset Detsset P. P_yy Dest
Tramett Detoset P. P.2yl Dist
Berneet Datanet -Tyy Dett
1200
1320
3000
1.000
1200
3520
3002
100
60C
2000
600
600
500
-00
1803
Category
Categpry
Tratnar Ply)
Halnnl Jalavel P (.y
Rair Delinat 11y) bisi
Iainnel Dslavet V' zlyl
bairaad Ltatinad dly) Dist
5003
2500
4023
2000
1250
3000
3043
2000
1000
500
750
1853
500-
3000
3000
3000
5000
-020
6000
4022
3000
3000
Gategayy
Table 4: Each row in the table shows the unbalanced class distribution for 4 clients, with the original balanced distribution of the left. Three different a values are shown: (top) 100, (mid) 1, (bottom) a = 0.1
G UNBALANCED DISTRIBUTION
In collaborative training between different entities, clients' data may be distributed unevenly. In medical data, for example, this may occur when clinics specialize in certain diseases different sensors. Thus, in addition to architectural differences, FLHA can also have data imbalance. Here we study the behavior of FLHA under such unbalanced data distributions.
We follow Yurochkin et al. 2019); Hsu et al. (011)) Lin et al. (2020) and use the (symmetric) Dirichlet distribution, parameterized by a concentration parameter a to split the CIFAR-IO training and test sets between the different clients. Table ssoows the resulting per-client class distributions. As can be seen, the smaller alpha is. the less balanced the distribution is.
Table Sisows the performance on CIFAR-IO under 3 different a values: 0,1, 100. The smaller a the more unbalanced the distribution is. Class distribution under the different values are shown in Table
When the training data is unevenly distributed. performance can either be measured in similar distributed test set or a balanced test set. In the former case, a client would like high performance on local, biased samples. Suppose a hospital specializes in a specific disease and hopes to improve model performance related to that disease through FL.. We refer to that "unbalanced" metric. Alternatively, a client might be interested in balancing its model bias, in which case the performance on the full ("balanced")) test set is of interest. We compare our FLHA-GHN against a local training. We also include the usual upper bound performance of standard FL with all clients using the same architecture. Table shows that FLHA-GHN outperforms local training in both "unbalanced" and "balanced" tasks and across all alpha values. Specifically, it can be seen that local training achieves high performance only on test sets of similar distribution, but severely sacrifices performance on sets of balanced
16

---

Under review as a conference paper at ICLR 2022
Client
Client
Client
Client
method
unbal.
bal.
unbal.
bal.
unbal.
bal.
unbal.
bal.
FLHI--GII
89.7
89.5
87.6
87.6
87.0
86.2
88.3
88.1
100
Local
81.8
75.1
74.2
80.5
80.5
81.1
80.3
Standard FL
92.9
93.7
93.2
93.7
94.1
93.7
94.5
93.7
FLHA-GHN
87.1
85.5
84.9
84.6
85.0
85.5
Local
87.8
75.8
87.3
83.5
84.4
84.0
84,4
83.5
Standard FL
93.8
93.0
92.9
93.0
92.6
93.0
92.9
93.0
FLHA-GHN
94.6
62.2
98.3
28.7
92.3
34.0
92.4
61.0
0.1 Standard FT
Local
93.6
51.1
96.6
27.8
90.9
25.2
90.7
54.1
69.3
53.8
58.2
53.8
42.3
53.8
49.6
53.8
Table 5: FLHA with unbalanced distribution. In the table, a corresponds to the level of unbalanced,e.g. a=100 (almost uniform), a=0.1 (extremelly unbalanced). unbal. and bal. are short for unbalanced and balanced and correspond to the distribution of the test set with unbalanced being the same distribution of each client's s training set.
GHN Init
From scratch
Arch (original)
86.1
83.9
Arch (No skip)
84.2
81.3
Arch (Skip first)
84.3
83.7
Arch 4 (Skip last)
85.6
83.6
Table 6: Generalization unseen architectures: seiv--on--rrchtecturee experiment CIFAR-10. Each row the accuracy on a held-out architecture while training on the other architectures.
distributions. The same architecture FL also shows trade-off. Despite being the most performant on the balanced test set, it sacrifices accuracy on local distributions. FLHA-GHN's capability to improve local training in both tasks can be attributed to the inherent personalization of the network. That is, beyond its ability to adapt to new architectures, FLHA-GHN can also learn a personalized weight prediction according to the client distributions. This result aligns well with the observation of Shamsian et al. (2021).
H GENERALIZATION ADDITIONAL RESULTS
Table 7 shows the performance of the 4 architectures used in our experiments, and how they are influenced when each is replaced by smaller architecture. The average drop in performance of 2.2 + 1.4 keep the performance well above the local-training alternative.
Replaced
Arch I
Arch 2
Arch 3
Arch 4
None
90.4
89.0
87.3
88.5
Arch
88.2
86.9
86.4
80.3
Arch
88.8
87.1
79.6
88.2
Arch 88.6
81.6
85,4
86.2
Arch 77.5
83.8
85.4
83.6
Table 7: Training with much smaller architecture shows an average performance drop by 2.2 + 1.4 pts. However, this well above the local training alternative.
17

---

Under review as a conference paper at ICLR 2022
Generalization to unseen small architecture
85
Accuracy
80
75
70
www
65
From scratch
60
FLHA
55
50
2500
5000
7500
10000
12500
15000 17500
iterations
Figure 7: Generalization smaller 4 layer CNN architecture. Our method (light blue) quickly ramps up to high performance and maintains a considerable gap compared to training from scratch until convergence(dark blue).
18


Reviews:
[{"review_id": "z1g38QODIL2", "reviewer": "Reviewer_Tk9o", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength\n\n1. This paper leverages the GHN techniques to solve the problem of system heterogeneity under FL settings.\n\n2. .The trained GHN model can be generalized to unseen architectures.\n\n3. Extensive evaluation on three real datasets demonstrate the superior performance of the proposed approach for tackling the heterogeneity issue in federated learning.\n\nWeakness\n\n1. The performance highly depends on the quality of trained GHN, and GHN is often affected by the distribution of client data. When the datasets on local devices are highly non-iid, the performance of this method might be significantly dropped.\n\n2. The dimension of the final output is fixed, even though it can be solved by concatenating or slicing the final output. It reduces the flexibility of the generated model. For example, if the output size is 64, then the dimensionality of each layer should be the multiplier of 64. Copying layers and concatenating them together to form a bigger layer is too brutal and may not be a good choice. In practice, the parameter of one layer is not mirrored.\n\n3. In the experiments, what does 'without graph' mean? The accuracy of the proposed method is similar to the 'without graph.'\n\n4. The authors fail to provide the convergence analysis for federated learning, especially the convergence of GHN under FL settings.\n\n5. The paper is more like a straightforward application of the Graph Hyper Network from the centralized training to the federate learning.", "summary_of_the_paper": "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.", "summary_of_the_review": "Overall, the well designed structure makes the workﬂow clear and easy to follow, but the further analysis and discussion are expected to clarify the contributions in the techniques as well as in the evaluation section.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "l1x04ADRkQG", "reviewer": "Reviewer_jPp2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1.  The proposed solution is novel. The authors adopt a graph neural network with a MLP as a hyper-network to generate model parameters  for each client. This enables each client to protect its own model design as privacy whiling enjoying the benefit of sharing knowledge.\n2.  The proposed framework is generalizable to unseen architecture. The authors have shown that when a new client with a different model architecture is encountered the model performs better than learning from scratch.\n3.  Experiment results on the x-ray dataset are promising. \n\nWeakness:\n1. The variation of selected models of clients is narrow. They are all variants of ResNet with close size.  \n2. Experiment results on CIFAR-10/100 do not show evident advantage of the proposed method over baseline method.\n3. Adding a hyper-network adds extra computation cost. It is suggested to compare training efficiency experimentally. ", "summary_of_the_paper": "This paper studies the problem of sharing knowledge among different organizations without disclosing their model architectures under the framework of federated learning. The authors addressed this problem by sharing a graph hyper-network that can populate model parameters based on different model architectures. ", "summary_of_the_review": "Previous works on federated learning either focus on the mechanism of parameter aggregation or the aspect of privacy. This paper opens a new direction in FL where clients may not be willing to share their unique model designs. From this perspective, I think this paper has promising impact on the research field of FL. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "WkFRKyaqB8K", "reviewer": "Reviewer_U48w", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths: \n- The paper tackles an interesting and well-needed problem: architecture-agnostic federated learning. \n- The paper is clearly organized and written. I can follow it without much effort. \n- The method to represent neural architectures as graphs is interesting. \n\nDrawbacks: \n- **Generalization experiments seem weak.** For realistic application scenarios, a need is that neural architectures should be chosen corresponding to the computational power (as mentioned in paragraph 2 in the introduction). The networks used in this paper do not differ very much in terms of size (maximally ResNet18), and is thus a little weak to support the motivation. It would be better if you can do the experiments on networks with larger variance in size (e.g. generalize a ResNet152 to a ResNet18 or even smaller). \n\nFurther, for figure 7, does this figure imply a <50% zero-shot accuracy of the 4-layer network, and that further adaptation can only be achieved after fine-tuning? Then, does the adaptation result depend on the data upon which to perform fine-tuning? The result does not seem very supportive of the authors' claim of generalization. \n\n- **Ability to deal with very deep networks**. Following the previous drawback, I am considering how the proposed FLHA-GHN can deal with very deep networks, and here is why. Suppose you use a 6-layer GHN, the GHN can distinguish between local graph structures within 6 hops. However, GNNs cannot tell the absolute position of nodes given the same local neighborhoods (You et al. 2019). Thus, for very deep networks (e.g. ResNet152) with highly repetitive structures, networks at different depths but with similar local connection patterns may get very similar weights. I wonder whether the above analysis is correct and would like to see the replies from the authors. \n\n(You et al. 2019) Position-aware Graph Neural Networks. Jiaxuan You et al. ICML 2019. \n\n- **Number of clients**. For the experimental settings, you use 4, 8, 16 clients, while pFedHN uses 50, 100, 500. I wonder what are the reasons behind such a big difference in the number of clients. Any justifications? \n- I am curious about how big the graph hypernetwork used in the experiments is. Intuitively, suppose a client wants a network with a linear layer with 16 inputs and 100 outputs. That is 1600 parameters. To generate weights from an MLP, we need an extra linear layer with 1600 output dimensions. Thus, we would need something like $16\\times 16\\times 100$ parameters for the layer. It is nice if you can show such an analysis. \n- **Heterogeneous data**. As shown in Appendix G, the approach to use GHN may lead to worse performance (compared to standard FL) on a balanced test set. This seems a drawback of the proposed method, as FL commonly needs to deal with heterogeneous data. The authors may want to justify the results. Specifically, in practice, we cannot guarantee that the test data distribution at each client is the same as its training data. \n\nMinor issues. \n- I fail to understand Table 7. ", "summary_of_the_paper": "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. ", "summary_of_the_review": "The paper has its merits. The paper tackles an important problem. The proposed method is novel and evaluated extensively. The writing is clear and easy to follow. However, the authors may want to justify some minor issues raised in the reviews. I am willing to raise my recommendation of this paper upon further clarifications. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Hc7UHYtFtmv", "reviewer": "Reviewer_7WPJ", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The proposed method addresses a key challenge in federated learned which is data heterogeneity. The approach is explained in good detail. \n\nThe main weaknesses in my opinion: \n\n- The novelty seems limited. The authors use a direct application of existing methods (graph hyper networks). \n\n- The modeling assumptions need to be explained in more detail. It is currently not very clear how the data distributions at clients is related to the neural architectures \\mathcal{A}_{c} \n\n- No theoretical analysis of the proposed method (what can you say about the sample size required to train a GHN?). \n\n- The description of the proposed method (Algorithm 1) needs to be expanded. Currently, there seems to be only one single sentence referring to Algorithm 1. \n\n- Numerical experiments use small-scale setups. It is not evident if the proposed method scales to millions of clients. ", "summary_of_the_paper": "The authors study methods for federated learning with clients using different neural network architectures. Graph hyper networks are used to predict useful weights of client-specific neural network architectures.  ", "summary_of_the_review": "* The clarity of the presentation needs improvement. First of all, the underlying modeling assumptions on the datasets are not clear. What do you assume about the data distributions $\\mathcal{P}_{c}$ and how are they related to the network architectures $\\mathcal{A}_{c}$. \n\n* It is nice that the authors write down two research questions at the end of the third paragraph Section 1. However, it is not very clear how and where these questions are answered later on. \n\n* Section 3 is very short. Maybe it could be expanded with a discussion of using graphs to represent neural architectures. \n\n* There needs to be more explanation of Algorithm 1 in the text. How can we tune the input parameter of Algorithm 1 ?  The steps of the inner loop in Algorithm 1 should be fleshed out to provide a sufficient level of detail for implementing Algorithm 1. Should the client architectures \\mathcal{A}_{c} be listed as input to Algorithm 1 ? \n\n* The literature review is missing a recent line of work on personalized FL (multi-task learning) using total variation minimization: \n\nD. Hallac, J. Leskovec, and S. Boyd, Network Lasso: Clustering and Optimization in Large Graphs, Proceedings SIGKDD, pages 387-396, 2015.\n\nA. Jung, A. O. Hero, III, A. C. Mara, S. Jahromi, A. Heimowitz and Y. C. Eldar, \"Semi-Supervised Learning in Network-Structured Data via Total Variation Minimization,\" in IEEE Transactions on Signal Processing, vol. 67, no. 24, pp. 6256-6269, 15 Dec.15, 2019, doi: 10.1109/TSP.2019.2953593.\n\nY. Sarcheshmehpour, Y. Tian, L. Zhang, A. Jung, “Networked Federated Multi-Task Learning,” CoRR abs/2105.12769, 2021\n\n* The following wordings are unclear: \n- \"... train a strong model ...\" what is a strong model ? \n- \"all clients need to share the same network architecture\" pls provide more justification for this claim. in general we can use arbitrary parametric models for clients as long as the parameter vectors are \"compatible\" e.g. have the same length. \n\n- \".. often run different networks ..\" what do you mean by \"running a network\" ? \n\n- \"... current FL approaches do not support mixtures of different architectures.\" im not sure if this is true. consider e.g. using different pre-trained nets as feature extractors for clients and then using the same last layers. as long as we use paramter vectors of same length for clients we can use e.g. network Lasso to implement FL. \n\n- \" Since different architectures have different layer compositions, representing layers as nodes allows meaningful knowledge aggregation across architectures. \" I do not understand this argument \n\n- \"...from some predefined family of models...\" what exactly is a family of models?\n\n- \"...for a suitable loss function ...\" what precisely makes a loss function suitable?\n\n- Figure 2 could include more details e.g. indicate hyper network parameters W_{G} , W_{H} and node representations h_{v}\n\n- \"... are learned from data using updates from all clients.\" unclear \n\n- consistently use abbreviation (e.g. HN) once introduced ! \n\n- \"... to represent non-parametric operation in the network ...\" what is a non-parametric operation ? \n\n- \"... see supplementary for more ... \" \n\n- \" ... and trained in a standard FL.\" \n\n- \" ... when local data is ...\" \n\n- \"...when local data percentage is at ...\" \n\n- \"...our GHN can immediately populate it ...\" \n\n- \" ... we have proposed a new setup ...\" pls refer to the relevant equation, algorithm or figure \n\n\n\n\n\n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Meta Review of Federated Learning with Heterogeneous Architectures using Graph HyperNetworks

This work investigates a method for federated learning in a neural architecture-agnostic setting. They do this by using a graph hypernetwork to predict the weights of given neural network architectures (which is not exactly known at the onset). The authors conduct federated learning experiments to demonstrate good performance on several real datasets, and also showed that the trained GHN model can generalize (somewhat) to unseen architectures (which are mainly in the ResNet family). Personally, as AC, I find the results very promising, and the experiments show that GHNs are highly applicable to real world applications. But the reviewers outline several weaknesses in the discussion that makes it difficult to recommend acceptance of this paper for ICLR 2022.

The main weaknesses of the work are that application is mainly focused on a narrow family of ResNet architectures (can it be shown to go beyond this? If not, can the writing be improved to show that this is useful enough for many applications?) Reviewer U48w suggested improvements to the generalization experiments, and other details that can be addressed in the writing. Reviewer Tk9o mentioned that this work can be seen as a straightforward application of GHNs (limited novelty), while other reviewers do acknowledge the novelty of the work. I recommend improving the writing to clearly address this and defend why this is not a straightforward application of previous work. With these improvements, I'm confident that this work will be accepted at a future ML conference or journal.

Even though I cannot recommend acceptance, both myself and other reviewers are looking forward to seeing improved versions of this work for publication in the future. As jPp2 also noted, “Previous works on federated learning either focus on the mechanism of parameter aggregation or the aspect of privacy. This paper opens a new direction in FL where clients may not be willing to share their unique model designs. From this perspective, I think this paper has promising impact on the research field of FL.” Good luck!
Decision: Reject