Federated averaging, local SGD, and SCAFFOLD

Let $N,d\ge 1$, let each client $i\in \{1,\dots ,N\}$ have a positive integer sample count $n_i$, and put $n={\sum }_{i=1}^N{n}_i$. For differentiable client objectives $F_i:{\mathbb{R}}^d\to \mathbb{R}$, the sample-count-weighted deployment objective is $f_{\mathrm{w}\mathrm{t}}(w)={\sum }_{i=1}^N\frac{n_i}{n}{F}_i(w)\mathrm{.}$ For differentiable client objectives $f_i:{\mathbb{R}}^d\to \mathbb{R}$, the equal-client objective is $f(w)=\frac{1}{N}{\sum }_{i=1}^N{f}_i(w)\mathrm{.}$ The symbols $F_i$ and $f_i$ denote separate families; they are not identified unless an application states $F_i=f_i$. In general, $f_{\mathrm{w}\mathrm{t}}$ and $f$ are different objectives.

For integers $1\le S\le N$, a uniform fixed-size client sample is a random subset $\mathcal{S}\subseteq \{1,\dots ,N\}$ of cardinality $S$ for which every $S$-element subset has probability ${\left(\genfrac{}{}{0px}{}{N}{S}\right)}^{-1}$. This is sampling without replacement. When several communication rounds are used, the sampled subsets are independent across rounds unless a theorem states a different sampling law.

McMahan et al., PMLR 54 (2017), §§1–4

Karimireddy et al., SCAFFOLD, §§2–3

Yuan, Optimization for Deep Learning, Lecture 14: Federated Learning, slides 7–8

Let $n,d,K,T,b,m$ be positive integers, let $f_j:{\mathbb{R}}^d\to \mathbb{R}$ be differentiable for $1\le j\le n$, and define $f(x)=n^{-1}{\sum }_{j=1}^n{f}_j(x)$. The inputs are a starting point $x_0\in {\mathbb{R}}^d$, positive step sizes ${\eta }_0,\dots ,{\eta }_{T-1}$, and integer synchronization times $0<t_1<\cdots <t_m=T,{\mathcal{I}}_T=\{t_1,\dots ,t_m\}\mathrm{.}$ Set $t_0=0$ and $H={\max }_{1\le j\le m}(t_j-t_{j-1})\mathrm{.}$ Initialize $x_0^k=x_0$ for each worker $k\in \{1,\dots ,K\}$. For each $t=0,\dots ,T-1$, each worker $k$ and each within-step draw $q=1,\dots ,b$, draw $i_{t,q}^k$ uniformly from $\{1,\dots ,n\}$, with all draws mutually independent, and form $g_t^k(x)=\frac{1}{b}{\sum }_{q=1}^b\mathrm{\nabla }{f}_{i_{t,q}^k}(x),z_{t+1}^k=x_t^k-{\eta }_t{g}_t^k(x_t^k)\mathrm{.}$ For $b=1$ this is a single stochastic gradient. If $t+1\mathrm{\notin }{\mathcal{I}}_T$, set $x_{t+1}^k=z_{t+1}^k$ for every worker. If $t+1\in {\mathcal{I}}_T$, set $x_{t+1}^k=\frac{1}{K}{\sum }_{\mathrm{\ell }=1}^K{z}_{t+1}^{\mathrm{\ell }}\text{for every }k\mathrm{.}$ The procedure returns the worker trajectories $(x_t^k{)}_{0\le t\le T,1\le k\le K}$. The common initialization at time $0$ is included when measuring the largest synchronization gap, and a convergence statement must specify any particular average or iterate it analyzes.

Independent random elements

Stich, Local SGD Converges Fast and Communicates Little, Algorithm 1 and Definition 2.1

Gorbunov et al., Local SGD: Unified Theory and New Efficient Methods, equation (4)

For vectors $u,v\in {\mathbb{R}}^d$ and $a>0$, $\mathrm{\parallel }u+v{\mathrm{\parallel }}_2^2\le (1+a)\mathrm{\parallel }u{\mathrm{\parallel }}_2^2+(1+a^{-1})\mathrm{\parallel }v{\mathrm{\parallel }}_2^2,2\mathrm{\mid }⟨u,v⟩\mathrm{\mid }\le a\mathrm{\parallel }u{\mathrm{\parallel }}_2^2+a^{-1}\mathrm{\parallel }v{\mathrm{\parallel }}_2^2\mathrm{.}$ For $m\ge 1$ vectors $u_1,\dots ,u_m\in {\mathbb{R}}^d$, ${\parallel {\sum }_{j=1}^m{u}_j\parallel }_2^2\le m{\sum }_{j=1}^m\mathrm{\parallel }{u}_j{\mathrm{\parallel }}_2^2\mathrm{.}$

Assume the Axiom of Choice. Let ${\mathcal{F}}_0\subseteq {\mathcal{F}}_1\subseteq \cdots \subseteq {\mathcal{F}}_m$ be a finite filtration and let $Z_j\in L^2(\mathrm{\Omega };{\mathbb{R}}^d)$ be ${\mathcal{F}}_j$-measurable for $1\le j\le m$, where $m\ge 1$. For each $0\le \mathrm{\ell }<d$, define $M_j^{\mathrm{\ell }}=\mathbb{E}[Z_j^{\mathrm{\ell }}\mid {\mathcal{F}}_{j-1}]$ by scalar conditional expectation and put $M_j=(M_j^{\mathrm{\ell }}{)}_{0\le \mathrm{\ell }<d}$ and ${\xi }_j=Z_j-M_j$. Assume ${\xi }_j\in L^2(\mathrm{\Omega };{\mathbb{R}}^d)$. If $\mathbb{E}[\mathrm{\parallel }{\xi }_j{\mathrm{\parallel }}_2^2\mid {\mathcal{F}}_{j-1}]\le v_j\text{almost surely},$ where each $v_j\ge 0$ is integrable and ${\mathcal{F}}_{j-1}$-measurable, then $\mathbb{E}{\parallel {\sum }_{j=1}^m{Z}_j\parallel }_2^2\le 2\mathbb{E}{\parallel {\sum }_{j=1}^m{M}_j\parallel }_2^2+2{\sum }_{j=1}^m\mathbb{E}[v_j]\mathrm{.}$

Karimireddy et al., SCAFFOLD, Appendix B, Lemma 4

Let $\mu >0$, $a\ge 4$, $T\ge 1$, and ${\eta }_t=\frac{4}{\mu (a+t)},w_t=(a+t{)}^2,S_T={\sum }_{t=0}^{T-1}{w}_t\mathrm{.}$ Suppose nonnegative sequences $u_t,e_t$ satisfy, for $0\le t<T$, $u_{t+1}\le (1-\mu {\eta }_t)u_t-A{\eta }_t{e}_t+B{\eta }_t^2+C{\eta }_t^3$ with $A>0$ and $B,C\ge 0$. Then $\frac{A}{S_T}{\sum }_{t=0}^{T-1}{w}_t{e}_t\le \frac{\mu a^3}{4S_T}{u}_0+\frac{2T(T+2a)}{\mu S_T}B+\frac{16T}{{\mu }^2{S}_T}C\mathrm{.}$ Also $S_T\ge T^3\mathrm{/}3$.

Let $R\ge 1$ and let $d_0,\dots ,d_R\ge 0$, $q_0,\dots ,q_{R-1}\ge 0$.

1. If $a,\eta >0$, $b,c\ge 0$, and $d_r+a\eta q_{r-1}\le d_{r-1}+b{\eta }^2+c{\eta }^3(1\le r\le R),$ then $\frac{1}{R}{\sum }_{r=1}^R{q}_{r-1}\le \frac{d_0}{a\eta R}+\frac{b\eta +c{\eta }^2}{a}\mathrm{.}$

2. If $0<\rho \eta <1$, $a>0$, $b,c\ge 0$, and $d_r\le (1-\rho \eta )d_{r-1}-a\eta q_{r-1}+b{\eta }^2+c{\eta }^3,$ then for $w_r=(1-\rho \eta {)}^{-r}$ and $W_R={\sum }_{r=1}^R{w}_r$, $\frac{1}{W_R}{\sum }_{r=1}^R{w}_r{q}_{r-1}\le \frac{d_0}{a\eta W_R}+\frac{b\eta +c{\eta }^2}{a}\mathrm{.}$ The left side in part 2 is the expectation of $q_{r-1}$ when an index $r\in \{1,\dots ,R\}$ is selected with probability $w_r\mathrm{/}{W}_R$.

Let $h_1,\dots ,h_N:{\mathbb{R}}^d\to \mathbb{R}$ be convex and continuously differentiable, with $\beta$-Lipschitz gradients for $\beta >0$. Put $f=N^{-1}{\sum }_{i=1}^N{h}_i$, and let $x^{\star }$ minimize $f$. Then $\frac{1}{N}{\sum }_{i=1}^N\mathrm{\parallel }\mathrm{\nabla }{h}_i(x)-\mathrm{\nabla }{h}_i(x^{\star }){\mathrm{\parallel }}_2^2\le 2\beta (f(x)-f(x^{\star }))(x\in {\mathbb{R}}^d)\mathrm{.}$

Karimireddy et al., SCAFFOLD, Appendix B, inequality (9)

Let $h:{\mathbb{R}}^d\to \mathbb{R}$ be convex and differentiable with $\beta$-Lipschitz gradient, where $\beta >0$. For every $0\le \eta \le 2\mathrm{/}\beta$ and $x,y\in {\mathbb{R}}^d$, $\mathrm{\parallel }x-y-\eta (\mathrm{\nabla }h(x)-\mathrm{\nabla }h(y)){\mathrm{\parallel }}_2\le \mathrm{\parallel }x-y{\mathrm{\parallel }}_2\mathrm{.}$

Karimireddy et al., SCAFFOLD, Appendix C, Lemma 6 proves a related contraction for $\mu$-strongly convex functions with $\eta \le 1\mathrm{/}\beta$ and squared-distance factor $1-\mu \eta$. The convex-only nonexpansiveness for the full range $0\le \eta \le 2\mathrm{/}\beta$ is established by the local proof below.

The inputs are positive integers $N,d,R,K,S$ with $1\le S\le N$, an initial server model $x_0\in {\mathbb{R}}^d$, client objectives for the chosen mode and their local oracle procedures $g_i$, positive step sizes ${\eta }_l,{\eta }_g$, and one aggregation mode from the two modes below. In count-weighted deployment mode the objectives are $F_i$ and the inputs also include positive integer record counts $n_1,\dots ,n_N$, with $n={\sum }_i{n}_i$. A call to $g_i$ returns a gradient vector according to its supplied deterministic or stochastic local-oracle rule; the client-sample randomness in a round is independent of that round's fresh oracle randomness. In uniform-client analysis mode the objectives are $f_i$.

For each round $r=1,\dots ,R$, independently draw a uniform fixed-size subset ${\mathcal{S}}_r$ of $S$ clients from $\{1,\dots ,N\}$. For every selected client $i$, initialize $y_{i,0}=x_{r-1}$ and for $k=1,\dots ,K$ call $g_i$ at $y_{i,k-1}$ and update $y_{i,k}=y_{i,k-1}-{\eta }_l{g}_i(y_{i,k-1}),1\le k\le K,$ The client sends its model difference $y_{i,K}-x_{r-1}$ to the server. Use one of the following aggregation branches in every round:

1. Count-weighted deployment mode. With positive client record counts $n_i$ and $n={\sum }_i{n}_i$, set $x_r=x_{r-1}+{\eta }_g{\sum }_{i\in {\mathcal{S}}_r}\frac{n_i}{n}(y_{i,K}-x_{r-1})\mathrm{.}$ Equivalently, each unselected client contributes its unchanged model $x_{r-1}$ to the full count-weighted average.

2. Uniform-client analysis mode. Set $x_r=x_{r-1}+\frac{{\eta }_g}{S}{\sum }_{i\in {\mathcal{S}}_r}(y_{i,K}-x_{r-1})\mathrm{.}$ The equal-client convergence result on this page uses this mode and the objective $f=N^{-1}{\sum }_i{f}_i$.

After the selected clients finish the round, apply the chosen branch to obtain $x_r$. Repeat through round $R$. Return the finite server trajectory $x_0,\dots ,x_R$.

Federated objectives and uniform client sampling

McMahan et al., PMLR 54 (2017), §2 and Algorithm 1

Karimireddy et al., SCAFFOLD, Algorithm 2

Chi and Li, Advances in Federated Optimization, slides 29–31

Yuan, Optimization for Deep Learning, Lecture 14: Federated Learning, slides 7–16 and 28–30

The count-weighted deployment mode and uniform-client analysis mode are different update rules when the $n_i$ are unequal. The strongly convex rate proved on this page applies only to the uniform-client mode and the equal-client objective. McMahan et al. give the algorithm and empirical experiments; those experiments do not prove a general convergence rate or a differential-privacy guarantee.

McMahan et al., PMLR 54 (2017), §§2–4

The inputs are positive integers $N,d,R,K,S$ with $1\le S\le N$, objective functions $f_i:{\mathbb{R}}^d\to \mathbb{R}$, an initial server model $x_0\in {\mathbb{R}}^d$, client controls $c_i^0\in {\mathbb{R}}^d$, positive step sizes ${\eta }_l,{\eta }_g$, a local stochastic-gradient oracle $g_i$, and a control-update option from $\{\mathrm{I},\mathrm{I}\mathrm{I}\}$. Set $c^0=N^{-1}{\sum }_i{c}_i^0$. The random client sample in each round is independent of that round's fresh local-oracle randomness; no independence among distinct clients' oracle draws is required here.

For each round $r=1,\dots ,R$, independently sample a uniform fixed-size subset ${\mathcal{S}}_r$ of $S$ clients. At the start of this round, write $x=x_{r-1}$, $c=c^{r-1}$, and $c_i=c_i^{r-1}$. For every selected client $i$, initialize $y_{i,0}=x$ and, for $k=1,\dots ,K$, call $g_i$ at $y_{i,k-1}$ and update $y_{i,k}=y_{i,k-1}-{\eta }_l(g_i(y_{i,k-1})-c_i+c)\mathrm{.}$ After its $K$ local steps, a selected client updates its control using either

• Option I: call $g_i$ at the round-start model $x_{r-1}$ and set $c_i^{+}=g_i(x)$; or
• Option II: set $c_i^{+}=c_i-c+(K{\eta }_l{)}^{-1}(x-y_{i,K})$.

The client sends $\mathrm{\Delta }{x}_i=y_{i,K}-x$ and $\mathrm{\Delta }{c}_i=c_i^{+}-c_i$. Unselected clients retain their controls. The server forms $\mathrm{\Delta }x=\frac{1}{S}{\sum }_{i\in {\mathcal{S}}_r}\mathrm{\Delta }{x}_i,\mathrm{\Delta }c=\frac{1}{S}{\sum }_{i\in {\mathcal{S}}_r}\mathrm{\Delta }{c}_i$ and updates $x_r=x+{\eta }_g\mathrm{\Delta }x,c^r=c^{r-1}+\frac{S}{N}\mathrm{\Delta }c\mathrm{.}$ The selected clients set $c_i^r=c_i^{+}$ and unselected clients set $c_i^r=c_i^{r-1}$. Thus $c^r=N^{-1}{\sum }_i{c}_i^r$ whenever $c^{r-1}=N^{-1}{\sum }_i{c}_i^{r-1}$ held before the round. The general procedure permits zero initial controls; the convergence rates stated below use the explicitly stated gradient-consistent initialization $c_i^0=\mathrm{\nabla }{f}_i(x_0)$. Repeat through round $R$ and return the model and control trajectories $((x_r,c^r,(c_i^r{)}_{i=1}^N){)}_{r=0}^R$.

The following schematic records only the communication flow. Its vertices are the server and the selected clients; the arrow labels are exactly the quantities sent in each direction. The local updates and server equations above, not the schematic, define the procedure.

diagram failed to render:
\begin{tikzpicture}[>=stealth]
\node[draw,rounded corners] (server) at (0,0) {Server: $(x,c)$};
\node[draw,rounded corners] (clients) at (5,0) {Selected clients $i\in\mathcal S$};
\draw[->] (server) to[bend left=18] node[above] {$(x,c)$} (clients);
\draw[->] (clients) to[bend left=18] node[below] {$(\Delta x_i,\Delta c_i)$} (server);
\end{tikzpicture}

Equal-client objective and fixed-size client sampling

Karimireddy et al., SCAFFOLD, §§4–5 and Algorithm 1

Gorbunov et al., Local SGD: Unified Theory and New Efficient Methods, §4.2

Chi and Li, Advances in Federated Optimization, slides 34–35

Yuan, Optimization for Deep Learning, Lecture 14: Federated Learning, slides 24–30

Theorem III's rates below do not assume bounded gradient dissimilarity. The separate Hessian-dissimilarity improvement in Theorem IV is proved only for quadratic client functions and is not asserted here for general smooth losses. The procedure transmits model and control updates but supplies no differential-privacy guarantee.

Karimireddy et al., SCAFFOLD, Theorems III–IV and §4

Let $v_1,\dots ,v_N\in {\mathbb{R}}^d$, let $1\le S\le N$, and choose $\mathcal{S}$ uniformly from the $S$-element subsets of $\{1,\dots ,N\}$. Put $\bar{v}=N^{-1}{\sum }_i{v}_i$ and ${\bar{v}}_{\mathcal{S}}=S^{-1}{\sum }_{i\in \mathcal{S}}{v}_i$. If $N>1$, then $\mathbb{E}[{\bar{v}}_{\mathcal{S}}]=\bar{v},\mathbb{E}\mathrm{\parallel }{\bar{v}}_{\mathcal{S}}-\bar{v}{\mathrm{\parallel }}_2^2=(1-\frac{S}{N})\frac{1}{S(N-1)}{\sum }_{i=1}^N\mathrm{\parallel }{v}_i-\bar{v}{\mathrm{\parallel }}_2^2\mathrm{.}$ In particular, $\mathbb{E}\mathrm{\parallel }{\bar{v}}_{\mathcal{S}}-\bar{v}{\mathrm{\parallel }}_2^2\le 2(1-\frac{S}{N})\frac{1}{S}\frac{1}{N}{\sum }_{i=1}^N\mathrm{\parallel }{v}_i{\mathrm{\parallel }}_2^2\mathrm{.}$ If $N=1$, necessarily $S=1$ and ${\bar{v}}_{\mathcal{S}}=\bar{v}$.

Let $f=n^{-1}{\sum }_{j=1}^n{f}_j$ have an attained minimum $f^{\star }=f(x^{\star })$, be $L$-smooth, and be $\mu$-strongly convex with $0<\mu \le L$. Let $x_t^1,\dots ,x_t^K$ be generated by local SGD, and define ${\bar{x}}_t=\frac{1}{K}{\sum }_{k=1}^K{x}_t^k,g_t=\frac{1}{K}{\sum }_{k=1}^K{g}_t^k(x_t^k),{\bar{g}}_t=\frac{1}{K}{\sum }_{k=1}^K\mathrm{\nabla }f(x_t^k),$ where $g_t^k$ is the batch gradient from the algorithm definition. If ${\eta }_t\le 1\mathrm{/}(4L)$, then $\begin{array}{rl}{\displaystyle \mathbb{E}\mathrm{\parallel }{\bar{x}}_{t+1}-x^{\star }{\mathrm{\parallel }}_2^2}& {\displaystyle \le (1-\mu {\eta }_t)\mathbb{E}\mathrm{\parallel }{\bar{x}}_t-x^{\star }{\mathrm{\parallel }}_2^2+{\eta }_t^2\mathbb{E}\mathrm{\parallel }{g}_t-{\bar{g}}_t{\mathrm{\parallel }}_2^2}\\ {\displaystyle }& {\displaystyle -\frac{{\eta }_t}{2}\mathbb{E}[f({\bar{x}}_t)-f^{\star }]+\frac{2{\eta }_{t}L}{K}\sum _{k=1}^K\mathbb{E}\mathrm{\parallel }{\bar{x}}_t-x_t^k{\mathrm{\parallel }}_2^2\mathrm{.}}\end{array}$ Expectations are over the finite product law of the sampled indices.

Stich, Local SGD Converges Fast and Communicates Little, Lemma 3.1 and Appendix A