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 outperforms popular state-of-the-art federated learning algorithms on deep CNN and LSTM architectures trained on real world datasets, while improving the communication efficiency. Edge devices such as mobile phones, sensor networks or vehicles have access to a wealth of data. However, due to concerns raised by data privacy, network bandwidth limitation, and device availability, it's unpractical to gather all local data to the data center and conduct centralized training. To address these concerns, federated learning is emerging (McMahan et al., 2017; Li et al., 2019; Smith et al., 2017; Caldas et al., 2018; Bonawitz et al., 2019) to allow local clients to collaboratively train a shared global model. The typical federated learning paradigm involves two stages: (i) clients train models over their datasets independently (ii) the data center uploads their locally trained models. The data center then aggregates the received models into a shared global model. One of the standard aggregation methods is FedAvg (McMahan et al., 2017) where parameters of local models are averaged element-wise with weights proportional to sizes of client datasets. FedProx (Sahu et al., 2018 ) adds a proximal term for client local cost functions, which limits the impact of local updates by restricting them to be close to the global model. Agnostic Federated Learning (AFL) (Mohri et al., 2019) , as another variant of FedAvg, optimizes a centralized distribution that is formed by a mixture of the client distributions. One shortcoming of the FedAvg algorithm is that coordinate-wise averaging of weights may have drastic detrimental effect on the performance and hence hinders the communication efficiency. This issue arises due to the permutation invariant nature of the neural network (NN) parameters, i.e. for any given NN there are many variations of it that only differ in the ordering of parameters and constitute local optima which are practically equivalent. Probabilistic Federated Neural Matching (PFNM) (Yurochkin et al., 2019) addresses this problem by finding permutation of the parameters of the NNs before averaging them. PFNM further utilizes Bayesian nonparametric machinery to adapt global model size to heterogeneity of the data. As a result, PFNM has better performance and communication efficiency, however it was only developed for fully connected NNs and tested on simple architectures. Our contribution In this work (i) we demonstrate how PFNM can be applied to CNNs and LSTMs, however we find that it gives very minor improvement over weight averaging when applied to modern deep neural network architectures; (ii) we propose Federated Matched Averaging (FedMA), a new layers-wise federated learning algorithm for modern CNNs and LSTMs utilizing matching and model size adaptation underpinnings of PFNM; (iii) We empirically study FedMA with real datasets under the federated learning constraints. In this paper, we presented FedMA, a new layer-wise federated learning algorithm designed for modern CNNs and LSTMs architectures utilizing probabilistic matching and model size adaptation. We demonstrate the convergence rate and communication efficiency of FedMA empirically. In the future, we would like to extend FedMA towards finding the optimal averaging strategy. Making FedMa support more building blocks e.g. residual structures in CNNs and batch normalization layers is also of interest. Table 4 : Detailed information of the VGG-9 architecture used in our experiments, all non-linear activation function in this architecture is ReLU; the shapes for convolution layers follows (Cin, Cout, c, c) In preprocessing the images in 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 µ r = 0.491372549, µ g = 0.482352941, µ b = 0.446666667, σ r = 0.247058824, σ g = 0.243529412, σ b = 0.261568627. Each channel pixel is normalized by subtracting the mean value in this color channel and then divided by the standard deviation of this color channel.
Communication efficient federated learning with layer-wise matching
{"src": "scitldr-AIC_train", "paper_id": "BkluqlSFDS"}