PaperlayerView exact revision on arXiv

Source-linked AI summary

Adaptive Federated Optimization

Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, H. Brendan McMahan

arXiv:2003.00295v5cs.LGcs.DCmath.OCstat.ML

TL;DR

Federated optimization must handle communication constraints, heterogeneous client data, and convergence difficulties in FedAvg. This paper introduces a client/server optimizer framework with adaptive server methods, analyzes convergence in nonconvex settings, and benchmarks the approach across diverse federated tasks, finding strong empirical performance. Its theoretical analysis assumes full client participation or incurs an additional variance term under limited participation, and it also assumes bounded variance and gradients.

  • Problem

    FedAvg can have convergence issues in federated learning, while standard optimization methods can incur high communication costs and federated data are heterogeneous.

  • Method

    The paper develops FEDOPT, a client/server optimizer framework that applies gradient-based server optimizers to averaged client model updates, including adaptive methods.

  • Results

    Adaptive optimizers improve convergence and show strong empirical performance across seven diverse federated learning tasks, improving upon commonly used baselines.

  • Takeaways & Limitations

    Adaptive server optimization provides a principled, theoretically justified way to incorporate adaptivity into federated learning without increasing client storage or communication costs.

  • Takeaways & Limitations

    The analysis assumes full client participation; extending it to limited participation adds a variance term depending on the sampled-client fraction, and it assumes bounded variance and gradients.

Abstract

from arXiv · show

Federated learning is a distributed machine learning paradigm in which a large number of clients coordinate with a central server to learn a model without sharing their own training data. Standard federated optimization methods such as Federated Averaging (FedAvg) are often difficult to tune and exhibit unfavorable convergence behavior. In non-federated settings, adaptive optimization methods have had notable success in combating such issues. In this work, we propose federated versions of adaptive optimizers, including Adagrad, Adam, and Yogi, and analyze their convergence in the presence of heterogeneous data for general non-convex settings. Our results highlight the interplay between client heterogeneity and communication efficiency. We also perform extensive experiments on these methods and show that the use of adaptive optimizers can significantly improve the performance of federated learning.

1 INTRODUCTION

Federated learning reduces communication by using local client updates, but heterogeneous data and FedAvg’s convergence issues motivate adaptive federated optimization. The paper introduces a general client/server optimizer framework and reports strong empirical performance for adaptive methods.

  • Federated learning setting: Federated learning keeps raw client data local while coordinating model learning through a central server.The paper focuses primarily on cross-device FL, where few clients participate per round and cannot maintain state across rounds.
  • Motivation: FedAvg reduces communication by having clients perform multiple local SGD epochs before the server averages their models.This local-update strategy can greatly reduce communication compared with standard distributed optimization.
  • Motivation: FedAvg can suffer from convergence issues because of client drift, limited adaptivity, and heavy-tailed stochastic gradient noise.Adaptive learning rates are described as beneficial in settings with heavy-tail gradient noise distributions.
  • Framework: The proposed framework combines local client optimizers with a gradient-based server optimizer applied to averaged client model updates.FedAvg is recovered when SGD is used for both optimizers and the server learning rate is 1.
  • Framework: Adaptive server optimization enables per-coordinate learning rates without increasing client storage or communication costs, while remaining compatible with cross-device FL.The framework can use adaptive optimizers on either clients or the server, with the paper emphasizing adaptive server optimization.
  • Empirical evaluation: The paper reports adaptive methods that improve commonly used baselines, are easier to tune, and perform strongly across seven diverse federated learning tasks.The benchmarks involve image and text data with varying heterogeneity and client counts.

2 FEDERATED LEARNING AND FEDAVG

The paper formalizes FedAvg as a server update using a pseudo-gradient derived from averaged client model differences, then generalizes it through the FEDOPT framework. It establishes convergence for adaptive server optimizers under nonconvex federated objectives while modeling client heterogeneity explicitly.

  • Assumptions: The analysis allows nonconvex client objectives and unbiased stochastic gradients while assuming Lipschitz gradients, bounded variance, and bounded gradients.The assumptions include local and global variance controls for stochastic gradients and client objectives.
  • Assumptions: Client heterogeneity is quantified by σg, with σg = 0 corresponding to the i.i.d. setting.The analysis makes no further assumptions about similarity between client datasets.
  • FEDAVG: FEDAVG broadcasts the global model, runs local SGD on selected clients, and averages their resulting models at the server.The method operates by selecting a subset of clients at each round.
  • FEDOPT: FEDOPT treats the negative average client model difference as a pseudo-gradient for a configurable server optimizer.FedAvg is the special case of SGD server optimization with learning rate 1; other server learning rates and optimizers are allowed.
  • FEDOPT: FEDOPT supports separate client and server optimizers, adaptive methods such as Adam and Yogi, server momentum, and learning-rate schedules.Client and server learning rates are denoted ηl and η, and may depend on the communication round.
  • Theory: The paper provides a theoretical basis for using averaged model differences as pseudo-gradients and guarantees convergence for server optimizers including Adagrad, Adam, and Yogi.This establishes principled adaptive federated optimizers within the FEDOPT framework.

3 ADAPTIVE FEDERATED OPTIMIZATION

The paper specializes federated optimization by using adaptive server optimizers with SGD clients, then analyzes convergence under heterogeneous data. Its analysis links learning-rate choices, local steps, heterogeneity, and communication efficiency.

  • Adaptive server optimization: Adaptive federated optimization uses SGD on clients and Adagrad, Adam, or Yogi on the server, preserving FedAvg-level communication and cross-device compatibility.Adaptive state is maintained on the server rather than increasing client storage or communication costs.
  • Convergence analysis: The framework supports full-participation convergence analysis for adaptive methods in general nonconvex settings, with extensions to partial participation and weighted averaging.The full-participation analysis can be generalized to partial participation with an additional variance term depending on the sampled-client fraction.
  • Client heterogeneity: Client and server learning rates can reduce, but not remove, the effect of client heterogeneity on convergence.The analysis notes that the server learning rate η = 1 used in FedAvg does not necessarily minimize the theoretical upper bound.
  • Convergence rates: O(1/√mKT) is obtained when T is sufficiently large compared with K, matching the best known rate for the analyzed general nonconvex setting.The same convergence rates are reported for FEDADAGRAD and FEDADAM, with the discussion applying to FEDYOGI as well.
  • Learning-rate decay: Convergence requires decaying the client learning rate rather than the server learning rate because client drift from local updates does not vanish when η_l remains constant.The paper also reports that learning-rate decay can improve empirical performance.
  • Communication efficiency: Larger K can reduce communication rounds when K = O(Tσ2_g), while highly heterogeneous settings may require mechanisms such as control variates.In the i.i.d. case where σ_g = 0, the analysis permits K to be very large.

4 EXPERIMENTAL EVALUATION: DATASETS, TASKS, AND METHODS

The experimental evaluation spans seven learning tasks across five federated datasets, including naturally arising client partitions. It compares adaptive federated optimizers with FedOpt variants, FedAvg, FedAvgM, and SCAFFOLD using communication rounds as a training-time proxy.

  • Datasets and tasks: Seven learning tasks across five datasets form the evaluation suite, including CIFAR-10, CIFAR-100, EMNIST, Shakespeare, and Stack Overflow.The suite contains image and text datasets, with three datasets using naturally arising client partitioning.
  • Implementation: The implementation samples clients uniformly across rounds and weights client outputs by each client’s number of training samples.Sampling is without replacement within a round and with replacement across rounds.
  • Compared methods: The comparison includes FEDADAGRAD, FEDADAM, FEDYOGI, FEDAVG, FEDAVGM, FEDOPT, and SCAFFOLD.FEDOPT uses SGD as both client and server optimizer, while FEDAVG and FEDAVGM use server momentum parameters of 0 and 0.9, respectively.
  • Tuning protocol: Hyperparameters η_l, η, and τ are selected by grid search using average training loss over the final 100 rounds rather than validation data.The paper notes that validation data is often inaccessible in federated, especially cross-device, settings.
  • Evaluation metrics: Communication rounds serve as a proxy for wall-clock training time because all algorithms exchange equal-sized objects between server and clients.Validation performance is measured throughout training, with task-specific validation-set construction.

5 EXPERIMENTAL EVALUATION: RESULTS

Experiments compare adaptive federated optimizers with non-adaptive methods and SCAFFOLD across multiple tasks. Adaptive methods often improve convergence, offer easier learning-rate tuning, and are relatively robust to τ.

  • Comparisons between methods: Adaptive methods are compared with FEDAVG, FEDAVGM, and SCAFFOLD using validation performance averaged over the last 100 rounds.The benchmark covers multiple image and text tasks, with additional results reported in appendices.
  • Comparisons between methods: FEDADAM and FEDYOGI improve initial convergence over FEDAVGM on CIFAR-10, CIFAR-100, and EMNIST AE, while remaining comparable to or better than non-adaptive methods.FEDADAGRAD converges faster on Shakespeare, whereas most optimizers perform well after sufficient tuning.
  • Comparisons between methods: SCAFFOLD performs comparably to or worse than FEDAVG and the adaptive methods on all evaluated tasks.On Stack Overflow, SCAFFOLD and FEDAVG are nearly identical because the large client population makes repeated sampling unlikely.
  • Ease of tuning: FEDADAM and FEDYOGI have many good client learning rates across server learning rates, whereas FEDAVGM has only a few good client learning rates for each server rate.This pattern is reported for Stack Overflow NWP and holds similarly across other tasks and optimizers.
  • Ease of tuning: τ = 10^-3 performs almost as well as other tested values for nearly all tasks and optimizers.FEDADAM and FEDYOGI show only small performance differences across τ except on Stack Overflow LR.

6 CONCLUSION

The conclusion presents adaptive optimization as a theoretically justified approach for improving federated-learning convergence and releases broad benchmarks and implementations. It also identifies differential privacy and fairness as future research directions.

  • Conclusion: Adaptive optimizers can be powerful tools for improving federated-learning convergence.The conclusion states this as the paper’s central empirical finding.
  • Conclusion: A client/server optimizer framework incorporates adaptivity into federated learning in a principled, intuitive, and theoretically justified manner.The framework supports the paper’s adaptive federated optimization methods.
  • Conclusion: The paper develops comprehensive benchmarks and an open-source framework containing models, datasets, and code.These resources are intended to support reproducibility and broad comparison.
  • Conclusion: Future research includes understanding how adaptivity affects differential privacy and fairness.These are explicitly identified as example directions for further study.

A.2 PROOF OF THEOREM 1

The proof of Theorem 1 bounds FEDADAGRAD’s optimization terms by combining smoothness, drift, variance, and adaptive-update arguments. The analysis also extends to limited client participation and derives the desired convergence rate.

  • Term bounds: FEDADAGRAD’s proof decomposes the descent argument into terms T1, T2, and T3, then bounds each using smoothness, drift, and the adaptive update.The proof uses telescoping sums and auxiliary lemmas to assemble the theorem bound.
  • Adaptive accumulator: The adaptive accumulator is bounded using its initialization v−1 ≥ τ and the monotonicity of v_t.This supports the control of terms involving √v_t + τ.
  • Drift bound: The proof bounds client drift under a step-size condition ηl ≤ 1/(8LK).Lemma 3 controls the deviation of local client iterates across multiple updates.
  • Limited participation: Under uniformly sampled partial participation, the expectation includes the sampled client set, and substituting the resulting bound yields the desired convergence rate.The proof uses the sampling structure and standard quadratic inequalities to control the participating-client terms.

B FEDERATED ALGORITHMS: IMPLEMENTATIONS AND PRACTICAL CONSIDERATIONS

The practical algorithms instantiate federated optimization through client and server optimizers, batched local training, and weighted client aggregation. Server-side adaptivity preserves cross-device compatibility without the extra client storage and communication required by client-side adaptive methods.

  • FEDOPT framework: FEDAVG is the special case of FEDOPT using SGD for both client and server optimizers with server learning rate η = 1.The framework also permits other client optimizers, server optimizers, and server learning rates.
  • Batched implementations: Practical implementations use E epochs over each client dataset and weight client updates by each client’s number of examples n_i.This produces batched versions of FEDOPT and the adaptive federated methods.
  • Optimizer implementations: FEDAVG and FEDAVGM use SGD on the server, with FEDAVGM adding momentum 0.9; both client and server learning rates are tuned.The experimental FEDAVG variant is more general than the original fixed-server-rate formulation.
  • Adaptive optimizers: FEDADAGRAD uses β1 = β2 = 0, while FEDADAM and FEDYOGI use β1 = 0.9 and β2 = 0.99 in the experiments.The paper notes that tuning these parameters could yield better results.
  • SCAFFOLD: SCAFFOLD is implemented with Option II, weighted client averaging, and control-variate state maintained across sampled clients.Its comparison is not fair in cross-device settings where clients cannot maintain state across rounds.
  • Client adaptivity: Adaptive server methods avoid the extra communication and client memory required by ADAALTER’s client accumulators.Client-side accumulators would double client-to-server and server-to-client communication relative to FEDAVG.

C DATASET & MODELS

The paper evaluates federated optimization across image and language datasets with heterogeneous client partitions and task-specific models. Experiments compare multiple server optimizers while tuning client, server, and adaptivity parameters.

  • Dataset construction: CIFAR-10 partitions 50,000 training examples across 500 clients using label distributions drawn from a symmetric Dirichlet distribution with parameter 0.1.Each client receives 100 examples, producing heterogeneous label distributions.
  • Dataset construction: CIFAR-100 uses a two-step LDA process over coarse and fine labels to create more realistic heterogeneous client datasets.The partitioning uses the Pachinko Allocation Method to reflect the dataset’s hierarchical label structure.
  • Optimization setup: The experiments compare SGD, SGDM, Adagrad, Adam, and Yogi as server optimizers, while clients consistently use mini-batch SGD.Client and server learning rates are selected through large grid searches, and τ is tuned for adaptive methods.
  • Tuning procedure: Hyperparameters are selected using average training loss over the last 100 communication rounds because each round observes only a fraction of users.The sampled training loss is therefore treated as a noisy estimate of population-level loss.

D.3 PER-TASK BATCH SIZES

The experimental setup fixes batch size per task while tuning learning rates and τ through validation-oriented hyperparameter selection. Reported results include task-specific validation and Stack Overflow test metrics.

  • Batch-size control: Batch size is fixed at a per-task level to avoid conflating it with the many other hyperparameters being tuned.Centralized and federated training use the same batch size when directly compared.
  • Hyperparameter selection: The selected hyperparameters are client learning rate η_l, server learning rate η, and adaptivity parameter τ that minimize average training loss over the last 100 communication rounds.The resulting values are used for the paper’s main figures and tables.
  • EMNIST CR evaluation: EMNIST CR validation accuracy is evaluated with constant η, η_l, and τ tuned using the final-100-round training criterion.The corresponding hyperparameter grids are described in Appendix D.
  • Stack Overflow metrics: Stack Overflow reports Accuracy for NWP and Recall@5 (×100) for LR after training on the test set.The table highlights performances within 0.5% of the best result.

E.1 RESULTS ON EMNIST CR

On EMNIST CR, the optimization methods perform similarly throughout training because the task is simple and client class heterogeneity is reduced. SCAFFOLD performs slightly worse than the other methods in this communication-limited setting.

  • Results: All evaluated methods are roughly comparable throughout EMNIST CR training.Validation accuracy is plotted for FEDADAGRAD, FEDADAM, FEDYOGI, FEDAVGM, FEDAVG, and SCAFFOLD.
  • Results: The similarity is attributed to EMNIST CR being simple, with most clients containing all classes in their local datasets.This reduces heterogeneity among client classes.
  • Results: SCAFFOLD performs slightly worse than FEDAVG and the other methods on EMNIST CR.The paper suggests stale client control variates and communication-limited experiments as possible reasons.
  • Evaluation protocol: Stack Overflow performance over training uses 10,000 randomly sampled test examples at each round because the full test set is prohibitively large.The completed model is also evaluated on the full test dataset.

E.3 LEARNING RATE ROBUSTNESS

Adaptive methods often show broader regions of effective client/server learning-rate combinations and greater stability in the best server rate as client rates vary. However, Stack Overflow LR and EMNIST AE are important exceptions, and client-rate decay can improve federated EMNIST CR accuracy while remaining communication-limited.

  • Per-task learning-rate grids: Adaptive methods often have rectangular high-performance regions, whereas FEDAVG and FEDAVGM often have triangular regions requiring simultaneous tuning of η and η_l.The patterns vary by optimizer and task.
  • Client/server-rate relation: Adaptive optimizers generally maintain a more stable best server learning rate η as client learning rate η_l varies across tasks.For FEDAVG and FEDAVGM, the best rates show a clear inverse relationship in virtually all tasks.
  • Exceptions: Stack Overflow LR is a clear exception, showing a definitive inverse relationship between client and server learning rates for all optimizers.EMNIST AE also differs, with a relatively small range of good client learning rates.
  • Exceptions: EMNIST AE differs qualitatively because escaping saddle points, rather than converging to critical points, is identified as its primary training obstacle.The paper expects this task to behave differently from the remaining tasks, including EMNIST CR.
  • Centralized comparison: 88% accuracy is achieved by centralized EMNIST CR training, significantly above the best federated results reported in Table 1.The paper relates this gap partly to theory requiring a small or decayed client learning rate.
  • Learning-rate decay: A staircase schedule reducing η_l by 0.1 every 500 rounds improves the accuracy of all federated optimizers and brings most close to the best centralized accuracy.The gap is not fully closed, and the added benefit appears modest when communication is limited.

F CREATING A FEDERATED CIFAR-100

The paper constructs a federated CIFAR-100 dataset by using CIFAR-100’s coarse-to-fine label hierarchy within a Pachinko Allocation Model. This produces heterogeneous client datasets with correlated label frequencies, limited label coverage, and no overlapping examples.

  • Dataset construction: CIFAR-100’s 20 coarse labels, each containing 5 fine labels, are represented as a rooted DAG for partitioning examples across clients.Fine labels are leaves and coarse labels are interior nodes.
  • Dataset construction: Each client is modeled as a document whose label distribution is sampled hierarchically from Dirichlet priors at the root and coarse-label nodes.Pachinko Allocation samples coarse and fine labels before selecting examples with the corresponding fine label.
  • Heterogeneity: A small root concentration α encourages clients to focus on a few coarse labels, while larger β makes fine-label distributions within each coarse group more concentrated.The hierarchical construction creates correlated fine-label frequencies within coarse-label groups.
  • Data allocation: Examples are sampled without replacement and removed from the dataset after assignment, ensuring that no two clients share an example.The procedure updates the relevant fine-label distribution when an example is exhausted.
  • Experimental setup: The partition uses N = 100 examples per client and M = 500 distinct training clients, with α = 0.1 and β = 10.These parameters are used in the federated CIFAR-100 experiments.
  • Heterogeneity: Clients contain only a fraction of the overall labels, and the number of unique labels varies across the 500 training clients.Most clients have between 2 and ... unique labels in the reported distribution.
Loading 2003.00295v5…