Federated Learning with Matched Averaging
Authors: Hongyi Wang
Mikhail Yurochkin
Yuekai Sun
Dimitris Papailiopoulos
Yasaman Khazaeni
Venue:  (2020)
Keywords: federated learning

Abstract: Federated learning allows edge devices to collaboratively learn a shared model while keeping the training data on device, decoupling the ability to do model training from the need to store the data in the cloud. We propose Federated matched averaging (FedMA) algorithm designed for federated learning of modern neural network architectures e.g. convolutional neural networks (CNNs) and LSTMs. FedMA constructs the shared global model in a layer-wise manner by matching and averaging hidden elements (i.e. channels for convolution layers; hidden states for LSTM; neurons for fully connected layers) with similar feature extraction signatures. Our experiments indicate that FedMA not only outperforms popular state-of-the-art federated learning algorithms on deep CNN and LSTM architectures trained on real world datasets, but also reduces the overall communication burden.

Published as a conference paper at ICLR 2020
FEDERATED LEARNING WITH MATCHED AVERAGING
Hongyi Wang"
Mikhail Yurochkin
Department of Computer Sciences
IBM Research
University of Wisconsin-Madison
MIT-IBM Watson Al Lab
hongyiwee wwn.gwwangssss... wisc. edu
mi miaiill.. wurochkiinn imm
Yuekai Sun
Dimitris Papailiopoulos
Department of Statistics University of Michigan yueka @umi ch. edu
Department of Electrical and Computer Engineering University of Maddionnn-Maaioon
io
Yasaman Khazaeni
IBM Research
yasaman. khazaeni Bus. ibm. com.
ABSTRACT
Federated learning allows edge devices to collaboratively learn shared model while keeping the training data on device, decoupling the ability to do model training from the need to store the data in the cloud. We propose the Federated matched averaging (FedMA) algorithm designed for federated learning of mod- ern neural network architectures e.g. convolutional neural networks (CNNs) and LSTMs. FedMA constructs the shared global model in layer-wise manner by matching and averaging hidden elements (i.e. channels for convolution layers: hidden states for LSTM; neurons for fully connected layers) with similar feature extraction signatures. Our experiments indicate that FedMA not only outperforms popular state-of-the-ar federated learning algorithms deep CNN and LSTM architectures trained on real world datasets, but also reduces the overall commu- nication burden
INTRODUCTION
Edge devices such mobile phones, sensors in sensor network, vehicles have access wealth of data. However, due to data privacy concerns, network bandwidth limitation, and device availability, it's impractical to gather all the data from the edge devices at the data center and conduct centralized training. To address these concerns, federated learning is emerging (McMahan et al.,
2017 Li al. 2019 Smith et al., 2017 Caldas et al. 2018 Bonawitz al.) 2019] Kairouz al., as a paradigm that allows local clients to collaboratively train a shared global model.
2019
The typical federated learning paradigm involves two stages: (i) clients train models with their local datasets independently, and (ii) the data center gathers the locally trained models and aggre- gates them to obtain a shared global model. One of the standard aggregation methods is FedAvg (McMahan et al. (017)) where parameters of local models are averaged element-wise with weights proportional to sizes of the client datasets. FedProx (Sahu et al.] (2018) adds a proximal term to the client cost functions, thereby limiting the impact of local updates by keeping them close to the global model. Agnostic Federated Learning (AFL) (Mohri et al. 2019p, as another variant of FedAvg, optimizes centralized distribution that is mixture of the client distributions.
One shortcoming of FedAvg is coordinate-on wise averaging of weights may have drastic detrimental effects the performance of the averaged model and adds significantly the communication bur- den. This issue arises due to the permutation invariance of neural network (NN) parameters, i.e. for
Work performed while doing an internship at IBM Rescarch.
Code available at athttps: / thub. com/ HMMA

---

Published as a conference paper at ICLR 2020
100
100
1-Round Feddug
Avg Local Ace
Testset Acc
80
2214
PFNM
60
60
FedMA
Laye--RRdddd FedAvg
Layer-Roudd FedProx
20
20
Training
LeWett on
VGCS on CIFAR-10
LeNet
VGGB CIFAR-10
LSTM Shakespeare
(a) Homogeneous
(b) Heterogeneous
Figure 1: Comparison among various federated learning methods with limited number of communications on LeNet trained on MNIST; VGG-9 trained CIFAR-I 10 dataset; L.STM trained on Shakespeare dataset over: (a) homogeneous data partition (h) heterogeneous data partition.
any given NN, there are many variants of it that only differ in the ordering of parameters. Prob- abilistic Federated Neural Matching (PFNM) (Yurochkin et al. (019b) addresses this problem by matching the neurons of client NNs before averaging them. PFNM further utilizes Bayesian non- parametric methods to adapt to global model size and to heterogeneity in the data. As result, PFNM has better performance and communication efficiency than FedAvg. Unfortunately, the method only works with simple architectures (e.g. fully connected feedforward networks).
Our contribution In this work, we demonstrate how PFNM can be applied CNNs and LSTMs, but we find that it only gives very minor improvements over weight averaging. To address this issue, we propose Federated Matched Averaging (FedMA), new layers-wise federated learning algorithm for modern CNNs and LSTMs that appeal to Bayesian nonparametric methods to adapt heterogeniety in the data. We show empirically that FedMA only reduces the communcations burden, but also outperforms state-of-the-ar federated learning algorithms.
2 FEDERATED MATCHED AVERAGING OF NEURAL NETWORKS
In this section we will discuss permutation invariance classes of prominent neural network archi- tectures and establish the appropriate notion of averaging in the parameter space of NNs. We will begin with the simplest case of a single hidden layer fully connected network, moving on to deep architectures and, finally, convolutional and recurrent architectures.
Permutation invariance of fully connected architectures A basic fully connected (FC) NN can be formulated as = a (x W22 (without loss of generality, biases are omitted to simplify no- tation), where a the non-linearity (applied entry-wise). Expanding the preceding expression 1 Et= o2..0(( (ar, Wi.:)), where and v denote ith row and column correspondingly and L. is the number of hidden units. Summation is permutation invariant operation, hence for any (W1, W2) there are L! practically equivalent parametrizations if this basic NN. is then more ap- propriate to write
-= (=), where IT is any L L permutation matrix.
Recall that permutation matrix is an orthogonal matrix that acts rows when applied the left and on columns when applied on the right. Suppose (W1, W2} are optimal weights, then weights obtained from training two homogeneous datasets X,, Xy {WilI,, I1 Wz) and (Wily,, IT W2).
now easy to see why naive averaging the parameter space not ap- propriate: with high probability II, * Hy and (W,ll, + W1 + WIll for any II. To meaningfully average neural networks in the weight space we should first undo the permutation IW1 + 1111y = W1.
2.1 MATCHED AVERAGING FORMULATION
In this section we formulate practical notion of parameter averaging under the permutation invari- ance. Let wiy be /th neuron learned on dataset j (i.e. 1th column of W in the previous example),
2

---

Published as a conference paper at ICLR 2020
8, denote the ith neuron in the global model, and c(:,-) be appropriate similarity function between a pair of neurons. Solution to the following optimization problem are the required permutations:
min (, (000), 0,) s.t. 8.
= Σd =
Then jui T/ and given weights {W W W provided by J clients, we compute the fed- erated neural network weights W1 ł, W W1,1liy and W22 + M3 W.2 We refer this approach matched averaging due to relation of equation 2 to the maximum bipartite matching problem. We note that if c(.,)) is squared Euclidean distance, we recover objective function similar to k-means clustering, however it has additional constraints on the "cluster assignments" T/. neces- sary to ensure that they form permutation matrices. In a special case where all local neural networks and the global model are assumed to have same number of hidden neurons, solving equation 2 is equivalent to finding a Wasserstein barycenter (Agueh & Carlier 201T) of the empirical distributions over the weights of local neural networks. Concurrent work of Singh & Jaggi (2019) explores the Wasserstein barycenter variant of equation
Solving matched averaging Objective function in equation 2 can be optimized using iterative procedure: applying the Hungarian matching algorithm (Kuhn 1955 to find permutation [ h corresponding to dataset j', holding other permutations { highy' fixed and iterating over the datasets. Important aspect of Federated Learning that we should consider here is the data hetero- geneity. Every client will learn collection of feature extractors, i.e. neural network weights, rep- resenting their individual data modality. As a consequence, feature extractors Iearned across clients may overlap only partially. To account for this we allow the size of the global model L to be un- known variable satisfying max; L,j I L < Σ, Ing where L is the number of neurons learned from dataset j. That is, global model is at least as big as the largest of the local models and at most as big as the concatenation of all the local models. Next we show that matched averaging with adaptive global model size remains amendable to iterative Hungarian algorithm with a special cost.
At each iteration, given current estimates of (mf (ms,,,, we find a corresponding global model (0; arg ming. Ejjyr, Thf 0=)] (this is typically closed-form expression or a simple optimization sub-problem, e.g. a mean if c(.,.) is Euclidean) and then we will use Hungarian al- gorithm to match this global model to neurons {wpthie of the dataset j' to obtain a new global model with L S L' L+ neurons. Due data heterogeneity, local model j' may have neurons not present in the global model built from other local models, therefore we want to avoid "poor" matches by saying that if the optimal match has cost larger than some threshold value c. instead of matching we create new global neuron from the corresponding local one. We also want modest size global model and therefore penalize its size with some increasing function f'(E'). This intuition is formalized in the following extended maximum bipartite matching formulation:
L+by Ly
min On has
ΣΣ
T C s.t. Σ -I1 Σ न E {0, 1) V i, where
j=1
1
c/ ce f(i),
i L L<i Lyr
The size of the new global model is then L' = max{i: म%( = 1 1- I,..,, Ly)., We note some technical details: after the optimization is done, each corresponding II? is of size L, x L and is not a permutation matrix in a classical sense when L3 # L. Its functionality is however similar: taking matrix product with weight matrix IIT implies permuting the weights to align with weights learned on the other datasets and padding with "dummy" neurons having zero weights (alternatively can pad weights w((1) first and complete II with missing rows to recover a proper permutation matrix). This "dummy" neurons should also be discounted when taking average. Without loss of generality, in the subsequent presentation we will ignore these technicalities to simplify the notation.
To complete the matched averaging optimization procedure it remains to specify similarity c(;)),, threshold and model size penalty f(.). Yurochkin et al. 2019a blc) studied fusion, i.e. aggregation,
3

---

Published as a conference paper at ICLR 2020
of model parameters in range of applications. The most relevant to our setting is Probabilistic Federated Neural Matching (PFNM) (urochkin et al., 2019b). They arrived at a special case of equation Bo compute maximum posterion estimate (MAP) of their Bayesian nonparametric model based on the Beta-Bernoulli process (BBP) (Thibaux & Jordan 2007), where similarity c(wra, 0,) is the corresponding posterior probability of jth client neuron generated from Gaussian with mean 0g. and and f()) are guided by the Indian Buffet Process prior (Ghahramani & Griffiths 2005). Instead of making heuristic choices, this formulation provides a model-based specification of equation 3 We refer to procedure for solving equation 2 with the setup from Yurochkin et al. (2019b) as BBP-MAP. We note that their PFNM is only applicable to fully connected architectures limiting its practicality. Our matched averaging perspective allows to formulate averaging of widely used architectures such as CNNs and LSTMs as instances of equation 2 and utilize the BBP-MAP as solver.
2.2 PERMUTATION INVARIANCE OF KEY ARCHITECTURES
Before moving onto the convolutional and recurrent architectures, we discuss permutation invariance in deep fully connected networks and corresponding matched averaging approach. We will utilize this as a building block for handling LSTMs and CNN architectures such as VGG (Simonyan &
Zisserman 2014) widely used in practice.
Permutation invariance of deep FCs We extend equation 1 to recursively define deep FC net- work:
Zn = -I=
(4)
where = 1,...,N N is the layer index, Ho is identity indicating non-ambiguity in the ordering of input features x = 2g and NT identity for the same in output classes. Conventionally (-)) is any non-linearity except for ÿ = xx where it is the identity function (or softmax if we want probabilities instead of logits). When N = 2, we recover a single hidden layer variant from equation [ To perform matched averaging of deep FCs obtained from J clients we need to find permutations for every layer of every client. Unfortunately, permutations within any consecutive pair of intermediate layers are coupled leading to NP-hard combinatorial optimization problem. Instead consider recursive (in layers) matched averaging formulation. Suppose we have {[lja-1}, then plugging W into equation 2 we find {T } and move onto next layer. The recursion base for this procedure is {T which we know is an identity permutation for any j.
Permutation invariance of CNNs The key observation in understanding permutation invariance of CNNs is that instead of neurons, channels define the invariance. To be more concrete, let Conv(z: x. W) define convolutional operation on input z with weights W E RC kxx where Cén Coutt are the numbers of input/output channels and w, h are the width and height of the filters. Applying any permutation to the output dimension of the weights and then same permutation to the input channel dimension of the subsequent layer will not change the corresponding CNN's forward pass. Analogous to equation 4 we can write:
In (CConv(En-1, II" W,I,,).
(5)
Note that this formulation permits pooling operations as those act within channels. To apply matched Couut
averaging for the with CNN layer we form inputs to equation 2 as (wit E R RDD . = 1, where D is the flattened Cir x w x h dimension of II. -.n-1 W; Wj.n- This result can be alternatively derived taking the IM2COL perspective. Similar to FCs, we can recursively perform matched averag- ing on deep CNNs. The immediate consequence of our result is the extension of PFNM (Yurochkin al. 2019b) to CNNs. Empirically, see Figure [ we found that this extension performs well on MNIST with simpler CNN architecture such as LeNet (LeCun et al., 1998) (4 layers) and signif- icantly outperforms coordinate weight averaging round FedAvg). However, breaks down for more complex architecture, e.g. VGG-9 (Simonyan & Zissermann 2014) (9 layers), needed to obtain good quality prediction on a more challenging CIFAR-10.
Permutation invariance of LSTMs Permutation invariance in the recurrent architectures is sociated with the ordering of the hidden states. At a first glance it appears similar to fully con- nected architecture, however the important difference is associated with the permutation invariance
4

---

Published as a conference paper at ICLR 2020
of the sidden-to-hiddee weights H R where L is the number of hidden states. In particu- lar. permutation of the hidden states affects both rows and columns of H. Consider a basic RNN h, o H. + W),, where W are the hopdten weights. To account for the permutation invariance of the hidden states, we notice that dimensions of li, should be permuted in the same way for any i, hence
h = o HII + WI).
(6)
To match RNNs, the basic sub-problem is to align hidden-to-hidden weights of two clients with Euclidean similarity, which requires minimizing 11I Hy over permutations II. This is quadratic assignment problem known to be NP-hard (Loiola et al., 2007). Fortunately, the same permutation appears in an already familiar context of input-to-hidden matching of WII. Our matched averaging RNN solution is utilize equation Plltgiinngl input-to- input-to-hidden hidden weights {W,) to find {11 Then federated hidden-to-hidee to--eeee weights are computed as H E, 11HH,111 and imput-to-hidde weights are computed before. We note that Gromov- Wasserstein distance (rromov 2007) from the optimal transport literature corresponds to a similar quadratic assignment problem. It may be possible to incorporate hidden-to- to-hideen weights H into the matching algorithm by exploring connections to approximate algorithms for computing Coomov-Waaserstern barycenter (Peyré et al. 2016). We leave this possibility for future work.
To finalize matched averaging of LSTMs, we discuss several specifics of the architecture. LSTMs have multiple cell states, each having its individual hidden-to-hidder and topu-tt-hiddenn weights. In out matched averaging we stack input-to-hidder weights into SD x L weight matrix (S is the number of cell states; D is input dimension and L is the number of hidden states) when computing the permutation matrices and then average all weights as described previously. LSTMs also often have an embedding layer, which we handle like fully connected layer. Finally, we process deep LSTMs in the recursive manner similar to deep FCs.
2.3 FEDERATED MATCHED AVERAGING (FEDMA) ALGORITHM
Defining the permutation invariance classes of CNNs and LSTMs allows to extend PFNM (Yurochkin et al. 2019b) to these architectures, however our empirical study in Figure [ demon- strates that such extension fails on deep architectures necessary solve more complex tasks. Our results suggest that recursive handling of layers with matched averaging may entail poor overall solution. To alleviate this problem and utilize the strength of matched averaging on "shallow" archi- tectures. we propose the following layer-wise matching scheme. First, data center gathers only the weights of the first layers from the clients and performs one-layer matching described previously to obtain the first layer weights of the federated model. Data center then broadcasts these weights to the clients, which proceed to train all consecutive layers on their datasets, keeping the matched federated layers frozen. This procedure is then repeated up to the last layer for which we conduct a weighted averaging based on the class proportions of data points per client. We summarize our Federated Matched Averaging (FedMA) in Algorithm[] T The FedMA approach requires communi- cation rounds equal to the number of layers in network. In Figure [ we show that with layer-wise wise matching FedMA performs well on the deeper VGG-9 CNN well LSTMs. In the more chal- lenging heterogeneous setting, FedMA outperforms FedAvg, FedProx trained with same number of communication rounds for LeNet and LSTM and for VGG-9) and other baselines, i.c. client individual CNNs and their ensemble.
FedMA with communication We' 've shown that in the heterogeneous data scenario FedMA out- performs other federated learning approaches, however it still lags in performance behind the entire data training. Of course the entire data training not possible under the federated learning con- straints, but it serves performance upper bound we should strive to achieve. To further improve the performance of our method, we propose FedMA with communication, where local clients receive the matched global model at the beginning of new round and reconstruct their local models with the size equal to the original local models (e.g. size of a VGG-9) based on the matching results of the previous round. This procedure allows to keep the size of the global model small in contrast to a naive strategy of utilizing full matched global model as a starting point across clients on every round.
5

---

Published as a conference paper at ICLR 2020
Algorithm 1: Federated Matched Averaging (FedMA)
Input : local weights of N-layer architectures (W,1. W 1 from J clients
Output: global weights (Wi.... Wx} 1;
while S N do
if n < N then
// call MBP to solve Eq.
n = E, W
else
Wn = Ek=1 E, Pyk sl,n where pkc is fraction of data points with label on worker j:
end
for
jE }l....,JJo do
W W5n++ < W.+1:
// permutate the next-layer weights
Train W jmn+....... W with W,, frozen;
end
end
3 EXPERIMENTS
We present an empirical study of FedMA with communication and compare it with state-of-the-ar methods i.e. FedAvg (McMahan al. 2017) and FedProx (Sahu al., 2018); analyze the perfor- mance under the growing number of clients and visualize the matching behavior of FedMA to study its interpretability. Our experimental studies are conducted over three real world datasets. Summary information about the datasets and associated models can be found in supplement Table
Experimental Setup
We implemented FedMA and the considered baseline methods in Py PyTch (Paszke al., 2017). We deploy our empirical study under simulated federated learning environ- ment where we treat one centralized node in the distributed cluster as the data center and the other nodes as local clients. All nodes our experiments are deployed on p3.2xlarge instances on Ama- zon EC2. We assume the data center samples all the clients join the training process for every communication round for simplicity.
For the CIFAR-10 dataset, individual image (details provided in the Supplement). We note that we ignore all batch normaliza- tion (Toffe & Szegedy 2015) layers in the VGG architecture and leave it for future work.
data augmentation (random crops, and flips) and normalize each
For CIFAR-10, we considered two data partition strategies to simulate federated learning scenario: (i) homogeneous partition where each local client has approximately equal proportion of each of the classes; (ii) heterogeneous partition for which number of data points and class proportions are unbalanced. We simulated heterogeneous partition into J clients by sampling Ps. (0.5) and allocating a Рk.j proportion of the training instances of class k to local client j. We use the original test
CIFAR-10 10 as our global test set for comparing performance of all methods. For the Shakespeare dataset, we treat each speaking role as a client (Caldas et al.) 2018) resulting in natural heterogeneous partition. We preprocess the Shakespeare dataset by filtering out the clients with less than 10k datapoints and sampling random subset of J = 66 clients. We allocate 80% of the data for training and amalgamate the remaining data into a global test set.
Communication Efficiency and Convergence Rate In this experiment we study performance of FedMA with communication. Our goal is to compare our method to FedAvg and FedProx in terms of the total message size exchanged between data center and clients (in Gigabytes) and the number of communication rounds (recall that completing one FedMA pass requires number of rounds equal to the number of layers in the local models) needed for the global model to achieve good perfor- mance the test data. We also compare to the performance of ensemble method. We evaluate all methods under the heterogencous federated learning scenario on CIFAR-10 10 with 16 clients with VGG-9 local models and Shakespeare dataset with J 66 clients with l-aerr L.STM net-

---

Published as a conference paper at ICLR 2020
0.5
0.5
0.4
0.4
Test Acc
Test Acc
FedAvg FedProx
FedAvg
0.3
0.3
FedProx
FedMA
FedMA
0.2
Ensemble Data Training
0.2
Ensemble
Entire Data
0.1
0.1
Training
0.2 044
96 0g 2. 222 2.4 2.6 Amount of Comm. (GB)
0
5
20
25
30
Round Comm.
(a) LSTM, Shakespeare; message size
(b) LSTM, Shakespeare; rounds
0.9
0.90
0.8
0.85
Test Acc
0.80
0.5
0.6
FedAvg
0.75
FedProx
FedAvg
FedProx
0.70
FedMA
FedMA
0.65
Ensemble
0.3
Ensemble
Entire Data
0.2
Entire Data
0.60
Training
Training
0.55
0
20
40
80
100
Amount Comm. (GB)
Round Comm.
(c) VGG-9, CIFAR- 10; message size
(d) VGG-9, CIFAR 10; rounds
Figure 2: Convergence rates various methods two federated learning scenarios: training VGG-9 on CIFAR-10 with J = 16 clients and training LSTM on Shakespeare dataset with J = 66 clients.
work. We fix the total rounds of communication allowed for FedMA, FedAvg, and FedProx í.e. 11 rounds for FedMA and 99/33 rounds for FedAvg and FedProx for the VGG-9/LSTM experiments respectively. We notice that number of local training epochs is common parameter shared by the three considered methods, presented in the next experiment) and report the convergence rate under E that yields the best final model accuracy over the global FedProx i.e. the coefficient, f associated with the proxy term, we also tune the parameter using grid search and report the best μ we found (0.001) for both VGG-9 and LSTM experiments. FedMA out- performs FedAvg and FedProx in all scenarios (Figure with its advantage especially pronounced when we evaluate convergence as a function of the message size Figures 2(a) and 2(c) Final performance of all trained models is summarized in Tables and
thus tune this parameter (denoted F; comprehensive analysis will be
set. We also notice that there another hyper-parameter in
Effeet of local training epochs As studied in previous work (McMahan et al. 2017; Caldas et al. 2018; Sahu et al. 2018). the number of local training epochs E can affect the performance of FedAvg and sometimes lead to diver- gence. We conduct experimental study on the effect of E over FedAvg, FedProx, and FedMA on VGG-9 trained on CIFAR-10 under heterogeneous setup. The candidate local epochs we consider are E' E {10, 20, 50, 70, 100, 150). For each of the candidate F, we run FedMA for rounds while FedAvg and FedProx for 54 rounds and report the final ac- curacy that each methods achieves. The result is shown Figure We observe that training longer benefits FedMA, supporting our assumption that FedMA performs best on lo- cal models with higher quality. For FedAvg. longer local training leads to deterioration of the final accuracy, which matches the observations made the previous literature (McMahan et al. 2017 Caldas al. 2018 Sahu al. 2018). FedProx only partially alleviates this problem. The result
0.86
0.84
0.83
FedAvg
FedProx
FedMA
0.81
0.80
20
40
60
80
100
12g
240
Local Epoch
Figure 3: The effect of number of local training epochs on various methods.

---

Published as a conference paper at ICLR 2020
of this experiment suggests that FedMA is the only method that local clients use train their model as long as they want.
Table I: Trained models summary for VGG-9 trained on CIFAR- as shown Figure
Method
FedAvg
FedProx
Ensemble
FedMA
Final Accuracy(%)
86.29
85.32
75.29
87.53
Best local epoch(E)
20
20
N/A
150
Model growth rate
1x
16x
1.11x
Table 2: Trained models summary for LSTM trained Shakespeare as shown in Figure
Method
FedAvg
FedProx
Ensemble
FedMA
Final Accuracy(%))
46.63
45.83
46.06
49.07
Best local epoch(E)
5
N/A
5
Model growth rate
1x
1x
66x
1.06x
Handling data bias Real world data often exhibit multimodality within each class, e.g. geo- diversity. has been shown that an observable amerocentric and eurocentric bias is present in the widely used ImageNet dataset (Shankar al. 2017 [uusakkvsky al., 2015). Classifiers trained on such data "learn" these biases and perform poorly on the under-represented domains (modalities) since correlation between the corresponding dominating domain and class can prevent the classifier from learning meaningful relations between features and classes. For example, classifier trained on amerocentric and eurocentric data may learn associate white color dress with "bride" class, therefore underperforming on the wedding images taken in countries where wedding traditions are different (Doshil 2018).
Local Acc:Cllento
Entire Data Training
The data bias scenario is an important aspect of federated learning, however it received little to no attention in the prior federated learning works. In this study we argue that FedMA can handle this type of problem. If we view each domain, e.g. geographic region, one client, local models will not be affected by the aggregate data biases and learn meaning- ful relations between features and classes. FedMA can then
Acc:Client
Oversampling
Ensemble
FedMA
FedAvg
Color Balanced
57.8 80.3
po.5 $3.6
80
Testset
60
2.1.6
40
be used to learn a good global model without biases. We have already demonstrated strong performance of FedMA federated learning problems with heterogeneous data across clients and this scenario is very similar. To verify this con- jecture we conduct the following experiment.
Methods
Figure 4:
Performance on skewed
CIFAR- dataset.
We simulate the skewed domain problem with CIFAR-10 dataset by randomly selecting classes and making 95% training images in those classes to be grayscale. For the remaining 5 we turn only 5% of the corresponding images into grayscale. By doing so, we create 5 grayscale images dominated classes and 5 colored images dominated classes. In the test set, there is half grayscale and half colored images for each class. We anticipate entire data training to pick up the uninformative correlations between greyscale and certain classes, leading to poor test performance without these correlations. In Figure we see that entire data training performs poorly in comparison to the regular (i.e. No Bias) training and testing on CIFAR- 10 dataset without any grayscaling. This experiment was motivated by Olga Russakovsky/ s talk at ICML 2019.
Next compare the federated learning based approaches. We split the images from color domi- nated classes and grayscale dominated classes into clients. We then conduct FedMA with com- munication, FedAvg, and FedProx with these clients. FedMA noticeably outperforms the entire
8

---

Published as a conference paper at ICLR 2020
Raw inputl
FedMA filter 0
(c) Clientl, filter 0
(d) Client2, filter 23
(c) FedAvg filter 0
(f) Raw input2
(g) FedMA filter 0 (h) Client filterO (i) Client 2, filter23
( FedAvg filterO
Figure 5: Representations generated by the first convolution layers of locally trained models, FedMA global model and the FedAvg global model.
data training and other federated learning approach as shown in Figure 4 This result suggests that FedMA may interest beyond learning under the federated learning constraints, where entire data training is the performance upper bound, but also to eliminate data biases and outperform entire data training.
We consider two additional approaches to eliminate data bias without the federated learning straints. One way to alleviate data bias to selectively collect more data to debias the dataset. In the context of our experiment, this means getting more colored images for grayscale dominated classes and more grayscale images for color dominated classes. We simulate this scenario by simply doing full data training where each class in both train and test images has equal amount of grayscale and color images. This procedure, Color Balanced, performs well, but selective collection of new data practice may expensive even not possible. Instead of collecting new data, one may consider oversampling from the available data to debias. In Oversampling, we sample the underrepresented domain (via sampling with replacement) make the proportion color and grayscale images to be equal for each class (oversampled images are also passed through the data augmentation pipeline, e.g. random flipping and cropping, to further enforce the data diversity). Such procedure may be prone to overfitting the oversampled images and we see that this approach only provides marginal improvement of the model accuracy compared to centralized training over the skewed dataset and performs noticeably worse than FedMA.
Data efficiency It is known that deep learning models per- form better when more training data is available. However, under the federated learning constraints, data efficiency has not been studied to the best of our knowledge. The challenge here is that when each bring their own version of the data distribution, which,
Local Aug Accs
9-Round FadAvg
Ensemble
Round FedProx
FodMA
Entire Data training
Testset Accuracy
80
clients join the federated system, they
70
60
if not handled properly, may deteriorate the performance de- spite the growing data size across the clients. To simulate
50
this scenario we first partition the entire training CIFAR-10 -0
dataset into 5 homogeneous pieces. We then partition each homogeneous data piece further into sub-pieces hetero- geneously. Using this strategy, we partition the CIFAR-10 training set into 25 heterogeneous small sub-datasets con- taining approximately 2k points each. We conduct a 5-step experimental study: starting from randomly selected homogeneous piece consisting of 5 associated heterogeneous sub-pieces, we simulate 5-client federated learning heterogeneous problem. For each consecutive step, add one of the remaining homogeneous data pieces consisting of 5 new clients with heterogeneous sub- datasets. Results presented in Figure Performance of FedMA (with single pass) improves
Number Clients
10
15
20
25
Figure 6: Data efficiency under the in- creasing number of clients.

---

Published as a conference paper at ICLR 2020
when new clients added to the federated learning system, while FedAvg with communication rounds deteriorates.
Interpretability One of the strengths of FedMA is that utilizes communication rounds more ef- ficiently than FedAvg. Instead of directly averaging weights element wise, FedMA identifies match- ing groups of convolutional filters and then averages them into the global convolutional filters. It's natural to ask "How does the matched filters look like?". In Figure [ we visualize the represen- tations generated by pair of matched local filters, aggregated global filter, and the filter returned by the FedAvg method over the same input image. Matched filters and the global filter found with FedMA are extracting the same feature of the input image, i.e. filter of client I and filter 23 of client are extracting the position of the legs of the horse, and the corresponding matched global filter 0 does the same. For the FedAvg, global filter 0 is the average of filter 0 of client I and filter O of client 2, which clearly tampers the leg extraction functionality of filter 0 of client 1.
4 CONCLUSION
We presented Federated Matched Averaging (FedMA), a layer-wise federated learning algorithm designed for modern CNNs and LSTMs architectures that accounts for permutation invariance of the neurons and permits global model size adaptation. Our method significantly outperforms prior federated learning algorithms in terms of its convergence when measured by the size of messages exchanged between server and the clients during training. We demonstrated that FedMA can effi- ciently utilize well-trained local modals, a property desired in many federated learning applications, but lacking in the prior approaches. We have also presented example where FedMA can help to resolve some the data biases and outperform aggregate data training.
In the future work want extend FedMA to improve federated learning of LSTMs using approxi- mate quadratic assignment solutions from the optimal transport literature, and enable additional deep learning building blocks, e.g. residual connections and batch normalization layers. We also believe it is important to explore fault tolerance of FedMA and study its performance on the larger datasets, particularly ones with biases preventing efficient training even when the data be aggregated, e.g. Inclusive Images (Dohhi 2018).
REFERENCES
Martial Aguch and Guillaume Carlier. Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis, 43(2):904-924, 2011.
Keith Bonawitz, Hubert Eichner, Wolfgang Gricskamp. Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloe Kiddon, Jakub Konecny, Stefano Mazzocchi, H Brendan McMahan, et al. Towards federated learning at scale: System design. arXiv preprint
2019.
Sebastian Caldas, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv:1812.01099 2018.
Tulsee Doshi. Introducing the inclusive images competition, 2018.
Zoubin Ghahramani and Thomas L Griffiths. Infinite latent feature models and the Indian buffet process. In Advances in Neural Information Processing Systems, pp. 475-482, 2005.
Mikhail Gromov. Metric structures for Riemannian and non non-Riemannian spaces. Springer Science & Business Media, 2007.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv: 2015.
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 arXiv: arXiv: 1912, 2019.
10

---

Published as a conference paper at ICLR 2020
Harold W Kuhn. The Hungarian method for the assignment problem. Naval Research Logisties (NRL), 2(1-)):83-97 1955.
Yann LeCun, Léon Bottou, Yoshua Bengio, Patrick Haffner, et al. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(1 1):2278-222 2324, 1998.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. arXiv preprint 1rXv::1908.07773 2019.
Eliane Maria Loiola, Nair Maria Maia Abreu, Paulo Oswaldo Boaventura-Neit Peter Hahn, and Tania Querido. A survey for the quadratic assignment problem. European journal of operational research, 176(2)657-690. 2007.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communicationn-effiienn learning of deep networks from decentralized data. In Artificial Intelli- gence and Statistics, pp. 1273-1282, 2017.
Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. In Interna- fional Conference on Machine Learning, pp. 4615 4615-4625, 2019.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. In NIPS-W: W, 2017.
Gabriel Peyré, Marco Cuturi, and Justin Solomon. Grosov-wasserstein averaging of kernel and distance matrices. In International Conference on Machine Learning, pp. 2664-2672, 2016.
Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence adam and beyond. In International Conference Learning Representations, 2018. URL http::ttt////eeeeeeeie
net//oorum:1===yya722
Olga Russakovsky. Strategies for mitigating social bias in deep learning systems. Invited talk at Identifying and Understanding Deep Learning Phenomena workshop, ICML 2019.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang. Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International journal of computer vision, 1123)22 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 preprins arXiv: 18/2.0661,. 2018.
Shreya Shankar, Yoni Halpern, Eric Breck, James Atwood, Jimbo Wilson, and D Sculley. No classi- fication without representation: Assessing geodiversity issues in open data sets for the developing world. arXiv preprint arXiv: 2017.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv: 1556, 2014.
Sidak Pal Singh and Martin Jaggi. 1rXiv:1910.0550 2019.
Model fusion via optimal transport.
arXiv preprint
Virginia Smith, Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar. Federated multi-task learning. Advances in Neural Information Processing Systems, pp. 4424 4434, 2017.
Romain Thibaux and Michael Jordan. Hierarchical Beta processes and the Indian buffet process. In Artificial Intelligence and Statisties, pp. 564-571, 2007.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, and Nghia Hoang. Sta- tistical model aggregation via parameter matching. Advances in Neural Information Processing Systems, pp.
2019a.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. Interna- fional Conference Machine Learning, pp. 7252-7261, 2019b.

---

Published as a conference paper at ICLR 2020
Mikhail Yurochkin, Zhiwei Fan, Aritra Guha, Paraschos Koutris, and XuanLong Nguyen. Scalable inference of topic evolution via models for latent geometric structures. In Advances in Neural Information Processing Systems, pp. 5949-5959, 2019c.
12

---

Published as a conference paper at ICLR 2020
A SUMMARY OF THE DATASETS USED IN THE EXPERIMENTS
The details of the datasets and eyper-parameters used our experiments are summarized in Table 3 In conducting the "freezing and retraining" process of FedMA, we notice when retraining the last FC layer while keeping all previous layers frozen, the initial learning rate we use for SGD doesn't t lead to good convergence (this is only for the VGG-9 architecture). To fix this issue, divide the initial learning rate by 10 ie. using 10-4 for the last FC layer retraining and allow the clients to retrain for times more epochs. We also switch off the l weight decay during the "freezing and retraining" process of FedMA except for the last FC layer where we use l2 weight decay of 10-4, For language task, we observe SGD with constant learning rate works well for all considered methods.
In our experiments, we use FedAvg and FedProx variants without the shared initialization since those would likely be more realistic when trying to aggregate locally pre-trained models. And FedMA still performs well in practical scenarios where local clients won't 't be able to share the random initialization.
Table The datasets used and their associated learning models and hyper- paper--arameters.
Method
MNIST
CIFAR-10
Shakespeare
McMMaan al. 2017)
# Data points
60, 000
50, 500
1, 0171 981
Model
LeNet
VGG-9
LSTM
Classes
10
10
80
# Parameters
431k
3, 491k
293k
Optimizer
SGD
SGD)
Hyper-params.
Init Ir: 0.01, 0.001 (last layer)
Ir: 0.8(const)
momentum: 0.9, la weight decay: 10-4
B DETAILS OF MODEL ARCHITECTURES AND HYPER-PARAMETERS
The details of the model architectures we used in the experiments Specifically, details of the VGG-9 model architecture we used can be found in Table details of the l-layer LSTM model used in our experimental study can be found in Table
summarized in this section.
C DATA AUGMENTATION AND NORMALIZATION DETAILS
In preprocessing the images CIFAR-10 dataset, we follow the standard data augmentation and normalization process. For data augmentation, random cropping and horizontal random flipping are used. Each color channels are normalized with mean and standard deviation by [4, 0.491372549, 1ig 0.482352941, 14 0.446666667, ar 0.247058824, og 0.243529412, a, = 0.261568621 Hach channel pixel is normalized by subtracting the mean value in this color channel and then divided by the standard deviation of this color channel.
D EXTRA EXPERIMENTAL DETAILS
D. SIIAPES OF FINAL GLOBAL MODEL
Here we report the shapes of final global VGG and LSTM models returned by FedMA with com- munication.
13

---

Published as a conference paper at ICLR 2020
Table 4:
Detailed information of the VGG-9 architecture used in experiments, all non- linear activation function in this architecture is ReLU: the shapes for convolution layers follows (Ci., Couts C,
Parameter
Shape
Layer hyper-parameter
layerI.convL..we Tweight
3 x 32 x 3 x 3
stride:1:padding:1
layerI.convl.bias
32
N/A
layer2.conv2.weight
32 x 64 x 3
stride: :paddin::.
layer2.conv2.bias
64
N/A
pooling.max
N/A
kernel 2ize:2;t
layer3.conv3.weight
64 X 128 x 3 x 3
stride:11padding:11
layer3.conv3.bias
128
N/A
layer4.conv4..eight
128 x 128 3 x 3
stride: 1 :padding:1
layer4.conv4.bias
128
N/A
pooling.max
N/A
kemel size:2:strid::2
dropout
N/A
p 5%
layers.convs.weight
128 x 256 3 x 3
stride:1:
layer5.conv5.bias
256
N/A
layer6.conv6.weigh
256 x 256 3 x 3
stride:1 :paddingg
layer6.conv6.bias
256
N/A
pooling.max
N/A
kerne size:2:stride:2
dropout
N/A
p 10%
leyer
4096 x 512
N/A
layer7.fe7.bias
512
N/A
wtyrr8.fc8.weight
512 x 512
N/A
layer8.fc8.bias
512
N/A
dropout
N/A
p 10%
weiigtt
512 x 10
N/A
layer9.fc9.bias
10
N/A
Table 5: Detailed information of the LSTM architecture in our experiment
Parameter
Shape
eniodtr.weight
80
Istm.weighttih.10
1024 8
wsti.weigtt.hb.10
1024 x 256
Istm.bias.ih.10
1024
Istm.bias.hh.10
1024
wecoder.weight
80 256
decoder.bias
80
14

---

Published as a conference paper at ICLR 2020
Table 6: Detailed information of the LSTM architecture in experiment
Parameter
Shape
Growth rate (#global #original params)
encoder.weight
80 x 21
2.63 x (1, 680/640)
lstm.weight.ih.10
1028 x 21
2.64 x (21,588/8 588/8, 192)
lstm.weight.hh.00
1028 x 257
1.01 x (264, 196/ 196/262, 144)
Istm.bias.ih.10
1028
1.004 (1, 028/ 024)
Istm.bias.hh.10
1028
1.004 11,08/11 024)
decoder.weight
80 x 257
1.004 x (20,560/22) 560/20,00000 480)
decoder.bias
80
1x
Total Number of Parameters 310, 160
1.06x (310, (6600293,550 584)
Table 7: Detailed information of the final global VGG-9 model returned by FRB; the shapes for convolution layers follows (Cins Cout, C, c)
Parameter
Shape
Growth rate (#global #original params)
layerI.convI..wiigh
x 47 x x
1.47 x (1, 269/864)
layerl.convl.bias
47
1.47 x (47/32)
layer2.comv2.weightt
47 x 79 X 3 x 3
1.81 X (33,
layer2.conv2.bias
79
1.23 x (79/64)
layer3.comv3.weeigtt
79 X 143 3 3
1.38 x (101, (01,,63/77 673/73, 728)
layer3.conv3.bias
143
1.12 (143/ (281))
layeri.conv4.weight
143 143 x x
1.24 x (184, 041/ /147, 45,4
layer4.conv4.bias
143
1.12 x (143/128)
layer5.conv5.weight
143 x 271 x 3 x 3
1.18 x (348,777/ 177/294,9112
layerS.conv5.bias
271
1.06 x (271/256)
layer6.conv6.weighh
271 271 x3x
1.12 x (660, 969/589, 824)
layer6.conv6.bias
271
1.06 (271/256)
wayert
4336 x 527
1.09 x (2, 285, 072/2 097, 152)
layer7.fc7.bias
527
1.02 x (527/512)
weiggt
527 x 527
1.05 (277, 7117,,79/221 729/262, 144)
layer8.fc8..r
527
1.02 x (527/512)
layer9.fe9.weight
527 x
1.02 x (5, 150/7 120)
layer9.fc9.bias
10
1x
Total Number of Parameters
3, 90900,222 235
1.11x (3, 900, 235/3, 231/3,491, 530)
D.2
FOR MBP-MAP
We follow FPNM Yurochkin et al. 2019b) to choose the hyper-parameters of BBP-MAP, which controls the choices of 0,, e, and the f(-) as discussed in Section More specifically, there are three parameters to choose i.e. 1) o the prior variance of weights of the global neural network; 2) Yuo which controls discovery of new hidden states. Increasing Yo leads to a larger final global model; 3) o² is the variance of the local neural network weights around corresponding global network weights. We empirically analyze the different choices of the three parammearsmeters and find the choice of
15

---

Published as a conference paper at ICLR 2020
Yu 0 1, for VGG-9 C1 dataset and To 1 o0 1, 1 for LSTM on Shakespeare dataset lead to good performance our experimental studies.
E PRACTICAL CONSIDERATIONS
Following from the discussion PFNM, here we briefly discuss the time complexity of FedMA. For simplicity, focus single-layer matching and assume all participating clients train the same model architecture. The complexity for matching the entire model follows trivially from this discussion. The worst complexity is achieved when no hidden states matched and is equal to O(D (JL)²) for building the cost matrix and ({)L))) for running the Hungarian algorithm where the definitions of D, and follow the discussion Section 2 The best complexity per layer is (achieved when all hidden states are matched) O(D- +²). Practically, when the number of participating clients i.e. is large and each client trains big model, the speed our algorithm can be relatively slow.
To seed up the Hungarian algorithm. Although there isn't any algorithm that achieves lower com- plexity, better implementation improves the constant significantly. In our experiments, we used an implementation based on shortest path augmentation i.e. lapso veer P Empirically, we observed that this implementation of the Hungarian algorithm leads to orders of magnitude speed ups over the vanilla implementation.
F EYPER-PARAMETEE FOR THE HANDLING DATA BIAS EXPERIMENTS
conducting the "handling data bias" experiments. We re-tune the local epoch E for both FedAvg and FedProx. The considered candidates of E are {5, 10, 20, 30}. We observe that relatively large choice of E can easily lead to poor convergence choices of E better, smaller E can always lead to good convergence. We use F. 5 for both FedAvg and FedProx in our experiments. For FedMA, we choose = 50 since it leads to a good convergence. For the "oversampling" baseline, found that using SGD train VGG-9 over oversampled dataset doesn't lead to good convergence, Moreover, when using constant learning rate, SGD can lead model divergence. Thus we use AMSGRAD (Reddi al. 2018) method for the "oversampling" baseline and train for 800 epochs. To make the comparison fair, we use AMSGRAD for all other centralized baselines get the reported results our experiments. Most of them converges when training for 200 epochs. We also test the performance of the "Entire Data Training". "Color Balanced"' and "No Bias" baselines over SGDD. We use learning rate 0.01 and & weight decay at 10-4 and train for 200 epochs for those three baselines. It seems the "Entire Data Training" " and "No Bias" baselines converges to slightly better accuracy i.e. 78.71% and 91.23% respectively (compared to 75.91% and 88.5% for AMSGRAD). But the "Color Balanced" doesn't seem to converge better accuracy (we get 87. 79% accuracy for SGD and 87.81% for AMSGRAD).
FedAvg. While FedProx tolerates larger
16


Reviews:
[{"review_id": "rkgJ17ZpFH", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Post Rebuttal Summary\n---------------------------------\nI have nudged my score up to an \"Accept\", based on my comments to the rebuttal below. I hope the authors continue to improve the readability of Sec. 2.1\n\nReview Summary\n--------------\nOverall I think this is almost above the bar to be accepted, and I could be persuaded with a strong rebuttal.  The strengths here are the extensive experiments and the easy-to-implement method. The primary weakness of this paper is that it is a \"straightfoward\" way to extend the BBP-MAP method to CNNs and RNNs, so the methodological novelty is weak relative to the BBP-MAP past work (Yurochkin et al. ICML 2019). Other technical weaknesses limit the ability to use this method on clients with diverse class distributions, which will be common in real deployments.\n\nPaper Summary\n-------------\nThis paper addresses the problem of federated learning, where J separate \"clients\" with disjoint datasets each train a neural network model for a supervised problem, and then try to aggregate all J individual client models into one \"global model\" in a coherent way. The natural problem is that due to hidden units being permutable within one network, naively taking parameter averages across two client models will lead to bad accuracy without first coming up with a consistent ordering of the units in each layer. \n\nPrevious work (Yurochkin et al. ICML 2019) has developed a Bayesian nonparametric model based on the Beta-Bernoulli Process (BBP) for the case of federated learning of multi-layer perceptrons. However, the extension to convolutional layers or recurrent layers has yet to be solved, which is the focus of this paper. \n\nThis paper's algorithm (Federated Matched Averaging (FedMA), see Alg 1), proceeds by iteratively stepping thru the CNN or RNN layer by layer greedily from input to output. At each layer, we first solve a BBP-MAP optimization (bipartite matching using a BBP maximum a-posteriori objective as cost function, a subprocedure taken direclty from Yurochkin et al.). This obtains a consistent low-cost permutation for each client model. Then, the global model weights for that layer is the average of the aligned client weights. After the current layer update, each client keeps training, keeping all layers up to the current frozen but revising later layers. This layer-by-layer training can be applied to both CNNs and RNNs.\n\nThe proposed approach is compared to FedAvg and FedProx on MNIST and CIFAR image classification tasks with CNNs, and Shakespeare text classification tasks with RNNs. Later experiments explore the effect of communication efficiency (MB transfered between client and master), effect of local training epochs, handling biased class distributions, and interpretabilty.\n\n\n\n\nNovelty & Significance\n-----------------------\nSolving federated learning problems is of increasing practical importance, and certainly trying to do so for CNNs and RNNs (more than just large MLPs) is important. So I like where the paper is going.\n\nAlthough the method is \"new\", it is more or less a straightforward extension of work by Yurochkin et al. (ICML 2019) to CNNs and RNNs. If you read the last few sentences of Yurochkin et al., you'll see \"Finally, it is of interest to extend our model-ing framework to other architectures such as Convolutional Neural Networks (CNNs) and Recurrent Neural Networks (RNNs). The permutation invariance necessitating matching inference also arises in CNNs since any permutation of the filters results in the same output, however additional bookkeeping is needed due to the pooling operations.\" I view this paper as a well-executed implementation of this \"bookkeeping\". Certainly not trivial, but to some readers perhaps not clearly \"above the bar\" for a top conference like ICLR.\n\n\nTechnical Concerns\n------------------\n\n## Concern 1: Client models will not always be alignable after permutation\n\nMy first concern is that there will not always be a one-to-one permutation of the neurons learned by two client models with different class distributions. Given fixed capacity at each layer, some clients may learn a filter for \"horse hooves\" (esp. if horse images are common to that client), while other clients may learn a filter for \"snake skin\" (if snakes are more common to that client). I wonder if we can quantify how well the aligned filters match in practice, and if there is any benefit to revising the alignment to allow some client-specific customizations (e.g. by having the global model can learn more units than the client model). \n\n## Concern 2: Use of the BBP-MAP subprocedure poorly motivated\n\nThe paper prioritizes a clean and easy-to-implement algorithm to resolve practical alignment issues between client CNN and RNN models. However, I was a bit underwhelmed that the BBP-MAP solution used by Yurochkin et al. was treated as a black-box subprocedure without much justification. I could see 2 preferable alternatives to the current use of BBP-MAP. Either a simpler approach using Eq. 2 with a squared error cost and the Munkres algorithm to solve bipartitite matching to obtain the permutation (which seems more in spirit of the rest of the paper). Or, a more sophisticated probabilistic approach (taking a Bayesian hierarchical model from Yurochkin et al. seriously and forming the estimated global weights from a weighted sums that includes both the clients (weighted by dataset size) and the assumed prior). As it is, I feel the BBP-MAP subprocedure in the current Algorithm 1 is poorly motivated for the task at hand.\n\n\n\nExperimental Evaluation\n-----------------------\n\nOverall the experiments were extensive and demonstrated several apparent advantages (reduced need to transfer large memory during communication, etc.). \n\n\nMinor Presentation Concerns\n---------------------\nBefore Eq. 2, you should introduce the \"\\theta\" notation\n\nI'm a bit confused about how \"FedMA\" differs from \"FedMA with communication\", even after reading Sec. 2.3. How exactly are communicate costs kept down? What are you sending from master to client at beginning of every \"round\" if not the full global model (all weights of the CNN)?", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "HJlE_doTKH", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Edit: Thanks for the thorough and responsive rebuttal! I'm particularly happy to see the additional background on BBP-MAP and the baselines you've added for handling data bias. You've comprehensively addressed my questions and I think this paper should be accepted.\n\nOriginal review:\n\nThe authors extend the recently proposed Probabilistic Federated Neural Matching (PFNM) algorithm of Yurochkin et al. ( 2019) to more kinds of neural networks, show that it isn't as effective for larger models as it is for LeNet-sized ones, and propose enhancements that lead to a state-of-the-art approach they call FedMA. I'm convinced that this represents a meaningful advance in federated learning, although the paper could use some tightening up, and the experiments are somewhat limited.\n\nSome feedback:\n\n- I'd like to see a little bit more description of BBP-MAP, as even though it's not one of the components of the algorithm you directly modify it's still the underlying mathematical primitive. How far is it from having the same effect that the \"best possible\" permutation would? How is it able to allow the number of neurons in the federated model to grow relative to the size of the client models?\n\n- Can you include the \"entire data\" baseline in more of the figures/plots (especially Figure 2)?\n\n- The models and datasets covered in the experiments are adequate to demonstrate that the presented technique is worth exploring, but probably not for someone considering applying it in the context of a deployed federated learning application. Since federated learning is a problem domain motivated more by applied concerns (privacy, edge vs. cloud compute, on-device ML) than other areas of machine learning theory, it would be particularly valuable to see experiments at larger scale (in particular, on larger or more realistic datasets).\n\n- The section that demonstrates how your model addresses skewed data domains is fascinating! That's one area in which your experiments are directly relevant to federated learning in practice, and it's a rapidly growing area of research in itself (e.g. in its relationship to causal learning that Leon Bottou has recently been exploring). Exploring this further could make for a whole separate paper. In the mean time, though, is there some kind of equivalent of the \"entire data\" baseline that would represent e.g. the best known technique for taking into account skewed domains outside the federated context?", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "BygfWezloB", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper offers a beautiful and simple method for federated learning. Strong empirical results. \n\nImportant area. \n                                                                                                                                                                                                                                                                                                                                                                                                .", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors presented a Federate Learning algorithm which constructs the global model layer-wise by matching and averaging hidden representations. They empirically demonstrate their method outperforms existing federated learning algorithms

This paper has received largely positive reviews. Unfortunately one reviewer wrote a very short review but was generally appreciative of the work. Fortunately, R1 wrote a detailed review with very specific questions and suggestions. The authors have addresses most of the concerns of the reviewers and I have no hesitation in recommending that this paper should be accepted. I request the authors to incorporate all suggestions made by the reviewers. 
Decision: Accept (Talk)