PaperlayerView exact revision on arXiv

Source-linked AI summary

FedBE: Making Bayesian Model Ensemble Applicable to Federated Learning

Hong-You Chen, Wei-Lun Chao

arXiv:2009.01974v4cs.LGstat.ML

TL;DR

Federated learning needs reliable aggregation of private client models despite non-i.i.d. data. FedBE samples and ensembles global models from a distribution fitted to client models, then distills the ensemble into a single model using unlabeled server data. Experiments report consistent gains over FEDAVG, especially with non-i.i.d. clients and deeper networks, while remaining compatible with other federated-learning improvements.

  • Problem

    Aggregating client models is challenging with non-i.i.d. data because FEDAVG can degrade and weight averaging does not fully use client-model information.

  • Method

    FedBE fits a distribution to client models, samples global models for Bayesian ensemble, and distills their predictions on unlabeled server data into one global model.

  • Results

    FedBE consistently outperforms FEDAVG across tested datasets, client conditions, and architectures, especially when neural networks are deeper.

  • Takeaways & Limitations

    FedBE can replace the aggregation method while remaining compatible with client-training regularization and server-momentum methods.

  • Takeaways & Limitations

    Multi-round Bayesian ensemble requires a single global model for clients, and FedBE assumes access to unlabeled server data.

Abstract

from arXiv · show

Federated learning aims to collaboratively train a strong global model by accessing users' locally trained models but not their own data. A crucial step is therefore to aggregate local models into a global model, which has been shown challenging when users have non-i.i.d. data. In this paper, we propose a novel aggregation algorithm named FedBE, which takes a Bayesian inference perspective by sampling higher-quality global models and combining them via Bayesian model Ensemble, leading to much robust aggregation. We show that an effective model distribution can be constructed by simply fitting a Gaussian or Dirichlet distribution to the local models. Our empirical studies validate FedBE's superior performance, especially when users' data are not i.i.d. and when the neural networks go deeper. Moreover, FedBE is compatible with recent efforts in regularizing users' model training, making it an easily applicable module: you only need to replace the aggregation method but leave other parts of your federated learning algorithm intact. Our code is publicly available at https://github.com/hongyouc/FedBE.

1 INTRODUCTION

Federated learning aggregates private client models, but non-i.i.d. data can make standard weight averaging unreliable. FedBE addresses this by sampling global models from a fitted distribution, ensembling their predictions, and distilling that ensemble into a single model for multi-round training.

  • Motivation: Federated learning lets clients collaboratively train a global model by sharing locally trained models rather than their private data.Training proceeds through repeated client updates followed by server aggregation and redistribution of a global initialization.
  • Motivation: FEDAVG averages client weights, but its performance can degrade drastically with non-i.i.d. client data and it may discard information such as model variances.The paper also notes potential problems for over-parameterized neural networks due to permutation-invariant properties.
  • Approach: FedBE treats client models as possible global models, fits a distribution over them, samples additional models, and combines their outputs through Bayesian model ensemble.The approach uses Bayesian inference to move beyond selecting or averaging a single weight vector.
  • Results: Bayesian model ensemble improves single-round predictions over FEDAVG, especially under non-i.i.d. client conditions.In the toy analysis, ensembling 10 sampled models reaches 69% test accuracy versus 64% for FEDAVG.
  • Limitations: A central limitation is that multi-round FedBE requires a single global model to represent the ensemble and assumes the server can access unlabeled data.The paper discusses this server-data assumption and associated privacy considerations separately.
  • Approach: FedBE distills ensemble predictions on unlabeled server data into one global model that can be sent to clients for subsequent communication rounds.Stochastic weight averaging is used during distillation because federated ensemble predictions can be noisy.

5 EXPERIMENT

Experiments across datasets, client conditions, architectures, and practical federated settings show that FEDBE robustly improves aggregation, especially for deeper networks and non-i.i.d. data. FEDBE also remains compatible with existing local-training and server-momentum methods, while requiring relatively little unlabeled data.

  • Experimental setup: Across CIFAR-10/100 and Tiny-ImageNet, FEDBE is evaluated under i.i.d. and non-i.i.d. conditions with ConvNet, ResNet, and MobileNetV2.The study covers multiple datasets, client conditions, and neural-network architectures.
  • Main results: FEDBE consistently outperforms FEDAVG, with a 2 ∼9% gain on non-i.i.d. CIFAR-10 that increases as networks deepen.The experiments use 40 communication rounds and compare FEDBE with FEDAVG and other baselines.
  • Compatibility: FEDBE improves FEDPROX and FEDAVGM, and combining FEDBE with these methods outperforms FEDBE alone in many cases.This supports compatibility with improved local training and server momentum.
  • Sampling: Sampling M = 10 ∼20 global models, in addition to the weight average and clients, is sufficient to make FEDBE effective.Bayesian model ensemble improves accuracy relative to alternative distillation targets.
  • Ensemble behavior: Bayesian model ensemble outperforms weight averaging at nearly all communication rounds, although its predictions remain noisy.This comparison uses FEDAVG-trained client models on a non-i.i.d. CIFAR-10 ResNet20 experiment.
  • Unlabeled data: FEDBE outperforms FEDAVG with 1K unlabeled examples, while 2K examples also suffice and additional data provide diminishing gains.The 1K dataset is 4% of total client data, while 2K is 5% of 40K CIFAR-10 labeled examples.
  • Unlabeled data: Out-of-domain unlabeled data from CIFAR-100 or Tiny-ImageNet match or outperform CIFAR-10 unlabeled data, supporting their use for distillation.The results suggest unlabeled data serve mainly as a distillation medium rather than revealing future test data.
  • Network depth: FEDBE is more robust than FEDAVG as network depth increases, but noisy distillation targets remain a key limitation for deeper models.Using ground-truth targets improves with added layers, whereas FEDBE suffers the least degradation among tested methods.

B.1 IMPLEMENTATION DETAILS

The implementation fixes training and sampling choices across experiments, including local epochs, weight decay, distillation epochs, augmentation, and batch-normalization handling.

  • ConvNet and ResNet use E = 20 local epochs, while MobileNet uses E = 10; 1-Ensemble baselines tune E from [10, 20, ..., 200].
  • Weight decay is searched separately for each network and method over [1e−3, 1e−4] using a validation set.The authors report that method-specific tuning is essential for fair comparisons.
  • Distillation epochs are tuned for v-Distillation, with 20 epochs found stable and also applied to FEDBE.
  • Server inference on unlabeled data omits augmentation, whereas local training and distillation use padded, flipped, and randomly cropped images.CIFAR images use 2-pixel padding and Tiny-ImageNet images use 4-pixel padding before cropping back to their original sizes.
  • For networks with batch normalization, the global distribution is constructed as in Section 3, with no experimental issues reported.The authors also cite stable SWAG behavior on very deep networks.

B.2 TRAINING FEDAVG

FEDAVG training is sensitive to the local learning rate in non-i.i.d. settings: within-round decay is unhelpful, whereas decay across communication rounds improves performance.

  • A local learning rate that is too large prevents convergence, while one that is too small prevents satisfactory performance within E local epochs.The authors identify an appropriate η_l as important for training on non-i.i.d. local data.
  • Decaying η_l across communication rounds improves performance, unlike decaying it within each round of local training.The schedule uses η_l = 0.01, 0.001, and 0.0001 over successive fractions of the total R rounds.
  • Figure 6 compares FEDAVG with ConvNet on Step-non-i.i.d. CIFAR-10 with and without learning-rate decay within each local-training round.
  • Figure 7 compares FEDAVG with ConvNet on Dirichlet-non-i.i.d. CIFAR-10 under learning-rate decay at later communication rounds and different α values.

B.3 EFFECTS OF WEIGHT DECAY IN LOCAL CLIENT TRAINING

Weight decay is used to reduce overfitting to local distributions, and experiments report higher accuracy and greater robustness to the number of local epochs for FEDAVG and FEDBE.

  • Non-i.i.d. federated learning is prone to model drift and sensitive to E, motivating ℓ2 regularization as weight decay during local training.
  • Table 7 reports FEDBE results using models sampled from a Dirichlet distribution on Step-non-i.i.d. CIFAR-10 for different α settings.
  • Weight decay increases test accuracy and makes both FEDAVG and FEDBE more robust to the choice of local epochs E.
  • Figure 8 compares FEDAVG and FEDBE with or without local-training weight decay across different numbers of local epochs E on Step-non-i.i.d. CIFAR-10.

C.1 GLOBAL MODEL SAMPLING IN FEDBE

FEDBE’s Dirichlet-based model sampling is largely insensitive to α, while clients’ models tend to be over-confident on incorrect predictions.

  • FEDBE accuracy is not sensitive to changes in the Dirichlet parameter α.
  • FEDBE with Dirichlet sampling is slightly worse than FEDBE with Gaussian sampling.
  • Client models tend to assign high confidence to incorrect predictions, whereas the study compares their confidence histograms with those of sampled models.

C.2 ANALYSIS ON WEIGHT AVERAGE, (BAYESIAN) MODEL ENSEMBLE, AND DISTILLATION

The one-round experiments compare weight averaging, model ensembles, Bayesian model ensembles, and distillation on Step-non-i.i.d. CIFAR-10. Bayesian model ensemble improves over weight averaging, while distillation can substantially reduce ensemble accuracy.

  • Bayesian model ensemble adds a 2% gain over model ensemble without distillation, which already clearly outperforms weight averaging.The comparison uses 10 sampled models beyond weight averaging and individual clients.
  • 60.5% and 62.5% accuracy from ensemble methods remained relatively higher than other methods but potentially weak as distillation targets.The passage reports these accuracies while discussing the limitations of summarizing ensembles into one model.
  • Distillation into a single global model largely degrades accuracy compared with retaining the ensemble.This finding highlights the difficulty of applying ensemble distillation directly in one-round federated learning.
  • SWA distillation outperforms SGD distillation for both model ensemble and Bayesian model ensemble.The result supports using SWA for the proposed distillation procedure.
  • With multiple communication rounds, FEDBE provides a better initialization for subsequent local training and eventually achieves much higher accuracy than 1-Ensemble.This contrasts with the one-round difficulty of summarizing ensemble predictions into a single model.

C.3 FEDBE VS. FEDAVG

Additional experiments show that FEDBE improves over FEDAVG across several settings, including feature analysis, normalization choices, distillation details, and compatibility with SCAFFOLD. The gains are especially evident for deeper models and non-i.i.d. data.

  • 77.5% accuracy with FEDBE exceeds 72.5% from end-of-training Bayesian ensemble and 70.2% from FEDAVG on Step-non-i.i.d. CIFAR-10 with ResNet20.The comparison demonstrates the importance of incorporating Bayesian model ensemble during multi-round federated learning.
  • Data augmentation enables knowledge distillation from FEDAVG predictions to improve accuracy to 72.6%/73.4% versus 72.0%/70.2% for ConvNet/ResNet20.Without augmentation, initializing the student with FEDAVG causes zero gradients in distillation.
  • FEDBE with SWA remains stable beyond 10 distillation epochs, achieving 77.5% at 20 epochs and 77.3% at 40 epochs.SGD-based FEDBE is more sensitive, reaching 74.9% at 20 epochs and 74.0% at 40 epochs for ResNet20.
  • FEDBE with GN/BN achieves 79.6%/80.2%, improving over FEDAVG with GN/BN at 76.4%/74.6% after 200 rounds.The experiments use ResNet20 on Step-non-i.i.d. CIFAR-10.
  • FEDBE+SCAFFOLD improves ResNet20/ResNet32 accuracy to 76.4%/72.7%, compared with 59.4%/55.3% for SCAFFOLD alone.SCAFFOLD with weight averaging stops improving after roughly 10 rounds, whereas FEDBE+SCAFFOLD continues improving.
  • FEDBE consistently outperforms FEDAVG in the reported non-i.i.d. CIFAR-100 experiment.The setting uses 100 clients, samples 10 clients per round, and runs for 100 rounds.

D.1 EXTRA COMPUTATION COST

FEDBE adds computation on the server rather than the clients. The paper argues this cost is usually acceptable because servers are computationally rich and the extra work can be parallelized.

  • FEDBE involves more computation than FEDAVG, with no additional client burden.The extra cost is incurred on the server.
  • On a 2080 Ti GPU for CIFAR-10 ConvNet, building distributions and sampling takes 0.2s, inference 2.4s, and distillation 10.4s.These measurements quantify the main additional server-side operations.
  • Constructing ensemble predictions requires evaluating each sampled model on the unlabeled set U, which can be parallelized on modern GPUs.The paper states that M = 10–20 samples are sufficient for effective Bayesian model ensemble.
  • Deeper models may be poorly calibrated and have non-convex loss surfaces, making weight averaging less effective and potentially requiring more communication rounds.This motivates the aggregation setting in which FEDBE incurs its additional server-side computation.

E FURTHER ANALYSIS

Across Step and Dirichlet non-i.i.d. CIFAR-10 settings with ResNet20 and ResNet32, FEDBE reaches the highest accuracy after roughly 10 rounds and continues improving with more rounds. Direct client-model ensemble methods can lead early but are later surpassed.

  • E.1 TEST ACCURACY AT DIFFERENT ROUNDS: FEDBE obtains the highest accuracy after roughly 10 communication rounds across the reported ResNet20/ResNet32 non-i.i.d. settings.The experiments cover both Step and Dirichlet settings over 40 rounds.
  • E.1 TEST ACCURACY AT DIFFERENT ROUNDS: FEDBE gradually improves as more communication rounds are added.The 40-round curves correspond to the final accuracies reported in the main experiments.
  • E.1 TEST ACCURACY AT DIFFERENT ROUNDS: v-Distillation is normally strongest during the first 10 rounds but is surpassed by FEDBE afterward.The paper hypothesizes that early client models make the fitted distributions unstable.
  • E.1 TEST ACCURACY AT DIFFERENT ROUNDS: Most ensemble methods outperform FEDAVG during the first 10 rounds, indicating robustness in aggregation.The comparison excludes 1-Ensemble from this statement.
Loading 2009.01974v4…