PaperlayerView exact revision on arXiv

Source-linked AI summary

Federated Learning via Posterior Averaging: A New Perspective and Practical Algorithms

Maruan Al-Shedivat, Jennifer Gillenwater, Eric Xing, Afshin Rostamizadeh

arXiv:2010.05273v4cs.LGcs.AIstat.ML

TL;DR

Federated learning seeks to train models across clients with heterogeneous local data, but optimization methods face challenges from communication limits and data heterogeneity. This paper formulates federated learning as posterior inference and develops FedPA to aggregate approximate local posteriors. FedPA generalizes federated averaging and achieves state-of-the-art results on challenging benchmarks.

  • Problem

    Federated learning must optimize a global objective across heterogeneous clients with limited communication, while additional local computation can bias FEDAVG updates and lead to inferior optima.

    Source: §1 · p. 1

  • Method

    FedPA uses approximate local posterior inference and aggregation to infer the global posterior mode with linear computation and communication costs.

    Source: §1 · p. 2, §4 · p. 4

  • Results

    FEDPA achieves state-of-the-art results on realistic federated-learning benchmarks with respect to multiple metrics of interest.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The posterior-inference perspective broadens the design space for federated-learning algorithms beyond purely optimization techniques.

    Source: §1 · p. 2

    The method has limitations arising from design choices, including its posterior-sampling and covariance-estimation techniques.

    Source: §6 · p. 9

Abstract

from arXiv · show

Federated learning is typically approached as an optimization problem, where the goal is to minimize a global loss function by distributing computation across client devices that possess local data and specify different parts of the global objective. We present an alternative perspective and formulate federated learning as a posterior inference problem, where the goal is to infer a global posterior distribution by having client devices each infer the posterior of their local data. While exact inference is often intractable, this perspective provides a principled way to search for global optima in federated settings. Further, starting with the analysis of federated quadratic objectives, we develop a computation- and communication-efficient approximate posterior inference algorithm -- federated posterior averaging (FedPA). Our algorithm uses MCMC for approximate inference of local posteriors on the clients and efficiently communicates their statistics to the server, where the latter uses them to refine a global estimate of the posterior mode. Finally, we show that FedPA generalizes federated averaging (FedAvg), can similarly benefit from adaptive optimizers, and yields state-of-the-art results on four realistic and challenging benchmarks, converging faster, to better optima.

1 INTRODUCTION

Federated learning must balance communication limits with heterogeneous client data, which can make federated averaging converge to inferior optima when clients do more local work. The paper reframes federated learning as posterior inference and introduces an efficient algorithm designed for stateless clients.

1 INTRODUCTION

  • Federated learning distributes learning across clients holding heterogeneous data, but communication costs and data heterogeneity can affect efficiency and convergence.FEDAVG reduces communication relative to mini-batch SGD by allowing multiple local SGD steps per round.Source: §1 · p. 1
  • Heterogeneous client data can make additional local computation bias FEDAVG updates and cause it to stagnate at inferior optima, while MB-SGD may dominate at convergence.Stateful-client remedies can help in cross-silo settings but are impractical in cross-device settings where an average client may participate only once.Source: §1 · p. 1
  • Posterior inference offers a route to fast training and consistent convergence with stateless clients, with the proposed algorithm designed for linear computation and communication costs.The approach aggregates information about local posteriors to infer modes of the global posterior.Source: §1 · p. 2
  • FEDPA generalizes FEDAVG, has complexity and memory footprint similar to FEDAVG, and achieves state-of-the-art results across multiple metrics on realistic federated-learning benchmarks.The paper also identifies FEDAVG with many local steps as a special case using identity estimates of local posterior covariances.Source: §1 · p. 2

2 RELATED WORK

Related work has studied federated averaging as an optimization method and sub-posterior aggregation as a way to scale MCMC. The paper connects these lines under federated-learning constraints.

Federated optimization

  • Prior work analyzes FEDAVG convergence in homogeneous and heterogeneous settings and develops adaptive optimization, while noting that FEDAVG needs careful learning-rate scheduling to reach the desired optimum.The authors describe their work as, to their knowledge, perhaps the first to connect and analyze federated optimization from a probabilistic-inference perspective.Source: §2 · p. 2

Distributed MCMC

  • Sub-posterior aggregation, including consensus Monte Carlo, was proposed to scale MCMC to large datasets; this paper adapts distributed inference ideas to federated-learning constraints.The cited methods are Neiswanger et al. (2013) and Scott et al. (2016).Source: §2 · p. 2

3 A POSTERIOR INFERENCE PERSPECTIVE ON FEDERATED LEARNING

The paper recasts the weighted federated objective as inference over a global posterior that decomposes into local posteriors. This provides an aggregation principle, while exact posterior statistics can be too costly to compute and communicate at realistic model scales.

Federated objective

  • The global objective is a weighted average of client losses, with client weights typically proportional to local dataset sizes so it matches the centralized training objective.Each client loss is computed from its local data.Source: §3 · p. 3

Posterior inference

  • With an uninformative uniform prior, global posterior modes coincide with the federated objective's optima, and the global posterior decomposes into local posteriors.The decomposition motivates inferring local posteriors and aggregating them at the server.Source: §3 · p. 3
  • Multiplicatively aggregating local posteriors can construct the global posterior and its mode at the server once clients communicate their local posteriors.The remaining challenge is making inference and communication efficient for federated settings.Source: §3 · p. 3

Federated least squares

  • For federated least squares, local posteriors are Gaussian, and their product is a Gaussian whose mean is also the global posterior mode.The client means and inverse covariances are sufficient statistics for determining that global mode.Source: §3 · p. 3
  • Communicating local means and inverse covariances requires O(d2) space and communication, while matrix inversion requires O(d3) computation; standard federated algorithms communicate O(d).These costs make direct posterior aggregation expensive for typical cross-device settings.Source: §3 · p. 3

Approximate federated posterior inference

  • For generally non-Gaussian posteriors, the paper proposes Laplace approximations using local means and covariances, then server-side computation of an approximate global posterior mode.The following practical algorithm is designed to make costs linear in the number of model parameters.Source: §3 · p. 4

4 FEDERATED POSTERIOR AVERAGING: A PRACTICAL ALGORITHM

Section 4 develops FedPA by converting global posterior-mode estimation into an equivalent optimization problem and using local posterior estimates to form client updates. Its design addresses efficient inference and communication, and the paper reports faster convergence to a better optimum than FEDAVG.

4 FEDERATED POSTERIOR AVERAGING: A PRACTICAL ALGORITHM

  • Direct global posterior-mode computation requires O(d^3) server work and O(d^2) client communication; an equivalent optimization enables O(d) communication and server computation per round when clients compute deltas.The optimization reformulation makes the global posterior mode the minimizer of a quadratic objective.Source: §4 · p. 4
  • Clients use SG-MCMC to estimate local posterior moments, with IASG mixing the chain before periodically producing samples by averaging intermediate SGD iterates.More local computation can produce more posterior samples and better local-moment estimates.Source: §4 · p. 5
  • A dynamic programming algorithm computes client deltas in O(d) time and memory, avoiding the O(d^3) computation and O(d^2) memory of naïve matrix inversion.It uses recursive covariance updates and can update the delta online as posterior samples arrive.Source: §4 · p. 5
  • With five local epochs, DP-based update computation approaches FEDAVG's complexity as dimensionality grows, leaving SGD steps as the majority of practical cost.Table 1 reports client-update times in milliseconds and relative increases for synthetic linear regressions.Source: §4 · p. 5, Table 1 · p. 5
  • FEDAVG can be viewed as posterior inference with identity covariance estimation; analysis of alternative samplers and exponential-family extensions remains future work.The paper uses a variant of SG-MCMC and notes HMC and NUTS as possible alternatives.Source: §4 · pp. 4–6
  • FEDPA converges faster and to a better optimum than FEDAVG, trades bias for slightly more variance near convergence, and benefits substantially from more local computation.The variance difference becomes visible only closer to convergence.Source: §4 · p. 6

5 EXPERIMENTS

The experiments compare FEDPA with adaptive FEDAVG baselines and MIME on four benchmark tasks. FEDPA achieves the best limited-round evaluation accuracy on both image tasks and also improves selected StackOverflow metrics.

5.1 THE SETUP

  • The four tasks use EMNIST, CIFAR100, and StackOverflow; comparisons include tuned adaptive FEDAVG variants and MIME.The task models are CNN, ResNet-18, LSTM, and multi-label logistic regression on bag-of-words vectors; FEDPA-ME uses IASG-produced approximate posterior samples and has the same local and global computational complexity as FEDAVG-ME, with two extra hyperparameters.Source: §5 · p. 6, §5.1 · pp. 6–7
  • Evaluation measures rounds to specified performance thresholds and best performance within round limits, which were selected partly for computational reproducibility.The authors also provide learning curves for a fuller view of algorithm behavior.Source: §5.1 · p. 7

5.2 RESULTS ON BENCHMARK TASKS

  • FEDPA-ME achieves the best limited-round evaluation accuracy on EMNIST-62 and CIFAR-100; its EMNIST-62 accuracy is 87.3% after 1500 rounds.That is a 0.5% absolute improvement and bridges 41.7% of the gap between the cited 88% centralized accuracy and 86.8% best prior federated accuracy.Source: Table 3 · p. 8, §5.2 · p. 9
  • Both multi-epoch methods reach 84% EMNIST-62 accuracy in under 100 rounds and 30% CIFAR-100 accuracy by round 350, but extra local computation eventually hurts FEDAVG.FEDAVG-ME does not consistently achieve 86% on EMNIST-62 within 1500 rounds and takes 350 extra rounds to reach 40% on CIFAR-100.Source: §5.2 · pp. 8–9
  • The experiments switch from FEDAVG-style burn-in to posterior sampling after a fixed number of rounds, a design choice the authors leave open to more adaptive strategies.They suggest local loss values as one possible trigger for beginning sampling.Source: §5.2 · p. 8
  • On StackOverflow NWP, FEDPA gains 0.4% accuracy over the best baseline; on LR, it improves micro- and macro-F1 while underperforming in precision and recall.The authors associate FEDPA's F1 advantage with better capture of rare labels, while FEDAVG better predicts high-frequency labels.Source: §5.2 · p. 9

6 CONCLUSION AND FUTURE DIRECTIONS

The conclusion presents posterior averaging as a practical, efficient generalization of federated averaging with state-of-the-art results on challenging benchmarks; the analysis examines its approximations and the relationship to FedAvg.

6 CONCLUSION AND FUTURE DIRECTIONS

  • FedPA generalizes federated averaging, is similarly practical and efficient, and achieves state-of-the-art results on multiple challenging benchmarks.Source: §6 · p. 9
  • The method’s limitations arise from design choices such as its posterior-sampling and covariance-estimation techniques.The authors identify alternative sampling strategies, covariance estimators, and variational inference as future directions.Source: §6 · p. 9

A PRELIMINARY ANALYSIS AND ABLATIONS

A.1 DISCUSSION OF THE CONVERGENCE OF FEDPA VS. FEDAVG

The analysis relates biased-gradient convergence to the bias of client deltas, then empirically contrasts how FEDAVG and FEDPA trade bias against variance as local computation increases.

A.1 DISCUSSION OF THE CONVERGENCE OF FEDPA VS. FEDAVG

  • FEDAVG’s uncorrected bias determines its distance from the global optimum, whereas FEDPA’s bias should diminish as local computation increases.Source: §A.1 · p. 13
  • The empirical comparison uses synthetic least-squares problems with 10, 100, or 1000 features and varies local steps and FEDPA’s shrinkage parameter.Source: §A.1 · p. 14
  • With more local steps, FEDAVG’s delta bias grows while its variance falls; FEDPA’s bias falls, and its variance eventually falls after an initial increase.The initial FEDPA variance increase occurs when the shrinkage parameter is held fixed; the parameter should be selected separately for each number of local steps.Source: §A.1 · p. 14, Figure 3 · p. 14

A.2 ANALYSIS OF THE QUALITY OF IASG-BASED SAMPLING AND COVARIANCE

The IASG analysis evaluates approximate posterior samples using effective sample size and reports how burn-in, sampling intervals, learning rate, and dimensionality affect sample quality.

A.2 ANALYSIS OF THE QUALITY OF IASG-BASED SAMPLING AND COVARIANCE

  • The experiments assess IASG samples by effective sample size on synthetic least-squares problems with 10, 100, and 1000 dimensions.The problems have 500 data points and batch size 10.Source: §A.2 · p. 15, Figure 4 · p. 15
  • Sample quality deteriorates as dimensionality increases, and learning rate is the most sensitive IASG hyperparameter.An overly large rate can cause divergence, while an overly small rate can produce correlated samples.Source: §A.2 · p. 15, Figure 4 · p. 15
  • More burn-in steps and more steps per sample generally improve IASG sample quality.Source: §A.2 · p. 15

B PROOFS

The appendix establishes posterior and quadratic optimization results, then gives a dynamic-programming approach for computing client deltas from posterior samples without the naive matrix costs.

B PROOFS

  • Under a uniform prior, an existing global posterior decomposes into a product of local posteriors, and its mode minimizes a quadratic objective.Source: B PROOFS · p. 16

C COMPUTATION OF CLIENT DELTAS VIA DYNAMIC PROGRAMMING

C.1 THE SHRINKAGE ESTIMATOR OF THE COVARIANCE

  • The covariance estimator combines identity and sample covariance, with its shrinkage parameter controlling the estimator’s bias-variance tradeoff.Source: §C.1 · p. 17
  • The shrinkage schedule is chosen so each new sample enables a rank-one covariance update, which supports incremental inverse updates via Sherman-Morrison.Source: C COMPUTATION OF CLIENT DELTAS VIA… · p. 16, §C.1 · p. 17

C.2 COMPUTING DELTAS USING SHERMAN-MORRISON AND DYNAMIC PROGRAMMING

  • The dynamic program updates the client delta as posterior samples arrive, providing an online, any-time estimate that improves with additional sampling.The method returns the FEDAVG delta during burn-in or after one posterior sample.Source: §C.2 · p. 19

D DETAILS ON THE EXPERIMENTAL SETUP

D.1 DATASETS, TASKS, AND MODELS

The empirical study uses EMNIST-62, CIFAR-100, and StackOverflow tasks, a subset of tasks introduced by Reddi et al. (2020). StackOverflow contributes next-word prediction and tag prediction tasks.

EMNIST-62

  • EMNIST-62 character recognition uses author-partitioned handwritten digits and letters, with a two-layer convolutional network.The dataset has 62 classes and writing-style heterogeneity across authors.Source: §D.1 · p. 19

CIFAR-100

  • CIFAR-100 is partitioned across 500 clients with 100 examples each and uses a modified ResNet-18 with group normalization.A two-step latent Dirichlet allocation over coarse and fine labels creates the client label distributions.Source: §D.1 · p. 19

StackOverflow

  • StackOverflow experiments cover next-word prediction with 10,000 words and multi-label tag prediction with 500 labels.The federated data are partitioned by user.Source: §D.1 · p. 20
  • Tag prediction represents each question or answer as a normalized bag-of-words vector and uses a simple linear regression model.Source: §D.1 · p. 20
  • Next-word prediction uses each client's first 128 sentences, padded or truncated to 20 words, and an embedding–LSTM model.Accuracy is measured on content words, excluding OOV, BOS, and EOS symbols.Source: §D.1 · p. 20

D.2 METHODS

The methods compare adaptive-server FEDAVG baselines with FEDPA under shared configurations, then select hyperparameters based on FEDAVG performance.

D.2 METHODS

  • FEDAVG baselines use adaptive server optimizers with one or multiple local epochs; client momentum is used for EMNIST-62, CIFAR-100, and StackOverflow LR.The baselines are called FEDAVGM for EMNIST-62 and CIFAR-100, FEDADAM for StackOverflow NWP, and FEDADAGRAD for StackOverflow LR in Reddi et al. (2020).Source: §D.2 · p. 20
  • FEDPA uses the FEDAVG configurations and matches FEDAVG during burn-in, differing only in client-side computation during sampling.Source: §D.2 · p. 20

D.3 HYPERPARAMETERS AND GRIDS

  • The best server and client learning rates are selected using FEDAVG performance, then used for FEDPA.Hyperparameter grids and selected values are reported in Tables 5 and 6, respectively.Source: §D.3 · p. 21, Table 6 · p. 21

E ADDITIONAL EXPERIMENTAL RESULTS

Additional learning curves show faster convergence with more local epochs and sampling-related improvements on several tasks, alongside a distinct StackOverflow LR precision–recall trade-off.

Additional experimental results

EMNIST-62

CIFAR-100 and StackOverflow

Related papers

Loading 2010.05273v4…