[Paper Review] Federated Learning with Matched Averaging

Hongyi Wang, Mikhail Yurochkin+3 more|arXiv (Cornell University)|Feb 15, 2020

Privacy-Preserving Technologies in Data18 references102 citations

TL;DR

FedMA is a layer-wise federated learning algorithm for CNNs and LSTMs that aligns and averages matched hidden elements across clients, improving accuracy and reducing communication.

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.

Motivation & Objective

• Motivate federated learning for modern neural architectures while preserving data privacy.
• Address permutation invariance in neural network parameters during aggregation.
• Propose a layer-wise matching and averaging algorithm that adapts global model size to heterogeneous client data.
• Demonstrate improved convergence and communication efficiency over FedAvg and FedProx on real datasets.

Proposed method

• Formulate permutation-invariance-aware parameter averaging via a matched averaging objective that aligns local neurons with global neurons (Equation 2).
• Extend PFNM to CNNs and LSTMs by handling layer-wise matching and adaptive global model size using BBP-MAP as a solver.
• Propose FedMA: a layer-wise algorithm that first matches and averages the first layer across clients, then freezes matched layers while clients train subsequent layers, repeating per layer.
• Use Hungarian algorithm to solve the assignment problem for neuron/channel matching and an adaptive growth strategy for the global model size.
• Incorporate a communication-efficient variant (FedMA with communication) where the global model size is kept small across rounds while reusing matched structures.

Experimental results

Research questions

• RQ1 Can permutation invariance in neural network parameters be effectively addressed in federated learning for CNNs and LSTMs?
• RQ2 Does layer-wise matched averaging (FedMA) improve accuracy and reduce communication compared to FedAvg and FedProx, especially under heterogeneity?
• RQ3 How should global model size adapt to heterogeneous client data in federated settings?
• RQ4 What is the impact of local training epochs on FedMA performance and communication efficiency?

Key findings

• FedMA substantially outperforms FedAvg and FedProx in heterogeneous settings across CNNs and LSTMs.
• Layer-wise matching with adaptive global model size enables better convergence with fewer communications.
• A communication-efficient FedMA variant maintains small global models while achieving strong performance.
• FedMA's layer-wise approach yields interpretability benefits by aligning and averaging matched filters/neurons.
• In experiments, FedMA achieves higher final accuracies than baselines in CIFAR-10 and Shakespeare tasks (CIFAR-10: FedMA 87.53% vs FedAvg 86.29%, FedProx 85.32%, Ensemble 75.29%; Shakespeare: FedMA 49.07% vs FedAvg 46.63%, FedProx 45.83%, Ensemble 46.06%).