Communication-Efficient Local SGD with
Age-Based Worker Selection

Consider the PS-based framework of distributed learning with heterogeneous data distribution. There are $M$ distributed workers in set $\mathcal{M}\triangleq\{1,...,M\}$. Each worker $m$ maintains a local dataset $\mathcal{D}_{m}$ of size $N_{m}$. It is drawn from the global dataset $\mathcal{D}=\{z_{i}\}_{i=1}^{N}$, i.e., we have $\mathcal{D}=\bigcup_{m\in\mathcal{M}}\mathcal{D}_{m}$. Our objective is to minimize the weighted loss function

where $\boldsymbol{\theta}\in\mathbb{R}^{d}$ is the $d$-dimension parameter to be optimized, and we have the definition $\mathcal{L}_{m}(\boldsymbol{\theta})\triangleq\mathbb{E}[\mathcal{L}_{m}(\boldsymbol{\theta};z_{m})]$, with $\mathcal{L}_{m}(\boldsymbol{\theta})$ being the local loss function of worker $m$ and $z_{m}$ being the sample drawn randomly from its local dataset $\mathcal{D}_{m}$.

To elaborate on local SGD, we define $\boldsymbol{\theta}^{j}$ as the global model parameter at training round $j$, and $\boldsymbol{\theta}_{m}^{j,0}$ as the local model of worker $m$ before operating local updates. At each training round $j$, a subset $\mathcal{M}_{D}^{j}\subseteq\mathcal{M}$ of workers are selected to download the global model $\boldsymbol{\theta}^{j}$ from the PS prior to computation. After that, each selected worker $m$ in $\mathcal{M}_{D}^{j}$ sets its local model as $\boldsymbol{\theta}_{m}^{j,0}=\boldsymbol{\theta}^{j}$ and starts operating $U$ local iterations, with the updating formula given as

for any local iteration $u=0,...,U-1$. Here $\eta$ is the stepsize and $\frac{1}{B}\sum_{b=1}^{B}\nabla\mathcal{L}(\boldsymbol{\theta}_{m}^{j,u};z_{m,b}^{j,u})$ is the minibatch gradient to be computed by worker $m$ at local iteration $u$, where $B$ is the minibatch size and $z_{m,b}^{j,u}$ is a sample drawn independently across iterations from the local dataset $\mathcal{D}_{m}$.

After all the workers in set $\mathcal{M}_{D}^{j}$ have completed their local computations, a subset of $S$ workers $\mathcal{M}_{U}^{j}\subseteq\mathcal{M}_{D}^{j}$ are selected to upload their latest models, and the PS aggregates the models

Note that here we employ the unbiased aggregation since the weight of each worker is reflected in the worker selection process, as will be specified later.

The training process ends when some stopping criterion is satisfied, with the total number of rounds denoted as $J$.

To gauge the efficiency of the proposed scheme, we are interested in two performance metrics, i.e., the number of training rounds and the communication cost required to reach a targeted training accuracy.

Firstly, the number of training rounds $J$ of the algorithm to achieve a targeted test accuracy is used to reflect the training speed, as well as the communication cost of the algorithm.

Secondly, we define the communication cost $C^{j}$ at training round $j$ as the number of communication rounds between the PS and the workers, given as

This is due to the fact that the size of the global parameter downloaded by the workers is the same as that of the latest parameter uploaded by each selected worker. As a result, the total communication cost is given as $C=\sum_{j=0}^{J-1}C^{j}$.

These two metrics will be used in Section IV to compare the performance of different schemes.

To elaborate on worker selection strategy in the proposed AgeSel, we define an $M$-length vector $\boldsymbol{\tau}_{M}$ to collect the ages of the workers, as in [14, 15]. Note that $\boldsymbol{\tau}_{M}$ is maintained by the PS and initialized to be a zero vector. Each element $\tau_{m}$ of vector $\boldsymbol{\tau}_{M}$ measures the number of consecutive rounds that worker $m$ has not been selected by the PS. Particularly, at each round $j$, $\tau_{m}$ is updated as follows:

To identify the workers with low participation frequency, we pre-define a threshold $\tau_{max}$. Accordingly, worker $m$ would be forced to participate when having its $\tau_{m}\geq\tau_{max}$. The specific procedure of AgeSel is delineated next.

At each round $j$, the PS selects $S$ workers to perform computation by checking the vector $\boldsymbol{\tau}_{M}$. More precisely, with vector $\boldsymbol{\tau}_{M}$, we can identify all the infrequent workers with their ages greater than $\tau_{max}$, which would be considered first. Let $S^{j}$ denote the number of these workers. The set $\mathcal{M}_{D}^{j}$ of selected workers can be determined as follows. For the first case with $S^{j}\geq S$, the PS simply picks the $S$ workers in an age-descending order. Note that the the workers with larger sizes of local datasets are prioritized when there are ties in ages. For the second case with $S^{j}<S$, the PS first picks all the $S^{j}$ infrequent workers and then chooses the rest $S-S^{j}$ without replacement from the workers having ages smaller than $\tau_{max}$ with the probabilities proportional to the sizes of their datasets.

Once the set $\mathcal{M}_{D}^{j}$ is determined, the PS then broadcasts the global parameter $\boldsymbol{\theta}^{j}$ to all the selected workers. By initiating its local model with $\boldsymbol{\theta}_{m}^{j,0}=\boldsymbol{\theta}^{j}$, each worker $m$ in $\mathcal{M}_{D}^{j}$ then starts $U$ iterations of local updates through (2) and sends its latest model $\boldsymbol{\theta}_{m}^{j,U}$ to the PS, i.e., we have $\mathcal{M}_{U}^{j}=\mathcal{M}_{D}^{j}$. Eventually, the PS updates the vector $\boldsymbol{\tau}_{M}$, along with the global model aggregated via (3).

Note that when $\tau_{max}$ is very large, we barely have infrequent workers. Then the set $\mathcal{M}_{D}^{j}$ is indeed chosen by weights (i.e., the sizes of datasets), as with the FedAvg algorithm[6]; i.e., in this case, AgeSel is reduced to FedAvg.

Merits of AgeSel: The benefit from the age-based mechanism used in AgeSel is two-fold. First, it has been shown that less participation of some workers can be detrimental to convergence rate due to the lack of gradient diversity [20], especially for heterogeneous local datasets. Our age-based selection strategy can balance the worker participation, thereby preserving gradient diversity and ensuring the fast convergence. Second, generation of age information $\boldsymbol{\tau}_{M}$ here incurs no extra communication and computation cost, in contrast to other information such as the (costly) norm of updates [15], used in existing alternatives. Hence, AgeSel is more communication- and computation-efficient.

We next establish the convergence of the proposed AgeSel algorithm under heterogeneous data distribution, with a general (not necessarily convex) objective function. Our analysis is based on the following two assumptions, which are widely adopted in related works such as [7, 8].

Assumption 1 (Smoothness and Lower Boundedness) Each local function $\mathcal{L}_{m}(\boldsymbol{\theta})$ is $L$-smooth, i.e.,

$\forall\boldsymbol{\theta}_{1},\boldsymbol{\theta}_{2}\in\mathbb{R}^{d}$. We also assume that the objective function $\mathcal{L}$ is bounded below by $\mathcal{L}^{*}$.

Assumption 2 (Unbiasedness and Bounded Variance) For the given model parameter $\boldsymbol{\theta}$, the local gradient estimator is unbiased, i.e.,

Moreover, both the variance of the local gradient estimator and the variance of the local gradient from the global one are bounded, i.e., there exist two constants $\sigma_{L},\sigma_{G}>0$, such that

With these assumptions, we can derive an upper bound for the expectation of the average squared gradient norm $\frac{1}{J}\mathbb{E}\left[\sum_{j=0}^{J-1}\left\|\nabla\mathcal{L}(\boldsymbol{\theta}^{j})\right\|^{2}\right]$, to prove the convergence of the proposed AgeSel. We start by presenting the following lemma.

Lemma 1: Under Assumptions 1-2 and $\eta$ is chosen such that $\eta\leq\frac{1}{8LU}$, there exists a positive constant $c<\frac{1}{2}-15U^{2}\eta^{2}L^{2}-L\eta(90U^{3}L^{2}\eta^{2}+3U)$, such that

where we have defined $Z_{1}=\frac{5U^{2}\eta^{3}L^{2}}{2}\left(\sigma_{L}^{2}+6U\sigma_{G}^{2}\right)$, $Z_{2}=15U^{3}L^{2}\eta^{2}(\sigma_{L}^{2}+6U\sigma_{G}^{2})+3U^{2}\sigma_{G}^{2}$; and $A^{j}$ is the cardinality of the set $\mathcal{A}^{j}$ of workers selected by age at round $j$, i.e., $A^{j}=\min\{S,S^{j}\}$.

Lemma 1 depicts the one-step difference of the objective function, from which we can see the impact of the age. Particularly, as $A^{j}$ grows larger, the variance term $-\frac{2L\eta^{2}A^{j}}{S}Z_{2}$ is reduced, while the variance term $\frac{2L\eta^{2}(A^{j})^{2}}{S^{2}}Z_{2}$ is increased. Therefore, the values of $\tau_{max}$ and $S$, deciding $A^{j}$ jointly, have a significant impact on training speed, which will be further demonstrated in Section IV. Based on Lemma 1, we are able to arrive at the final convergence result.

Theorem 1: With the same conditions as in Lemma 1, we have

where we have defined the constant $V=\frac{1}{c}\Big{[}\frac{\eta L}{2SB}\sigma_{L}^{2}+\frac{Z_{1}}{\eta U}+(3\eta L-\frac{2L\eta M}{SR})\frac{Z_{2}}{U}\Big{]}$, and $R$ denotes the maximum number of rounds to traverse all the workers in $\mathcal{M}$ due to the age-based mechanism.

Theorem 1 states that the proposed AgeSel scheme can achieve a sublinear convergence rate $\mathcal{O}(\frac{1}{J})$ as local SGD with partial worker participation[7, 8]. The advantage of age-based worker selection is reflected in the expression of $V$, i.e., a smaller $\tau_{max}$ leads to a smaller $R$, and then implies a reduced $V$.

With the $L$-smoothness of the objective function $\mathcal{L}$ in Assumption 1, by taking expectation of all the randomness over $\mathcal{L}(\boldsymbol{\theta}^{j+1})$ we have:

where $(a1)$ is due to the definitions $g^{j}=\frac{1}{S}\sum_{m\in\mathcal{M}_{U}^{j}}g_{m}^{j}$ where $p_{m}=N_{m}/N$ and $g_{m}^{j}=-\eta\sum_{u=0}^{U-1}\frac{1}{B}\sum_{b=1}^{B}\nabla\mathcal{L}(\boldsymbol{\theta}_{m}^{j,u};z_{m,b}^{j,u})$; and $(b1)$ can be obtained directly from (a1).

We then bound the terms $T_{1}$ and $T_{2}$ respectively. For the term $T_{1}$, with the definition of the weighted global update $\bar{g}^{j}=\sum_{m\in\mathcal{M}}p_{m}g_{m}^{j}$ and the unbiased global update $\tilde{g}^{j}=\frac{1}{M}\sum_{m\in\mathcal{M}}g_{m}^{j}$, we have

where $(a2)$ splits the selected workers at round $j$ into the ones selected by ages in set $\mathcal{A}^{j}$ with $|\mathcal{A}^{j}|=A^{j}=\min\{S,S^{j}\}$ and the ones selected by weights in set $\mathcal{M}_{U}^{j}\backslash\mathcal{A}^{j}$; $(b2)$ is because the selection by ages is essentially unweighted random selection in expectation and the selection by weights is equivalent to $\bar{g}^{j}$ in expectation. The two terms in (13) are then bounded separately.

To bound the first term, we have:

where $(a3)$ is derived from the definition of $\bar{g}^{j}$; $(b3)$ and $(c3)$ come from direct computation; $(d3)$ uses the fact that $\langle\mathbf{x},\mathbf{y}\rangle\leq\frac{1}{2}[\|\mathbf{x}\|^{2}+\|\mathbf{y}\|^{2}]$; $(e3)$ is due to Jensen inequality and Cauchy-Schwartz inequality; $(f3)$ follows from Assumption 1; and $(g3)$ is due to the fact that $\sum_{m=1}^{M}p_{m}=1$ and [24, Lemma 3], which proves that

under the condition that $\eta\leq\frac{1}{8LU}$, where $\sigma_{L}$ and $\sigma_{G}$ are two constants defined in Assumption 2.

Likewise, the second term in (13) can be bounded as below, with $p_{m}$ replaced by $\frac{1}{M}$:

Substituting (14) and (16) into (13), we have

With $\mathbb{I}\{\cdot\}$ denoting the indicator function, and $\mathcal{A}\backslash\mathcal{B}$ denotes the complementary set of set $\mathcal{B}$ in set $\mathcal{A}$, the term $T_{2}$ can be bounded as

where $(a4)$ is due to the definition of $g^{j}$; $(b4)$ comes directly from $(a3)$; $(c4)$ follows from the fact that $\mathbb{E}[\|x\|^{2}]=\mathbb{E}[\|x-\mathbb{E}[x]\|^{2}+\|\mathbb{E}[x]\|^{2}]$ and $\mathbb{E}[\nabla\mathcal{L}_{m}(\boldsymbol{\theta}_{m}^{j,u};z_{m,b}^{j,u})]=\nabla\mathcal{L}_{m}(\boldsymbol{\theta}_{m}^{j,u})$; $(d4)$ follows from Cauchy-Schwartz inequality; $(e4)$ is due to the fact that $\mathbb{E}[\|x_{1}+...+x_{n}\|^{2}]=\mathbb{E}[\|x_{1}\|^{2}+...+\|x_{n}\|^{2}]$ if $x_{i}^{\prime}$s are independent with zero mean; $(f4)$ is due to the definition that the subset of workers selected by ages at round $j$ is $\mathcal{A}^{j}$ and $|\mathcal{A}^{j}|=A^{j}=\min\{S,S^{j}\}$; $(g4)$ is due to Cauchy-Schwartz inequality.

Next, we bound the term $\mathbb{E}\left[\left\|\sum_{u=0}^{U-1}\nabla\mathcal{L}_{m}(\boldsymbol{\theta}_{m}^{j,u})\right\|^{2}\right]$ in (18) and (14) as follows:

where $(a5)$ is due to Cauchy-Schwartz inequality and the bounded variance assumption; $(b5)$ follows from (15); $(c5)$ is due to the definition that $C_{1}=90U^{4}L^{2}\eta^{2}+3U^{2}$ and $C_{2}=15U^{3}L^{2}\eta^{2}(\sigma_{L}^{2}+6U\sigma_{G}^{2})+3U^{2}\sigma_{G}^{2}$.

By substituting the upper bounds (17), (18) of the terms $T_{1}$ and $T_{2}$ in (12), we readily have:

where $(a6)$ comes from direct substitution; $(b6)$ uses the result in (19); $(c6)$ follows from direct computation; $(d6)$ uses the fact that $0\leq A^{j}\leq S$; and $(e6)$ follows from the fact that there exists a constant $c$ such that $0<c<\frac{1}{2}-15U^{2}\eta^{2}L^{2}-L\eta(90U^{3}L^{2}\eta^{2}+3U)$. The proof of Lemma 1 is then complete.