---
abstract: 'We study a family of algorithms, which we refer to as *local update* methods, that generalize many federated learning and meta-learning algorithms. We prove that for quadratic objectives, local update methods perform stochastic gradient descent on a *surrogate* loss function which we exactly characterize. We show that the choice of client learning rate controls the condition number of that surrogate loss, as well as the distance between the minimizers of the surrogate and true loss functions. We use this theory to derive novel convergence rates for federated averaging that showcase this trade-off between the condition number of the surrogate loss and its alignment with the true loss function. We validate our results empirically, showing that in communication-limited settings, proper learning rate tuning is often sufficient to reach near-optimal behavior. We also present a practical method for automatic learning rate decay in local update methods that helps reduce the need for learning rate tuning, and highlight its empirical performance on a variety of tasks and datasets.'
author:
- |
    Zachary Charles kyleglover@example.com\
    Jakub Konečný kyleglover@example.com\
    Google Research
bibliography:
- 'main.bib'
title: |
    On the Outsized Importance of Learning Rates\
    in Local Update Methods
---

Local Update Methods, Local SGD, Federated Averaging, Federated Learning, Meta Learning, MAML

Introduction
============

Historically, machine learning was analyzed from a "centralized" perspective, in which a model is trained on a single central source of data. In recent years, there has been a shift away from centralized machine learning, due in part to the increase of user data and the increasing awareness of the risks to privacy that can accompany centralized data collection.

Federated learning (FL) [@kairouz2019advances] is a distributed framework for learning models without directly sharing user data. In this framework, heterogeneous clients all use their own data to perform local training. In the popular FedAvg algorithm [@mcmahan17fedavg], the client models are then averaged at a central server, broadcast to a (possibly different) sample of clients, and the process is repeated. The core tenet is that instead of having clients share data, we instead share the results of *local updates* the clients perform on their own datasets using an optimization algorithm.

While there has been growing interest in FL in research communities (see [@kairouz2019advances] and [@li2019federated] for surveys of many recent works and open problems), this general paradigm of performing local updates on heterogeneous datasets has a storied history in machine learning. In particular, much of the work on meta-learning has focused on trying to learn models that perform well (or can quickly learn how to perform well) on a large number of heterogeneous tasks. This similarity with FL is even more clear in work on model-agnostic meta-learning (MAML) [@finn2017model], in which local client gradient updates are used to learn a global model. Connections between these two areas were noted by @jiang2019improving and have since been explored in many other works [@khodak2019adaptive; @fallah2020personalized].

While there is a wide variety of theoretical and empirical analyses of the aforementioned methods, it is generally difficult to understand their behavior in *heterogeneous settings*. There is enough evidence that these methods are useful in practice in complex scenarios [@hard2018federated; @yang2018applied; @hard2020training], yet on a theoretical level, many works derive results comparable to, or worse than, that of mini-batch SGD in heterogeneous or even homogeneous settings; See [@kairouz2019advances] for a discussion of homogeneity and heterogeneity, and see [@woodworth2020local] for a detailed discussion of comparisons to mini-batch SGD. Unfortunately, these results shed little light onto how methods such as FedAvg improve (or degrade) convergence.

In this work, we analyze a generalized local update paradigm that encompasses many FL and MAML methods, as well as other popular optimization methods such as mini-batch SGD. In order to better understand the structure of these methods in heterogeneous settings without an abundance of assumptions, we focus on the special case of quadratic loss functions. We are generally concerned with understanding the following questions that bridge both theory and practice.

-   How do local update methods improve or hinder convergence?

-   Why, despite a relative paucity of theoretical evidence, do these methods often perform better in practice than theoretically established methods such as mini-batch SGD?

-   What obstacles are there to the performance of local update methods, and how do we mitigate these issues?

As a partial answer to these questions, we highlight the main findings of our work.

1.  We show that in the quadratic case, local update methods are equivalent to the stochastic gradient method on a surrogate loss function which we exactly characterize. Thus, we can view local update methods that use multiple heterogeneous datasets as instead performing SGD on a single "central" loss function.

2.  We show that methods such as FedAvg and many incarnations of MAML implicitly regularize the condition number of this surrogate loss function, allowing for improved convergence of the surrogate loss. On the other hand, we show that this condition number reduction comes at the cost of increasing the discrepancy between minimizers of the surrogate and the true loss function. Notably, this trade-off is controlled by fundamental algorithmic choices, especially the choice of learning rate.

3.  We give explicit convergence rates for FedAvg that exhibit the trade-off between the condition number and the discrepancy between the surrogate and true loss functions above. Our results are similar in scope to work by @woodworth2020local (showing that local SGD can outperform mini-batch SGD), but work under heterogeneous data settings.

4.  We use our theoretical insights to design practical improvements to federated learning methods. First, we show that *decoupling client and server learning rates* has significant implications for improving convergence to better models. We show that despite the non-optimality of critical points of FedAvg, combining this learning rate decoupling with proper tuning can result in near-optimal performance in settings with limited communication. Finally, we detail a simple, practical method for automatic learning rate decay in federated learning that helps reduce the burden of learning rate tuning. We show empirically that this method improves the convergence of FedAvg, without requiring manually crafted learning rate schedules, across a suite of realistic and challenging non-convex tasks.

Related work
------------

#### Federated learning

Federated learning is a distributed machine learning paradigm in which training is done locally on clients, without any centralized data aggregation. Federated learning has enabled privacy-aware learning in a variety of applications [@hard2018federated; @chen2019federated; @brisimi2018federated; @samarakoon2018federated; @hard2020training], and has seen a large volume of work on the intersection of federated learning with topics including differential privacy [@mcmahan2017learning; @augenstein2020generative], fairness [@mohri2019agnostic; @li2019fair], robustness [@ghosh2019robust; @bagdasaryan2018backdoor; @sun2019can], and communication-efficiency [@konevcny2016federated; @sattler2019robust; @basu2019qsparse; @reisizadeh2019fedpaq]. For a more detailed discussion of federated learning, we defer to surveys by @kairouz2019advances and @li2019federated.

#### Meta-learning

In meta-learning (aka *learning to learn*), the objective is to use a collection of tasks to learn how to learn a new task efficiently [@vanschoren2019meta]. A particularly influential recent approach is model-agnostic meta-learning (MAML) proposed by @finn2017model. The core idea has inspired a number of extensions [@antoniou2018how; @nichol2018first; @rusu2018metalearning; @grant2018recasting; @rajeswaran2019meta; @raghu2019rapid], which broadly use a two-level optimization structure to perform meta-learning. Convergence properties of some of these optimization algorithms were recently studied by @fallah2019convergence, who also highlight differences in convergence of MAML and first-order approximations to MAML.

#### Federated optimization

One of the most common approaches to optimization in the setting of federated learning is the FedAvg method [@mcmahan17fedavg]. While designed for heterogeneous sources of data, the study of FedAvg has roots in that of Local SGD [@zinkevich2010parallelized; @stich2018local; @wang2018cooperative; @stich2019error; @yu2019parallel; @khaled2020tighter], a communication-efficient optimization method for homogeneous clients. As interest in federated learning has grown, so too has the number of proposed federated optimization methods. These can often be seen as variants of FedAvg, that incorporate techniques such as momentum [@hsu2019measuring], adaptive optimization [@reddi2020adaptive; @xie2019local], proximal updates [@li2018federated; @pathak2020fedsplit] and control variates [@karimireddy2019scaffold]. We again defer to @kairouz2019advances and @li2019federated for more detailed references.

#### Convergence (and non-convergence) of FedAvg

While we defer to @kairouz2019advances [Section 3.2] for a complete discussion of federated optimization, we discuss a few important connections. First, while there has been huge progress in theoretical understandings of FedAvg, existing works generally have not been able to show that these methods consistently improve upon mini-batch SGD [@woodworth2020local]. Even theoretically and empirically successful techniques such as SCAFFOLD [@karimireddy2019scaffold] have only been shown to converge faster than mini-batch SGD on quadratic objectives.

This failure of convergence was noted by @li2019convergence, who showed that without learning rate decay, FedAvg is not guaranteed to converge. Later, @karimireddy2019scaffold and @woodworth2020local showed that there are settings where FedAvg converges provably slower than mini-batch SGD. Similarly, @malinovsky2020local and @pathak2020fedsplit showed that in heterogeneous settings, FedAvg can converge to sub-optimal points, even in non-stochastic, strongly convex settings. @pathak2020fedsplit further give a proximal version of federated gradient descent that converges to the empirical risk minimizer in convex settings.

#### Comparisons to our work

Our work is most closely related to that of @malinovsky2020local, @pathak2020fedsplit. We also evince the non-convergence of FedAvg. However, we extend the analysis to stochastic settings, and to a more general class of algorithms that encompasses many meta-learning algorithms. As such, our work is also closely related to that of @fallah2019convergence, who demonstrated differences (and non-convergence issues) of various MAML algorithms. Our work takes this a step further, where we give a unified view of both MAML and federated learning methods, and give a broader characterization of the sub-optimal convergence of these methods in the case of quadratic losses. Our work is also novel in its focus on the interplay between convergence, suboptimality, and algorithmic choices, especially learning rates.

#### Notation

For a vector $v \in {\mathbb{R}}^d$, we let $\norm{v}$ denote its $\ell_2$ norm. For a matrix $A \in {\mathbb{R}}^{n\times m}$, we let $\norm{A}$ denote its operator norm with respect to the $\ell_2$ vector norm. For a symmetric positive semi-definite matrix $A$, we will let $A^{1/2}$ denote its matrix square root. For any real symmetric matrix $A$ (therefore with real eigenvalues), we will let $\lambda_{\max}(A)$ and $\lambda_{\min}(A)$ denote its largest and smallest eigenvalues, respectively.

Preliminaries
=============

Suppose we wish to learn a model $x \in {\mathbb{R}}^d$. Let ${\mathcal{I}}$ denote some collection of *clients*, and let $\mathcal{P}$ be a distribution on ${\mathcal{I}}$. For each $i \in {\mathcal{I}}$, we assume that there is an associated *data distribution* ${\mathcal{D}}_i$ on some example space ${\mathcal{Z}}$. For any $z \in {\mathcal{Z}}$, we assume there is a unique corresponding symmetric matrix $A_z \in {\mathbb{R}}^{d \times d}$ and vector $c_z \in {\mathbb{R}}^d$, and define a quadratic loss function $$\label{eq:quadratic_loss}
    f(x; z) := \frac{1}{2}\|A_z^{1/2}(x-c_z)\|^2.$$ We let $\nabla f(x; z)$ denote the gradient of the function $x \mapsto f(x; z)$ with respect to $x$. For $i \in {\mathcal{I}}$, we define the client loss function $f_i$ and the overall loss function $f$ as follows: $$\label{eq:objective}
f_i(x) := \operatorname*{\mathbb{E}}_{z \sim {\mathcal{D}}_i} [f(x ; z)],~~~f(x) := \operatorname*{\mathbb{E}}_{i \sim {\mathcal{P}}}[f_i(x)].$$ One common objective in our setup is to minimize $f(x)$, though this is often not the direct goal of MAML methods. Note that the joint distribution over $({\mathcal{I}}, {\mathcal{Z}})$ implicitly defines a (marginal) distribution over ${\mathcal{Z}}$, recovering standard risk minimization frameworks. This framework also encompasses distributed risk minimization in which ${\mathcal{P}}$ is a uniform distribution over a finite set of nodes $i \in {\mathcal{I}}$ and ${\mathcal{D}}_i$ is the uniform distribution over the (finite) dataset stored at node $i$. However, we take a more general approach and do not assume ${\mathcal{I}}$ or ${\mathcal{Z}}$ to be finite throughout. We also focus on the *heterogeneous* setting, where the client distributions ${\mathcal{D}}_i$ are not all identical, as opposed to the *homogeneous* setting, where all ${\mathcal{D}}_i$ are identical.

#### Modelling assumptions and relevance

As FL has matured, it has become more evident that there are two varieties, with distinct system-imposed constraints, recently termed by @kairouz2019advances as *cross-device federated learning* and *cross-silo federated learning*.[^1] The primary distinction between these two frameworks that is relevant to our work is that in cross-silo FL, there are relatively few participating clients. Moreover, these clients are typically reliable and almost always available. By contrast, in cross-device FL there are potentially very large numbers of clients, only a small fraction of which are available at any given point in time. Furthermore, the clients cannot be addressed directly or re-identified if participating multiple times. For a more detailed summary, see [@kairouz2019advances Table 1].

In cross-device FL, a client $i$ sampled from ${\mathcal{I}}$ corresponds to a single device, and ${\mathcal{D}}_i$ corresponds to the data available on that device. In many practical cross-device FL systems (see @bonawitz2019towards [@hard2018federated]), the server does not control the selection of clients from the global population ${\mathcal{I}}$. Instead, participation is initiated by the clients, based on pre-defined eligibility criteria, such as whether the device is charging and on unmetered wifi. Thus, the client distribution ${\mathcal{P}}$ can be considered as fixed, with only minor possibilities for it to be shaped by the server (e.g. whether to enforce sampling without replacement).

On the other hand, in many examples of cross-silo FL, participating clients correspond to various medical or financial organizations, or different geographical regions of the same organization [@nvidia_clara; @yang2019federated]. The participating clients are typically fixed in advance, and often all of them participate in every communication round. Thus, while cross-silo FL may be accurately described by a finite-sum optimization problem, this framework is less useful for cross-device FL.

To see this, consider the task of next word prediction on mobile devices. The FL training described by @hard2018federated runs for $3000$ communication rounds, with up to $500$ clients participating in each round. That is at most $1.5$ million distinct clients, a small fraction of the total number of possible clients[^2]. This also implies that it is nearly impossible to compute exact values of the loss $f(x)$. Instead, evaluation of a model's quality is done using the same mechanism as the training – by using a subset of the clients eligible at a given time – which has significant implications for algorithm design (as we discuss in Sections \[sec:exp\_valid\] and \[sec:lr\_decay\]). These issues are exacerbated by heterogeneity; Under extreme heterogeneity, finite-sum modelling approaches may lead to theory that does not accurately represent practical FL systems. Thus, our modelling assumptions are designed to encompass both cross-silo and cross-device setting.

Our setup is also relevant to that of model-agnostic meta-learning (MAML), first proposed by @finn2017model. In MAML, the main objective is to find a gradient-based mechanism, which given a task $i$ sampled from ${\mathcal{P}}$, adapts to have good performance on the distribution ${\mathcal{D}}_i$. Unlike cross-device FL, where we generally cannot quantify ${\mathcal{P}}$ directly because of data restrictions, the distribution ${\mathcal{P}}$ is the primary object of interest in MAML. However, much like cross-device FL this distribution is generally not known a priori, but instead is problem-dependent.

<span><span style="font-variant:small-caps;">LocalUpdate</span></span>algorithms {#sec:local_update_algs}
--------------------------------------------------------------------------------

In the following, we will consider a broad class of algorithms that attempt to minimize $f(x)$ (such as in FL methods) or attempt to learn a model that personalizes well with respect to ${\mathcal{P}}$ (such as in meta learning algorithms). We refer to these as <span><span style="font-variant:small-caps;">LocalUpdate</span></span>algorithms. In such methods, at each round, a central coordinator (which we will refer to as a *server*) works with $M$ *clients* (or in the language of MAML, tasks) sampled from ${\mathcal{P}}$ and broadcasts its global model to the clients. Each client $i$ optimizes its loss function $f_i$ (initializing at the broadcast model) by iteratively applying mini-batch SGD with batch size $B$ and *client* (inner) learning rate $\gamma$. The mini-batch gradients are computed by taking samples from the client's local dataset ${\mathcal{D}}_i$. The client then sends a linear combination (parameterized by $\Theta = (\theta_1, \theta_2, \dots)$ where $\theta_i \in {\mathbb{R}}_{\geq 0}$) of its gradients to the server. We will only consider $\Theta$ with finite support. For such $\Theta$, we define $$K(\Theta) = \max\{i~|~\theta_i > 0\}.$$ Throughout our work, we will omit the trailing zeros in any $\Theta$ with finite support. The server averages the available updates, and, treating this average as a stochastic gradient of the loss function $f(x)$, performs a gradient step with a *server* (outer) learning rate $\eta$. Algorithms \[alg:outerloop\] and \[alg:innerloop\] give pseudo-code for <span><span style="font-variant:small-caps;">LocalUpdate</span></span>.

\[t\]<span>.48</span>

\[H\]

**OuterLoop$(x, \{\eta_t\}_{t\geq 1}, \{\gamma_t\}_{t \geq 1}, \Theta)$:**

$x_1 = x$ $I_t \leftarrow$ (random set of $M$ clients) $q_t^i \leftarrow \text{InnerLoop}(i, x_t, \gamma_t, \Theta)$ $q_t \leftarrow (\nicefrac{1}{M})\sum_{i \in I_t} q_{t}^i$ $x_{t+1} = x_t-\eta_tq_t$ return $x_{T+1}$

\[t\]<span>.48</span>

\[H\]

**${\textup{InnerLoop}}(i, x, \gamma, \Theta$):**

$x_1 = x$ sample a set $S_k$ of size $B$ from ${\mathcal{D}}_i$ $g_k = (\nicefrac{1}{B}) \sum_{z \in S_k} \nabla f(x_k ; z)$ $x_{k+1} \leftarrow x_k - \gamma g_k$ return $\sum_{k = 1}^{K(\Theta)} \theta_k g_k$

This method recovers some well-known algorithms for specific choice of $\gamma, \eta$, and $\Theta$. For convenience of notation, we define $$\label{eq:theta_maml}
    \Theta_K = (\underbrace{0, \dots, 0}_{\text{K-1 times}}, 1),$$ so in particular $\Theta_1 = (1)$ and $K(\Theta_1) = 1$, and similarly $$\label{eq:theta_fedavg}
    \Theta_{1:K} = (\underbrace{1, \dots, 1}_{\text{K times}}).$$

Many existing training algorithms can be expressed as special cases of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>. We give a non-exhaustive list below.

-   The simplest setting is mini-batch SGD. This can be recovered in multiple ways. For example, suppose each client $i$ corresponds to a single example $z_i$. Then, <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\Theta = \Theta_1$ is equivalent to mini-batch SGD with batch size $M$ and learning rate $\eta$.

-   Alternatively, if there is only a single client ($|{\mathcal{I}}|=1$), then <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\Theta = \Theta_1$ becomes mini-batch SGD with batch size $B$ and learning rate $\eta$. As expected, the choice of $\gamma$ has no impact in either instance of mini-batch SGD.

-   More generally, setting $\Theta = \Theta_1$ recovers distributed mini-batch SGD, with total batch size $MB$ and learning rate $\eta$. Again, $\gamma$ has no impact on the global model.

-   When there is a single client and $\Theta = \Theta_{1:K}$, then <span><span style="font-variant:small-caps;">LocalUpdate</span></span>recovers the Lookahead optimizer [@zhang2019lookahead] with $K$ "fast weights".

-   In the homogeneous setting, if $\Theta = \Theta_{1:K}$, and $\gamma = \eta$, then <span><span style="font-variant:small-caps;">LocalUpdate</span></span>is equivalent to Local SGD with $K$ local steps. For further details, see Appendix \[sec:special\_cases\].

-   In the heterogeneous setting, if we set $\Theta = \Theta_{1:K}$ and $\gamma = \eta$, then <span><span style="font-variant:small-caps;">LocalUpdate</span></span>is equivalent to FedAvg with $K$ local steps (see Appendix \[sec:special\_cases\] for details). When $\gamma$ and $\eta$ are not necessarily equal, we actually recover Reptile [@nichol2018first], as well as the Generalized FedAvg algorithm in [@reddi2020adaptive]. For convenience of notation, we will refer to this algorithm as FedAvg/Reptile throughout. This equivalence between FedAvg and Reptile was first noted by @jiang2019improving. In fact, we show in Section \[sec:exp\_valid\] that this decoupling of client and server learning rates is critical to understanding and improving the convergence of FedAvg.

-   When $\Theta = \Theta_K$, we recover the first-order MAML (FOMAML) algorithm of @finn2017model. A similar functional relation between FOMAML and Reptile was previously described by @nichol2018first.

-   In the MAML algorithm, @finn2017model use $K$ local update steps for each "task" (in our vocabulary, client). We will refer to this as $\mathcal{K}$-MAML throughout. As we show in Section \[sec:maml\], when the underlying loss functions are quadratic and the clients perform gradient descent updates, $\mathcal{K}$-MAML is recovered by setting $\Theta = \Theta_{2K+1}$. This gives a previously unknown connection between FL and MAML algorithms. As we discuss in Section \[sec:maml\], this does not hold when the clients use SGD due to potential biases in estimating Hessian-gradient products via stochastic gradients.

As written, both clients and server use SGD as their optimizer in <span><span style="font-variant:small-caps;">LocalUpdate</span></span>. However, one could use techniques such as momentum or adaptive learning rates on either the server (as explored by @reddi2020adaptive) or the client (as explored by @xie2019local). While our results can be extended to these settings, we leave this to future work. Our goal is not to derive convergence results for as broad a class of algorithms as possible. Rather, we wish to understand how the choice of $\gamma$ and $\Theta$ impact the dynamics of optimization, especially in heterogeneous settings.

We note that in Algorithm \[alg:innerloop\], each clients performs a designated number of steps of mini-batch SGD, with samples taken from some underlying client distribution $\mathcal{D}_i$. When $\mathcal{D}_i$ is the uniform distribution over some finite set $\mathcal{S}_i$, we could instead write Algorithm \[alg:innerloop\] in terms of performing some number of epochs $E$ of mini-batch SGD over $\mathcal{S}_i$, as is done in [@mcmahan17fedavg] and many other works on federated learning. In this case, the batch size $B$ dictates the number of client gradient steps (as the client roughly take $E|\mathcal{S}_i|/B$ steps). Thus, in such settings, the choice of $B$ has an analogous impact as the choice of the number of local steps $K(\Theta)$ in Algorithm \[alg:innerloop\]. For simplicity of analysis, we will analyze the latter throughout, but our results can be easily extended to the former.

Outline
-------

The rest of this paper is organized as follows. In Section \[sec:local\_update\_as\_sgd\], we show that a round of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>method is equivalent to performing a *single* (stochastic) gradient step with respect to a surrogate objective, which we exactly characterize.

In Section \[sec:example\], we use simple examples to show that the surrogate loss and the original loss can vary substantially. Moreover, we show how choices of $\gamma$ and $\Theta$ affect the discrepancy between the two losses. In particular, we highlight how the choice of $\gamma$ is crucial to the performance of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>. In Section \[sec:properties\], we analyze spectral properties of the surrogate loss, and show that <span><span style="font-variant:small-caps;">LocalUpdate</span></span>can be viewed as implicit regularization on the condition number, where the amount of regularization is controlled by $\gamma$ and $\Theta$.

The next sections present to the best of our knowledge a novel proof technique, characterizing the convergence of FedAvg/Reptile in heterogeneous settings.[^3] In Section \[sec:distance\], we bound the distance between the minimizers of the surrogate and the true loss function in terms of the client learning rate $\gamma$. We use these results in Section \[sec:convergence\] to derive convergence rates for FedAvg/Reptile that highlight how the choice of client learning rate $\gamma$ gives rise to a trade-off between local and global optimization. In particular, we show that learning rate decay is both sufficient and necessary for convergence to the true risk minimizer.

While our theoretical results are valid only for quadratic loss functions, in Section \[sec:exp\_valid\] we show empirically that our conclusions carry over to more general settings, including non-convex objectives. Our empirical results highlight the importance of learning rate tuning in federated learning. In Section \[sec:lr\_decay\], we combine our theoretical insights with important systems-level constraints to design a method for automatic learning rate decay methods for local update methods. In particular, we present a simple, easy to implement method for automatic learning rate decay, and show its efficacy in improving accuracy and reducing the need for client learning rate tuning.

<span><span style="font-variant:small-caps;">LocalUpdate</span></span>as SGD {#sec:local_update_as_sgd}
============================================================================

When $K(\Theta) > 1$ and $\gamma > 0$, the dynamics of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>may be very different than those of mini-batch SGD. We will show that for quadratic functions, these dynamics are related but distinct. In particular, we will show that any local update method on a quadratic function can be viewed as SGD on some appropriately defined surrogate loss function. Moreover, the discrepancy between the true loss function and the surrogate loss function is dictated by the choice of client learning rate $\gamma$ and $\Theta$.

For $i \in {\mathcal{I}}$, define: $$A_i := \operatorname*{\mathbb{E}}_{z \sim {\mathcal{D}}_i}[A_z].$$ We assume throughout that $A_i$ is finite and invertible. We also define $$c_i := A_i^{-1}\operatorname*{\mathbb{E}}_{z \sim {\mathcal{D}}_i}[A_zc_z].$$ Again, we assume this is finite. We then have the following lemma.

\[lem:f\_i\_loss\] For all $i \in {\mathcal{I}}$, there is some constant $\tau_i$ such that $$f_i(x) = \dfrac{1}{2}\norm{A_i^{1/2}(x-c_i)}^2 + \tau_i.$$

In the sequel, we will omit the constant term $\tau_i$, and let $$f_i(x) = \dfrac{1}{2}\norm{A_i^{1/2}(x-c_i)}^2$$ as this does not change the gradients of the loss $f_i$. Since each $A_z$ is symmetric and positive definite, so is $A_i$. We define the following: $$A := \operatorname*{\mathbb{E}}_{i\sim{\mathcal{P}}}[A_i],~~~c:=\operatorname*{\mathbb{E}}_{i\sim{\mathcal{P}}}[c_i].$$ We will assume that these expectations exist and are finite throughout. We will also utilize the following mild assumptions at different times.

\[assm0\] $K(\Theta) > 0$.

\[assm1\] There are $\mu, L > 0$ such that for all $i$, $$\mu I \preceq A_i \preceq L I.$$

\[assm2\]There are finite $\sigma_A$ and $\sigma_c$ such that $$\operatorname*{\mathbb{E}}_{i\sim{\mathcal{P}}}\sbr{\norm{A_i-A}^2} \leq \sigma_A^2.$$ $$\operatorname*{\mathbb{E}}_{i\sim{\mathcal{P}}}\sbr{\norm{c_i-c}^2} \leq \sigma_c^2.$$

Assumption \[assm0\] prevents pathologically bad choices of $\Theta$ in which clients simply send 0 to the server at every round. Assumption \[assm1\] amounts to assuming upper and lower bounds on the Lipschitz and strong convexity parameters of each loss function $f_i$. This is satisfied if there are a finite number of clients, and for each, $A_i$ is positive definite. However, it is often true in more generality if the underlying matrices $A_z$ satisfy some kind of bounded eigenvalue condition. Moreover, when the number of clients is finite, we can always ensure that $\mu I \preceq A_i$ for all $i$ by adding $\ell_2$ regularization to our objective function.

Assumption \[assm2\] assumes that the matrices $A_i$ and optimal points $c_i$ for each loss function have bounded variance. We do not assume that the gradients computed by the clients have bounded norm. Intuitively, as $\sigma_c \to 0$, local update methods should provide more benefit, as the clients are taking more steps towards a shared optimum. While $\sigma_c = 0$ in the case of homogeneous data distributions (i.e. ${\mathcal{D}}_i$ are the same for all $i\in{\mathcal{I}}$), these two conditions are not equivalent. There are heterogeneous data distributions which still yield $\sigma_c = 0$. Also, note that $c$ is in general not the minimizer of the objective $f$.

Fix $i \in {\mathcal{I}}$, and consider Algorithm \[alg:innerloop\]. We initialize $x_1 = x$, and then at each iteration $k$ we sample a set $S_k$ uniformly at random (with replacement) from ${\mathcal{D}}_i$, then update via $$\label{eq:inner_g_update}
    g_{k} = \frac{1}{B}\sum_{z \in S_k} \nabla f(x_k ; z),$$ $$\label{eq:inner_x_update}
    x_{k+1} = x_k- \gamma g_k.$$ We first prove a basic recurrence relation concerning the local gradients $g_k$ for task $i$.

\[lem:grad\_recurrence\] For $i \in {\mathcal{I}}$, suppose that $A_i$ is invertible and $g_k$ as in Algorithm \[alg:innerloop\], for all $k \geq 1$, $$\label{eq:grad_recurrence}
    \operatorname*{\mathbb{E}}[g_{k+1}] = (I-\gamma A_i)\operatorname*{\mathbb{E}}[g_k].$$

#### Defining the surrogate loss

Using Lemma \[lem:grad\_recurrence\], we will show that Algorithm \[alg:innerloop\] can be viewed as performing SGD on a *surrogate loss*. This surrogate loss will be parameterized by the inputs $\gamma$ and $\Theta$ to Algorithm \[alg:innerloop\]. To define the surrogate loss, we first define, for each client $i \in {\mathcal{I}}$, a *distortion matrix* $Q_i(\gamma, \Theta)$ as follows: $$\label{eq:Q_matrix}
Q_i(\gamma, \Theta) := \sum_{k=1}^{K(\theta)} \theta_k (I-\gamma A_i)^{k-1}.$$ We can then define, for each $i \in {\mathcal{I}}$, the client's surrogate loss function: $$\label{eq:surrogate_i}
\tilde{f}_i(x, \gamma, \Theta) := \frac{1}{2}\|(Q_i(\gamma, \Theta)A_i)^{1/2}(x-c_i)\|.$$ The overall surrogate loss function is then given by $$\label{eq:surrogate_loss}
    \tilde{f}(x, \gamma, \Theta) := \operatorname*{\mathbb{E}}_{i \sim {\mathcal{P}}} [\tilde{f}_i(x, \gamma, \Theta)].$$ Informally, the matrix $Q_i(\gamma, \Theta)$ can be viewed as causing a distortion to the matrix $A_i$. When $\Theta = \Theta_1$, one can see that $Q_i(\gamma, \Theta) = I$, in which case there is no distortion. For other $\Theta$, $Q_i(\gamma, \Theta)$ may significantly distort $A_i$, and can amplify heterogeneity of the $A_i$. Using Lemma \[lem:grad\_recurrence\], we derive the following property of the output of the Algorithm \[alg:innerloop\].

\[thm:sgd\_objective\] Suppose that $A_i$ is invertible. Then $$\label{eq:thm_1}
\operatorname*{\mathbb{E}}[{\textup{InnerLoop}}(i, x, \gamma, \Theta)] = \nabla_x \tilde{f}_i(x, \gamma, \Theta) = Q_i(\gamma, \Theta)\nabla f_i(x).$$

By direct computation, $$\nabla \tilde{f}_i(x, \gamma, \Theta) = Q_i(\gamma, \Theta)A_i(x-c_i).$$ On the other hand, by Lemma \[lem:grad\_recurrence\], $$\begin{aligned}
\operatorname*{\mathbb{E}}[{\textup{InnerLoop}}(i, w, \gamma, \Theta)] &= \sum_{k=1}^{K(\theta)} \theta_k(1-\gamma A_i)^{k-1}\operatorname*{\mathbb{E}}[g_1]\\
&= \sum_{k=1}^{K(\theta)} \theta_k(1-\gamma A_i)^{k-1}A_i(x-c_i)\\
&= Q_i(\gamma, \Theta)A_i(x-c_i).\end{aligned}$$ This proves the first equality. The second follows from noting that $\nabla f_i(x) = A_i(x-c_i)$.

Let $q_t$ be as in Algorithm \[alg:outerloop\]. Then Theorem \[thm:sgd\_objective\] implies $\operatorname*{\mathbb{E}}[q_t] = \nabla \tilde{f}(x, \gamma, \Theta)$. In particular, one round of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with a given $\gamma, \Theta$ is equivalent to performing one step of SGD on the surrogate loss function $\tilde{f}(x, \gamma, \Theta)$.

We note that a version of Theorem was first shown for the case $\Theta = \Theta_{2}$ by @fallah2020personalized, and was used to compare the behavior of FOMAML and MAML. We will take this comparison a step further, by showing in Section \[sec:maml\] that in the non-stochastic client setting, MAML can also be viewed as performing SGD on a similarly-defined surrogate loss.

Theorem \[thm:sgd\_objective\] has important consequences regarding the impact of other "outer optimizers" in Algorithm \[alg:outerloop\], such as the adaptive server optimization [@reddi2020adaptive]. If we treat the the output of Algorithm \[alg:innerloop\] simply as a stochastic gradient oracle of $\tilde{f}(x, \gamma, \Theta)$, we can apply existing convergence guarantees of any gradient based methods to understand the behavior of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>method with different outer optimizers. In particular, this implies that the choice of outer optimization method primarily impacts the *speed* of convergence to $\operatorname*{argmin}_x \tilde{f}(x, \gamma, \Theta)$, but not the point <span><span style="font-variant:small-caps;">LocalUpdate</span></span>actually converges to. We empirically analyze the use of adaptive server methods in <span><span style="font-variant:small-caps;">LocalUpdate</span></span>in Sections \[sec:exp\_valid\] and \[sec:lr\_decay\].

MAML {#sec:maml}
----

As previously discussed, in the setting above, one can actually view MAML as a special case of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>. In this section we elaborate on the claim, using a similar presentation of MAML as in [@nichol2018first]. MAML with $K$ local steps (which we refer to as $\mathcal{K}$-MAML) can be viewed as a simple modification of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>. Algorithm \[alg:outerloop\] proceeds in the same manner. In Algorithm \[alg:innerloop\], each client still executes $K$ mini-batch SGD steps. However, what each client sends to the server differs from Algorithm \[alg:outerloop\].

For simplicity, we define $X_{K}^i(x)$ as the function that runs $K$ steps of mini-batch SGD, starting from $x$, for some fixed mini-batches $S_1, \dots, S_K$ of size $B$ drawn independently from $\mathcal{D}_i$. For convenience, we let $X_0^i(x) = x$. We then define $$m^i_K(x ; z) := f(X_{K}^i(x); z),$$ $$m_{K}^i(x) := \operatorname*{\mathbb{E}}_{z \sim {\mathcal{D}}_i} m^i_K(x ; z).$$ Note that these are implicitly functions of the mini-batches $S_1, \ldots, S_K$ sampled from ${\mathcal{D}}_i$. The output $q^i(x)$ of client $i$ (as a function of its initial model $x$) is a stochastic estimate of $\nabla m_K^i(x)$, so that $$\label{eq:maml_output}
\operatorname*{\mathbb{E}}[q^i(x)] = \nabla m_K(x ; z),$$ The remainder of the MAML algorithm proceeds in the same way as <span><span style="font-variant:small-caps;">LocalUpdate</span></span>. Namely, the server averages the client outputs, and uses this as a gradient estimate with learning rate $\eta$. That is, $$\label{eq:maml_server_update}
q(x) = \sum_{i \in I} q^i(x),$$ $$\label{eq:maml_server_update_2}
x' = x - \eta q(x).$$

We now show that when the clients use gradient descent to perform their local update, $\mathcal{K}$-MAML is in expectation equivalent to performing <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\Theta = \Theta_{2K+1}$.

\[thm:maml\]If $X_K^i(x)$ is the function that runs $K$ steps of gradient descent, starting from $x$, on the client dataset ${\mathcal{D}}_i$, then $$\nabla m^i_K(x) = \operatorname*{\mathbb{E}}[{\textup{InnerLoop}}(x, \gamma, \Theta_{2K+1})] = \nabla_{x} \tilde{f}_i(x, \gamma, \Theta_{2K+1}).$$

It is fruitful to reflect on what this means. Informally, this result shows that for quadratic functions, the gradient of the loss after $K$ steps of gradient descent, taken with respect to the initial point $x$, is in expectation the gradient of the loss function after $K+1$ additional SGD steps. In particular, given $q(x)$ as in , we have $$\operatorname*{\mathbb{E}}[q(x)] = \nabla_x \tilde{f}(x, \gamma, \Theta_{2K+1}).$$ Thus, the MAML update in amounts to a single stochastic gradient update on the surrogate loss function $\tilde{f}(x, \gamma, \Theta_{2K+1}).$

We note that this result relied on the clients using gradient descent to compute $X_k^i(x)$. For computational efficiency, this is often instead done using mini-batch SGD. However, computing $\nabla m_K^i(x)$ then involves computing unbiased estimates of the Hessian and gradient using the same batches of data. By the chain rule, estimating $\nabla m_K^i(x)$ involves multiplying these Hessian and gradient estimates. However, the product of these unbiased estimators need not be unbiased since they were computed with respect to the same batch of data. Thus, this correspondence between $\mathcal{K}$-MAML and <span><span style="font-variant:small-caps;">LocalUpdate</span></span>may break down in computationally-efficient (but biased) MAML implementations. For more detailed discussion on this bias, see [@fallah2019convergence]. We also note that @fallah2019convergence analyze the convergence properties of MAML and FOMAML, and independently observe that MAML and FOMAML need not share stationary points for quadratic objectives.

Local update methods tend towards different global minima {#sec:example}
=========================================================

As we will show in Section \[sec:properties\], $\tilde{f}_i(x, \gamma, \Theta)$ shares many properties with $f_i(x)$, including having the same global minima. However, we first highlight that a crucial correspondence breaks down when considering the population-level global minima in heterogeneous settings. That is, in general $$\operatorname*{argmin}_x f(x) \neq \operatorname*{argmin}_x \tilde{f}(x, \gamma, \Theta).$$ In order to enhance our understanding of the surrogate loss function, we first give both analytic and empirical examples.

#### Analytic examples

Let $\mathcal{P}$ have support $\{1, 2\}$, with each option equally likely. Suppose that ${\mathcal{D}}_1, {\mathcal{D}}_2$ have support only on the points $z = 1$ and $z = 2$ respectively, and suppose $A_z = z, c_z = 1/z$, so that $f(x; z) = \frac{1}{2}z(x - 1/z)^2$. We see that $f(x)$ is minimized at $x = 2/3$. On the other hand, let $\Theta = \Theta_{1:2}$. For $\gamma \geq 0$, we can compute the $i$-th distortion matrix by $$\begin{aligned}
    Q_i(\gamma, \Theta) = 1 + (1-\gamma i) = 2-\gamma i.\end{aligned}$$ Note that this is positive definite for $i = 1, 2$ as long as $\gamma \in [0, 1/2)$. By , $$\begin{aligned}
\tilde{f}(x, \gamma, \Theta) = \dfrac{(2-\gamma)(x-1)^2}{4} + (1-\gamma)(x-1/2)^2.\end{aligned}$$ For $\gamma \in [0, 6/5)$, this is a positive definite quadratic function, with minimum given by $$\operatorname*{argmin}_{x} \tilde{f}(x, \gamma, \Theta) = \dfrac{4-3\gamma}{6-5\gamma},$$ implying $$\norm{\operatorname*{argmin}_{x} \tilde{f}(x, \gamma, \Theta) - \operatorname*{argmin}_{x} f(x)} = \dfrac{\gamma}{3(6-5\gamma)}.$$ Therefore, even if we run <span><span style="font-variant:small-caps;">LocalUpdate</span></span>until convergence, it would not converge to the true risk minimizer. This holds even though there are only two clients, each with a single data point. In other words, some form of learning rate decay is *necessary* for convergence to the risk minimizer. The necessity of learning rate decay was first shown by @li2019convergence, and later shown in [@malinovsky2020local] and [@pathak2020fedsplit]. We take this analysis further, by showing how this sub-optimal behavior is explicitly governed by algorithmic choices, especially learning rate and the number of local steps taken.

For instance, in the example above, as $\gamma \to 0$, the surrogate risk minimizer converges to the true risk minimizer. Conversely, as $\gamma$ increases, the two minima become further apart. In fact, as $\gamma \to 6/5$, the distance between the two optima grows arbitrarily large, even though $\tilde{f}(x, \gamma, \Theta)$ is a positive definite quadratic function. The critical issue here is that while $\gamma < 6/5$ ensures that $f(x, \gamma, \Theta_{1:2})$ is well-behaved (ie. that it is a positive definite quadratic function), it is not small enough to ensure that each client's surrogate loss function $\tilde{f}_i(x, \gamma, \Theta)$ is well-behaved. Note $f_2(x, \gamma, \Theta_{1:2})$ is a negative definite quadratic function for $\gamma > 1/2$, which causes this divergence. Even though the iterates of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>are converging, there may be some client which is not converging in any meaningful sense. Therefore:

**The client learning rate should be set sufficiently small so that all client loss functions are well-behaved, even if the overall loss function is well-behaved.**

A similar analysis shows that if we instead fix $\gamma$, the distance between the surrogate risk and true risk minimizers depends on $K$. Let $\gamma = 1/2$. By , $$Q_1(1/2, \Theta_{1:K}) = \sum_{k=0}^{K-1} 2^{-k} = 2-2^{-(K-1)},~~Q_2(1/2, \Theta_{1:K}) = 1,$$ implying $$\begin{aligned}
    f(x, 1/2, \Theta_{1:K}) =\dfrac{1}{2}\left((1-2^{-K})(x-1)^2 + (x-1/2)^2\right).\end{aligned}$$ This is a positive definite quadratic with minima given by $$\operatorname*{argmin}_x f(x, 1/2, \Theta_{1:K}) = \dfrac{3\cdot 2^{K} - 2}{2^{K+2}-2}.$$ This then implies $$\norm{\operatorname*{argmin}_{x} f(x, 1/2, \Theta_{1:K}) - \operatorname*{argmin}_{x} f(x)} = \dfrac{2^K-2}{6(2^{K+1}-1)}.$$

When $K = 1$, this gap is 0, while as $K \to \infty$, the distance increases monotonically to $1/12$. In fact, as $K \to \infty$, the minimizer of the surrogate loss function converges to the expected value of the client loss minimizers. In Lemma \[lem:asymptotic\] we prove an even stronger statement, and show that it holds for all positive definite quadratic loss functions.

#### Empirical examples

Next, we give an empirical generalization of the above example for further illustration. Let $f(x; z) = (1/2)z x^2 - x$ for $z > 0$ (ie. $A_z = z$ and $c_z = 1/z$). We let ${\mathcal{P}}$ have support $[0.5, 2]$ and density function $q(i) = 8i/15$. For each $i \in \operatorname*{supp}({\mathcal{P}})$, we let the client distribution ${\mathcal{D}}_i$ be supported on a single point $z = i$, so that $f_i(x) = (1/2)ix^2 - x$. Again, deterministic $f_i$ will still be sufficient observe discrepancies between the true and surrogate loss functions.

In Figure \[fig:1d-example\], we plot the behavior of FedAvg/Reptile with $K$ local steps (<span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\Theta = \Theta_{1:K}$) on this problem. We fix $\eta = 10^{-4}, M = B = 1$, and illustrate the change in behavior as either $\gamma$ varies and $K$ is fixed, or $K$ varies and $\gamma$ is fixed. We also plot the true minimizer and the average minimizer, $\operatorname*{\mathbb{E}}_{i\sim\mathcal{P}} [\operatorname*{argmin}_x f_i(x)]$.

![<span><span style="font-variant:small-caps;">LocalUpdate</span></span>convergence behavior with $\Theta = \Theta_{1:K}$ on an example $1$-dimensional problem. (Left) We fix $\gamma = 0.5$ and vary $K$. (Right) We fix $K = 20$ and vary $\gamma$.<span data-label="fig:1d-example"></span>](figures/local_update_example.pdf){width="0.9\linewidth"}

If $\gamma = 0$ or $K = 1$, we converge to $x^*$. As $\gamma$ or $K$ increases, we converge to a point further from $x^*$. As $K \to \infty$ we converge to the average minimizer. We also see that *decreasing* the local stepsize *increases* the variance. This is to be expected: When $\gamma=0$, <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $M=1$ reduces to mini-batch SGD with batch size $K$, but the gradients in the batch are *summed* rather than averaged. For $\gamma > 0$, the magnitude of the gradients being summed decreases as the client converges to its minimizer. If we set $\gamma$ to be larger than $1.0$, we see an even greater gap between the surrogate minimizer and the true minimizer (due to the presence of negative-definite clients).

In Section \[sec:distance\], we derive general bounds on the distance between surrogate and true minimizers. To do so, we will use spectral properties of the matrix $Q_i(\gamma, \Theta)$, which we derive in the next section.

Surrogate loss properties {#sec:properties}
=========================

In order to understand the surrogate loss function $\tilde{f}(x, \gamma, \Theta)$, we will analyze properties of the distortion matrix $Q_i(\gamma, \Theta)$. We first note that in certain special cases, $Q_i(\gamma, \Theta)$ is some scaled version of the identity matrix. The following result is a simple consequence of the definition in .

\[lem:Q\_special\_case\] Let $a = \sum_{k=1}^{K(\theta)} \theta_k$. If $\gamma = 0$ or $K(\theta) = 1$, then $Q_i(\gamma, \Theta) = aI$ and $\tilde{f}_i(x, \gamma, \Theta) = af_i(x)$.

Under these settings, we therefore have $\tilde{f}(x, \gamma, \Theta) = af(x)$, and so the surrogate and true loss functions only differ by a constant. More generally, the functions $\tilde{f}_i(x, \gamma, \Theta)$ inherit many properties from $f_i(x)$, as captured in the following lemma.

\[lem:Q\_structure\] Let $L_i = \lambda_{\max}(A_i), \mu_i = \lambda_{\min}(A_i) > 0$ and suppose that $\gamma < L_i^{-1}$. Then

1.  $Q_i(\gamma, \Theta)$ is symmetric and positive definite.

2.  $\tilde{f}_i(x, \gamma, \Theta)$ and $f_i(x)$ have the same unique global minimizer $x^* = c_i$.

3.  For each eigenvalue $\lambda$ of $A_i$, $Q_i$ has an eigenvalue $$\sum_{k = 1}^{K(\Theta)} \theta_k(1-\gamma \lambda)^{k-1}.$$ with the same multiplicity.

4.  $Q_i(\gamma, \Theta)$ satisfies $$\lambda_{\max}(Q_i(\gamma, \Theta)) = \sum_{k = 1}^{K(\Theta)} \theta_k(1-\gamma \mu_i)^{k-1}.$$ $$\lambda_{\min}(Q_i(\gamma, \Theta)) = \sum_{k = 1}^{K(\Theta)} \theta_k(1-\gamma L_i)^{k-1}.$$

For convenience, we note some special cases of Lemma \[lem:Q\_structure\] for $\Theta = \Theta_{K}, \Theta_{1:K}$.

\[lem:Q\_structure\_fomaml\] Let $L_i = \lambda_{\max}(A_i), \mu_i = \lambda_{\min}(A_i)$ and suppose that $\gamma < L_i^{-1}$. Then $$\lambda_{\max}(Q_i(\gamma, \Theta_K)) = (1-\gamma\mu_i)^K,~~~
        \lambda_{\min}(Q_i(\gamma, \Theta_K)) = (1-\gamma L_i)^K.$$

If Assumption \[assm1\] holds and $\gamma < L^{-1}$, then $$\lambda_{\max}(Q_i(\gamma, \Theta_K)) \leq (1-\gamma\mu)^K,~~~
        \lambda_{\min}(Q_i(\gamma, \Theta_K)) \geq (1-\gamma L)^K.$$

\[lem:Q\_structure\_fedavg\] Let $L_i = \lambda_{\max}(A_i), \mu_i = \lambda_{\min}(A_i)$ and suppose that $\gamma < L_i^{-1}$. Then $$\lambda_{\max}(Q_i(\gamma, \Theta_{1:K})) = \dfrac{(1-(1-\gamma \mu_i)^K)}{\gamma \mu_i},~~~
        \lambda_{\min}(Q_i(\gamma, \Theta_{1:K})) = \dfrac{(1-(1-\gamma L_i)^K)}{\gamma L_i}.$$

If Assumption \[assm1\] holds and $\gamma < L^{-1}$, then $$\lambda_{\max}(Q_i(\gamma, \Theta_{1:K})) \leq \dfrac{(1-(1-\gamma \mu)^K)}{\gamma \mu},~~~
        \lambda_{\min}(Q_i(\gamma, \Theta_{1:K})) \geq \dfrac{(1-(1-\gamma L)^K)}{\gamma L}.$$

As the eigenvalue bounds above suggest, ${Q_i(\gamma, \Theta_{1:K})}$ actually has a relatively simple form, as we show in the next lemma.

\[lem:QA\_fedavg\] Suppose $0 < \gamma < L_i^{-1}$. Then $${Q_i(\gamma, \Theta_{1:K})}= (I-(I-\gamma A_i))(\gamma A_i)^{-1}.$$

Note that when $\gamma = 0$, implies ${Q_i(\gamma, \Theta_{1:K})}= KI$.

We are also interested in the matrix $Q_i(\gamma, \Theta)A_i$, as its eigenvalues govern the Lipschitz and strong convexity parameters of the function $\tilde{f}_i(x, \gamma, \Theta)$. We have the following result.

\[lem:QA\_structure\] Let $L_i = \lambda_{\max}(A_i), \mu_i = \lambda_{\min}(A_i)$.

1.  For each eigenvector and eigenvalue pair $(v, \lambda)$ of $A_i$, $v$ is an eigenvector of $Q_i(\gamma, \Theta)A_i$ with eigenvalue $$\sum_{k =1}^{K(\Theta)} \theta_k(1-\gamma\lambda)^{k-1}\lambda.$$

2.  If $\gamma < L_i^{-1}$, $Q_i(\gamma, \Theta)A_i$ is symmetric and positive definite with eigenvalues satisfying $$\lambda_{\max}(Q_i(\gamma, \Theta)A_i) \leq \sum_{k = 1}^{K(\Theta)} \theta_k(1-\gamma \mu_i)^{k-1}L_i.$$ $$\lambda_{\min}(Q_i(\gamma, \Theta)A_i) \geq \sum_{k = 1}^{K(\Theta)} \theta_k(1-\gamma L_i)^{k-1}\mu_i.$$

The bounds in Lemma \[lem:QA\_structure\] can be refined for specific $\Theta$. We first consider FedAvg/Reptile, when $\Theta = \Theta_{1:K}$. By Lemma \[lem:QA\_fedavg\], when $0 < \gamma < L_i^{-1}$, we have $${Q_i(\gamma, \Theta_{1:K})}A_i = \dfrac{I-(I-\gamma A_i)^K}{\gamma}.$$ We will therefore be able to compute the eigenvalues of ${Q_i(\gamma, \Theta_{1:K})}A_i$ in terms of the function $$\label{eq:phi}
\phi_{K, \lambda}(\gamma) := \gamma^{-1}(1-(1-\gamma\lambda)^K).$$

In fact, $\phi_{K, \lambda}(\gamma)$ is actually continuous at $0$, with its value being given by $\phi_{K, \lambda}(0) = K\lambda$. One way to see this is by noting that by basic properties of geometric sums, for $0 \leq \gamma \leq \lambda^{-1}$, $$\label{eq:phi_alt}
\phi_{K, \lambda}(\gamma) = \sum_{k=1}^K (1-\gamma\lambda)^{k-1}\lambda.$$

We can now give strong bounds on the spectrum of ${Q_i(\gamma, \Theta_{1:K})}A_i$. We get the following:

\[lem:QA\_structure\_fedavg\] Let $L_i = \lambda_{\max}(A_i), \mu_i = \lambda_{\min}(A_i)$.

1.  For each eigenvector, eigenvalue pair $(v, \lambda)$ of $A_i$, $v$ is an eigenvector of ${Q_i(\gamma, \Theta_{1:K})}A_i$ with eigenvalue $\phi_{K, \lambda}(\gamma)$.

2.  If $\gamma < L_i^{-1}$, the maximum and minimum eigenvalues of ${Q_i(\gamma, \Theta_{1:K})}A_i$ are given by $$\lambda_{\max}(Q_i(\gamma, \Theta_{1:K})A_i) = \phi_{K, L_i}(\gamma),$$ $$\lambda_{\min}(Q_i(\gamma, \Theta_{1:K})A_i) = \phi_{K, \mu_i}(\gamma).$$

3.  If Assumption \[assm1\] holds and $\gamma < L^{-1}$, then $$\lambda_{\max}(Q_i(\gamma, \Theta_{1:K})A_i) \leq \phi_{K, L}(\gamma),$$ $$\lambda_{\min}(Q_i(\gamma, \Theta_{1:K})A_i) \geq \phi_{K, \mu}(\gamma).$$

We can also tighten the bounds in Lemma \[lem:QA\_structure\] for the MAML-style algorithms, where $\Theta = \Theta_K$, as long as the learning rate is set appropriately.

\[lem:QA\_structure\_maml\] Let $L_i = \lambda_{\max}(A_i), \mu_i = \lambda_{\min}(A_i)$.

1.  For each eigenvector, eigenvalue pair $(v, \lambda)$ of $A_i$, $v$ is an eigenvector of $Q_i(\gamma, \Theta_K) A_i$ with eigenvalue $(1-\gamma\lambda)^{K-1}\lambda$.

2.  If $\gamma < (KL_i)^{-1}$, the maximum and minimum eigenvalues of $Q_i(\gamma, \Theta_K) A_i$ are given by $$\lambda_{\max}(Q_i(\gamma, \Theta_{K})A_i) = (1-\gamma L)^{K-1}L,$$ $$\lambda_{\min}(Q_i(\gamma, \Theta_{K})A_i) = (1-\gamma\mu)^{K-1}\mu.$$

3.  If Assumption \[assm1\] holds and $\gamma < (KL)^{-1}$, then $$\lambda_{\max}(Q_i(\gamma, \Theta_{K})A_i) \leq (1-\gamma L)^{K-1}L,$$ $$\lambda_{\min}(Q_i(\gamma, \Theta_{K})A_i) \geq (1-\gamma\mu)^{K-1}\mu.$$

Lemmas \[lem:QA\_structure\_fedavg\] and \[lem:QA\_structure\_maml\] imply the following results regarding the condition number of the surrogate loss.

\[cor:cond\_fedavg\] If Assumption \[assm1\] holds and $\gamma < L^{-1}$, then $\tilde{f}(x, \gamma, \Theta_{1:K})$ is $\tilde{L}$-smooth and $\tilde{\mu}$-strongly convex where $$\tilde{L} \leq \phi_{K, L}(\gamma),$$ $$\tilde{\mu} \geq \phi_{K, \mu}(\gamma).$$ Therefore, $\tilde{f}(x, \gamma, \Theta_{1:K})$ has condition number $\tilde{\kappa}$ satisfying $$\tilde{\kappa} \leq \dfrac{\phi_{K, L}(\gamma)}{\phi_{K, \mu}(\gamma)} = \dfrac{1-(1-\gamma L)^K}{1-(1-\gamma\mu)^K}.$$

\[cor:cond\_maml\] If Assumption \[assm1\] holds and $\gamma < (KL)^{-1}$, then $\tilde{f}(x, \gamma, \Theta_{K})$ is $\tilde{L}$-smooth and $\tilde{\mu}$-strongly convex where $$\tilde{L} \leq (1-\gamma L)^{K-1}L,$$ $$\tilde{\mu} \geq (1-\gamma \mu)^{K-1}\mu.$$ Therefore, $\tilde{f}(x, \gamma, \Theta_{1:K})$ has condition number $\tilde{\kappa}$ satisfying $$\tilde{\kappa} \leq \left(\dfrac{1-\gamma L}{1-\gamma\mu}\right)^{K-1}\dfrac{L}{\mu}.$$

#### Local computation as implicit regularization

For $\Theta = \Theta_K$, we clearly see that as $K \to 1$ or $\gamma \to 0$, then $\tilde{\kappa} \to L/\mu$, the condition number of the true loss function. However, if $\gamma$ is not close to 0, we see an exponential reduction (in terms of $K$) of the condition number. While the analysis is not quite as clear for $\Theta = \Theta_{1:K}$, one can show that for all $\gamma$, $\tilde{\kappa} \leq L/\mu$, with equality if and only if $\gamma = 0$ or $K = 1$. Moreover, the condition number decreases as $K \to \infty$ or $\gamma \to L^{-1}$.

It is well known that the condition number measures how quickly methods such as gradient descent can find a minimizer of a strongly convex function (see Chapter 3 of [@bubeck2014convex] for reference). By performing more local computations on the clients, and with larger learning rate, we actually reduce the condition number of the surrogate loss. We see that intuitive notions about methods such as FedAvg (e.g. that more local computation improves convergence) can be made formal by analyzing properties of the surrogate loss. Thus, we have the following important takeaway:

**Methods such as MAML, FedAvg, and Reptile perform implicit regularization on the condition number of the surrogate loss function they are actually optimizing.**

When $\Theta = \Theta_K$ or $\Theta_{1:K}$, we see that <span><span style="font-variant:small-caps;">LocalUpdate</span></span>may be able to optimize the surrogate loss more quickly (due to the condition number reduction). However, as shown in Section \[sec:example\], the surrogate loss may differ drastically from the true loss. In the next section, we use the spectral properties of the surrogate loss derived above to quantify the distance between the minimizers of these two functions.

Bounding the distance between global minima {#sec:distance}
===========================================

As discussed above, <span><span style="font-variant:small-caps;">LocalUpdate</span></span>is not optimizing the desired loss function $f(x)$, but a surrogate loss function $\tilde{f}(x, \gamma, \Theta)$. In this section, we will bound the distance between the minima of these two functions. We will assume that Assumptions \[assm0\], \[assm1\], and \[assm2\] hold throughout. Recall that by Lemma \[cor:cond\_fedavg\], as long as $\gamma < L^{-1}$, $\tilde{f}(x, \gamma, \Theta)$ is strongly convex and therefore has a unique minimizer.

For simplicity of notation, we will fix $\gamma, \Theta$ and let $Q_i$ refer to $Q_i(\gamma, \Theta)$ throughout this section. We will also define matrices $\tilde{A}_i$ and $\tilde{A}$, depending on $\Theta$, as follows: $$\label{eq:distance_quantities}
\tau := \rbr*{\sum_{k=1}^{K(\Theta)} \theta_k}^{-1},~~~\tilde{A}_i := \tau Q_iA_i,~~~\tilde{A} := \operatorname*{\mathbb{E}}_i\sbr*{\tilde{A}_i}.$$

We also define the following quantities: $$x^*(\gamma, \Theta) := \operatorname*{argmin}_{x} \tilde{f}(x, \gamma, \Theta),~~x^* := \operatorname*{argmin}_x f(x).$$

By Lemma \[lem:Q\_special\_case\], if $\gamma = 0$ or $\theta_k = 0$ for $k \geq 2$, then $Q_i = \tau^{-1}I$. Therefore, $\tilde{f}(x, \gamma, \Theta) = \tau^{-1}f(x)$, and so $x^*(0, \Theta) = x^*$. However, for general $\gamma$ and $\Theta$, $x^*(\gamma, \Theta)\neq x^*$.

We are first interested in how far apart the two minimizers can possibly be. We first consider the asymptotic affect of $K$ when setting $\Theta = \Theta_{1:K}$. In fact, varying $K$ can only change the distance between the two by a fixed amount, as shown in the following.

\[lem:asymptotic\] Suppose $0 < \gamma < L^{-1}$. Then for all $K \geq 1$, $$\tilde{f}(x, \gamma, \Theta_{1:K}) \leq \frac{1}{2\gamma}\operatorname*{\mathbb{E}}_{i\sim{\mathcal{P}}}\norm{x-c_i}^2.$$ Moreover, $\tilde{f}(x, \gamma, \Theta_{1:K})$ converges pointwise to this function, ie. $$\lim_{K \to \infty} \tilde{f}(x, \gamma, \Theta_{1:K}) = \dfrac{1}{2\gamma}\operatorname*{\mathbb{E}}_{i \sim {\mathcal{P}}}[\norm{x-c_i}^2]$$ which has a unique minimizer at $x = \operatorname*{\mathbb{E}}_{i \sim {\mathcal{P}}}\sbr{c_i} = c$.

Intuitively, we see that as long as the client learning rate is not too high, the worst possible surrogate loss is the one defined by the average distance to the client optimizers. Intuitively, as $K \to \infty$, <span><span style="font-variant:small-caps;">LocalUpdate</span></span>will take steps oriented more and more towards the average of the client minimizers (one-shot averaging), which is reflected in the experiment in Figure \[fig:1d-example\].

We now wish to understand the non-asymptotic regime, especially the distance between the surrogate risk minimizer and the true risk minimizer, as this will help inform us how to set $\gamma$ in <span><span style="font-variant:small-caps;">LocalUpdate</span></span>. We have the following result.

\[thm:optima\_distance\] If $\gamma < L^{-1}$, then $$\label{eq:optima_distance}
\norm{x^*(\gamma, \Theta) - x^*} \leq \dfrac{L\sigma_c}{\mu}\rbr*{1+\dfrac{\sigma_A}{\mu}}\dfrac{1- \chi(\gamma, \Theta)}{ \chi(\gamma, \Theta)}$$ where $$\label{eq:chi}
\chi(\gamma, \Theta) := \dfrac{\sum_{k = 1}^{K(\Theta)} \theta_k (1-\gamma L)^{k-1}}{\sum_{k = 1}^{K(\Theta)} \theta_k}.$$

When $\Theta = \Theta_{1:K}$ (as in FedAvg/Reptile), we can derive an even tighter bound that omits the direct dependency on $L/\mu$. Recall that by Corollary \[cor:cond\_fedavg\], the condition number of $\tilde{f}(x, \gamma, \Theta_{1:K})$ is bounded above by $$\tilde{\kappa} = \dfrac{\phi_{K, L}(\gamma)}{\phi_{K, \mu}(\gamma)}$$ where $\phi_{K, \lambda}(\gamma)$ is as in . We will see in the following theorem that the distance between minimizers is controlled by $\phi_{K, L}(\gamma)$ and $\phi_{K, \mu}(\gamma)$.

\[thm:optima\_distance\_fedavg\] If $\gamma < L^{-1}$, then $$\norm{x^*(\gamma, \Theta_{1:K}) - x^*} \leq \sigma_c\left(1+\frac{\sigma_A}{\mu}\right)\dfrac{LK - \phi_{K, L}(\gamma)}{\phi_{K, \mu}(\gamma)}.$$

Informally, the term $KL-\phi_{K, L}(\gamma)$ measures the discrepancy between the surrogate loss function and the true loss function; as $\gamma \to 0$, $\phi_{K, L}(\gamma) \to KL$.

In order to get better control on Theorem \[thm:optima\_distance\_fedavg\] for $\gamma > 0$, we will show that when $\gamma$ is sufficiently small, $\phi_{K, \lambda}(\gamma)$ is close to $K\lambda$.

\[lem:gamma\_cond\] Let $0 \leq \epsilon \leq 1-e^{-K}$ and suppose $$\gamma \leq  \dfrac{\ln(1/(1-\epsilon))}{K\lambda}.$$ Then $$\phi_{K, \lambda}(\gamma) \geq (1-\epsilon)K\lambda.$$

This results in the immediate corollary.

Suppose $\epsilon \leq 1-e^{-K}$ and $$\gamma \leq \dfrac{\ln(1/(1-\epsilon))}{KL}.$$ Then $$\phi_{K, L}(\gamma) \geq (1-\epsilon)KL~~\text{and}~~
    \phi_{K, \mu}(\gamma) \geq (1-\epsilon)K\mu.$$

Plugging these into Theorem \[thm:optima\_distance\_fedavg\], we get the following.

\[cor:optima\_distance\_fedavg\] For $\epsilon \leq 1-e^{-K}$, suppose $$\gamma \leq \dfrac{\ln(1/(1-\epsilon))}{KL}.$$ Then $$\norm{x^*(\gamma, \Theta_{1:K}) - x^*} \leq \sigma_c\left(1+\frac{\sigma_A}{\mu}\right)\dfrac{L}{\mu}\dfrac{\epsilon}{1-\epsilon}.$$

We can also use a similar analysis to bound the distance between $x^*(\gamma, \Theta)$ for different values of $\gamma$. We will focus on the FedAvg/Reptile case. We will show that this distance depends on the discrepancy between the eigenvalues of the matrices ${Q_i(\gamma, \Theta_{1:K})}A_i$.

\[thm:distance\_optima\_gamma\] Let $\gamma_1 \leq \gamma_2 < L^{-1}$. Then $$\norm{x^*(\gamma_1, \Theta_{1:K}) - x^*(\gamma_2, \Theta_{1:K})} \leq \sigma_c\left(1 + \dfrac{\phi_{K, L}(\gamma_2)}{\phi_{K, \mu}(\gamma_2)}\right)\left(\dfrac{\phi_{K, L}(\gamma_1) - \phi_{K, L}(\gamma_2)}{\phi_{K, \mu}(\gamma_1)}\right).$$

The presence of the $\phi$ terms makes the dependence on $|\gamma_1-\gamma_2|$ somewhat opaque. In fact, we have the following simpler (though looser) bound.

\[cor:distance\_optima\_gamma\] Let $\gamma_1 \leq \gamma_2 < L^{-1}$. Then $$\norm{x^*(\gamma_1, \Theta_{1:K}) - x^*(\gamma_2, \Theta_{1:K})} \leq 2\sigma_c\dfrac{L^3}{\mu^2}(\gamma_2-\gamma_1).$$

One particularly useful consequence is that if the client learning rate satisfies $\gamma_t = c/t$ for some constant $c$, then $$\norm{x^*(\gamma_{t+1}, \Theta_{1:K}) - x^*(\gamma_{t}, \Theta_{1:K})} \leq O\left(\dfrac{1}{t^2}\right).$$ We will use this later to show that by decaying the client learning rate in this manner, successive model updates in <span><span style="font-variant:small-caps;">LocalUpdate</span></span>will be closely aligned.

Convergence of FedAvg/Reptile {#sec:convergence}
=============================

We now wish to use Theorem \[thm:optima\_distance\_fedavg\] to understand how quickly FedAvg/Reptile with $K$ local steps converges to $x^*$. However, by Lemma \[thm:sgd\_objective\], we know that performing FedAvg/Reptile with a fixed client learning rate of $\gamma$ will only result in convergence to $x^*(\gamma, \Theta_{1:K})$, the minima of the surrogate loss $\tilde{f}(x, \gamma, \Theta_{1:K})$, not to the true risk minimizer $x^*$. We will therefore analyze the convergence behavior of these algorithms with and without learning rate decay.

Our goal in this section is two-fold. First, we wish to understand the trade-offs incurred by performing local computation instead of mini-batch SGD. Second, we wish to show that by Theorem \[thm:sgd\_objective\], we can analyze federated learning and meta-learning algorithms using classical optimization techniques. There are a large number of important works on federated optimization, that consider more general cases than ours. Unfortunately, the proof techniques behind many of these are relatively opaque, and require careful accounting of the bias incurred by performing local computation. This sometimes leads to either proof errors, or else omitted critical assumptions [@woodworth2020local Appendix A]. Both of these can hinder understanding or make comparisons between convergence rates difficult. By contrast, while limited to a much narrower range of loss functions, our analysis uses essentially standard convex optimization analyses (such as by @rakhlin2011making and @bottou2018optimization), combined with the results from Section \[sec:distance\]. We also emphasize that while we focus on FedAvg/Reptile, our results can be easily extended to more general instances of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>.

Fixed client learning rate
--------------------------

We first wish to understand the setting where the client learning rate $\gamma$ is fixed. Fix a client step-size $\gamma$ and $K$, and for notational convenience, define $$\tilde{f}_\gamma(x) := \tilde{f}(x, \gamma, \Theta_{1:K}),$$ $$\tilde{f}_\gamma^* := \min_x \tilde{f}_\gamma(x),$$ $$x^*_\gamma :=\operatorname*{argmin}_x \tilde{f}_\gamma(x).$$ As $\gamma$ tends towards $0$, $\tilde{f}_\gamma(x)$ converges to the true loss function $f(x)$ (under different notions of convergence depending on the set of assumptions made). For example, under Assumption \[assm1\] this convergence will occur uniformly on ${\mathbb{R}}^d$.

Fix $\gamma_t = \gamma$ for all $t$ in <span><span style="font-variant:small-caps;">LocalUpdate</span></span>. Then, at each iteration $t$ the server starts at a point $x_t$ which it broadcasts to some number of clients. The clients compute local updates $q_t^i$ via ${\textup{InnerLoop}}(\gamma, \Theta)$ (Algorithm \[alg:innerloop\]), and send these values to the server. The server then computes the average $q_t$ of the $q_t^i$ and updates its model via $x_{t+1} = x_t-\eta_tq_t$.

Given a starting point $x$, we let $q_\gamma(x)$ denote the random vector computed by averaging $M$ vectors of the form $q^i = {\textup{InnerLoop}}(i, x, \gamma, \Theta)$ where $i \sim {\mathcal{P}}$. Thus, in Algorithm \[alg:outerloop\] with a constant client learning rate $\gamma$, $x_{t+1} = x_t - \eta_tq_\gamma(x_t)$. Recall that by Theorem \[thm:sgd\_objective\], $$\operatorname*{\mathbb{E}}[q_\gamma(x)] = \nabla \tilde{f}_\gamma(x).$$

Throughout this section, we will assume Assumptions \[assm1\] and \[assm2\], as well as the following "bounded variance" condition.

\[assm3\] For all $x$ and $i$, $\operatorname*{\mathbb{E}}_{z \sim {\mathcal{D}}_i}\norm{\nabla_x f(x ; z) - \nabla f_i(x)}^2 \leq G^2$.

We can then translate this into a bound on the variance of $q(x, \gamma, \Theta_{1:K})$.

\[lem:var\_bound\] Suppose Assumption \[assm3\] holds. Then for all $\gamma \geq 0$, $$\operatorname*{\mathbb{E}}\norm{q_\gamma(x) - \nabla \tilde{f}_\gamma(x)}^2 \leq \dfrac{KG^2}{MB}.$$

We will also use the following bound on the strong convexity parameter of our surrogate loss functions.

\[lem:ft\_params\] Suppose that Assumption \[assm1\] holds. Then for all $t$, $\tilde{f}_\gamma$ is $\mu_\gamma$-strongly convex where $$\mu_\gamma := \phi_{K, \mu}(\gamma) = \dfrac{1-(1-\gamma \mu)^K}{\gamma}.$$ and $L_\gamma$-smooth where $$L_\gamma := \phi_{K, L}(\gamma) = \dfrac{1-(1-\gamma L)^K}{\gamma}.$$

Note that these results follow directly from Lemma \[lem:QA\_structure\_fedavg\] and the fact that the strong convexity and smoothness parameters are governed by the maximum and minimum eigenvalues of ${Q_i(\gamma, \Theta_{1:K})}A_i$.

Using techniques similar to those in [@rakhlin2011making], we arrive at the following descent lemma.

\[lem:descent\_step\] Suppose that Assumptions \[assm1\] and \[assm3\] hold, and that we have step sizes $\{\eta_t\}_{t \geq 1}$ satisfying $$\label{eq:eta_t_bound}
\eta_t \leq \dfrac{\mu_\gamma}{L_\gamma^2}.$$ Then $$\operatorname*{\mathbb{E}}[\norm{x_{t+1} - x_\gamma^*}^2] \leq (1-\eta_t\mu_\gamma)\operatorname*{\mathbb{E}}[\norm{x_t - x_\gamma^*}^2] + \eta_t^2\dfrac{KG^2}{MB}.$$

#### On upper bounds for server learning rates

Note that in Lemma \[lem:descent\_step\], we assume a slightly stronger condition than is often assumed in optimization literature, namely that $\eta \leq 1/L_\gamma\kappa_\gamma$ where $\kappa_\gamma$ is the condition number. Typically, works on optimization would only require $\eta$ to be at most the inverse of the Lipschitz constant. While it is an open question as to whether this condition is necessary, there are a few relevant factors. First, we note that we can relax to $\eta_t \leq L_\gamma^{-1}$ if we strengthen Assumption \[assm3\] to a bounded gradient assumption instead of a bounded variance assumption. Second, when $\gamma$ is moderately large with respect to $\mu$, Lemma \[cor:cond\_fedavg\] implies that $\mu_\gamma \approx L_\gamma$, so gives a similar condition to assuming $\eta_t \leq L_\gamma^{-1}$. Finally, we note that a similar bound on the learning rate was used by @reisizadeh2019fedpaq in conjunction with a bounded variance assumption. While we conjecture that this condition can be relaxed, we leave this for future work.

Applying Lemma \[lem:descent\_step\] repeatedly, we derive at the following.

\[Fixed client LR, fixed server LR\]\[thm:fix\_client\_fix\_server\] Suppose that $\eta \leq \mu_\gamma/L_\gamma^2$. Then the outputs of ${\textsc{LocalUpdate}\xspace}(\eta, \gamma, \Theta_{1:K})$ satisfy $$\operatorname*{\mathbb{E}}[\norm{x_{t+1}-x_\gamma^*}^2] \leq (1-\eta\mu_\gamma)^t\norm{x_1-x_\gamma^*}^2 + \dfrac{\eta KG^2}{\mu_\gamma MB}.$$

Many prior results for fixed client learning rate provide a bound of the same general form as Theorem \[thm:fix\_client\_fix\_server\] (ie. a sum of a decaying term and a constant error term), but bound the distance from $x^*$, rather than from $x_\gamma^*$. This makes the error term's significance more opaque. While non-federated optimization results often have constant error terms due to stochasticity, in federated convergence results (eg. @khaled2020tighter [Theorem 5]), the constant term often does not disappear in deterministic setting ($G=0$). To the reader, it may not be immediately clear why this is the case. This could be due to actual convergence properties, or due to the analysis not being tight. By contrast, our result shows that this error term is an inherent property of the algorithm, as in general, $x_\gamma^* \neq x^*$.

As is the case in general stochastic optimization, a constant learning rate is only sufficient to arrive in a neighborhood of the critical point of the underlying loss. However, this critical point is not the true risk minimizer $x^*$. As in Section \[sec:example\], we see that in heterogeneous settings, client learning rate decay is necessary for convergence to the true risk minimizer.

To make the suboptimality gap tend towards zero, we must decay the server learning rate $\eta$ over time, as in the following theorem.

\[Fixed client LR, decaying server LR\]\[thm:fix\_client\_decay\_server\] Suppose that for all $t \geq 1$, $$\eta_t = \dfrac{a_\gamma}{b_\gamma+t}~~\text{for}~~a_\gamma = \frac{2}{\mu_\gamma}~~\text{and}~~b_\gamma \geq \frac{2L_\gamma^2}{\mu_\gamma^2}.$$ Then the outputs of ${\textsc{LocalUpdate}\xspace}(\{\eta_t\}, \gamma, \Theta_{1:K})$ satisfy $$\operatorname*{\mathbb{E}}[\norm{x_{t}-x_\gamma^*}^2] \leq \dfrac{\nu_\gamma}{b_\gamma + t}$$ where $$\nu_\gamma := \max\left\{\dfrac{4KG^2}{\mu_\gamma^2MB}, (b_\gamma + 1)\norm{x_1-x_\gamma^*}^2 \right\}.$$

We can now derive a convergence rate towards the true risk minimizer, $x^*$.

\[cor:fix\_client\_decay\_server\] Fix $\epsilon$ satisfying $\epsilon \leq 1-e^{-K}$. Suppose $$\gamma \leq \dfrac{\ln(1/(1-\epsilon))}{KL}$$ and that $\eta_t$ is as in Theorem \[thm:fix\_client\_decay\_server\]. Then the outputs of ${\textsc{LocalUpdate}\xspace}(\{\eta_t\}, \gamma, \Theta_{1:K})$ satisfy $$\operatorname*{\mathbb{E}}[\norm{x_t-x^*}^2] \leq \dfrac{2\nu_\gamma}{b_\gamma + t} + 2\sigma_c^2\left(1 + \dfrac{\sigma_a}{\mu}\right)^2\dfrac{L^2}{\mu^2}\dfrac{\epsilon^2}{(1-\epsilon)^2}$$ where $$\nu_\gamma \leq \max\left\{\dfrac{4G^2}{(1-\epsilon)^2\mu^2 KMB} , (b_\gamma+1)\norm{x_1-x_\gamma^*}^2\right\}.$$

While notationally complex, this result has a few important facets. Define $$\beta_1 = \dfrac{G^2}{(1-\epsilon)^2\mu^2 KMB},~~\beta_2 = (b_\gamma+1)\norm{x_1-x_\gamma^*},~~\beta_3 = \sigma_c^2\left(1 + \dfrac{\sigma_a}{\mu}\right)^2\dfrac{L^2}{\mu^2}\dfrac{\epsilon^2}{(1-\epsilon)^2}.$$ Thus, this result shows that, $$\operatorname*{\mathbb{E}}[\norm{x_t-x^*}^2] \leq \mathcal{O}\rbr*{ \frac{\max\{\beta_1, \beta_2\}}{t} + \beta_3}.$$

As $\epsilon \to 0$, $\gamma \to 0$, <span><span style="font-variant:small-caps;">LocalUpdate</span></span>becomes roughly equivalent to mini-batch SGD with batches of size $KB$ on $M$ clients. In fact, $\beta_1$ is the convergence term we would derive from performing mini-batch SGD with batches of size $KB$ on $M$ clients per round. In particular, we get a variance reduction of $(KMB)^{-1}$. This variance reduction is similar in nature to work by @woodworth2020local in the homogeneous setting ($\sigma_A = \sigma_c = 0$), which shows an analogous improvement in convergence rates. Note that $\beta_1$ gets larger as $\epsilon$ gets larger. Thus, the variance reduction is reduced in heterogeneous settings.

The term $\beta_2$ measures the initial suboptimality to the surrogate loss function. It is here that having a small $\gamma$ may incur a price. Note that $b_\gamma$ is twice the condition number of $\tilde{f}_\gamma$. As shown in Corollary \[cor:cond\_fedavg\] and in the discussion thereafter, larger $\gamma$ lead to smaller $b_\gamma$ due to the implicit regularization of local computation. Conversely, at $\gamma = 0$, $b_\gamma = 2 L^2/\mu^2$, which may dominate $\beta_1$. Thus, smaller $\gamma$ lead to a larger effective initial suboptimality gap.

The last term $\beta_3$ measures the discrepancy between the surrogate loss function and the true loss function. If $\sigma_c = 0$, for instance, when the data is completely homogeneous, this discrepancy is 0. We then recover similar results to that of @woodworth2020local, which shows a $1/K$ improvement in convergence for Local SGD with $K$ steps, but in the heterogeneous setting. If $\sigma_c \neq 0$, we can still remove the effect of heterogeneity by setting $\epsilon = 0$ (which means setting $\gamma = 0$). In this case, FedAvg reduces to mini-batch SGD with batches of size $KB$ on each client.

Decaying client and server learning rates
-----------------------------------------

In this section, we will show that by decaying the client and server learning rates appropriately, we can derive a bound on $\norm{x_t-x^*}^2$ that does not require setting the learning rates in terms of the desired optimality gap $\epsilon$.

Throughout this section, we again focus on the FedAvg/Reptile setting. At every iteration $t$, we will use client learning rate $\gamma_t$ and server learning rate $\eta_t$ that decay at a $1/t$ rate. Let $\kappa = L/\mu$ denote the condition number of $f(x)$. We have the following theorem.

\[Decaying client LR, decaying server LR\]\[thm:convergence\_overall\] Let $$\eta_t = \dfrac{a_t}{b+t}~~\text{for}~~a_t = \frac{3}{\mu_t}~~\text{and}~~b = 3\kappa^2$$ and let $$\gamma_t = \min\left\{\dfrac{1}{L(b + t)}, \dfrac{\ln(2)}{K\mu}\right\}.$$ Suppose we run <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\gamma_t, \eta_t$ as above and $\Theta = \Theta_{1:K}$ to produce iterates $x_t$. Then for all $t \geq 1$, $$\operatorname*{\mathbb{E}}[\norm{x_t-x^*}^2] \leq \dfrac{2\nu + 16\sigma_c^2\kappa^2}{b + t}.$$ where $$\nu := \max\left\{\dfrac{18G^2}{\mu^2 KMB}, (b+1)\norm{x_1-x_1^*}^2\right\}.$$

No attempt was made to optimize constants. Rather, the point was to show that we can derive bounds on the distance to the true risk minimizer $x^*$ that hold for all $t$ and that help illustrate the implicit trade-offs in <span><span style="font-variant:small-caps;">LocalUpdate</span></span>. In particular, we do not require that $\gamma$, $\eta$, or $t$ be defined in terms of the desired suboptimality gap $\epsilon$. Rather, all that was required was learning rate decay on the order of $O(1/t)$ at both the server and the client.

Recall that by Theorem \[thm:fix\_client\_fix\_server\], some form of learning rate decay is necessary for convergence to the true risk minimizer. We therefore see that in some settings, learning rate decay is sufficient for <span><span style="font-variant:small-caps;">LocalUpdate</span></span>to converge to the true risk minimizer. While this was previously shown in the finite-sum setting for strongly convex functions by @li2019convergence, this result required bounded gradients, and did not illustrate the trade-offs in using FedAvg/Reptile over mini-batch SGD.

#### Comparison to mini-batch SGD

Similar to Theorem 1 of [@woodworth2020local], we see that performing <span><span style="font-variant:small-caps;">LocalUpdate</span></span>incurs a kind of variance reduction of $1/MB$ when compared to performing vanilla SGD over a shuffled version of the entire dataset. We also see that using <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\gamma > 0$ does incur a potential benefit: Rather than having the convergence depend on the suboptimality gap $\norm{x_1-x^*}^2$ (as is the case for mini-batch SGD), it depends on $\norm{x_1-x_1^*}^2$. This may be much smaller depending on the initialization. For example, recall that by Lemma \[lem:asymptotic\], if $\gamma_1 > 0$, then as $K \to \infty$, $x_1^*$ tends to the "one-shot average" of the client minimizers, which typically requires many fewer communication rounds to estimate than the true risk minimizer. However, this reduction only benefits the convergence up to a point, in which case it becomes beneficial to use smaller $\gamma$.

As suggested by lower bounds on FedAvg in [@karimireddy2019scaffold] on FedAvg, our bounds do not show that FedAvg always converges faster than mini-batch SGD, and in fact, doing so may not be possible without further assumptions (such as a bound on the heterogeneity among clients) or more sophisticated optimization techniques (such as the use of control variates in SCAFFOLD [@karimireddy2019scaffold]).

While our theoretical results are only valid for the case of quadratic loss functions, we conjecture that even in much broader settings, the choice of learning rate $\gamma$ still dictates a trade-off between accuracy and initial convergence. we will show in the next section that this holds empirically, even in non-convex settings.

Experimental results {#sec:exp_valid}
====================

In this section, we analyze <span><span style="font-variant:small-caps;">LocalUpdate</span></span>empirically, in order to understand how the choice of client and server learning rates impact convergence in more realistic machine learning tasks. In particular, we focus on (not necessarily convex) tasks and datasets that reflect federated learning in practice. We will show that both the choice of client learning rate and tuning of the corresponding server learning rate can be vital to attain the best performance of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>, especially in limited communication settings.

#### Datasets and models

We use four different datasets: the federated extended MNIST dataset (FEMNIST) [@caldas2018leaf], the federated version of CIFAR-100 created by @reddi2020adaptive, the Shakespeare dataset [@caldas2018leaf] and the Stack Overflow dataset [@stackoverflow]. The first two are image datasets, the second two are text datasets. All datasets are publicly available. We specifically use the versions available in TensorFlow Federated [@ingerman2019tff]. All four datasets contain training and test clients. For the purposes of our experiments, we only use the training clients, as our work only concerns the loss of clients in the training population. Notably, our work does not broach the subject of generalization, which we leave to future work. The number of clients and examples in each dataset is presented in Table \[table:datasets\].

  Dataset            \# of Clients   Total \# of Examples
  ---------------- --------------- ----------------------
  CIFAR-100                    500                 50,000
  FEMNIST                    3,400                671,585
  Shakespeare                  715                 16,068
  Stack Overflow           342,477            135,818,730

  : Dataset statistics.<span data-label="table:datasets"></span>

For FEMNIST, we train a moderately-sized CNN (the same as used in [@mcmahan17fedavg]) to perform character recognition. For CIFAR-100, we train a ResNet-18 (where we replace the batch norm layers with group norm, as suggested by [@hsieh2019non] and used by [@reddi2020adaptive]). For Shakespeare, we train an RNN with 2 LSTM layers to perform next-character-prediction. For Stack Overflow, we perform two distinct task: tag prediction (TP) and next-word-prediction (NWP). For Stack Overflow TP, we use a logistic regression classifier with one-versus-all classification. Note that this implies that Stack Overflow TP is a *convex* task. For Stack Overflow NWP, we train an RNN with 1 LSTM layer to perform next-word-prediction. These five tasks were previously analyzed by @reddi2020adaptive for the purposes of comparing adaptive and non-adaptive federated optimization methods. For further details on datasets and models, see Appendix \[sec:models\].

#### Implementation and hyperparameters

We implement <span><span style="font-variant:small-caps;">LocalUpdate</span></span>in TensorFlow Federated [@ingerman2019tff]. In all experiments, the set of clients ${\mathcal{I}}$ is finite and we let ${\mathcal{P}}$ be the uniform distribution over these clients. Each ${\mathcal{D}}_i$ is the uniform distribution over some finite set of client examples. We analyze the performance of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\Theta = \Theta_{1:10}$ (ie. FedAvg/Reptile with $K = 10$ local steps) across the tasks discussed above. For FEMNIST, CIFAR-100, and Shakespeare, we sample $M = 10$ clients per round, while for Stack Overflow, we sample $M = 50$ clients (due to its much larger number of clients). In order to derive fair comparisons for different hyperparameter settings, we use a random seed to determine which clients are sampled at each round from ${\mathcal{P}}$. All plots are made using the same seed. We sample clients without replacement within each round, but with replacement across rounds. We use a batch size of $B = 20$ for FEMNIST and CIFAR-100, $B = 4$ for Shakespeare, and $B = 16$ for Stack Overflow.

Fixed server learning rates
---------------------------

We first perform a comparable analysis to that in Section \[sec:example\] above for FEMNIST, CIFAR-100, and Shakespeare. Namely, we fix the server learning rate $\eta = 0.01$, and see how the training loss varies as a function of the client learning rate $\gamma$. We vary $\gamma$ over $$\gamma \in \{0, 10^{-3}, \dots, 10^{-1}, 1, 10, 10^2\}$$ and omit results that diverged due to the client learning rate being set too large. We plot the true loss function $f(x)$ (defined in ) in Figure \[fig:constant\_server\_lr\].

<span>.31</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with fixed server learning rate $\eta = 0.01$ and $\Theta = \Theta_{1:10}$, and varying client learning rate $\gamma$.<span data-label="fig:constant_server_lr"></span>](figures/emnist_cnn_constant_server_train_loss.pdf "fig:"){width="\linewidth"}

<span>.31</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with fixed server learning rate $\eta = 0.01$ and $\Theta = \Theta_{1:10}$, and varying client learning rate $\gamma$.<span data-label="fig:constant_server_lr"></span>](figures/cifar100_constant_server_train_loss.pdf "fig:"){width="\linewidth"}

<span>.31</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with fixed server learning rate $\eta = 0.01$ and $\Theta = \Theta_{1:10}$, and varying client learning rate $\gamma$.<span data-label="fig:constant_server_lr"></span>](figures/shakespeare_constant_server_train_loss.pdf "fig:"){width="\linewidth"}

We see that on all three tasks, especially CIFAR-100, the choice of client learning rate can impact not just the speed of convergence, but what point the algorithm converges to. In general, we see very similar behavior to that in Section \[sec:example\], despite the non-convex loss functions involved in all three tasks. For both FEMNIST and CIFAR-100, smaller client learning rates eventually reach lower training losses than higher learning rates. This is particularly evident in the results for CIFAR-100. While $\gamma = 10^{-2}$ initially performs better than all other methods, it is eventually surpassed by $\gamma = 10^{-3}$, and $\gamma = 0$ ends up obtaining a comparable accuracy. Conversely, setting $\gamma = 1.0$ results in a sub-optimal training loss.

We also see that while suboptimality gaps exist for FEMNIST and Shakespeare, they are much smaller than for CIFAR-100. Thus, our results also suggest that the theoretical suboptimality of larger $\gamma$ may not be as important a facet in practice. We see here that despite the asymptotic suboptimality of FedAvg with $\gamma > 0$, in realistic settings where the number of communication rounds is limited, there may be little to no disadvantage to using $\gamma > 0$. In fact, as we see for CIFAR-100, there may be advantages to using larger $\gamma$ in settings with relatively few communication rounds.

Deriving fair comparisons between different client learning rates
-----------------------------------------------------------------

While the results in the previous section indicate the importance of $\gamma$, these comparisons are in some sense unfair. In particular, there is no reason that we need to fix the server learning rate $\eta$ across different values of $\gamma$. In fact, this ignores differences in the Lipschitz constant of the associated surrogate loss functions. As shown in Corollary \[cor:cond\_fedavg\], choosing smaller $\gamma$ increases the Lipschitz constant $L_\gamma$ of the surrogate loss function, which in turn generally necessitates a smaller server learning rate. Thus, choosing larger $\gamma$ enable higher values of $\eta$. To justify this further, we plot the $\ell_2$ norm of the update applied to the server model at each round. That is, recall that in Algorithm \[alg:outerloop\], we update the model via $$x_{t+1} = x_t - \eta q_t.$$ In Figure \[fig:constant\_server\_lr\_update\_norm\], we plot $\norm{q_t}$ for different choices of $\gamma$ in the CIFAR-100 task. Specifically, we plot the mean value of $\norm{q_t}$ for each consecutive 1000 rounds, as well as the standard deviation within those rounds.

<span>.6</span> ![The $\ell_2$ norm of the model update of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>on CIFAR-100 at each round with fixed server learning rate $\eta = 0.01$ and varying client learning rate $\gamma$. We plot the mean $\ell_2$ norm across every consecutive 1000 rounds (bold lines), as well as the standard deviation across these rounds (pale regions).<span data-label="fig:constant_server_lr_update_norm"></span>](figures/cifar100_constant_server_update_norm.pdf "fig:"){width="\linewidth"}

As expected, the $\ell_2$ norm of the updates is highly dependent on the choice of $\gamma$. In particular, smaller learning rates lead to larger model updates. We also see that the absolute variance of the model update norm increases as well. As discussed after Corollary \[cor:fix\_client\_decay\_server\], the variance reduction offered by setting $\gamma$ small SGD is potentially at odds with the increase in the condition number increase, a phenomenon reflected in Figure \[fig:constant\_server\_lr\_update\_norm\]. Thus, we see that to derive a fair comparison between client learning rates, we must also allow the server learning rate to increase as the client learning rate increases.

One important caveat to this observation is that it is not enough to simply directly tie the client and server learning rate together. This is in fact what is done in the original incarnation of FedAvg in [@mcmahan17fedavg]. As discussed in Section \[sec:local\_update\_algs\], the original "vanilla" FedAvg algorithm corresponds to <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\Theta = \Theta_{1:K}$ and $\gamma = \eta$. This does allow for the use of larger server learning rates with smaller client learning rates. However, this is not enough to necessarily derive optimal performance of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>.

To demonstrate this, we plot the performance of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\eta = \gamma$ and $\Theta = \Theta_{1:10}$ for the CIFAR-100 task in Figure \[fig:cifar100\_vanilla\_fedavg\]. We see results in stark contrast to Figure \[fig:constant\_server\_lr\](b). In particular, $\gamma = 1.0$ actually results in an increased training loss, and $\gamma = 0.1$ outperforms $\gamma = 0.01$ for most rounds. Moreover, by setting $\gamma = \eta$, we do not allow $\gamma = 0$, even though $\gamma > 0$ necessarily results in a surrogate loss that does not match the true training loss. While this version of FedAvg is convenient from an implementation and hyperparameter tuning point of view, it does not result in the best performance of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>. As we show in the next section, we can greatly improve performance by tuning client and server learning rates separately.

<span>.6</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\Theta = \Theta_{1:10}$ and $\gamma = \eta$ for varying values of $\gamma$ on the CIFAR-100 task.<span data-label="fig:cifar100_vanilla_fedavg"></span>](figures/cifar100_vanilla_fedavg_train_loss.pdf "fig:"){width="\linewidth"}

On the importance of tuned server learning rates {#sec:tuned_server}
------------------------------------------------

Based on the discussion in the section above, to give the most fair comparisons between different values of $\gamma$ in <span><span style="font-variant:small-caps;">LocalUpdate</span></span>, we must also tune the server learning rate $\eta$. We perform the same experiments as in Figure \[fig:constant\_server\_lr\], but where we tune the server learning rate $\eta$ for each choice of $\gamma$. We vary $\eta$ over $$\label{eq:server_grid}
\eta \in \{10^{-3}, 10^{-2.5}, \dots, 10^{1}\}$$ and select $\eta$ that results in the smallest average training loss over the last 100 rounds. We note that we use this averaging method as a single round of federated learning only samples a small number of clients. We plot the loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>under these settings in Figures \[fig:tuned\_server\_lr\]. For a list of all client learning rates and the corresponding best server learning rate for each task, see Appendix \[sec:tuned\_lrs\].

Notably, we still see the same general trend discussed in our convergence rates in Section \[sec:convergence\]. While large values of $\gamma$ may obtain a smaller training loss initially, eventually smaller values of $\gamma$ perform comparably, if not better. However, it is instructive to note that this may take many thousands of communication rounds. In particular, for all three tasks, we see that $\gamma = 0$ does not perform comparably larger $\gamma$ until near the end of our training procedure. We thus see the following:

**By performing appropriate server and client learning rate tuning, we can mitigate the suboptimality of FedAvg/Reptile for $\gamma > 0$, especially when the number of communication rounds is limited.**

In particular, our empirical results suggest that in realistic federated learning training tasks, the suboptimality of FedAvg/Reptile may only be an asymptotic concern. If we can only perform a limited number of training rounds, it often does benefit us to use larger values of $\gamma$. We note that this aligns with our discussion of how $\gamma > 0$ leads to condition number regularization (see Corollary \[cor:cond\_fedavg\]).

<span>.31</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\Theta = \Theta_{1:10}$, varying client learning rate $\gamma$, and tuned server learning rate $\eta$.<span data-label="fig:tuned_server_lr"></span>](figures/emnist_cnn_tuned_server_train_loss.pdf "fig:"){width="\linewidth"}

<span>.31</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\Theta = \Theta_{1:10}$, varying client learning rate $\gamma$, and tuned server learning rate $\eta$.<span data-label="fig:tuned_server_lr"></span>](figures/cifar100_tuned_server_train_loss.pdf "fig:"){width="\linewidth"}

<span>.31</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\Theta = \Theta_{1:10}$, varying client learning rate $\gamma$, and tuned server learning rate $\eta$.<span data-label="fig:tuned_server_lr"></span>](figures/shakespeare_tuned_server_train_loss.pdf "fig:"){width="\linewidth"}

Adaptive optimization and <span><span style="font-variant:small-caps;">LocalUpdate</span></span> {#sec:adaptive_localupdate}
------------------------------------------------------------------------------------------------

In order to understand the behavior of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>more generally, we perform similar experiments on the Stack Overflow dataset. However, as shown by @reddi2020adaptive, performance of FedAvg on next-word-prediction tasks can be greatly improved by the use of adaptive optimization. In particular, @reddi2020adaptive found that the use of the <span><span style="font-variant:small-caps;">Yogi</span></span>optimizer [@zaheer2018adaptive] improved performance in a wide variety of settings, including on the same tag-prediction and next-word-prediction tasks. Note that the <span><span style="font-variant:small-caps;">Yogi</span></span>optimizer is similar to the <span><span style="font-variant:small-caps;">Adam</span></span>optimizer [@kingma2014adam], except that it uses a kind of *additive* adaptive update that can improve convergence by making more controlled progress. For more details, see [@zaheer2018adaptive].

Thus, for these tasks, we use a modified version of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>in which the server uses the client update $q_t$ as an estimate of the gradient of the loss function, and applies <span><span style="font-variant:small-caps;">Yogi</span></span>to this gradient. That is, we use Algorithm \[alg:outerloop\], but update the model via $$x_{t+1} = {\textsc{Yogi}\xspace}(x_t, q_t, \eta_t).$$ We use the version of <span><span style="font-variant:small-caps;">Yogi</span></span>proposed by @zaheer2018adaptive, with first momentum term of $\beta_1 = 0.9$, second momentum parameter of $\beta_2 = 0.99$, an initial accumulator value of $0$, and an $\epsilon$ value of $10^{-5}$. We then perform analogous experiments to those above, where we first fix $\eta = 0.01$ and vary $\gamma$. However, due to the size of the Stack Oveflow dataset (see Table \[table:datasets\]), we did not compute the total loss $f(x)$ over all clients. Instead, we plot the average loss of the clients that participated in a given round, *before* local training occurs. This "loss at current round" can be viewed as a stochastic estimate of the true loss $f(x)$. This loss is plotted in Figure \[fig:stack\_overflow\_constant\]. We also perform experiments where we vary $\gamma$ and tune $\eta$. We select the value of $\eta$ with the smallest average training loss over the last 100 communication rounds. The result is given in Figure \[fig:stack\_overflow\_tuned\]. For a list of all client learning rates and the corresponding best server learning rate for each task, see Appendix \[sec:tuned\_lrs\].

<span>.45</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with server <span><span style="font-variant:small-caps;">Yogi</span></span>and $\Theta = \Theta_{1:10}$ on Stack Overflow with varying client learning rate $\gamma$ and fixed server learning rate $\eta = 0.01$.<span data-label="fig:stack_overflow_constant"></span>](figures/so_lr_constant_server_train_loss.pdf "fig:"){width="\linewidth"}

<span>.45</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with server <span><span style="font-variant:small-caps;">Yogi</span></span>and $\Theta = \Theta_{1:10}$ on Stack Overflow with varying client learning rate $\gamma$ and fixed server learning rate $\eta = 0.01$.<span data-label="fig:stack_overflow_constant"></span>](figures/so_nwp_constant_server_train_loss.pdf "fig:"){width="\linewidth"}

<span>.45</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with server <span><span style="font-variant:small-caps;">Yogi</span></span>and $\Theta = \Theta_{1:10}$ on Stack Overflow with varying client learning rate $\gamma$ and tuned server learning rate $\eta$.<span data-label="fig:stack_overflow_tuned"></span>](figures/so_lr_tuned_server_train_loss.pdf "fig:"){width="\linewidth"}

<span>.45</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with server <span><span style="font-variant:small-caps;">Yogi</span></span>and $\Theta = \Theta_{1:10}$ on Stack Overflow with varying client learning rate $\gamma$ and tuned server learning rate $\eta$.<span data-label="fig:stack_overflow_tuned"></span>](figures/so_nwp_tuned_server_train_loss.pdf "fig:"){width="\linewidth"}

For fixed $\eta$, large $\gamma$ leads to an initially smaller loss that is eventually beaten or matched by smaller $\gamma$. However, we see similar behavior among all but the largest $\gamma$ in both tasks, potentially due to the adaptivity in <span><span style="font-variant:small-caps;">Yogi</span></span>. Looking at Figure \[fig:stack\_overflow\_tuned\], we see that tuning the server learning rate has slightly different effects on the two tasks: While Stack Overflow TP with tuned $\eta$ leads to the smaller $\gamma$ performing better throughout the training process, Stack Overflow NWP with tuned $\eta$ actually allows larger $\gamma$ to achieve lower loss throughout. Notably, we see that the gap between large $\gamma$ and small $\gamma$ in Stack Overflow NWP winnows as the number of rounds increases. Our results reinforce the notion that in communication-limited settings, server learning rate tuning is critical to ensure the best performance possible. This holds even when not using SGD on the server, but instead using an adaptive optimizer.

Automatic learning rate decay {#sec:lr_decay}
=============================

In short, our results in the section above suggest that depending on the desired number of communication rounds, we may wish to use different client learning rates. Unfortunately, this requires a large degree of hyperparameter tuning. Both the client and server learning rate must be tuned. However, if the number of communication rounds in a federated learning system is limited, it may not be feasible to conduct extensive hyperparameter tuning, as the communication rounds required to do so may be better utilized by training your model for more rounds.

To help reduce the amount of learning rate tuning required, we propose a method for automatic learning rate decay that helps mitigate the need for *client* learning rate tuning. Our method will utilize our theoretical and empirical observations above showing that large client learning rates should be used initially, while smaller learning rates eventually reach a lower training loss. By decaying the client learning rate automatically over time, we mitigate the need to tune it. Instead, it can be set to any moderate value that does not result in divergent behavior on clients. We will be particularly concerned with systems-level constraints (such as those encountered in federated learning) when describing our method.

One particularly important restriction in many local update settings is the ability to compute the training loss $f(x)$. Recall that we defined $$f(x) := \operatorname*{\mathbb{E}}_{i \sim {\mathcal{P}}} [f_i(x)],~~~f_i(x) := \operatorname*{\mathbb{E}}_{z \sim {\mathcal{D}}_i} [f(x; z)]$$ where ${\mathcal{P}}$ is a distribution over clients, and ${\mathcal{D}}_i$ is the client's data distribution. In settings with limited communication, it may not be feasible to sample most or even a moderate fraction of the clients from ${\mathcal{P}}$. Moreover, sampling a client $i \sim {\mathcal{P}}$ purely for the purposes of estimating the loss may be a waste of resources, as that client could be used for training purposes.

Thus, we propose a version of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>that simultaneously computes local updates and estimates the loss function. Specifically, at each round $t$, we sample a set $I_t$ of clients of size $M$. For each $i \in I_t$, we first compute an estimate $\ell_t^i$ of the local loss function $f_i(x_t)$, as well as the client's model update $q_t^i$ as defined in Algorithm \[alg:innerloop\]. Note that it is important that $f_i(x_t)$ is estimated *before* computing any local updates, as otherwise we risk overfitting to the client's dataset. The server then computes $$\label{eq:f_t_eq}
\ell_t := \frac{1}{M} \sum_{i \in I_t} \ell_t^i,~~~q_t := \frac{1}{M}\sum_{i \in I_t} q_t^i.$$

As in <span><span style="font-variant:small-caps;">LocalUpdate</span></span>, we will use this $q_t$ to update the model. The loss estimate $\ell_t$ will be used to determine whether to decay the client or server learning rates. We will use the strategy of decaying learning rates on plateaus. If the loss $f_t$ is not sufficiently small, we will decay the client and server learning rates. We record the minimum loss seen up to round $t$, and then check if $$\ell_t < \min_{j < t} \ell_j - \Delta$$ for some small $\Delta > 0$. Intuitively, if the loss $\ell_t$ has not decreased sufficiently, we may be at a suboptimal point (due to a large client learning rate $\gamma$) and therefore should decrease $\gamma$, which necessitates a decrease in $\eta$. We propose decaying $\gamma$ and $\eta$ by fixed factors $\alpha, \beta \in (0, 1)$, as this kind of "staircase" learning rate schedule has been repeatedly demonstrated to perform well across many tasks [@krizhevsky2014one; @goyal2017accurate]. A full version of this algorithm, which we refer to as <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>, is given in Algorithms \[alg:outerloop\_decay\] and \[alg:innerloop\_decay\].

\[t\]<span>.52</span>

\[H\]

**OuterLoop$(x, \eta, \gamma, \Theta, \Delta > 0, \alpha, \beta \in (0, 1))$:**

$x_1 = x$ $I_t \leftarrow$ (random set of $M$ clients) $\ell_t^i, q_t^i \leftarrow \text{InnerLoop}(i, x_t, \gamma, \Theta)$ $\ell_t \leftarrow (\nicefrac{1}{M})\sum_{i \in I_t} \ell_t^i$ $q_t \leftarrow (\nicefrac{1}{M})\sum_{i \in I_t} q_{t}^i$ $x_{t+1} = x_t-\eta_tq_t$ $\gamma \leftarrow \alpha\gamma,~~\eta \leftarrow \beta\eta$ return $x_{T+1}$

\[t\]<span>.46</span>

\[H\]

**${\textup{InnerLoop}}(i, x, \gamma, \Theta$):**

$x_1 = x$ $\ell \leftarrow \operatorname*{\mathbb{E}}_{z\sim{\mathcal{D}}_i} [f(x; z)]$ sample a set $S_k$ of size $B$ from ${\mathcal{D}}_i$ $g_k = (\nicefrac{1}{B}) \sum_{z \in S_k} \nabla f(x_k ; z)$ $x_{k+1} \leftarrow x_k - \gamma g_k$ return $\ell,~\sum_{k = 1}^{K(\Theta)} \theta_k g_k$

Practical refinements to <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span> {#sec:practical_refinements}
----------------------------------------------------------------------------------------------------

In this section we describe three practical refinements to <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>that can improve convergence behavior. First, in settings where the number of clients sampled per round $M$ is sufficiently small compared to $|{\mathcal{I}}|$, we can instead estimate the loss by using a moving window average across rounds. That is, given some window size $W > 0$, we can compute an estimate of the loss $\hat{\ell}_t$ with reduced variance via $$\hat{\ell}_t = \dfrac{1}{W} \sum_{j = 0}^{W-1} \ell_{t-j}$$ where $f_t$ is as in . This is the average loss of all clients in the last $W$ rounds. Note that Algorithm \[alg:outerloop\_decay\] corresponds to $W = 1$. We would then decay the client and server learning rates if $$\label{eq:fhat_condition}
\hat{\ell}_t > \min_{j < t} \hat{\ell}_j - \Delta.$$

Second, even using moving windows to estimate $f(x)$, the heterogeneity of clients can still cause problematic variance. Thus, one can instead decay the learning rates if there has been no progress (relative to $\Delta$) for $P$ *consecutive* rounds. That is, we keep a counter for how many consecutive rounds the condition holds. If this counter ever reaches $P$, we then decay $\gamma$ and $\eta$. If $$\hat{\ell}_t \leq \min_{j < t} \hat{\ell}_j - \Delta$$ then we reset the counter to 0. Note that Algorithm \[alg:outerloop\_decay\] corresponds to $P = 1$.

Last, it is often useful to have a *cooldown period* $C$, where after decaying the learning rates $\gamma$ and $\eta$, we do not decay the learning rate for the next $C$ rounds. This is beneficial both for recovering a new estimate of the loss function, and for ensuring that the learning rate does not decay too frequently. Additionally, we recommend using this cooldown period for the first $C$ rounds as well, as this allows one to develop a better estimate of $f(x)$ before decaying the learning rate.

While in practice, setting $W, P$, and $C$ may seem difficult, we found that setting them all to be the same value led to good behavior across datasets and tasks. Moreover, the value can be estimated using simple heuristics. For example, suppose that we think that $N$ randomly sampled clients are sufficient to give a good representation of ${\mathcal{P}}$. Then, $P = W = C = \lceil N/M \rceil$ should serve as a default value for these parameters. In practice, we found that fixing all three values to some moderate constant (ex. $P = W = C = 100$) was sufficient, even across datasets with widely varying numbers of clients.

Empirical evaluation
--------------------

The primary motivation for <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>is removing the need for client learning rate tuning. Intuitively, we can use any moderately large value of $\gamma$ that results in non-divergent client behavior, and this will be gradually scaled back over time as we reach suboptimal critical points of the corresponding surrogate loss. To validate this, we compare <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>to <span><span style="font-variant:small-caps;">LocalUpdate</span></span>on FEMNIST, CIFAR-100, Shakespeare, and Stack Overflow. For Stack Overflow, we again use a modified version where we apply <span><span style="font-variant:small-caps;">Yogi</span></span>on the server. In particular, we compare tuned but constant $\gamma$ and $\eta$, to <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>.

#### Implementation and hyperparameters

We implement <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>in TensorFlow Federated as well. When computing the client's loss estimate $\ell$, we compute the average loss over the entire client dataset. Unlike in the previous section, we do not tune the client learning rate $\gamma$. We instead vary $\gamma$ over $\{10^2, 10, 1, 10^{-1}, 10^{-2}\}$ and select the largest $\gamma$ that results in a non-divergent training loss, as we intend to test the hypothesis that <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>does not need $\gamma$ to be tuned in the same way that <span><span style="font-variant:small-caps;">LocalUpdate</span></span>does. In all settings, we still tune the server learning rate, as this is can be vital for getting accurate and fair comparisons (as discussed in Section \[sec:tuned\_server\]). We use the same server learning rate grid as with <span><span style="font-variant:small-caps;">LocalUpdate</span></span>, given in , and select the value that reaches the lowest average training loss over the last 100 communication rounds. A table of learning rates used by <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>for each task is given in Table \[table:decay\_eta\].

\[t\]<span>l\*<span>2</span><span>c</span></span> Task & $\gamma$ & $\eta$\
CIFAR-100 & 0 & -1\
FEMNIST & -1 & -1\
Shakespeare & 0 &\
StackOverflow NWP & -1 & -\
StackOverflow TP & 2 & 0\

We tune no other hyperparameters. In <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>, specifically Algorithm \[alg:outerloop\_decay\], we set $\Delta = 10^{-4}, \alpha = 0.1, \beta = 0.9$. We also use the practical refinements discussed in Section \[sec:practical\_refinements\], setting $P = W = C = 100$. All other parameters are identical to that of standard <span><span style="font-variant:small-caps;">LocalUpdate</span></span>.

<span>.31</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\Theta = \Theta_{1:10}$, varying client learning rate $\gamma$, and tuned server learning rate $\eta$. We also plot the cross-entropy loss of <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>with tuned server learning rate (ReduceLR).<span data-label="fig:compare_decay"></span>](figures/emnist_cnn_compare_reduce_train_loss.pdf "fig:"){width="\linewidth"}

<span>.31</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\Theta = \Theta_{1:10}$, varying client learning rate $\gamma$, and tuned server learning rate $\eta$. We also plot the cross-entropy loss of <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>with tuned server learning rate (ReduceLR).<span data-label="fig:compare_decay"></span>](figures/cifar100_compare_reduce_train_loss.pdf "fig:"){width="\linewidth"}

<span>.31</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with $\Theta = \Theta_{1:10}$, varying client learning rate $\gamma$, and tuned server learning rate $\eta$. We also plot the cross-entropy loss of <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>with tuned server learning rate (ReduceLR).<span data-label="fig:compare_decay"></span>](figures/shakespeare_compare_reduce_train_loss.pdf "fig:"){width="\linewidth"}

<span>.31</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with server <span><span style="font-variant:small-caps;">Yogi</span></span>and $\Theta = \Theta_{1:10}$. We vary the client learning rate $\gamma$ and tune the server learning rate $\eta$. We also plot the cross-entropy loss of <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>with server <span><span style="font-variant:small-caps;">Yogi</span></span>and tuned server learning rate (ReduceLR).<span data-label="fig:stackoverflow_decay"></span>](figures/so_lr_compare_reduce_train_loss.pdf "fig:"){width="\linewidth"}

<span>.31</span> ![Cross-entropy training loss of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with server <span><span style="font-variant:small-caps;">Yogi</span></span>and $\Theta = \Theta_{1:10}$. We vary the client learning rate $\gamma$ and tune the server learning rate $\eta$. We also plot the cross-entropy loss of <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>with server <span><span style="font-variant:small-caps;">Yogi</span></span>and tuned server learning rate (ReduceLR).<span data-label="fig:stackoverflow_decay"></span>](figures/so_nwp_compare_reduce_train_loss.pdf "fig:"){width="\linewidth"}

#### Results

For FEMNIST, CIFAR-100, and Shakespeare, we plot the results in Figure \[fig:compare\_decay\], where we also compare to the tuned results in Section \[sec:tuned\_server\]. We find that in all three tasks, <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>eventually does as well as <span><span style="font-variant:small-caps;">LocalUpdate</span></span>with tuned $\gamma$, without the need for client learning rate tuning. For FEMNIST and Shakespeare, we find that <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>almost immediately does better than <span><span style="font-variant:small-caps;">LocalUpdate</span></span>and continues to do at least as well throughout the course of training, often better. While this is not true for CIFAR-100, the results are still instructive. We see that while the client learning rate is initially set to a suboptimal value ($\gamma = 1.0$), the automatic learning rate decay enables us to move away from this suboptimal basin and towards something comparable to the best tuned client learning rate after enough rounds. We see that despite initializing with a bad $\gamma$, the non-divergence in earlier rounds is sufficient to allow eventually near-optimal performance in the later rounds.

For both Stack Overflow tasks we plot an analogous results in Figure \[fig:stackoverflow\_decay\]. As discussed in Section \[sec:adaptive\_localupdate\], due to the size of the Stack Overflow dataset, we do not plot the loss $f(x)$ over all clients. Instead, we plot the average loss among all clients in each round before training. For Stack Overflow NWP, we see that <span><span style="font-variant:small-caps;">LocalUpdateDecay</span></span>performs comparably to the best tuned client learning rate. However, for Stack Overflow TP, we see that the decay actually helps significantly. While we initialize with a suboptimal client learning rate $\gamma = 100$, (which clearly achieves higher loss throughout), by decaying the client learning rate we are able to obtain comparable or lower loss than all other fixed $\gamma$.

Open questions
==============

Our work above opens up a number of possible follow-up directions in the area of federated optimization. First, we expect the same kind of analysis obtained above to apply to methods similar to <span><span style="font-variant:small-caps;">LocalUpdate</span></span>. For example, while the FedProx algorithm [@li2019federated] does not fit the format of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>as presented in this work, we believe that it can be analyzed in a similar way. As shown by @pathak2020fedsplit, even in the non-stochastic setting, FedProx is not optimizing the true loss function. Thus, a natural question is to understand exactly what loss function is being optimized, and how the structure of FedProx encourages convergence over FedAvg. Another natural algorithm for analysis is a more general version of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>in which the client learning rate is decayed during a client's local computation. This may help combat adverse effects incurred by setting $K$ to be too large.

More generally, we would also like to understand the behavior of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>on non-quadratic functions. Even generalizing the analysis above to the strongly convex case would be substantial progress towards understanding federated learning in heterogeneous settings. While there may be no surrogate loss that <span><span style="font-variant:small-caps;">LocalUpdate</span></span>is directly optimizing through SGD, we believe that the intuition behind our work can still be carried forward for more general loss functions. In particular, one might expect that the optimization dynamics of <span><span style="font-variant:small-caps;">LocalUpdate</span></span>can be parameterized in terms of $\gamma$ and $\Theta$, even in non-quadratic settings, and that the selection of these parameters governs a trade-off between the speed of convergence and the accuracy of the resulting critical point.

Another interesting open direction is determining how to best set $\Theta$ for a given problem. As seen above, the choice of $\Theta$ drastically alters optimization dynamics. While it is often chosen in an ad hoc manner (based in part on the cost of communication), one could imagine attempting to minimizing the number of rounds need to obtain a given accuracy level with respect to $\Theta$. Even for FedAvg/Reptile, it is not clear how to set the number of local steps $K$. Insights into this could greatly improve the performance of federated learning algorithms.

Finally, we note that our work is fundamentally concerned with the training dynamics of local update methods. In practice, we are often instead interested in the generalization ability of a model. We suspect that the choice of parameters in <span><span style="font-variant:small-caps;">LocalUpdate</span></span>can have large implications for generalization ability, the study of which we leave to future work.

0.2in

[^1]: A different categorization, *vertical* and *horizontal*, was proposed by @yang2019federated, which is based on modelling constraints, rather than on system constraints. The setup in this work applies primarily to horizontal FL, though we expect that much of our framework carries over to the vertical setting.

[^2]: As of May 26, 2020, the Google Play Store reports "1,000,000,000+ installs" for the GBoard application.

[^3]: While we focus on FedAvg and Reptile, we note that a similar analysis can be performed for any of the special cases listed above, using a similar proof strategy.

{'source': 'arXiv_out/2007'}