PaperlayerView exact revision on arXiv

Source-linked AI summary

Federated Learning with Matched Averaging

Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, Yasaman Khazaeni

arXiv:2002.06440v1cs.LGstat.ML

TL;DR

Federated learning must combine models trained on private, distributed data despite communication constraints and neural-network permutation invariance. The paper proposes FedMA, which matches and averages model elements layer by layer across modern architectures; experiments report stronger convergence and reduced communication burden than prior federated methods. Its scope includes assumptions about heterogeneous feature extractors, and incorporating hidden-to-hidden LSTM weights into matching remains future work.

  • Problem

    Federated learning needs to aggregate independently trained client models without centralizing data, but coordinate-wise averaging can be ineffective under parameter permutation and prior matching methods support limited architectures.

  • Method

    FedMA performs layer-wise matched averaging for CNNs and LSTMs, matching hidden elements by feature-extraction similarity while allowing the global model size to adapt to heterogeneous clients.

  • Results

    FedMA significantly outperforms prior federated algorithms in convergence measured by the size of messages exchanged between server and clients during training.

  • Takeaways & Limitations

    FedMA can efficiently use well-trained local models and, in an evaluated example, help resolve data biases while outperforming aggregate-data training.

  • Takeaways & Limitations

    The method leaves incorporating hidden-to-hidden LSTM weights into the matching algorithm for future work.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Federated learning lets clients collaboratively train a shared model without centralizing their data, addressing privacy, bandwidth, and availability constraints. The introduction motivates FedMA as a response to the limitations of coordinate-wise averaging and prior matching methods on modern architectures.

  • Motivation: Federated learning enables local clients to collaboratively train a shared global model without gathering their data at a central data center.The motivation includes data privacy concerns, limited network bandwidth, and device availability.
  • Existing federated learning methods: FedAvg averages local-model parameters element-wise, weighted by client dataset sizes, while FedProx constrains local updates with a proximal term.Agnostic Federated Learning instead optimizes a centralized mixture of client distributions.
  • Limitation of weight averaging: Coordinate-wise weight averaging can harm performance because neural-network parameters are permutation invariant.Different parameter orderings can represent equivalent networks, making direct averaging inappropriate.
  • Prior matching approach: PFNM matches client neurons before averaging and improves performance and communication efficiency over FedAvg, but it supports only simple architectures.The introduction identifies fully connected feedforward networks as an example of its supported scope.
  • Contribution: FedMA extends matched averaging to modern CNNs and LSTMs through a layer-wise algorithm that adapts global model size to heterogeneous data.The authors report reduced communication burden and improved performance over state-of-the-art federated learning algorithms.

2 FEDERATED MATCHED AVERAGING OF NEURAL NETWORKS

The paper formulates matched averaging to address permutation-invariant neural-network parameters and extends it to deep fully connected, convolutional, and recurrent architectures. FedMA instead matches layers sequentially, adapting global-model size to heterogeneous feature extractors while reducing communication burden.

  • Matched averaging formulation: Neural networks have permutation-equivalent parameterizations, so coordinate-wise averaging can produce weights that do not correspond to any equivalent well-performing network.The paper motivates undoing client-specific permutations before averaging hidden elements.
  • Matched averaging formulation: Matched averaging aligns client neurons with global neurons using a similarity function and optimizes the resulting assignments with iterative Hungarian matching.The formulation can use adaptive global-model size and penalize poor matches or excessive growth.
  • Matched averaging formulation: Under heterogeneous data, the global model size is allowed to range from the largest local model to the concatenation of all local models.This accommodates client-specific feature extractors that overlap only partially across datasets.
  • Architecture extensions: Extending PFNM to CNNs performs well on shallow LeNet trained on MNIST but breaks down on deeper VGG-9 needed for CIFAR-10.The failure motivates a layer-wise rather than recursively assembled matching procedure.
  • Architecture extensions: For recurrent networks, FedMA matches input-to-hidden weights while computing hidden-to-hidden weights after alignment, leaving incorporation of those recurrent weights into matching for future work.Direct hidden-to-hidden alignment is a quadratic assignment problem known to be NP-hard.
  • Federated Matched Averaging: FedMA matches and freezes one layer at a time, broadcasts each matched layer to clients, and trains subsequent layers before repeating through the network.The final layer uses weighted averaging based on client class proportions, and communication rounds equal the number of layers.
  • Federated Matched Averaging: FedMA with communication reconstructs full-sized local models from previous matching results while keeping the global model small across rounds.This variant addresses the gap between FedMA and training on all data while preserving the federated setting’s communication constraints.

3 EXPERIMENTS

Experiments evaluate FedMA against federated baselines across CNN, LSTM, skewed-data, data-efficiency, and interpretability settings. FedMA generally shows stronger convergence, robustness to local training epochs, improved performance with additional clients, and more coherent matched-filter representations.

  • Experimental setup: Experiments compare FedMA with FedAvg and FedProx on three real-world datasets using simulated federated environments.The study includes heterogeneous data, skewed CIFAR-10, growing client counts, convergence, and representation analyses.
  • Communication efficiency and convergence: FedMA outperforms FedAvg and FedProx in all evaluated convergence scenarios, especially when convergence is measured by exchanged message size.The scenarios use VGG-9 on CIFAR-10 with 16 clients and a one-layer LSTM on Shakespeare with 66 clients.
  • Effect of local training epochs: FedMA is the only evaluated method whose clients can train for arbitrarily many local epochs without the reported deterioration or divergence behavior.The experiment evaluates local epochs E ∈ {10, 20, 50, 70, 100, 150} on heterogeneous CIFAR-10 with VGG-9.
  • Handling data bias: On skewed CIFAR-10, FedMA noticeably outperforms entire-data training and other federated approaches.The setup creates grayscale-dominated and color-dominated class groups across two clients, with balanced grayscale and color representation in the test set.
  • Data efficiency: With a single pass, FedMA performance improves as clients are added, whereas FedAvg with 9 communication rounds deteriorates.The data-efficiency experiment partitions CIFAR-10 into homogeneous pieces and then distributes them across clients.
  • Interpretability: FedMA matches convolutional filters before averaging, producing global-filter representations that can be compared with locally trained and FedAvg representations.Figure 5 visualizes matched local filters, the aggregated FedMA filter, and the FedAvg filter on the same input image.

4 CONCLUSION

FedMA is presented as a layer-wise algorithm for modern CNN and LSTM architectures that accounts for neuron permutation invariance and adapts global model size. The paper reports improved communication-efficient convergence, effective use of trained local models, and potential benefits for data-bias settings, while identifying extensions to additional architectures and larger biased datasets.

  • FedMA is a layer-wise federated learning algorithm for modern CNNs and LSTMs that accounts for neuron permutation invariance and permits global model size adaptation.
  • FedMA significantly outperforms prior federated learning algorithms when convergence is measured by the message size exchanged during training.
  • FedMA efficiently utilizes well-trained local models, a property the authors describe as lacking in prior approaches.
  • Future work includes improving LSTM federated learning, supporting residual connections and batch normalization, studying fault tolerance, and evaluating larger biased datasets.

A SUMMARY OF THE DATASETS USED IN THE EXPERIMENTS

The experiments use specified datasets, learning models, and hyper-parameters summarized in Table 3, with practical choices for local-model aggregation and VGG-9 retraining.

  • For VGG-9, the last fully connected layer is retrained with an initial learning rate of 10^-4 after dividing the original rate by 10.
  • The experiments use FedAvg and FedProx variants without shared initialization to reflect aggregation of locally pretrained models.
  • FedMA still performs well when local clients cannot share the random initialization.
  • Dataset, learning-model, and hyper-parameter details are summarized in Table 3.

B DETAILS OF MODEL ARCHITECTURES AND HYPER-PARAMETERS

The experiments use a VGG-9 model and a one-layer LSTM, with their architectural details provided in separate tables.

  • The VGG-9 architecture used in the experiments is detailed in Table 4.
  • The one-layer LSTM architecture used in the experiments is detailed in Table 5.

C DATA AUGMENTATION AND NORMALIZATION DETAILS

CIFAR-10 images are preprocessed with random cropping, horizontal flipping, and per-channel normalization using specified means and standard deviations.

  • CIFAR-10 preprocessing uses random cropping and horizontal random flipping for data augmentation.
  • Each color channel is normalized by subtracting its channel mean and dividing by its channel standard deviation.
  • The channel means are µr = 0.491372549, µg = 0.482352941, and µb = 0.446666667.
  • The channel standard deviations are σr = 0.247058824, σg = 0.243529412, and σb = 0.261568627.

D.1 SHAPES OF FINAL GLOBAL MODEL

This section reports the final global VGG and LSTM model shapes returned by FedMA and the hyper-parameter choices used in the experiments.

  • The section reports final global VGG and LSTM model shapes returned by FedMA with communication.
  • VGG-9 uses ReLU nonlinear activation functions, with convolution-layer shapes represented as (Cin, Cout, c, c).
  • The experiments include detailed LSTM architecture information in Tables 5 and 6.
  • The final global VGG-9 architecture is separately documented with convolution-layer shapes represented as (Cin, Cout, c, c).
  • Hyper-parameters: For BBP-MAP, γ0 controls discovery of new hidden states, while σ2 is the variance of local weights around corresponding global weights.
  • Hyper-parameters: γ0 = 10^-3, σ2_0 = 1, and σ2 = 1 led to good performance for the LSTM on the Shakespeare dataset.

E PRACTICAL CONSIDERATIONS

FedMA’s matching complexity depends on the number of clients, model size, and matched hidden states; optimized Hungarian-algorithm implementations substantially improve practical speed.

  • The worst-case matching complexity is O(D · (JL)^2) for cost-matrix construction plus O((JL)^3) for the Hungarian algorithm.
  • When all hidden states are matched, the best per-layer complexity is O(D·L^2+L^3).
  • FedMA can be relatively slow when many clients participate and each client trains a large model.
  • The lapsolver implementation of the Hungarian algorithm produced orders-of-magnitude speedups over the vanilla implementation in the authors’ experiments.

F HYPER-PARAMETERS FOR THE HANDLING DATA BIAS EXPERIMENTS

The data-bias experiments retuned local epochs for FedAvg, FedProx, and FedMA, selecting values associated with good convergence in the experiments.

  • FedAvg and FedProx were evaluated with local epoch candidates E ∈ {5, 10, 20, 30}.
  • A relatively large E can easily lead to poor FedAvg convergence, while FedProx tolerates larger E better.
  • The experiments used E = 5 for FedAvg and FedProx, and E = 50 for FedMA because these choices led to good convergence.
Loading 2002.06440v1…