# Multi-task Federated Edge Learning (MtFEEL) in Wireless Networks

Sawan Singh Mahara, Shruti M., B. N. Bharath, Akash Murthy111The authors are
with the Department of electrical engineering at IIT Dharwad, Dharwad,
Karnataka<EMAIL_ADDRESS><EMAIL_ADDRESS>
Akash Murthy is at Euprime Pvt. Ltd<EMAIL_ADDRESS>
###### Abstract

Federated Learning (FL) has evolved as a promising technique to handle
distributed machine learning across edge devices. A single neural network (NN)
that optimises a global objective is generally learned in most work in FL,
which could be suboptimal for edge devices. Although works finding a NN
personalised for edge device specific tasks exist, they lack generalisation
and/or convergence guarantees. In this paper, a novel communication efficient
FL algorithm for personalised learning in a wireless setting with guarantees
is presented. The algorithm relies on finding a “better“ empirical estimate of
losses at each device, using a weighted average of the losses across different
devices. It is devised from a Probably Approximately Correct (PAC) bound on
the true loss in terms of the proposed empirical loss and is bounded by (i)
the Rademacher complexity, (ii) the discrepancy, (iii) and a penalty term.
Using a signed gradient feedback to find a personalised NN at each device, it
is also proven to converge in a Rayleigh flat fading (in the uplink) channel,
at a rate of the order
$\max\left\\{\frac{1}{SNR},\frac{1}{\sqrt{T}}\right\\}$. Experimental results
show that the proposed algorithm outperforms locally trained devices as well
as the conventionally used FedAvg and FedSGD algorithms under practical SNR
regimes.

###### Index Terms:

Federated Learning, Multi-Task-Learning, SignSGD, Deep Learning, PAC bound,
Distributed ML.

##  I Introduction

The wide spread adoption of smartphones and internet services with
considerable computing capabilities has enabled machine learning (ML)
algorithms to work in a distributed fashion (see [1]). Since the data and the
edge devices are heterogeneous, a new paradigm called _Federated Learning
(FL)_ (see [2, 3, 4, 5]) has emerged, where the data is distributed while the
central node controls the exchange of the data. FL is faced with several
challenges such as (i) _stragglers_ , where edge devices leave the network and
stops contributing to the FL process, without warning, (ii) untimely updates
from heterogeneous edge devices with varying computation power, (iii)
statistical heterogeneity in data, (iv) privacy concerns, and (v)
communication between edge devices and a central node or a Base Station (BS)
being expensive [6]. A de facto version of FL algorithm used in applications
such as next word prediction in mobile keyboards is called Federated Averaging
(FedAvg) [7] which aims to optimise a global objective. It has been shown that
FedAvg outperforms models trained using data from individual devices [8].
FedAvg involves performing multiple rounds of SGD on a subset of users in the
network, and communicating the resulting neural network to a central BS, where
they will be averaged and communicated back to devices. This process is
repeated until convergence. Although successful, FedAvg does not fully address
the issues of data heterogeneity [9], and is known to diverge in non-i.i.d
settings due to model drift [6]. To circumvent this, a modified version of the
FedAvg algorithm called FedProx is proposed in [6]. Subsequently, numerous
algorithms such as SCAFFOLD [10], MIME [11], SlowMo [12], QG-DSGDm [13] have
been proposed to handle the model drift problem. Other extensions include FL
algorithms to handle differential privacy [14], secure aggregation [15], and
scenarios with less number of devices participating in the FL process [5]. In
all of the above work (except [6]), a single model is returned that optimizes
a global objective, and hence does not perform well in the case of
heterogeneous data. Therefore, there is a pressing need towards designing a
personalized FL system in wireless scenarios, with less communication between
the BS and edge devices. This problem is addressed in this paper.

###  I-A Related Work and Motivation

A naive implementation of the FL using SGD would require repeated exchanges of
gradients of the losses (typically a vector of a million or more entries),
which leads to huge radio resource requirements. This communication overhead
can be overcome by compressing the gradient information before being
transmitted [16]. A simple way of compressing the gradient is to use the sign
of the gradient called SIGNSGD. A detailed theoretical analysis of a majority
vote based SIGNSGD with non-convex loss is provided in [17]. In wireless edge
devices, the number of updates to each device can be further reduced by
exploiting the nature of the wireless medium leading to a solution called
_over-the-air aggregation_ also called _over-the-air computation_ [18, 19,
20]. An extension of this called one-bit broadband digital aggregation (OBDA)
is proposed in [21]. Further, a similar work with client scheduling and
resource block allocation with improper channel state information (CSI) is
considered in [22]. The authors in [23] study the impact of wireless channel
hostilities under some assumptions on the CSI. Mobile edge computing devices
that have resource constraints and spotty wireless communication links pose
their own set of challenges. Providing some guarantees on the model accuracy
achievable in such situations has been looked at by the authors of [24]. In
all of the above literature, an estimate of the average loss is minimised in a
Federated fashion to obtain a single neural network.

The above work fails to provide generalisation guarantees and doesn’t work
well on device specific tasks. In non-IID settings, the performance at a
device with heterogeneous data distributions is difficult to improve with a
single model. This issue can be tackled by learning multiple models for
different target distributions. In this light, a distributed multi-task
learning (MTL) algorithm called Mocha tries to learn a single neural network
optimised to its task (see [25]). Here, each task is learned by solving a
primal-dual optimisation problem in a convex setting. The mixture methods FL
framework (see [26], [27], [28], [29]) achieve some personalisation by combing
the model parameters obtained by training a local model and a global model.
Global and local parameter mixing can be done across neural network layers, by
incorporating the lower layers to adapt to each device’s data while having the
higher layers shared among other devices (see [30]). Adapting the existing
federated averaging algorithm to mixture methods, the authors of [31] utilise
meta learning to personalise a global model to each device. Alternatively, if
the assumption that the non-IID data are partitioned into groups and can be
clustered, Clustered FL ([32] ,[33], [34], [28]) addresses these challenges by
grouping devices with similar distributions to improve model accuracies.
Another approach to improve model accuracies via personalisation is by a
weighted combination method. For example, FedFomo [35] uses the information of
how much any device could benefit from another device’s model. Many issues in
the above work are as follows (i) a lack of a personalised model, (ii) poor
performance due to inaccurate estimates of loss functions using local data,
(iii) a lack of generalisation guarantees and convergence guarantees. This
work address all these issues in a systematic manner.

###  I-B Contributions of the paper

The setup studied in the paper consists of a wireless network of edge devices
(like smart phones) connected to a BS with the goal of learning optimal neural
network at each of the edge devices to perform some supervised learning tasks
such as classification, prediction, regression, to name a few. In particular,
an improved estimate of the loss at each device is used to optimize the neural
network weights. An improved estimate at each device is obtained by using
weighted loss across all the devices. Naturally, the devices with similar data
should be given higher weights. Finding these weights and subsequently the
neural network tailored towards the task of each device is a challenging
problem of multi-task learning that is addressed in the paper. This paper
presents a systematic approach to finding the weights backed by theory. In
particular, a Probably Approximately Correct (PAC) [36] bound on the
performance of the weighted average losses across devices with respect to the
true loss is presented. It is shown that the bound depends on (i) Rademacher
complexity; a measure of the complexity of the learning task, (ii)
discrepancy; a measure of statistical “closeness” of the data between any two
devices, and (iii) a regularization term on the weights. Based on the insights
provided by the bound, a distributed learning algorithm to find (i) an
estimate of the discrepancy, (ii) a device importance weighting metric and
(iii) the weights of neural networks, is presented. In the absence of wireless
abnormalities, within $T$ communication rounds, the algorithm is shown to
converge at a rate of $1/\sqrt{T}$. At the end of the training, all devices
are provided with a custom neural network, which characterises the multi-task
nature of the algorithm. In a Rayleigh flat fading channel scenario, the
algorithm is shown to converge provided the $SNR$ is reasonably high. In
particular, the convergence as a function of $SNR$ and $T$ is shown to be
$\mathcal{O}\left(\max\left\\{\frac{1}{SNR},\frac{1}{\sqrt{T}}\right\\}\right)$.
In other words, for a fixed $SNR$, higher rounds of training $T$ does not
guarantee to yield better convergence performance. This system is simulated
using python and tensorflow and the proposed algorithm (MtFEEL) was shown to
outperform locally trained neural networks as well as existing state-of-the-
art federated algorithms. The simulations were also performed in a wireless
setting and shown to outperform classical approaches such as FedAvg, FedSGD
and local training in SNRs of practical interests. However, at lower $SNR$s,
performance of the MtFEEL algorithm degrades, in which case it is better to
use classical federated algorithms, as indicated by our convergence results.
These insights can prove useful in designing next generation wireless
networks. The paper is organized as follows. In Sec. II, the system model and
the problem considered in the paper are presented. Section III presents the
first main result based off of which, the algorithm in the noiseless scenario
is presented in Sec. IV. The convergence results of this algorithm in an
ideal, error free regime is presented in Sec. V. Section VI extends the
convergence result to noisy channel. The performance of the proposed algorithm
on real data set is presented in Sec. VII. Finally, the paper is concluded in
Sec. VIII.

##  II System Model

The paper considers the problem of federated multi-task learning with $N$
devices (example, mobiles) and a BS, as shown in Fig. 1.

Figure 1: Federated learning scenario

Each device has a certain task like next word prediction, and the user of the
device provides data (supervised) to learn the ML model. Assume that the
$k$-th user has $n_{k}$ training data denoted by
$S_{k}=\\{(\bm{x}_{k1},y_{k1}),(\bm{x}_{k2},y_{k2}),\ldots,(\bm{x}_{kn_{k}},y_{kn_{k}})\\}$,
where $\bm{x}_{ij}\in\mathcal{X}$ is the feature vector corresponding to the
$j$-th training example at the $i$-th edge device, and $y_{ij}\in\mathcal{Y}$
is the corresponding label. Let $\bm{S}:=\\{S_{1},S_{2},\ldots S_{N}\\}$ be
the set of all samples present. Here, $\mathcal{X}\subseteq\mathbb{R}^{d}$ and
$\mathcal{Y}$ represent the feature space and the output (label) space,
respectively. The data at the $k$-th device is drawn in an i.i.d. fashion with
distribution denoted by $\mathcal{D}_{k}$, $k=1,2,\ldots,N$. Further, the data
across devices is assumed to be independent but not necessarily identical
distributed.222This assumption is made for the sake of clarity of
presentation. Typically, the devices would not like to communicate the raw
samples due to privacy concerns. Therefore, the learning should happen in a
federated fashion. The learning rule/hypothesis considered is of the form
$h_{\bm{w}_{k}}:\mathcal{X}\rightarrow\mathcal{Y}$, $k=1,2,\ldots,N$.333Note
that instead of $\mathcal{Y}$, one can consider $\Delta_{y}$, the simplex over
$\mathcal{Y}$ as well. It is important to note that any neural network
architecture can be characterised in this way. Let
$W=\\{\bm{w}_{1}\ldots\bm{w}_{N}\\}\in\mathbb{R}^{d\times N}$. Given a feature
vector $\bm{x}_{k}\in\mathcal{X}$, and the corresponding label
$y_{k}\in\mathcal{Y}$, $k=1,2,\ldots,N$, the performance of the neural network
at each device $k$ is measured using a loss function
$l:\mathcal{Y}\times\mathcal{Y}\rightarrow\mathbb{R}^{+}$. For the ease of
notation, let
$\mathcal{L}_{k}(\bm{w}):=\underset{(\bm{x},y)\sim\mathcal{D}_{k}}{\mathbb{E}}l\left(h_{\bm{w}}(\bm{x}),y\right)$
and the corresponding estimate by
$\hat{\mathcal{L}}(\bm{w},S_{k}):=\frac{1}{n_{k}}\sum_{i=1}^{n_{k}}l\left(h_{\bm{w}}(\bm{x}_{ki}),y_{ki}\right)$.
Note that the estimate is unbiased, i.e.,
$\mathcal{L}_{k}(\bm{w})=\underset{S_{k}\sim\mathcal{D}_{k}}{\mathbb{E}}\hat{\mathcal{L}}(\bm{w},S_{k})$
for any $\bm{w}\in\mathbb{R}^{d}$. In the classical federated setting, the
goal is to solve the following optimization problem, i.e., find one neural
network $\bm{w}\in\mathbb{R}^{d}$ across all the devices:

$\min_{\bm{w}\in
W}\left\\{\Phi_{\bm{w}}:=\sum_{k=1}^{N}\frac{n_{k}}{n}\mathcal{L}_{k}(\bm{w})\right\\},$
(1)

where the total number of samples $n:=\sum_{k=1}^{N}n_{k}$. Note that in
practice, the solution to the above can be obtained by using an estimate of
the gradient of $\mathcal{L}_{k}(\bm{w})$ in the SGD algorithm. The challenge
in the FL setting is that the overall gradient is the sum of the gradients of
individual average losses. A number of solutions such as FedAvg [9], one bit
gradient based majority vote [17], and many more have been proposed in the
literature [16] The co-efficient $n_{k}/n$ will determine the importance of
the loss corresponding to the user $k$. However, if the number of samples in
the future from user $k$ is less, then the model returned by solving the above
problem may result in poor performance across several devices. A solution to
the above is to look at the worst case scenario as described below [3]:

$\min_{\bm{w}\in
W}\sup_{\bm{\lambda}\in\bm{\Lambda}}\left\\{\Phi_{\bm{w},\bm{\lambda}}:=\sum_{k=1}^{N}\lambda_{k}\mathcal{L}_{k}(\bm{w})\right\\},$
(2)

where $\bm{\lambda}:=(\lambda_{1},\lambda_{2},\ldots,\lambda_{N})$, and the
constraint set $\bm{\Lambda}\subseteq\Delta_{N}$ incorporates the prior
knowledge on devices that may drop off from the network (see [3]). Note that
the above problem uses the same neural network $\bm{w}$ across all the
devices, unlike the scheme proposed in this work. The neural network can be
made more personalised by using different neural network weights, as done by
the authors of [37]. This work however, explicitly considers the statistical
heterogeneity across devices as well as weighting every loss metric in
accordance with the heterogeneity, which leads to solving the following
optimization problem.

$\min_{\bm{w}_{1},\bm{w}_{2},\ldots,\bm{w}_{N}}\sup_{\bm{\lambda}\in\bm{\Lambda}}\left\\{\Phi_{W,\bm{\lambda}}:=\sum_{k=1}^{N}\lambda_{k}\mathcal{L}_{k}(\bm{w}_{k})\right\\},$
(3)

In order to solve the above problem, the devices should compute an estimate of
$\mathcal{L}_{k}(\bm{w}_{k})$ denoted by
$\hat{\mathcal{L}}(\bm{w}_{k},S_{k}):=\frac{1}{n_{k}}\sum_{i=1}^{n_{k}}l\left(h_{\bm{w}_{k}}(\bm{x}_{ki}),y_{ki}\right)$,
which can be used as a proxy in (3). Note that the estimate of the average
loss of the $k$-th device depends only on its data. However, it is natural to
include neighbors’ empirical estimate of the average losses while estimating
the average loss of the $k$-th device if the neighboring data distribution is
“close” to the distribution of the data of the device. An extreme scenario is
that of an i.i.d. data across devices where a simple averaging works well. One
approach is to take the average of the empirical loss across all the devices.
This may lead to a bad estimate of the average loss since the neighboring data
are given equal weights. A way around this problem is to “optimally” allocate
weights across users data. This leads to the following optimization problem
that needs to be solved in a federated manner

$\min_{\bm{w}_{1},\bm{w}_{2},\ldots,\bm{w}_{N}}\sup_{\bm{\lambda}\in\bm{\Lambda}}\left\\{\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S}):=\sum_{k=1}^{N}\lambda_{k}\sum_{j=1}^{N}\alpha_{kj}\hat{\mathcal{L}}(\bm{w}_{k},S_{j})\right\\}.$
(4)

It is natural to constraint $\alpha_{kj}$ as $\sum_{j=1}^{N}\alpha_{kj}=1$. In
the above, let the weights be denoted by
$\bm{\alpha}:=\\{\bm{\alpha}_{1},\bm{\alpha}_{2},\ldots,\bm{\alpha}_{N}\\}$
and $\bm{\alpha}_{m}:=\\{\alpha_{m1},\alpha_{m2},\ldots,\alpha_{mN}\\}$ for
all $m=1,2,\ldots,N$. The neural network weights computed using (4) will be
close to the one computed in (2) provided the gap between
$\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})$ and
${\Phi}_{W,\bm{\lambda}}$ is small. It is now evident that the federated
learning requires users to communicate with the BS and vice-versa. In the
following subsection, the channel model for the BS and users to communicate is
presented.

###  II-A Channel Model

The communication model consists of a single BS Single Input Single Output
(SISO) system with multiple users or edge devices. In particular, the time is
assumed to be slotted, and the channel between any device $k$ and the BS is
assumed to be a wireless Rayleigh flat fading channel. The complex baseband
received signal $y(t)\in\mathcal{C}$ at the BS for an input
$x_{k}(t)\in\mathcal{C}$ by a device $k$ is given by

$\displaystyle y_{k}(t)=h_{k}(t)x_{k}(t)+z_{k}(t),~{}t=1,2,\ldots,\text{ and
}k=1,2,\ldots,N,$ (5)

where
$h_{k}(t)\stackrel{{\scriptstyle\text{i.i.d.}}}{{\sim}}\mathcal{CN}(0,1)$ is
the fading channel coefficient between device $k$ and the BS. The noise
$z_{k}(t)\;{\sim}\;\mathcal{CN}(0,\sigma^{2})$ is a circularly symmetric
complex Gaussian random variable. Typically, the devices are mobile phones,
and hence are power limited. Therefore, the power constraint at the device $k$
is given by $\mathbb{E}{|x_{k}(t)|^{2}}\leq P_{k}$. On the other hand, the BS
is assumed to have enough power to communicate without any errors. This
assumption is made for the sake of simplicity and the ease of exposition. At
any given time slot, it is assumed that the scheduler will assign the channel
of bandwidth $B$ to any edge device that wishes to communicate. For example,
in an OFDMA system, a resource block is allocated to a user who wishes to
transmit. For the sake of simplicity, the impact of the scheduling scheme on
the convergence of the proposed algorithm is ignored. In order to device an
algorithm, and prove convergence under a noisy communication channel, first a
noiseless scenario is considered. Subsequently, the results are extended to
study the impact of wireless communication channels on the convergence of the
algorithm. The first main result of this paper is to prove a PAC bound [36] on
$\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})-\Phi_{W,\bm{\lambda}}$ when
the channels between devices and BS are ideal, i.e., without any error. This
PAC bound will be used to device a distributed federated algorithm that is
shown to outperform state-of-the-art federated algorithms. Subsequently, a
convergence guarantee on the proposed algorithm is also provided. The impact
of the wireless channel model explained above on the convergence is provided
in Sec. VI. The following section presents the main result.

##  III Main Result - I

To state the first main result of the paper, the following three quantities
will be required (i) Rademacher complexity, (ii) Minimum $\epsilon$-cover, and
(iii) discrepancy, which are defined below.

######  Definition 1.

(Minimax weighted Rademacher complexity [3]) The Rademacher complexity for the
class of neural networks $W$ for a given $\bm{\lambda}\in\bm{\Lambda}$ is
defined as

$\mathcal{R}_{\bm{\lambda}}\left(W\right):=\underset{\bm{S},\bm{\sigma}}{\mathbb{E}}\left[\sup_{\begin{subarray}{c}\bm{w}_{1}\bm{w}_{2},\ldots\bm{w}_{N}\\\
\alpha\in\Delta_{N}\end{subarray}}\sum_{k,j=1}^{N}\frac{\lambda_{k}\alpha_{kj}}{n_{j}}\sum_{i=1}^{n_{j}}\sigma_{kj,i}l\left(h_{\bm{w}_{k}}\left(\bm{x}_{ji}\right),y_{ji}\right)\right],$

where the Rademacher random variables $\sigma_{kj,i}\in\\{1,-1\\}$ for
$k,j=1,2,\ldots,N$ occur with equal probability. The max weighted Rademacher
complexity is defined as
${\mathcal{R}_{\bm{\Lambda}}\left(W\right)=\max_{\bm{\lambda}\in\bm{\Lambda}}\mathcal{R}_{\bm{\lambda}}\left(W\right).}$

This complexity is a measure of how well any members of a class of real valued
hypotheses can approximate random noise. The expressibility of the neural
network is measured based on how well the hypothesis (model) class fits the
noise. The class can therefore be expected to learn more intricate decision
boundaries. In order to find the $\sup$ over $\bm{\lambda}\in\bm{\Lambda}$, it
is useful to quantize the set $\bm{\Lambda}$ for which the following
definition of minimum $\epsilon$-cover comes in handy [38].

######  Definition 2.

(Minimum $\epsilon$-cover [38]) The set
$\left\\{\bm{v}_{1},\ldots,\bm{v}_{p}\right\\}$ is said to be an
$\epsilon$-cover of $\bm{\Lambda}$ with respect to $\ell_{1}$-distance if
$\bm{\Lambda}\subseteq\cup_{i=1}^{p}B\left(\bm{v}_{i},\epsilon\right)$, where
the $L_{1}$ ball is defined as
$B(\bm{v}_{i},\epsilon):=\\{\bm{z}\in\bm{\Lambda}:\|\bm{z}-\bm{v}_{i}\|_{1}<\epsilon\\}$.
The minimum $\epsilon$-cover $\bm{\Lambda}_{\epsilon}$ of a set $\bm{\Lambda}$
is any $\epsilon$-cover with the smallest $p$.

It is expected that the optimal weights $\alpha_{kj}$ will depend on
“closeness” of the distributions of data across the devices. The following
definition provides a measure of the difference in two distributions with
respect to a loss function.

######  Definition 3.

(Discrepancy [3]) Given two data distributions $\mathcal{D}_{k}$ and
$\mathcal{D}_{j}$ of the devices $k$ and $j$ respectively, the corresponding
discrepancy with respect to the loss
$l:\mathcal{Y}\times\mathcal{Y}\rightarrow\mathbb{R}^{+}$ is defined as
$d_{kj}:=\sup_{\bm{w}\in\mathbb{R}^{d}}\Delta_{kj}(\bm{w})$, where
$\Delta_{kj}(\bm{w}):=\left|\mathcal{L}_{k}(\bm{w})-\mathcal{L}_{j}(\bm{w})\right|$.

Recall that
$\mathcal{L}_{k}(\bm{w}):=\underset{(\bm{x},y)\sim\mathcal{D}_{k}}{\mathbb{E}}l\left(h_{\bm{w}}(\bm{x}),y\right)$.
The first term above corresponds to the average loss at the $k$-th device
while the second term corresponds to the average loss at the $j$-th device.
The difference provides a measure of how different the two data distributions
are with respect to the neural network $\bm{w}$. Maximising over
$\bm{w}\in\mathbb{R}^{d}$ gives the worst case difference. If the data at the
two devices are i.i.d., then, it is easy to see that the discrepancy is zero.
On the other hand, if the distributions of data at the two devices are
different, and the same neural network is not able to classify the data, then
the discrepancy is high. A naive way of using this while devising an algorithm
is to allocate higher weights $\alpha_{kj}$ if $d_{kj}$ is small. In the
following, a theoretical result that provides insights on how to choose these
weights in a systematic way is provided. In particular, using the definitions
above, a PAC bound on the difference between the true average loss in (3) and
its estimate in (4) is provided in the following theorem. The proof is
relegated to Appendix [A].

######  Theorem 1.

(PAC bound) Assuming that the loss is bounded, i.e., $l(a,b)\leq
M\;\forall\;a,b\in\mathcal{Y}$, for every $\epsilon>0$, with a probability of
at least $1-\delta$, $\delta>0$, the following holds

$\displaystyle{\Phi}_{W,\bm{\lambda}}\leq\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})+2\mathcal{R}_{\bm{\Lambda}}\left(W\right)+M\texttt{Pen}(\bm{\lambda},\bm{\alpha})+MN\epsilon,$
(6)

where
$\texttt{Pen}(\bm{\lambda},\bm{\alpha}):=\sqrt{\frac{N}{2}\sum_{j=1}^{N}\sum_{k=1}^{N}\left(\frac{\lambda_{k}\alpha_{kj}}{n_{j}}\right)^{2}\log\left(\frac{|\Lambda_{\epsilon}|}{\delta}\right)}+\frac{1}{M}\sum_{k,j=1}^{N}\lambda_{k}\alpha_{kj}{d}_{kj}$,
$d_{kj}$ is the discrepancy, and $\bm{\Lambda}_{\epsilon}$ is the minimum
$\epsilon$-cover of $\bm{\Lambda}$.

The above guarantee suggests that the neural network weights
$\\{\bm{w}_{1},\bm{w}_{2},\ldots\bm{w}_{N}\\}$ and ${\alpha}_{kj}$,
$k,j=1,2\ldots,N$, henceforth known as the _importance coefficients_ , can be
chosen in such a way that the error term in the theorem (6) is minimized.
Inspired by this, the following optimisation problem is proposed:

$\displaystyle\min_{\bm{w}_{1},\bm{w}_{2},\ldots,\bm{w}_{N}}$
$\displaystyle\min_{\bm{\alpha}}\left\\{\hat{\Psi}_{W,\bm{\lambda},\bm{\alpha}}:=\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})+\sum_{k=1}^{N}\gamma_{k}\parallel\bm{w}_{k}\parallel_{2}+M\texttt{Pen}(\bm{\lambda},\bm{\alpha})\right\\},$
(7)

where the second term above is used to regularize the neural network
coefficients. The above needs to be solved in a distributed fashion ensuring
that the communication overhead is low. In the following section, a federated
algorithm is proposed to solve (7).

##  IV Distributed FEEL (DFL) Algorithm

It is important to note that in order to solve (7), the knowledge of the
discrepancy is required. However, the devices have access to data, and hence
the discrepancy needs to be estimated in a distributed fashion. This estimate
will be used as a proxy in (7) to design the federated algorithm. In the
following subsection, an algorithm to estimate the discrepancy is proposed.

###  IV-A Distributed Discrepancy Estimation (DDE)

From the definition of the discrepancy, the DDE algorithm amounts to solving
$\sup_{\bm{w}\in\mathbb{R}^{d}}\\\
\left|\mathcal{L}_{k}(\bm{w})-\mathcal{L}_{j}(\bm{w})\right|$ for all
$k,j=1,2,\ldots,N$. Since the true average in the expression for the
discrepancy is unknown, estimates of those two terms denoted by
$\hat{\mathcal{L}}(\bm{w},S_{k})$ and $\hat{\mathcal{L}}(\bm{w},S_{j})$ will
be used. A natural approach to solving the problem is to use a gradient ascent
algorithm in a distributed manner. Algorithm $1$ shows the distributed
implementation of the gradient ascent to compute an estimate of the
discrepancy. Since the discrepancy involves an absolute value, the gradient
does not exist at all points. However, to circumvent this problem, a
generalized (sub) gradient is used in place of the gradient (see step $8$ of
the algorithm)

$\displaystyle\partial_{\bm{w}}(\Delta_{kj}(\bm{w})):=$
$\displaystyle(\nabla(\hat{\mathcal{L}}(\bm{w},S_{j}))-\nabla(\hat{\mathcal{L}}(\bm{w},S_{k})))\mathbbm{1}\\{(\hat{\mathcal{L}}(\bm{w},S_{k})>\hat{\mathcal{L}}(\bm{w},S_{j}))\\}-$
$\displaystyle(\nabla(\hat{\mathcal{L}}(\bm{w},S_{k}))-\nabla(\hat{\mathcal{L}}(\bm{w},S_{j}))\mathbbm{1}\\{(\hat{\mathcal{L}}(\bm{w},S_{j})\leq\hat{\mathcal{L}}(\bm{w},S_{k}))\\}.$
(8)

This follows directly from the sub-derivative of $|x|$ which is $-1$ if $x<0$,
$1$ if $x>0$ and if $x=0$, it is any point in the interval $[-1,1]$ (see
[39]). As the losses are continuous random variables, the probability of the
event that the losses are equal is zero, and hence the gradient in (IV-A) is
sufficient.

1Initialise discrepancies $\hat{d}_{jk}$ for $k,j\in\\{1,2,\ldots N\\}$ as $1$
and $\bm{w}^{0}\sim\mathcal{N}(0,I),\;I\in\mathbb{R}^{d\times d}$

2 for  _$j\in\\{1,2,\ldots,N\\}$_ do

3 for  _$k\in\\{j,\ldots,N\\}$_ do

4 for  _$t\in\\{1,2,\ldots,T\\}$ and $j\neq k$_ do

5 Broadcast $\bm{w}^{t}$ to devices ${j}$ and ${k}$

6 Receive (sub)gradients $\nabla\mathcal{\hat{L}}(\bm{w}^{t},S_{j})$,
$\nabla\mathcal{\hat{L}}(\bm{w}^{t},S_{k})$ and losses
$\mathcal{\hat{L}}(\bm{w}^{t},S_{j})$ and
$\mathcal{\hat{L}}(\bm{w}^{t},S_{k})$ from devices ${j}$ and ${k}$,
respectively

7 Set
$\hat{d}_{jk}=\hat{d}_{kj}:=|(\mathcal{\hat{L}}(\bm{w}^{t},S_{j})-\mathcal{\hat{L}}(\bm{w}^{t},S_{k}))|$

8 (Sub)Gradient Ascent using
$\bm{w}^{t+1}=\bm{w}^{t}+\eta\partial_{\bm{w}}(\Delta_{kj}(\bm{w}^{t}))$

9 end for

10

11 end for

12

13 end for

Output all $\hat{d}_{jk}$ for $k,j\in\\{1,2,\ldots,N\\}$

Algorithm 1 DDE Algorithm

In general, the problem is non-convex, and hence the above algorithm need not
converge to the global maximum. Since this is a gradient ascent algorithm, the
convergence to a local maximum follows from the standard argument [39]. The
time complexity of Algorithm $1$ is polynomial of the order
$\mathcal{O}(N^{2}dT)$. In the next subsection, using estimates of
discrepancies, a distributed federated learning algorithm is developed.

_Note:_ In the simulations, the number of iterations were fixed and it was
observed that when the data across devices were i.i.d, $T$ was smaller than
the threshold number of iterations. Although there is a scope for improvement
in terms of communication complexity while estimating the discrepancy, the
focus of this paper is to show that the performance of the federated algorithm
can be improved using discrepancy estimate while maintaining the communication
complexity to be nominal.

###  IV-B Proposed DFL Algorithm

The discrepancy estimates obtained from the DDE algorithm can be used as
proxies for the true discrepancies while solving the problem in (7). One can
use the classical federated algorithm to solve the problem. However, this
requires exchange of gradients, which in many problems can be of very high
dimension leading to a communication bottleneck. Therefore, in this section, a
signed gradient method is proposed, which is different from [17]. Although a
general approach of quantized gradients can used in this context, for the sake
of simplicity, a simple signed gradient will be used, and relegate the
analysis of quantized gradients to future work. It is important to note that
the mathematical tools used here can be extended to handle quantized gradients
scenario. In the signed gradient scenario, computing the gradient of the
objective in (7) involves finding the sign of the gradient with respect to
$\\{\bm{w}_{1},\bm{w}_{2},\ldots\bm{w}_{N}\\}$ instead of the full gradient,
and the full gradient with respect to $\bm{\alpha}$. This will be a function
of a relatively much lesser dimensional information. The gradient with respect
to the importance coefficients $\bm{\alpha}$ does not depend on the neural
network weights but depends on the discrepancy and the losses. This involves
sending $\mathcal{O}(N)$ parameters while the gradient with respect to neural
network weights $\bm{w}$ can potentially involve millions
($\gg\mathcal{O}(N)$) of parameters. The signed gradient of the objective
function (7) with respect to device $k$’s weights $\bm{w}_{k}$ turns out to be

$\nabla_{\bm{w}_{k},\operatorname{sign}}\hat{\Psi}_{W,\bm{\lambda},\bm{\alpha}}=\lambda_{k}\sum_{m=1}^{N}\alpha_{km}\texttt{sign}(\hat{\bm{g}}_{km})+\gamma_{k}\operatorname{sign}(\bm{w}_{k}),$
(9)

where the gradient at the $m^{th}$ device running model $k$ denoted by
$\hat{\bm{g}}_{km}:=\nabla_{\bm{w}_{k}}\mathcal{\hat{L}}(\bm{w}_{k},S_{m})$,
and sign$(\hat{\bm{g}}_{km})$ represents the sign of the estimated gradient
$\hat{\bm{g}}_{km}$. The full gradient of the loss
$\hat{\Psi}_{W,\bm{\lambda},\bm{\alpha}}$ for a device $k$, with respect to
the importance coefficients
$\bm{\alpha}_{k}=[\alpha_{k1}\;\alpha_{k2}\dots\;\alpha_{kN}]$ is given by
$\nabla_{\bm{\alpha}_{k}}\hat{\Psi}_{W,\bm{\lambda},\bm{\alpha}}:=\left[\frac{\partial\hat{\Psi}_{W,\bm{\lambda},\bm{\alpha}}}{\partial{\alpha_{k1}}},\;\frac{\partial\hat{\Psi}_{W,\bm{\lambda},\bm{\alpha}}}{\partial{\alpha_{k2}}},\ldots,\frac{\partial\hat{\Psi}_{W,\bm{\lambda},\bm{\alpha}}}{\partial{\alpha_{kN}}}\right]$,
where

$\frac{\partial\hat{\Psi}_{W,\bm{\lambda},\bm{\alpha}}}{\partial{\alpha_{km}}}=\lambda_{k}\mathcal{\hat{L}}(\bm{w}_{k},\bm{S}_{m})+\lambda_{k}\hat{d}_{km}+\frac{M\sqrt{\frac{N}{2}\displaystyle\log\left(\frac{\mid\bm{\Lambda}_{\epsilon}\mid}{\delta}\right)}\frac{\left(\displaystyle\lambda_{k}^{2}\alpha_{km}\right)}{n_{m}^{2}}}{\displaystyle\sqrt{\left(\displaystyle\frac{1}{2}\displaystyle\sum_{k,j=1}^{N}\left(\frac{\lambda_{k}\alpha_{kj}}{n_{j}}\right)^{2}\right)}}.$
(10)

In particular, each device $m$ sends the signed gradients
$\texttt{sign}(\hat{\bm{g}}_{km})$ (from equation (9)) and losses
$\mathcal{\hat{L}}(w_{m},S_{k})\;\forall\;k=1,2,\ldots,N$ to the BS. The BS
aggregates these according to (9) and (10) to get the estimated gradients
$\nabla_{\bm{w}_{k},\operatorname{sign}}\hat{\Psi}_{W,\bm{\lambda},\bm{\alpha}}$
and $\nabla_{\bm{\alpha}_{k}}\hat{\Psi}_{W,\bm{\lambda},\bm{\alpha}}$. These
estimates are used in obtaining an improved estimate of the optimal weights
$\bm{w}_{k}$ and importance coefficients $\alpha_{kj}$ for all
$k,j=1,2,\ldots,N$. In order to satisfy the constraint that
$\sum_{j=1}^{N}\alpha_{kj}=1$, the importance coefficient needs to be
projected onto the simplex $\Delta_{N}$ (see step $8$ of the Algorithm $2$).
Note that the gradient of (7) with respect to $\bm{\alpha}$ depends on the
estimated discrepancy. Computing this gradient is straightforward, and hence
not explicitly mentioned in the algorithm. The convergence analysis of
Algorithm $2$ is discussed in the next section.

1Initialise $\bm{\alpha}_{k}^{0}\in\Delta_{N}$, $\bm{w}_{k}\in\mathbb{R}^{d}$
for $k=1,2,\ldots,N$

2 for  _$t=1,2,\ldots,T$_ do

3 Broadcast $\bm{w}_{j}^{t}$ to the device $j=1,2,\ldots,N$

4 for  _devices $k=1,2,\ldots,N$_ do

5 Get $\texttt{sign}(\hat{\bm{g}}_{mk})\;\text{ for }m=1,2,\ldots,N$

6 Get $\hat{\mathcal{L}}(\bm{w}_{m},S_{k})\;\text{ for }m=1,2,\ldots,N$

7

8 end for

9 Gradient Descent step on $\bm{w}_{k}\;\text{ for }k=1,2,\ldots,N$:
$\bm{w}_{k}^{t+1}=\bm{w}_{k}^{t}-\eta\nabla_{\bm{w}_{k},\operatorname{sign}}\hat{\Psi}_{W^{t},\bm{\lambda},\bm{\alpha}^{t}}$

10 Gradient Descent with Projection step on
$\bm{\alpha}_{k}\;\text{for}\;k=1,2,\ldots,N$:
$\bm{\alpha}=\bm{\alpha}_{k}^{t}-\mu\nabla_{\bm{\alpha}_{k}}\hat{\Psi}_{W^{t+1},\bm{\lambda},\bm{\alpha}^{t}}$
$\bm{\alpha}_{k}^{t+1}=\underset{{\bm{x}\in\Delta_{N}}}{\operatorname{argmin}}\parallel\bm{x}-\bm{\alpha}\parallel_{1}$

11 end for

Algorithm 2 Proposed DFL (MtFEEL): Input discrepancy from Algorithm $1$

##  V Convergence Analysis

In this section, the convergence analysis of Algorithm $2$ is presented. In
order to prove one of the main results on the convergence, the following
standard assumptions are made.

######  Assumption 1.

(Boundedness [40]): The loss function $l(h_{\bm{w}}(\bm{x}),y)$ is assumed to
be bounded i.e., $l(h_{\bm{w}}(\bm{x}),y)\leq B<\infty,$ for all
$\bm{w}\in\mathbb{R}^{d}$ and any $\bm{x}\in\mathcal{X}$ and
$y\in\mathcal{Y}$.

Assumption 1 implies that $\Phi_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})\leq B$,
which follows from the facts that
$\overset{N}{\underset{{k=1}}{\sum}}\lambda_{k}=1$ and
$\overset{N}{\underset{{j=1}}{\sum}}\alpha_{kj}=1,$ $k=1,2,\ldots,N$. Next,
the gradient is assumed to be smooth (see [17]).

######  Assumption 2.

($\beta-$ Smoothness): The function $l(h_{\bm{w}}(\bm{x}),y)$ is assumed to be
$\beta-$smooth in $\bm{w}$, i.e., $\mid\left(\nabla
l(h_{\bm{w}_{1}}(\bm{x}),y)_{i}-\nabla
l(h_{\bm{w}_{2}}(\bm{x}),y)_{i}\right)\mid\leq
L\mid(\bm{w}_{1}-\bm{w}_{2})_{i}\mid$ for any
$\bm{w}_{1},\bm{w}_{2}\in\mathbb{R}^{d},\bm{x}\in\mathcal{X},y\in\mathcal{Y},i=1,2,\ldots,d,$
and $0<L_{i}<\infty$ for all $i$.

Note that the above assumption implies that

$\left|l(h_{\bm{w}_{1}}(\bm{x}),y)-\left[l(h_{\bm{w}_{2}}(\bm{x}),y)+{\nabla
l(h_{\bm{w}_{1}}(\bm{x}),y)}^{T}(\bm{w}_{1}-\bm{w}_{2})\right]\right|\leq\frac{1}{2}\sum_{i}L_{i}\left(w_{1i}-w_{2i}\right)^{2},$

for any $\bm{w}_{1},\bm{w}_{2}\in\mathbb{R}^{d},\bm{x}\in\mathcal{X}$ and
$y\in\mathcal{Y}.$

######  Assumption 3.

(Bounded Variance): Assume that every component of the estimated gradient,
i.e., $\hat{g}_{kj,i}$ is an unbiased estimate of the true gradient
$g_{kj,i}$, i.e., $\mathbb{E}\left[\hat{g}_{kj,i}\right]=g_{kj,i}$, for
$k,j=1,2,\ldots,N$ and $i=1,2,\ldots,n_{j}$. Further, the variance of every
component is bounded by some $\sigma_{kj,i}^{2}<\infty$, i.e.,
$\mathbb{E}\left[{(\hat{g}_{kj,i}-g_{kj,i})}^{2}\right]\leq\sigma_{kj,i}^{2}.$

Since $\bm{\alpha}_{k}$ should satisfy the constraint that
$\sum_{i=1}^{N}\alpha_{ki}=1$, the gradient descent step is followed by a
projection. The following definition comes in handy while deriving the
convergence results for the proposed algorithm.

######  Definition 4.

(Projected gradient [40]) Let $\Psi:\mathcal{K}\rightarrow\mathbb{R}$ be a
function on a closed convex set $\mathcal{K}\subseteq\Delta_{N}$. The
projected gradient of $\bm{z}\in\mathbb{R}^{d}$ with respect to $\Psi$ denoted
$\nabla_{\mathcal{K},\bm{z}}\Psi:\mathcal{K}\rightarrow\mathbb{R}^{N}$ is
defined as

$\nabla_{\mathcal{K},\bm{z}}\Psi:=\frac{1}{\mu}(\bm{z}-\Pi_{\mathcal{K}}[\bm{z}-\mu\nabla_{\bm{z}}\Psi(\bm{z})]),$
(11)

where
$\Pi_{\mathcal{K}}(\bm{z})=\underset{{\bm{x}\in\Delta_{N}}}{\operatorname{argmin}}\parallel\bm{x}-\bm{z}\parallel$
is the projection operator, and any $\mu>0$.

Note that if $\Pi_{\mathcal{K}}(\bm{z})=\bm{z}$, then the above coincides with
the gradient. Further, gradient update using the above ensures that
$\sum_{j=1}^{N}\alpha_{kj}=1$. In order to prove convergence of the Algorithm
$2$, it suffices to prove that $\bm{w}_{k}$ and $\bm{\alpha}_{k}$ converges
for all $k=1,2,\ldots,N$. Therefore $\Psi_{W,\bm{\lambda},\bm{\alpha}}$ for
the $k$-th component is written as,

$\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}=\Phi_{\bm{w}_{k},{\lambda_{k}},\bm{\alpha_{k}}}(\bm{S})+\texttt{Reg}({\lambda_{k}},\bm{\alpha_{k}})+\lambda_{k}\sum_{j=1}^{N}\alpha_{kj}{d}_{kj},$
(12)

where
$\Phi_{\bm{w}_{k},{\lambda_{k}},\bm{\alpha_{k}}}(\bm{S}):=\lambda_{k}\sum_{j=1}^{N}\alpha_{kj}\mathcal{L}_{j}(\bm{w}_{k}),$
and

$\texttt{Reg}({\lambda_{k}},\bm{\alpha_{k}}):=\frac{M}{N}\sqrt{\frac{N}{2}\sum_{j=1}^{N}\sum_{k=1}^{N}\left(\frac{\lambda_{k}\alpha_{kj}}{n_{j}}\right)^{2}\log\left(\frac{\mid\bm{\Lambda}_{\epsilon}\mid}{\delta}\right)}.$

It is important to note that the term corresponding to the $L_{2}$-regularizer
$\gamma_{k}\left|\bm{w}_{k}\right|$ is ignored, i.e., $\gamma_{k}=0$ in order
to prove the convergence result. However, the proof can be easily extended to
the case of $\gamma_{k}\neq 0$. The proof of convergence requires the
objective $\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}$ in (12) to be
Lipschitz function of $\bm{\alpha}_{k}$ and $\bm{w}_{k}$. This is the essence
of the following Lemmas.

######  Proposition 1.

The function
$\texttt{Reg}(\lambda_{k},\bm{\alpha_{k}}):=\frac{M}{N}\sqrt{\frac{N}{2}\sum_{j=1}^{N}\sum_{k=1}^{N}\left(\frac{\lambda_{k}\alpha_{kj}}{n_{j}}\right)^{2}\log\left(\frac{\mid\bm{\Lambda}_{\epsilon}\mid}{\delta}\right)}$
is Lipschitz in $\bm{\alpha}_{k}$ with Lipschitz constant
$\beta^{\prime}:=\frac{M}{\sqrt{2N}}\sqrt{\log\left(\frac{\mid\bm{\Lambda}_{\epsilon}\mid}{\delta}\right)}$.

_Proof:_ The proof is provided in Appendix B. ∎

######  Proposition 2.

The function $\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}$ is Lipschitz in
$\bm{\alpha}_{k}$ with Lipschitz constant
$\beta:=\beta^{\prime}+2\lambda_{k}M$ and Lipschitz in $\bm{w}_{k}$ with
Lipschitz constant $\lambda_{k}d$.

_Proof:_ The proof is provided in Appendix C. ∎

The following theorem uses the above results and definitions to show that
Algorithm $2$ converges.

######  Theorem 2.

(Convergence of Algorithm $2$). After $T$ iterations, choosing the learning
rates $\eta^{t}=\frac{1}{\sqrt{T}},\;\mu^{t}=\frac{1}{\sqrt{T}}$ and the batch
size $n_{t}=T$, the following holds:

$\displaystyle\mathbb{E}\left[\frac{1}{T}\overset{T-1}{\underset{{t=0}}{\sum}}\Delta_{k}^{t}\right]\leq\frac{1}{\sqrt{T}}\Bigg{(}2\lambda_{k}^{2}\sum_{m=1}^{N}\parallel\bm{\sigma}_{km}\parallel_{1}+\frac{\parallel
L\parallel_{1}\lambda_{k}^{2}}{2}+\Psi_{\bm{w}_{k}^{0},\lambda_{k},\bm{\alpha}_{k}^{0}}-\Psi_{k}^{*}\Bigg{)},$
(13)

where
$\Delta_{k}^{t}:=\left(\lambda_{k}^{2}\sum_{m=1}^{N}\alpha_{km}^{t+1}\left\|\bm{g}_{km}^{t}\right\|_{1}\right)+\left(1-\frac{\beta}{2\sqrt{T}}\right)\left[\left\|\nabla_{\mathcal{K},\bm{\alpha}_{k}}\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}}\right\|_{2}^{2}\right]$.

_Proof:_ The proof is provided in Appendix D.
Note that as $T\rightarrow\infty$, the right hand side goes to zero. In other
words, each term on the right hand side of (13) can be made arbitrarily small
by choosing appropriately large $T$. Since the first term corresponds to the
average gradient of $\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}$ with
respect to $\bm{w}_{k}$ scaled by $\lambda_{k}$, and the second term
corresponds to the gradient of $\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}$
with respect to $\bm{\alpha}_{k}$, there exists a time $t$ beyond which the
gradient of $\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}$ is small. This
shows that rate of convergence is $\mathcal{O}\left(\frac{1}{\sqrt{T}}\right)$
similar to [17]. In the next section, using the results of this section, a
convergence result of the proposed algorithm when the channel between the edge
device and the BS is a noisy wireless channel is presented.

##  VI Convergence analysis in a Wireless Channel

The above analysis holds good in the presence of an error free channel. For a
more pragmatic approach, the following analysis on a noisy channel is
presented. The channels between each device $k$ and the BS is considered to be
a single tap Rayleigh flat fading channel, with channel coefficient $h_{k}$ as
defined in section II-A.444The time slot index $t$ in $h_{k}(t)$ is ignored,
and is understood from the context. It is assumed that the entries of
$\operatorname{sign}(\hat{g}_{ki})$ is i.i.d with the probability of $1$ being
$q$. It is clear that if $d$ bits are sent without any errors, then the
convergence is guaranteed, as in Theorem 3. However, in the wireless channel,
it is expected to have errors, which depends on the SNR of the channel.
Therefore, it is important to investigate the impact of the SNR on the
convergence. Towards this, define the outage event as
$\mathcal{O}:=\left\\{d\leq\log_{2}\left(1+\displaystyle\frac{|h_{k}|^{2}{P}_{k}}{B\sigma^{2}}\right)\right\\}$,
where the maximum transmissible power by a user is ${P}_{k}$ and the channel
noise variance is $\sigma^{2}$ for channel bandwidth $B$.555The symbol
$\mathcal{O}$ is used to represent both outage as well as “order of”. It
should be clear depending on the context. It is assumed that in the case of
outage event, the errors are bound to happen, and hence the communication is
said to have failed. Strictly speaking, even in the case of outage, it is
possible that a few bits will get through without any errors, which can be
used to move roughly in the direction of gradient. This can potentially help
in improving the convergence. However, for the sake of simplicity, the above
case is ignored. Further, the discrepancy estimates assume error free channel.
As mentioned earlier, the impact of scheduling on the convergence is ignored.
In this setting, the following theorem characterizes the convergence of the
proposed algorithm under fading channel.

######  Theorem 3.

In a Rayleigh fading up-link channel with $SNR_{k}:=\frac{P_{k}}{B\sigma^{2}}$
and bandwidth $B$ at the device $k$, by choosing the learning rates, and batch
size as in Theorem 2, the following bound holds good

$\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}[\Delta_{k}^{t}]\leq\frac{U(\displaystyle
2^{\frac{d}{B}}-1)}{SNR_{k}}+\frac{1}{\sqrt{T}}\Bigg{(}2\lambda_{k}^{2}\sum_{m=1}^{N}\parallel\bm{\sigma}_{km}\parallel_{1}+\frac{\parallel
L\parallel_{1}\lambda_{k}^{2}}{2}+\Psi_{\bm{w}_{k}^{0},\lambda_{k},\bm{\alpha}_{k}^{0}}-\Psi_{k}^{*}\Bigg{)},$

where $U:=\beta^{\prime}+2\lambda_{k}M+\lambda_{k}d$.

It can be observed from the above theorem that the average is small provided
that the quantity
$\mathcal{O}\left(\max\left\\{\frac{1}{SNR_{k}},\frac{1}{\sqrt{T}}\right\\}\right)$
is small. This ensures that there exist a $t$ for which $\Delta_{k}^{t}$ is
small, making the sum of the gradients of
$\Psi_{\bm{w}_{k},\bm{\alpha}_{k},\lambda_{k}}$ with respect to
$\bm{\alpha}_{k}$ and the neural network weights $\bm{w}_{k}$ small, ensuring
that the solution is close to a sub-optimal minimum. However, for a fixed
$SNR_{k}$, it is useless to train for more that $\mathcal{O}(SNR_{k}^{2})$
number of iterations. Thus, the SNR of the transmission acts as a bottleneck
while training in the fading channel scenario. In addition, the higher the
Bandwidth, lower the number of iterations required. These observations are
made in the experimental results as well, which is detailed in the next
section.

##  VII Experimental Results

The experimental setup uses the MNIST handwriting data set to emulate a
federated learning scenario with $N=30$ devices. More specifically, three
cohorts of devices denoted $A$, $B$ and $C$ were constructed. Further, $12$
devices were assigned to each cohorts $A$ and $B$ while $6$ devices were
allocated to the cohort $C$. All the devices were equipped with a neural
networks capable of performing a $10$ class prediction on the images. As a
part of training, all the devices in cohort $A$ had the digits $0-5$, devices
in $B$ were given digits $6-9$ and devices in $C$ were given digits $3-7$.
This emulates the heterogeneous data set across cohorts while having
homogeneous data within each cohort. It is important to note that the cohorts
are _unknown_. There were around $100$ samples in every device out of which
$20$ were used for training and $80$ for testing. This mimics edge devices
with very less data and the larger number of samples for testing was set to
gain insights on the generalisation capabilities of the algorithm. In order to
compare the proposed algorithm with some baseline performance, each device was
also trained using just local data. The proposed algorithm is also compared
with the popular FedSGD and FedAVG (see [41] and [42]) algorithms. In the
following section, the performance of the proposed algorithm when the channel
is noise free is presented.

###  VII-A Perfect Uplink-Channel

In this scenario, the communication between devices and the BS occurs in an
error free regime. To begin with, the discrepancy is computed for each pair of
devices using Algorithm 1 described in Sec. IV-A. Using this in Algorithm 2,
the respective neural network weights are computed. Figures 2(a) and 2(b)
present the accuracy of various algorithms during training and testing phases
versus communication rounds. Here, in each communication round, the gradient
is updated using step $8$ of the Algorithm 2. It is observed that both FedSGD
and sign FedSGD perform similarly albeit the latter being more unstable in the
initial few rounds of communication. MtFEEL also acts unstably, but to a much
lesser degree as depicted in figure 2(a). This instability can be attributed
to the fact that the quantization (binarization to be precise) errors in the
initial stages will be large as the magnitude of the gradients are much larger
than $1$. The fluctuations eventually cease after around $150$ communication
rounds as the gradient approaches the local minimum. It is important to note
that the testing accuracy of the proposed algorithm is better than local
training, FedAvg and FedSGD. On the other hand, the training accuracy of the
proposed algorithm is inferior compared to FedAvg and FedSGD; this is
attributed to the over-fitting of the algorithms. The superior performance of
the proposed algorithm is due to the fact that the discrepancy for all the
devices in a cluster is small. Hence the solution to the optimization problem
results in close to equal weights being allocated to the devices within a
cluster and approximately zero weights across clusters during the learning
phase. In summary, it was observed that for devices in a cluster, the
algorithm converged to FedSGD, as expected. Recall that the cluster devices
are unknown, and the algorithm managed to learn them quite well.

(a) Performance from start to finish.

(b) Performance towards the end.

Figure 2: Accuracy in the case of Error free channel.

The MtFEEL average loss (see (7)) versus communication rounds is depicted in
Fig. 3 (a). It reaches its minimum at around $150$ iterations, which is when
no more significant gradient descent steps are being taken. This is also shown
in Figs. 3 (b) and (c) which depict
$\parallel\bm{w}_{k}^{t+1}-\bm{w}_{k}^{t}\parallel_{2}^{2}$ and
$\parallel\bm{\alpha}_{k}^{t+1}-\bm{\alpha}_{k}^{t}\parallel_{2}^{2}$ averaged
across all devices $k$, varying across communication rounds. This confirms the
convergence of the MtFEEL algorithm. These experiments demonstrate a proof of
concept and more elaborate experiments using different data sets is relegated
to future work. In the next subsection, the performance of the proposed
algorithm under fading channel is presented.

(a) MtFEEL loss

(b) Neural Network Weights

(c) Importance Coefficients

Figure 3: MtFEEL parameters across communication rounds in an error free
channel.

###  VII-B Noisy Uplink Channel

The MtFEEL algorithm is sensitive to the loss values incurred by a model when
run on different devices. The estimates of the importance coefficients
$\bm{\alpha}$ being sensitive to these loss values, can be unreliable if the
initial gradients are erroneous. It is observed that for a given uplink
$SNR_{k}$ for device $k$, increasing the number of communication rounds may
not help in improving the performance. In fact, beyond some threshold on $T$
for a given $SNR$, the performance can degrade due to the fact that the
gradient is bounded away from zero as $T$ increases, as indicated in Theorem
3. This is empirically observed in Fig. 4 (a). For $SNR_{k}<-10$dB, the
performance degrades, and results in very poor accuracy. In addition to the
above, a bit flipping model is also considered, where each component of the
gradient is independently flipped to $1$ (or $-1$) with a probability of $p$
(or $1-p$). The accuracy versus $p$ is plotted in Fig. 4 (b). It is clear that
beyond a threshold on $p$ ($p>0.2$), the performance of the proposed algorithm
improves, an observation in line with fading channel case. In summary, it is
better to use the proposed algorithm in most practical regimes of interest,
while the classical FedSGD approaches are better for very low SNR due to its
robustness for large number of errors. The reason behind this is that the
signed FedSGD doesn’t attempt to find similar devices to aggregate, and
naively considers all devices homogeneous, and can maintain its robustness for
relatively higher error probabilities.

(a) Rayeigh Fading Channel

(b) Bit flipping channel

Figure 4: Channel with Error.

##  VIII Conclusions

This work proposed a distributed FL algorithm across multiple devices that
results in custom neural network for each device. In particular, every device
learns its model with the help of other “similar” devices by sending signed
gradient information to a central BS. Each device aims to minimise an estimate
of a proposed loss using a weighted average of the empirical losses across
devices. These weights, called the importance coefficients, are dependent on
the similarity of data distributions between any pair of clients in the
network. This loss function is minimised by computing neural network weights
tailor made for each device. Theoretical guarantees on the proposed estimation
method are provided, and an algorithm is devised to compute the importance
coefficients and the neural network weights across devices. The guarantee is
depends on the weighted average of the losses, a notion called discrepancy
which is a measure of the dependency of the data across devices with respect
to the loss function, and a penalty term. An algorithm is proposed to estimate
this discrepancy in a distributed fashion. The FL algorithm was shown to
converge at the rate of $1/\sqrt{T}$, where $T$ is the number of communication
rounds when no errors are present in the communication links in the network.
In the case of a Rayleigh flat fading channel, the convergence of the
algorithm is shown to depend on the SNR and $1/\sqrt{T}$. In particular, it
was shown that the convergence is limited by the SNR, i.e., at low SNR, the
increase in communication rounds $T$ would not help in convergence.
Empirically, using the MNIST data set, the proposed algorithm was compared
with FedSGD, local training and FedAvg, and was shown to outperform these
algorithms.

## References

  * [1] T. Kraska, A. Talwalkar, J. C. Duchi, R. Griffith, M. J. Franklin, and M. I. Jordan, “Mlbase: A distributed machine-learning system.” in _Cidr_ , vol. 1, 2013, pp. 2–1.
  * [2] J. Konečnỳ, H. B. McMahan, D. Ramage, and P. Richtárik, “Federated optimization: Distributed machine learning for on-device intelligence,” _arXiv preprint arXiv:1610.02527_ , 2016.
  * [3] M. Mohri, G. Sivek, and A. T. Suresh, “Agnostic federated learning,” in _International Conference on Machine Learning_. PMLR, 2019, pp. 4615–4625.
  * [4] K. Bonawitz, H. Eichner, W. Grieskamp, D. Huba, A. Ingerman, V. Ivanov, C. Kiddon, J. Konecny, S. Mazzocchi, H. B. McMahan _et al._ , “Towards federated learning at scale: System design,” _arXiv preprint arXiv:1902.01046_ , 2019.
  * [5] Q. Yang, Y. Liu, T. Chen, and Y. Tong, “Federated machine learning: Concept and applications,” _ACM Transactions on Intelligent Systems and Technology (TIST)_ , vol. 10, no. 2, pp. 1–19, 2019.
  * [6] T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith, “Federated optimization in heterogeneous networks,” _arXiv preprint arXiv:1812.06127_ , 2018.
  * [7] A. Hard, K. Rao, R. Mathews, S. Ramaswamy, F. Beaufays, S. Augenstein, H. Eichner, C. Kiddon, and D. Ramage, “Federated learning for mobile keyboard prediction,” _arXiv preprint arXiv:1811.03604_ , 2018.
  * [8] B. McMahan and D. Ramage, “Federated learning: Collaborative machine learning without centralized training data,” _Google Research Blog_ , vol. 3, 2017\.
  * [9] H. B. McMahan, E. Moore, D. Ramage, S. Hampson _et al._ , “Communication-efficient learning of deep networks from decentralized data,” _arXiv preprint arXiv:1602.05629_ , 2016.
  * [10] S. P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. T. Suresh, “Scaffold: Stochastic controlled averaging for federated learning,” in _International Conference on Machine Learning_. PMLR, 2020, pp. 5132–5143.
  * [11] S. P. Karimireddy, M. Jaggi, S. Kale, M. Mohri, S. J. Reddi, S. U. Stich, and A. T. Suresh, “Mime: Mimicking centralized stochastic algorithms in federated learning,” _arXiv preprint arXiv:2008.03606_ , 2020.
  * [12] J. Wang, V. Tantia, N. Ballas, and M. Rabbat, “Slowmo: Improving communication-efficient distributed sgd with slow momentum,” _arXiv preprint arXiv:1910.00643_ , 2019.
  * [13] T. Lin, S. P. Karimireddy, S. U. Stich, and M. Jaggi, “Quasi-global momentum: Accelerating decentralized deep learning on heterogeneous data,” _arXiv preprint arXiv:2102.04761_ , 2021.
  * [14] H. B. McMahan, D. Ramage, K. Talwar, and L. Zhang, “Learning differentially private recurrent language models,” _arXiv preprint arXiv:1710.06963_ , 2017\.
  * [15] K. Bonawitz, V. Ivanov, B. Kreuter, A. Marcedone, H. B. McMahan, S. Patel, D. Ramage, A. Segal, and K. Seth, “Practical secure aggregation for privacy-preserving machine learning,” in _Proceedings of the 2017 ACM SIGSAC Conference on Computer and Communications Security_ , 2017, pp. 1175–1191.
  * [16] J. Konečnỳ, H. B. McMahan, F. X. Yu, P. Richtárik, A. T. Suresh, and D. Bacon, “Federated learning: Strategies for improving communication efficiency,” _arXiv preprint arXiv:1610.05492_ , 2016.
  * [17] J. Bernstein, Y.-X. Wang, K. Azizzadenesheli, and A. Anandkumar, “signsgd: Compressed optimisation for non-convex problems,” _arXiv preprint arXiv:1802.04434_ , 2018.
  * [18] G. Zhu, Y. Wang, and K. Huang, “Broadband analog aggregation for low-latency federated edge learning,” _IEEE Transactions on Wireless Communications_ , 2019.
  * [19] M. M. Amiri and D. Gündüz, “Machine learning at the wireless edge: Distributed stochastic gradient descent over-the-air,” in _2019 IEEE International Symposium on Information Theory (ISIT)_. IEEE, 2019, pp. 1432–1436.
  * [20] K. Yang, T. Jiang, Y. Shi, and Z. Ding, “Federated learning via over-the-air computation,” _IEEE Transactions on Wireless Communications_ , 2020.
  * [21] G. Zhu, Y. Du, D. Gunduz, and K. Huang, “One-bit over-the-air aggregation for communication-efficient federated edge learning: Design and convergence analysis,” _arXiv preprint arXiv:2001.05713_ , 2020.
  * [22] M. M. Wadu, S. Samarakoon, and M. Bennis, “Joint client scheduling and resource allocation under channel uncertainty in federated learning,” _IEEE Transactions on Communications_ , pp. 1–1, 2021.
  * [23] S. Samarakoon, M. Bennis, W. Saad, and M. Debbah, “Federated learning for ultra-reliable low-latency v2v communications,” in _2018 IEEE Global Communications Conference (GLOBECOM)_. IEEE, 2018, pp. 1–7.
  * [24] C. Feng, Z. Zhao, Y. Wang, T. Q. S. Quek, and M. Peng, “On the design of federated learning in the mobile edge computing systems,” _IEEE Transactions on Communications_ , pp. 1–1, 2021.
  * [25] V. Smith, C.-K. Chiang, M. Sanjabi, and A. S. Talwalkar, “Federated multi-task learning,” in _Advances in Neural Information Processing Systems_ , I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, Eds., vol. 30. Curran Associates, Inc., 2017. [Online]. Available: https://proceedings.neurips.cc/paper/2017/file/6211080fa89981f66b1a0c9d55c61d0f-Paper.pdf
  * [26] Y. Deng, M. M. Kamani, and M. Mahdavi, “Adaptive personalized federated learning,” _arXiv preprint arXiv:2003.13461_ , 2020.
  * [27] F. Hanzely and P. Richtárik, “Federated learning of a mixture of global and local models,” _arXiv preprint arXiv:2002.05516_ , 2020.
  * [28] Y. Mansour, M. Mohri, J. Ro, and A. T. Suresh, “Three approaches for personalization with applications to federated learning,” _arXiv preprint arXiv:2002.10619_ , 2020.
  * [29] T. Li, S. Hu, A. Beirami, and V. Smith, “Ditto: Fair and robust federated learning through personalization,” 2021.
  * [30] P. P. Liang, T. Liu, L. Ziyin, N. B. Allen, R. P. Auerbach, D. Brent, R. Salakhutdinov, and L.-P. Morency, “Think locally, act globally: Federated learning with local and global representations,” _arXiv preprint arXiv:2001.01523_ , 2020.
  * [31] A. Fallah, A. Mokhtari, and A. Ozdaglar, “Personalized federated learning: A meta-learning approach,” _arXiv preprint arXiv:2002.07948_ , 2020.
  * [32] F. Sattler, K.-R. Müller, and W. Samek, “Clustered federated learning: Model-agnostic distributed multi-task optimization under privacy constraints,” 2019.
  * [33] A. Ghosh, J. Chung, D. Yin, and K. Ramchandran, “An efficient framework for clustered federated learning,” in _Advances in Neural Information Processing Systems_ , H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, Eds., vol. 33. Curran Associates, Inc., 2020, pp. 19 586–19 597. [Online]. Available: https://proceedings.neurips.cc/paper/2020/file/e32cc80bf07915058ce90722ee17bb71-Paper.pdf
  * [34] C. Briggs, Z. Fan, and P. Andras, “Federated learning with hierarchical clustering of local updates to improve training on non-iid data,” in _2020 International Joint Conference on Neural Networks (IJCNN)_. IEEE, 2020, pp. 1–9.
  * [35] M. Zhang, K. Sapra, S. Fidler, S. Yeung, and J. M. Alvarez, “Personalized federated learning with first order model optimization,” _arXiv preprint arXiv:2012.08565_ , 2020.
  * [36] B. Guedj, “A primer on pac-bayesian learning,” 2019.
  * [37] V. Smith, C.-K. Chiang, M. Sanjabi, and A. S. Talwalkar, “Federated multi-task learning,” in _Advances in Neural Information Processing Systems_ , 2017, pp. 4424–4434.
  * [38] Mehryar, A. Rostamizadeh, and A. Talwalkar, _The Foundations of Machine Learning_ , 2012. [Online]. Available: http://mitpress.mit.edu/books/foundations-machine-learning-0
  * [39] S. Boyd, L. Xiao, and A. Mutapcic, “Subgradient methods,” _lecture notes of EE392o, Stanford University, Autumn Quarter_ , vol. 2004, pp. 2004–2005, 2003\.
  * [40] E. Hazan, K. Singh, and C. Zhang, “Efficient regret minimization in non-convex games,” _arXiv preprint arXiv:1708.00075_ , 2017.
  * [41] H. B. McMahan, E. Moore, D. Ramage, and B. A. y Arcas, “Federated learning of deep networks using model averaging,” _CoRR_ , vol. abs/1602.05629, 2016\. [Online]. Available: http://arxiv.org/abs/1602.05629
  * [42] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas, “Communication-efficient learning of deep networks from decentralized data,” in _Artificial Intelligence and Statistics_ , 2017, pp. 1273–1282.

##  Appendix A Proof Of Theorem 1

Recall that $\displaystyle
S_{k}:=\\{(\bm{x}_{k1},y_{k1}),(\bm{x}_{k2},y_{k2}),\ldots,(\bm{x}_{kn_{k}},y_{kn_{k}})\\}$
is the set of data samples present at device $k$. Each device $k$ is assumed
to have a neural network with weights
$\bm{w}_{k}\in\mathbb{R}^{d},k=1,2,\ldots N.$ Let
$W:=\\{\bm{w}_{1},\bm{w}_{2},\ldots\bm{w}_{N}\\}\subseteq\mathcal{R}^{dXN}$
denote the set of neural networks. The neural network weights $\bm{w}_{k}$ at
the device $k$ is optimized with respect to the loss function
$\mathcal{L}_{k}(\bm{w}_{j}):=\underset{(\bm{x}_{k},y_{k})\sim\mathcal{D}_{k}}{\mathbb{E}}\left[l\left(h_{\bm{w}_{j}}\left(\bm{x}\right),y\right)\right]$.
The weighted sum of the average losses of all the devices in the network is
given by

$\displaystyle\Phi_{W,\bm{\lambda}}=$
$\displaystyle\sum_{k=1}^{N}\lambda_{k}\mathcal{L}_{k}(\bm{w}_{k})=\sum_{k,j=1}^{N}\lambda_{k}\alpha_{kj}\mathcal{L}_{j}(\bm{w}_{k})+\sum_{k,j=1}^{N}\lambda_{k}\alpha_{kj}\Big{[}\mathcal{L}_{k}(\bm{w}_{k})-\mathcal{L}_{j}(\bm{w}_{k})\Big{]},$

where the above is obtained by using the facts that
$\sum_{j=1}^{N}\alpha_{kj}=1$ and $\sum_{k=1}^{N}\lambda_{k}=1$. Now using the
definitions of
$\Phi_{W,\bm{\lambda},\bm{\alpha}}:=\sum_{k,j=1}^{N}\lambda_{k}\alpha_{kj}\mathcal{L}_{j}(\bm{w}_{k})$
and the discrepancy between two devices $k$ and $j$ as $\displaystyle
d_{kj}:=\sup_{\bm{w}}\left|\mathcal{L}_{k}(\bm{w})-\mathcal{L}_{j}(\bm{w})\right|$,
the above can be upper bounded by

$\Phi_{W,\bm{\lambda}}\leq\Phi_{W,\bm{\lambda},\bm{\alpha}}+\sum_{k,j=1}^{N}\lambda_{k}\alpha_{kj}d_{kj}.$
(14)

Recall that
$\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})=\sum_{k=1}^{N}\lambda_{k}\sum_{j=1}^{N}\alpha_{kj}\hat{\mathcal{L}}(\bm{w}_{k},S_{j})$
is an estimate of ${\Phi}_{W,\bm{\lambda},\bm{\alpha}}$ Towards relating
$\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})$ to
$\Phi_{W,\bm{\lambda},\bm{\alpha}}$, define
$\psi\left(\bm{S}\right):=\sup_{\bm{w}\in
W}\left(\Phi_{W,\bm{\lambda},\bm{\alpha}}-\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})\right)$.
Consider two set of samples $\bm{S}^{\prime}:=\left\\{S_{1}^{\prime},\ldots
S_{N}^{\prime}\right\\}$ and $\bm{S}=\left\\{S_{1},\ldots S_{N}\right\\}$ that
differ only in a single element, say
$(\bm{x}^{\prime}_{ki},y^{\prime}_{ki})\in S^{\prime}_{k}$ and
$(\bm{x}_{ki},y_{ki})\in S_{k}$ where $k\in\\{1\dots N\\}$ and $i\in\\{1\dots
n_{k}\\}$. In order to apply McDiarmid’s inequality (see [38]), one needs to
bound the following difference

$\displaystyle\psi\left(\bm{S}^{\prime}\right)-\psi(\bm{S})$
$\displaystyle=\sup_{\bm{w}\in
W}\left(\Phi_{W,\bm{\lambda},\bm{\alpha}}-\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S}^{\prime})\right)-\sup_{\bm{w}\in
W}\left(\Phi_{W,\bm{\lambda},\bm{\alpha}}-\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})\right)$
$\displaystyle\stackrel{{\scriptstyle\text{(a)}}}{{\leq}}\sup_{\bm{w}\in
W}\left[\left(\Phi_{W,\bm{\lambda},\bm{\alpha}}-\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S}^{\prime})\right)-\left(\Phi_{W,\bm{\lambda},\bm{\alpha}}-\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})\right)\right]$
$\displaystyle{\leq}\sup_{\bm{w}\in
W}\left(\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})-\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S}^{\prime})\right)$
$\displaystyle=\sup_{\bm{w}\in
W}\left(\sum_{k,j=1}^{N}\lambda_{k}\alpha_{kj}\left[\hat{\mathcal{L}}(\bm{w}_{k},S_{j})-\hat{\mathcal{L}}(\bm{w}_{k},S_{j}^{\prime})\right]\right)\stackrel{{\scriptstyle\text{(b)}}}{{\leq}}\sum_{k=1}^{N}\frac{\lambda_{k}}{n_{j}}\alpha_{kj}M,$

where $(a)$ follows from the property of supremum, and $(b)$ follows from the
facts that the loss is assumed to be bounded, i.e.,
$\hat{\mathcal{L}}(\bm{w}_{k},S_{j})<M<\infty$, and $(\bm{x}_{ji},y_{ji})$ and
$(\bm{x}^{\prime}_{ji},y^{\prime}_{ji})$ differ in only one index $i$. Using
McDiarmid’s inequality for some $\delta>0$ and $\bm{w}\in W,$ it is easy to
see that the following holds with a probability of at-least $1-\delta$ (see
[38])

$\displaystyle\Phi_{W,\bm{\lambda},\bm{\alpha}}$
$\displaystyle\leq\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})+\underset{\bm{S}}{\mathbb{E}}[\psi(\bm{S})]+M\sqrt{\frac{1}{2}\sum_{j=1}^{N}\left(\sum_{k=1}^{N}\frac{\lambda_{k}\alpha_{kj}}{n_{j}}\right)^{2}\log\left(\frac{1}{\delta}\right)}.$

Let $\bm{\Lambda}_{\epsilon}$ be an $\epsilon$-cover of
$\bm{\Lambda}\subseteq\mathbb{R}^{n}$. By the definition of $\epsilon$-cover
(see definition 2), for any
$\bm{\lambda}\in\bm{\Lambda},\exists\;\bm{\lambda}_{\epsilon}\in\bm{\Lambda}_{\epsilon}$
such that
$\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})\leq\hat{\Phi}_{W,\bm{\lambda}_{\epsilon},\bm{\alpha}}(\bm{S})+\displaystyle
MN\epsilon$. Using this in (A) with the union bound, the following holds with
a probability of at-least $1-\delta$ (see [3])

$\displaystyle\Phi_{W,\bm{\lambda},\bm{\alpha}}\leq\hat{\Phi}_{W,\bm{\lambda}_{\epsilon},\bm{\alpha}}(\bm{S})+\underset{\bm{S}}{\mathbb{E}}[\psi(\bm{S})]+MN\epsilon+M\sqrt{\frac{1}{2}\sum_{j=1}^{N}\left(\sum_{k=1}^{N}\frac{\lambda_{k}\alpha_{kj}}{n_{j}}\right)^{2}\log\left(\frac{\mid\bm{\Lambda}_{\epsilon}\mid}{\delta}\right)}.$
(15)

Now, it remains to bound the term
$\underset{\bm{S}}{\mathbb{E}}[\psi(\bm{S})]$

$\displaystyle\underset{\bm{S}}{\mathbb{E}}[\psi(\bm{S})]$
$\displaystyle=\underset{\bm{S}}{\mathbb{E}}\left[\sup_{\bm{w}\in
W}\left(\Phi_{W,\bm{\lambda},\bm{\alpha}}-\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})\right)\right]=\underset{\bm{S}}{\mathbb{E}}\left[\sup_{\bm{w}\in
W}\underset{\bm{S}^{\prime}}{\mathbb{E}}\left(\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S}^{\prime})-\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})\right)\right]$
$\displaystyle\stackrel{{\scriptstyle\text{(a)}}}{{\leq}}\underset{\bm{S},\bm{S}^{\prime}}{\mathbb{E}}\left[\sup_{\bm{w}\in
W}\left(\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S}^{\prime})-\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})\right)\right]$
$\displaystyle=\underset{\bm{S},\bm{S}^{\prime}}{\mathbb{E}}\left[\sup_{\bm{w}\in
W}\left(\sum_{k,j,i=1}^{N,N,n_{j}}\frac{\lambda_{k}\alpha_{kj}}{n_{j}}\Big{[}l(h_{\bm{w}_{k}}(\bm{x}_{ji}^{\prime}),y_{ji}^{\prime})-l(h_{\bm{w}_{k}}(\bm{x}_{ji}),y_{ji})\Big{]}\right)\right],$

where $(a)$ follows from the Jensen’s inequality. Since
$\left(l(h_{\bm{w}_{k}}(\bm{x}_{ji}^{\prime}),y_{ji}^{\prime})-l(h_{\bm{w}_{k}}(\bm{x}_{ji}),y_{ji})\right)$
and
$\left(l(h_{\bm{w}_{k}}(\bm{x}_{ji}),y_{ji})-l(h_{\bm{w}_{k}}(\bm{x}_{ji}^{\prime}),y_{ji}^{\prime})\right)$
have the same distribution, the above can be written as

$\displaystyle\underset{\bm{S}}{\mathbb{E}}[\psi(\bm{S})]=$
$\displaystyle\underset{\bm{S},\bm{S}^{\prime},\bm{\sigma}}{\mathbb{E}}\left[\sup_{\bm{w}\in
W}\left(\sum_{k,j,i=1}^{N,N,n_{j}}\frac{\sigma_{kji}\lambda_{k}\alpha_{kj}}{n_{j}}\left(l(h_{\bm{w}_{k}}(\bm{x}_{ji}^{\prime}),y_{ji}^{\prime})-l(h_{\bm{w}_{k}}(\bm{x}_{ji}),y_{ji})\right)\right)\right]$
$\displaystyle\stackrel{{\scriptstyle\text{(a)}}}{{\leq}}$
$\displaystyle\underset{\bm{S}^{\prime},\bm{\sigma}}{\mathbb{E}}\left[\sup_{\begin{subarray}{c}\bm{w}\in
W\\\
\alpha\in\Delta_{N}\end{subarray}}\left(\sum_{k,j,i=1}^{N,N,n_{j}}\mathcal{(}B)^{{}^{\prime}}_{k,j,i}\right)\right]+\underset{\bm{S},\bm{\sigma}}{\mathbb{E}}\left[\sup_{\begin{subarray}{c}\bm{w}\in
W\\\
\alpha\in\Delta_{N}\end{subarray}}\left(-\sum_{k,j,i=1}^{N,N,n_{j}}\mathcal{(}B)_{kji}\right)\right]=2\mathcal{R}_{\bm{\Lambda}}\left(W\right),$

where the Rademacher random variable
$\sigma_{kj,i}\stackrel{{\scriptstyle\text{iid}}}{{\sim}}\\{-1,1\\}$,
$\mathcal{(}B)^{{}^{\prime}}_{k,j,i}:=\frac{\sigma_{kji}\lambda_{k}\alpha_{kj}}{n_{j}}l(h_{\bm{w}_{k}}(\bm{x}_{ji}^{\prime}),y_{ji}^{\prime})$,
and
$\mathcal{(}B)_{k,j,i}:=\frac{\sigma_{kji}\lambda_{k}\alpha_{kj}}{n_{j}}l(h_{\bm{w}_{k}}(\bm{x}_{ji}),y_{ji})$.
In the above, $(a)$ follows from the fact that $-\sigma_{kji}$ and
$\sigma_{kji}$ have the same distribution and the last inequality above
follows from the definition of minimax weighted Rademacher complexity, stated
in Definition (1). Using the result above, i.e.,
$\underset{\bm{S}}{\mathbb{E}}[\psi(\bm{S})]\leq
2\mathcal{R}_{\bm{\Lambda}}\left(W\right)$ in (15) results in

$\displaystyle\Phi_{W,\bm{\lambda},\bm{\alpha}}$
$\displaystyle\leq\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})+2\mathcal{R}_{\bm{\Lambda}}\left(W\right)+MN\epsilon+M\sqrt{\frac{1}{2}\sum_{j=1}^{N}\left(\sum_{k=1}^{N}\frac{\lambda_{k}\alpha_{kj}}{n_{j}}\right)^{2}\log\left(\frac{|\bm{\Lambda}_{\epsilon}|}{\delta}\right)}.$
(16)

Substituting (16) in (14) results in

$\displaystyle\Phi_{W,\bm{\lambda}}$
$\displaystyle\leq\hat{\Phi}_{W,\bm{\lambda},\bm{\alpha}}(\bm{S})+2\mathcal{R}_{\bm{\Lambda}}\left(W\right)+MN{\epsilon}+M\sqrt{\frac{N}{2}\log\left(\frac{\mid\bm{\Lambda}_{\epsilon}\mid}{\delta}\right)\sum_{k,j=1}^{N}\left(\frac{\lambda_{k}\alpha_{kj}}{n_{j}}\right)^{2}}+\sum_{k,j=1}^{N}\lambda_{k}\alpha_{kj}d_{kj},$

where the term in the square root is upper bounded using the property of norms
on vectors in $\mathbb{R}^{n}$;
$\parallel\cdot\parallel_{1}\leq\sqrt{n}\parallel\cdot\parallel_{2}$. This
completes the proof. ∎

##  Appendix B Proof of Proposition 1

The objective here is to show that the following is a Lipschitz continuous
function in $\alpha_{kj}$

$\displaystyle\texttt{Reg}(\lambda_{k},\bm{\alpha_{k}})=\frac{M}{N}\sqrt{\frac{N}{2}\log\left(\frac{\mid\bm{\Lambda}_{\epsilon}\mid}{\delta}\right)\sum_{j=1}^{N}\sum_{k=1}^{N}\left(\frac{\lambda_{k}\alpha_{kj}}{n_{j}}\right)^{2}}.$

Using the mean value theorem, it is sufficient to prove that the norm of the
gradient is finite, i.e., each component of
$\nabla_{\bm{\alpha}_{k}}\texttt{Reg}(\lambda_{k},\bm{\alpha_{k}})$ is finite.
The partial derivative of $\texttt{Reg}(\lambda_{k},\bm{\alpha_{k}})$ with
respect to $\alpha_{lm}$ for any $l$ and $m$ is given by

$\displaystyle\frac{\partial\;{\texttt{Reg}(\lambda_{k},\bm{\alpha_{k}})}}{\partial\alpha_{lm}}=\frac{M\sqrt{\frac{N}{2}\displaystyle\log\left(\frac{\mid\bm{\Lambda}_{\epsilon}\mid}{\delta}\right)}\frac{\left(\displaystyle\lambda_{k}^{2}\alpha_{km}\right)}{n_{m}^{2}}}{N\displaystyle\sqrt{\left(\displaystyle\frac{1}{2}\displaystyle\sum_{k,j=1}^{N}\left(\frac{\lambda_{k}\alpha_{kj}}{n_{j}}\right)^{2}\right)}}.$
(17)

Using the fact that $\displaystyle\frac{1}{\sqrt{1+x}}\leq 1$ for $x\geq 0$,
$\displaystyle\frac{\partial\;{\texttt{Reg}(\lambda_{k},\bm{\alpha_{k}})}}{\partial\alpha_{lm}}\leq\frac{M}{N}\sqrt{\frac{N}{2}\log\left(\frac{\mid\bm{\Lambda}_{\epsilon}\mid}{\delta}\right)}\frac{\lambda_{l}}{n_{m}}<\infty,$
for all $l$ and $m$. Now, recalling that $\lambda_{k}\in[0,1]$ and $n_{m}\geq
1$, it is easy to see that the $L_{2}$-norm of
$\nabla_{\bm{\alpha}_{k}}\texttt{Reg}(\lambda_{k},\bm{\alpha_{k}})$ is given
by

$\displaystyle\|\nabla_{\bm{\alpha}_{k}}\texttt{Reg}(\lambda_{k},\bm{\alpha_{k}})\|_{2}\leq\sqrt{\frac{N}{2}\log\left(\frac{\mid\bm{\Lambda}_{\epsilon}\mid}{\delta}\right)\frac{M^{2}\lambda_{k}}{N^{2}}\sum_{m=1}^{N}\frac{1}{n_{m}^{2}}}\leq{\beta^{\prime}}:=\frac{M}{\sqrt{2N}}\sqrt{\log\left(\frac{\mid\bm{\Lambda}_{\epsilon}\mid}{\delta}\right)}<\infty,$

where $\beta^{\prime}$ is the Lipschitz constant. This completes the proof. ∎

##  Appendix C Proof Of Proposition 2

Consider the objective function
$\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}$ in (12). To prove
$\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}$ is Lipschitz in
$\bm{\alpha}_{k}$, it suffices to prove
$\|\nabla_{\bm{\alpha}_{k}}\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}}\|_{2}<\infty$.
Taking the partial derivative with respect to $\bm{\alpha}_{km}$

$\displaystyle\frac{\partial\;\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}}{\partial\alpha_{km}}$
$\displaystyle=\frac{\partial\lambda_{k}\overset{N}{\underset{{j=1}}{\sum}}\alpha_{kj}\mathcal{L}_{j}(\bm{w}_{k})}{\partial\;\alpha_{km}}+\frac{\partial\;\gamma_{k}\parallel\bm{w}_{k}\parallel_{2}}{\partial\alpha_{km}}+\frac{\partial\;\texttt{Reg}({\lambda_{k}},\bm{\alpha_{k}})}{\partial\alpha_{km}}+\frac{\partial\;\lambda_{k}\sum_{j=1}^{N}\alpha_{kj}{d}_{kj}}{\partial\alpha_{km}}$
$\displaystyle=\lambda_{k}\mathcal{L}_{m}(\bm{w}_{k})+\frac{\partial\;\texttt{Reg}({\lambda_{k}},\bm{\alpha_{k}})}{\partial\alpha_{km}}+\lambda_{k}{d}_{km}\stackrel{{\scriptstyle\text{(a)}}}{{\leq}}\lambda_{k}M+\frac{\partial\;\texttt{Reg}({\lambda_{k}},\bm{\alpha_{k}})}{\partial\alpha_{km}}+\lambda_{k}M,$

where $(a)$ follows by upper bounding the loss, i.e.,
$\mathcal{L}_{m}(\bm{w}_{k})$ and discrepancy, i.e., $d_{km}$ by $M$ for all
$k$ and $m$. Now, it is easy to see that the $L_{2}$-norm of
$\nabla_{\bm{\alpha}_{k}}\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}$ is
given by

$\|\nabla_{\bm{\alpha}_{k}}\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}\|_{2}\leq\lambda_{k}M+\|\nabla_{\bm{\alpha}_{k}}\texttt{Reg}(\lambda_{k},\bm{\alpha_{k}})\|_{2}+\lambda_{k}M\stackrel{{\scriptstyle\text{(b)}}}{{=}}\beta^{\prime}+2\lambda_{k}M<\infty,$

where
${\beta^{\prime}}:=\frac{M}{\sqrt{2N}}\sqrt{\log\left(\frac{\mid\bm{\Lambda}_{\epsilon}\mid}{\delta}\right)}$,
and $(b)$ follows from Proposition 1. Next it remains to prove that
$\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}$ is Lipschitz in $\bm{w}_{k}$.
The gradient with respect to $\bm{w}_{k}$ is given by 666For simplicity
$\gamma_{k}$ is assumed to be $0$.
$\nabla_{\bm{w}_{k}}\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}=\lambda_{k}\overset{N}{\underset{{j=1}}{\sum}}\alpha_{kj}\texttt{sign}({\bm{g}}_{kj})$
and the $L_{2}$-norm of
$\nabla_{\bm{w}_{k}}\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}$ is
$\|\nabla_{\bm{w}_{k}}\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}\|_{2}\stackrel{{\scriptstyle\text{(c)}}}{{\leq}}\lambda_{k}d.$
Here, $(c)$ follows from the fact that
$\overset{N}{\underset{{j=1}}{\sum}}\alpha_{kj}=1$ and
$\texttt{sign}({\bm{g}}_{kj})\leq 1$ for all $k$ and $j$. The Lipschitz
constant of $\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}$ will thus be
$U:=\lambda_{k}d+\beta^{\prime}+2\lambda_{k}M$.

##  Appendix D Proof Of Theorem 2

Consider the $\beta$-smoothness assumption (see 2) of
$\Psi_{\bm{w}_{k},\lambda_{k},\bm{\alpha}_{k}}$ with respect to $\bm{w}_{k}$
for all $k=1,2,\ldots,N$ in (12)

$\displaystyle\Psi_{\bm{w}_{k}^{t+1},\lambda_{k},\bm{\alpha}_{k}^{t+1}}-\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t+1}}$
$\displaystyle\leq\langle\nabla_{\bm{w}_{k}}\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t+1}},\bm{w}_{k}^{t+1}-\bm{w}_{k}^{t}\rangle+\sum_{i=1}^{d}\frac{L_{i}}{2}\left(\bm{w}_{k}^{t+1}-\bm{w}_{k}^{t}\right)_{i}^{2},$

where $(\bm{x})_{i}$ is the $i^{th}$ component of the vector
$\bm{x}\in\mathbb{R}^{d}$. Define
$\Delta_{\bm{w}_{k}}\Psi_{\bm{w}_{k}^{t+1},\lambda_{k},\bm{\alpha}_{k}^{t+1}}:=\Psi_{\bm{w}_{k}^{t+1},\lambda_{k},\bm{\alpha}_{k}^{t+1}}-\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t+1}}$.
Now, using
$\bm{w}_{k}^{t+1}-\bm{w}_{k}^{t}=-\eta^{t}\nabla_{\bm{w}_{k},\operatorname{sign}}\hat{\Psi}_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t+1}}$
from step 7 of the Algorithm $2$, the above becomes

$\displaystyle\Delta_{\bm{w}_{k}}\Psi_{\bm{w}_{k}^{t+1},\lambda_{k},\bm{\alpha}_{k}^{t+1}}$
$\displaystyle\leq-\eta^{t}\nabla_{\bm{w}_{k}}\Psi_{{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t+1}}}^{T}\nabla_{\bm{w}_{k},\operatorname{sign}}\hat{\Psi}_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t+1}}+\sum_{i=1}^{d}\frac{L_{i}}{2}\left(-\eta^{t}\nabla_{\bm{w}_{k},\operatorname{sign}}\hat{\Psi}_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t+1}}\right)_{i}^{2}.$

Substituting for the true gradient
$\nabla_{\bm{w}_{k}}\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t+1}}=\lambda_{k}\sum_{m=1}^{N}\alpha_{km}^{t+1}\bm{g}_{km}^{t},$
and its estimate
$\nabla_{\bm{w}_{k},\operatorname{sign}}\hat{\Psi}_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t+1}}=\lambda_{k}\sum_{j=1}^{N}\alpha_{kj}^{t+1}\texttt{sign}\left(\hat{\bm{g}}_{kj}^{t}\right)$,
and after some algebraic manipulations, the above becomes

$\displaystyle\Delta_{\bm{w}_{k}}\Psi_{\bm{w}_{k}^{t+1},\lambda_{k},\bm{\alpha}_{k}^{t+1}}$
$\displaystyle\leq-\eta^{t}\left(\lambda_{k}^{2}\sum_{m=1}^{N}\alpha_{km}^{t+1}{\parallel\bm{g}_{km}^{t}\parallel}_{1}\right)+2\eta^{t}\Bigg{(}\lambda_{k}^{2}\sum_{m=1}^{N}\sum_{i=1}^{d}{\alpha_{km}^{t+1}}|g_{km,i}^{t}|\mathcal{F}_{k,i}^{t}\Bigg{)}$
$\displaystyle+\frac{(\eta^{t})^{2}}{2}\sum_{i=1}^{d}L_{i}\left(\lambda_{k}\sum_{j=1}^{N}\alpha_{kj}^{t+1}\operatorname{sign}\left(\hat{\bm{g}}_{kj}^{t}\right)\right)_{i}^{2},$

where
$\mathcal{F}_{k,i}^{t}:=\sum_{j=1}^{N}\alpha_{kj}^{t+1}\mathbbm{1}\left[\texttt{sign}\left(\hat{g}_{kj,i}^{t}\right)\neq\texttt{sign}\left(g_{km,i}^{t}\right)\right]\Bigg{)}$.
Using the fact that the $i^{th}$ component of the signed gradient is less than
or equal to $1$, i.e., $\texttt{sign}\left(\hat{\bm{g}}_{kj,i}^{t}\right)\leq
1$, the above can be further upper bounded as

$\displaystyle\Delta_{\bm{w}_{k}}\Psi_{\bm{w}_{k}^{t+1},\lambda_{k},\bm{\alpha}_{k}^{t+1}}$
$\displaystyle\leq-\eta^{t}\left(\lambda_{k}^{2}\sum_{m=1}^{N}\alpha_{km}^{t+1}{\parallel\bm{g}_{km}^{t}\parallel}_{1}\right)+2\eta^{t}\Bigg{(}\lambda_{k}^{2}\sum_{m=1}^{N}\sum_{i=1}^{d}{\alpha_{km}^{t+1}}|g_{km,i}^{t}|\mathcal{F}_{k,i}^{t}\Bigg{)}$
$\displaystyle+$ $\displaystyle\frac{(\eta^{t})^{2}}{2}\parallel
L\parallel_{1}\lambda_{k}^{2}.$

Consider the expected improvement at $t+1$ conditioned on the previous
iterate, i.e.,

$\displaystyle\mathbb{E}\left[\Delta_{\bm{w}_{k}}\Psi_{\bm{w}_{k}^{t+1},\lambda_{k},\bm{\alpha}_{k}^{t+1}}\mid\bm{w}_{k}^{t}\right]\leq-\eta^{t}\left(\lambda_{k}^{2}\sum_{m=1}^{N}\alpha_{km}^{t+1}{\parallel\bm{g}_{km}^{t}\parallel}_{1}\right)+$
$\displaystyle
2\eta^{t}\Bigg{(}\lambda_{k}^{2}\sum_{m,j=1}^{N}\sum_{i=1}^{d}{\alpha_{km}}\mid
g_{km,i}^{t}\mid\mathbb{P}\left[\texttt{sign}\left(\hat{g}_{kj,i}^{t}\right)\neq\texttt{sign}\left(g_{km,i}^{t}\right)\right]\Bigg{)}+\frac{(\eta^{t})^{2}}{2}\parallel
L\parallel_{1}\lambda_{k}^{2}.$ (18)

In order to bound the above further, consider the term

$\displaystyle\mathbb{P}\left[\texttt{sign}\left(\hat{g}_{kj,i}^{t}\right)\neq\texttt{sign}\left(g_{km,i}^{t}\right)\right]$
$\displaystyle\stackrel{{\scriptstyle\text{(a)}}}{{\leq}}\mathbb{P}\left[\mid\hat{g}_{kj,i}^{t}-g_{km,i}^{t}\mid\geq\mid
g_{km,i}^{t}\mid\right]\stackrel{{\scriptstyle\text{(b)}}}{{\leq}}$
$\displaystyle\frac{\mathbb{E}\left[\left|\hat{g}_{kj,i}^{t}-g_{km,i}^{t}\right|\right]}{\left|g_{km,i}^{t}\right|}$
(19)
$\displaystyle\stackrel{{\scriptstyle(c)}}{{\leq}}\frac{\sqrt{\mathbb{E}\left[{\left(\hat{g}_{kj,i}^{t}-g_{km,i}^{t}\right)}^{2}\right]}}{\left|g_{km,i}^{t}\right|}\stackrel{{\scriptstyle(d)}}{{\leq}}\frac{\sigma_{km,i}^{t}}{\left|g_{km,i}^{t}\right|}.$

In the above, $(a)$ follows from the fact that
$\\{\texttt{sign}\left(\hat{g}_{kj,i}^{t}\right)\neq\texttt{sign}\left(g_{km,i}^{t}\right)\\}\subseteq\\{\mid\hat{g}_{kj,i}^{t}-g_{km,i}^{t}\mid\geq\mid
g_{km,i}^{t}\mid\\}$, $(b)$ follows from the Markov’s inequality, $(c)$ is
obtained from the Jensen’s inequality, and $(d)$ follows since
$\hat{g}_{kj,i}^{t}$ is an unbiased estimate of $g_{km,i}^{t}$ and using the
definition of variance. The following is obtained by substituting (19) in (18)
and using $\sigma_{km,i}^{t}\leq\frac{\sigma_{km,i}}{\sqrt{n_{t}}}$

$\displaystyle\mathbb{E}\left[\Delta_{\bm{w}_{k}}\Psi_{\bm{w}_{k}^{t+1},\lambda_{k},\bm{\alpha}_{k}^{t+1}}\mid\bm{w}_{k}^{t+1}\right]\leq-\eta^{t}$
$\displaystyle\left(\lambda_{k}^{2}\sum_{m=1}^{N}\alpha_{km}^{t+1}\left\|\bm{g}_{km}^{t}\right\|_{1}\right)+2\eta^{t}\left(\lambda_{k}^{2}\sum_{m=1}^{N}\frac{{\alpha_{km}^{t+1}}\parallel\bm{\sigma}_{km}\parallel_{1}}{\sqrt{n_{t}}}\right)$
$\displaystyle+$ $\displaystyle\frac{(\eta^{t})^{2}}{2}\parallel
L\parallel_{1}\lambda_{k}^{2}.$

Since $\alpha_{km}\leq 1$ for all $k$ and $m$, the above can be further
bounded as

$\displaystyle\leq-\eta^{t}\left(\lambda_{k}^{2}\sum_{m=1}^{N}\alpha_{km}^{t+1}\left\|\bm{g}_{km}^{t}\right\|_{1}\right)+\left(2\eta^{t}\lambda_{k}^{2}\sum_{m=1}^{N}\frac{\parallel\bm{\sigma}_{km}\parallel_{1}}{\sqrt{n_{t}}}\right)+\frac{(\eta^{t})^{2}}{2}\parallel
L\parallel_{1}\lambda_{k}^{2}.$ (20)

Now, bounding the difference of the objective function when $\bm{w}_{k}^{t}$
is fixed while $\bm{\alpha}_{k}^{t}$ is a variable

$\displaystyle\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t+1}}-\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}}$
$\displaystyle\leq\left\langle\nabla_{\bm{\alpha}_{k}}\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}},\bm{\alpha}_{k}^{t+1}-\bm{\alpha}_{k}^{t}\right\rangle+\frac{\beta}{2}{\parallel\bm{\alpha}_{k}^{t+1}-\bm{\alpha}_{k}^{t}\parallel}_{2}^{2},$

where
$\beta:=M\Big{(}\sqrt{\frac{1}{2N}\log\left(\frac{\mid\bm{\Lambda}_{\epsilon}\mid}{\delta}\right)}+2\lambda_{k}\Big{)}$
is as defined in Lemma 2. Substituting
$\bm{\alpha}_{k}^{t+1}-\bm{\alpha}_{k}^{t}$ =
$-\mu^{t}\nabla_{\mathcal{K},\bm{\alpha}_{k}}\hat{\Psi}_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}}$
from step 8 of the Algorithm $2$, where
$\nabla_{\mathcal{K},\bm{\alpha}_{k}}\hat{\Psi}_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t+1}}$
denotes the projected gradient with respect to $\bm{\alpha}_{k}$, the
difference
$\Delta_{\bm{\alpha}_{k}}\Psi_{\bm{w}_{k}^{t+1},\lambda_{k},\bm{\alpha}_{k}^{t+1}}:=\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t+1}}-\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}}$
becomes
$\Delta_{\bm{\alpha}_{k}}\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}}=-\mu^{t}\left\langle\nabla_{\bm{\alpha}_{k}}\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}},\nabla_{\mathcal{K},\bm{\alpha}_{k}}\hat{\Psi}_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}}\right\rangle+\frac{\beta(\mu^{t})^{2}}{2}\left\|\nabla_{\mathcal{K},\bm{\alpha}_{k}}\hat{\Psi}_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}}\right\|_{2}^{2}$.
Using
$\left\langle\nabla_{\bm{\alpha}_{k}}\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}},\nabla_{\mathcal{K},\bm{\alpha}_{k}}\hat{\Psi}_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}}\right\rangle\geq\|\nabla_{\mathcal{K},\bm{\alpha}_{k}}\hat{\Psi}_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}}\|_{2}^{2}$
from Lemma 3.2 of [40], the above can be further bounded as

$\Delta_{\bm{\alpha}_{k}}\Psi_{\bm{w}_{k}^{t+1},\lambda_{k},\bm{\alpha}_{k}^{t}}\leq-\mu^{t}\left\|\nabla_{\mathcal{K},\bm{\alpha}_{k}}\hat{\Psi}_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}}\right\|_{2}^{2}+\frac{\beta(\mu^{t})^{2}}{2}\left\|\nabla_{\mathcal{K},\bm{\alpha}_{k}}\hat{\Psi}_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}}\right\|_{2}^{2}.$
(21)

Let
$\Psi_{k}^{*}:=\underset{\bm{w}_{k},\bm{\alpha}_{k}}{\min}\Psi_{\bm{w}_{k},\lambda_{k}\bm{\alpha}_{k}}$
such that $\sum_{j=1}^{N}\alpha_{kj}=1$. Towards completing the proof,
consider

$\displaystyle\Psi_{\bm{w}_{k}^{0},\lambda_{k},\bm{\alpha}_{k}^{0}}-\Psi_{k}^{*}$
$\displaystyle\stackrel{{\scriptstyle\text{(a)}}}{{\geq}}$
$\displaystyle\Psi_{\bm{w}_{k}^{0},\lambda_{k},\bm{\alpha}_{k}^{0}}-\mathbb{E}[\Psi_{\bm{w}_{k}^{t+1},\lambda_{k},\bm{\alpha}_{k}^{t+1}}]\stackrel{{\scriptstyle\text{(b)}}}{{=}}\mathbb{E}\sum_{t=0}^{T-1}[\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}}-\Psi_{\bm{w}_{k}^{t+1},\lambda_{k},\bm{\alpha}_{k}^{t+1}}]$
$\displaystyle\stackrel{{\scriptstyle\text{(c)}}}{{=}}$
$\displaystyle\mathbb{E}\sum_{t=0}^{T-1}\Big{[}-\Big{(}\Psi_{\bm{w}_{k}^{t+1},\lambda_{k},\bm{\alpha}_{k}^{t+1})}-\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t+1}}\Big{)}-\Big{(}\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t+1}}-\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}}\Big{)}\Big{]}$
$\displaystyle=$
$\displaystyle-\mathbb{E}\sum_{t=0}^{T-1}\Delta_{\bm{w}_{k}}\Psi_{\bm{w}_{k}^{t+1},\lambda_{k},\bm{\alpha}_{k}^{t+1}}-\mathbb{E}\sum_{t=0}^{T-1}\Delta_{\bm{\alpha}_{k}}\Psi_{\bm{w}_{k}^{t+1},\lambda_{k},\bm{\alpha}_{k}^{t+1}},$

where $(a)$ follows from the optimality of $\Psi_{k}^{*}$, $(b)$ follows from
the telescopic sum and $(c)$ follows by adding and subtracting
$\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t+1}}$. Substituting (20)
and (21) in (D) and choosing the learning rates $\eta^{t}=\frac{1}{\sqrt{T}}$
and $\mu^{t}=\frac{1}{\sqrt{T}}$, and batch size $n_{t}=T$ results in (13). ∎

##  Appendix E Proof Of Theorem 3

It suffices to show that
$\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}[\Delta_{k}^{t}]\to 0$ as $T\to\infty$ in
a Rayleigh fading channel. Here,
$\Delta_{k}^{t}:=\Bigg{[}\left(\lambda_{k}^{2}\sum_{m=1}^{N}\alpha_{km}^{t+1}\left\|\bm{g}_{km}^{t}\right\|_{1}\right)+\left(1-\frac{\beta}{2\sqrt{T}}\right)\left[\left\|\nabla_{\mathcal{K},\bm{\alpha}_{k}}\Psi_{\bm{w}_{k}^{t},\lambda_{k},\bm{\alpha}_{k}^{t}}\right\|_{2}^{2}\right]\Bigg{]}$.
From the total expectation law, it follows that

$\displaystyle\mathbb{E}[\Delta_{k}^{t}]=$
$\displaystyle\mathbb{E}[\Delta_{k}^{t}|\mathcal{O}]\mathbb{P}[O]+\mathbb{E}\left[\Delta_{k}^{t}|\mathcal{O}^{\complement}\right]\mathbb{P}[\mathcal{O}^{\complement}]{\leq}U\mathbb{P}[\mathcal{O}]+\mathbb{E}\left[\Delta_{k}^{t}|\mathcal{O}^{\complement}\right]\mathbb{P}[\mathcal{O}^{\complement}],$

where $\mathcal{O}$ is the outage event. The inequality above follows from the
fact that $\Delta_{k}^{t}$ is bounded above by the Lipschitz constant
$U:=\beta^{\prime}+2\lambda_{k}M+\lambda_{k}d$ from lemmas 1 and 2. Now from
the definition of outage event, and using the fact that
$\mathbb{P}[O^{\complement}]\leq 1$, the above can be bounded as

$\displaystyle\mathbb{E}[\Delta_{k}^{t}]\leq$ $\displaystyle
U\mathbb{P}\Bigg{\\{}d\geq
B\ln\left(1+\frac{\mathcal{P}_{k}|h_{k}|^{2}}{B\sigma^{2}}\right)\Bigg{\\}}+\mathbb{E}\left[\Delta_{k}^{t}|O^{\complement}\right]$
$\displaystyle=$ $\displaystyle
U\left(1-\exp\left\\{-\displaystyle\frac{(2^{\frac{d}{B}}-1)}{SNR_{k}}\right\\}\right)+\mathbb{E}\left[\Delta_{k}^{t}|O^{\complement}\right].$

Using the above, the average of $\Delta_{k}^{t}$ is bounded as follows

$\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}\Delta_{k}^{t}\stackrel{{\scriptstyle\text{(a)}}}{{\leq}}U\frac{(\displaystyle
2^{\frac{d}{B}}-1)}{SNR_{k}}+\frac{1}{\sqrt{T}}\Bigg{(}2\lambda_{k}^{2}\sum_{m=1}^{N}\parallel\bm{\sigma}_{km}\parallel_{1}+\frac{\parallel
L\parallel_{1}\lambda_{k}^{2}}{2}+\Psi_{\bm{w}_{k}^{0},\lambda_{k},\bm{\alpha}_{k}^{0}}-\Psi_{k}^{*}\Bigg{)},$

where the above follows from the facts that $1-e^{-x}\leq x$ and that when
there is no outage, the analysis is identical to the case in Appendix D. This
completes the proof of the theorem. ∎