# Federated Learning under Importance Sampling

Elsa Rizk, Stefan Vlaski, and Ali H. Sayed The authors are with the School of
Engineering, École Polytechnique Fédérale de Lausanne, 1015 Lausanne,
Switzerland (e-mail<EMAIL_ADDRESS><EMAIL_ADDRESS>ali.sayed@epfl.ch). A short conference article dealing with an earlier version
of this work without extended arguments and proofs appears in [1].
This work has been submitted for review.

###### Abstract

Federated learning encapsulates distributed learning strategies that are
managed by a central unit. Since it relies on using a selected number of
agents at each iteration, and since each agent, in turn, taps into its local
data, it is only natural to study optimal sampling policies for selecting
agents and their data in federated learning implementations. Usually, only
uniform sampling schemes are used. However, in this work, we examine the
effect of importance sampling and devise schemes for sampling agents and data
non-uniformly guided by a performance measure. We find that in schemes
involving sampling without replacement, the performance of the resulting
architecture is controlled by two factors related to data variability at each
agent, and model variability across agents. We illustrate the theoretical
findings with experiments on simulated and real data and show the improvement
in performance that results from the proposed strategies.

###### Index Terms:

federated learning, distributed learning, importance sampling, asynchronous
SGD, non-IID data, heterogeneous agents

##  I Introduction

In this work, we focus on algorithms that fall into the broad class of
stochastic gradient descent (SGD). We consider a collection of $K$
heterogeneous agents that may have different computational powers. Each agent
$k$ has locally $N_{k}$ data points, which we denote by $\\{x_{k,n}\\}$; the
subscript $k$ refers to the agent, while the subscript $n$ denotes the sample
index within agent $k$’s dataset. The goal of the agents is to find an
optimizer for the aggregate risk function:

$w^{o}\>\overset{\Delta}{=}\>\operatorname*{argmin}_{w\in\mathbb{R}^{M}}\frac{1}{K}\sum_{k=1}^{K}P_{k}(w),$
(1)

where each $P_{k}(\cdot)$ is an empirical risk defined in terms of a loss
function $Q_{k}(\cdot)$:

$P_{k}(w)\>\overset{\Delta}{=}\>\frac{1}{N_{k}}\sum_{n=1}^{N_{k}}Q_{k}(w;x_{k,n}).$
(2)

Multiple strategies exist for solving such problems. They can be categorized
into two main classes: a) partially decentralized strategies, which include a
central process with access to all data and which controls the distribution of
the data into the nodes for processing [2, 3, 4]; and b) fully decentralized
strategies, which consist of multiple agents connected by a graph topology and
operating locally without oversight by a central processor [5, 6, 7, 8].
Federated learning [9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21] offers
a midterm solution, which consists of several agents collecting and processing
local data that are then aggregated at the central processor.

When implementing SGD, most strategies choose the samples according to a
uniform distribution. In this work, we shall consider more general non-uniform
sampling schemes, where the agents are sampled according to some distribution
$\pi_{k}$ and the local data at agent $k$ are in turn sampled according to
some other distribution $\pi_{n}^{(k)}$. In this setting, the central
processor selects the subset of agents for processing according to $\pi_{k}$
and, once selected, an agent $k$ will sample its data according to
$\pi_{n}^{(k)}$. The importance sampling process in this article therefore
involves two layers. We use the superscript $(k)$ to denote the sampling
distribution of the data at agent $k$. The sampling distributions
$\\{\pi_{k},\pi_{n}^{(k)}\\}$ are not fixed; instead we will show how to adapt
them in order to enhance performance. At the same time, we will provide a
detailed convergence analysis and establish performance limits.

###  I-A Related Work

Several works studied the convergence of the federated learning algorithm or
distributed SGD under differing assumptions. These assumptions usually relate
to the nature of the data (IID or non-IID), nature of the cost function
(convex or non-convex), agent participation (full or partial), and operation
(synchronous or asynchronous) [22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 10, 11,
32]. Other works examine the convergence behavior of variations of the
traditional FedAvg algorithm, such as FedProx [10], hierarchical version of
FedAvg [32], multi-task federated learning [11], and dynamic FedAvg [1] – see
Table 1.

TABLE I: List of references on the convergence analysis of federated learning under different assumptions. This work along with our previous work [1] are the only ones to tackle the 3 challenges of federated learning(non-IID data, asynchronous mode of operation, partial agent participation). References | Algorithm | Function Type | Data Heterogeneity | Operation | Agent Participation | Other Assumptions
---|---|---|---|---|---|---
[23] | dist. gradient descent | convex | non-IID | synchronous | full | smooth
[24] | dist. SGD | convex | IID | synchronous | full | smooth
[25, 26] | dist. SGD | non-convex | IID | synchronous | full | smooth
[27] | dist. SGD | non-convex |  | non-IID
---
IID
| synchronous
---
asynchronous
full | -
[28] | dist. SGD | convex | non-IID | synchronous | full | bounded gradients
[29] | dist. momentum SGD | non-convex | non-IID | synchronous | full | -
[30] | FedAvg |  | convex
---
some non-convex
non-IID | asynchronous | full | -
[31] | FedAvg | convex | non-IID | synchronous | partial | bounded gradients
[10] | FedProx | non-convex | non-IID | asynchronous | partial | -
[32] | HierFAVG |  | convex
---
non-convex
non-IID | synchronous | full | -
[11] | MOCHA | convex | non-IID | synchronous | full | -
[1] | Dynamic FedAvg | convex | non-IID | asynchronous | partial | model drift
_this work_ | ISFedAvg | convex | non-IID | asynchronous | partial | importance sampling

By contrast, not much work has been done on selection schemes for agents and
data in federated learning. Given the architecture of a federated learning
solution, this is a natural and important question to consider. The existing
works in this domain can be split into two categories: those seeking better
accuracy, and those seeking fairness. Of the works pertaining to the first
category, reference [33] develops a new client selection scheme, called FedCS,
where the goal of the central server is to choose as many agents as possible
that can complete an iteration by a required deadline, after acquiring
information about the agents’ resources. Reference [34] builds on this
previous work to deal with non-IID data, and allows the server to collect some
of the data from the agents and participate in the training of the model. The
authors of [35] consider non-uniform sampling of agents and suggest
approximate sampling probabilities that maximize the average inner product of
the local gradient with the global gradient. References [36, 37] fall under
the second category; in agnostic federated learning [36], the data
distribution is assumed to be a mixture of the local distributions, and a
minimax problem for agent selection is solved. Reference [37] generalizes the
previous work by reweighting the cost function and assigning higher weights to
agents with higher loss.

While there exist works that study the effect of importance sampling in
distributed learning [38, 39, 40, 41, 42], all of these works apply importance
sampling to the data at each agent. To our knowledge, there are no works that
examine the combined effect of two hierarchical layers of sampling: one for
the nodes and another for their data. By introducing a two-layer importance
sampling scheme to the federated learning paradigm, we can tackle the problem
of importance sampling both in relation to agents and also in relation to
data.

###  I-B Sampling and Inclusion Probabilities

Before describing the problem setting, we need to clarify the difference
between two notions: (a) sampling probability and (b) inclusion probability.
Consider the following illustrative example. Consider $N=4$ balls of which we
wish to choose $B=2$ balls non-uniformly and without replacement. Let the
sampling probabilities be $\pi_{n}=\\{1/3,1/6,1/3,1/6\\}$. This means that,
initially, balls 1 and 3 are twice as likely to be selected compared to balls
2 and 4. For the first trial, all the inclusion probabilities are equal to the
sampling probabilities, i.e., $\mathbb{P}(n\text{ chosen on $1^{st}$
trial})=\pi_{n}$. However, since we are sampling without replacement, the
inclusion probabilities for the second trial depend on the outcome of the
first trial, i.e., $\mathbb{P}(n\text{ chosen on $2^{nd}$ trial}|m$ chosen on
$1^{st}$ trial$)=\pi_{n}/(1-\pi_{m})$. Using the sampling probabilities, we
can evaluate the likelihood that each ball will end up belonging to the
selected set of $2$ balls. In particular, the probability that ball $1$ is
chosen either in the first or second trial is given by:

$\displaystyle\mathbb{P}(\text{1 chosen})$
$\displaystyle=\sum_{n=2}^{4}\mathbb{P}\big{(}\text{$1$ chosen on $1^{st}$
trial \& $n$ chosen on $2^{nd}$ trial)}$
$\displaystyle\quad+\mathbb{P}\big{(}n\text{ chosen on $1^{st}$ trial \& $1$
chosen on $2^{nd}$ trial})$
$\displaystyle=\sum_{n=2}^{4}\pi_{1}\frac{\pi_{n}}{1-\pi_{1}}+\pi_{n}\frac{\pi_{1}}{1-\pi_{n}}.$
(3)

Thus, the sampling probability is the working probability. It is the
probability used to actually choose the samples, while the inclusion
probability is a descriptive probability that indicates the likelihood of a
ball being included in the final selected subset. Observe that the inclusion
probabilities depend on the sampling scheme, while the sampling probabilities
do not. When considering uniform sampling without replacement, the inclusion
probability is a multiple of the sampling probability. For example, sampling
$B$ numbers from $\\{1,2,\cdots,N\\}$ with sampling probabilities $1/N$, the
inclusion probability is found to be $\mathbb{P}(n\in\mathcal{B})=B/N$. Note
further that while the sampling probabilities sum to 1 over all the sampling
space, the inclusion probabilities $\mathbb{P}(n\in\mathcal{B})$ sum to $B$.
In our derivations, we will be relying frequently on the inclusion
probabilities.

Next, we consider a total number of $K$ agents. At each iteration $i$ of the
algorithm, a subset of agents $\mathcal{L}_{i}$ of size $L$ is chosen randomly
without replacement. We denote the probability that agent $k$ is included in
the sample by $Lp_{k}$ [43], i.e.,

$p_{k}\>\overset{\Delta}{=}\>\frac{\mathbb{P}\big{(}k\in\mathcal{L}_{i}\big{)}}{L}.$
(4)

In addition, each sampled agent $k$ will run a mini-batch SGD by sampling
$B_{k}$ data points $\mathcal{B}_{k,i}$ without replacement from its local
data. We denote the probability of inclusion of data point $n$ by
$B_{k}p_{n}^{(k)}$, i.e.,

$p_{n}^{(k)}\>\overset{\Delta}{=}\>\frac{\mathbb{P}\big{(}n\in\mathcal{B}_{k,i}\big{)}}{B_{k}}.$
(5)

We refer to $p_{k}$ and $p_{n}^{(k)}$ as the normalized inclusion
probabilities. They sum to $1$ over the sampling space; $p_{k}$ sums to $1$
over all agents and $p_{n}^{(k)}$ over the data at each agent.

##  II Algorithm Derivation

The goal of the federated learning algorithm is to approximate the centralized
solution $w^{o}$ while dealing with the constraint of distributed data. The
goal is achieved by using an unbiased estimate of the gradient of the cost
function, $\frac{1}{K}\sum_{k=1}^{K}\nabla_{w^{{\sf T}}}P_{k}(w).$ As
explained in [1] for the case of uniform sampling, if we assume each agent $k$
runs $E_{k}$ epochs per iteration $i$ (with each epoch using $B_{k}$ samples
in $\mathcal{B}_{k,i,e}$), then we can construct an unbiased estimate for the
true gradient by considering the following estimator:

$\frac{1}{L}\sum_{k\in\mathcal{L}_{i}}\frac{1}{E_{k}B_{k}}\sum_{e=1}^{E_{k}}\sum_{b\in\mathcal{B}_{k,i,e}}\nabla_{w^{{\sf
T}}}Q_{k}(w;\bm{x}_{k,b}),$ (6)

as opposed to the original estimator from [9], where the main difference is
the scaling by the epoch size $E_{k}$. This correction is important for the
performance of the averaged model. Since the number of epochs $E_{k}$ can be
non-uniform across the agents, then, without correction, agents with large
epoch sizes will bias the solution by driving it towards their local model and
away from $w^{o}$.

Expression (6) is still not sufficient for our purposes in this article, since
agents and data are allowed to be sampled non-uniformly without replacement.
In this case, we need to adjust (6) by including the inclusion probabilities
[38]. The inclusion probabilities are necessary to ensure the estimate is
unbiased, as will later be seen in Lemma 1. The local estimate of the gradient
at agent $k$ becomes $\frac{1}{Kp_{k}}\widehat{\nabla_{w^{{\sf T}}}P_{k}}(w)$,
with:

$\widehat{\nabla_{w^{{\sf
T}}}P_{k}}(w)\>\overset{\Delta}{=}\>\frac{1}{E_{k}B_{k}}\sum_{e=1}^{E_{k}}\sum_{b\in\mathcal{B}_{k,i,e}}\frac{1}{N_{k}p_{b}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w;\bm{x}_{k,b}).$ (7)

Motivated by (7), we can write down a stochastic gradient update at each agent
$k$ at epoch $e$, and at the central processor at iteration $i$:

$\displaystyle\bm{w}_{k,e}=$ $\displaystyle\>\bm{w}_{k,e-1}$
$\displaystyle-\frac{\mu}{Kp_{k}E_{k}B_{k}}\sum_{b\in\mathcal{B}_{k,i,e}}\frac{1}{N_{k}p_{b}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,b}),$ (8) $\displaystyle\bm{w}_{i}=$
$\displaystyle\frac{1}{L}\sum_{k\in\mathcal{L}_{i}}\bm{w}_{k,E_{k}},$ (9)

where at each iteration $i$, step (8) is repeated for $e=1,2,\cdots,E_{k}$. We
arrive at the Algorithm 1, which we refer to as Importance Sampling Federated
Averaging (ISFedAvg).

initialize $w_{0}$

for each iteration $i=1,2,\cdots$ do

Select the set of participating agents $\mathcal{L}_{i}$ by sampling $L$ times
from $\\{1,\ldots,K\\}$ without replacement according to the sampling
probabilities $\pi_{k}$.

for each agent $k\in\mathcal{L}_{i}$ do

initialize $\bm{w}_{k,0}=\bm{w}_{i-1}$

for each epoch $e=1,2,\cdots E_{k}$ do

Find indices of the mini-batch sample $\mathcal{B}_{k,i,e}$ by sampling
$B_{k}$ times from $\\{1,\ldots,N_{k}\\}$ without replacement according to the
sampling probabilities
$\pi_{n}^{(k)}$.$\bm{g}=\dfrac{1}{B_{k}}\sum\limits_{b\in\mathcal{B}_{k,i,e}}\dfrac{1}{N_{k}p_{b}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,b})$
$\bm{w}_{{k},e}=\bm{w}_{{k},e-1}-\mu\dfrac{1}{E_{k}Kp_{k}}\bm{g}$

end for

end
for$\bm{w}_{i}=\dfrac{1}{L}\sum\limits_{k\in\mathcal{L}_{i}}\bm{w}_{k,E_{k}}$

end for

Algorithm 1 (Importance Sampling Federated Averaging)

##  III Convergence Analysis

###  III-A Modeling Conditions

To facilitate the analysis of the algorithm, we list some common assumptions
on the nature of the local risk functions and their respective minimizers.
Specifically, we assume convex cost functions with smooth gradients.

######  Assumption 1.

The functions $P_{k}(\cdot)$ are $\nu-$strongly convex, and
$Q_{k}(\cdot;x_{k,n})$ are convex, namely:

$\displaystyle P_{k}(w_{2})\geq P_{k}(w_{1})+\nabla_{w^{{\sf
T}}}P_{k}(w_{1})(w_{2}-w_{1})+\frac{\nu}{2}\|w_{2}-w_{1}\|^{2},$ (10)
$\displaystyle Q_{k}(w_{2};x_{k,n})\geq Q_{k}(w_{1};x_{k,n})+\nabla_{w^{{\sf
T}}}Q_{k}(w_{1};x_{k,n})(w_{2}-w_{1}).$ (11)

Also, the functions $Q_{k}(\cdot;x_{k,n})$ have $\delta-$Lipschitz gradients:

$\displaystyle\|\nabla_{w^{{\sf T}}}Q_{k}(w_{2};x_{k,n})-\nabla_{w^{{\sf
T}}}Q_{k}(w_{1};x_{k,n})\|$ $\displaystyle\leq\delta\|w_{2}-w_{1}\|.$ (12)

∎

We further assume that the individual minimizers
$w_{k}^{o}=\operatorname*{argmin}_{w\in\mathbb{R}^{M}}P_{k}(w)$ do not drift
too far away from $w^{o}$.

######  Assumption 2.

The distance of each local model $w^{o}_{k}$ to the global model $w^{o}$ is
uniformly bounded, $\|w^{o}_{k}-w^{o}\|\leq\xi.$ ∎

###  III-B Error Recursion

Iterating the local update (8) over multiple epochs and combining according to
(9), we obtain the following update for the central iterate:

$\displaystyle\bm{w}_{i}=\>\bm{w}_{i-1}-\mu\frac{1}{L}\sum_{k\in\mathcal{L}_{i}}$
$\displaystyle\frac{1}{Kp_{k}E_{k}B_{k}}\sum_{e=1}^{E_{k}}\sum_{b\in\mathcal{B}_{k,i,e}}\frac{1}{N_{k}p_{b}^{(k)}}$
$\displaystyle\times\nabla_{w^{{\sf T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,b}).$
(13)

To simplify the notation, we introduce the error terms:

$\displaystyle\bm{s}_{i}$
$\displaystyle\>\overset{\Delta}{=}\>\frac{1}{L}\sum_{\ell\in\mathcal{L}_{i}}\frac{1}{Kp_{\ell}}\widehat{\nabla_{w^{{\sf
T}}}P_{\ell}}(\bm{w}_{i-1})-\frac{1}{K}\sum_{k=1}^{K}\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1}),$ (14) $\displaystyle\bm{q}_{i}$
$\displaystyle\>\overset{\Delta}{=}\>\frac{1}{L}\sum_{\ell\in\mathcal{L}_{i}}\frac{1}{Kp_{\ell}E_{\ell}B_{\ell}}\sum_{e=1}^{E_{\ell}}\sum_{b\in\mathcal{B}_{\ell,i,e}}\frac{1}{N_{\ell}p_{b}^{(\ell)}}$
$\displaystyle\qquad\qquad\times\big{(}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{\ell,e-1};\bm{x}_{\ell,b})-\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};\bm{x}_{\ell,b})\big{)}.$ (15)

The first error term $\bm{s}_{i}$, which we call gradient error, captures the
error from approximating the true gradient by using subsets of agents and
data; while, the second error term $\bm{q}_{i}$, which we call incremental
error, captures the error resulting from the incremental implementation, where
at each epoch during one iteration, the gradient is calculated at the local
iterate $w_{k,e-1}$. Note that this second error evaluates the loss function
at the local and global iterates. As we will show later, the incremental error
will fade away, and the dominant factor will be the gradient error. Before
establishing the main result in Theorem 1 on the convergence of ISFedAvg
algorithm, we present preliminary results that will lead to it. Thus, to show
the convergence of the algorithm, we must assure the gradient noise
$\bm{s}_{i}$ has zero mean and bounded variance, and the incremental noise
$\bm{q}_{i}$ has bounded variance. Furthermore, since we split the noise due
to the stochastic gradient into incremental and gradient noise, we can split
the analysis into that of the centralized steps and the local epochs. By
proving that both the centralized and local steps converge, we show the global
algorithm converges too.

Replacing the two error terms (14) and (15) into recursion (III-B) and
subtracting $w^{o}$ from both sides of the equation, we get the following
error recursion:

$\widetilde{\bm{w}}_{i}=\widetilde{\bm{w}}_{i-1}+\mu\frac{1}{K}\sum_{k=1}^{K}\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})+\mu\bm{s}_{i}+\mu\bm{q}_{i}.$ (16)

To bound the $\ell_{2}-$norm of the error, we split it into two terms,
centralized and incremental, using Jensen’s inequality with some constant
$\alpha\in(0,1)$ to be defined later:

$\displaystyle\|\widetilde{\bm{w}}_{i}\|^{2}\leq$
$\displaystyle\frac{1}{\alpha}\left\|\widetilde{\bm{w}}_{i-1}+\mu\frac{1}{K}\sum_{k=1}^{K}\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})+\mu\bm{s}_{i}\right\|^{2}$
$\displaystyle+\frac{1}{1-\alpha}\mu^{2}\|\bm{q}_{i}\|^{2}.$ (17)

We start with the first term that represents the centralized solution. We need
to show that it converges. To do so, we start with the gradient noise, and we
establish in Lemma 1 that it remains bounded. We bound the gradient noise
$\bm{s}_{i}$ under two constructions: sampling with replacement, and sampling
without replacement.

######  Lemma 1 (Estimation of first and second order moments of the gradient
noise).

The gradient noise defined in (14) has zero mean:

$\mathbb{E}\\{\bm{s}_{i}|\bm{w}_{i-1}\\}=0,$ (18)

with bounded variance, regardless of the sampling scheme. More specifically,
sampling agents and data with replacement, results in the following bound:

$\displaystyle\mathbb{E}\\{\|\bm{s}_{i}\|^{2}|\bm{w}_{i-1}\\}\leq$
$\displaystyle\beta_{s}^{2}\|\widetilde{\bm{w}}_{i-1}\|^{2}+\sigma^{2}_{s},$
(19)

where $\widetilde{\bm{w}}_{i-1}=w^{o}-\bm{w}_{i-1}$ and the constants:

$\displaystyle\beta_{s}^{2}$
$\displaystyle\>\overset{\Delta}{=}\>\frac{3\delta^{2}}{L}+\frac{1}{LK^{2}}\sum_{k=1}^{K}\frac{1}{p_{k}}\left(\beta_{s,k}^{2}+3\delta^{2}\right),$
(20) $\displaystyle\sigma_{s}^{2}$
$\displaystyle\>\overset{\Delta}{=}\>\frac{1}{LK^{2}}\sum_{k=1}^{K}\frac{1}{p_{k}}\left\\{\sigma_{s,k}^{2}+\left(3+\frac{6}{E_{k}B_{k}}\right)\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\|^{2}\right\\},$ (21) $\displaystyle\beta_{s,k}^{2}$
$\displaystyle\>\overset{\Delta}{=}\>\frac{3\delta^{2}}{E_{k}B_{k}}\left(1+\frac{1}{N_{k}^{2}}\sum_{n=1}^{N_{k}}\frac{1}{p_{n}^{(k)}}\right),$
(22) $\displaystyle\sigma_{s,k}^{2}$
$\displaystyle\>\overset{\Delta}{=}\>\frac{6}{E_{k}B_{k}N_{k}^{2}}\sum_{n=1}^{N_{k}}\frac{1}{p_{n}^{(k)}}\|\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})\|^{2}.$ (23)

On the other hand, sampling agents and data without replacement results in the
same bound but without the scaling by $L$ in the constants $\beta_{s}^{2}$ and
$\sigma_{s}^{2}$.

###### Proof.

See Appendix B. ∎

The term $\sigma_{s,k}^{2}$ in the bound captures what we call data
variability. It is controlled by the mini-batch size $B_{k}$; as the mini-
batch increases the effect of this term is reduced. The $\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\|$ term quantifies the suboptimality of the global model
locally; we call its effect model variability. It is reduced when the data and
agents are more heterogeneous. From Assumption 2, we can bound it uniformly.

Now that we have identified the mean and variance of the gradient noise, we
can proceed to establish the important conclusion that the following
centralized solution:

$\bm{w}_{i}=\bm{w}_{i-1}-\mu\frac{1}{L}\sum_{\ell\in\mathcal{L}_{i}}\widehat{\nabla_{w^{{\sf
T}}}P_{\ell}(\bm{w}_{i-1})},$ (24)

converges exponentially to an $O(\mu)-$neighbourhood of the optimizer. In this
implementation, the center processor aggregates the approximate gradients of
the selected agents. We will subsequently call upon this result to examine the
convergence behavior of the proposed federated learning solution.

######  Lemma 2 (Mean-square-error convergence of the centralized solution).

Consider the centralized recursion (24) where the cost functions satisfy
Assumption 1, and where the first and second order moments of the gradient
noise process satisfy the conditions in Lemma 1. Also, the samples are chosen
without replacement. For step-size values satisfying
$\mu<2\nu/(\delta^{2}+\beta_{s}^{2}),$ it holds that
$\mathbb{E}\|\widetilde{\bm{w}}_{i}\|^{2}$ converges exponentially fast
according to the recursion (25), where
$\lambda=1-2\mu\nu+\mu^{2}(\delta^{2}+\beta_{s}^{2})\in[0,1).$

$\mathbb{E}\|\widetilde{\bm{w}}_{i}\|^{2}\leq\lambda\mathbb{E}\|\widetilde{\bm{w}}_{i-1}\|^{2}+\mu^{2}\sigma_{s}^{2}$
(25)

It follows from (25) that, for sufficiently small step-sizes:

$\mathbb{E}\|\widetilde{\bm{w}}_{i}\|^{2}\leq\lambda^{i}\mathbb{E}\|\widetilde{\bm{w}}_{0}\|^{2}+\frac{1-\lambda^{i}}{1-\lambda}\mu^{2}\sigma_{s}^{2}.$
(26)

###### Proof.

see Appendix C. ∎

We next bound the incremental noise $\bm{q}_{i}$. To do so, we introduce the
local terms:

$\displaystyle\bm{q}_{k,i,e}\>\overset{\Delta}{=}\>$
$\displaystyle\frac{1}{Kp_{k}}\Bigg{(}\frac{1}{B_{k}}\sum_{b\in\mathcal{B}_{k,i,e}}\frac{1}{N_{k}p_{b}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,b})$ $\displaystyle-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{k,e-1})\Bigg{)}.$ (27)

We show in the next lemma that the local gradient noise $\bm{q}_{k,i,e}$ has
zero mean and bounded variance. This result is useful for showing that the
local SGD steps converge in the mean-square-error sense towards their local
models $w^{o}_{k}$.

######  Lemma 3 (Estimation of first and second order moments of the local
gradient noise).

The local gradient noise defined in (III-B) has zero mean:

$\mathbb{E}\left\\{\bm{q}_{k,i,e}\big{|}\mathcal{F}_{e-1},\mathcal{L}_{i}\right\\}=0,$
(28)

and bounded variance, regardless of the sampling scheme:

$\displaystyle\mathbb{E}\left\\{\|\bm{q}_{k,i,e}\|^{2}\big{|}\mathcal{F}_{e-1},\mathcal{L}_{i}\right\\}\leq$
$\displaystyle\frac{E_{k}}{K^{2}p_{k}^{2}}\beta_{s,k}\|\widetilde{\bm{w}}_{k,e-1}\|^{2}+\frac{1}{K^{2}p_{k}^{2}}\sigma_{q,k}^{2},$
(29)

where $\mathcal{F}_{e-1}=\\{w_{k,0},w_{k,1},\cdots,w_{k,e-1}\\}$ is the
filtration describing all sources of randomness due to the previous iterates,
$\widetilde{\bm{w}}_{k,e}=w_{k}^{o}-\bm{w}_{k,e}$, and the constants are as
defined in (22) and:

$\sigma_{q,k}^{2}=\frac{3}{B_{k}N_{k}^{2}}\sum_{n=1}^{N_{k}}\frac{1}{p_{n}^{(k)}}\|\nabla_{w^{{\sf
T}}}Q_{k}(w_{k}^{o};x_{k,n})\|^{2}.$ (30)

###### Proof.

see Appendix D. ∎

Now that we have showed that the local gradient noise of the incremental step
has bounded variance, we can study the mean square deviation of the local SGD.

######  Lemma 4 (Mean-square-error convergence of the local incremental step).

For every agent $k$, consider the local stochastic gradient recursion (8)
where the cost function is subject to Assumption 1, and where the first and
second order moments of the gradient noise process satisfy the conditions in
Lemma 3. For step-size values satisfying:

$\mu<\frac{2\nu}{\delta^{2}+\frac{E_{k}}{K^{2}p_{k}^{2}}\beta_{s,k}^{2}},$
(31)

it holds that $\mathbb{E}\|\widetilde{\bm{w}}_{k,e}\|^{2}$ converges
exponentially fast according to the recursion:

$\mathbb{E}\|\widetilde{\bm{w}}_{k,e}\|^{2}\leq\lambda_{k}\mathbb{E}\|\widetilde{\bm{w}}_{k,e-1}\|^{2}+\mu^{2}\sigma_{q,k}^{2},$
(32)

where:

$\lambda_{k}=1-2\nu\mu+\mu^{2}\left(\delta^{2}+\frac{E_{k}}{K^{2}p_{k}^{2}}\beta^{2}_{s,k}\right)\in[0,1).$
(33)

It follows from (32) that, for sufficiently small step-sizes:

$\displaystyle\mathbb{E}\|\widetilde{\bm{w}}_{k,e}\|^{2}\leq\lambda_{k}^{e}\mathbb{E}\|\widetilde{\bm{w}}_{k,0}\|^{2}+\frac{1-\lambda_{k}^{e}}{1-\lambda_{k}}\mu^{2}\sigma_{q,k}^{2}.$
(34)

###### Proof.

see Appendix E ∎

We can finally bound the incremental noise in the following lemma.

######  Lemma 5 (Estimation of the second order moments of the incremental
noise).

The incremental noise defined in (15) has bounded variance:

$\displaystyle\mathbb{E}\|\bm{q}_{i}\|^{2}$ $\displaystyle\leq
O(\mu)\mathbb{E}\|\widetilde{\bm{w}}_{i-1}\|^{2}+O(\mu)\xi^{2}+O(\mu^{2})\frac{1}{K}\sum_{k=1}^{K}\sigma_{q,k}^{2},$
(35)

where the $O(\cdot)$ terms depend on epoch sizes, local convergence rates,
total number of data samples, number of agents, Lipschitz constant, and data
and agent normalized inclusion probabilities. Thus,
$\mathbb{E}\|\bm{q}_{i}\|^{2}=O(\mu)$.

###### Proof.

see Appendix F. ∎

We observe an average data variability term across agents $\sigma_{q,k}^{2}$,
and a model variability term $\xi^{2}$. However, the effect of the latter
dominates since it is multiplied by an $O(\mu)$ term as opposed to
$O(\mu^{2})$.

###  III-C Main Theorem

Now that we have bounded each term of (III-B) and using Lemma 2 and 5, we
find:

$\displaystyle\mathbb{E}\|\widetilde{\bm{w}}_{i}\|^{2}$
$\displaystyle\leq\frac{1}{\alpha}\left(\lambda\mathbb{E}\|\widetilde{\bm{w}}_{i-1}\|^{2}+\mu^{2}\sigma_{s}^{2}\right)+\frac{O(\mu^{3})}{1-\alpha}\left(\mathbb{E}\|\widetilde{\bm{w}}_{i-1}\|^{2}+\xi^{2}\right)$
$\displaystyle\quad+\frac{O(\mu^{4})}{1-\alpha}\frac{1}{K}\sum_{k=1}^{K}\sigma_{q,k}^{2},$
$\displaystyle=\lambda^{\prime}\mathbb{E}\|\widetilde{\bm{w}}_{i-1}\|^{2}+\frac{\mu^{2}\sigma_{s}^{2}}{\alpha}+\frac{O(\mu^{3})}{1-\alpha}\xi^{2}+\frac{O(\mu^{4})}{1-\alpha}\frac{1}{K}\sum_{k=1}^{K}\sigma_{q,k}^{2},$
(36)

where:

$\displaystyle\lambda^{\prime}$
$\displaystyle\>\overset{\Delta}{=}\>\frac{\lambda}{\alpha}+\frac{O(\mu^{3})}{1-\alpha},$
$\displaystyle=\frac{1-2\mu\nu+\mu^{2}(\delta^{2}+\beta_{s}^{2})}{\alpha}+\frac{O(\mu^{3})}{1-\alpha}$
$\displaystyle=O\left(\frac{1}{\alpha}\right)+O\left(\frac{\mu}{\alpha}\right)+O\left(\frac{\mu^{2}}{\alpha}\right)+O\left(\frac{\mu^{3}}{1-\alpha}\right).$
(37)

Applying this bound recursively we obtain:

$\displaystyle\mathbb{E}\|\widetilde{\bm{w}}_{i}\|^{2}$
$\displaystyle\leq\left(\lambda^{\prime}\right)^{i}\mathbb{E}\|\widetilde{\bm{w}}_{0}\|^{2}+\frac{1-\left(\lambda^{\prime}\right)^{i}}{1-\lambda^{\prime}}\Bigg{(}\frac{\mu^{2}\sigma_{s}^{2}}{\alpha}+\frac{O(\mu^{3})}{1-\alpha}\xi^{2}$
$\displaystyle\quad+\frac{O(\mu^{4})}{1-\alpha}\frac{1}{K}\sum_{k=1}^{K}\sigma_{q,k}^{2}\Bigg{)},$
(38)

and then taking the limit $i\to\infty$:

$\displaystyle\lim_{i\to\infty}\mathbb{E}\|\widetilde{\bm{w}}_{i}\|^{2}$
$\displaystyle\leq\frac{1}{1-\lambda^{\prime}}\left(\frac{\mu^{2}\sigma_{s}^{2}}{\alpha}+\frac{O(\mu^{3})}{1-\alpha}\xi^{2}+\frac{O(\mu^{4})}{1-\alpha}\frac{1}{K}\sum_{k=1}^{K}\sigma_{q,k}^{2}\right).$
(39)

for $\lambda^{\prime}<1$, which for $\alpha=\sqrt{\lambda}$ is achieved when:

$\displaystyle\mu<\min\left\\{\frac{2\nu}{\delta^{2}+\beta_{s}^{2}},\frac{2\nu}{\delta^{2}+\frac{E_{k}}{K^{2}p_{k}^{2}}\beta_{s,k}^{2}}\right\\},$
(40) $\displaystyle O(\mu^{3})<(1-\sqrt{\lambda})^{2}.$ (41)

Thus, since $\alpha=O(1)$, $1-\alpha=O(\mu)$ and $1-\lambda^{\prime}=O(\mu)$:

$\displaystyle\lim_{i\to\infty}\mathbb{E}\|\widetilde{\bm{w}}_{i}\|^{2}\leq
O(\mu)\sigma_{s}^{2}+O(\mu)\xi^{2}+O(\mu^{2})\frac{1}{K}\sum_{k=1}^{K}\sigma_{q,k}^{2}.$
(42)

The result is summarized in the theorem.

######  Theorem 1 (Mean-square-error convergence of federated learning under
importance sampling).

Consider the iterates $\bm{w}_{i}$ generated by the importance sampling
federated averaging algorithm. For sufficiently small step-size $\mu$, it
holds that the mean-square-error converges exponentially fast:

$\displaystyle\mathbb{E}\|\widetilde{\bm{w}}_{i}\|^{2}\leq$ $\displaystyle
O\left((\lambda^{\prime})^{i}\right)+O(\mu)\left(\sigma_{s}^{2}+\xi^{2}\right)+O(\mu^{2})\frac{1}{K}\sum_{k=1}^{K}\sigma_{q,k}^{2},$
(43)

where $\lambda^{\prime}=1-O(\mu)+O(\mu^{2})\in[0,1).$ ∎

##  IV Importance Sampling

Due to the heterogeneity of nodes which arise from their data and
computational capabilities, it is important to guide the algorithm based on
the potential contribution that each agent can have on the overall
performance. By allowing asynchronicity, i.e., different epoch sizes among
agents, we can take advantage of the varying computational capabilities. From
[38], we know that the choice of samples at each iteration affects the
solution. Therefore, instead of choosing the samples uniformly, we consider
importance sampling where samples are chosen according to some distribution to
be determined. A similar scheme can be enforced on the participating agents.
In what follows, we show that using importance sampling enhances the overall
performance.

###  IV-A Agent Level: Importance Sampling of Data

Every agent $k$ at each epoch must select a mini-batch of data based on the
normalized inclusion probabilities $p_{n}^{(k)}$. To find the optimal
probabilities, we minimize the bound on the variance of the local gradient
noise $\sigma_{s,k}^{2}$. We solve the problem for both cases when sampling is
done with replacement and without replacement. The results are the same for
both sampling schemes.

######  Lemma 6 (Optimal local data inclusion probabilities).

The optimal local data normalized inclusion probabilities are given by:

$p_{n}^{(k),o}\>\overset{\Delta}{=}\>\frac{\|\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})\|}{\sum_{m=1}^{N_{k}}\|\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,m})\|}.$ (44)

###### Proof.

By introducing a Lagrange multiplier $\lambda$, the optimization problem can
be reformulated as:

$\min_{p_{n}^{(k)},\lambda}\sum_{n=1}^{N_{k}}\frac{1}{p_{n}^{(k)}}\|\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})\|^{2}+\lambda\left(\sum_{n=1}^{N_{k}}p_{n}^{(k)}-1\right).$
(45)

Then, taking the derivative with respect to $p_{n}^{(k)}$ and setting it to
zero we get (46). Next, substituting $p_{n}^{(k),o}$ into the condition, we
find (47).

$\displaystyle p_{n}^{(k),o}$ $\displaystyle=\frac{\|\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})\|}{\sqrt{\lambda}}$ (46)
$\displaystyle\sqrt{\lambda}$
$\displaystyle=\sum_{m=1}^{N_{k}}\|\nabla_{w^{{\sf T}}}Q_{k}(w^{o};x_{k,m})\|$
(47)

∎

As seen in Lemma 6, more weight is given to a data point that has a greater
gradient norm, thus increasing its chances of being sampled and resulting in a
faster convergence rate. In addition, we observe that the more homogeneous the
data is the more uniform the inclusion probability is.

###  IV-B Cloud Level: Importance Sampling of Agents

At each iteration, the cloud must select a subset of agents to participate.
The agents are selected in accordance with the normalized inclusion
probabilities $p_{k}$. To find the optimal probabilities, we minimize the
bound on the variance of the gradient noise $\sigma_{s}^{2}$ The following
result holds for sampling with and without replacement, since the gradient
noise only differ by a multiplicative factor.

######  Lemma 7 (Optimal agent inclusion probabilities for sampling with
replacement).

The optimal agent normalized inclusion probabilities are given by (48), where
$\alpha_{k}=\left(3+\frac{6}{E_{k}B_{k}}\right)$:

$p_{k}^{o}\>\overset{\Delta}{=}\>\frac{\sqrt{\sigma_{s,k}^{2}+\alpha_{k}\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\|^{2}}}{\sum_{\ell=1}^{K}\sqrt{\sigma_{s,\ell}^{2}+\alpha_{\ell}\|\nabla_{w^{{\sf
T}}}P_{\ell}(w^{o})\|^{2}}}.$ (48)

###### Proof.

The proof follows similarly to that of Lemma 6. ∎

We observe that the normalized inclusion probabilities will be closer to a
uniform distribution the more the data and model variability terms are similar
across agents.

###  IV-C Practical Issues

In the previous subsections, we focused on finding the optimal inclusion
probabilities for the agents and data. However, several practical issues
arise. The first is that all probabilities are calculated based on the optimal
model $w^{o}$, which we do not have access to. To overcome this issue, we
estimate the probabilities at each iteration by calculating them according to
the current model $\bm{w}_{i-1}$. Thus,

$\displaystyle\widehat{p}_{n}^{(k),o}$ $\displaystyle=\frac{\|\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};x_{k,n})\|}{\sum_{m=1}^{N_{k}}\|\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};x_{m})\|},$ (49) $\displaystyle\widehat{p}_{k}^{o}$
$\displaystyle=\frac{\sqrt{\sigma^{2}_{s,k}+\alpha_{k}\|\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\|^{2}}}{\sum_{\ell=1}^{K}\sqrt{\sigma^{2}_{s,\ell}+\alpha_{\ell}\|\nabla_{w^{{\sf
T}}}P_{\ell}(\bm{w}_{i-1})\|^{2}}}.$ (50)

In addition, since calculating the true gradient of the local loss function is
costly, we replace it with the mini-batch approximation when calculating
$p_{k}^{o}$:

$\widehat{p}_{k}^{o}=\frac{\sqrt{\sigma^{2}_{s,k}+\alpha_{k}\left\|\widehat{\nabla_{w^{{\sf
T}}}P_{k}}(\bm{w}_{i-1})\right\|^{2}}}{\sum_{\ell=1}^{K}\sqrt{\sigma^{2}_{s,\ell}+\alpha_{\ell}\left\|\widehat{\nabla_{w^{{\sf
T}}}P_{\ell}}(\bm{w}_{i-1})\right\|^{2}}}.$ (51)

Furthermore, every agent has access to all of its data and consequently to all
of the gradients. However, the cloud does not have access to the gradients of
all agents, and in turn cannot calculate the denominator of $p_{k}$. Instead,
we propose the following solution: at iteration 0, all probabilities are set
to $p_{k}=\frac{1}{K}$; then, during the $i^{th}$ iteration, after the
participating agents $\ell\in\mathcal{L}_{i}$ send the cloud their stochastic
gradients $\widehat{\nabla_{w^{{\sf T}}}P_{\ell}}(\bm{w}_{i-1})$, the
probabilities are updated as follows:

$\displaystyle\widehat{p}_{k}^{o}=$
$\displaystyle\frac{\sqrt{\sigma^{2}_{s,k}+\alpha_{k}\left\|\widehat{\nabla_{w^{{\sf
T}}}P_{k}}(\bm{w}_{i-1})\right\|^{2}}}{\sum_{\ell\in\mathcal{L}_{i}}\sqrt{\sigma^{2}_{s,\ell}+\alpha_{\ell}\left\|\widehat{\nabla_{w^{{\sf
T}}}P_{\ell}}(\bm{w}_{i-1})\right\|^{2}}}\left(1-\sum_{\ell\in\mathcal{L}_{i}^{c}}\widehat{p}_{\ell}^{o}\right),$
(52)

where the multiplicative factor follows from ensuring all the probabilities
$\widehat{p}_{k}^{o}$ sum to $1$. Similarly for the local probabilities, since
we are implementing mini-batch SGD, we only update the probabilities of the
data points that were sampled:

$\displaystyle\widehat{p}_{n}^{(k),o}=$ $\displaystyle\frac{\|\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};x_{k,n})\|}{\sum_{b\in\mathcal{B}_{k,i,e}}\|\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};x_{b})\|}\left(1-\sum_{b\in\mathcal{B}_{k,i,e}^{c}}\widehat{p}_{b}^{(k),o}\right).$
(53)

Finally, the last problem arises when sampling without replacement. We have
found the optimal inclusion probabilities and not the optimal sampling
probabilities, and moving from the former to the latter is not trivial. Thus,
we rely on the literature under sampling without replacement with unequal
probabilities. Multiple sampling schemes exist such that the sampling
probabilities do not need to be calculated explicitly. In general, there are
multiple non-uniform sampling without replacement schemes that guarantee the
same inclusion probabilities. We choose to implement the sampling scheme
proposed in [44], which ensures the inclusion probabilities are $p_{k}^{o}$
and $p_{n}^{(k),o}$ for the agents and data, respectively. More explicitly, we
first calculate the progressive totals of the inclusion probabilities
$\Pi_{k}=\sum_{\ell=1}^{k}Lp_{k}$, for $k=1,2,\cdots,K$, and we set
$\Pi_{0}=0$. Then, we select uniformly at random a uniform variate
$d\in[0,1)$. Then, we select the $L$ agents that satisfy $\Pi_{k-1}\leq
q+\ell<\Pi_{k}$, for some $\ell=0,1,\cdots,L-1$.

##  V Experimental Section

To validate the theoretical results, we devise two experiments. The first
consists of simulated data with quadratic risk functions, and the second
consists of a real dataset with logistic risk functions.

###  V-A Regression

We first validate the theory on a regression problem. We consider $K=300$
agents, for which we generate $N_{k}=100$ data points for each agent $k$ as
follow: Let $\bm{u}_{k,n}$ denote an independent streaming sequence of two-
dimensional random vectors with zero mean and covariance matrix
$R_{u_{k}}=\mathbb{E}\>\bm{u}_{k,i}\bm{u}_{k,i}^{{\sf T}}$. Let
$\bm{d}_{k}(n)$ denote a streaming sequence of random variables that have zero
mean and variance $\sigma_{d_{k}}^{2}=\mathbb{E}\>\bm{d}_{k}^{2}(n)$. Let
$r_{d_{k}u_{k}}=\mathbb{E}\>\bm{d}_{k}(n)\bm{u}_{k,n}$ be the cross-variance
vector. The data $\\{\bm{d}_{k}(n),\bm{u}_{k,n}\\}$ are related by the
following linear regression model:

$\bm{d}_{k}(n)=\bm{u}_{k,n}w^{\star}+\bm{v}_{k}(n),$ (54)

for some randomly generated parameter vector $w^{\star}$ and where
$\bm{v}_{k}(n)$ is a zero mean white noise process with variance
$\sigma_{v_{k}}^{2}=\mathbb{E}\bm{v}_{k}^{2}(n)$, independent of
$\bm{u}_{k,n}$. The local risk is given by:

$P_{k}(w)=\frac{1}{N_{k}}\sum_{n=1}^{N_{k}}\|\bm{d}_{k}(n)-\bm{u}_{k,n}^{{\sf
T}}w\|^{2}+\rho\|w\|^{2}.$ (55)

We set $\rho=0.001$, while the batch sizes $B_{k}$ and the epoch sizes $E_{k}$
are chosen uniformly at random from the range $[1,10]$ and $[1,5]$,
respectively. During each iteration, there are $L=6$ active agents. To test
the performance of the algorithm, we calculate at each iteration the mean-
square-deviation (MSD) of the parameter vector $\bm{w}_{i}$ with respect to
the true model $w^{o}$:

$\mbox{\rm MSD}_{i}=\|\bm{w}_{i}-w^{o}\|^{2}.$ (56)

The optimization problem has the closed form expression:

$w^{o}=\left(\widehat{R}_{u}+\rho
I\right)^{-1}\widehat{R}_{u}w^{\star}+\left(\widehat{R}_{u}+\rho
I\right)^{-1}\widehat{r}_{uv},$ (57)

where:

$\displaystyle\widehat{R}_{u}$
$\displaystyle\>\overset{\Delta}{=}\>\frac{1}{K}\sum_{k=1}^{K}\frac{1}{N_{k}}\sum_{n=1}^{N_{k}}\bm{u}_{k,n}^{\sf
T}\bm{u}_{k,n},$ (58) $\displaystyle\widehat{r}_{uv}$
$\displaystyle\>\overset{\Delta}{=}\>\frac{1}{K}\sum_{k=1}^{K}\frac{1}{N_{k}}\sum_{n=1}^{N_{k}}\bm{v}_{k}(n)\bm{u}_{k,n}.$
(59)

We run four tests: we first run the standard FedAvg algorithm where the mini-
batches are chosen uniformly with replacement. We then run Algorithm 1, once
with the optimal probabilities $p_{n}^{(k)}$ and $p_{k}$ in (44)–(48), once
with the approximate probabilities (49)–(50), and once with (52)–(53). We
implement the sampling scheme from [44]. We set the step-size $\mu=0.01$. Each
test is repeated 100 times, and the resulting MSD is averaged. We get the
curves as shown in Figure 1. We see that the importance sampling scheme does
better than the standard sampling algorithm. This comes as no surprise, since
the probabilities were chosen to minimize the bound on the MSD. The importance
sampling scheme (green curve) improved the MSD bound by 23.1 dB compared to
the standard federated learning scheme (blue curve). Furthermore, we observe
that approximate probabilities do not degrade the performance of the
algorithm. Our proposed approximate solution (52)–(53) (purple curve) performs
just as well as using the true probabilities. In fact, we observe that the
approximate probabilities converge to the true ones,
$\|p_{k}^{o}-\widehat{p}_{k}^{o}\|=1.22e-2$ and
$\frac{1}{K}\sum_{k=1}^{K}\|p_{n}^{(k),o}-\widehat{p}_{n}^{(k),o}\|=1.54e-2$.
Similarly, the approximate probabilities (50)–(49) do not degrade the overall
performance. They, in fact, outperform the other solutions and converge faster
(red curve). This is not surprising, since, at each iteration, we are
attributing higher probabilities to agents and data points that have greater
gradients. We are increasing their chances of being selected and thus taking
steeper steps towards the true model.

Figure 1: MSD plots of the regression problem: blue curve is the standard
mini-batch implementation, green curve is the importance sampling
implementation with the true probabilities, red curve is the importance
sampling implementation with approximate probabilities (49)–(50), purple curve
is the importance sampling implementation with approximate probabilities
(52)–(53).

###  V-B Classification

We next study the theory in a classification context. We consider the ijcnn1
dataset [45]. The dataset consists of $35000$ training samples and $91701$
testing samples of $M=22$ attributes. We distribute the data randomly in a
non-IID fashion to $K=100$ agents. Each agent receives a random number $N_{k}$
of data points, where $N_{k}$ ranges from 79 to 688. We run the two algorithms
FedAvg and ISFedAvg. We set $\mu=0.25$, $\rho=0.0001$, $L=10$, $B_{k}=1$, and
$E_{k}=1$. We plot the testing error in Figure 2. We observe that importance
sampling improves the testing error from $22.45\%$ to $18.46\%$. This is
because importance sampling is more sample efficient.

Figure 2: Testing error plots of the classification problem.

##  VI Conclusion

This work incorporates two levels of importance sampling into the operation of
federated learning: one for selecting agents and another for selecting data
batches at the agents. Optimal dynamical choices for the sampling
probabilities are derived, and a detailed convergence analysis is performed.
We also provided approximate expressions for the optimal sampling policies and
illustrate the theoretical findings and the performance enhancement by means
of simulations.

##  Appendix A Result on the Variance of the Mini-batch Estimate

We introduce the following auxiliary result that is a generalization of [46].
Let $\\{\mathcal{S}=\bm{x}_{n}\in\mathds{R}^{M}\\}_{n=1}^{N}$ denote a set of
$N$ _independent_ random variables, each with mean
$\mathbb{E}\bm{x}_{n}=\overline{x}_{n}$ and variance
$\sigma_{n}^{2}=\mathbb{E}{\|\bm{x}_{n}-\mathbb{E}\bm{x}_{n}\|}^{2}$. We
consider the problem of estimating the expected value of the sample mean:

$\overline{x}\triangleq\mathbb{E}\left(\frac{1}{N}\sum_{n=1}^{N}\bm{x}_{n}\right)$
(60)

We consider two estimators for $\overline{x}$, both constructed by considering
a mini-batch of samples, where $\bm{x}_{b}^{\mathrm{r}}$ is constructed by
sampling from $\mathcal{S}$ _with replacement_ , and
$\bm{x}_{b}^{\mathrm{nr}}$ is sampled from $\mathcal{S}$ _without
replacement_. Let $p_{n}$ be the normalized inclusion probability of
$\bm{x}_{n}$. We then define the two estimators (61) and (62), and we would
like to quantify the efficacy of these estimators in estimating
$\overline{x}$.

$\displaystyle\widehat{\bm{x}}^{\mathrm{r}}$
$\displaystyle\triangleq\frac{1}{B}\sum_{b=1}^{B}\frac{1}{Np_{b}}\bm{x}_{b}^{\mathrm{r}}$
(61) $\displaystyle\widehat{\bm{x}}^{\mathrm{nr}}$
$\displaystyle\triangleq\frac{1}{B}\sum_{b=1}^{B}\frac{1}{Np_{b}}\bm{x}_{b}^{\mathrm{nr}}$
(62)

######  Lemma 8 (Variance of the mini-batch mean with and without
replacement).

Both estimators are unbiased and it holds that:

$\displaystyle\mathbb{E}{\widehat{\bm{x}}^{\mathrm{r}}}=\mathbb{E}{\widehat{\bm{x}}^{\mathrm{nr}}}$
$\displaystyle=\overline{x},$ (63)
$\displaystyle\mathbb{E}{\|\widehat{\bm{x}}^{\mathrm{r}}-\overline{x}\|}^{2}$
$\displaystyle=\frac{1}{B}\sum_{n=1}^{N}p_{n}\left(\frac{1}{N^{2}p_{n}^{2}}\sigma_{n}^{2}+{\left\|\frac{1}{Np_{n}}\overline{x}_{n}-\overline{x}\right\|}^{2}\right),$
(64)
$\displaystyle\mathbb{E}{\left\|\widehat{\bm{x}}^{\mathrm{nr}}-\overline{x}\right\|}^{2}$
$\displaystyle=\frac{1}{B}\sum_{n=1}^{N}p_{n}\left(\frac{1}{N^{2}p_{n}^{2}}\sigma_{n}^{2}+\left\|\frac{1}{Np_{n}}\overline{x}_{n}-\overline{x}\right\|^{2}\right)$
$\displaystyle+\frac{1}{B^{2}}\sum_{n_{1}\neq
n_{2}}\mathbb{P}(\mathbb{I}_{n_{1}}=1,\mathbb{I}_{n_{2}}=1)$
$\displaystyle\times\left(\frac{1}{Np_{n_{1}}}\overline{x}_{n_{1}}-\overline{x}\right)\left(\frac{1}{Np_{n_{2}}}\overline{x}_{n_{2}}-\overline{x}\right),$
(65)

where the notation $\mathbb{I}_{n}$ evaluates to one if $x_{n}\in{\cal
B}^{nr}$ and is zero otherwise.

###### Proof.

We begin with the with-replacement setting. The randomness of the samples
introduces some intricacies that need to be accounted for in the notation. For
the mean, we have:

$\displaystyle\mathbb{E}\widehat{\bm{x}}^{\mathrm{r}}$
$\displaystyle=\frac{1}{B}\sum_{b=1}^{B}\mathbb{E}\left(\frac{1}{Np_{b}}\bm{x}_{b}^{\mathrm{r}}\right)=\frac{1}{B}\sum_{b=1}^{B}\mathbb{E}\left\\{\mathbb{E}\left\\{\frac{1}{Np_{b}}\bm{x}_{b}^{\mathrm{r}}\bigg{|}\mathcal{S}\right\\}\right\\}$
$\displaystyle=\frac{1}{B}\sum_{b=1}^{B}\mathbb{E}\left\\{\sum_{n=1}^{N}p_{n}\frac{1}{Np_{n}}\bm{x}_{n}\right\\}=\frac{1}{B}\sum_{b=1}^{B}\overline{x}=\overline{x}.$
(66)

For the variance we find:

$\displaystyle\mathbb{E}{\left\|\widehat{\bm{x}}^{\mathrm{r}}-\overline{x}\right\|}^{2}$
$\displaystyle=\>\mathbb{E}{\left\|\frac{1}{B}\sum_{b=1}^{B}\frac{1}{Np_{b}}\bm{x}_{b}^{\mathrm{r}}-\overline{x}\right\|}^{2},$
$\displaystyle=\>\mathbb{E}{\left\|\frac{1}{B}\sum_{b=1}^{B}\left(\frac{1}{Np_{b}}\bm{x}_{b}^{\mathrm{r}}-\overline{x}\right)\right\|}^{2},$
$\displaystyle=\>\frac{1}{B^{2}}\sum_{b=1}^{B}\mathbb{E}{\left\|\frac{1}{Np_{b}}\bm{x}_{b}^{\mathrm{r}}-\overline{x}\right\|}^{2}$
$\displaystyle\quad+\frac{1}{B^{2}}\sum_{b_{1}\neq
b_{2}}\mathbb{E}\left\\{{\left(\frac{1}{Np_{b_{1}}}\bm{x}_{b_{1}}-\overline{x}\right)}{\left(\frac{1}{Np_{b_{2}}}\bm{x}_{b_{2}}-\overline{x}\right)}\right\\},$
$\displaystyle\stackrel{{\scriptstyle(a)}}{{=}}\>\frac{1}{B^{2}}\sum_{b=1}^{B}\mathbb{E}{\left\|\frac{1}{Np_{b}}\bm{x}_{b}^{\mathrm{r}}-\overline{x}\right\|}^{2}$
$\displaystyle\quad+\frac{1}{B^{2}}\sum_{b_{1}\neq
b_{2}}\mathbb{E}{\left\\{\frac{1}{Np_{b_{1}}}\bm{x}_{b_{1}}-\overline{x}\right\\}}\mathbb{E}{\left\\{\frac{1}{Np_{b_{2}}}\bm{x}_{b_{2}}-\overline{x}\right\\}},$
$\displaystyle\stackrel{{\scriptstyle(b)}}{{=}}\>\frac{1}{B^{2}}\sum_{b=1}^{B}\mathbb{E}{\left\|\frac{1}{Np_{b}}\bm{x}_{b}^{\mathrm{r}}-\overline{x}\right\|}^{2},$
(67)

where $(a)$ is a result of the fact that the elements of $\mathcal{S}$ are
independent and $\bm{x}_{b}^{\mathrm{r}}$ is sampled from $\mathcal{S}$
independently, and hence $\bm{x}_{b_{1}}$ and $\bm{x}_{b_{2}}$ are
independent. Step $(b)$ then follows from:

$\mathbb{E}\left(\frac{1}{Np_{b}}\bm{x}_{b}\right)=\mathbb{E}\left(\frac{1}{N}\sum_{n=1}^{N}\bm{x}_{n}\right)=\overline{x}.$
(68)

Then,

$\displaystyle\mathbb{E}{\left\|\widehat{\bm{x}}^{\mathrm{r}}-\overline{x}\right\|}^{2}$
$\displaystyle{=}\>\frac{1}{B^{2}}\sum_{b=1}^{B}\mathbb{E}{\left\|\frac{1}{Np_{b}}\bm{x}_{b}^{\mathrm{r}}-\overline{x}\right\|}^{2},$
$\displaystyle{=}\>\frac{1}{B^{2}}\sum_{b=1}^{B}\mathbb{E}\left\\{\mathbb{E}{\left\|\frac{1}{Np_{b}}\bm{x}_{b}^{\mathrm{r}}-\overline{x}\right\|}^{2}\bigg{|}\mathcal{S}\right\\},$
$\displaystyle{=}\>\frac{1}{B^{2}}\sum_{b=1}^{B}\mathbb{E}\left\\{\sum_{n=1}^{N}{p_{n}\left\|\frac{1}{Np_{n}}\bm{x}_{n}-\overline{x}\right\|}^{2}\right\\},$
$\displaystyle{=}\>\frac{1}{B^{2}}\sum_{b=1}^{B}\sum_{n=1}^{N}p_{n}\mathbb{E}{\left\|\frac{1}{Np_{n}}\bm{x}_{n}-\overline{x}\right\|}^{2},$
$\displaystyle{=}\>\frac{1}{B}\sum_{n=1}^{N}p_{n}\mathbb{E}{\left\|\frac{1}{Np_{n}}\bm{x}_{n}-\frac{1}{Np_{n}}\overline{x}_{n}+\frac{1}{Np_{n}}\overline{x}_{n}-\overline{x}\right\|}^{2},$
$\displaystyle{=}\>\frac{1}{B}\sum_{n=1}^{N}p_{n}\left(\mathbb{E}{\left\|\frac{1}{Np_{n}}\bm{x}_{n}-\frac{1}{Np_{n}}\overline{x}_{n}\right\|}^{2}+{\left\|\frac{1}{Np_{n}}\overline{x}_{n}-\overline{x}\right\|}^{2}\right),$
$\displaystyle{=}\>\frac{1}{B}\sum_{n=1}^{N}p_{n}\left(\frac{1}{N^{2}p_{n}^{2}}\sigma_{n}^{2}+{\left\|\frac{1}{Np_{n}}\overline{x}_{n}-\overline{x}\right\|}^{2}\right).$
(69)

We now proceed to study the efficiency of the without replacement mini-batch
mean. The fact that the $\bm{x}_{b}$ are sampled from $\mathcal{S}$ without
replacement causes pairs $\bm{x}_{b_{1}},\bm{x}_{b_{2}}$ to no longer be
independent. We denote the set of points sampled from $\mathcal{S}$ without
replacement by $\mathcal{B}^{\mathrm{nr}}$ and introduce the activation
function by:

$\mathbb{I}_{n}\triangleq\begin{cases}1,\ \mathrm{if}\
\bm{x}_{n}\in\mathcal{B}^{\mathrm{nr}},\\\ 0,\ \mathrm{if}\
\bm{x}_{n}\notin\mathcal{B}^{\mathrm{nr}}.\end{cases}$ (70)

Then, the estimator $\widehat{\bm{x}}^{\mathrm{nr}}$ can be written
equivalently as:

$\widehat{\bm{x}}^{\mathrm{nr}}=\frac{1}{B}\sum_{n=1}^{N}\mathbb{I}_{n}\frac{1}{Np_{n}}\bm{x}_{n}.$
(71)

For the mean, we have:

$\displaystyle\mathbb{E}\widehat{\bm{x}}^{\mathrm{nr}}$
$\displaystyle=\frac{1}{B}\sum_{n=1}^{N}\mathbb{E}\left\\{\mathbb{I}_{n}\frac{1}{Np_{n}}\bm{x}_{n}\right\\}=\frac{1}{B}\sum_{n=1}^{N}\mathbb{E}\mathbb{I}_{n}\times\mathbb{E}\frac{1}{Np_{n}}\bm{x}_{n}$
$\displaystyle=\frac{1}{B}\sum_{n=1}^{N}Bp_{n}\times\frac{1}{Np_{n}}\overline{x}_{n}=\frac{1}{N}\sum_{n=1}^{N}\overline{x}_{n}=\overline{x}.$
(72)

For the variance, we have:

$\displaystyle\mathbb{E}\left\|\widehat{\bm{x}}^{\mathrm{nr}}-\overline{x}\right\|^{2}=$
$\displaystyle\>\mathbb{E}\left\|\frac{1}{B}\sum_{n=1}^{N}\mathbb{I}_{n}\left(\frac{1}{Np_{n}}\bm{x}_{n}-\overline{x}\right)\right\|^{2},$
$\displaystyle=$
$\displaystyle\>\frac{1}{B^{2}}\sum_{n=1}^{N}\mathbb{E}\left\|\mathbb{I}_{n}\left(\frac{1}{Np_{n}}\bm{x}_{n}-\overline{x}\right)\right\|^{2}$
$\displaystyle+\frac{1}{B^{2}}\sum_{n_{1}\neq
n_{2}}\mathbb{E}\Bigg{\\{}\mathbb{I}_{n_{1}}\left(\frac{1}{Np_{n_{1}}}\bm{x}_{n_{1}}-\overline{x}\right)\mathbb{I}_{n_{2}}$
$\displaystyle\qquad\qquad\qquad\qquad\times\left(\frac{1}{Np_{n_{2}}}\bm{x}_{n_{2}}-\overline{x}\right)\Bigg{\\}}.$
(73)

We begin with:

$\displaystyle\mathbb{E}\left\|\mathbb{I}_{n}\left(\frac{1}{Np_{n}}\bm{x}_{n}-\overline{x}\right)\right\|^{2}$
$\displaystyle=$
$\displaystyle\>\mathbb{E}\left\\{\left\|\mathbb{I}_{n}\left(\frac{1}{Np_{n}}\bm{x}_{n}-\overline{x}\right)\right\|^{2}\bigg{|}\mathbb{I}_{n}=1\right\\}\times\mathds{P}\left(\mathbb{I}_{n}=1\right)$
$\displaystyle+\mathbb{E}\left\\{\left\|\mathbb{I}_{n}\left(\frac{1}{Np_{n}}\bm{x}_{n}-\overline{x}\right)\right\|^{2}\bigg{|}\mathbb{I}_{n}=0\right\\}\times\mathds{P}\left(\mathbb{I}_{n}=0\right),$
$\displaystyle=$
$\displaystyle\>Bp_{n}\left(\mathbb{E}\left\|\frac{1}{Np_{n}}\bm{x}_{n}-\frac{1}{Np_{n}}\overline{x}_{n}+\frac{1}{Np_{n}}\overline{x}_{n}-\overline{x}\right\|^{2}\right),$
$\displaystyle=$
$\displaystyle\>Bp_{n}\left(\frac{1}{N^{2}p^{2}_{n}}\mathbb{E}\|\bm{x}_{n}-\overline{x}_{n}\|^{2}+\left\|\frac{1}{Np_{n}}\overline{x}_{n}-\overline{x}\right\|^{2}\right),$
$\displaystyle=$
$\displaystyle\>Bp_{n}\left(\frac{1}{N^{2}p^{2}_{n}}\sigma_{n}^{2}+\left\|\frac{1}{Np_{n}}\overline{x}_{n}-\overline{x}\right\|^{2}\right).$
(74)

For the cross-term we have:

$\displaystyle\>\mathbb{E}\left\\{\mathbb{I}_{n_{1}}\left(\frac{1}{Np_{n_{1}}}\bm{x}_{n_{1}}-\overline{x}\right)\mathbb{I}_{n_{2}}\left(\frac{1}{Np_{n_{2}}}\bm{x}_{n_{2}}-\overline{x}\right)\right\\}$
$\displaystyle=$
$\displaystyle\>\mathbb{E}\bigg{\\{}\left(\frac{1}{Np_{n_{1}}}\bm{x}_{n_{1}}-\overline{x}\right)\left(\frac{1}{Np_{n_{2}}}\bm{x}_{n_{2}}-\overline{x}\right)\bigg{|}\mathbb{I}_{n_{1}}=1,\mathbb{I}_{n_{2}}=1\bigg{\\}}$
$\displaystyle\times\mathds{P}\left(\mathbb{I}_{n_{1}}=1,\mathbb{I}_{n_{2}}=1\right),$
$\displaystyle=$
$\displaystyle\>\mathbb{P}\left(\mathbb{I}_{n_{2}}=1,\mathbb{I}_{n_{1}}=1\right)\left(\frac{1}{Np_{n_{1}}}\mathbb{E}\bm{x}_{n_{1}}-\overline{x}\right)\left(\frac{1}{Np_{n_{2}}}\mathbb{E}\bm{x}_{n_{2}}-\overline{x}\right),$
$\displaystyle=$
$\displaystyle\>\mathbb{P}\left(\mathbb{I}_{n_{2}}=1,\mathbb{I}_{n_{1}}=1\right)\left(\frac{1}{Np_{n_{1}}}\overline{x}_{n_{1}}-\overline{x}\right)\left(\frac{1}{Np_{n_{2}}}\overline{x}_{n_{2}}-\overline{x}\right).$
(75)

We then get the desired result. ∎

We note the following bound on the variance of sampling without replacement
estimator. Using Jensen’s inequality, we can get rid of the cross-term, and
deduce from (26) the following inequality:

$\mathbb{E}\left\|\widehat{\bm{x}}^{\mathrm{nr}}-\overline{x}\right\|^{2}\leq\sum_{n=1}^{N}p_{n}\left(\frac{1}{N^{2}p_{n}^{2}}\sigma_{n}^{2}+\left\|\frac{1}{Np_{n}}\overline{x}_{n}-\overline{x}\right\|^{2}\right).$
(76)

##  Appendix B Proof of Lemma 1

###### Proof.

We start with the sampling with-replacement construction. We have $K$ agents
from which we sample $L$. Thus, $N$ and $B$ in Lemma 8 are $K$ and $L$,
respectively. Also:

$\displaystyle\bm{x}_{k}$ $\displaystyle=\widehat{\nabla_{w^{{\sf
T}}}P_{k}}(\bm{w}_{i-1}),$ (77) $\displaystyle\overline{x}_{k}$
$\displaystyle=\nabla_{w^{{\sf T}}}P_{k}(\bm{w}_{i-1}),$ (78)
$\displaystyle\overline{x}$
$\displaystyle=\frac{1}{K}\sum_{k=1}^{K}\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1}).$ (79)

Then $\sigma_{k}^{2}$, which quantifies the second order moment of the local
gradient noise, becomes:

$\displaystyle\sigma_{k}^{2}$
$\displaystyle=\mathbb{E}\left\\{\left\|\widehat{\nabla_{w^{{\sf
T}}}P_{k}}(\bm{w}_{i-1})-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\right\|^{2}\bigg{|}\bm{w}_{i-1}\right\\},$
$\displaystyle=\frac{1}{E_{k}^{2}B_{k}^{2}}\sum_{e=1}^{E_{k}}\sum_{b\in\mathcal{B}_{k,i,e}}\mathbb{E}\Bigg{\\{}\bigg{\|}\frac{1}{N_{k}p_{b}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};\bm{x}_{k,b})$
$\displaystyle\qquad\qquad\qquad\qquad\qquad\quad-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\bigg{\|}^{2}\bigg{|}\bm{w}_{i-1}\Bigg{\\}},$
$\displaystyle=\frac{1}{E_{k}B_{k}^{2}}\sum_{b\in\mathcal{B}_{k,i,e}}\mathbb{E}\Bigg{\\{}\bigg{\|}\frac{1}{N_{k}p_{b}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};\bm{x}_{k,b})$
$\displaystyle-\frac{1}{N_{k}p_{b}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};\bm{x}_{k,b})+\frac{1}{N_{k}p_{b}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};\bm{x}_{k,b})$ $\displaystyle-\nabla_{w^{{\sf
T}}}P_{k}(w^{o})+\nabla_{w^{{\sf T}}}P_{k}(w^{o})-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\bigg{\|}^{2}\bigg{|}\bm{w}_{i-1}\Bigg{\\}},$
$\displaystyle\stackrel{{\scriptstyle(a)}}{{\leq}}\frac{3}{E_{k}B_{k}^{2}}\sum_{b\in\mathcal{B}_{k,i,e}}\Bigg{\\{}\mathbb{E}\Bigg{\\{}\bigg{\|}\frac{1}{N_{k}p_{b}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};\bm{x}_{k,b})$
$\displaystyle\qquad\qquad\qquad\qquad-\frac{1}{N_{k}p_{b}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};\bm{x}_{k,b})\bigg{\|}^{2}\bigg{|}\bm{w}_{i-1}\Bigg{\\}}$
$\displaystyle+\mathbb{E}\left\\{\left\|\frac{1}{N_{k}p_{b}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};\bm{x}_{k,b})-\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\right\|^{2}\bigg{|}\bm{w}_{i-1}\right\\}$
$\displaystyle+\mathbb{E}\left\\{\left\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\right\|^{2}\bigg{|}\bm{w}_{i-1}\right\\}\Bigg{\\}},$
$\displaystyle=\frac{3}{E_{k}B_{k}^{2}}\sum_{b\in\mathcal{B}_{k,i,e}}\Bigg{\\{}\sum_{n=1}^{N_{k}}p_{n}^{(k)}\bigg{\|}\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};x_{k,n})$
$\displaystyle\qquad\qquad\qquad\qquad\qquad-\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})\bigg{\|}^{2}$
$\displaystyle+\sum_{n=1}^{N_{k}}p_{n}^{(k)}\left\|\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})-\nabla_{w^{{\sf T}}}P_{k}(w^{o})\right\|^{2}$
$\displaystyle+\left\|\nabla_{w^{{\sf T}}}P_{k}(w^{o})-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\right\|^{2}\Bigg{\\}},$
$\displaystyle\stackrel{{\scriptstyle(b)}}{{\leq}}\frac{3}{E_{k}B_{k}^{2}}\sum_{b\in\mathcal{B}_{k,i,e}}\Bigg{\\{}\left(1+\sum_{n=1}^{N_{k}}\frac{1}{N_{k}^{2}p_{n}^{(k)}}\right)\delta^{2}\|\widetilde{\bm{w}}_{i-1}\|^{2}$
$\displaystyle+\sum_{n=1}^{N_{k}}p_{n}^{(k)}\left\|\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})-\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\right\|^{2}\Bigg{\\}},$
$\displaystyle\stackrel{{\scriptstyle(c)}}{{\leq}}\frac{3\delta^{2}}{E_{k}B_{k}}\left(1+\frac{1}{N_{k}^{2}}\sum_{n=1}^{N_{k}}\frac{1}{p_{n}^{(k)}}\right)\|\widetilde{\bm{w}}_{i-1}\|^{2}$
$\displaystyle+\frac{6}{E_{k}B_{k}N_{k}^{2}}\sum_{n=1}^{N_{k}}\frac{1}{p_{n}^{(k)}}\left\|\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})\right\|^{2}$
$\displaystyle+\frac{6}{E_{k}B_{k}}\left\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\right\|^{2}$
$\displaystyle=\beta_{s,k}^{2}\|\widetilde{\bm{w}}_{i-1}\|^{2}+\sigma_{s,k}^{2}+\frac{6}{E_{k}B_{k}}\left\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\right\|^{2},$ (80)

where $(a)$ and $(c)$ follow from using Jensen’s inequality, and $(b)$ follows
from using the $\delta-$Lipschitz property of the gradients. Thus, using Lemma
8, we bound the stochastic noise variance as follows:

$\displaystyle\mathbb{E}\left\\{\|\bm{s}_{i}\|^{2}|\bm{w}_{i-1}\right\\}=\frac{1}{L}\sum_{k=1}^{K}p_{k}\Bigg{\\{}\frac{1}{K^{2}p_{k}^{2}}\sigma_{k}^{2}$
$\displaystyle+\bigg{\|}\frac{1}{Kp_{k}}\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})-\frac{1}{K}\sum_{\ell=1}^{K}\nabla_{w^{{\sf
T}}}P_{\ell}(\bm{w}_{i-1})\bigg{\|}^{2}\Bigg{\\}}.$ (81)

We focus on the second term since the first term has already been bounded.
Using Jensen’s inequality in $(a)$ and Lipschitz condition of the gradients in
$(b)$, we get:

$\displaystyle\left\|\frac{1}{Kp_{k}}\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})-\frac{1}{K}\sum_{\ell=1}^{K}\nabla_{w^{{\sf
T}}}P_{\ell}(\bm{w}_{i-1})\right\|^{2}$
$\displaystyle=\frac{1}{K^{2}}\bigg{\|}\frac{1}{p_{k}}\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})-\frac{1}{p_{k}}\nabla_{w^{{\sf
T}}}P_{k}(w^{o})+\frac{1}{p_{k}}\nabla_{w^{{\sf T}}}P_{k}(w^{o})$
$\displaystyle+\sum_{\ell=1}^{K}\nabla_{w^{{\sf
T}}}P_{\ell}(w^{o})-\sum_{\ell=1}^{K}\nabla_{w^{{\sf
T}}}P_{\ell}(\bm{w}_{i-1})\bigg{\|}^{2},$
$\displaystyle\stackrel{{\scriptstyle(a)}}{{\leq}}\frac{3}{K^{2}p^{2}_{k}}\left\|\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})-\nabla_{w^{{\sf T}}}P_{k}(w^{o})\right\|^{2}$
$\displaystyle+\frac{3}{K^{2}p^{2}_{k}}\left\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\right\|^{2}$
$\displaystyle+\frac{3}{K}\sum_{\ell=1}^{K}\left\|\nabla_{w^{{\sf
T}}}P_{\ell}(w^{o})-\nabla_{w^{{\sf T}}}P_{\ell}(\bm{w}_{i-1})\right\|^{2},$
$\displaystyle\stackrel{{\scriptstyle(b)}}{{\leq}}3\delta^{2}\left(1+\frac{1}{K^{2}p_{k}^{2}}\right)\|\widetilde{\bm{w}}_{i-1}\|^{2}+\frac{3}{K^{2}p^{2}_{k}}\left\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\right\|^{2}.$ (82)

Then, putting things together, we get:

$\displaystyle\mathbb{E}\left\\{\|\bm{s}_{i}\|^{2}|\bm{w}_{i-1}\right\\}$
$\displaystyle\leq\frac{1}{L}\sum_{k=1}^{K}p_{k}\Bigg{\\{}\frac{\beta_{s,k}^{2}}{K^{2}p_{k}^{2}}\|\widetilde{\bm{w}}_{i-1}\|^{2}+\frac{1}{K^{2}p_{k}^{2}}\sigma_{s,k}^{2}$
$\displaystyle\qquad\frac{6}{K^{2}p_{k}^{2}E_{k}B_{k}}\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\|^{2}+\frac{3}{K^{2}p^{2}_{k}}\left\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\right\|^{2}$
$\displaystyle\qquad+3\delta^{2}\left(1+\frac{1}{K^{2}p_{k}^{2}}\right)\|\widetilde{\bm{w}}_{i-1}\|^{2}\Bigg{\\}},$
$\displaystyle=\frac{1}{L}\sum_{k=1}^{K}\left(\frac{\beta_{s,k}^{2}}{K^{2}p_{k}}+3\delta^{2}p_{k}+\frac{3\delta^{2}}{K^{2}p_{k}}\right)\|\widetilde{\bm{w}}_{i-1}\|^{2}$
$\displaystyle+\frac{1}{LK^{2}}\sum_{k=1}^{K}\frac{1}{p_{k}}\left\\{\sigma_{s,k}^{2}+\left(3+\frac{6}{E_{k}B_{k}}\right)\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\|^{2}\right\\},$
$\displaystyle=\left(\frac{3\delta^{2}}{L}+\frac{1}{LK^{2}}\sum_{k=1}^{K}\frac{1}{p_{k}}\left(\beta_{s,k}^{2}+3\delta^{2}\right)\right)\|\widetilde{\bm{w}}_{i-1}\|^{2}$
$\displaystyle+\frac{1}{LK^{2}}\sum_{k=1}^{K}\frac{1}{p_{k}}\left\\{\sigma_{s,k}^{2}+\left(3+\frac{6}{E_{k}B_{k}}\right)\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\|^{2}\right\\},$
$\displaystyle=\beta_{s}^{2}\|\widetilde{\bm{w}}_{i-1}\|^{2}+\sigma_{s}^{2}.$
(83)

Next, we move to the sampling without replacement construction. The variance
$\sigma_{k}^{2}$ becomes:

$\displaystyle\sigma_{k}^{2}$
$\displaystyle=\mathbb{E}\left\\{\left\|\widehat{\nabla_{w^{{\sf
T}}}P_{k}}(\bm{w}_{i-1})-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\right\|^{2}\bigg{|}\bm{w}_{i-1}\right\\},$
$\displaystyle=\mathbb{E}\Bigg{\\{}\bigg{\|}\frac{1}{E_{k}B_{k}}\sum_{e=1}^{E_{k}}\sum_{n=1}^{N_{k}}\mathbb{I}_{n}\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};\bm{x}_{n})$ $\displaystyle\qquad\quad-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\bigg{\|}^{2}\bigg{|}\bm{w}_{i-1}\Bigg{\\}},$
$\displaystyle=\frac{1}{E_{k}^{2}}\sum_{e=1}^{E_{k}}\mathbb{E}\Bigg{\\{}\bigg{\|}\frac{1}{B_{k}}\sum_{n=1}^{N_{k}}\mathbb{I}_{n}\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};\bm{x}_{n})$
$\displaystyle\qquad\qquad\qquad\quad-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\bigg{\|}^{2}\bigg{|}\bm{w}_{i-1}\Bigg{\\}},$
$\displaystyle=\frac{1}{E_{k}B_{k}^{2}}\sum_{n=1}^{N_{k}}\mathbb{E}\Bigg{\\{}\bigg{\|}\mathbb{I}_{n}\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};\bm{x}_{n})$
$\displaystyle\qquad\qquad\qquad\quad-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\bigg{\|}^{2}\bigg{|}\bm{w}_{i-1}\Bigg{\\}},$
$\displaystyle+\frac{1}{E_{k}B_{k}^{2}}\sum_{n_{1}\neq
n_{2}}\mathbb{E}\Bigg{\\{}\mathbb{I}_{n_{1}}\bigg{(}\frac{1}{N_{k}p_{n_{1}}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};\bm{x}_{n_{1}})$ $\displaystyle-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\bigg{)}\mathbb{I}_{n_{2}}\bigg{(}\frac{1}{N_{k}p_{n_{2}}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};\bm{x}_{n_{2}})$ $\displaystyle-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\bigg{)}\bigg{|}\bm{w}_{i-1}\Bigg{\\}}.$ (84)

Starting with the first term, we use Jensen’s inequality in $(a)$ and $(c)$
and the Lipschitz condition in $(b)$ to get:

$\displaystyle\frac{1}{E_{k}B_{k}^{2}}\sum_{n=1}^{N_{k}}\mathbb{P}(\mathbb{I}_{n}=1)\mathbb{E}\Bigg{\\{}\bigg{\|}\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};x_{k,n})$ $\displaystyle\qquad\qquad-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\bigg{\|}^{2}\bigg{|}\bm{w}_{i-1},\mathbb{I}_{n}=1\Bigg{\\}},$
$\displaystyle=\frac{1}{E_{k}B_{k}}\sum_{n=1}^{N_{k}}p_{n}^{(k)}\bigg{\|}\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};x_{k,n})$
$\displaystyle\qquad\qquad\qquad\quad-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\bigg{\|}^{2},$
$\displaystyle=\frac{1}{E_{k}B_{k}}\sum_{n=1}^{N_{k}}p_{n}^{(k)}\bigg{\|}\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};x_{k,n})$
$\displaystyle-\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})+\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})$ $\displaystyle-\nabla_{w^{{\sf
T}}}P_{k}(w^{o})+\nabla_{w^{{\sf T}}}P_{k}(w^{o})-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\bigg{\|}^{2},$
$\displaystyle\stackrel{{\scriptstyle(a)}}{{\leq}}\frac{3}{E_{k}B_{k}}\sum_{n=1}^{N_{k}}\frac{1}{N_{k}^{2}p_{n}^{(k)}}\|\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};x_{k,n})$ $\displaystyle-\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})\|^{2}$
$\displaystyle+p_{n}^{(k)}\left\|\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})-\nabla_{w^{{\sf T}}}P_{k}(w^{o})\right\|^{2}$
$\displaystyle+p_{n}^{(k)}\left\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})-\nabla_{w^{{\sf T}}}P_{k}(\bm{w}_{i-1})\right\|^{2},$
$\displaystyle=\frac{3}{E_{k}B_{k}}\sum_{n=1}^{N_{k}}\Bigg{\\{}\frac{1}{N_{k}^{2}p_{n}^{(k)}}\|\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{i-1};x_{k,n})$ $\displaystyle-\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})\|^{2}$
$\displaystyle+p_{n}^{(k)}\left\|\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})-\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\right\|^{2}\Bigg{\\}}$
$\displaystyle+\frac{3}{E_{k}B_{k}}\left\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})-\nabla_{w^{{\sf T}}}P_{k}(\bm{w}_{i-1})\right\|^{2},$
$\displaystyle\stackrel{{\scriptstyle(b)}}{{\leq}}\frac{3\delta^{2}}{E_{k}B_{k}}\left(1+\frac{1}{N_{k}^{2}}\sum_{n=1}^{N_{k}}\frac{1}{p_{n}^{(k)}}\right)\|\widetilde{\bm{w}}_{i-1}\|^{2}$
$\displaystyle+\frac{3}{E_{k}B_{k}}\sum_{n=1}^{N_{k}}p_{n}^{(k)}\left\|\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})-\nabla_{w^{{\sf T}}}P_{k}(w^{o})\right\|^{2},$
$\displaystyle\stackrel{{\scriptstyle(c)}}{{\leq}}\frac{3\delta^{2}}{E_{k}B_{k}}\left(1+\frac{1}{N_{k}^{2}}\sum_{n=1}^{N_{k}}\frac{1}{p_{n}^{(k)}}\right)\|\widetilde{\bm{w}}_{i-1}\|^{2}$
$\displaystyle+\frac{6}{E_{k}B_{k}}\sum_{n=1}^{N_{k}}\frac{1}{N_{k}^{2}p_{n}^{(k)}}\|\nabla_{w^{{\sf
T}}}Q_{k}(w^{o};x_{k,n})\|^{2}$
$\displaystyle+\frac{6}{E_{k}B_{k}}\|\nabla_{w^{{\sf T}}}P_{k}(w^{o})\|^{2},$
$\displaystyle=\beta_{s,k}^{2}\|\widetilde{\bm{w}}_{i-1}\|^{2}+\sigma_{s,k}^{2}+\frac{6}{E_{k}B_{k}}\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\|^{2},$ (85)

The cross-term reduces to 0 by first conditionig over
$\mathbb{I}_{n_{1}}=1,\mathbb{I}_{n_{2}}=1$ and then splittinng the
expectation. Each of the two terms are zero. Thus, putting everything
together, we get:

$\displaystyle\sigma_{k}^{2}\leq\beta_{s,k}^{2}\|\widetilde{\bm{w}}_{i-1}\|^{2}+\sigma_{s,k}^{2}+\frac{6}{E_{k}B_{k}}\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\|^{2}.$ (86)

Next, to bound the second order moment of the gradient noise, we use (76):

$\displaystyle\mathbb{E}\left\\{\|\bm{s}_{i}\|^{2}|\bm{w}_{i-1}\right\\}\leq\sum_{k=1}^{K}p_{k}\Bigg{(}\frac{1}{K^{2}p_{k}^{2}}\sigma_{k}^{2}$
$\displaystyle+\bigg{\|}\frac{1}{Kp_{k}}\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})-\frac{1}{K}\sum_{\ell=1}^{K}\nabla_{w^{{\sf
T}}}P_{\ell}(\bm{w}_{i-1})\bigg{\|}^{2}\Bigg{)}.$ (87)

The first term is of the same form as for sampling with replacement, and thus
can be bounded similarly:

$\displaystyle\mathbb{E}\\{\|\bm{s}_{i}\|^{2}|\bm{w}_{i-1}\\}\leq\sum_{k=1}^{K}p_{k}\Bigg{\\{}\frac{\beta_{s,k}^{2}}{K^{2}p_{k}^{2}}\|\widetilde{\bm{w}}_{i-1}\|^{2}+\frac{1}{K^{2}p_{k}^{2}}\sigma_{s,k}^{2}$
$\displaystyle+\frac{6}{K^{2}p_{k}^{2}E_{k}B_{k}}\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\|^{2}+\frac{3}{K^{2}p_{k}^{2}}\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})\|^{2}$
$\displaystyle+3\delta^{2}\left(1+\frac{1}{K^{2}p_{k}^{2}}\right)\|\widetilde{\bm{w}}_{i-1}\|^{2}\Bigg{\\}}$
$\displaystyle=\beta_{s}^{2}\|\widetilde{\bm{w}}_{i-1}\|^{2}+\sigma_{s}^{2}.$
(88)

∎

##  Appendix C Proof of Lemma 2

###### Proof.

We first note the following result by using
$\frac{1}{K}\sum_{k=1}^{K}\nabla_{w^{{\sf T}}}P_{k}(w^{o})=0$:

$\displaystyle\left\|\widetilde{\bm{w}}_{i-1}+\mu\frac{1}{K}\sum_{k=1}^{K}\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\right\|^{2}$
$\displaystyle=\|\widetilde{\bm{w}}_{i-1}\|^{2}+\mu^{2}\left\|\frac{1}{K}\sum_{k=1}^{K}\nabla_{w^{{\sf
T}}}P_{k}(w^{o})-\nabla_{w^{{\sf T}}}P_{k}(\bm{w}_{i-1})\right\|^{2}$
$\displaystyle\quad+2\mu\widetilde{\bm{w}}_{i-1}^{\sf
T}\frac{1}{K}\sum_{k=1}^{K}\nabla_{w^{{\sf T}}}P_{k}(\bm{w}_{i-1}),$
$\displaystyle\stackrel{{\scriptstyle(a)}}{{\leq}}\|\widetilde{\bm{w}}_{i-1}\|^{2}+\mu^{2}\frac{1}{K}\sum_{k=1}^{K}\left\|\nabla_{w^{{\sf
T}}}P_{k}(w^{o})-\nabla_{w^{{\sf T}}}P_{k}(\bm{w}_{i-1})\right\|^{2}$
$\displaystyle\quad+2\mu\widetilde{\bm{w}}_{i-1}^{\sf
T}\frac{1}{K}\sum_{k=1}^{K}\nabla_{w^{{\sf T}}}P_{k}(\bm{w}_{i-1}),$
$\displaystyle\stackrel{{\scriptstyle(b)}}{{\leq}}(1+\mu^{2}\delta^{2})\|\widetilde{\bm{w}}_{i-1}\|^{2}+2\mu\widetilde{\bm{w}}_{i-1}^{\sf
T}\frac{1}{K}\sum_{k=1}^{K}\nabla_{w^{{\sf T}}}P_{k}(\bm{w}_{i-1}),$
$\displaystyle\stackrel{{\scriptstyle(c)}}{{\leq}}(1+\mu^{2}\delta^{2})\|\widetilde{\bm{w}}_{i-1}\|^{2}$
$\displaystyle\quad+2\mu\frac{1}{K}\sum_{k=1}^{K}\left(P_{k}(w^{o})-P_{k}(\bm{w}_{i-1})-\frac{\nu}{2}\|\widetilde{\bm{w}}_{i-1}\|^{2}\right),$
$\displaystyle\stackrel{{\scriptstyle(d)}}{{\leq}}(1+\mu^{2}\delta^{2})\|\widetilde{\bm{w}}_{i-1}\|^{2}-2\mu\frac{1}{K}\sum_{k=1}^{K}\nu\|\widetilde{\bm{w}}_{i-1}\|^{2},$
$\displaystyle=(1-2\mu\nu+\mu^{2}\delta^{2})\|\widetilde{\bm{w}}_{i-1}\|^{2},$
(89)

where $(a)$ follows from Jensen’s inequality, $(b)$ from the Lipschitz
condition, and $(c)$ and $(d)$ from strong convexity condition.

Returning to the main expression:

$\widetilde{\bm{w}}_{i-1}+\mu\frac{1}{K}\sum_{k=1}^{K}\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})+\mu\bm{s}_{i},$ (90)

and taking conditional expectations, we obtain:

$\displaystyle\mathbb{E}\left\\{\left\|\widetilde{\bm{w}}_{i-1}+\mu\frac{1}{K}\sum_{k=1}^{K}\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})+\mu\bm{s}_{i}\right\|^{2}\Bigg{|}\bm{w}_{i-1}\right\\},$
$\displaystyle\stackrel{{\scriptstyle(a)}}{{=}}\mathbb{E}\left\\{\left\|\widetilde{\bm{w}}_{i-1}+\mu\frac{1}{K}\sum_{k=1}^{K}\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})\right\|^{2}\big{|}\bm{w}_{i-1}\right\\}$
$\displaystyle\quad+\mu^{2}\mathbb{E}\left\\{\|\bm{s}_{i}\|^{2}\big{|}\bm{w}_{i-1}\right\\},$
$\displaystyle\stackrel{{\scriptstyle(b)}}{{\leq}}(1-2\mu\nu+\mu^{2}\delta^{2})\|\widetilde{\bm{w}}_{i-1}\|^{2}$
$\displaystyle\quad+\mu^{2}\left(\beta_{s}^{2}\|\widetilde{\bm{w}}_{i-1}\|^{2}+\eta_{s}\|\widetilde{\bm{w}}_{i-1}\|+\sigma_{s}^{2}\right),$
(91)

where the cross-term in $(a)$ is zero because of the zero mean property of the
gradient noise, and $(b)$ follows from (C) and using the bound on the second
order moment of the gradient noise.

Next, taking expectation again to remove the conditioning we get:

$\displaystyle\mathbb{E}\left\|\widetilde{\bm{w}}_{i-1}+\mu\frac{1}{K}\sum_{k=1}^{K}\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{i-1})+\mu\bm{s}_{i}\right\|^{2}$
$\displaystyle\leq\left(1-2\mu\nu+\mu^{2}(\delta^{2}+\beta_{s}^{2})\right)\mathbb{E}\|\widetilde{\bm{w}}_{i-1}\|^{2}+\mu^{2}\eta_{s}\mathbb{E}\|\widetilde{\bm{w}}_{i-1}\|$
$\displaystyle\quad+\mu^{2}\sigma_{s}^{2}.$ (92)

∎

##  Appendix D Proof of Lemma 3

###### Proof.

To show the mean is zero, it is enough to calculate the mean of the
approximate gradient. We start with the sampling with replacement scheme where
the samples are chosen independently from each other:

$\displaystyle\mathbb{E}\left\\{\frac{1}{B_{k}}\sum_{b\in\mathcal{B}_{k,i,e}}\frac{1}{N_{k}p_{b}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,b})\big{|}\mathcal{F}_{e-1},\mathcal{L}_{i}\right\\}$
$\displaystyle=\frac{1}{B_{k}}\sum_{b\in\mathcal{B}_{k,i,e}}\mathbb{E}\left\\{\frac{1}{N_{k}p_{b}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,b})\big{|}\mathcal{F}_{e-1},\mathcal{L}_{i}\right\\},$
$\displaystyle=\frac{1}{B_{k}}\sum_{b\in\mathcal{B}_{k,i,e}}\sum_{n=1}^{N_{k}}\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,n}),$
$\displaystyle=\frac{1}{N_{k}}\sum_{n=1}^{N_{k}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,n}).$ (93)

As for the sampling without replacement scheme, since the samples are now
dependent, we introduce the indicator function $\mathbb{I}_{n}$ and the
derivation goes as follows:

$\displaystyle\mathbb{E}\left\\{\frac{1}{B_{k}}\sum_{b\in\mathcal{B}_{k,i,e}}\frac{1}{N_{k}p_{b}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,b})\big{|}\mathcal{F}_{e-1},\mathcal{L}_{i}\right\\}$
$\displaystyle=\mathbb{E}\left\\{\frac{1}{B_{k}}\sum_{n=1}^{N_{k}}\frac{\mathbb{I}_{n}}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,n})\big{|}\mathcal{F}_{e-1},\mathcal{L}_{i}\right\\},$
$\displaystyle=\frac{1}{B_{k}}\sum_{n=1}^{N_{k}}\frac{\mathbb{P}(\mathbb{I}_{n}=1)}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,n}),$
$\displaystyle=\frac{1}{N_{k}}\sum_{n=1}^{N_{k}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,n}).$ (94)

Next, to bound the second order moment, we start with an intermediate step and
bound the second order moment of the individual gradient noise of one sample.
The derivation below holds regardless of the sampling scheme. By adding and
subtracting $\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w_{k}^{o};\bm{x}_{k,n})$, adding $\nabla_{w^{{\sf
T}}}P_{k}(w_{k}^{o})=0$, and then using Jensen’s inequality and Lipschitz
condition, we get:

$\displaystyle\left\|\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,n})-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{k,e-1})\right\|^{2}$ $\displaystyle\leq
3\Bigg{\|}\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,n})$
$\displaystyle\quad-\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w_{k}^{o};\bm{x}_{k,n})\Bigg{\|}^{2}$
$\displaystyle\quad+3\left\|\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(w_{k}^{o};\bm{x}_{k,n})\right\|^{2}+3\delta^{2}\|\widetilde{\bm{w}}_{k,e-1}\|^{2}.$
(95)

Then, taking the conditional expectation and using the Lipschitz property, we
get:

$\displaystyle\mathbb{E}\Bigg{\\{}\Bigg{\|}\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,n})$ $\displaystyle\quad-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{k,e-1})\Bigg{\|}^{2}\Bigg{|}\mathcal{F}_{e-1},\mathcal{L}_{i}\Bigg{\\}}$
$\displaystyle\leq\sum_{n=1}^{N_{k}}\frac{3p_{n}^{(k)}}{N_{k}^{2}\left(p_{n}^{(k)}\right)^{2}}\big{(}\|\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};x_{k,n})$ $\displaystyle\quad-\nabla_{w^{{\sf
T}}}Q_{k}(w_{k}^{o};x_{k,n})\|^{2}+\left\|\nabla_{w^{{\sf
T}}}Q_{k}(w_{k}^{o};x_{k,n})\right\|^{2}\big{)}$
$\displaystyle\quad+3\delta^{2}\|\widetilde{\bm{w}}_{k,e-1}\|^{2},$
$\displaystyle\leq\sum_{n=1}^{N_{k}}\frac{3}{N_{k}^{2}p_{n}^{(k)}}\left(\delta^{2}\|\widetilde{\bm{w}}_{k,e-1}\|^{2}+\left\|\nabla_{w^{{\sf
T}}}Q_{k}(w_{k}^{o};x_{k,n})\right\|^{2}\right)$
$\displaystyle\quad+3\delta^{2}\|\widetilde{\bm{w}}_{k,e-1}\|^{2}.$ (96)

Now going back to calculating the second order moment of the local incremental
gradient noise, we first start with the sampling with replacement. Using the
fact that the samples are independent we get:

$\displaystyle\mathbb{E}\left\\{\|\bm{q}_{k,i,e}\|^{2}\big{|}\mathcal{F}_{e-1},\mathcal{L}_{i}\right\\}$
$\displaystyle=\frac{1}{K^{2}p_{k}^{2}B_{k}^{2}}\sum_{b\in\mathcal{B}_{k,i,e}}\mathbb{E}\Bigg{\\{}\Bigg{\|}\frac{1}{N_{k}p_{b}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,b})$ $\displaystyle\quad-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{k,e-1})\Bigg{\|}^{2}\Bigg{|}\mathcal{F}_{e-1},\mathcal{L}_{i}\Bigg{\\}}$
$\displaystyle\leq\frac{3\delta^{2}}{K^{2}p_{k}^{2}B_{k}}\left(1+\frac{1}{N_{k}^{2}}\sum_{n=1}^{N_{k}}\frac{1}{p_{n}^{(k)}}\right)\|\widetilde{\bm{w}}_{k,e-1}\|^{2}$
$\displaystyle+\frac{3}{K^{2}p_{k}^{2}B_{k}N_{k}^{2}}\sum_{n=1}^{N_{k}}\frac{1}{p_{n}^{(k)}}\|\nabla_{w^{{\sf
T}}}Q_{k}(w_{k}^{o};x_{k,n})\|^{2}.$ (97)

As for the sampling without replacement, we also introduce the indicator
function and write out the square of sums. The cross-terms disappear since
each term has zero mean. The derivation then follows similarly to that of the
sampling with replacement. More formally:

$\displaystyle\mathbb{E}\left\\{\|\bm{q}_{k,i,e}\|^{2}\big{|}\mathcal{F}_{e-1},\mathcal{L}_{i}\right\\}$
$\displaystyle=\frac{1}{K^{2}p_{k}^{2}B_{k}^{2}}\sum_{n=1}^{N_{k}}\mathbb{P}(\mathbb{I}_{n}=1)\mathbb{E}\Bigg{\\{}\Bigg{\|}\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};\bm{x}_{k,n})$ $\displaystyle\quad-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{k,e-1})\Bigg{\|}^{2}\Bigg{|}\mathbb{I}_{n}=1,\mathcal{F}_{e-1},\mathcal{L}_{i}\Bigg{\\}}$
$\displaystyle=\frac{1}{K^{2}p_{k}^{2}B_{k}}\sum_{n=1}^{N_{k}}p_{n}^{(k)}\Bigg{\|}\frac{1}{N_{k}p_{n}^{(k)}}\nabla_{w^{{\sf
T}}}Q_{k}(\bm{w}_{k,e-1};x_{k,n})$ $\displaystyle\quad-\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{k,e-1})\Bigg{\|}^{2}$
$\displaystyle\leq\frac{3\delta^{2}}{K^{2}p_{k}^{2}B_{k}}\left(1+\frac{1}{N_{k}^{2}}\sum_{n=1}^{N_{k}}\frac{1}{p_{n}^{(k)}}\right)\|\widetilde{\bm{w}}_{k,e-1}\|^{2}$
$\displaystyle+\frac{3}{K^{2}p_{k}^{2}B_{k}N_{k}^{2}}\sum_{n=1}^{N_{k}}\frac{1}{p_{n}^{(k)}}\|\nabla_{w^{{\sf
T}}}Q_{k}(w_{k}^{o};x_{k,n})\|^{2}.$ (98)

∎

##  Appendix E Proof of Lemma 4

###### Proof.

We subtract $w_{k}^{o}$ from both sides of (8) and use (III-B) to get:

$\displaystyle\widetilde{\bm{w}}_{k,e}$
$\displaystyle=\widetilde{\bm{w}}_{k,e-1}+\mu\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{k,e-1})+\mu\bm{q}_{k,i,e}.$ (99)

We bound the first two terms and use the fact that $\nabla_{w^{{\sf
T}}}P_{k}(w_{k}^{o})=0$, Lipschitz condition, and the convexity of the cost
function:

$\displaystyle\left\|\widetilde{\bm{w}}_{k,e-1}+\mu\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{k,e-1})\right\|^{2}$
$\displaystyle\quad+\mu^{2}\|\nabla_{w^{{\sf T}}}P_{k}(\bm{w}_{k,e-1})\|^{2},$
$\displaystyle=\|\widetilde{\bm{w}}_{k,e-1}\|^{2}+2\mu\widetilde{\bm{w}}_{k,e-1}^{\sf
T}\nabla_{w^{{\sf T}}}P_{k}(\bm{w}_{k,e-1})$
$\displaystyle\quad+\mu^{2}\|\nabla_{w^{{\sf
T}}}P_{k}(w_{k}^{o})-\nabla_{w^{{\sf T}}}P_{k}(\bm{w}_{k,e-1})\|^{2},$
$\displaystyle\leq(1+\mu^{2}\delta^{2})\|\widetilde{\bm{w}}_{k,e-1}\|^{2}+2\mu\widetilde{\bm{w}}_{k,e-1}^{\sf
T}\nabla_{w^{{\sf T}}}P_{k}(\bm{w}_{k,e-1}),$
$\displaystyle\leq(1-2\nu\mu+\mu^{2}\delta^{2})\|\widetilde{\bm{w}}_{k,e-1}\|^{2}.$
(100)

Returning to (99), squaring both sides, conditioning on the filtration
$\mathcal{F}_{e-1}$, and taking expectations we obtain:

$\displaystyle\mathbb{E}\left\\{\|\widetilde{\bm{w}}_{k,e}\|^{2}\big{|}\mathcal{F}_{e-1}\right\\}$
$\displaystyle\stackrel{{\scriptstyle(a)}}{{=}}\mathbb{E}\left\\{\|\widetilde{\bm{w}}_{k,e-1}+\mu\nabla_{w^{{\sf
T}}}P_{k}(\bm{w}_{k,e-1})\|^{2}\big{|}\mathcal{F}_{e-1}\right\\}$
$\displaystyle\quad+\mu^{2}\mathbb{E}\left\\{\|\bm{q}_{k,i,e}^{2}\|^{2}\big{|}\mathcal{F}_{e-1}\right\\},$
$\displaystyle\leq\left(1-2\nu\mu+\mu^{2}\left(\delta^{2}+\frac{E_{k}}{K^{2}p_{k}^{2}}\beta^{2}_{s,k}\right)\right)\|\widetilde{\bm{w}}_{k,e-1}\|^{2}$
$\displaystyle\quad+\mu^{2}\frac{1}{K^{2}p_{k}^{2}}\sigma_{q,k}^{2},$ (101)

where the cross term in $(a)$ is zero because of the zero mean property of the
local incremental gradient noise. Taking expectations on both sides again
removes the condition on the filtration and leads to the desired result. By
further iterating recursion (32) we obtain:

$\mathbb{E}\|\widetilde{\bm{w}}_{k,e}\|^{2}\leq\lambda_{k}^{e}\mathbb{E}\|\widetilde{\bm{w}}_{k,0}\|^{2}+\frac{1-\lambda_{k}^{e}}{1-\lambda_{k}}\mu^{2}\sigma_{q,k}^{2}.$
(102)

∎

##  Appendix F Proof of Lemma 5

###### Proof.

First, using Jensen’s inequality $(a)$ and Lipschitz continuity $(b)$, we
obtain:

$\displaystyle\|\bm{q}_{i}\|^{2}$
$\displaystyle\stackrel{{\scriptstyle(a)}}{{\leq}}\frac{1}{L}\sum_{\ell\in\mathcal{L}_{i}}\frac{1}{K^{2}p_{\ell}^{2}E_{\ell}B_{\ell}}\sum_{e=1}^{E_{\ell}}\sum_{b\in\mathcal{B}_{\ell,i,e}}\frac{1}{N^{2}_{\ell}\left(p_{b}^{(\ell)}\right)^{2}}$
$\displaystyle\quad\times\left\|\nabla_{w^{{\sf
T}}}Q_{\ell}(\bm{w}_{\ell,e-1};\bm{x}_{\ell,b})-\nabla_{w^{{\sf
T}}}Q_{\ell}(\bm{w}_{i-1};\bm{x}_{\ell,b})\right\|^{2},$
$\displaystyle\stackrel{{\scriptstyle(b)}}{{\leq}}\frac{\delta^{2}}{L}\sum_{\ell\in\mathcal{L}_{i}}\frac{1}{K^{2}p_{\ell}^{2}E_{\ell}B_{\ell}}\sum_{e=1}^{E_{\ell}}\sum_{b\in\mathcal{B}_{\ell,i,e}}\frac{\|\bm{w}_{i-1}-\bm{w}_{\ell,e-1}\|^{2}}{N_{\ell}^{2}\left(p_{b}^{(\ell)}\right)^{2}}.$
(103)

Next, we focus on $\|\bm{w}_{i-1}-\bm{w}_{\ell,e-1}\|^{2}$, and by applying
Jensen’s inequality in $(a)$ and $(b)$ and Lipschitz condition in $(c)$ we
obtain:

$\displaystyle\|\bm{w}_{i-1}-\bm{w}_{\ell,e-1}\|^{2}$
$\displaystyle=\mu^{2}\left\|\frac{1}{E_{\ell}B_{\ell}}\sum_{f=0}^{e-2}\sum_{b\in\mathcal{B}_{\ell,i,f}}\frac{1}{N_{\ell}p_{b}^{(\ell)}}\nabla_{w^{{\sf
T}}}Q_{\ell}(\bm{w}_{\ell,f};\bm{x}_{\ell,b})\right\|^{2},$
$\displaystyle\stackrel{{\scriptstyle(a)}}{{\leq}}\frac{\mu^{2}}{E_{\ell}B_{\ell}}\sum_{f=0}^{e-2}\sum_{b\in\mathcal{B}_{\ell,i,f}}\frac{1}{N_{\ell}^{2}\left(p_{b}^{(\ell)}\right)^{2}}\|\nabla_{w^{{\sf
T}}}Q_{\ell}(\bm{w}_{\ell,f};\bm{x}_{\ell,b})$
$\displaystyle\quad-\nabla_{w^{{\sf
T}}}Q_{\ell}(w_{\ell}^{o};\bm{x}_{\ell,b})+\nabla_{w^{{\sf
T}}}Q_{\ell}(w_{\ell}^{o};\bm{x}_{\ell,b})\|^{2},$
$\displaystyle\stackrel{{\scriptstyle(b)}}{{\leq}}\frac{2\mu^{2}}{E_{\ell}B_{\ell}}\sum_{f=0}^{e-2}\sum_{b\in\mathcal{B}_{\ell,i,f}}\frac{1}{N_{\ell}^{2}\left(p_{b}^{(\ell)}\right)^{2}}\big{(}\|\nabla_{w^{{\sf
T}}}Q_{\ell}(\bm{w}_{\ell,f};\bm{x}_{\ell,b})$
$\displaystyle\quad-\nabla_{w^{{\sf
T}}}Q_{\ell}(w_{\ell}^{o};\bm{x}_{\ell,b})\|^{2}+\|\nabla_{w^{{\sf
T}}}Q_{\ell}(w_{\ell}^{o};\bm{x}_{\ell,b})\|^{2}\big{)},$
$\displaystyle\stackrel{{\scriptstyle(c)}}{{\leq}}\frac{2\mu^{2}}{E_{\ell}B_{\ell}}\sum_{f=0}^{e-2}\sum_{b\in\mathcal{B}_{\ell,i,f}}\frac{1}{N_{\ell}^{2}\left(p_{b}^{(\ell)}\right)^{2}}\Big{(}\delta^{2}\|\widetilde{\bm{w}}_{\ell,f}\|^{2}$
$\displaystyle\quad+\|\nabla_{w^{{\sf
T}}}Q_{\ell}(w_{\ell}^{o};\bm{x}_{\ell,b})\|^{2}\Big{)}.$

Then, taking the expectation given the previous filtration $\mathcal{F}_{e-2}$
and the participating agents $\mathcal{L}_{i}$, we see that:

$\displaystyle\mathbb{E}\left\\{\|\bm{w}_{i-1}-\bm{w}_{\ell,e-1}\|^{2}\big{|}\mathcal{F}_{e-2},\mathcal{L}_{i}\right\\}$
$\displaystyle\leq\frac{2\mu^{2}\delta^{2}}{E_{\ell}B_{\ell}}\sum_{f=0}^{e-2}\|\widetilde{\bm{w}}_{\ell,f}\|^{2}\mathbb{E}\left\\{\sum_{b\in\mathcal{B}_{\ell,i,f}}\frac{1}{N_{\ell}^{2}\left(p_{b}^{(\ell)}\right)^{2}}\Bigg{|}\mathcal{F}_{e-2},\mathcal{L}_{i}\right\\}$
$\displaystyle\quad+\frac{2\mu^{2}}{E_{\ell}B_{\ell}}\sum_{f=0}^{e-2}\mathbb{E}\left\\{\sum_{b\in\mathcal{B}_{\ell,i,f}}\frac{\|\nabla_{w^{{\sf
T}}}Q_{\ell}(w_{\ell}^{o};\bm{x}_{\ell,b})\|^{2}}{N_{\ell}^{2}\left(p_{b}^{(\ell)}\right)^{2}}\Bigg{|}\mathcal{F}_{e-2},\mathcal{L}_{i}\right\\}$
(104)
$\displaystyle=\frac{2\mu^{2}\delta^{2}}{E_{\ell}B_{\ell}}\sum_{f=0}^{e-2}\|\widetilde{\bm{w}}_{\ell,f}\|^{2}\sum_{n=1}^{N_{\ell}}\frac{\mathbb{P}(\mathbb{I}_{n}=1)}{N_{\ell}^{2}\left(p_{n}^{(\ell)}\right)^{2}}$
$\displaystyle\quad+\frac{2\mu^{2}}{E_{\ell}B_{\ell}}\sum_{f=0}^{e-2}\sum_{n=1}^{N_{\ell}}\frac{\mathbb{P}(\mathbb{I}_{n}=1)}{N_{\ell}^{2}\left(p_{n}^{(\ell)}\right)^{2}}\|\nabla_{w^{{\sf
T}}}Q_{\ell}(w_{\ell}^{o};x_{\ell,n})\|^{2},$
$\displaystyle=\frac{2\mu^{2}\delta^{2}}{E_{\ell}}\sum_{f=0}^{e-2}\|\widetilde{\bm{w}}_{\ell,f}\|^{2}\sum_{n=1}^{N_{\ell}}\frac{1}{N_{\ell}^{2}p_{n}^{(\ell)}}+\frac{2\mu^{2}(e-1)}{3E_{\ell}}\sigma_{q,\ell}^{2}.$
(105)

Then, taking expectation again over the filtration, we obtain:

$\displaystyle\mathbb{E}\left\\{\|\bm{w}_{i-1}-\bm{w}_{\ell,e-1}\|^{2}\big{|}\mathcal{L}_{i}\right\\}$
$\displaystyle\leq\frac{2\mu^{2}\delta^{2}}{E_{\ell}}\sum_{f=0}^{e-1}\mathbb{E}\left\\{\|\widetilde{\bm{w}}_{\ell,f}\|^{2}\big{|}\mathcal{L}_{i}\right\\}\sum_{n=1}^{N_{\ell}}\frac{1}{N_{\ell}^{2}p_{n}^{(\ell)}}+\frac{2\mu^{2}(e-1)}{3E_{\ell}}\sigma_{q,\ell}^{2},$
$\displaystyle\stackrel{{\scriptstyle(a)}}{{\leq}}\frac{2\mu^{2}\delta^{2}}{E_{\ell}}\sum_{n=1}^{N_{\ell}}\frac{1}{N_{\ell}^{2}p_{n}^{(\ell)}}\sum_{f=0}^{e-1}\Bigg{(}\lambda_{\ell}^{f}\mathbb{E}\left\\{\|\widetilde{\bm{w}}_{\ell,0}\|^{2}\big{|}\mathcal{L}_{i}\right\\}$
$\displaystyle\quad+\frac{\mu^{2}}{K^{2}p_{\ell}^{2}}\frac{1-\lambda_{\ell}^{f}}{1-\lambda_{\ell}}\sigma_{q,\ell}^{2}\Bigg{)}+\frac{2\mu^{2}(e-1)}{3E_{\ell}}\sigma_{q,\ell}^{2},$
$\displaystyle=\frac{2\mu^{2}\delta^{2}}{E_{\ell}}\sum_{n=1}^{N_{\ell}}\frac{1}{N_{\ell}^{2}p_{n}^{(\ell)}}\Bigg{(}\frac{1-\lambda_{\ell}^{e}}{1-\lambda_{\ell}}\mathbb{E}\left\\{\|\widetilde{\bm{w}}_{\ell,0}\|^{2}\big{|}\mathcal{L}_{i}\right\\}$
$\displaystyle\quad+\frac{\mu^{2}}{K^{2}p_{k}^{2}}\frac{e(1-\lambda_{\ell})-1+\lambda_{\ell}^{e}}{(1-\lambda_{\ell})^{2}}\sigma_{q,\ell}^{2}\Bigg{)}+\frac{2\mu^{2}(e-1)}{3E_{\ell}}\sigma_{q,\ell}^{2},$
$\displaystyle\stackrel{{\scriptstyle(b)}}{{\leq}}\frac{2\mu^{2}\delta^{2}}{E_{\ell}}\sum_{n=1}^{N_{\ell}}\frac{1}{N_{\ell}^{2}p_{n}^{(\ell)}}\Bigg{(}2\frac{1-\lambda_{\ell}^{e}}{1-\lambda_{\ell}}\mathbb{E}\left\\{\|\widetilde{\bm{w}}_{i-1}\|^{2}\big{|}\mathcal{L}_{i}\right\\}$
$\displaystyle\quad+2\frac{1-\lambda_{\ell}^{e}}{1-\lambda_{\ell}}\mathbb{E}\left\\{\|w^{o}-w_{\ell}^{o}\|^{2}\big{|}\mathcal{L}_{i}\right\\}$
$\displaystyle\quad+\frac{\mu^{2}}{K^{2}p_{k}^{2}}\frac{e(1-\lambda_{\ell})-1+\lambda_{\ell}^{e}}{(1-\lambda_{\ell})^{2}}\sigma_{q,\ell}^{2}\Bigg{)}+\frac{2\mu^{2}(e-1)}{3E_{\ell}}\sigma_{q,\ell}^{2},$
(106)

where we used Lemma 4 in $(a)$, and in $(b)$ we added and subtracted $w^{o}$
and used Jensen’s inequality. Then, summing over $e$ results in:

$\displaystyle\frac{1}{E_{\ell}}\sum_{e=1}^{E_{\ell}}\mathbb{E}\left\\{\|\bm{w}_{i-1}-\bm{w}_{\ell,e-1}\|^{2}\big{|}\mathcal{L}_{i}\right\\}$
$\displaystyle\leq\frac{2\mu^{2}\delta^{2}}{E_{\ell}^{2}}\sum_{n=1}^{N_{\ell}}\frac{1}{N_{\ell}^{2}p_{n}^{(\ell)}}\Bigg{(}2\frac{(E_{\ell}+1)(1-\lambda_{\ell})-1+\lambda_{\ell}^{E_{\ell}+1}}{(1-\lambda_{\ell})^{2}}$
$\displaystyle\quad\times\left(\mathbb{E}\left\\{\|\widetilde{\bm{w}}_{i-1}\|^{2}\big{|}\mathcal{L}_{i}\right\\}+\mathbb{E}\left\\{\|w^{o}-w_{\ell}^{o}\|^{2}\big{|}\mathcal{L}_{i}\right\\}\right)$
$\displaystyle\quad+\frac{E_{\ell}(E_{\ell}+1)(1-\lambda_{\ell})^{2}-2E_{\ell}(1-\lambda_{\ell})+2\lambda_{\ell}-2\lambda_{\ell}^{E_{\ell}+1}}{(1-\lambda_{\ell})^{3}}$
$\displaystyle\quad\times\frac{\mu^{2}}{K^{2}p_{k}^{2}}\sigma_{q,\ell}^{2}\Bigg{)}+\frac{E_{\ell}(E_{\ell}-1)\mu^{2}}{3E_{\ell}^{2}}\sigma_{q,\ell}^{2}.$
(107)

Taking the expectation of (F) given the choice of the agents and plugging the
above expression, we get:

$\displaystyle\mathbb{E}\left\\{\|\bm{q}_{i}\|^{2}\big{|}\mathcal{L}_{i}\right\\}$
$\displaystyle\leq\frac{\delta^{2}}{L}\sum_{\ell\in\mathcal{L}_{i}}\frac{1}{K^{2}p_{\ell}^{2}}\Big{(}a\mu^{2}\mathbb{E}\left\\{\|\widetilde{\bm{w}}_{i-1}\|^{2}\big{|}\mathcal{L}_{i}\right\\}$
$\displaystyle\quad+a\mu^{2}\mathbb{E}\left\\{\|w^{o}-w_{\ell}^{o}\|^{2}\big{|}\mathcal{L}_{i}\right\\}$
$\displaystyle\quad+(b\mu^{4}+c\mu^{2})\sigma_{q,\ell}^{2}\Big{)}\frac{1}{B_{\ell}N_{\ell}^{2}}\mathbb{E}\left\\{\sum_{b\in\mathcal{B}_{\ell,i,e}}\frac{1}{\left(p_{b}^{(\ell)}\right)^{2}}\Bigg{|}\mathcal{L}_{i}\right\\},$
$\displaystyle=\frac{\delta^{2}}{L}\sum_{\ell\in\mathcal{L}_{i}}\frac{1}{K^{2}p_{\ell}^{2}}\Big{(}a\mu^{2}\mathbb{E}\left\\{\|\widetilde{\bm{w}}_{i-1}\|^{2}\big{|}\mathcal{L}_{i}\right\\}$
$\displaystyle\quad+a\mu^{2}\mathbb{E}\left\\{\|w^{o}-w_{\ell}^{o}\|^{2}\big{|}\mathcal{L}_{i}\right\\}+(b\mu^{4}+c\mu^{2})\sigma_{q,\ell}^{2}\Big{)}\sum_{n=1}^{N_{\ell}}\frac{1}{N_{\ell}^{2}p_{n}^{(\ell)}},$
(108)

where we introduced constants $a,b,c$ to make the notation simpler. Then,
taking again the expectation to remove the conditioning and using Assumption
2:

$\displaystyle\mathbb{E}\|\bm{q}_{i}\|^{2}$
$\displaystyle\leq\delta^{2}\sum_{k=1}^{K}\frac{1}{K^{2}p_{k}}\sum_{n=1}^{N_{k}}\frac{1}{N_{k}^{2}p_{n}^{(k)}}\Big{(}a\mu^{2}\mathbb{E}\|\widetilde{\bm{w}}_{i-1}\|^{2}$
$\displaystyle\quad+a\mu^{2}\|w^{o}-w_{k}^{o}\|^{2}+(b\mu^{4}+c\mu^{2})\sigma_{q,k}^{2}\Big{)}$
$\displaystyle\leq\delta^{2}\sum_{k=1}^{K}\frac{1}{K^{2}p_{k}}\sum_{n=1}^{N_{k}}\frac{1}{N_{k}^{2}p_{n}^{(k)}}\Big{(}a\mu^{2}\mathbb{E}\|\widetilde{\bm{w}}_{i-1}\|^{2}+a\mu^{2}\xi^{2}$
$\displaystyle\quad+(b\mu^{4}+c\mu^{2})\sigma_{q,k}^{2}\Big{)}.$ (109)

Further simplifying the notation gives us the desired result. Thus, since
$a=O(\mu^{-1})$, $b=O(\mu^{-2})$ and $c=O(1)$, we get
$\mathbb{E}\|\bm{q}_{i}\|^{2}=O(\mu)$. ∎

## References

  * [1] E. Rizk, S. Vlaski, and A. H. Sayed, “Dynamic federated learning,” in _Proc. IEEE SPAWC_ , Atlanta, Georgia, USA, 26–29 May 2020, pp. 1–5.
  * [2] A. Agarwal and J. C. Duchi, “Distributed delayed stochastic optimization,” in _Advances in Neural Information Processing Systems 24_ , J. Shawe-Taylor, R. S. Zemel, P. L. Bartlett, F. Pereira, and K. Q. Weinberger, Eds. Curran Associates, Inc., 2011, pp. 873–881.
  * [3] M. Zinkevich, M. Weimer, L. Li, and A. J. Smola, “Parallelized stochastic gradient descent,” in _Proc. Advances in Neural Information Processing Systems_ , 2010, pp. 2595–2603.
  * [4] D. P. Bertsekas, “A new class of incremental gradient methods for least squares problems,” _SIAM J. Optim_ , vol. 7, pp. 913–926, 1996.
  * [5] A. H. Sayed, “Adaptation, learning, and optimization over networks,” _Foundations and Trends® in Machine Learning_ , vol. 7, no. 4-5, pp. 311–801, 2014.
  * [6] J. Chen and A. H. Sayed, “On the learning behavior of adaptive networks — part i: Transient analysis,” _IEEE Trans. Information Theory_ , vol. 61, no. 6, pp. 3487–3517, June 2015.
  * [7] A. Nedic and A. Ozdaglar, “Distributed subgradient methods for multi-agent optimization,” _IEEE Transactions on Automatic Control_ , vol. 54, no. 1, pp. 48–61, Jan 2009.
  * [8] J. C. Duchi, A. Agarwal, and M. J. Wainwright, “Dual averaging for distributed optimization: Convergence analysis and network scaling,” _IEEE Transactions on Automatic Control_ , vol. 57, no. 3, pp. 592–606, 2012\.
  * [9] H. B. McMahan, E. Moore, D. Ramage, S. Hampson _et al._ , “Communication-efficient learning of deep networks from decentralized data,” _Proc. International Conference on Artificial Intelligence and Statistics_ , vol. 54, pp. 1273–1282, 20–22 April 2017.
  * [10] T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith, “Federated optimization in heterogeneous networks,” in _ICML 2019 Workshop AMTL Workshop_ , Long Beach, CA, USA, June 2019, pp. 1–28.
  * [11] V. Smith, C.-K. Chiang, M. Sanjabi, and A. Talwalkar, “Federated multi-task learning,” in _Proc. Advances in Neural Information Processing Systems_ , Long Beach, California, December 2017, pp. 4424–4434.
  * [12] S. Caldas, J. Konečny, H. B. McMahan, and A. Talwalkar, “Expanding the reach of federated learning by reducing client resource requirements,” _arXiv:1812.07210_ , 2018.
  * [13] F. Sattler, S. Wiedemann, K. Müller, and W. Samek, “Robust and communication-efficient federated learning from non-i.i.d. data,” _IEEE Transactions on Neural Networks and Learning Systems_ , pp. 1–14, November 2019\.
  * [14] J. Konečný, H. B. McMahan, F. X. Yu, P. Richtárik, A. T. Suresh, and D. Bacon, “Federated learning: Strategies for improving communication efficiency,” _arXiv:1610.05492_ , 2016.
  * [15] M. Mohri, G. Sivek, and A. T. Suresh, “Agnostic federated learning,” in _Proc. International Conference on Machine Learning_ , Long Beach, California, Jun 2019, pp. 4615–4625.
  * [16] L. Corinzia and J. M. Buhmann, “Variational federated multi-task learning,” _arXiv:1906.06268_ , 2019.
  * [17] M. Khodak, M.-F. F. Balcan, and A. S. Talwalkar, “Adaptive gradient-based meta-learning methods,” in _Proc. Advances in Neural Information Processing Systems_ , Vancouver, Canada, December 2019, pp. 5915–5926.
  * [18] F. Chen, M. Luo, Z. Dong, Z. Li, and X. He, “Federated meta-learning with fast convergence and efficient communication,” _arXiv:1802.07876_ , 2019.
  * [19] K. Bonawitz, V. Ivanov, B. Kreuter, A. Marcedone, H. B. McMahan, S. Patel, D. Ramage, A. Segal, and K. Seth, “Practical secure aggregation for privacy-preserving machine learning,” in _Proc. ACM SIGSAC Conference on Computer and Communications Security_ , New York, NY, Oct 2017, p. 1175–1191.
  * [20] R. C. Geyer, T. Klein, and M. Nabi, “Differentially private federated learning: A client level perspective,” _arXiv:1712.07557_ , 2017.
  * [21] H. B. McMahan, D. Ramage, K. Talwar, and L. Zhang, “Learning differentially private recurrent language models,” in _Proc. International Conference on Learning Representations_ , Vancouver, BC, Canada, April 2018, pp. 1–14.
  * [22] P. Jiang and G. Agrawal, “A linear speedup analysis of distributed deep learning with sparse and quantized communication,” in _Proc. Advances in Neural Information Processing Systems_ , Montreal, Canada, December 2018, pp. 2525–2536.
  * [23] A. Khaled, K. Mishchenko, and P. Richtárik, “First analysis of local gd on heterogeneous data,” _arXiv:1909.04715_ , 2019.
  * [24] S. U. Stich, “Local SGD converges fast and communicates little,” in _ICLR 2019 - International Conference on Learning Representations_ , New Orleans, USA, May 2019, pp. 1–19.
  * [25] J. Wang and G. Joshi, “Cooperative sgd: A unified framework for the design and analysis of communication-efficient sgd algorithms,” _ArXiv:1808.07576_ , 2018.
  * [26] F. Zhou and G. Cong, “On the convergence properties of a k-step averaging stochastic gradient descent algorithm for nonconvex optimization,” _Proc. International Joint Conference on Artificial Intelligence_ , pp. 3219–3227, July 2018.
  * [27] H. Yu, S. Yang, and S. Zhu, “Parallel restarted sgd with faster convergence and less communication: Demystifying why model averaging works for deep learning,” in _Proc. AAAI Conference on Artificial Intelligence_ , Honolulu, Hawaii, USA, July 2019, pp. 5693–5700.
  * [28] S. Wang, T. Tuor, T. Salonidis, K. K. Leung, C. Makaya, T. He, and K. Chan, “Adaptive federated learning in resource constrained edge computing systems,” _IEEE Journal on Selected Areas in Communications_ , vol. 37, no. 6, pp. 1205–1221, 2019.
  * [29] H. Yu, R. Jin, and S. Yang, “On the linear speedup analysis of communication efficient momentum SGD for distributed non-convex optimization,” in _Proc. International Conference on Machine Learning_ , Long Beach, California, Jun 2019, pp. 7184–7193.
  * [30] C. Xie, S. Koyejo, and I. Gupta, “Asynchronous federated optimization,” _arXiv:1903.03934_ , 2019.
  * [31] X. Li, K. Huang, W. Yang, S. Wang, and Z. Zhang, “On the convergence of fedavg on non-iid data,” in _International Conference on Learning Representations_ , Addis Ababa, Ethiopia, 2020, pp. 1–26.
  * [32] L. Liu, J. Zhang, S. H. Song, and K. B. Letaief, “Client-edge-cloud hierarchical federated learning,” _arXiv:1905.06641_ , 2019.
  * [33] T. Nishio and R. Yonetani, “Client selection for federated learning with heterogeneous resources in mobile edge,” in _Proc. IEEE International Conference on Communications (ICC)_ , Shanghai, China, 2019, pp. 1–7.
  * [34] N. Yoshida, T. Nishio, M. Morikura, K. Yamamoto, and R. Yonetani, “Hybrid-fl for wireless networks: Cooperative learning mechanism using non-iid data,” in _Proc. IEEE International Conference on Communications (ICC)_ , Dublin, Ireland, 2020, pp. 1–7.
  * [35] H. T. Nguyen, V. Sehwag, S. Hosseinalipour, C. G. Brinton, M. Chiang, and H. V. Poor, “Fast-convergent federated learning,” _arXiv:2007.13137_ , 2020.
  * [36] M. Mohri, G. Sivek, and A. T. Suresh, “Agnostic federated learning,” in _Proc. International Con-ference on Machine Learning (ICML)_ , vol. 97, Long Beach, California, USA, 09–15 Jun 2019, pp. 4615–4625.
  * [37] T. Li, M. Sanjabi, A. Beirami, and V. Smith, “Fair resource allocation in federated learning,” in _Proc. International Conference on Learning Representations_ , Addis Ababa, Ethiopia, 2020, pp. 1–27.
  * [38] K. Yuan, B. Ying, S. Vlaski, and A. H. Sayed, “Stochastic gradient descent with finite samples sizes,” in _Proc. International Workshop on Machine Learning for Signal Processing (MLSP)_ , Salerno, Italy, 2016, pp. 1–6.
  * [39] G. Alain, A. Lamb, C. Sankar, A. Courville, and Y. Bengio, “Variance reduction in SGD by distributed importance sampling,” _arXiv:1511.06481_ , 2015.
  * [40] D. Needell, R. Ward, and N. Srebro, “Stochastic gradient descent, weighted sampling, and the randomized kaczmarz algorithm,” in _Proc. Advances in Neural Information Processing Systems_ , Montreal, Canada, 2014, pp. 1017–1025.
  * [41] P. Zhao and T. Zhang, “Stochastic optimization with importance sampling for regularized loss minimization,” in _Proc. International Con-ference on Machine Learning (ICML)_ , Lille, France, 2015, pp. 1355–1363.
  * [42] S. U. Stich, A. Raj, and M. Jaggi, “Safe adaptive importance sampling,” in _Proc. Advances in Neural Information Processing Systems_ , Long Beach, California, USA, 2017, pp. 4381–4391.
  * [43] D. G. Horvitz and D. J. Thompson, “A generalization of sampling without replacement from a finite universe,” _Journal of the American statistical Association_ , vol. 47, no. 260, pp. 663–685, 1952.
  * [44] H. O. Hartley and J. N. K. Rao, “Sampling with unequal probabilities and without replacement,” _Ann. Math. Statist._ , vol. 33, no. 2, pp. 350–374, 06 1962.
  * [45] D. Prokhorov, “Ijcnn 2001 neural network competition,” 2001. [Online]. Available: http://www.csie.ntu.edu.tw/-cj1in/libsvmtools/
  * [46] K. R. W. Brewer and M. Hanif, _Sampling with Unequal Probabilities_. New York: Springer-Verlag, 1983.