Adaptive Federated Optimization

Sashank ReddiZachary CharlesManzil ZaheerZachary GarrettKeith RushJakub KonečnýSanjiv KumarH. Brendan McMahan

article2020ICLR2,126 citations

Develops federated versions of adaptive optimizers including Adam, Adagrad, and Yogi, establishing non-convex convergence guarantees and demonstrating significant performance improvements over standard FedAvg on heterogeneous client data.

Listen

Federated learning is an increasingly important machine learning framework that enables many distributed clients, such as mobile phones or enterprise nodes, to collaboratively train a shared model without transmitting their raw private data to a central server. However, standard optimization algorithms like Federated Averaging often face severe convergence issues, high sensitivity to hyperparameter tuning, and performance degradation. These challenges are driven by data heterogeneity across devices and the presence of sparse or heavy-tailed gradient noise, which is common in complex text and image tasks.

The article establishes a generalized optimization framework that incorporates adaptive optimization methods on the server side to accelerate model convergence and improve performance across heterogeneous client networks without inflating device-level communication or memory burdens.

To evaluate this framework, the authors developed federated adaptations of standard adaptive optimizers—including Federated Adagrad, Federated Adam, and Federated Yogi—by applying coordinate-wise adaptive updates at the server while maintaining standard stochastic gradient descent on local clients. The authors supported this approach with theoretical convergence proofs in general non-convex settings and conducted empirical simulations across seven benchmark tasks and five datasets, encompassing both image and text domains under realistic, non-uniform data distributions.

The analysis produced several critical findings:

Server-side adaptive optimizers consistently outperformed traditional non-adaptive methods, achieving superior accuracy across all tested dense and sparse tasks.
In sparse text-processing tasks, adaptive methods achieved dramatic performance gains; for example, in Stack Overflow tag prediction, Federated Adagrad reached a Recall@5 of 67.1% compared to 30.0% for standard Federated Averaging.
For complex tasks such as image classification and next-word prediction, combining adaptivity with server-side momentum was essential to achieving fast initial convergence and high final accuracy.
Server-side adaptivity significantly streamlined the tuning process by remaining robust across broad ranges of client and server learning rates, whereas traditional methods required precise, simultaneous tuning of multiple hyperparameters.
The state-based variance reduction baseline, SCAFFOLD, performed poorly in cross-device settings where clients participate infrequently, confirming that maintaining client state is impractical in large-scale mobile regimes.

These findings demonstrate that placing adaptive algorithms entirely on the central server captures the performance and convergence benefits of adaptive optimization while preserving the lightweight communication and storage footprints essential for cross-device deployments. Furthermore, the framework lowers operational risks and deployment costs by reducing the extensive trial-and-error traditionally required to tune federated learning rates.

Organizations developing federated machine learning systems should adopt server-side adaptive optimizers—such as Federated Adam or Federated Yogi—over standard Federated Averaging, especially for applications dealing with language processing or heterogeneous user behaviors. Engineering teams should pair these optimizers with decaying client learning rates over training rounds to maximize final accuracy. Next steps include conducting pilot testing within operational edge networks and researching how server-side adaptivity interacts with user differential privacy constraints and model fairness.

Confidence in these findings is high for standard supervised image and language tasks under moderate client heterogeneity, as results are grounded in formal mathematical proofs and diverse benchmarks. However, decision-makers should note that theoretical bounds rely on standard smoothness and bounded variance assumptions, and practical performance in production will remain subject to real-world edge network dropouts, variable compute capacities, and severe data non-uniformity.

Google communication federated learning optimization Learning Rate Client Client Learning Rate

arXiv: 2003.00295 tensorflow/federated

Paper: Adam: A Method for Stochastic Optimization, Diederik P. Kingma et al. (2015). Introduces Adam, the foundational adaptive moment estimation optimizer whose per-coordinate scaling principles are directly adapted to the server-side update in federated learning.
Paper: Adaptive Subgradient Methods for Online Learning and Stochastic Optimization, John Duchi et al. (2011). Presents AdaGrad, establishing the core adaptive subgradient paradigm that serves as a baseline optimizer adapted in this work for federated settings.
Paper: On the Convergence of Adam and Beyond, Sashank J. Reddi et al. (2018). Analyzes the convergence pitfalls of standard adaptive methods like Adam and introduces Yogi/AMSGrad principles that the source builds on for stable non-convex federated optimization.
Paper: Federated Optimization in Heterogeneous Networks, Tian Li et al. (2018). Provides the foundational analysis of federated optimization under device and data heterogeneity (FedProx), which the source aims to improve using adaptive server methods.
Paper: SCAFFOLD: Stochastic Controlled Averaging for Federated Learning, Sai Praneeth Karimireddy et al. (2019). Analyzes client drift and convergence bottlenecks in standard federated averaging under non-IID data, providing key motivation for developing adaptive federated algorithms.
Paper: On the Convergence of FedAvg on Non-IID Data, Xiang Li et al. (2019). Establishes convergence guarantees and rate limits of Federated Averaging (FedAvg) on heterogeneous data, illustrating the optimization hurdles addressed by the source.
Paper: Federated Optimization: Distributed Machine Learning for On-Device Intelligence, Jakub Konečný et al. (2016). Introduces core federated optimization concepts, non-IID data challenges, and communication-efficient local update frameworks.
Paper: An overview of gradient descent optimization algorithms, Sebastian Ruder (2016). Provides a comprehensive survey of first-order adaptive gradient descent methods (such as Adagrad, RMSprop, and Adam) that form the building blocks of federated adaptive optimization.

Paper: Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization, Jianyu Wang et al. (2020). Extends federated optimization theory by analyzing objective inconsistency across heterogeneous clients when using various local solvers, including momentum and adaptive updates.
Paper: The future of digital health with federated learning, Nicola Rieke et al. (2020). Surveys real-world medical applications and challenges of federated learning, demonstrating where advanced federated optimization techniques are deployed across heterogeneous clinical silos.

Cover for Adaptive Federated Optimization

Abstract

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.

Topics

communication federated learning optimization optimization problems application servers

Keywords

Learning Rate Client Client Learning Rate Server Learning Rate log10 Adaptive Federated Clients FEDAVG FEDADAM

1 Introduction
2 Federated Learning and FedAvg
3 Adaptive Federated Optimization
4 Experimental Evaluation: Datasets, Tasks, and Methods
5 Experimental Evaluation: Results
5.1 Comparisons between methods
5.2 Ease of tuning
5.3 Other Findings
6 Conclusion
References
A Proof of results
A.1 Main Challenges
A.2 Proof of Theorem
A.2.1 Limited Participation
A.3 Proof of Theorem
A.4 Auxiliary Lemmatta
B Federated Algorithms: Implementations and Practical Considerations
B.1 FedAvg and FedOpt
B.2 SCAFFOLD
B.3 Lookahead, AdaAlter, and client adaptivity
C Dataset & Models
C.1 CIFAR-10/CIFAR-100
C.2 EMNIST
C.3 Shakespeare
C.4 Stack Overflow
D Experiment Hyperparameters
D.1 Hyperparameter tuning
D.2 Hyperparameter grids
D.3 Per-task batch sizes
D.4 Best performing hyperparameters
E Additional Experimental Results
E.1 Results on EMNIST CR
E.2 Stack Overflow test set performance
E.3 Learning rate robustness
E.4 On the relation between client and server learning rates
E.5 Robustness of the adaptivity parameter
E.6 Improving performance with learning rate decay
F Creating a Federated CIFAR-100

Knowls

Knowl 1 — FedOpt Algorithmic Framework
model/method
The $\text{FedOpt}$ framework generalizes federated optimization by decoupling local client optimization from global server aggregation. In standard federated learning, the goal is to solve:

$\min_{x \in \mathbb{R}^d} f(x) = \frac{1}{m} \sum_{i=1}^m F_i(x)$

where $m$ is the total number of clients, $F_i(x) = \mathbb{E}_{z \sim \mathcal{D}_i}[f_i(x, z)]$ is the local objective function for client $i$ , and $\mathcal{D}_i$ is the local data distribution.

In $\text{FedOpt}$ , at each communication round $t \in \{0, \dots, T-1\}$ :
1. The server broadcasts the current global model $x_t$ to an active subset of sampled clients $S \subseteq \{1, \dots, m\}$ .
2. Each client $i \in S$ initializes its local model $x_{i,0}^t = x_t$ and runs $K$ iterations of a client optimizer $\text{ClientOpt}$ (such as SGD) with local learning rate $\eta_l$ on its local data:
$x_{i,k+1}^t = \text{ClientOpt}(x_{i,k}^t, g_{i,k}^t, \eta_l, t)$

where $g_{i,k}^t$ is an unbiased stochastic gradient estimate of $\nabla F_i(x_{i,k}^t)$ . 3. Each client computes its total model update $\Delta_i^t = x_{i,K}^t - x_t$ and transmits it back to the server. 4. The server computes the average client update (or a sample-weighted average):

$\Delta_t = \frac{1}{|S|} \sum_{i \in S} \Delta_i^t$
1. The server treats $-\Delta_t$ as a global pseudo-gradient and updates the global model using a gradient-based server optimizer $\text{ServerOpt}$ with server learning rate $\eta$ :
$x_{t+1} = \text{ServerOpt}(x_t, -\Delta_t, \eta, t)$

Standard $\text{FedAvg}$ is recovered as a special case where $\text{ClientOpt}$ is SGD, $\text{ServerOpt}$ is SGD, and $\eta = 1$ .
Knowl 2 — Adaptive Federated Optimization Algorithms (FedAdagrad, FedYogi, FedAdam)
algorithm
Adaptive federated optimization algorithms instantiate $\text{FedOpt}$ by using mini-batch SGD on the clients and an adaptive coordinate-wise optimizer on the server. Because the optimizer state (such as momentum and second-moment accumulators) is maintained solely on the central server, client memory and communication bandwidth requirements match those of $\text{FedAvg}$ , ensuring compatibility with cross-device federated learning.
Input: Initial global model $x_0 \in \mathbb{R}^d$ , initial second-moment accumulator $v_{-1} \ge \tau^2 \mathbf{1}$ , adaptivity parameter $\tau > 0$ , client learning rate $\eta_l$ , server learning rate $\eta$ , decay parameters $\beta_1, \beta_2 \in [0, 1)$
for $t = 0, \dots, T - 1$ do
    Sample a subset $S$ of clients
    for each client $i \in S$ in parallel do
         $x_{i,0}^t = x_t$
        for $k = 0, \dots, K - 1$ do
            Compute unbiased stochastic gradient estimate $g_{i,k}^t$ of $\nabla F_i(x_{i,k}^t)$
             $x_{i,k+1}^t = x_{i,k}^t - \eta_l g_{i,k}^t$
         $\Delta_i^t = x_{i,K}^t - x_t$
     $\Delta_t = \frac{1}{|S|} \sum_{i \in S} \Delta_i^t$
     $m_t = \beta_1 m_{t-1} + (1 - \beta_1) \Delta_t$
    if FedAdagrad then
         $v_t = v_{t-1} + \Delta_t^2$
    else if FedYogi then
         $v_t = v_{t-1} - (1 - \beta_2) \Delta_t^2 \odot \text{sign}(v_{t-1} - \Delta_t^2)$
    else if FedAdam then
         $v_t = \beta_2 v_{t-1} + (1 - \beta_2) \Delta_t^2$
     $x_{t+1} = x_t + \eta \frac{m_t}{\sqrt{v_t} + \tau}$
In practical implementations with non-uniform client data sizes, the server update computes $\Delta_t = \sum_{i \in S} \frac{n_i}{n} \Delta_i^t$ with $n = \sum_{i \in S} n_i$ , where $n_i$ is the number of examples on client $i$ . Standard hyperparameters are $\beta_1 = 0.9$ , $\beta_2 = 0.99$ for $\text{FedAdam}$ and $\text{FedYogi}$ , and $\beta_1 = 0$ for $\text{FedAdagrad}$ .
Knowl 3 — Nonconvex Smooth Heterogeneous Optimization Assumptions
assumption
The theoretical convergence analysis of adaptive federated optimization relies on the following standard assumptions on client loss functions $F_i(x) = \mathbb{E}_{z \sim \mathcal{D}_i}[f_i(x, z)]$ and the global objective $f(x) = \frac{1}{m} \sum_{i=1}^m F_i(x)$ for all $x, y \in \mathbb{R}^d$ , $j \in \{1, \dots, d\}$ , and $i \in \{1, \dots, m\}$ :
1. Lipschitz Gradient ( $L$ -Smoothness): Each client loss function $F_i$ is $L$ -smooth, satisfying
$\|\nabla F_i(x) - \nabla F_i(y)\| \le L \|x - y\|$
1. Bounded Local Variance: The stochastic gradients on each client have bounded coordinate variance $\sigma_{l,j}^2$ , satisfying
$\mathbb{E}\left[ \left( \nabla [f_i(x, z)]_j - [\nabla F_i(x)]_j \right)^2 \right] = \sigma_{l,j}^2$

with total local variance $\sigma_l^2 = \sum_{j=1}^d \sigma_{l,j}^2$ .
1. Bounded Global Variance (Client Heterogeneity): The variation between client gradient coordinates and the global objective gradient coordinate is bounded by $\sigma_{g,j}^2$ , satisfying
$\frac{1}{m} \sum_{i=1}^m \left( [\nabla F_i(x)]_j - [\nabla f(x)]_j \right)^2 \le \sigma_{g,j}^2$

with total global variance $\sigma_g^2 = \sum_{j=1}^d \sigma_{g,j}^2$ . Here, $\sigma_g = 0$ corresponds to the i.i.d. client data setting.
1. Bounded Gradients: For any client $i$ , parameter $x \in \mathbb{R}^d$ , sample $z$ , and coordinate $j \in \{1, \dots, d\}$ , the gradient magnitude is bounded:
$|[\nabla f_i(x, z)]_j| \le G$
Knowl 4 — Nonconvex Convergence of FedAdagrad
theoretical result
Under the $L$ -smoothness, bounded variance, and bounded gradient assumptions with full client participation ( $S = \{1, \dots, m\}$ ), $\beta_1 = 0$ , and $\sigma^2 = \sigma_l^2 + 6K\sigma_g^2$ , if the client learning rate satisfies:

$\eta_l \le \frac{1}{16K} \min \left\{ \frac{1}{L}, \frac{1}{T^{1/6}} \left( \frac{\tau}{120 L^2 G} \right)^{1/3} \right\}$

and

$\eta_l \le \frac{1}{16K} \min \left\{ \frac{\tau \eta L}{2G^2}, \frac{\tau}{4L\eta}, \frac{1}{T^{1/4}} \left( \frac{\tau^2}{G L \eta} \right)^{1/2} \right\}$

then the iterates of $\text{FedAdagrad}$ satisfy:

$\min_{0 \le t \le T - 1} \mathbb{E}\|\nabla f(x_t)\|^2 \le \mathcal{O}\left( \left( \frac{G}{\sqrt{T}} + \frac{\tau}{\eta_l K T} \right) (\Psi + \widetilde{\Psi}_{\text{var}}) \right)$

where

$\Psi = \frac{f(x_0) - f(x^*)}{\eta} + \frac{5\eta_l^3 K^2 L^2 T \sigma^2}{2\tau}$

$\widetilde{\Psi}_{\text{var}} = \frac{2\eta_l K G^2 + \tau \eta L}{\tau^2} \left( \frac{2\eta_l^2 K T \sigma_l^2}{m} + 10\eta_l^4 K^3 L^2 T \sigma^2 \right)$

Setting $\eta_l = \Theta(1 / (KL\sqrt{T}))$ , $\eta = \Theta(\sqrt{Km})$ , and $\tau = G/L$ , the convergence bound for sufficiently large $T$ evaluates to:

$\min_{0 \le t \le T-1} \mathbb{E}\|\nabla f(x_t)\|^2 = \mathcal{O}\left( \frac{f(x_0) - f(x^*)}{\sqrt{mKT}} + \frac{2\sigma_l^2 L}{G^2 \sqrt{mKT}} + \frac{\sigma^2}{GKT} + \frac{\sigma^2 L \sqrt{m}}{G^2 \sqrt{K T^3}} \right)$

For large $T$ , the dominant term is $\mathcal{O}(1/\sqrt{mKT})$ , matching the optimal rate for distributed non-convex optimization.
Knowl 5 — Nonconvex Convergence of FedAdam and FedYogi
theoretical result
Under the $L$ -smoothness, bounded variance, and bounded gradient assumptions with full client participation ( $S = \{1, \dots, m\}$ ), $\beta_1 = 0$ , and $\sigma^2 = \sigma_l^2 + 6K\sigma_g^2$ , if the client learning rate satisfies $\eta_l \le \frac{1}{16LK}$ and:

$\eta_l \le \frac{1}{16K} \min \left\{ \left( \frac{\tau}{120 L^2 G} \right)^{1/3}, \frac{\tau}{2(2G + \eta L)} \right\}$

then the iterates of $\text{FedAdam}$ (and similarly $\text{FedYogi}$ ) satisfy:

$\min_{0 \le t \le T-1} \mathbb{E}\|\nabla f(x_t)\|^2 = \mathcal{O}\left( \frac{\sqrt{\beta_2} \eta_l K G + \tau}{\eta_l K T} (\Psi + \Psi_{\text{var}}) \right)$

where

$\Psi = \frac{f(x_0) - f(x^*)}{\eta} + \frac{5\eta_l^3 K^2 L^2 T \sigma^2}{2\tau}$

$\Psi_{\text{var}} = \left( G + \frac{\eta L}{2} \right) \left[ \frac{4\eta_l^2 K T \sigma_l^2}{m\tau^2} + \frac{20\eta_l^4 K^3 L^2 T \sigma^2}{\tau^2} \right]$

When parameterized with $\eta_l = \Theta(1 / (KL\sqrt{T}))$ , $\eta = \Theta(\sqrt{Km})$ , and $\tau = G/L$ , the convergence rate for sufficiently large $T$ is:

$\min_{0 \le t \le T-1} \mathbb{E}\|\nabla f(x_t)\|^2 = \mathcal{O}\left( \frac{f(x_0) - f(x^*)}{\sqrt{mKT}} + \frac{2\sigma_l^2 L}{G^2 \sqrt{mKT}} + \frac{\sigma^2}{GKT} + \frac{\sigma^2 L \sqrt{m}}{G^2 \sqrt{K T^3}} \right)$
Knowl 6 — Interplay Between Client Heterogeneity, Local Steps, and Step Sizes
theoretical result
The theoretical analysis reveals three structural dynamics in adaptive federated optimization:
1. Communication Efficiency vs. Local Steps: To preserve the asymptotic optimal convergence rate $\mathcal{O}(1/\sqrt{mKT})$ , the number of local steps $K$ per communication round can be scaled as long as:
$K = \mathcal{O}\left( T \frac{\sigma_l^2}{\sigma_g^2} \right)$

When data is homogeneous ( $\sigma_g = 0$ ), $K$ can be arbitrarily large without degrading convergence. In heterogeneous settings ( $\sigma_g > 0$ ), client drift constrains $K$ .
1. Learning Rate Decay: Client drift introduced by local steps does not vanish if $\eta_l$ is constant as $T \to \infty$ . Therefore, the client learning rate $\eta_l$ must be decayed at a rate of $\mathcal{O}(1/\sqrt{t})$ to ensure asymptotic convergence. In contrast, the server learning rate $\eta$ does not need to decay.
2. Client vs. Server Learning Rate Balancing: The server learning rate $\eta = 1$ typically used in $\text{FedAvg}$ is suboptimal in heterogeneous settings. Setting a larger server learning rate $\eta = \Theta(\sqrt{Km})$ alongside a smaller client step size $\eta_l = \Theta(1/(KL\sqrt{T}))$ reduces the negative impact of client heterogeneity $\sigma_g$ on optimization convergence.

Knowl 7 — Empirical Performance Across Federated Benchmark Tasks

data/table

Experimental evaluation across seven federated learning tasks demonstrates that server-side adaptive optimizers ( $\text{FedAdagrad}$ , $\text{FedAdam}$ , $\text{FedYogi}$ ) consistently match or outperform non-adaptive baselines ( $\text{FedAvg}$ , $\text{FedAvgM}$ ) and client control variate methods ( $\text{SCAFFOLD}$ ).

Task	FedAdagrad	FedAdam	FedYogi	FedAvgM	FedAvg
CIFAR-10 (% Acc)	72.1	77.4	78.0	77.4	72.8
CIFAR-100 (% Acc)	47.9	52.5	52.4	52.4	44.7
EMNIST CR (% Acc)	85.1	85.6	85.5	85.2	84.9
Shakespeare (% Acc)	57.5	57.0	57.2	57.3	56.9
Stack Overflow NWP (% Acc)	23.8	25.2	25.2	23.8	19.5
Stack Overflow LR (Recall@5 $\times 100$ )	67.1	65.8	65.9	36.9	30.0
EMNIST AE (MSE $\times 1000$ )	4.20	1.01	0.98	1.65	6.47

Performance is evaluated as the average validation metric over the final 100 training rounds, with bold numbers indicating performance within 0.5% of the best result. Key takeaways include:

Sparse-Gradient Tasks: On Stack Overflow Tag Prediction (SO LR) and Next Word Prediction (SO NWP), adaptive methods achieve dramatic improvements (e.g., Recall@5 improves from 30.0/36.9 to 67.1). Infrequent words produce sparse updates across clients; adaptive coordinate scaling allows meaningful updates on rare features while dampening high-frequency noise.
Vision and Autoencoder Tasks: $\text{FedAdam}$ and $\text{FedYogi}$ outperform non-adaptive methods by large margins on CIFAR-100 (52.5% vs 44.7%) and EMNIST Autoencoder reconstruction error (0.98 vs 6.47 MSE $\times 1000$ ).
Cross-Device Incompatibility of Control Variates: $\text{SCAFFOLD}$ underperforms in realistic cross-device regimes because client control variates become stale when only a small fraction of clients participate per round.

Knowl 8 — Hyperparameter Robustness and Ease of Tuning in Adaptive FL
empirical result
Systematic grid-search evaluations over client learning rate $\eta_l$ , server learning rate $\eta$ , and adaptivity parameter $\tau$ demonstrate distinct tuning properties between adaptive and non-adaptive federated optimizers:
1. Rectangular vs. Triangular Learning Rate Regions: In grid evaluations plotting validation accuracy across combinations of $(\log_{10} \eta_l, \log_{10} \eta)$ , non-adaptive methods ( $\text{FedAvg}$ , $\text{FedAvgM}$ ) exhibit narrow triangular regions of viable hyperparameters, indicating that $\eta_l$ and $\eta$ are tightly coupled and must be tuned simultaneously. In contrast, adaptive server methods ( $\text{FedAdam}$ , $\text{FedYogi}$ ) exhibit broad rectangular regions of high performance, demonstrating robustness across varied client learning rates for a fixed server learning rate.
2. Universal Adaptivity Threshold $\tau$ : Across image and text tasks, performance of $\text{FedAdam}$ , $\text{FedYogi}$ , and $\text{FedAdagrad}$ is stable across several orders of magnitude of $\tau \in [10^{-5}, 10^{-1}]$ . Setting $\tau = 10^{-3}$ achieves near-optimal performance across virtually all tasks, removing the necessity of tuning $\tau$ extensively.
Knowl 9 — Creating Heterogeneous Federated Datasets Using Pachinko Allocation
algorithm
To generate realistic hierarchical non-i.i.d. label distributions for federated benchmarking (such as federated CIFAR-100), the Pachinko Allocation Method (PAM) is applied to partition examples without replacement over a directed acyclic graph (DAG) representing coarse and fine label hierarchies.
Input: Target clients $M$ , samples per client $N$ , Dirichlet parameters $\alpha, \beta \in \mathbb{R}_{\ge 0}$ , dataset $S = \{(x, c, y)\}$ with coarse labels $c \in C$ and fine labels $y \in Y$ , DAG $G$ rooted at $r$ with children $G[r] = C$ and $G[c] = \{y \in Y : y \in c\}$
for $m = 1, \dots, M$ do
    Sample root coarse multinomial $\theta_r \sim \text{Dir}(\alpha, |G[r]|)$
    for each coarse topic $c \in C \cap G[r]$ do
        Sample fine multinomial $\theta_c \sim \text{Dir}(\beta, |G[c]|)$
     $D_m = \emptyset$
    for $n = 1, \dots, N$ do
        Sample coarse label $c \sim \text{Multinomial}(\theta_r)$
        Sample fine label $y \sim \text{Multinomial}(\theta_c)$
        Select $(x, c, y) \in S$ uniformly at random with label $y$
         $D_m = D_m \cup \{(x, c, y)\}$
         $S = S \setminus \{(x, c, y)\}$
        if $S_y = \emptyset$ then
             $G = G \setminus \{y\}$
            Renormalize $\theta_c$ over remaining active fine children in $G[c]$
            if $S_c = \emptyset$ then
                 $G = G \setminus \{c\}$
                Renormalize $\theta_r$ over remaining active coarse topics in $G[r]$
Using $\alpha = 0.1$ , $\beta = 10$ , $M = 500$ , and $N = 100$ , PAM creates correlated label partitions where each client concentrates on a few coarse topics while drawing multiple correlated fine labels within each topic.

Knowl 10 — Client Learning Rate Schedules Narrow the Centralized-Federated Gap

empirical result

In federated optimization, constant client step sizes $\eta_l$ produce persistent asymptotic client drift that prevents federated models from matching centralized model performance.

On the EMNIST Character Recognition task, centralized training achieves $88.0\%$ accuracy. Under constant client step sizes $\eta_l$ , federated optimizers saturate between $84.9\%$ ( $\text{FedAvg}$ ) and $85.6\%$ ( $\text{FedAdam}$ ). Introducing an exponential staircase decay schedule ( $\text{ExpDecay}$ ), where the client step size $\eta_l$ is decayed by a factor of $0.1$ every $500$ communication rounds, significantly narrows this gap:

Setting	FedAdagrad	FedAdam	FedYogi	FedAvgM	FedAvg
Centralized Baseline	88.0	87.9	88.0	87.7	87.7
Constant $\eta_l$	85.1	85.6	85.5	85.2	84.9
ExpDecay $\eta_l$	85.3	86.2	86.2	85.8	85.2

Decaying the client learning rate over rounds directly dampens the local update variance and drift terms, improving test accuracy across all optimizers.

Coverage note — None was omitted; all central theoretical rates, algorithms (FedOpt framework, FedAdagrad/FedAdam/FedYogi, PAM dataset construction), empirical comparisons, tuning landscapes, and learning rate schedule findings have been fully captured.

References

1.The TensorFlow Federated Authors. TensorFlow Federated Stack Overflow dataset, 2019. URL https://www.tensorflow.org/federated/api_docs/python/tff/simulation/datasets/stackoverflow/load_data.
2.Debraj Basu, Deepesh Data, Can Karakus, and Suhas Diggavi. Qsparse-local-SGD: Distributed SGD with quantization, sparsification and local computations. In Advances in Neural Information Processing Systems, pp. 14668–14679, 2019.
3.Keith Bonawitz, Hubert Eichner, Wolfgang Grieskamp, Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloé Kiddon, Jakub Konečný, Stefano Mazzocchi, Brendan McMahan, Timon Van Overveldt, David Petrou, Daniel Ramage, and Jason Roselander. Towards federated learning at scale: System design. In A. Talwalkar, V. Smith, and M. Zaharia (eds.), Proceedings of Machine Learning and Systems, volume 1, pp. 374–388. Proceedings of MLSys, 2019. URL https://proceedings.mlsys.org/paper/2019/file/bd686fd640be98efaae0091fa301e613-Paper.pdf.
4.Sebastian Caldas, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. LEAF: A benchmark for federated settings. arXiv preprint arXiv:1812.01097, 2018.
5.Gregory Cohen, Saeed Afshar, Jonathan Tapson, and Andre Van Schaik. EMNIST: Extending MNIST to handwritten letters. In 2017 International Joint Conference on Neural Networks (IJCNN), pp. 2921–2926. IEEE, 2017.
6.Aaron Defazio and Léon Bottou. On the ineffectiveness of variance reduced optimization for deep learning. arXiv preprint arXiv:1812.04529, 2018.
7.Aaron Defazio, Francis Bach, and Simon Lacoste-Julien. SAGA: A fast incremental gradient method with support for non-strongly convex composite objectives. In NIPS, pp. 1646–1654, 2014.
8.John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(Jul):2121–2159, 2011.
9.Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training ImageNet in 1 hour. arXiv preprint arXiv:1706.02677, 2017.
10.Kevin Hsieh, Amar Phanishayee, Onur Mutlu, and Phillip B Gibbons. The non-IID data quagmire of decentralized machine learning. arXiv preprint arXiv:1910.00189, 2019.
11.Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.06335, 2019.
12.Alex Ingerman and Krzys Ostrowski. Introducing TensorFlow Federated, 2019. URL https://medium.com/tensorflow/introducing-tensorflow-federated-a4147aa20041.
13.Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. In Advances in Neural Information Processing Systems, pp. 315–323, 2013.
14.Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. arXiv preprint arXiv:1912.04977, 2019.
15.Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for on-device federated learning. arXiv preprint arXiv:1910.06378, 2019.
16.Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. First analysis of local GD on heterogeneous data. arXiv preprint arXiv:1909.04715, 2019.
17.Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
18.Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical report, Citeseer, 2009.
19.Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. arXiv preprint arXiv:1812.06127, 2018.
20.Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. arXiv preprint arXiv:1908.07873, 2019a.
21.Wei Li and Andrew McCallum. Pachinko allocation: DAG-structured mixture models of topic correlations. In Proceedings of the 23rd international conference on Machine learning, pp. 577–584, 2006.
22.Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of FedAvg on non-IID data. arXiv preprint arXiv:1907.02189, 2019b.
23.Xiaoyu Li and Francesco Orabona. On the convergence of stochastic gradient descent with adaptive stepsizes. arXiv preprint arXiv:1805.08114, 2018.
24.Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun. Adaptive gradient methods with dynamic bound of learning rate. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL https://openreview.net/forum?id=Bkg3g2R9FX.
25.Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, AISTATS 2017, 20-22 April 2017, Fort Lauderdale, FL, USA, pp. 1273–1282, 2017. URL http://proceedings.mlr.press/v54/mcmahan17a.html.
26.H. Brendan McMahan and Matthew Streeter. Adaptive bound optimization for online convex optimization. In COLT, 2010a.
27.H. Brendan McMahan and Matthew J. Streeter. Adaptive bound optimization for online convex optimization. In COLT The 23rd Conference on Learning Theory, 2010b.
28.Sashank J Reddi, Ahmed Hefny, Suvrit Sra, Barnabás Póczós, and Alex Smola. Stochastic variance reduction for nonconvex optimization. arXiv:1603.06160, 2016.
29.Sashank J Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of ADAM and beyond. arXiv preprint arXiv:1904.09237, 2019.
30.Sebastian U. Stich. Local SGD converges fast and communicates little. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id=S1g2JnRcFX.
31.Sebastian U Stich and Sai Praneeth Karimireddy. The error-feedback framework: Better rates for SGD with delayed gradients and compressed communication. arXiv preprint arXiv:1909.05350, 2019.
32.Jianyu Wang and Gauri Joshi. Cooperative SGD: A unified framework for the design and analysis of communication-efficient SGD algorithms. arXiv preprint arXiv:1808.07576, 2018.
33.Shiqiang Wang, Tiffany Tuor, Theodoros Salonidis, Kin K Leung, Christian Makaya, Ting He, and Kevin Chan. Adaptive federated learning in resource constrained edge computing systems. IEEE Journal on Selected Areas in Communications, 37(6):1205–1221, 2019.
34.Rachel Ward, Xiaoxia Wu, and Leon Bottou. Adagrad stepsizes: Sharp convergence over nonconvex landscapes, from any initialization. arXiv preprint arXiv:1806.01811, 2018.
35.Xiaoxia Wu, Simon S Du, and Rachel Ward. Global convergence of adaptive gradient methods for an over-parameterized neural network. arXiv preprint arXiv:1902.07111, 2019.
36.Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 3–19, 2018.
37.Cong Xie, Oluwasanmi Koyejo, Indranil Gupta, and Haibin Lin. Local AdaAlter: Communication-efficient stochastic gradient descent with adaptive learning rates. arXiv preprint arXiv:1911.09030, 2019.
38.Hao Yu, Sen Yang, and Shenghuo Zhu. Parallel restarted SGD with faster convergence and less communication: Demystifying why model averaging works for deep learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 5693–5700, 2019.
39.Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. In Advances in Neural Information Processing Systems, pp. 9815–9825, 2018.
40.Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank J. Reddi, Sanjiv Kumar, and Suvrit Sra. Why ADAM beats SGD for attention models. arXiv preprint arxiv:1912.03194, 2019a.
41.Michael R. Zhang, James Lucas, Jimmy Ba, and Geoffrey E. Hinton. Lookahead optimizer: k steps forward, 1 step back. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d’Alché-Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, 8-14 December 2019, Vancouver, BC, Canada, pp. 9593–9604, 2019b.
42.Martin Zinkevich, Markus Weimer, Lihong Li, and Alex J Smola. Parallelized stochastic gradient descent. In Advances in neural information processing systems, pp. 2595–2603, 2010.

Citation

MLA

Reddi, S., et al. “Adaptive Federated Optimization”. arXiv, 2020, http://arxiv.org/abs/2003.00295v5.

APA

Reddi, S., Charles, Z., Zaheer, M., Garrett, Z., Rush, K., Konečný, J., Kumar, S., & McMahan, H. B. (2020). Adaptive Federated Optimization. arXiv. http://arxiv.org/abs/2003.00295v5

Chicago

Reddi, S., Z. Charles, M. Zaheer, et al. 2020. “Adaptive Federated Optimization”. arXiv. http://arxiv.org/abs/2003.00295v5.

Harvard

Reddi, S. et al. (2020) “Adaptive Federated Optimization”, arXiv [Preprint]. Available at: http://arxiv.org/abs/2003.00295v5.

Vancouver

1. Reddi S, Charles Z, Zaheer M, Garrett Z, Rush K, Konečný J, Kumar S, McMahan HB (2020) Adaptive Federated Optimization. arXiv

BibTeX

@article{reddi2020adaptive,
  title = {Adaptive Federated Optimization},
  author = {Reddi, Sashank and Charles, Zachary and Zaheer, Manzil and Garrett, Zachary and Rush, Keith and Konečný, Jakub and Kumar, Sanjiv and McMahan, H. Brendan},
  year = {2020},
  journal = {arXiv},
  url = {http://arxiv.org/abs/2003.00295v5},
  eprint = {2003.00295}
}

Metadata:arXiv

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: Published with permission

Adaptive Federated Optimization

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — FedOpt Algorithmic Framework

Knowl 2 — Adaptive Federated Optimization Algorithms (FedAdagrad, FedYogi, FedAdam)

Knowl 3 — Nonconvex Smooth Heterogeneous Optimization Assumptions

Knowl 4 — Nonconvex Convergence of FedAdagrad

Knowl 5 — Nonconvex Convergence of FedAdam and FedYogi

Knowl 6 — Interplay Between Client Heterogeneity, Local Steps, and Step Sizes