# GIFAIR-FL: An Approach for Group and Individual Fairness in Federated
Learning

Xubo Yue  Maher Nouiehed  Raed Al Kontar
Industrial and Operations Engineering
University of Michigan  Ann Arbor

###### Abstract

In this paper we propose GIFAIR-FL: an approach that imposes group and
individual fairness to federated learning settings. By adding a regularization
term, our algorithm penalizes the spread in the loss of client groups to drive
the optimizer to fair solutions. Theoretically, we show convergence in non-
convex and strongly convex settings. Our convergence guarantees hold for both
$i.i.d.$ and non-$i.i.d.$ data. To demonstrate the empirical performance of
our algorithm, we apply our method on image classification and text prediction
tasks. Compared to existing algorithms, our method shows improved fairness
results while retaining superior or similar prediction accuracy.

##  1 Introduction

Federated learning (FL) is a machine learning approach that allows
decentralized training of models without the need to access local private data
from a central location. By virtue of FL, devices can collaboratively extract
knowledge and build models while keeping their personal data stored locally.
Over the last few years, tremendous work have been proposed to improve the
performance of FL algorithms; be it speeding up FL algorithms to enable faster
convergence (karimireddy2020scaffold; yuan2020federated), tackling
heterogeneous data both in size and distribution (zhao2018federated;
li2018federated; sattler2019robust; ghosh2019robust; li2019fedmd), improving
the parameter aggregation strategies at the central server
(wang2020federated), designing personalized FL algorithms
(fallah2020personalized), protecting federated systems from adversarial
attacks (bhagoji2019analyzing; wang2020attack), and promoting fairness
(li2019fair). Among those advances, fairness is a critical yet under-
investigated area.

In the training phase of FL algorithms, devices with few data, limited
bandwidth/memory, or unreliable connection may not be favored by conventional
FL algorithms. Such devices can potentially incur higher error rates. This
vicious circle often continues and causes those devices to relinquish their
opportunities for participating in future training rounds. Besides this
aforementioned notion of individual fairness, group fairness also deserves
attention in FL. As FL penetrates practical applications, it is important to
achieve fair performance across groups of clients characterized by their
gender, ethnicity, socio-economic status, geographic location, etc. Despite
the importance of this notion of group fairness, unfortunately, very limited
work exist along this line in FL.

#### Contribution:

We propose an algorithm, GIFAIR-FL, that aims for fairness in FL. GIFAIR-FL
resorts to regularization techniques by penalizing the spread in the loss of
clients to drive the optimizer to fair solutions. We show that our regularized
formulation can be viewed as a dynamic client re-weighting technique that
adaptively gives higher weights to low-performing individuals or groups. Our
proposed method adapts the client weights at every communication round
accordingly. One key feature of GIFAIR-FL is that it can handle both group-
level and individual-level fairness. We then prove that, under some
conditions, our algorithm will converge to an optimal solution for strongly
convex function and to a stationary solution for non-convex function under
non-$\bm{i.i.d}$. settings. Through empirical results on image classification
and text prediction datasets, we demonstrate that GIFAIR-FL can promote
fairness while achieving superior or similar prediction accuracy relative to
recent state-of-the-art fair FL algorithms. Besides that, GIFAIR-FL can be
easily plugged into other FL algorithms for different purposes including
personalized FL.

The rest of the paper is organized as follows. In Sec. 2, we introduce
important notations/definitions and briefly review FL. In Sec. 3, we present
GIFAIR-FL; a simple yet effective approach for fairness in FL. We then provide
convergence guarantees for GIFAIR-FL in Sec. 4. Related work is highlighted in
Sec. 5. Experiments on image classification dataset and text prediction task
are then presented in Sec. 6. Finally, Sec. 7 concludes the paper with a brief
discussion.

##  2 Background

Notations: Suppose there are $K\geq 2$ local devices and each device has
$N_{k}$ number of datapoints. Denote by
$D_{k}=\big{(}(x_{k,1},y_{k,1}),(x_{k,2},y_{k,2}),\ldots,(x_{k,N_{k}},y_{k,N_{k}})\big{)}$
the data stored in device $k$ where $x\in\mathcal{X}$ is the input,
$\mathcal{X}$ is the input space, $y\in\mathcal{Y}$ is the output/label and
$\mathcal{Y}$ is the output space. Denote by $\Delta_{\mathcal{Y}}$ the
simplex over $\mathcal{Y}$, $h:\mathcal{X}\mapsto\Delta_{\mathcal{Y}}$ the
hypothesis and $\mathcal{H}$ a family of hypotheses $h$. Let $\ell$ be a loss
function defined over $\Delta_{\mathcal{Y}}\times\mathcal{Y}$. Without loss of
generality, assume $\ell\geq 0$. The loss of $h$ is therefore given by
$\ell(h(x),y)$. Let
$\mathcal{L}_{\mathcal{D}}(h)=\mathbb{E}_{\mathcal{D}}[\ell(h(x),y)]$ be the
expected loss of $h$ with respect to a distribution $\mathcal{D}$. Our goal is
to find the optimal $h_{\mathcal{D}}$ such that
$h_{\mathcal{D}}=\operatorname*{arg\,min}_{h}\mathcal{L}_{\mathcal{D}}(h)$.

Federated Learning: In the FL setting, data distribution $\mathcal{D}$ is
typically a mixture of client data distribution. Denote by $\mathcal{D}_{k}$
the data distribution of client $k$. We define
$\mathcal{D}_{\bm{p}}\coloneqq\sum_{k=1}^{K}p_{k}\mathcal{D}_{k}$, where
$\bm{p}=(p_{1},\ldots,p_{K})$ and $p_{k}=\frac{N_{k}}{\sum_{k}N_{k}}$. In
practice, since $\mathcal{D}_{\bm{p}}$ is unknown,
$\mathcal{L}_{\mathcal{D}_{\bm{p}}}$ is also unknown. Therefore, we resort to
minimizing the empirical loss. Let $\bm{\theta}\in\bm{\Theta}$ be a vector of
parameters defining a hypothesis $h$. Therefore, in the FL setting, we can
write the vanilla global objective function as
$F(\bm{\theta})\coloneqq\sum_{k=1}^{K}p_{k}F_{k}(\bm{\theta})$, where
$F_{k}(\bm{\theta})=\frac{1}{N_{k}}\sum_{n=1}^{N_{k}}\ell(h(x_{k,n}),y_{k,n})$
is the local empirical risk. FedAvg (mcmahan2017communication) is one of the
most popular algorithms in the FL community: a central server selects a subset
of local devices and let each device run $E$ steps of local SGD. During each
communication round, each selected device $k$ sends its parameter
$\bm{\theta}_{k}$ to the server and server averages those model parameters.
Afterwards, the server broadcasts this averaged parameter to all local
devices. This cycle is repeated several times till convergence.

Fairness: Suppose there are $d\in[2,K]$ groups and each client can be assigned
to one of those groups $s\in[d]\coloneqq\\{1,\ldots,d\\}$. Note that clients
from different groups are typically non-$\bm{i.i.d.}$ Denote by
$k^{i},k\in[K],i\in[d]$ the index of $k$-th local device in group $i$.
Throughout this paper, we drop the superscript $i$ unless we want to emphasize
$i$ explicitly. The group fairness can be formally defined as follows.

######  Definition 1.

Denote by $\\{a^{i}_{\bm{w}}\\}_{1\leq i\leq d}$ the set of performance
measures (e.g., testing accuracy) of a trained model $\bm{w}$. For trained
models $\bm{\theta}$ and $\tilde{\bm{\theta}}$, we say $\bm{\theta}$ is more
fair than $\tilde{\bm{\theta}}$ if $Var(\\{a^{i}_{\bm{\theta}}\\}_{1\leq i\leq
d})<Var(\\{a^{i}_{\tilde{\bm{\theta}}}\\}_{1\leq i\leq d})$, where $Var$ is
variance.

It can be seen that when $d=K$, Definition 1 is equivalent to individual
fairness (li2019fair).

######  Remark 1.

The variance metric in Definition 1 can be replaced by other metrics such as
Cosine similarity or Entropy. Similar to the definition in li2018federated,
our goal is to ensure a uniform performance among all groups or devices.

##  3 Fairness Formulation

Our fair FL formulation aims at imposing group fairness while minimizing the
training error. More specifically, our goal is to minimize the discrepancies
in the average group losses while achieving a low training error. By
penalizing the spread in the loss among client groups, we propose a
regularization framework for computing optimal parameters $\bm{\theta}$ that
balances learning accuracy and fairness. This translates to solving the
following optimization problem

$\displaystyle\min_{\bm{\theta}}\;H(\bm{\theta})\triangleq\sum_{k=1}^{K}p_{k}F_{k}(\bm{\theta})+\lambda\sum_{1\leq
i<j\leq d}|L_{i}(\bm{\theta})-L_{j}(\bm{\theta})|,$ (1)

where $\lambda$ is a positive scalar that balances fairness and goodness-of-
fit, and

$L_{i}(\bm{\theta})\triangleq\dfrac{1}{|{\cal A}_{i}|}\sum_{k\in{\cal
A}_{i}}F_{k}(\bm{\theta})$

is the average loss for client group $i$, ${\cal A}_{i}$ is the set of indices
of devices who belong to group $i$, and $|\mathcal{A}|$ is the cardinality of
the set $\mathcal{A}$. In typical FL settings, each client uses local data to
optimize a surrogate of the global objective function. For instance, FedAvg
approach simply uses the local objective function
$H_{k}(\bm{\theta})=F_{k}(\bm{\theta})$ for a given client $k$. We thus obtain

$H(\bm{\theta})=\sum_{k=1}^{K}\,p_{k}H_{k}(\bm{\theta}).$

Using a similar approach and to account for the fairness term in Eq. (1), we
construct the following local objective function for a given client $k$

$\displaystyle\bm{\theta}_{k}\triangleq\operatorname*{arg\,min}_{\bm{\theta}}\;H_{k}(\bm{\theta})\triangleq\left(1+\dfrac{\lambda}{p_{k}\,|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta})\right)F_{k}(\bm{\theta}),$ (2)

where

$\displaystyle r_{k}(\bm{\theta})\triangleq\sum_{1\leq j\neq s_{k}\leq
d}\mbox{sign}(L_{s_{k}}(\bm{\theta})-L_{j}(\bm{\theta})),$ (3)

and $s_{k}\in[d]$ is the group index of device $k$. In simple words,
$r_{k}(\bm{\theta})\in\\{-d+1,-d+3,\ldots,d-3,d-1\\}$ is a scalar directly
related to the statistical ordering of $L_{s_{k}}$ among client group losses.
To illustrate that in a simple example, suppose that at a given $\bm{\theta}$,
we have $L_{1}(\bm{\theta})\geq L_{2}(\bm{\theta})\geq\ldots\geq
L_{d}(\bm{\theta})$. Then,

$r_{k}(\bm{\theta})=\begin{cases}\begin{array}[]{ll}d-1&\mbox{if }s_{k}=1\\\
d-3&\mbox{if }s_{k}=2\\\ \quad\vdots&\\\ -d+1&\mbox{if
}s_{k}=d.\end{array}\end{cases}$

Notice that in Eq. (2) $H_{k}(\bm{\theta})$ is not differentiable due to the
$r_{k}(\bm{\theta})$ component. Our approach, however, fixes
$r_{k}(\bm{\theta})$ during the communication round. Specifically, at the end
of each communication, the server computes an aggregated model parameter at
which $r_{k}(\bar{\bm{\theta}})$ is evaluated. These values are broadcasted to
local devices that perform updates on the aggregated parameter
$\bar{\bm{\theta}}$ while keeping $r_{k}$ fixed. The intuition is that each
local device does not have any information about other devices and hence
cannot update the values of $r_{k}$ during local training.

Our next lemma provides another interesting interpretation for the our global
objective and shows that, similar to FedAvg,
$H(\bm{\theta})=\sum_{k=1}^{K}\,p_{k}\,H_{k}(\bm{\theta})$.

######  Lemma 1.

For any given $\bm{\theta}$, the global objective function $H(\bm{\theta})$
defined in (1) can be expressed as

$H(\bm{\theta})=\sum_{k=1}^{K}p_{k}\,\left(1+\dfrac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta})\right)F_{k}(\bm{\theta}),$ (4)

where

$r_{k}(\bm{\theta})=\sum_{1\leq j\neq s_{k}\leq
d}{\mbox{sign}}(L_{s_{k}}(\bm{\theta})-L_{j}(\bm{\theta})).$

Consequently,

$H(\bm{\theta})=\sum_{k=1}^{K}\,p_{k}H_{k}(\bm{\theta}).$

According to Lemma 1, one can view our global objective as a parameter-based
weighted sum of the client loss functions. Particularly, rather than using
uniform weighting for clients, our assigned weights are functions of the
parameter $\bm{\theta}$. More specifically, for a given parameter
$\bm{\theta}$, our objective yields higher weights for groups with higher
average group loss; hence, imposing group fairness. To illustrate this idea,
we provide a simple concrete example.

######  Example 1.

Without loss of generality and for a given $\bm{\theta}$, consider four
different groups each having 10 clients with
$L_{1}(\bm{\theta})>L_{2}(\bm{\theta})>L_{3}(\bm{\theta})>L_{4}(\bm{\theta})$.
Then, our global objective function (8) equals

$\displaystyle\sum_{k\in{\cal
A}_{1}}p_{k}\left(1+\dfrac{3\lambda}{10p_{k}}\right)F_{k}(\bm{\theta})+\sum_{k\in{\cal
A}_{2}}p_{k}\left(1+\dfrac{\lambda}{10p_{k}}\right)F_{k}(\bm{\theta})$
$\displaystyle\qquad+\sum_{k\in{\cal
A}_{3}}p_{k}\left(1-\dfrac{\lambda}{10p_{k}}\right)F_{k}(\bm{\theta})+\sum_{k\in{\cal
A}_{4}}p_{k}\left(1-\dfrac{3\lambda}{10p_{k}}\right)F_{k}(\bm{\theta}).$

The objective clearly demonstrates a higher weight applied to clients that
belong to a group with a higher average loss.

According to (2), the optimization problem solved by every selected client is
a weighted version of the local objective in FedAvg. This will impose a higher
weight for clients that belong to groups with higher average loss. These
weights will be dynamically updated at every communication round. To assure
positive weights for clients, we require the following bounds on $\lambda$

$0\leq\lambda<\lambda_{max}\triangleq\min_{k}\left\\{\dfrac{p_{k}|{\cal
A}_{s_{k}}|}{d-1}\right\\}$

When $\lambda=0$, our approach is exactly FedAvg. Moreover, a higher value of
$\lambda$ imposes more emphasis on the fairness term. When imposing individual
fairness, each client is considered to be a group. Hence,

$H(\bm{\theta})=\sum_{k=1}^{K}\left(1+\dfrac{\lambda}{p_{k}}r_{k}(\bm{\theta})\right)F_{k}(\bm{\theta}).$

Our approach can be compared to methods that use re-weighting of clients. For
instance, AFL proposed by mohri2019agnostic computes at every communication
round the worst-case distribution of weights among clients. This approach
promotes robustness but may be overly conservative in the sense that it
focuses on the largest loss and thus causes very pessimistic performance to
other clients. Our algorithm, however, adaptively updates the weight of
clients at every communication round based on the statistical ordering of
client/group losses. Moreover, the dynamic update of the weights can
potentially avoid over-fitting by impeding updates for clients with low loss.
We will further demonstrate the advantages of our algorithm in Sec. 6. In the
next subsection, we will further detail our algorithm.

###  3.1 Algorithm

In this section, we describe our proposed algorithm GIFAIR-FL which is
detailed in Algorithm 1. At every communication round $c$, our algorithm
selects a set of clients to participate in the training and shares $r_{k}^{c}$
with each selected client. For each client, multiple (stochastic) gradient
descent steps are then applied to a weighted client loss function. The updated
parameters are then passed to the server that aggregates these results and
computes $r_{k}^{c+1}$.

Computationally, our approach requires evaluating the client loss function at
every communication round to compute $r_{k}^{c}$. Compared to existing fair FL
approaches, GIFAIR-FL is simple and computationally efficient. For instance,
q-FFL proposed by li2019fair first runs FedAvg to obtain a well-tuned learning
rate and uses this learning rate to roughly estimate the Lipschitz constant
$L$. Another example is AFL which requires running two gradient calls at each
iteration to estimate the gradients of model and weight parameters. Similarly,
Ditto requires running additional steps of SGD, at each communication round,
to generate personalized solutions. In contrast, our proposed method can be
seen as a fairness-aware weighted version of FedAvg.

Data: number of devices $K$, fraction $\alpha$, number of communication rounds
$C$, number of local updates $E$, SGD learning rate schedule
$\\{\eta^{(t)}\\}_{t=1}^{E}$, initial model parameter $\bm{\theta}$,
regularization parameter $\lambda$, initial loss $\\{L_{i}\\}_{1\leq i\leq d}$

for  _$c=0:(C-1)$_ do

Select $\left\lceil\alpha K\right\rceil$ clients by sampling probability
$p_{k}$ and denote by ${\cal S}_{c}$ the indices of these clients;

Server broadcasts $\left(\bm{\theta},\left\\{\dfrac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r^{c}_{k}(\bm{\theta})\right\\}_{k\in{\cal S}_{c}}\right)$;

for  _$k\in\mathcal{S}_{c}$_ do

$\bm{\theta}_{k}^{(0)}=\bm{\theta}$;

for  _$t=0:(E-1)$_ do

$\bm{\theta}^{(t+1)}_{k}=\bm{\theta}^{(t)}_{k}-\eta^{(t)}\left(1+\dfrac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r^{c}_{k}(\bm{\theta})\right)\nabla F_{k}(\bm{\theta}^{(t)}_{k})$
;

// Note that $r^{c}_{k}(\bm{\theta})$ is fixed during local update

end for

end for

Aggregation $\bm{\bar{\theta}}_{c}=\frac{1}{\left\lceil\alpha
K\right\rceil}\sum_{k\in\mathcal{S}_{c}}\bm{\theta}^{(E)}_{k}$, Set
$\bm{\theta}=\bm{\bar{\theta}}_{c}$;

Calculate $L_{i}=\frac{1}{|{\cal A}_{i}|}\sum_{k\in{\cal
A}_{i}}F_{k}(\bm{\theta}^{(E)}_{k})$ for all $i\in[d]$ and update
$r_{k}^{c+1}(\bm{\theta})$;

$c\leftarrow c+1$;

end for

Return $\bm{\bar{\theta}}_{C}$.

Algorithm 1 GIFAIR-FL Algorithm

In algorithm 1, we sample local devices by sampling probability $p_{k}$ and
aggregate model parameters by unweighted average. If devices are instead
uniformly sampled, then the aggregation strategy should be replaced by
$\bm{\bar{\theta}}_{c}=\frac{K}{\left\lceil\alpha
K\right\rceil}\sum_{k\in\mathcal{S}_{c}}p_{k}\bm{\theta}^{(E)}_{k}$
(li2019convergence).

######  Remark 2.

Instead of broadcasting $p_{k}$ and $|{\cal A}_{s_{k}}|$ separately to local
devices, the central server broadcasts the product $\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}^{c}(\bm{\theta})$ to client $k$. Hence, the local device $k$
cannot obtain any information about $p_{k}$, $|{\cal A}_{s_{k}}|$ and
$r_{k}^{c}(\bm{\theta})$. This strategy can protect privacy of other devices.

In addition to the algorithm, we also derive a generalization bound for our
learning model GIFAIR-FL. Due to space limitation, we defer this Theorem to
Appendix 2.3.

##  4 Convergence Results with non-$i.i.d.$ data

In this section, we show that, under mild conditions, GIFAIR-FL converges to
the global optimal solution at a rate of $\mathcal{O}(\frac{1}{T})$ for
strongly convex function and to a stationary point at a rate of
$\mathcal{O}(\frac{1}{\sqrt{T}})$, up to a logarithmic factor, for non-convex
function. Our theorems hold for both $\bm{i.i.d.}$ and non-$\bm{i.i.d.}$ data.

###  4.1 Strongly Convex Functions

We assume each device performed $E$ steps of local updates and make the
following assumptions. Here, our assumptions are based on $F_{k}$ rather than
$H_{k}$. These assumptions are very common in many FL papers
(li2019communication; li2018federated; li2019convergence).

######  Assumption 1.

$F_{k}$ is $L$-smooth and $\mu$-strongly convex for all $k\in[K]$.

######  Assumption 2.

Denote by $D^{(t)}_{k}$ the batched data from device $k$ and $\nabla
F_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k})$ the stochastic gradient calculated
on this batched data. The variance of stochastic gradient is bounded.
Specifically,

$\displaystyle\mathbb{E}\bigg{\\{}\left\lVert\nabla
F_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k})-\nabla
F_{k}(\bm{\theta}^{(t)}_{k})\right\rVert^{2}\bigg{\\}}\leq\sigma_{k}^{2},\;\;\forall
k\in[K].$

######  Assumption 3.

The expected squared norm of stochastic gradient is bounded. Specifically,

$\displaystyle\mathbb{E}\bigg{\\{}\left\lVert\nabla
F_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k})\right\rVert^{2}\bigg{\\}}\leq
G^{2},\forall k\in[K].$

Typically, data from different groups are non-$\bm{i.i.d.}$. We modify the
definition in li2019convergence to roughly quantify the degree of
non-$i.i.d.$-ness. Specifically,

$\displaystyle\Gamma_{K}=H^{*}-\sum_{k=1}^{K}p_{k}H^{*}_{k}=\sum_{k=1}^{K}p_{k}(H^{*}-H^{*}_{k}),$

where $H^{*}\triangleq
H(\bm{\theta}^{*})=\sum_{k=1}^{K}H_{k}(\bm{\theta}^{*})$ is the optimal value
of the global objective function and $H_{k}^{*}\triangleq
H_{k}(\bm{\theta}_{k}^{*})$ is the optimal value of the local loss function.
If data are $i.i.d.$, then $\Gamma_{K}\to 0$ as the number of samples grows.
Otherwise, $\Gamma_{K}\neq 0$ (li2019convergence). Given all aforementioned
assumptions, we next prove the convergence of our proposed algorithm. We first
assume all devices participate in each communication round (i.e.,
$|\mathcal{S}_{c}|=K,\forall c$).

######  Theorem 2.

Assume Assumptions 1-3 hold and $|\mathcal{S}_{c}|=K$. If $\eta^{(t)}$ is
decreasing in a rate of $\mathcal{O}(\frac{1}{t})$ and
$\eta\leq\mathcal{O}(\frac{1}{L})$, then for $\gamma,\mu,\epsilon>0$, we have

$\displaystyle\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(T)})\bigg{\\}}-H^{*}\leq\frac{L}{2}\frac{1}{\gamma+T}\bigg{\\{}\frac{4\xi}{\epsilon^{2}\mu^{2}}+(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}},$

where
$\xi=8(E-1)^{2}G^{2}+4L\Gamma+2\frac{\Gamma_{max}}{\eta^{(t)}}+4\sum_{k=1}^{K}p_{k}^{2}\sigma_{k}^{2}$
and
$\Gamma_{max}\coloneqq\sum_{k=1}^{K}p_{k}|(H^{*}-H^{*}_{k})|\geq|\sum_{k=1}^{K}p_{k}(H^{*}-H^{*}_{k})|=|\Gamma_{K}|$.

This expression depicts a $\mathcal{O}(\frac{1}{T})$ convergence rate which
agrees with FedAvg. Note that the convergence rate is also affected by $\xi$,
which contains the degree of non-$i.i.d.$-ness. Let
$\Gamma_{K}=\Gamma_{max}=0$, then we can obtain the convergence bound for
$i.i.d.$ data.

Next, we assume only a fraction of devices participate in each communication
round (i.e., $|\mathcal{S}_{c}|=\alpha K,\forall c,\alpha\in(0,1)$). Along
this line, we assume local devices are sampled according to the sampling
probability $p_{k}$ (li2018federated). Our Theorem can also be extended to the
scenario where devices are sampled uniformly (i.e., with the same
probability). In this scenario, the aggregation strategy is
$\bm{\bar{\theta}}_{c}=\frac{K}{\left\lceil\alpha
K\right\rceil}\sum_{k\in\mathcal{S}_{c}}p_{k}\bm{\theta}_{k}$
(li2019convergence). The proof will be similar.

######  Theorem 3.

Assume at each communication round, central server sampled a fraction $\alpha$
of devices and those local devices are sampled according to the sampling
probability $p_{k}$. Additionally, assume Assumption 1-3 hold. If $\eta^{(t)}$
is decreasing in a rate of $\mathcal{O}(\frac{1}{t})$ and
$\eta\leq\mathcal{O}(\frac{1}{L})$, then for $\gamma,\mu,\epsilon>0$, we have

$\displaystyle\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(T)})\bigg{\\}}-H^{*}\leq\frac{L}{2}\frac{1}{\gamma+T}\bigg{\\{}\frac{4(\xi+\tau^{\prime})}{\epsilon^{2}\mu^{2}}+(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}},$

where $\tau^{\prime}=\frac{4G^{2}E^{2}}{\left\lceil\alpha K\right\rceil}$.

###  4.2 Non-convex Functions

To prove the convergence result on non-convex functions, we replace Assumption
1 by the following assumption.

######  Assumption 4.

$F_{k}$ is $L$-smooth for all $k\in[K]$.

######  Theorem 4.

Assume Assumptions 2-4 hold and $|\mathcal{S}_{c}|=K$. If
$\eta^{(t)}=\mathcal{O}(\frac{1}{\sqrt{t}})$ and
$\eta^{(t)}\leq\mathcal{O}(\frac{1}{L})$, then our algorithm converges to a
stationary point. Specifically,

$\displaystyle\min_{t=1,\ldots,T}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}\leq\frac{1}{\sqrt{T}}\bigg{\\{}2\big{(}1+2KL^{2}\log(T+1)\big{)}\mathbb{E}\big{\\{}H(\bm{\bar{\theta}}^{(1)})-H^{*}\big{\\}}+2\xi_{\Gamma_{K}}\bigg{\\}},$

where
$\xi_{\Gamma_{K}}=\mathcal{O}\bigg{(}\big{(}2KL^{2}\Gamma_{K}+8KLG^{2}+(2L+8KL)\sum_{k=1}^{K}p_{k}\sigma_{k}^{2}\big{)}\log(T+1)\bigg{)}$.

It can be seen that GIFAIR-FL will converge to a stationary point at a rate of
$\widetilde{\mathcal{O}}(\frac{1}{\sqrt{T}})$. Here, we note that a decay
learning rate schedule is necessary in Theorems 6-8. Otherwise, our algorithm
cannot converge due to the noise in SGD.

##  5 Related Work

Despite the rich literature for fairness in the non-federated machine learning
setting, traditional fairness measures such as demographic disparity
(feldman2015certifying), equal opportunity and equalized odds
(hardt2016equality) do not naturally extend to FL as there is no clear notion
of an outcome which is “good” for a device (kairouz2019advances). Instead,
fairness in the FL can be reframed as equal access to effective models.
Specifically, the goal is to train a global model that incurs a uniformly good
performance across all devices (kairouz2019advances). Indeed, most existing
work is thriving to develop fair FL algorithms along this line. Below we
highlight this literature.

Fair Federated Learning: mohri2019agnostic propose a minimax optimization
framework Agnostic federated learning (AFL). AFL optimizes the worst weighted
combination of local devices and is demonstrated to be robust to unseen
testing data. du2020fairness further refine the notation of AFL and propose an
AgnosticFair algorithm. Specifically, they linearly parametrize weight
parameters by kernel functions and show that AFL can be viewed as a special
case of AgnosticFair. Upon that hu2020fedmgda+ combine minimax optimization
with gradient normalization techniques to produce a fair algorithm FedMGDA+.
Motivated by fair resource allocation problems, li2019fair propose $q$-Fair
federated learning (q-FFL). q-FFL reweights loss functions such that devices
with poor performance will be given relatively higher weights. The q-FFL
objective is proved to encourage individual fairness in FL. However, this
algorithm requires accurate estimation of a local Lipschitz constant $L$.
Later, li2020tilted developed a tilted empirical risk minimization (TERM)
algorithm to handle outliers and class imbalance in statistical estimation
procedures. TERM has been shown to be superior to q-FFL in many FL
applications. Along this line, huang2020fairness propose to use training
accuracy and frequency to adjust weights of devices to promote fairness.
zhang2020fairfl develop an algorithm to minimize the discrimination index of
global model to encourage fairness. Here we note that some work study the
collaborative fairness in the FL setting (zhang2020hierarchically;
xu2020towards; lyu2020collaborative). The goal of those work, which is
perpendicular to our purpose, is to provide more rewards to high-contributing
participants while penalize free riders. Personalized Federated Learning:
Though personalization techniques do not directly target fairness, recent
papers have shown that personalized FL algorithms may improve fairness.
(arivazhagan2019federated; liang2020think) use different layers of a network
to represent global and personalized solutions. Specifically, they fit
personalized layers to each local device such that each device will return a
task-dependent solutions based on its own local data. wang2019federated;
yu2020salvaging; dinh2020personalized; li2021ditto resort to fine-tuning
technique to learn personalized models. Notably, li2021ditto develop a multi-
task personalized FL algorithm Ditto. After optimizing a global objective
function, Ditto allows local devices running more steps of SGD, subject to
some constraints, to minimize their own losses. li2021ditto have shown that
Ditto can significantly improve testing accuracies among local devices and
encourage fairness.

##  6 Experiments

In this section, we test GIFAIR-FL on image classification and the text
prediction datasets. We benchmark our model with the following algorithms:
q-FFL (li2019fair), TERM (li2020tilted), FedMGDA+ (hu2020fedmgda+), AFL
(mohri2019agnostic), and FedMGDA+ (hu2020fedmgda+). To the best of our
knowledge, those are well-known state-of-the-art FL algorithms aim for
fairness. We also benchmark our model with Ditto (li2021ditto) which is a
personalized FL approach using multi-task learning. Here we note that our
algorithm, GIFAIR-FL, and all other benchmarks (except Ditto) learn one single
global model rather than generating device-specific personalized solutions.
However, we include Ditto, to shed light on the ability of GIFAIR-FL to
improve fairness even compared to personalized FL and also on GIFAIR-FL’s
amenability to be readily extended to personalized FL settings as shown in the
GIFAIR-FL \+ Ditto benchmark. To combine GIFAIR-FL with Ditto, we replace the
global objective function of Ditto by our fairness objective function (1). We
detail this combined algorithm in the Appendix.

###  6.1 Image Classification

We start by considering a widely-used image classification dataset FEMNIST
(Federated Extended MNIST) (caldas2018leaf). FEMNIST consists of images of
digits (0-9) and English characters (A-Z, a-z) with 62 classes.

Individual Fairness (FEMNIST-skewed, $d=100$) Following the setting in
(li2018federated), we first sample 10 lower case characters (‘a’-‘j’) from
Extended MNIST (EMNIST) (cohen2017emnist) and distribute 5 classes of images
to each device. Each local device has 500 images. There are 100 devices in
total. Results are reported in Table 2. (FEMNIST-original, $d=500$) Following
the setting in (li2021ditto), we sample 500 devices and train models using the
default data stored in each device. Results are reported in Table 3.

Group Fairness (FEMNIST-3-groups, $d=3$) We manually divide FEMNIST data into
three groups. See Table 1 for the detailed assignment. This assignment is
inspired by the statistic that most people prefer to write in lowercase
letters while a small amount of people use capital letters or a mixed of two
types (jones2004case). In such cases, it is important to assure that an FL
algorithm is capable of achieving similar performance between such groups.
Results are reported in Table 4.

Group | Data Type | Number of Images | Number of Devices
---|---|---|---
Group 1 | Capital Letters + Digits | 800 | 60
Group 2 | Lowercase Letters + Digits | 1,000 | 100
Group 3 | Capital/Lowercase Letters + Digits | 600 | 40
Table 1: Data Structure of FEMNIST-3-groups

Implementation: For all tasks, we randomly split the data on each local device
into a $70\%$ training set, a $10\%$ validation set and a $20\%$ testing set.
This is a common data splitting strategy used in many FL papers
(li2018federated; chen2018federated; reddi2020adaptive). The batch size is set
to be 32. We use the tuned initial learning rate $0.1$ and decay rate $0.99$
for each method. During each communication round, $10$ devices are randomly
selected and each device will run 2 epochs of SGD. We use a CNN model with 2
convolution layers followed by 2 fully connected layers. All benchmark models
are well-tuned. Specifically, we solve q-FFL with
$q\in\\{0,0.001,0.01,0.1,1,2,5,10\\}$ (li2019fair) in parallel and select the
best $q$. Here, the best $q$ is defined as the $q$ value where the variance
decreases the most while the averaged testing accuracy is superior or similar
to FedAvg. This definition is borrowed from the original q-FFL paper
li2019fair. Similarly, we train TERM with $t\in\\{1,2,5\\}$ and select the
best $t$ (li2020tilted). For Ditto, we tune the regularization parameter
$\lambda_{Ditto}\in\\{0.01,0.05,0.1,0.5,1,2,5\\}$. In GIFAIR-FL, we tune the
parameter
$\lambda\in\\{0,0.1\lambda_{max},0.2\lambda_{max},\ldots,0.8\lambda_{max},0.9\lambda_{max}\\}$.
Here kindly note that $\lambda_{max}$ is a function of $p_{k},|{\cal
A}_{s_{k}}|$ and $d$ (i.e., data-dependent).

Performance metrics: Denote by $a_{k}$ the prediction accuracy on device $k$.
We define (1) individual-level mean accuracy as
$\bar{a}\coloneqq\frac{1}{K}\sum_{k=1}^{K}a_{k}$ and (2) individual-level
variance as $Var(a)\coloneqq\frac{1}{K}\sum_{k=1}^{K}(a_{k}-\bar{a})^{2}$.

Algorithm | FedAvg | q-FFL | TERM | FedMGDA+ | Ditto | GIFAIR-FL |  GIFAIR-FL \+ Ditto
---|---|---|---|---|---|---|---
$\bar{a}$ | 79.2 (1.0) | 84.6 (1.9) | 84.2 (1.3) | 85.0 (1.7) | 92.5 (3.1) | 87.9 (0.9) |  93.4 (1.4)
$\sqrt{Var(a)}$ | 22.3 (1.1) | 18.5 (1.2) | 13.8 (1.0) | 14.9 (1.6) | 14.3 (1.0) |  5.7 (0.8) | 7.6 (1.0)
Table 2: Empirical results on FEMNIST-skewed. Each experiment is repeated 5 times. Algorithm | FedAvg | q-FFL | TERM | AFL | Ditto | GIFAIR-FL |  GIFAIR-FL \+ Ditto
---|---|---|---|---|---|---|---
$\bar{a}$ | 80.4 (1.3) | 80.9 (1.1) | 81.0 (1.0) | 82.4 (1.0) | 83.7 (1.9) | 83.2 (0.7) | 84.2 (1.1)
$\sqrt{Var(a)}$ | 11.1 (1.4) | 10.6 (1.3) | 10.3 (1.2) | 9.85 (0.9) | 10.1 (1.6) | 5.2 (0.8) | 4.3 (1.0)
Table 3: Test accuracy on FEMNIST-original. Each experiment is repeated 5 times. Algorithm | FedAvg | q-FFL | TERM | FedMGDA+ | Ditto | GIFAIR-FL
---|---|---|---|---|---|---
Group 1 | 79.72 (2.08) | 81.15 (1.97) | 81.29 (1.45) | 81.03 (2.28) | 82.37 (2.06) | 83.41 (1.34)
Group 2 | 90.93 (2.35) | 88.24 (2.13) | 88.08 (1.09) | 89.12 (1.74) | 92.05 (2.00) | 88.29 (1.22)
Group 3 | 80.21 (2.91) | 80.93 (1.86) | 81.84 (1.44) | 81.33 (1.59) | 83.03 (2.18) | 84.37 (1.85)
Discrepancy | 11.21 | 7.31 | 6.79 | 8.09 | 9.02 | 3.92
Table 4: Test accuracy on FEMNIST-3-groups. Each experiment is repeated 5
times. Discrepancy is the difference between the largest accuracy and the
smallest accuracy.

###  6.2 Text Data

Individual Fairness: We train a RNN to predict the next character using text
data built from The Complete Works of William Shakespeare. In this dataset,
each speaking role in the plays is treated as a device. Following the setting
in mcmahan2017communication; li2019fair, we subsample 31 roles ($d=31$). The
RNN model takes an 80-character sequence as input and outputs one character
after two LSTM layers and one densely-connected layer. For FedAvg, q-FFL and
Ditto, the best initial learning rate is 0.8 and decay rate is 0.95
(li2021ditto). We also adopt this setting to GIFAIR-FL. The batch size is set
to be 10. The number of local epochs is fixed to be 1 and all models are
trained for 500 epochs. Results are reported in Table 5.

Algorithm | FedAvg | q-FFL | AFL | Ditto | GIFAIR-FL |  GIFAIR-FL \+ Ditto
---|---|---|---|---|---|---
$\bar{a}$ | 53.21 (0.31) | 53.90 (0.30) | 54.58 (0.14) | 60.74 (0.42) | 57.04 (0.23) | 61.81 (0.19)
$\sqrt{Var(a)}$ | 9.25 (6.17) | 7.52 (5.10) | 8.44 (5.65) | 8.32 (4.77) | 3.14 (1.25) | 5.04 (1.85)
Table 5: Mean and standard deviation of test accuracy on Shakespeare $(d=31)$.
Each experiment is repeated 5 times.

Group Fairness: We obtain the gender information from
https://shakespeare.folger.edu/ and group speaking roles based on gender
($d=2$). It is known that the majority of characters in Shakespearean drama
are males. Simply training a FedAvg model on this dataset will cause implicit
bias towards male characters. On a par with this observation, we subsample 25
males and 10 females from The Complete Works of William Shakespeare. Here we
note that each device in the male group implicitly has more text data. The
setting of hyperparameters is same as that of individual fairness. Results are
reported in Table 6.

Algorithm | FedAvg | q-FFL | FedMGDA+ | Ditto | GIFAIR-FL
---|---|---|---|---|---
Male | 72.95 (1.70) | 67.14 (2.18) | 67.07 (2.11) | 74.19 (3.75) | 67.42(0.98)
Female | 40.39 (1.49) | 43.26 (2.05) | 43.85 (2.32) | 45.73 (4.01) | 52.04 (1.10)
Discrepancy | 32.56 | 23.88 | 23.22 | 28.46 | 15.38
Table 6: Test accuracy on Shakespeare ($d=2$). Each experiment is repeated 5
times.

###  6.3 Analysis of Results

Based on Table 2-6, we can obtain some important insights. First, compared to
other benchmark models, GIFAIR-FL leads to significantly more fair solutions.
As shown in Table 2, 3 and 5, our algorithm significantly reduces the variance
of testing accuracy of all devices (i.e., $Var(a)$) while the average testing
accuracy remains consistent. Second, from Tables 4 and 6, it can be seen that
GIFAIR-FL boosted the performance of the group with the worst testing accuracy
and achieved the smallest discrepancy. Notably, this boost did not effect the
performance of other groups. This indicates that GIFAIR-FL is capable of
ensuring fairness among different groups while retaining a superior or similar
prediction accuracy compared to existing benchmark models. Finally, we note
that GIFAIR-FL sometimes achieves lower prediction performance than Ditto.
This is understandable as Ditto provides a personalized solution $\bm{v}_{k}$
to each device $k$ while our model only returns a global parameter
$\bm{\bar{\theta}}$. Yet, as shown in the last column, if we combine our FL
formulation with Ditto by changing the global model of Ditto to Eq. (1), then
the prediction performance can be significantly improved without sacrificing
fairness. This also reveals that our formulation can be easily incorporated to
other FL models. However, even without combining with personalized algorithm,
GIFAIR-FL achieves superior testing performance than existing fair FL
benchmark models.

###  6.4 Sensitivity Analysis

Figure 1: Sensitivity with respect to $\lambda$ (Shakespeare Dataset).

In this section, we study the effect of the tuning parameter
$\lambda\in[0,\lambda_{max})$ using the Shakespeare dataset. Results are
reported in Figure 1. It can be seen that as $\lambda$ increases, the
discrepancy between male and female groups decreases accordingly. However,
after $\lambda$ passing a certain threshold, the averaged testing accuracy of
female group remained flat yet the performance of male group significantly
dropped. Therefore, in practice, it is recommended to consider a moderate
$\lambda$ value. Intuitively, when $\lambda=0$, GIFAIR-FL becomes FedAvg. When
$\lambda$ is close to $\lambda_{max}$, the coefficient (i.e.,
$(1+\lambda\frac{1}{p_{k}|{\cal A}_{s_{k}}|}r_{k})$) of device with good
performance will be close to zero and the updating is impeded. A moderate
$\lambda$ balances those two situations well. Besides this example, we also
conducted additional sensitivity analysis. Due to space limitation, we defer
those results into Appendix.

##  7 Conclusion

In this paper, we propose GIFAIR-FL: an approach that imposes group and
individual fairness to federated learning scenario. Experiments show that
GIFAIR-FL can lead to more fair solutions compared to recent state-of-the-art
fair and personalized FL algorithms while retaining similar testing
performance. To the best of our knowledge, fairness in FL is an under
underinvestigated area and we hope our work will help inspire continued
exploration into fair FL algorithms.

One possible future research direction is considering sharpness-aware learning
rate scheduler that automatically adjusts learning rate for each device
(yue2020salr). We believe this approach can further enable personalization in
the FL.

##  8 Appendix

In Sec. 9, we restate our main assumptions. In Sec. 10, we provide the
detailed proofs of Lemma and Theorems in our main paper. Finally, in Sec. 12,
we present some additional empirical results.

##  9 Assumptions

We make the following assumptions.

######  Assumption 5.

$F_{k}$ is $L$-smooth and $\mu$-strongly convex for all $k\in[K]$.

######  Assumption 6.

Denote by $D^{(t)}_{k}$ the batched data from client $k$ and $\nabla
F_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k})$ the stochastic gradient calculated
on this batched data. The variance of stochastic gradients are bounded.
Specifically,

$\displaystyle\mathbb{E}\bigg{\\{}\left\lVert\nabla
F_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k})-\nabla
F_{k}(\bm{\theta}^{(t)}_{k})\right\rVert^{2}\bigg{\\}}\leq\sigma_{k}^{2},\forall
k\in[K].$

It can be shown that, at local iteration $t$ during communication round $c$,

$\displaystyle\mathbb{E}\bigg{\\{}\left\lVert\nabla
H_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k})-\nabla
H_{k}(\bm{\theta}^{(t)}_{k})\right\rVert^{2}\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}\left\lVert(1+\frac{\lambda
r^{c}_{k}}{p_{k}|{\cal A}_{s_{k}}|})\nabla
F_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k})-(1+\frac{\lambda
r_{k}^{c}}{p_{k}|{\cal A}_{s_{k}}|})\nabla
F_{k}(\bm{\theta}^{(t)}_{k})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}^{c})^{2}\sigma_{k}^{2},\forall k\in[K].$

######  Assumption 7.

The expected squared norm of stochastic gradient is bounded. Specifically,

$\displaystyle\mathbb{E}\bigg{\\{}\left\lVert\nabla
F_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k})\right\rVert^{2}\bigg{\\}}\leq
G^{2},\forall k\in[K].$

It can be shown that, at local iteration $t$ during communication round $c$,

$\displaystyle\mathbb{E}\bigg{\\{}\left\lVert\nabla
H_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k})\right\rVert\bigg{\\}}=\mathbb{E}\bigg{\\{}\left\lVert(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}^{c})\nabla
F_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}^{c})^{2}G^{2},\forall k\in[K].$

For the non-convex setting, we replace Assumption 5 by the following
assumption.

######  Assumption 8.

$F_{k}$ is $L$-smooth for all $k\in[K]$.

In our proof, for the sake of neatness, we drop the superscript of
$r^{c}_{k}$.

We use the definition in li2019convergence to roughly quantify the degree of
non-$i.i.d.$-ness. Specifically,

$\displaystyle\Gamma_{K}=H^{*}-\sum_{k=1}^{K}p_{k}H^{*}_{k}=\sum_{k=1}^{K}p_{k}(H^{*}-H^{*}_{k}).$

If data from all sensitive attributes are $i.i.d.$, then $\Gamma_{K}=0$ as
number of clients grows. Otherwise, $\Gamma_{K}\neq 0$ (li2019convergence).

##  10 Detailed Proof

###  10.1 Proof of Lemma

######  Lemma 2.

For any given $\bm{\theta}$, the global objective function $H(\bm{\theta})$
defined in the main paper can be expressed as

$H(\bm{\theta})=\sum_{k=1}^{K}\left(p_{k}+\dfrac{\lambda}{|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta})\right)F_{k}(\bm{\theta}),$

where

$\displaystyle r_{k}(\bm{\theta})\triangleq\sum_{1\leq j\neq s_{k}\leq
d}{\mbox{sign}}(L_{s_{k}}(\bm{\theta})-L_{j}(\bm{\theta}))$

and $s_{k}\in[d]$ is the group index of device $k$. Consequently,

$H(\bm{\theta})=\sum_{k=1}^{K}\,p_{k}H_{k}(\bm{\theta}).$

###### Proof.

By definition, at communication round $c$,

$\displaystyle H(\bm{\theta})$
$\displaystyle=\sum_{k=1}^{K}p_{k}F_{k}(\bm{\theta})+\lambda\sum_{1\leq
i<j\leq d}|L_{i}(\bm{\theta})-L_{j}(\bm{\theta})|$
$\displaystyle=\sum_{k=1}^{K}p_{k}F_{k}(\bm{\theta})+\lambda\sum_{1\leq
i<j\leq d}\bigg{|}\dfrac{1}{|{\cal A}_{i}|}\sum_{k\in{\cal
A}_{i}}F_{k}(\bm{\theta})-\dfrac{1}{|{\cal A}_{j}|}\sum_{k\in{\cal
A}_{j}}F_{k}(\bm{\theta})\bigg{|}$
$\displaystyle=\sum_{k=1}^{K}p_{k}F_{k}(\bm{\theta})+\lambda\sum_{1\leq
i<j\leq
d}\mbox{sign}(L_{i}(\bm{\theta})-L_{j}(\bm{\theta}))\bigg{(}\dfrac{1}{|{\cal
A}_{i}|}\sum_{k\in{\cal A}_{i}}F_{k}(\bm{\theta})-\dfrac{1}{|{\cal
A}_{j}|}\sum_{k\in{\cal A}_{j}}F_{k}(\bm{\theta})\bigg{)}$
$\displaystyle=\sum_{k=1}^{K}p_{k}F_{k}(\bm{\theta})+\lambda\sum_{u=1}^{d-1}\sum_{u<j\leq
d}\mbox{sign}(L_{u}(\bm{\theta})-L_{j}(\bm{\theta}))\bigg{(}\dfrac{1}{|{\cal
A}_{u}|}\sum_{k\in{\cal A}_{u}}F_{k}(\bm{\theta})-\dfrac{1}{|{\cal
A}_{j}|}\sum_{k\in{\cal A}_{j}}F_{k}(\bm{\theta})\bigg{)}$
$\displaystyle=\sum_{k=1}^{K}p_{k}F_{k}(\bm{\theta})+\lambda\sum_{u=1}^{d}\sum_{k\in{\cal
A}_{u}}\sum_{u\neq j\leq
d}\mbox{sign}(L_{u}(\bm{\theta})-L_{j}(\bm{\theta}))\frac{F_{k}(\bm{\theta})}{|{\cal
A}_{u}|}$
$\displaystyle=\sum_{k=1}^{K}p_{k}F_{k}(\bm{\theta})+\sum_{k=1}^{K}\frac{\lambda}{|{\cal
A}_{s_{k}}|}\sum_{1\leq j\neq s_{k}\leq
d}\mbox{sign}(L_{s_{k}}(\bm{\theta})-L_{j}(\bm{\theta}))F_{k}(\bm{\theta})$
$\displaystyle=\sum_{k=1}^{K}\left(p_{k}+\dfrac{\lambda}{|{\cal
A}_{s_{k}}|}r^{c}_{k}(\bm{\theta})\right)F_{k}(\bm{\theta}).$

The fifth equality is achieved by rearranging the equation and merging items
with the same group label. By definition of $H_{k}$, we thus proved

$\displaystyle H(\bm{\theta})=\sum_{k=1}^{K}\,p_{k}H_{k}(\bm{\theta}).$

∎

###  10.2 Learning Bound

We present a generalization bound for our learning model. Denote by
$\mathcal{G}$ the family of the losses associated to a hypothesis set
$\mathcal{H}:\mathcal{G}=\\{(x,y)\mapsto\ell(h(x),y):h\in\mathcal{H}\\}$. The
weighted Rademacher complexity (mohri2019agnostic) is defined as

$\displaystyle\mathfrak{R}_{\bm{m}}(\mathcal{G},\bm{p})\coloneqq\mathop{\mathbb{E}}_{\bm{\sigma}}\bigg{[}\sup_{h\in\mathcal{H}}\sum_{k=1}^{K}\frac{p_{k}}{N_{k}}\sum_{n=1}^{N_{k}}\sigma_{k,n}\ell(h(x_{k,n}),y_{k,n})\bigg{]}$

where $\bm{m}=(N_{1},N_{2},\ldots,N_{k})$, $\bm{p}=(p_{1},\ldots,p_{K})$ and
$\bm{\sigma}=(\sigma_{k,n})_{k\in[K],n\in[N_{k}]}$ is a collection of
Rademacher variables taking values in $\\{-1,+1\\}$. Denote by
$\mathcal{L}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)$ the expected loss according
to our fairness formulation. Denote by
$\hat{\mathcal{L}}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)$ the expected empirical
loss (See Appendix for a detailed expression).

######  Theorem 5.

Assume that the loss $\ell$ is bounded above by $M>0$. Fix $\epsilon_{0}>0$
and $\bm{m}$. Then, for any $\delta_{0}>0$, with probability at least
$1-\delta_{0}$ over samples $D_{k}\sim\mathcal{D}_{k}$, the following holds
for all $h\in\mathcal{H}$:

$\displaystyle\mathcal{L}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)\leq\hat{\mathcal{L}}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)+\sqrt{\frac{1}{2}\sum_{k=1}^{K}(\frac{p_{k}}{N_{k}}M+\lambda\frac{d(d-1)}{2}M)^{2}\log\frac{1}{\delta_{0}}}+2\mathfrak{R}_{\bm{m}}(\mathcal{G},\bm{p})+\lambda\frac{d(d-1)}{2}M.$

It can be seen that, given a sample of data, we can bound the generalization
error
$\mathcal{L}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)-\hat{\mathcal{L}}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)$
with high probability. When $\lambda=0$, the bound is same as the
generalization bound in FedAvg (mohri2018foundations). When we consider the
worst combination of $p_{k}$ by taking the supremum of the upper bound in
Theorem 5 and let $\lambda=0$, then our generalization bound is same as the
one in AFL (mohri2019agnostic).

###### Proof.

Define

$\displaystyle\Phi(D_{1},\ldots,D_{K})=\sup_{h\in\mathcal{H}}\bigg{(}\mathcal{L}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)-\hat{\mathcal{L}}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)\bigg{)}.$

Let $D^{\prime}=(D_{1}^{\prime},\ldots,D_{K}^{\prime})$ be a sample differing
from $D=(D_{1},\ldots,D_{K})$ only by one point $x_{k,n}^{\prime}$. Therefore,
we have

$\displaystyle\Phi(D^{\prime})-\Phi(D)$
$\displaystyle=\sup_{h\in\mathcal{H}}\bigg{(}\mathcal{L}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)-\hat{\mathcal{L}}_{\mathcal{D}^{{}^{\prime}\lambda}_{\bm{p}}}(h)\bigg{)}-\sup_{h\in\mathcal{H}}\bigg{(}\mathcal{L}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)-\hat{\mathcal{L}}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)\bigg{)}$
$\displaystyle\leq\sup_{h\in\mathcal{H}}\bigg{(}\mathcal{L}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)-\hat{\mathcal{L}}_{\mathcal{D}^{{}^{\prime}\lambda}_{\bm{p}}}(h)\bigg{)}-\bigg{(}\mathcal{L}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)-\hat{\mathcal{L}}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)\bigg{)}$
$\displaystyle\leq\sup_{h\in\mathcal{H}}\big{\\{}\sup_{h\in\mathcal{H}}\mathcal{L}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)-\sup_{h\in\mathcal{H}}\hat{\mathcal{L}}_{\mathcal{D}^{{}^{\prime}\lambda}_{\bm{p}}}(h)-\mathcal{L}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)+\hat{\mathcal{L}}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)\big{\\}}$
$\displaystyle=\sup_{h\in\mathcal{H}}\big{\\{}\hat{\mathcal{L}}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)-\hat{\mathcal{L}}_{\mathcal{D}^{{}^{\prime}\lambda}_{\bm{p}}}(h)\big{\\}}$

By definition,

$\displaystyle\hat{\mathcal{L}}_{\mathcal{D}^{{}^{\prime}\lambda}_{\bm{p}}}(h)$
$\displaystyle=\sum_{k=1}^{K}\frac{p_{k}}{N_{k}}\sum_{n=1}^{N_{k}}\ell(h(x^{\prime}_{k,n}),y^{\prime}_{k,n})+$
$\displaystyle\lambda\sum_{1\leq i<j\leq d}|\frac{\sum_{k\in{\cal
A}_{i}}\frac{1}{N_{k}}\sum_{n=1}^{N_{k}}\ell(h(x^{\prime}_{k,n}),y^{\prime}_{k,n})}{|{\cal
A}_{i}|}-\frac{\sum_{k\in{\cal
A}_{j}}\frac{1}{N_{k}}\sum_{n=1}^{N_{k}}\ell(h(x^{\prime}_{k,n}),y^{\prime}_{k,n})}{|{\cal
A}_{j}|}|.$

Therefore,

$\displaystyle\sup_{h\in\mathcal{H}}\big{\\{}\hat{\mathcal{L}}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)-\hat{\mathcal{L}}_{\mathcal{D}^{{}^{\prime}\lambda}_{\bm{p}}}(h)\big{\\}}$
$\displaystyle\leq\sup_{h\in\mathcal{H}}\bigg{[}\frac{p_{k}}{N_{k}}(\ell(h(x^{\prime}_{k,n}),y^{\prime}_{k,n})-\ell(h(x_{k,n}),y_{k,n}))+\lambda\frac{d(d-1)}{2}M\bigg{]}$
$\displaystyle\leq\frac{p_{k}}{N_{k}}M+\lambda\frac{d(d-1)}{2}M.$

By McDiarmid’s inequality, for
$\delta_{0}=\exp\bigg{(}\frac{-2\epsilon_{0}^{2}}{\sum_{k=1}^{K}(\frac{p_{k}}{N_{k}}M+\lambda\frac{d(d-1)}{2}M)^{2}}\bigg{)}$,
the following holds with probability at least $1-\delta_{0}$

$\displaystyle\Phi(D)-\mathbb{E}_{D}[\Phi(D)]\leq\epsilon_{0}=\sqrt{\frac{1}{2}\sum_{k=1}^{K}(\frac{p_{k}}{N_{k}}M+\lambda\frac{d(d-1)}{2}M)^{2}\log\frac{1}{\delta_{0}}}.$

Our next goal is to bound $\mathbb{E}[\Phi(D)]$. We have

$\displaystyle\mathbb{E}_{D}[\Phi(D)]$
$\displaystyle=\mathbb{E}_{D}\bigg{[}\sup_{h\in\mathcal{H}}\bigg{(}\mathcal{L}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)-\hat{\mathcal{L}}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)\bigg{)}\bigg{]}$
$\displaystyle=\mathbb{E}_{D}\bigg{[}\sup_{h\in\mathcal{H}}\mathbb{E}_{D^{\prime}}\bigg{(}\hat{\mathcal{L}}_{\mathcal{D}^{{}^{\prime}\lambda}_{\bm{p}}}(h)-\hat{\mathcal{L}}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)\bigg{)}\bigg{]}$
$\displaystyle\leq\mathbb{E}_{D}\mathbb{E}_{D^{\prime}}\sup_{h\in\mathcal{H}}\bigg{(}\hat{\mathcal{L}}_{\mathcal{D}^{{}^{\prime}\lambda}_{\bm{p}}}(h)-\hat{\mathcal{L}}_{\mathcal{D}^{\lambda}_{\bm{p}}}(h)\bigg{)}$
$\displaystyle\leq\mathbb{E}_{D}\mathbb{E}_{D^{\prime}}\sup_{h\in\mathcal{H}}\bigg{[}\sum_{k=1}^{K}\frac{p_{k}}{N_{k}}\sum_{n=1}^{N_{k}}\ell(h(x^{\prime}_{k,n}),y^{\prime}_{k,n})-\sum_{k=1}^{K}\frac{p_{k}}{N_{k}}\sum_{n=1}^{N_{k}}\ell(h(x_{k,n}),y_{k,n})+\lambda\frac{d(d-1)}{2}M\bigg{]}$
$\displaystyle\leq\mathbb{E}_{D}\mathbb{E}_{D^{\prime}}\mathbb{E}_{\bm{\sigma}}\sup_{h\in\mathcal{H}}\bigg{[}\sum_{k=1}^{K}\frac{p_{k}}{N_{k}}\sum_{n=1}^{N_{k}}\sigma_{k,n}\ell(h(x^{\prime}_{k,n}),y^{\prime}_{k,n})-\sum_{k=1}^{K}\frac{p_{k}}{N_{k}}\sum_{n=1}^{N_{k}}\sigma_{k,n}\ell(h(x_{k,n}),y_{k,n})+\lambda\frac{d(d-1)}{2}M\bigg{]}$
$\displaystyle\leq
2\mathfrak{R}_{\bm{m}}(\mathcal{G},\bm{p})+\lambda\frac{d(d-1)}{2}M.$

Therefore,

$\displaystyle\Phi(D)\leq\sqrt{\frac{1}{2}\sum_{k=1}^{K}(\frac{p_{k}}{N_{k}}M+\lambda\frac{d(d-1)}{2}M)^{2}\log\frac{1}{\delta_{0}}}+2\mathfrak{R}_{\bm{m}}(\mathcal{G},\bm{p})+\lambda\frac{d(d-1)}{2}M.$

∎

###  10.3 Convergence (Strongly Convex)

Our proof is based on the convergence result of FedAvg (li2019convergence).

######  Theorem 6.

Assume Assumptions in the main paper hold and $|\mathcal{S}_{c}|=K$. For
$\gamma,\mu>0$ and $\eta^{(t)}$ is decreasing in a rate of
$\mathcal{O}(\frac{1}{t})$. If $\eta^{(t)}\leq\mathcal{O}(\frac{1}{L})$, we
have

$\displaystyle\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(T)})\bigg{\\}}-H^{*}\leq\frac{L}{2}\frac{1}{\gamma+T}\bigg{\\{}\frac{4\xi}{\epsilon^{2}\mu^{2}}+(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}},$

where
$\xi=8(E-1)^{2}G^{2}+4L\Gamma_{K}+2\frac{\Gamma_{max}}{\eta^{(t)}}+4\sum_{k=1}^{K}p_{k}^{2}\sigma_{k}^{2}$
and
$\Gamma_{max}\coloneqq\sum_{k=1}^{K}p_{k}|(H^{*}-H^{*}_{k})|\geq|\sum_{k=1}^{K}p_{k}(H^{*}-H^{*}_{k})|=|\Gamma_{K}|$.

###### Proof.

For each device $k$, we introduce an intermediate model parameter
$\bm{w}^{(t+1)}_{k}=\bm{\theta}^{(t)}_{k}-\eta^{(t)}\nabla
H_{k}(\bm{\theta}^{(t)}_{k})$. If iteration $t+1$ is in the communication
round, then $\bm{\theta}^{(t+1)}_{k}=\sum_{k=1}^{K}p_{k}\bm{w}^{(t+1)}_{k}$
(i.e., aggregation). Otherwise, $\bm{\theta}^{(t+1)}_{k}=\bm{w}^{(t+1)}_{k}$.
Define $\bm{\bar{w}}^{(t)}=\sum_{k=1}^{K}p_{k}\bm{w}^{(t)}_{k}$ and
$\bm{\bar{\theta}}^{(t)}=\sum_{k=1}^{K}p_{k}\bm{\theta}^{(t)}_{k}$. Also,
define $\bm{g}^{(t)}=\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k})$ and
$\bm{\bar{g}}^{(t)}=\mathbb{E}(\bm{g}^{(t)})=\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\theta}^{(t)}_{k})$.

Denote by $\bm{\theta}^{*}$ the optimal model parameter of the global
objective function $H(\cdot)$. At iteration $t$, we have

$\displaystyle\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t+1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}=\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\eta^{(t)}\bm{g}^{(t)}-\bm{\theta}^{*}-\eta^{(t)}\bm{\bar{g}}^{(t)}+\eta^{(t)}\bm{\bar{g}}^{(t)}\right\rVert^{2}\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}-\eta^{(t)}\bm{\bar{g}}^{(t)}\right\rVert^{2}\bigg{\\}}+\mathbb{E}\bigg{\\{}2\eta^{(t)}\langle\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}-\eta^{(t)}\bm{\bar{g}}^{(t)},\bm{\bar{g}}^{(t)}-\bm{g}^{(t)}\rangle\bigg{\\}}+\mathbb{E}\bigg{\\{}\eta^{(t)2}\left\lVert\bm{g}^{(t)}-\bm{\bar{g}}^{(t)}\right\rVert^{2}\bigg{\\}}$
$\displaystyle=\underbrace{\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}-\eta^{(t)}\bm{\bar{g}}^{(t)}\right\rVert^{2}\bigg{\\}}}_{A}+\underbrace{\mathbb{E}\bigg{\\{}\eta^{(t)2}\left\lVert\bm{g}^{(t)}-\bm{\bar{g}}^{(t)}\right\rVert^{2}\bigg{\\}}}_{B},$

since
$\mathbb{E}\bigg{\\{}2\eta^{(t)}\langle\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}-\eta^{(t)}\bm{\bar{g}}^{(t)},\bm{\bar{g}}^{(t)}-\bm{g}^{(t)}\rangle\bigg{\\}}=0$.
Our remaining work is to bound term $A$ and term $B$.

#### Part I: Bounding Term $A$

We can split term $A$ above into three parts:

$\displaystyle\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}-\eta^{(t)}\bm{\bar{g}}^{(t)}\right\rVert^{2}\bigg{\\}}=\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}\underbrace{-2\eta^{(t)}\mathbb{E}\bigg{\\{}\langle\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*},\bm{\bar{g}}^{(t)}\rangle\bigg{\\}}}_{\text{C}}+\underbrace{\eta^{(t)2}\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{g}}^{(t)}\right\rVert^{2}\bigg{\\}}}_{\text{D}}.$

For part C, We have

C
$\displaystyle=-2\eta^{(t)}\mathbb{E}\bigg{\\{}\langle\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*},\bm{\bar{g}}^{(t)}\rangle\bigg{\\}}=-2\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\langle\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*},\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\rangle\bigg{\\}}$
$\displaystyle=-2\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\langle\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)},\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\rangle\bigg{\\}}-2\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\langle\bm{\theta}_{k}^{(t)}-\bm{\theta}^{*},\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\rangle\bigg{\\}}$

To bound C, we need to use Cauchy-Schwarz inequality, inequality of arithmetic
and geometric means. Specifically, the Cauchy-Schwarz inequality indicates
that

$\displaystyle\langle\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)},\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\rangle\geq-\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)}\right\rVert\left\lVert\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert$

and inequality of arithmetic and geometric means further implies

$\displaystyle-\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)}\right\rVert\left\lVert\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert\geq-\frac{\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)}\right\rVert^{2}+\left\lVert\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}}{2}.$

Therefore, we obtain

C
$\displaystyle=-2\eta^{(t)}\mathbb{E}\bigg{\\{}\langle\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*},\bm{\bar{g}}^{(t)}\rangle\bigg{\\}}=-2\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\langle\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*},\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\rangle\bigg{\\}}$
$\displaystyle=-2\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\langle\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)},\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\rangle\bigg{\\}}-2\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\langle\bm{\theta}_{k}^{(t)}-\bm{\theta}^{*},\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\rangle\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}\eta^{(t)}\sum_{k=1}^{K}p_{k}\frac{1}{\eta^{(t)}}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)}\right\rVert^{2}+\eta^{(t)^{2}}\sum_{k=1}^{K}p_{k}\left\lVert\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}$
$\displaystyle\qquad-2\eta^{(t)}\sum_{k=1}^{K}p_{k}(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}(\bm{\theta}^{*}))-2\eta^{(t)}\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))\mu}{2}\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}},$

where
$-2\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\langle\bm{\theta}_{k}^{(t)}-\bm{\theta}^{*},\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\rangle\bigg{\\}}$ is bounded by the property of
strong convexity of $H_{k}$.

Since $H_{k}$ is $(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L$-smooth, we know

$\displaystyle\left\lVert\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\leq
2(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}^{*})$

and therefore

D
$\displaystyle=\eta^{(t)2}\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{g}}^{(t)}\right\rVert^{2}\bigg{\\}}\leq\eta^{(t)2}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\left\lVert\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\bigg{\\}}$ $\displaystyle\leq
2\eta^{(t)2}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}^{*})\bigg{\\}}$

by convexity of norm.

Therefore, combining C and D, we have

$\displaystyle A$
$\displaystyle=\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}-\eta^{(t)}\bm{\bar{g}}^{(t)}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}+2\eta^{(t)2}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}^{*})\bigg{\\}}$
$\displaystyle\qquad+\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\frac{1}{\eta^{(t)}}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)}\right\rVert^{2}\bigg{\\}}+\eta^{(t)^{2}}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\left\lVert\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\qquad-2\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}(\bm{\theta}^{*}))\bigg{\\}}-2\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))\mu}{2}\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}-\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))\mu\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}+\sum_{k=1}^{K}p_{k}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)}\right\rVert^{2}$
$\displaystyle\qquad+\underbrace{4\eta^{(t)2}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}^{*})\bigg{\\}}-2\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}(\bm{\theta}^{*}))\bigg{\\}}}_{\text{E}}.$

In the last inequality, we simply rearrange other terms and use the fact that
$\left\lVert\nabla H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\leq
2(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}^{*})$ as
aforementioned.

To bound E, we define
$\gamma^{(t)}_{k}=2\eta^{(t)}(1-2(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L\eta^{(t)})$. Assume
$\eta^{(t)}\leq\frac{1}{4(1+\frac{(d-1)}{\min\\{p_{k}|{\cal
A}_{s_{k}}|\\}}\lambda)L}$, then we know $\eta^{(t)}\leq\gamma_{k}^{(t)}\leq
2\eta^{(t)}$.

Therefore, we have

E
$\displaystyle=4\eta^{(t)2}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}^{*})\bigg{\\}}-2\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}(\bm{\theta}^{*}))\bigg{\\}}$
$\displaystyle=4\eta^{(t)2}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}^{*})\bigg{\\}}-2\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}^{*}+H_{k}^{*}-H_{k}(\bm{\theta}^{*}))\bigg{\\}}$
$\displaystyle=-2\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(1-2(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L\eta^{(t)})(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}^{*})\bigg{\\}}+2\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(H_{k}(\bm{\theta}^{*})-H_{k}^{*})\bigg{\\}}$
$\displaystyle=-\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}\gamma^{(t)}_{k}p_{k}(H_{k}(\bm{\theta}_{k}^{(t)})-H^{*}+H^{*}-H^{*}_{k})\bigg{\\}}+2\eta^{(t)}\mathbb{E}\bigg{\\{}H^{*}-\sum_{k=1}^{K}p_{k}H_{k}^{*}\bigg{\\}}$
$\displaystyle=\underbrace{-\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}\gamma_{k}^{(t)}p_{k}(H_{k}(\bm{\theta}_{k}^{(t)})-H^{*})\bigg{\\}}}_{\text{F}}+\underbrace{\
\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}(2\eta^{(t)}-\gamma_{k}^{(t)})p_{k}(H^{*}-H^{*}_{k})\bigg{\\}}}_{\text{G}}.$

If $H^{*}-H^{*}_{k}\geq 0$ for some $k$, then
$(2\eta^{(t)}-\gamma_{k}^{(t)})p_{k}(H^{*}-H^{*}_{k})\leq
2\eta^{(t)}p_{k}(H^{*}-H^{*}_{k})$. If $H^{*}-H^{*}_{k}<0$ otherwise, then
$(2\eta^{(t)}-\gamma_{k}^{(t)})p_{k}(H^{*}-H^{*}_{k})$ is negative and
$(2\eta^{(t)}-\gamma_{k}^{(t)})p_{k}(H^{*}-H^{*}_{k})\leq-2\eta^{(t)}p_{k}(H^{*}-H^{*}_{k})$.
Therefore, by definition of $\Gamma_{max}$,

$\displaystyle\text{G}\leq
2\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}|H^{*}-H^{*}_{k}|\bigg{\\}}=2\eta^{(t)}\Gamma_{max}.$

The remaining goal of Part I is to bound term F. Note that

F
$\displaystyle=-\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}\gamma_{k}^{(t)}p_{k}(H_{k}(\bm{\theta}_{k}^{(t)})-H^{*})\bigg{\\}}$
$\displaystyle=-\mathbb{E}\bigg{\\{}\bigg{(}\sum_{k=1}^{K}p_{k}\gamma_{k}^{(t)}(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}(\bm{\bar{\theta}}^{(t)}))+\sum_{k=1}^{K}p_{k}\gamma_{k}^{(t)}(H_{k}(\bm{\bar{\theta}}^{(t)})-H^{*})\bigg{)}\bigg{\\}}$
$\displaystyle\leq-\mathbb{E}\bigg{\\{}\bigg{(}\sum_{k=1}^{K}p_{k}\gamma_{k}^{(t)}\langle\nabla
H_{k}(\bm{\bar{\theta}}^{(t)}),\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t)}\rangle+\sum_{k=1}^{K}p_{k}\gamma_{k}^{(t)}(H_{k}(\bm{\bar{\theta}}^{(t)})-H^{*})\bigg{)}\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}\frac{1}{2}\gamma_{k}^{(t)}p_{k}\bigg{[}\eta^{(t)}\left\lVert\nabla
H_{k}(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}+\frac{1}{\eta^{(t)}}\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t)}\right\rVert^{2}\bigg{]}-\sum_{k=1}^{K}p_{k}\gamma_{k}^{(t)}(H_{k}(\bm{\bar{\theta}}^{(t)})-H^{*})\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}\gamma_{k}^{(t)}p_{k}\bigg{[}\eta^{(t)}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H_{k}(\bm{\bar{\theta}}^{(t)})-H_{k}^{*})+\frac{1}{2\eta^{(t)}}\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t)}\right\rVert^{2}\bigg{]}$
$\displaystyle\qquad-\sum_{k=1}^{K}p_{k}\gamma_{k}^{(t)}(H_{k}(\bm{\bar{\theta}}^{(t)})-H^{*})\bigg{\\}}.$

In the second inequality, we again use the Cauchy–Schwarz inequality and
Inequality of arithmetic and geometric means. In the last inequality, we use
the fact that $\left\lVert\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\leq
2(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}^{*})$.

Since $\eta^{(t)}\leq\gamma_{k}^{(t)}\leq 2\eta^{(t)}$, we can bound E as

$\displaystyle\text{E}\leq\text{F}+\mathbb{E}\bigg{\\{}2\eta^{(t)}\Gamma_{max}\bigg{\\}}$
$\displaystyle=(\eta^{(t)}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L-1)\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}\gamma_{k}^{(t)}p_{k}\bigg{[}(H_{k}(\bm{\bar{\theta}}^{(t)})-H^{*})\bigg{]}\bigg{\\}}$
$\displaystyle\qquad+\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}\eta^{(t)}\gamma_{k}^{(t)}p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H^{*}-H^{*}_{k})\bigg{\\}}$
$\displaystyle\qquad+\frac{1}{2\eta^{(t)}}\sum_{k=1}^{K}\gamma_{k}^{(t)}p_{k}\bigg{\\{}\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t)}\right\rVert^{2}\bigg{\\}}+2\eta^{(t)}\Gamma_{max}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}\eta^{(t)}\gamma_{k}^{(t)}p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H^{*}-H^{*}_{k})\bigg{\\}}+\frac{1}{2\eta^{(t)}}\sum_{k=1}^{K}\gamma_{k}^{(t)}p_{k}\mathbb{E}\bigg{\\{}\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t)}\right\rVert^{2}\bigg{\\}}+2\eta^{(t)}\Gamma_{max}$
$\displaystyle\leq\sum_{k=1}^{K}p_{k}\mathbb{E}\bigg{\\{}\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t)}\right\rVert^{2}\bigg{\\}}+\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}\eta^{(t)}\gamma_{k}^{(t)}p_{k}(1+\frac{d-1}{p_{k}|{\cal
A}_{s_{k}}|}\lambda)L(H^{*}-H^{*}_{k})\bigg{\\}}+2\eta^{(t)}\Gamma_{max}$
$\displaystyle\leq\sum_{k=1}^{K}p_{k}\mathbb{E}\bigg{\\{}\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t)}\right\rVert^{2}\bigg{\\}}+4\eta^{(t)2}L\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(H^{*}-H^{*}_{k})\bigg{\\}}+2\eta^{(t)}\Gamma_{max}$
$\displaystyle=\sum_{k=1}^{K}p_{k}\mathbb{E}\bigg{\\{}\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t)}\right\rVert^{2}\bigg{\\}}+4\eta^{(t)2}L\Gamma_{K}+2\eta^{(t)}\Gamma_{max}$

The second inequality holds because $(\eta^{(t)}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L-1)\leq 0$ and the fourth inequality uses the
fact that $1+\frac{d-1}{p_{k}|{\cal A}_{s_{k}}|}\lambda\leq 2$ based on the
constraint of $\lambda$.

Therefore,

$\displaystyle A$
$\displaystyle\leq\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}-\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))\mu\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}+\sum_{k=1}^{K}p_{k}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)}\right\rVert^{2}+\text{E}$
$\displaystyle\leq
2\sum_{k=1}^{K}p_{k}\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)}\right\rVert^{2}\bigg{\\}}+4\eta^{(t)2}L\Gamma_{K}+2\eta^{(t)}\Gamma_{max}+\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\qquad-\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(1-\frac{d-1}{p_{k}|{\cal
A}_{s_{k}}|}\lambda)\mu\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq
2\sum_{k=1}^{K}p_{k}\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)}\right\rVert^{2}\bigg{\\}}+4\eta^{(t)2}L\Gamma_{K}+2\eta^{(t)}\Gamma_{max}+\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\qquad-\eta^{(t)}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p^{2}_{k}(1-\frac{d-1}{p_{k}|{\cal
A}_{s_{k}}|}\lambda)\mu\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq
2\sum_{k=1}^{K}p_{k}\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)}\right\rVert^{2}\bigg{\\}}+4\eta^{(t)2}L\Gamma_{K}+2\eta^{(t)}\Gamma_{max}+\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\qquad-\eta^{(t)}\mathbb{E}\bigg{\\{}(1-\frac{d-1}{\min\\{p_{k}|{\cal
A}_{s_{k}}|\\}}\lambda)\mu\frac{1}{K}\left\lVert\sum_{k=1}^{K}p_{k}\bm{\theta}_{k}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle=2\sum_{k=1}^{K}p_{k}\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)}\right\rVert^{2}\bigg{\\}}+4\eta^{(t)2}L\Gamma_{K}+2\eta^{(t)}\Gamma_{max}+(1-\eta^{(t)}(1-\frac{d-1}{\min\\{p_{k}|{\cal
A}_{s_{k}}|\\}}\lambda)\frac{\mu}{K})\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$

The third inequality uses the fact that $0\leq p_{k}\leq 1$ and
$-p_{k}^{2}\geq-p_{k}$. The last inequality uses the fact that
$\left\lVert\sum_{k=1}^{K}p_{k}\bm{\theta}_{k}\right\rVert^{2}\leq
K\sum_{k=1}^{K}\left\lVert
p_{k}\bm{\theta}_{k}\right\rVert^{2}=K\sum_{k=1}^{K}p^{2}_{k}\left\lVert\bm{\theta}_{k}\right\rVert^{2}$
and $1-\frac{d-1}{p_{k}|{\cal A}_{s_{k}}|}\lambda\geq
1-\frac{d-1}{\min\\{p_{k}|{\cal A}_{s_{k}}|\\}}\lambda$.

#### Part II: Bounding Term
$\sum_{k=1}^{K}p_{k}\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)}\right\rVert^{2}\bigg{\\}}$
in Term A

For any iteration $t\geq 0$, denote by $t_{0}\leq t$ the index of previous
communication iteration before $t$. Since the FL algorithm requires one
communication each $E$ steps, we know $t-t_{0}\leq E-1$ and
$\bm{\theta}_{k}^{(t_{0})}=\bm{\bar{\theta}}^{(t_{0})}$. Assume
$\eta^{(t)}\leq 2\eta^{(t+E)}$. Since $\eta^{(t)}$ is decreasing, we have

$\displaystyle\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)}\right\rVert^{2}\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\left\lVert(\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t_{0})})-(\bm{\bar{\theta}}^{(t)}-\bm{\bar{\theta}}^{(t_{0})})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t_{0})}\right\rVert^{2}\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\left\lVert\sum_{t=0}^{t-1}\eta^{(t)}\nabla
F_{k}(\bm{\theta}_{k}^{(t)},D_{k}^{(t)})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(t-t_{0})\sum_{t=0}^{t-1}\eta^{(t)2}\left\lVert\nabla
F_{k}(\bm{\theta}_{k}^{(t)},D_{k}^{(t)})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\sum_{k=1}^{K}p_{k}\sum_{t=t_{0}}^{t-1}(E-1)\eta^{(t)2}G^{2}\leq\sum_{k=1}^{K}p_{k}\sum_{t=t_{0}}^{t-1}(E-1)\eta^{(t_{0})2}G^{2}$
$\displaystyle\leq\sum_{k=1}^{K}p_{k}(E-1)^{2}\eta^{(t_{0})2}G^{2}\leq
4\eta^{(t)2}(E-1)^{2}G^{2}.$

#### Part III: Bounding Term B

By assumption, it is easy to show

$\displaystyle\mathbb{E}\bigg{\\{}\eta^{(t)2}\left\lVert\bm{g}^{(t)}-\bm{\bar{g}}^{(t)}\right\rVert^{2}\bigg{\\}}\leq\eta^{(t)2}\sum_{k=1}^{K}p_{k}^{2}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}\sigma_{k}^{2}.$

#### Part IV: Proving Convergence

So far, we have shown that

$\displaystyle\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t+1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}\leq\text{A}+\text{B}$
$\displaystyle\leq
8\eta^{(t)2}(E-1)^{2}G^{2}+4\eta^{(t)2}L\Gamma_{K}+2\eta^{(t)}\Gamma_{max}+(1-\eta^{(t)}(1-\frac{d-1}{p_{k}|{\cal
A}_{s_{k}}|}\lambda)\mu)\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\qquad+\eta^{(t)2}\sum_{k=1}^{K}p_{k}^{2}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}\sigma_{k}^{2}$
$\displaystyle=(1-\eta^{(t)}(1-\frac{d-1}{\min\\{p_{k}|{\cal
A}_{s_{k}}|\\}}\lambda)\frac{\mu}{K})\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}+\eta^{(t)^{2}}\xi$

where
$\xi=8(E-1)^{2}G^{2}+4L\Gamma_{K}+2\frac{\Gamma_{max}}{\eta^{(t)}}+\sum_{k=1}^{K}p_{k}^{2}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}\sigma_{k}^{2}$.

Let $\eta^{(t)}=\frac{\beta}{t+\gamma}$ with
$\beta>\frac{1}{(1-\frac{d-1}{min\\{p_{k}|{\cal
A}_{s_{k}}|\\}}\lambda)\frac{\mu}{K}}$ and $\gamma>0$. Define
$\epsilon\coloneqq(1-\frac{d-1}{min\\{p_{k}|{\cal A}_{s_{k}}|\\}}\lambda)$.
Let
$v=\max\\{\frac{\beta^{2}\xi}{\beta\epsilon\mu-1},(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\\}$.
We will show that
$\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\leq\frac{v}{\gamma+t}$
by induction. For $t=1$, we have
$\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\leq(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\leq\frac{v}{\gamma+1}$.
Now assume this is true for some $t$, then

$\displaystyle\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t+1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq(1-\eta^{(t)}\epsilon\mu)\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}+\eta^{(t)^{2}}\xi$
$\displaystyle\leq(1-\frac{\beta\epsilon\mu}{t+\gamma})\frac{v}{t+\gamma}+\frac{\beta^{2}\xi}{(t+\gamma)^{2}}$
$\displaystyle=\frac{t+\gamma-1}{(t+\gamma)^{2}}v+\frac{\beta^{2}\xi}{(t+\gamma)^{2}}-\frac{\beta\epsilon\mu-1}{(t+\gamma)^{2}}v.$

It is easy to show
$\frac{t+\gamma-1}{(t+\gamma)^{2}}v+\frac{\beta^{2}\xi}{(t+\gamma)^{2}}-\frac{\beta\epsilon\mu-1}{(t+\gamma)^{2}}v\leq\frac{v}{t+\gamma+1}$
by definition of $v$. Therefore, we proved
$\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\leq\frac{v}{\gamma+t}$.

By definition, we know $H$ is
$\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))}{2}L$-smooth. Therefore,

$\displaystyle\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})\bigg{\\}}-H^{*}$
$\displaystyle\leq\frac{\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))}{2}L}{2}\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\qquad\leq\frac{\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))}{2}L}{2}\frac{v}{\gamma+t}.$

By choosing $\beta=\frac{2}{\epsilon\frac{\mu}{K}}$ We have

$\displaystyle
v=\max\\{\frac{\beta^{2}\xi}{\beta\epsilon\mu-1},(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\\}\leq\frac{\beta^{2}\xi}{\beta\epsilon\mu-1}+(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\leq\frac{4\xi}{\epsilon^{2}\mu^{2}}+(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}.$

Therefore,

$\displaystyle\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(T)})\bigg{\\}}-H^{*}$
$\displaystyle\leq\frac{\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))}{2}L}{2}\frac{1}{\gamma+T}\bigg{\\{}\frac{4\xi}{\epsilon^{2}\mu^{2}}+(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\frac{\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda(d-1)}{p_{k}|{\cal
A}_{s_{k}}|})}{2}L}{2}\frac{1}{\gamma+T}\bigg{\\{}\frac{4\xi}{\epsilon^{2}\mu^{2}}+(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\frac{L}{2}\frac{1}{\gamma+T}\bigg{\\{}\frac{4\xi}{\epsilon^{2}\mu^{2}}+(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}.$

We thus proved our convergence result. ∎

######  Theorem 7.

Assume at each communication round, central server sampled a fraction $\alpha$
of devices and those local devices are sampled according to the sampling
probability $p_{k}$. Additionally, assume Assumptions in the main paper hold.
For $\gamma,\mu,\epsilon>0$, we have

$\displaystyle\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(T)})\bigg{\\}}-H^{*}\leq\frac{L}{2}\frac{1}{\gamma+T}\bigg{\\{}\frac{4(\xi+\tau)}{\epsilon^{2}\mu^{2}}+(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}},$

$\tau=\frac{E^{2}}{\left\lceil\alpha
K\right\rceil}\sum_{k=1}^{K}p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}G^{2}$.

###### Proof.

$\displaystyle\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t+1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t+1)}-\bar{\bm{w}}^{(t+1)}+\bar{\bm{w}}^{(t+1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t+1)}-\bar{\bm{w}}^{(t+1)}\right\rVert^{2}+\left\lVert\bar{\bm{w}}^{(t+1)}-\bm{\theta}^{*}\right\rVert^{2}+2\langle\bm{\bar{\theta}}^{(t+1)}-\bar{\bm{w}}^{(t+1)},\bar{\bm{w}}^{(t+1)}-\bm{\theta}^{*}\rangle\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t+1)}-\bar{\bm{w}}^{(t+1)}\right\rVert^{2}+\left\lVert\bar{\bm{w}}^{(t+1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}.$

Note that the expectation is taken over subset $\mathcal{S}_{c}$.

#### Part I: Bounding Term
$\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t+1)}-\bar{\bm{w}}^{(t+1)}\right\rVert^{2}\bigg{\\}}$

Assume $\left\lceil\alpha K\right\rceil$ number of local devices are sampled
according to sampling probability $p_{k}$. During the communication round, we
have $\bar{\bm{\theta}}^{t+1}=\frac{1}{\left\lceil\alpha
K\right\rceil}\sum_{l=1}^{\left\lceil\alpha K\right\rceil}\bm{w}^{(t+1)}_{l}$.
Therefore,

$\displaystyle\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t+1)}-\bar{\bm{w}}^{(t+1)}\right\rVert^{2}\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}\frac{1}{\left\lceil\alpha
K\right\rceil^{2}}\left\lVert\sum_{l=1}^{\left\lceil\alpha
K\right\rceil}\bm{w}^{(t+1)}_{l}-\bar{\bm{w}}^{(t+1)}\right\rVert^{2}\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}\frac{1}{\left\lceil\alpha
K\right\rceil^{2}}\sum_{l=1}^{\left\lceil\alpha
K\right\rceil}\left\lVert\bm{w}^{(t+1)}_{l}-\bar{\bm{w}}^{(t+1)}\right\rVert^{2}\bigg{\\}}$
$\displaystyle=\frac{1}{\left\lceil\alpha
K\right\rceil}\sum_{k=1}^{K}p_{k}\left\lVert\bm{w}^{(t+1)}_{k}-\bar{\bm{w}}^{(t+1)}\right\rVert^{2}.$

We know

$\displaystyle\sum_{k=1}^{K}p_{k}\left\lVert\bm{w}^{(t+1)}_{k}-\bar{\bm{w}}^{(t+1)}\right\rVert^{2}=\sum_{k=1}^{K}p_{k}\left\lVert(\bm{w}^{(t+1)}_{k}-\bar{\bm{\theta}}^{(t_{0})})-(\bar{\bm{w}}^{(t+1)}-\bar{\bm{\theta}}^{(t_{0})})\right\rVert^{2}\leq\sum_{k=1}^{K}p_{k}\left\lVert(\bm{w}^{(t+1)}_{k}-\bar{\bm{\theta}}^{(t_{0})})\right\rVert^{2},$

where $t_{0}=t-E+1$. Similarly,

$\displaystyle\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t+1)}-\bar{\bm{w}}^{(t+1)}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\frac{1}{\left\lceil\alpha
K\right\rceil}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\left\lVert(\bm{w}^{(t+1)}_{k}-\bar{\bm{\theta}}^{(t_{0})})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\frac{1}{\left\lceil\alpha
K\right\rceil}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\left\lVert(\bm{w}^{(t+1)}_{k}-\bm{\theta}_{k}^{(t_{0})})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\frac{1}{\left\lceil\alpha
K\right\rceil}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}E\sum_{m=t_{o}}^{t}\left\lVert\eta^{(m)}\nabla
H_{k}(\bm{\theta}_{k}^{(m)},D_{k}^{(m)})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\frac{E^{2}\eta^{(t_{0})2}}{\left\lceil\alpha
K\right\rceil}\sum_{k=1}^{K}p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}G^{2}$
$\displaystyle\leq\frac{E^{2}\eta^{(t)2}}{\left\lceil\alpha
K\right\rceil}\sum_{k=1}^{K}p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}G^{2}$

using the fact that $\eta^{(t)}$ is non-increasing in $t$.

#### Part II: Convergence Result

As aforementioned,

$\displaystyle\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t+1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t+1)}-\bar{\bm{w}}^{(t+1)}\right\rVert^{2}+\left\lVert\bar{\bm{w}}^{(t+1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\frac{E^{2}\eta^{(t)2}}{\left\lceil\alpha
K\right\rceil}\sum_{k=1}^{K}p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}G^{2}+(1-\eta^{(t)}\epsilon\frac{\mu}{K})\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}+\eta^{(t)^{2}}\xi$
$\displaystyle=(1-\eta^{(t)}\epsilon\frac{\mu}{K})\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}+\eta^{(t)^{2}}\bigg{(}\xi+\frac{E^{2}}{\left\lceil\alpha
K\right\rceil}\sum_{k=1}^{K}p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}G^{2}\bigg{)}.$

Let $\tau=\frac{E^{2}}{\left\lceil\alpha
K\right\rceil}\sum_{k=1}^{K}p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}G^{2}$. Let
$\eta^{(t)}=\frac{\beta}{t+\gamma}$ with
$\beta>\frac{1}{\epsilon\frac{\mu}{K}}$ and $\gamma>0$. Let
$v=\max\\{\frac{\beta^{2}(\xi+\tau)}{\beta\epsilon\mu-1},(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\\}$.
Similar to the full device participation scenario, we can show that
$\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}\leq\frac{v}{\gamma+t}$
by induction.

By definition, we know $H$ is
$\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))}{2}L$-smooth. Therefore,

$\displaystyle\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})\bigg{\\}}-H^{*}$
$\displaystyle\leq\frac{\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))}{2}L}{2}\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\frac{\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))}{2}L}{2}\frac{v}{\gamma+t}.$

By choosing $\beta=\frac{2}{\epsilon\frac{\mu}{K}}$ We have

$\displaystyle
v=\max\\{\frac{\beta^{2}\xi}{\beta\epsilon\mu-1},(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\\}\leq\frac{\beta^{2}\xi}{\beta\epsilon\mu-1}+(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\leq\frac{4\xi}{\epsilon^{2}\mu^{2}}+(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}.$

Therefore,

$\displaystyle\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(T)})\bigg{\\}}-H^{*}$
$\displaystyle\leq\frac{\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))}{2}L}{2}\frac{1}{\gamma+T}\bigg{\\{}\frac{4(\xi+\tau)}{\epsilon^{2}\mu^{2}}+(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\frac{L}{2}\frac{1}{\gamma+T}\bigg{\\{}\frac{4(\xi+\tau)}{\epsilon^{2}\mu^{2}}+(\gamma+1)\left\lVert\bm{\bar{\theta}}^{(1)}-\bm{\theta}^{*}\right\rVert^{2}\bigg{\\}}$

∎

###  10.4 Convergence (Non-convex)

######  Lemma 3.

If $\eta^{(t)}\leq\frac{2}{L}$, then
$\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})\bigg{\\}}\leq\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(1)})\bigg{\\}}$.

###### Proof.

$\displaystyle\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t+1)})\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)}-\eta^{(t)}\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k}))\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)}-\eta^{(t)}\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\bar{\theta}}^{(t)},D^{(t)}_{k}))\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)}-\eta^{(t)}g^{(t)}(\bm{\bar{\theta}}^{(t)}))\bigg{\\}}$

Here we used the fact that $\bm{\bar{\theta}}^{(t)}=\bm{\theta}^{(t)}_{k}$
since the aggregated model parameter has been distributed to local devices. By
Taylor’s theorem, there exists a $\bm{w}^{(t)}$ such that

$\displaystyle\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t+1)})\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})-\eta^{(t)}g^{(t)}(\bm{\bar{\theta}}^{(t)})^{T}g^{(t)}(\bm{\bar{\theta}}^{(t)})+\frac{1}{2}(\eta^{(t)}g^{(t)}(\bm{\bar{\theta}}^{(t)}))^{T}g^{(t)}(\bm{w}^{(t)})(\eta^{(t)}g^{(t)}(\bm{\bar{\theta}}^{(t)}))\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})-\eta^{(t)}g^{(t)}(\bm{\bar{\theta}}^{(t)})^{T}g^{(t)}(\bm{\bar{\theta}}^{(t)})+\eta^{(t)2}\frac{\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))}{2}L}{2}\left\lVert
g^{(t)}(\bm{\bar{\theta}}^{(t)}\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})\bigg{\\}}-\eta^{(t)}\left\lVert
g^{(t)}(\bm{\bar{\theta}}^{(t)}\right\rVert^{2}+\eta^{(t)2}\frac{L}{2}\left\lVert
g^{(t)}(\bm{\bar{\theta}}^{(t)}\right\rVert^{2}$

since $H$ is $\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))}{2}L$-smooth. It can be shown that if
$\eta^{(t)}\leq\frac{2}{L}$, we have

$\displaystyle-\eta^{(t)}\left\lVert
g^{(t)}(\bm{\bar{\theta}}^{(t)}\right\rVert^{2}+\eta^{(t)2}\frac{L}{2}\left\lVert
g^{(t)}(\bm{\bar{\theta}}^{(t)}\right\rVert^{2}\leq 0.$

Therefore, By choosing $\eta^{(t)}\leq\frac{2}{L}$, we proved
$\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})\bigg{\\}}\leq\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(1)})\bigg{\\}}$.
∎

######  Theorem 8.

Assume Assumptions in the main paper hold and $|\mathcal{S}_{c}|=K$. If
$\eta^{(t)}=\mathcal{O}(\frac{1}{\sqrt{t}})$ and
$\eta^{(t)}\leq\mathcal{O}(\frac{1}{L})$, then for $>0$

$\displaystyle\min_{t=1,\ldots,T}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}\leq\frac{1}{\sqrt{T}}\bigg{\\{}2(1+2KL^{2}\sum_{t=1}^{T}\eta^{(t)2})\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(1)})-H^{*}\bigg{\\}}+2\sum_{t=1}^{T}\xi^{(t)}\bigg{\\}},$

where
$\xi^{(t)}=2KL^{2}\eta^{(t)2}\Gamma_{K}+(8\eta^{(t)3}KL^{2}(E-1)+8KL\eta^{(t)2}+4(2+4L)KL\eta^{(t)4}(E-1))G^{2}+(2L\eta^{(t)2}+8KL\eta^{(t)2})\sum_{k=1}^{K}p_{k}\sigma_{k}^{2}$

###### Proof.

Since $H$ is $\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))}{2}L$-smooth, we have

$\displaystyle\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t+1)})\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})\bigg{\\}}+\underbrace{\mathbb{E}\bigg{\\{}\langle\nabla
H(\bm{\bar{\theta}}^{(t)}),\bm{\bar{\theta}}^{(t+1)}-\bm{\bar{\theta}}^{(t)}\rangle\bigg{\\}}}_{\text{A}}+$
$\displaystyle\qquad\frac{\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))}{2}L}{2}\underbrace{\mathbb{E}\bigg{\\{}\left\lVert\bm{\bar{\theta}}^{(t+1)}-\bm{\bar{\theta}}^{(t)}\right\rVert^{2}\bigg{\\}}}_{\text{B}}.$

#### Part I: Bounding Term A

We have

A $\displaystyle=-\eta^{(t)}\mathbb{E}\bigg{\\{}\langle\nabla
H(\bm{\bar{\theta}}^{(t)}),\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\theta}_{k}^{(t)},D_{k}^{(t)})\rangle\bigg{\\}}=-\eta^{(t)}\mathbb{E}\bigg{\\{}\langle\nabla
H(\bm{\bar{\theta}}^{(t)}),\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\rangle\bigg{\\}}$
$\displaystyle=-\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}-\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\bigg{\\}}+\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})-\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\bigg{\\}}$
$\displaystyle=-\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}-\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\bigg{\\}}+\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\bar{\theta}}^{(t)})-\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq-\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}-\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\bigg{\\}}+\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}K\sum_{k=1}^{K}p_{k}\left\lVert\nabla
H_{k}(\bm{\bar{\theta}}^{(t)})-\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq-\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}-\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\bigg{\\}}+$
$\displaystyle\qquad\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}K\sum_{k=1}^{K}p_{k}((1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L)^{2}\underbrace{\left\lVert\bm{\bar{\theta}}^{(t)}-\bm{\theta}_{k}^{(t)}\right\rVert^{2}}_{\text{C}}\bigg{\\}}.$

In the convex setting, we proved that

$\displaystyle\text{C}\leq 4\eta^{(t)2}(E-1)G^{2}.$

This is also true for the non-convex setting since we do not use any property
of convex functions.

#### Part II: Bounding Term B

We have

B
$\displaystyle=\mathbb{E}\bigg{\\{}\left\lVert\eta^{(t)}g^{(t)}\right\rVert^{2}\bigg{\\}}=\mathbb{E}\bigg{\\{}\left\lVert\eta^{(t)}\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\theta}_{k}^{(t)},D_{k}^{(t)})\right\rVert^{2}\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}\left\lVert\eta^{(t)}\sum_{k=1}^{K}p_{k}(\nabla
H_{k}(\bm{\theta}_{k}^{(t)},D_{k}^{(t)})-\nabla
H_{k}(\bm{\theta}_{k}^{(t)}))\right\rVert^{2}\bigg{\\}}+\mathbb{E}\bigg{\\{}\left\lVert\eta^{(t)}\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\bigg{\\}}$
$\displaystyle=\eta^{(t)2}\sum_{k=1}^{K}p_{k}^{2}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H_{k}(\bm{\theta}_{k}^{(t)},D_{k}^{(t)})-\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\bigg{\\}}+\mathbb{E}\bigg{\\{}\left\lVert\eta^{(t)}\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\eta^{(t)2}\sum_{k=1}^{K}p_{k}^{2}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}\sigma_{k}^{2}+\eta^{(t)2}\mathbb{E}\bigg{\\{}K\sum_{k=1}^{K}p_{k}^{2}\left\lVert\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\bigg{\\}}.$

Since $H_{k}$ is $(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L$-smooth, we know

$\displaystyle\left\lVert\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\leq
2(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}^{*}).$

Therefore,

B
$\displaystyle\leq\eta^{(t)2}\sum_{k=1}^{K}p_{k}^{2}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}\sigma_{k}^{2}+$
$\displaystyle\qquad\eta^{(t)2}\mathbb{E}\bigg{\\{}K\sum_{k=1}^{K}2p_{k}^{2}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}^{*})\bigg{\\}}$
$\displaystyle=\eta^{(t)2}\sum_{k=1}^{K}p_{k}^{2}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}\sigma_{k}^{2}+$
$\displaystyle\qquad\eta^{(t)2}\mathbb{E}\bigg{\\{}K\sum_{k=1}^{K}2p_{k}^{2}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H_{k}(\bm{\theta}_{k}^{(t)})-H^{*}+H^{*}-H_{k}^{*})\bigg{\\}}$
$\displaystyle\leq\eta^{(t)2}\sum_{k=1}^{K}p_{k}^{2}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}\sigma_{k}^{2}+$
$\displaystyle\qquad\eta^{(t)2}\mathbb{E}\bigg{\\{}K\sum_{k=1}^{K}2p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H_{k}(\bm{\theta}_{k}^{(t)})-H^{*}+H^{*}-H_{k}^{*})\bigg{\\}}$

since $0\leq p_{k}\leq 1$ and $p_{k}^{2}\leq p_{k}$.

Therefore,

$\displaystyle\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t+1)})\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})\bigg{\\}}-\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}\underbrace{-\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\sum_{k=1}^{K}p_{k}\nabla
H_{k}(\bm{\theta}_{k}^{(t)})\right\rVert^{2}\bigg{\\}}}_{\text{D}<0}+$
$\displaystyle\qquad\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}K\sum_{k=1}^{K}p_{k}((1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L)^{2}4\eta^{(t)2}(E-1)G^{2}\bigg{\\}}$
$\displaystyle\qquad+\frac{\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))}{2}L}{2}\bigg{[}\eta^{(t)2}\sum_{k=1}^{K}p_{k}^{2}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}\sigma_{k}^{2}$
$\displaystyle\qquad+\eta^{(t)2}\mathbb{E}\bigg{\\{}K\sum_{k=1}^{K}2p_{k}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L(H_{k}(\bm{\theta}_{k}^{(t)})-H^{*}+H^{*}-H_{k}^{*})\bigg{\\}}\bigg{]}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})\bigg{\\}}-\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\qquad\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}K\sum_{k=1}^{K}p_{k}((1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L)^{2}4\eta^{(t)2}(E-1)G^{2}\bigg{\\}}$
$\displaystyle\qquad+\frac{\sum_{k=1}^{K}p_{k}\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))}{2}L}{2}\bigg{[}\eta^{(t)2}\sum_{k=1}^{K}p_{k}^{2}(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}\sigma_{k}^{2}+$
$\displaystyle\qquad\underbrace{4KL\eta^{(t)2}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(H_{k}(\bm{\theta}_{k}^{(t)})-H^{*})+\sum_{k=1}^{K}p_{k}(H^{*}-H_{k}^{*})\bigg{\\}}}_{\text{E}}\bigg{]}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})\bigg{\\}}-\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}+\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}K\sum_{k=1}^{K}4p_{k}L^{2}4\eta^{(t)2}(E-1)G^{2}\bigg{\\}}$
$\displaystyle\qquad+\frac{L}{2}\bigg{[}\eta^{(t)2}\sum_{k=1}^{K}4p_{k}^{2}\sigma_{k}^{2}+\underbrace{4KL\eta^{(t)2}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(H_{k}(\bm{\theta}_{k}^{(t)})-H^{*})+\sum_{k=1}^{K}p_{k}(H^{*}-H_{k}^{*})\bigg{\\}}}_{\text{E}}\bigg{]}$

Here

E
$\displaystyle=4KL\eta^{(t)2}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(H_{k}(\bm{\theta}_{k}^{(t)})-H^{*})\bigg{\\}}+4KL\eta^{(t)2}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(H^{*}-H_{k}^{*})\bigg{\\}}$
$\displaystyle=4KL\eta^{(t)2}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}(\bm{\bar{\theta}}^{(t)}))\bigg{\\}}+4KL\eta^{(t)2}\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(H_{k}(\bm{\bar{\theta}}^{(t)})-H^{*})\bigg{\\}}+4KL\eta^{(t)2}\Gamma_{K}$
$\displaystyle=4KL\eta^{(t)2}\underbrace{\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}(\bm{\bar{\theta}}^{(t)}))\bigg{\\}}}_{\text{F}}+4KL\eta^{(t)2}\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})-H^{*}\bigg{\\}}+4KL\eta^{(t)2}\Gamma_{K}.$

We can bound term F as

F
$\displaystyle=\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(H_{k}(\bm{\theta}_{k}^{(t)})-H_{k}(\bm{\bar{\theta}}^{(t)}))\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(\langle\nabla
H_{k}(\bm{\bar{\theta}}^{(t)}),\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t)}\rangle+\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L}{2}\underbrace{\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t)}\right\rVert^{2}}_{\leq
4\eta^{(t)2}(E-1)G^{2}})\bigg{\\}}$

where we use the fact that $H_{k}$ is $(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L$-smooth. To bound the inner product, we again
use the inequality of arithmetic and geometric means and Cauchy–Schwarz
inequality:

$\displaystyle\langle\nabla
H_{k}(\bm{\bar{\theta}}^{(t)}),\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t)}\rangle\leq\left\lVert\nabla
H_{k}(\bm{\bar{\theta}}^{(t)})\right\rVert\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t)}\right\rVert\leq\frac{\left\lVert\nabla
H_{k}(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}+\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t)}\right\rVert^{2}}{2}.$

It can be shown that

$\displaystyle\mathbb{E}\bigg{\\{}\left\lVert\nabla
H_{k}(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}$
$\displaystyle=\mathbb{E}\bigg{\\{}\left\lVert\nabla
F_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k})\right\rVert\bigg{\\}}^{2}+\mathbb{E}\bigg{\\{}\left\lVert\nabla
F_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k})-\nabla
F_{k}(\bm{\theta}^{(t)}_{k})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}\left\lVert\nabla
F_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k})\right\rVert^{2}\bigg{\\}}+\mathbb{E}\bigg{\\{}\left\lVert\nabla
F_{k}(\bm{\theta}^{(t)}_{k},D^{(t)}_{k})-\nabla
F_{k}(\bm{\theta}^{(t)}_{k})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))^{2}(G^{2}+\sigma_{k}^{2})\leq
4(G^{2}+\sigma_{k}^{2})$

Therefore, we can simplify F as

F
$\displaystyle\leq\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(\frac{\left\lVert\nabla
H_{k}(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}+\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t)}\right\rVert^{2}}{2}+\frac{(1+\frac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}))L}{2}\underbrace{\left\lVert\bm{\theta}_{k}^{(t)}-\bm{\bar{\theta}}^{(t)}\right\rVert^{2}}_{\leq
4\eta^{(t)2}(E-1)G^{2}})\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}(\frac{4(G^{2}+\sigma_{k}^{2})+4\eta^{(t)2}(E-1)G^{2}}{2}+4L\eta^{(t)2}(E-1)G^{2})\bigg{\\}}$
$\displaystyle=2\mathbb{E}\bigg{\\{}\sum_{k=1}^{K}p_{k}\sigma_{k}^{2}\bigg{\\}}+2G^{2}+(2+4L)\eta^{(t)2}(E-1)G^{2}$

Combining with E, we obtain

E $\displaystyle\leq
4KL\eta^{(t)2}\bigg{(}2\sum_{k=1}^{K}p_{k}\sigma_{k}^{2}+2G^{2}+(2+4L)\eta^{(t)2}(E-1)G^{2}\bigg{)}+4KL\eta^{(t)2}\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})-H^{*}\bigg{\\}}+4KL\eta^{(t)2}\Gamma_{K}$

#### Part III: Proving Convergence

Therefore,

$\displaystyle\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})\bigg{\\}}-\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t+1)})\bigg{\\}}+\frac{1}{2}\eta^{(t)3}\mathbb{E}\bigg{\\{}K\sum_{k=1}^{K}4p_{k}L^{2}4(E-1)G^{2}\bigg{\\}}+$
$\displaystyle\frac{L}{2}\bigg{[}\eta^{(t)2}\sum_{k=1}^{K}4p_{k}^{2}\sigma_{k}^{2}+4KL\eta^{(t)2}\bigg{(}2\sum_{k=1}^{K}p_{k}\sigma_{k}^{2}+2G^{2}+(2+4L)\eta^{(t)2}(E-1)G^{2}\bigg{)}+4KL\eta^{(t)2}\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})-H^{*}\bigg{\\}}$
$\displaystyle\qquad+4KL\eta^{(t)2}\Gamma_{K}\bigg{]}$
$\displaystyle=\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})\bigg{\\}}-\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t+1)})\bigg{\\}}+2KL^{2}\eta^{(t)2}\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})-H^{*}\bigg{\\}}+2KL^{2}\eta^{(t)2}\Gamma_{K}$
$\displaystyle\qquad+(8\eta^{(t)3}KL^{2}(E-1)+8KL\eta^{(t)2}+4(2+4L)KL\eta^{(t)4}(E-1))G^{2}+(2L\eta^{(t)2}+8KL\eta^{(t)2})\sum_{k=1}^{K}p_{k}\sigma_{k}^{2}.$

Let
$\xi^{(t)}=2KL^{2}\eta^{(t)2}\Gamma_{K}+(8\eta^{(t)3}KL^{2}(E-1)+8KL\eta^{(t)2}+4(2+4L)KL\eta^{(t)4}(E-1))G^{2}+(2L\eta^{(t)2}+8KL\eta^{(t)2})\sum_{k=1}^{K}p_{k}\sigma_{k}^{2}$,
then

$\displaystyle\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})\bigg{\\}}-\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t+1)})\bigg{\\}}+2KL^{2}\eta^{(t)2}\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})-H^{*}\bigg{\\}}+\xi^{(t)}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})\bigg{\\}}-\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t+1)})\bigg{\\}}+2KL^{2}\eta^{(t)2}\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(1)})-H^{*}\bigg{\\}}+\xi^{(t)}$

since $\eta^{(t)}\leq\frac{1}{\sqrt{2K}L}$ and
$\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t)})\bigg{\\}}\leq\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(1)})\bigg{\\}}$
by Lemma 3. By taking summation on both side, we obtain

$\displaystyle\sum_{t=1}^{T}\frac{1}{2}\eta^{(t)}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(1)})\bigg{\\}}-\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(t+1)})\bigg{\\}}+2KL^{2}\sum_{t=1}^{T}\eta^{(t)2}\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(1)})-H^{*}\bigg{\\}}+\sum_{t=1}^{T}\xi^{(t)}$
$\displaystyle\leq\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(1)})\bigg{\\}}-\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{*})\bigg{\\}}+2KL^{2}\sum_{t=1}^{T}\eta^{(t)2}\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(1)})-H^{*}\bigg{\\}}+\sum_{t=1}^{T}\xi^{(t)}$
$\displaystyle=(1+2KL^{2}\sum_{t=1}^{T}\eta^{(t)2})\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(1)})-H^{*}\bigg{\\}}+\sum_{t=1}^{T}\xi^{(t)}.$

This implies

$\displaystyle\min_{t=1,\ldots,T}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}\sum_{t=1}^{T}\eta^{(t)}\leq
2(1+2KL^{2}\sum_{t=1}^{T}\eta^{(t)2})\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(1)})-H^{*}\bigg{\\}}+2\sum_{t=1}^{T}\xi^{(t)}$

and therefore

$\displaystyle\min_{t=1,\ldots,T}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}\leq\frac{1}{\sum_{t=1}^{T}\eta^{(t)}}\bigg{\\{}2(1+2KL^{2}\sum_{t=1}^{T}\eta^{(t)2})\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(1)})-H^{*}\bigg{\\}}+2\sum_{t=1}^{T}\xi^{(t)}\bigg{\\}}.$

Let $\eta^{(t)}=\frac{1}{\sqrt{t}}$, then we have
$\sum_{t=1}^{T}\eta^{(t)}\geq\sqrt{T}$ and
$\sum_{t=1}^{T}\eta^{(t)2}\leq\log(T+1)$. Therefore,

$\displaystyle\min_{t=1,\ldots,T}\mathbb{E}\bigg{\\{}\left\lVert\nabla
H(\bm{\bar{\theta}}^{(t)})\right\rVert^{2}\bigg{\\}}\leq\frac{1}{\sqrt{T}}\bigg{\\{}2(1+2KL^{2}\sum_{t=1}^{T}\eta^{(t)2})\mathbb{E}\bigg{\\{}H(\bm{\bar{\theta}}^{(1)})-H^{*}\bigg{\\}}+2\sum_{t=1}^{T}\xi^{(t)}\bigg{\\}}.$

∎

##  11 GIFAIR-FL \+ Ditto

Ditto is iteratively solving a bi-level optimization problem. During each
communication round, the serve broadcast a model parameter $\bm{\theta}$. In
the first level, each device runs $E_{1}$ steps of SGD to collaboratively
optimize a global objective function $\sum_{k=1}^{K}p_{k}F_{k}(\bm{\theta})$.
In the second level, Ditto uses $\bm{\theta}$ to solve the following
optimization problem for each device $k$:

$\displaystyle\min_{\bm{v}_{k}}h_{k}(\bm{v}_{k};\bm{\theta})\coloneqq
F_{k}(\bm{v}_{k})+\frac{\lambda_{Ditto}}{2}\left\lVert\bm{v}_{k}-\bm{\theta}\right\rVert^{2}.$

When combining GIFAIR-FL with Ditto, we replace the objective function in the
first level by our fair objective function

$\displaystyle\sum_{k=1}^{K}p_{k}F_{k}(\bm{\theta})+\lambda\sum_{1\leq i<j\leq
d}|L_{i}(\bm{\theta})-L_{j}(\bm{\theta})|.$ (5)

A detailed algorithm is listed below.

Data: number of devices $K$, fraction $\alpha$, number of communication rounds
$C$, number of local updates $E_{1},E_{2}$, SGD learning rate schedule
$\\{\eta_{1}^{(t)}\\}_{t=1}^{E}$ and $\\{\eta_{2}^{(t)}\\}_{t=1}^{E}$, initial
model parameter $\bm{\theta},\\{\bm{v}_{k}\\}_{k\in[K]}$, regularization
parameter $\lambda,\lambda_{Ditto}$, initial loss $\\{L_{i}\\}_{1\leq i\leq
d}$

Result: model parameter $\bm{\bar{\theta}}_{C},\\{\bm{v}_{k}\\}_{k\in[K]}$.

for  _$c=0:(C-1)$_ do

Select $\left\lceil\alpha K\right\rceil$ clients by sampling probability
$p_{k}$ and denote by ${\cal S}_{c}$ the indices of these clients;

Server broadcasts $\left(\bm{\theta},\left\\{\dfrac{\lambda}{p_{k}|{\cal
A}_{k}|}r^{c}_{k}(\bm{\theta})\right\\}_{k\in{\cal S}_{c}}\right)$;

for  _$k\in\mathcal{S}_{c}$_ do

$\bm{\theta}_{k}^{(0)}=\bm{\theta}$;

for  _$t=0:(E_{1}-1)$_ do

$\bm{\theta}^{(t+1)}_{k}=\bm{\theta}^{(t)}_{k}-\eta_{1}^{(t)}\left(1+\dfrac{\lambda}{p_{k}|{\cal
A}_{k}|}r^{c}_{k}(\bm{\theta})\right)\nabla F_{k}(\bm{\theta}^{(t)}_{k})$;

end for

end for

for  _$k\in\mathcal{S}_{c}$_ do

$\bm{v}_{k}^{(0)}=\bm{v}_{k}$;

for  _$t=0:(E_{2}-1)$_ do

$\bm{v}_{k}^{(t+1)}=\bm{v}_{k}^{(t)}-\eta_{2}^{(t)}\left(\nabla
F_{k}(\bm{v}_{k}^{(t)})+\lambda_{Ditto}(\bm{v}_{k}^{(t)}-\bm{\theta})\right)$;

end for

end for

Aggregation $\bm{\bar{\theta}}_{c}=\frac{1}{\left\lceil\alpha
K\right\rceil}\sum_{k\in\mathcal{S}_{c}}\bm{\theta}^{(E_{1})}_{k}$, Set
$\bm{\theta}=\bm{\bar{\theta}}_{c}$;

Set $\bm{v}_{k}=\bm{v}_{k}^{(E_{2})}$ for all $k\in[K]$;

Calculate $L_{i}=\frac{1}{|{\cal A}_{i}|}\sum_{k\in{\cal
A}_{i}}F_{k}(\bm{\theta}^{(E_{1})}_{k})$ for all $i\in[d]$;

$c\leftarrow c+1$;

end for

Return $\bm{\bar{\theta}}_{C}$.

Algorithm 2 GIFAIR-FL \+ Ditto Algorithm

##  12 Additional Experiments

We conduct a sensitivity analysis using the FEMNIST-3-groups setting. Results
are reported in Figure 2. Similar to the observation in the main paper, it can
be seen that as $\lambda$ increases, the discrepancy between two groups
decreases accordingly. Here kindly note that we did not plot group 3 for the
sake of neatness. The line of group should stay in the middle of two lines.

Figure 2: Sensitivity analysis on FEMNIST

Claim: Let $f(x,y)=h(x)+g(y)$ where both $h,g$ are $\mu_{1},\mu_{2}$-strongly
convex. Then $f(x,y)$ is $\frac{\mu_{1}+\mu_{2}}{2}$-strongly convex.

$\displaystyle f(x_{1},y_{1})$ $\displaystyle=h(x_{1})+g(y_{1})$
$\displaystyle\geq h(x_{2})+\nabla
h(x_{2})^{T}(x_{1}-x_{2})+\frac{\mu_{1}}{2}\left\lVert
x_{1}-x_{2}\right\rVert^{2}+g(y_{2})+\nabla
g(y_{2})^{T}(y_{1}-y_{2})+\frac{\mu_{2}}{2}\left\lVert
y_{1}-y_{2}\right\rVert^{2}$
$\displaystyle=\underbrace{h(x_{2})+g(y_{2})}_{f(x_{2},y_{2})}+\underbrace{[\nabla
h(x_{2})^{T},\nabla g(y_{2})^{T}]\begin{bmatrix}(x_{1}-x_{2})\\\
(y_{1}-y_{2})\end{bmatrix}}_{\nabla
f(x_{2},y_{2})^{T}\begin{bmatrix}(x_{1}-x_{2})\\\
(y_{1}-y_{2})\end{bmatrix}}+\underbrace{\frac{\mu_{1}}{2}\left\lVert
x_{1}-x_{2}\right\rVert^{2}+\frac{\mu_{2}}{2}\left\lVert
y_{1}-y_{2}\right\rVert^{2}}_{\frac{\mu_{1}+\mu_{2}}{2}\left\lVert\begin{bmatrix}(x_{1}-x_{2})\\\
(y_{1}-y_{2})\end{bmatrix}\right\rVert^{2}_{2}}$

Our objective function is

$\displaystyle\min_{\bm{\theta}}\;H(\bm{\theta},\bm{\theta}_{1},\ldots,\bm{\theta}_{K})\triangleq\sum_{k=1}^{K}p_{k}F_{k}(\bm{\theta})+\lambda\sum_{1\leq
i<j\leq
d}|L_{i}(\\{\bm{\theta}_{k}\\}_{k\in\mathcal{A}_{i}})-L_{j}(\\{\bm{\theta}_{k}\\}_{k\in\mathcal{A}_{j}})|,$
(6) $\displaystyle\text{ s.t. }\sum_{k=1}^{K}p_{k}\bm{\theta}_{k}=\bm{\theta}$
(7)

where

$L_{i}(\\{\bm{\theta}\\}_{k\in{\cal A}_{i}})\triangleq\dfrac{1}{|{\cal
A}_{i}|}\sum_{k\in{\cal A}_{i}}F_{k}(\bm{\theta}_{k}).$

This is the group loss based on the local parameters.

To proceed, we have

$H(\bm{\theta},\bm{\theta}_{1},\ldots,\bm{\theta}_{K})=\sum_{k=1}^{K}p_{k}\underbrace{F_{k}(\bm{\theta})}_{\mu_{k}\text{-strongly
convex}}+\sum_{k=1}^{K}p_{k}\underbrace{\dfrac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}_{k})F_{k}(\bm{\theta}_{k})}_{\dfrac{\lambda}{p_{k}|{\cal
A}_{s_{k}}|}r_{k}(\bm{\theta}_{k})\mu_{k}\text{-strongly convex}},$ (8)

where

$r_{k}(\bm{\theta}_{k})=\sum_{1\leq j\neq s_{k}\leq
d}{\mbox{sign}}(L_{s_{k}}(\\{\bm{\theta}_{m}\\}_{m\in\mathcal{A}_{s_{k}}})-L_{j}(\\{\bm{\theta}_{m}\\}_{m\in\mathcal{A}_{j}})).$

In our algorithm, we treat $r_{k}(\bm{\theta}_{k})$ as fixed during local
updating.