In Federated Learning (FL), the distributed nature and heterogeneity of client data present both opportunities and challenges. While collaboration among clients can significantly enhance the learning process, not all collaborations are beneficial; some may even be detrimental. In this study, we introduce a novel algorithm that assigns adaptive aggregation weights to clients participating in FL training, identifying those with data distributions most conducive to a specific learning objective. We demonstrate that our aggregation method converges no worse than the method that aggregates only the updates received from clients with the same data distribution. Furthermore, empirical evaluations consistently reveal that collaborations guided by our algorithm outperform traditional FL approaches. This underscores the critical role of judicious client selection and lays the foundation for more streamlined and effective FL implementations in the coming years.

§ INTRODUCTION

Federated Learning (FL) introduces an innovative paradigm redefining traditional machine learning workflow. Instead of centrally pooling sensitive client data, FL allows for model training on decentralized data sources stored directly on client devices
[Konečný et al., 2016, Zhang et al., 2021, Li et al., 2020]. In this approach, rather than training Machine Learning (ML) models in a centralized manner, a shared model is distributed to all clients. Each client then performs local training, and model updates are exchanged between clients and the FL orchestrator (often referred to as the master server) [McMahan et al., 2017, Shokri & Shmatikov, 2015, Karimireddy et al., 2020].

Personalized Federated Learning (PFL).
The concept of PFL [Collins et al., 2021, Hanzely et al., 2020, Sadiev et al., 2022] has been gaining traction. In this framework, each client, often referred to as an agent, takes part in developing their own personalized model variant. This tailored training approach leverages local data distributions, aiming to design models that cater to the distinct attributes of each client's dataset [Fallah et al., 2020].
However, a prominent challenge arises in this decentralized training landscape due to the data's non-IID (independent and identically distributed) nature across various clients. Data distributions that differ considerably can have a pronounced impact on the convergence and generalization capabilities of the trained models.
While certain client-specific data distributions might strengthen model performance, others could prove detrimental, introducing biases or potential adversarial patterns.
Additionally, within the personalized federated learning paradigm, the emphasis on crafting individualized models could inadvertently heighten these data disparities
[Kairouz et al., 2021]. Consequently, this may lead to models that deliver subpar or, in some cases, incorrect results when applied to wider or diverse datasets [Kulkarni et al., 2020].

Collaboration as a service. In this paper, we introduce a modified protocol for FL that deviates from a strictly personalized approach. Rather than focusing solely on refining individualized models, our approach seeks to harness the advantages of distinct data distributions, curb the detrimental effects of outlier clients, and promote collaborative learning. Through this innovative training mechanism, our algorithm discerns which clients are optimal collaborators to ensure faster convergence and potentially better generalization.

§.§ Setup

Setup. We assume that there are $n$ clients participating in the training and consider the first one as a target client. The goal is to train the model for this client, i.e., we focus on the stochastic optimization problem

\begin{equation}
\textstyle{\min}_{x\in\R^d} \left\{f(x) \equiv f_1(x) := \E_{\xi_1 \sim \cD_1} [f_{\xi_1}(x)]\right\}, \label{eq:main_problem}
\end{equation}

where $f_{\xi_1}:\R^d \to \R$ is the loss function on sample $\xi_1$ and $f:\R^d \to \R$ is an expected loss. Other clients can also have data sampled from similar distributions, but we also allow adversarial participants, e.g., Byzantines [Lamport et al., 1982, Lyu et al., 2020]. That is, some clients can be beneficial for the training in certain stages, but they are not assumed to be known apriori.

§.§ Contribution

Our main contributions are listed below.

* New method: . We proposed a new method called Merit-based Federated Averaging for Diverse Datasets () that aims to solve (<ref>). The key idea is to use the stochastic gradients received from the clients to adjust the weights of averaging through the inexact solving of the auxiliary problem.

* Provable convergence under mild assumptions. We prove that converges not worse than SGD that averages only the stochastic gradients received from clients having the same data distribution (these clients are not known apriori) for smooth non-convex and Polyak-Lojasiewicz functions under standard bounded variance assumption.

* Utilizing all possible benefits. We numerically show that can even benefit from collaboration with clients having different data distributions when these distributions are close to the target one.

§.§ Related work

Federated optimization. Standard results in distributed/federated optimization focus on the  problem:

\begin{equation}
\textstyle{\min}_{x\in \R^d}\frac{1}{n}\sum_{i=1}^n f_i(x), \label{eq:standard_ERM}
\end{equation}

where $f_i(x)$ represents either expected or empirical loss on the client $i$. This problem significantly differs from (<ref>), since one cannot completely ignore the updates from some clients to achieve a better solution. Typically, in this case, communication is the main bottleneck of the methods for solving such problems. To address this issue one can use communication compression [Alistarh et al., 2017, Stich et al., 2018, Mishchenko et al., 2019], local steps [Stich, 2018, Khaled et al., 2020, Kairouz et al., 2021, Wang et al., 2021, Mishchenko et al., 2022], client importance sampling [Cho et al., 2020, Nguyen et al., 2020, Ribero & Vikalo, 2020, Lai et al., 2021, Luo et al., 2022, Chen et al., 2022], or decentralized protocols [Lian et al., 2017, Song et al., 2022]. However, these techniques are orthogonal to what we focus on in our paper, though incorporating them into our algorithm is a prominent direction for future research.

Clustered FL. Another way of utilizing benefits from the other clients is the clustering of clients based on some information about their data or personalized models. Tang et al., 2021 propose a personalized formulation with $\ell_2$-regularization that attracts a personalized model of a worker to the center of the cluster that this worker belongs to. A similar objective is studied by Ma et al., 2022. Ghosh et al., 2020 develop an algorithm that updates clusters's centers using the gradients of those clients that have the smallest loss functions at the considered cluster's center. It is worth mentioning that, in contrast to our work, the mentioned works modify the personalized objective to illustrate some benefits of collaboration while we focus on the pure personalized problem of the target client. Under the assumption that the data distributions of each client are mixtures of some finite set of underlying distributions, Marfoq et al., 2021 derive the convergence result for the Federated Expectation-Maximization algorithm. This is the closest work to our setup in the Clustered FL literature. However, in contrast to Marfoq et al., 2021, we do not assume that the gradients are bounded, and the local loss functions have bounded gradient dissimilarity. Another close work to ours is [Fraboni et al., 2021], where the authors consider so-called clustered-based sampling. However, Fraboni et al., 2021 also make a non-standard assumption on the bounded dissimilarity of the local loss functions, while one of the key properties of our approach is its robustness to arbitrary clients' heterogeneity.

Non-uniform averaging. There are also works studying the convergence of distributed SGD-type methods that use non-uniform (but fixed) weights of averaging. Ding & Wang, 2022 propose a method to detect collaboration partners and adaptively learn "several" models for numerous heterogeneous clients. Directed graph edge weights are used to calculate group partitioning. Since the calculation of optimal weights in their approach is based on similarity measures between clients' data, it is unclear how to compute them in practice without sacrificing the users's data privacy. Even et al., 2022 develop and analyze another approach for personalized aggregation, where each client filters gradients and aggregates them using fixed weights. The optimal weights also require estimating the distance between distributions (or communicating empirical means among all clients and estimating effective dimensions).
Both works do not consider weights evolving in time, which is one of the key features of our method.

Non-fixed weights are considered in [Wu & Wang, 2021], but the authors focus on non-personalized problem formulation. In particular, Wu & Wang, 2021 propose the method called FedAdp that uses cosine similarity between gradients and the Gompertz function for updating aggregation weights. Under the strong bounded local gradient dissimilarity assumption[Wu & Wang, 2021 assume that there exist constants $A,B > 0$ such that $A\|\nabla f(x)\| \leq \|\nabla f_i(x)\| \leq B\|\nabla f(x)\|$ for every client $i \in [n]$ and any $x$, where $f(x) = \frac{1}{n}\sum_{i=1}^n f_i(x)$.], Wu & Wang, 2021 derive a non-conventional upper bound (for the loss function at the last iterate of their algorithm) that does not necessarily imply convergence of the method.

Bi-level optimization.

Taking into account that we want to solve problem (<ref>) using the information coming from not only the target client, it is natural to consider the following bi-level optimization (BLO) problem formulation:

\begin{eqnarray}
\min_{\w\in \Delta^\n_1}&& \f(\x^*(\w)), \label{eq:blo} \\ \text{s.t.}&& \x^*(\w) \in \textstyle{\argmin}_{\x \in \R^d} \textstyle{\sum}_{i=1}^n \w_i \f_{i}(\x), \label{eq:LL_problem}
\end{eqnarray}

where $\Delta_1^n$ is a unit simplex in $\R^n$: $\Delta_1^n =  \{\w \in \R^n\mid \sum_{i=1}^n w_i = 1,\; w_i \geq 0 \;\forall i\in [n]\}$. The problem in (<ref>) is usually called the upper-level problem (UL), while the problem in (<ref>) is the lower-level (LL) one. Since in our case $f(x) \equiv f_1(x)$, (<ref>)-(<ref>) is equivalent to (<ref>). In the general case, this equivalence does not always hold and, in addition, function $f$ is allowed to depend on $w$ not only through $x^*$. All these factors make the general BLO problem hard to solve. The literature for this general class of problems is quite rich, and we cover only closely related works.

The closest works to ours are [Chen et al., 2021], which propose so-called Target-Aware  Weighted Training (TAWT), and its extension to the federated setup [Huang et al., 2022]. Their analysis relies on the existence of weights $\w$, such that $\dist\rbr*{\sum_{i=1}^n\w_i\cD_i , \cD_{\text{target}}} = 0$ in terms so-called representation-based distance [Chen et al., 2021], which is also zero in our case, or existence of identical neighbors. However, the analysis is based on BLO's techniques and requires a hypergradient estimation, i.e., $\nabla_w f(x^*(w), w)$, which is usually hard to compute. To avoid the hypergradient calculation, Chen et al., 2021 also propose a heuristic based on the usage of cosine similarity between the clients' gradients, which makes the implementation of the algorithm similar to FedAdp [Wu & Wang, 2021].

In fact, there are two major difficulties in estimating hypergradient. The first one is that the optimal solution $\x^*(\w)$ of the lower problem for every given $\w$ needs to be estimated. The known approaches iteratively update the lower variable $\x$ multiple times before updating $\w$, which causes high communication costs in a distributed setup. A lot of methods [Ghadimi & Wang, 2018, Hong et al., 2020, Chen et al., 2021, Ji et al., 2021, Ji et al., 2022] are proposed to effectively estimate $\x^*(\w)$ before updating $\w$, but anyway the less precise estimate slowdowns the convergence. The second obstacle is that hypergradient calculation requires second-order
derivatives of $\f_i(\w,\x)$. Many existing methods [Chen et al., 2022, Dagréou et al., 2022] use an explicit second-order derivation of $\f_i(\w,\x)$ with a major focus on efficiently estimating its Jacobian and inverse Hessian, which is computationally expensive itself, but also dramatically increases the communication cost in a distributed setup.
A number of methods [Chen et al., 2022, Li et al., 2022, Dagréou et al., 2022] avoid directly estimating its second-order computation and only use the first-order information of both upper and lower objectives, but they still have high communication costs and do not exploit our assumptions.

For a more detailed review of BLO, we refer to [Zhang et al., 2023, Liu et al., 2021, Chen et al., 2022].

§ : MERIT-BASED FEDERATED LEARNING FOR DIVERSE DATASETS

Recall that the primary objective the target client seeks to solve is given by (<ref>) where $n$ workers are connected with a parameter-server. Standard Parallel SGD

\begin{equation}
\textstyle  x^{t+1} = x^t - \frac{\gamma}{n}\sum_{i=1}^n g_i(x^t, \xiv_i), \label{eq:SGD}
\end{equation}

where $g_i(x^t, \xiv_i)$ denotes a stochastic gradient (unbiased estimate of $\nabla f_i(x^t)$) received from client $i$, cannot solve problem (<ref>) in general, since workers $\{2,\ldots, n\}$ do not necessarily have the same data distribution as the target client. This issue can be solved if we modify the method as follows:

\begin{equation}
\textstyle  x^{t+1} = x^t - \frac{\gamma}{|\cG|}\sum_{i\in \cG} g_i(x^t, \xiv_i), \label{eq:SGD_ideal}
\end{equation}

where $\cG$ denotes the set of workers that have the same data distribution as the target worker. However, the group $\cG$ is not known in advance. This aspect makes the method from (<ref>) impractical. Moreover, this method ignores potentially useful vectors received from the workers having different yet similar data distributions.

: Merit-based Federated Learning for Diverse Datasets
[1]
Input: Starting point $\x^0 \in \R^d$, stepsize $\gamma > 0$
server sends $\x^t$ to each worker
workers $i= 1,\dots,\n$ in parallel
compute stoch. gradient
$\g_i(\x^t, \xiv_i)$  from local data
send $\g_i(\x^t, \xiv_i)$ to the server

$\w^{t+1} \approx \argmin\limits_{w\in \Delta_1^n}f\left(x^t - \gamma \sum\limits_{i=1}^n \w_i  \g_i(\x^t, \xiv_i)\right)$

$\x^{t+1} = \x^t - \gamma \sum\limits_{i=1}^n \w^{t+1}_i  \g_i(\x^t, \xiv_i)$.

§.§ The Proposed Method

We develop Merit-based Federated Learning for Diverse Datasets (; see Algorithm <ref>) aimed at solving (<ref>) and safely gathering all potential benefits from collaboration with other clients. As in Parallel SGD all clients are required to send the stochastic gradients to the server. However, in contrast to uniform averaging of the received stochastic gradients, uses the weights $w^t$ from the unit simplex $\Delta_1^n$ that are updated at each iteration. In particular, the new vector of weights $w^{t+1} \in \R^n$ at iteration $t$ approximates $\argmin_{w\in \Delta_1^n}f(x^t - \gamma \sum_{i=1}^n \w_i  \g_i(\x^t, \xiv_i))$. Then, the server uses the obtained weights for averaging the stochastic gradients and updating $x^t$.

§.§ Auxiliary Problem in Line <ref>

From the theoretical perspective, solving the problem in Line <ref> is, in general, not easier than solving the original problem (<ref>). Therefore, we present two particular approaches for efficiently addressing this problem.

Approach 1: use fresh data. Let us assume that the target client can obtain new samples from distribution $\cD_1$ at any moment in time. To avoid any risk of compromising clients' privacy, the dataset of the target client should be stored on the target client only, and stochastic gradients received from other clients cannot be directly sent to the target client. To satisfy these requirements, one can approximate

\begin{equation}\textstyle
\argmin_{\w \in \Delta_1^n} \left\{\varphi_t(w) \equiv f \left(x^t - \gamma \sum_{i=1}^n \w_i  \g_i(\x^t, \xiv_i)\right)\right\}
\end{equation}

using zeroth-order[In this case, the server can ask the target client to evaluate loss values at the required points without sending the stochastic gradients received from other workers.] Mirror Descent (or its accelerated version) [Duchi et al., 2015, Shamir, 2017, Gasnikov et al., 2022]:

\begin{equation}\textstyle
w^{k+1} = \argmin_{w \in \Delta_1^n}\left\{\alpha\langle \tilde{g}^k, w\rangle + D_r(w, w^k)\right\}, \label{eq:zo_MD}
\end{equation}

where $\alpha > 0$ is the stepsize, $\tilde{g}^k$ is a finite-difference approximation of the directional derivative of sampled function

\begin{equation}\textstyle
\varphi_{t,\xi^k}(w) \eqdef f_{\xi^k} \left(x^t - \gamma \sum_{i=1}^n \w_i  \g_i(\x^t, \xiv_i)\right),
\end{equation}

where $\xi^k$ is a fresh sample from the distribution $\cD_1$ independent from all previous steps of the method, e.g., one can use

\begin{equation*}\textstyle
\tilde{g}^k = \frac{n(\varphi_{t,\xi^k}(w^k + he) - \varphi_{t,\xi^k}(w^k - he))}{2h}
\end{equation*}

for $h > 0$ and $e$ being sampled from the uniform distribution on the unit Euclidean sphere, and $D_r(w,w^k) = r(w) - r(w^k) - \langle \nabla r(w^k), w - w^k \rangle$ is the Bregman divergence associated with a $1$-strongly convex function $r$. Although, typically, the oracle complexity bounds for gradient-free methods have $\cO(n)$ dependence on the problem dimension [Gasnikov et al., 2022], one can get just $\cO(\log^2(n))$, in the case of the optimization over the probability simplex [Shamir, 2017, Gasnikov et al., 2022]. More precisely, if $f$ is $M_2$-Lipschitz w.r.t. $\ell_2$-norm and convex, then one can achieve $\E[\varphi_t(w) - \varphi_t(w^*)] \leq \delta$ using $\cO(\nicefrac{M_2^2\log^2(n)}{\delta^2})$ computations of $\varphi$, where $R$ is $\ell_1$-distance between the starting point and the solution [Gasnikov et al., 2022] and prox-function is chosen as $r(w) = \sum_{i=1}^n w_i\log(w_i)$, which is $1$-strongly convex w.r.t. $\ell_1$-norm.

Approach 2: use additional validation data. Alternatively, one can assume that the target client has an additional validation dataset $\widehat{D}$ sampled from $\cD_1$. Then, instead of function $f$ in Line <ref>, one can approximately minimize

\begin{equation}\textstyle
\widehat f(x) = \frac{1}{|\widehat D|}\sum_{\xi \in \widehat D} f_{\xi}(x), \label{eq:validation_loss}
\end{equation}

which under certain conditions provably approximates the original function $f(x)$ with any predefined accuracy if the dataset $\widehat D$ is sufficiently large [Shalev-Shwartz et al., 2009, Feldman & Vondrak, 2019].
We consider this case experimentally.

§ CONVERGENCE ANALYSIS

Mean Estimation: $\mu = 0.001$, MD learning rate = 3.5.

Mean Estimation: $\mu = 0.01$, MD learning rate = 4.5.

Mean Estimation: $\mu = 0.1$, MD learning rate = 12.5.

CIFAR10: $\alpha = 0.5$.

CIFAR10: $\alpha = 0.7$.

CIFAR10: $\alpha = 0.9$.

CIFAR10: $\alpha = 0.99$.

CIFAR10 (extra validation): $\alpha = 0.5$

CIFAR10 (extra validation): $\alpha = 0.7$

CIFAR10 (extra validation): $\alpha = 0.9$

CIFAR10 (extra validation): $\alpha = 0.99$.

GoEmotions: $\alpha = 0.5$

GoEmotions: $\alpha = 0.7$

GoEmotions: $\alpha = 0.9$

GoEmotions: $\alpha = 0.99$

GoEmotions (extra validation): $\alpha = 0.5$

GoEmotions (extra validation): $\alpha = 0.7$

GoEmotions (extra validation): $\alpha = 0.9$

GoEmotions (extra validation): $\alpha = 0.99$

In this section, we provide a convergence analysis of under the following standard assumptions.

§.§ Assumptions

For all $i\in \cG$ the stochastic gradient $\g_i(\x, \xiv_i)$ is an unbiased estimator of $\nabla\f_i(\x)$ with bounded variance, i.e., $\E_{\xiv_i} [\g_i(\x, \xiv_i)] = \nabla\f_i(\x)$ and for some $\sigma \geq 0$

\begin{equation}  \label{eq:xi-var}
\textstyle
\E_{\xiv_i} \left[\norm*{ \g_i(\x, \xiv_i) -  \nabla\f_i(\x)}^2\right] \leq \sigma^2.
\end{equation}

The above assumption is known as the bounded variance assumption. It is classical for the analysis of stochastic optimization methods, e.g., see [Nemirovski et al., 2009, Juditsky et al., 2011].

Next, we assume the smoothness of the objective.

$\f$ is $L$-smooth, i.e., $\forall\; \x,\y \in \R^d$

\begin{equation}\label{eq:lipschitzness}\textstyle
\f\rbr{\x} \le \f\rbr{\y} + \inp*{\nabla \f\rbr{\y}}{\x - \y} + \frac{L}{2}\norm*{\x - \y}^2. \tag{Lip}
\end{equation}

We also make the following (optional) assumption called Polyak-Łojasiewicz (PŁ) condition [Polyak, 1963, Lojasiewicz, 1963].

$\f$ satisfies Polyak-Łojasiewicz (PŁ) condition with parameter $\mu$, i.e., for $\mu \ge 0$

\begin{equation} \label{eq:pl}\textstyle
\f^* \ge \f\rbr*{\x} - \frac{1}{2\mu}\norm*{\nabla \f(\x)}^{2},\quad \forall\;\x \in \R^d. \tag{PL}
\end{equation}

This assumption belongs to the class of structured non-convexity assumptions allowing linear convergence for first-order methods such as Gradient Descent [Necoara et al., 2019].

§.§ Main Result

The main result for is given below.

Let Assumptions <ref> and <ref> hold. Then after $T$ iterations, with $\gamma \leq \frac{1}{2L}$ outputs $\x^{i}$, $i=0,\cdots, T-1$ such that

\begin{equation*}\textstyle
\frac{1}{T}\sum_{t=0}^{T-1}\E \norm*{\nabla\f\rbr*{\x^{t}}}^{2}
\le \frac{2\rbr*{\f\rbr*{\x^{0}} - \f\rbr*{\x^{*}}}}{T\gamma} + \frac{2\sigma^2 \gamma L}{\gn} + \frac{2 \delta}{\gamma},
\end{equation*}

where $\delta$ is the accuracy of solving the problem in Line <ref> and $G = |\cG|$.
Moreover if Assumption <ref> additionally holds, then after $T$ iterations of with $\gamma \leq \frac{1}{2L}$ outputs $\x^{T}$ such that

\begin{equation*}\textstyle
\E\f\rbr*{\x^{T}} - \f^* \le \rbr*{1 - \gamma\mu}^T \rbr*{\f\rbr*{\x^{0}} - \f^*}  + \frac{\sigma^2 \gamma L}{\mu\gn} + \frac{\delta}{\gamma\mu}.
\end{equation*}

We notice that

\textstyle
\min_{w \in \Delta_1^n} \varphi_t(w) \leq f (x^t - \frac{\gamma}{|\cG|}\sum_{i\in \cG} g_i(x^t, \xiv_i) ).$$

This implies that

\begin{align*}\textstyle
\E\big[f(x^{t+1})\mid x^t, &\{g_i(x^t, \xiv_i)\}_{i\in[n]}\big] \\
\textstyle
\leq f(x^t - \frac{\gamma}{|\cG|}\sum _{i\in \cG} g_i(x^t, \xiv_i)) + \delta.
\end{align*}

The rest of the proof is similar to the standard analysis of SGD; see the complete proof in the appendix.

If $\delta$ is sufficiently small, then the above result matches the known results for Parallel SGD [Ghadimi & Lan, 2013, Karimi et al., 2016, Khaled & Richtárik, 2022] that uniformly averages only the workers from the group $\cG$, i.e., those workers that have data distribution $\cD_1$ (see the method in (<ref>)). More precisely, we see a linear speed-up of $\nicefrac{1}{G}$ in the obtained convergence rates. However, does not require the knowledge of the workers that have the same distribution. Moreover, as we see in our numerical experiments, can converge even better when there exist workers having different but close data distribution, and it is not necessary to solve the auxiliary problem in Line <ref> with high precision.

§ EXPERIMENTS

To ensure transparency and facilitate reproducibility, we have provided access to the source code for all of our experiments at <https://github.com/nazya/MeritFed>.

We evaluate the performance of proposed methods in comparison with FedAdp [Wu & Wang, 2021] and TAWT [Chen et al., 2021]. We also compare standard SGD with uniform weights (labeled as SGD Full), SGD that accumulates only gradients from clients with the target distribution (SGD Ideal) and two versions of our algorithm. The first one, labeled as MeritFed SMD, samples gradient for the Mirror Descent subroutine in contrast to the other one, labeled as MeritFed MD, that uses the full dataset (additional or train) to calculate gradient for Mirror Descent step.

We present evolution of weighs (if applicable) using heat-map plots.

§.§ Mean Estimation

First of all, we consider the mean estimation problem, i.e., finding such a vector that minimizes the mean squared distance to the data samples. More formally, the goal of the target client is to solve

\begin{equation}\textstyle
\min_{x\in \R^d} \E_{\xi \sim \cD_1}\|x - \xi\|^2, \notag
\end{equation}

that has the optimum at $x^* = \E_{\xi \sim \cD_1}[\xi]$. We consider $\cD_1 = \cN(0, \mI)$ and also two other distributions from where some clients also get samples: $\cD_2 = \cN(\mu\mathbf{1}, \mI)$ and $\cD_3 = \cN(e, \mI)$, where $\mathbf{1} = (1,1,\ldots,1)^\top \in \R^d$, $\mu > 0$ is a parameter, and $e$ is some vector that we obtain in advance via sampling uniformly at random from the unit Euclidean sphere. We consider $150$ clients with data distributed as follows: the first $5$ workers have data from $\cD_1$ (the first group of clients), the next $95$ workers have data from $\cD_2$ (the second group of clients), and the remaining $50$ clients have data from $\cD_3$ (the third group of clients). Each client has $1000$ samples from the corresponding distribution, and the target client has additional $1000$ samples for validation, i.e., for solving the problem in Line <ref>. The dimension of the problem is $d = 10$.  Parameters that are the same for all experiments: number of peers $= 150$, number of samples $= 1000$, batch size $= 100$, learning rate $= 0.01$, number of steps for Mirror Descent $= 50$.

We consider three cases: $\mu = 0.001$, $0.01$, $0.1$. The smaller $\mu$ is, the closer $\cD_2$ is to $\cD_1$ and, thus, the more beneficial the samples from the second group are. Therefore, for small $\mu$, we expect to see that outperforms SGD Ideal. Moreover, since the workers from the third group have quite different data distribution, SGD Full is expected to work worse than other baselines.
The results are presented in Figures <ref>,  <ref> and <ref>.

§.§ Image Classification: CIFAR10 + ResNet18

This part is devoted to image classification on the CIFAR10 [Krizhevsky et al., 2009] dataset using ResNet18 [He et al., 2016] model and cross-entropy loss. Specifically, the target client's objective is to classify the first three classes: 0, 1 and 2. This client possesses data with these three labels.
The following ten workers also have datasets where a proportion, denoted by $\alpha \in (0,1]$, consists of classes from the set ${0,1,2}$, while the remaining $1-\alpha$ portion includes classes from the set ${3,4,5}$.
The remaining clients have data from the rest, e.g. ${6,7,8,9}$ labeled. The data is randomly distributed among clients without overlaps, adhering to the aforementioned label restrictions.
We use the same methods with the following parameters: a batch-size of 75, a learning rate of 0.01, 10 steps for Mirror Descent; and a learning rate of 0.1 and a batch-size of 90 for Mirror Descent.
We normalize images (similarly to [Horváth & Richtárik, 2020]).
Since an additional validation dataset can be used by  , we cut 300 samples of each target class (${0, 1, 2}$) off from the test data. So that accuracy and loss are calculated on the rest of test data and includes includes only labels ${0, 1, 2}$, modeling the case, when the target client is aimed at only classifying its own labels, e.g. ${0, 1, 2}$.

The results without additional validation dataset are presented in Figures <ref>-<ref> (evolution of weights presented in Figures <ref>-<ref>) and with in Figures <ref>-<ref>.
The plots are averaged over 7 runs with different seeds. Additionally accuracy plots present a standard variation area.

§.§ Texts Classification: GoEmotions + BERT

The next problem we consider is devoted to fine-tuning pretrained BERT [Devlin et al., 2018]  model for emotions classification on the GoEmotions dataset [Demszky et al., 2020].
The dataset consist of texts labeled with one or more of 28 emotions.
First of all, we form "truncated dataset" by cutting the dataset so that its each entry has the only label.

Then we use Ekman mapping [Ekman, 1992] to split the data between clients. According to the mapping 28 emotions can be mapped to 7 basic emotions. So that, we simulate a situation when the target client aimed at classification only emotions belonging to the only basic emotion, e.g. the target client has only emotions belonging to "joy" class and namely includes only "joy", "amusement", "approval", "excitement", "gratitude",  "love", "optimism", "relief", "pride", "admiration", "desire", "caring". The distribution of these sub-emotions is kept to be the same as the distribution of the truncated train dataset.

Clients, that data are suppose to have similar distribution, also has texts from base class "joy" and are labeled as one of the sub-emotion belonging to "joy".  The distribution of sub-emotions is also the same as the distribution of the truncated train dataset. These texts constitute an $\alpha$ portion of the total client's data. The other $1-\alpha$ portion of the texts is taken from "neutral" class.

The rest of clients are supposed to have different distribution and their data consist of either texts belonging to one of the other basic emotion, either mixed with neutral (if there is not enough texts to have a desired number of samples) or texts from "neutral" class only. Again, the distribution of sub-emotions is the same as the distribution of the truncated train dataset.

We use the same methods with the following parameters: a batch size of 40, a learning rate of 0.01, 10 steps for Mirror Descent; a batch-size and a learning rate of 0.1 for Mirror Descent.
The plots are averaged over 3 runs with different seeds. Additionally, accuracy plots present a standard variation area.

As well as for previous experiments, the target client benefits from collaborating with clients from the second group and achieves better accuracy using . The results without additional validation dataset are presented in Figures <ref>-<ref> and with in Figures <ref>-<ref>.
Evolution of weights presented in Figures <ref>-<ref>.

§.§ Closer look at experimental results

Experimental results fit the described intuition and our theory quite well:

the workers from the second group are beneficial

(since their distributions are close enough to the distribution of the target client). Indeed, achieves better optimization error (due to the smaller variance because of the averaging with more workers), and weights for aggregation are

much greater over the first two groups.
However, when

the dissimilarity between distributions is large
the second group becomes less useful for the training, and has comparable performance to SGD Ideal.

FedAdp is biased to the target client and assigns almost identical weights to either clients with similar either dissimilar distributions, that results in accuracy decrease at the end of the training, in contrast to which tracks and maintain less weights to non-beneficial clients.

TAWT is much more biased to the target client, that makes its almost identical to SGD Ideal.

It worth noting that seeks to assign greater weights to clients with similar distributions. Moreover, one can observe, that is able to effectively exploit data heterogeneity at initial iteration when it is beneficial.

§ CONCLUSION

In this paper, we introduced a novel algorithm called Merit-based Federated Learning () to address the challenges posed by the heterogeneous data distributions in federated learning (FL). We demonstrated that can effectively harness the advantages of distinct data distributions, control the detrimental effects of outlier clients, and promote collaborative learning. Our approach assigns adaptive aggregation weights to clients participating in FL training, allowing for faster convergence and potentially better generalization. The key contributions of this paper include the development of , provable convergence under mild assumptions, and the ability to utilize benefits from collaborating with clients having different but similar data distributions.

However, it is important to acknowledge certain limitations in our current work. Firstly, relies on the availability of validation data on the target client to solve the auxiliary problem. In practice, collecting and maintaining such validation data might not always be feasible or efficient. Secondly, the experiments in this paper were conducted with a limited number of clients and a moderate dataset size. Extending to large-scale FL scenarios with a substantial number of clients and massive datasets may pose scalability challenges. Addressing these limitations is part of our plan for future work.

Furthermore,  serves as a foundation for numerous extensions and enhancements. Future research can explore topics such as acceleration techniques, adaptive or scaled optimization methods (e.g., variants akin to Adam) on the server side, communication compression strategies, and the efficient implementation of similar collaborative learning approaches for all clients simultaneously. These directions will contribute to the continued development of federated learning methods, making them more efficient, robust, and applicable to a wide range of practical scenarios.

[Alistarh et al., 2017]
Alistarh, D., Grubic, D., Li, J., Tomioka, R., and Vojnovic, M.
Qsgd: Communication-efficient sgd via gradient quantization and encoding.
Advances in neural information processing systems, 30, 2017.

[Chen et al., 2022]
Chen, C., Chen, X., Ma, C., Liu, Z., and Liu, X.
Gradient-based bi-level optimization for deep learning: A survey.
arXiv preprint arXiv:2207.11719, 2022a.

[Chen et al., 2021]
Chen, S., Crammer, K., He, H., Roth, D., and Su, W. J.
Weighted training for cross-task learning.
arXiv preprint arXiv:2105.14095, 2021a.

[Chen et al., 2021]
Chen, T., Sun, Y., and Yin, W.
Closing the gap: Tighter analysis of alternating stochastic gradient methods for bilevel problems.
Advances in Neural Information Processing Systems, 34:0 25294–25307, 2021b.

[Chen et al., 2022]
Chen, T., Sun, Y., Xiao, Q., and Yin, W.
A single-timescale method for stochastic bilevel optimization.
In International Conference on Artificial Intelligence and Statistics, pp.  2466–2488. PMLR, 2022b.

[Chen et al., 2022]
Chen, W., Horváth, S., and Richtárik, P.
Optimal client sampling for federated learning.
Transactions on Machine Learning Research, 2022c.
ISSN 2835-8856.
URL <https://openreview.net/forum?id=8GvRCWKHIL>.

[Cho et al., 2020]
Cho, Y. J., Wang, J., and Joshi, G.
Client selection in federated learning: Convergence analysis and power-of-choice selection strategies.
arXiv preprint arXiv:2010.01243, 2020.

[Collins et al., 2021]
Collins, L., Hassani, H., Mokhtari, A., and Shakkottai, S.
Exploiting shared representations for personalized federated learning.
In International conference on machine learning, pp.  2089–2099. PMLR, 2021.

[Dagréou et al., 2022]
Dagréou, M., Ablin, P., Vaiter, S., and Moreau, T.
A framework for bilevel optimization that enables stochastic and global variance reduction algorithms.
arXiv preprint arXiv:2201.13409, 2022.

[Demszky et al., 2020]
Demszky, D., Movshovitz-Attias, D., Ko, J., Cowen, A., Nemade, G., and Ravi, S.
GoEmotions: A Dataset of Fine-Grained Emotions.
In 58th Annual Meeting of the Association for Computational Linguistics (ACL), 2020.

[Devlin et al., 2018]
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K.
BERT: Pre-training of deep bidirectional transformers for language understanding.
arXiv preprint arXiv:1810.04805, 2018.

[Ding & Wang, 2022]
Ding, S. and Wang, W.
Collaborative learning by detecting collaboration partners.
Advances in Neural Information Processing Systems, 35:0 15629–15641, 2022.

[Duchi et al., 2015]
Duchi, J. C., Jordan, M. I., Wainwright, M. J., and Wibisono, A.
Optimal rates for zero-order convex optimization: The power of two function evaluations.
IEEE Transactions on Information Theory, 610 (5):0 2788–2806, 2015.

[Ekman, 1992]
Ekman, P.
An argument for basic emotions.
Cognition & emotion, 60 (3-4):0 169–200, 1992.

[Even et al., 2022]
Even, M., Massoulié, L., and Scaman, K.
On sample optimality in personalized collaborative and federated learning.
Advances in Neural Information Processing Systems, 35:0 212–225, 2022.

[Fallah et al., 2020]
Fallah, A., Mokhtari, A., and Ozdaglar, A.
Personalized federated learning: A meta-learning approach.
arXiv preprint arXiv:2002.07948, 2020.

[Feldman & Vondrak, 2019]
Feldman, V. and Vondrak, J.
High probability generalization bounds for uniformly stable algorithms with nearly optimal rate.
In Conference on Learning Theory, pp.  1270–1279. PMLR, 2019.

[Fraboni et al., 2021]
Fraboni, Y., Vidal, R., Kameni, L., and Lorenzi, M.
Clustered sampling: Low-variance and improved representativity for clients selection in federated learning.
In International Conference on Machine Learning, pp.  3407–3416. PMLR, 2021.

[Gasnikov et al., 2022]
Gasnikov, A., Dvinskikh, D., Dvurechensky, P., Gorbunov, E., Beznosikov, A., and Lobanov, A.
Randomized gradient-free methods in convex optimization.
arXiv preprint arXiv:2211.13566, 2022a.

[Gasnikov et al., 2022]
Gasnikov, A., Novitskii, A., Novitskii, V., Abdukhakimov, F., Kamzolov, D., Beznosikov, A., Takac, M., Dvurechensky, P., and Gu, B.
The power of first-order smooth optimization for black-box non-smooth problems.
In International Conference on Machine Learning, pp.  7241–7265. PMLR, 2022b.

[Ghadimi & Lan, 2013]
Ghadimi, S. and Lan, G.
Stochastic first-and zeroth-order methods for nonconvex stochastic programming.
SIAM Journal on Optimization, 230 (4):0 2341–2368, 2013.

[Ghadimi & Wang, 2018]
Ghadimi, S. and Wang, M.
Approximation methods for bilevel programming.
arXiv preprint arXiv:1802.02246, 2018.

[Ghosh et al., 2020]
Ghosh, A., Chung, J., Yin, D., and Ramchandran, K.
An efficient framework for clustered federated learning.
Advances in Neural Information Processing Systems, 33:0 19586–19597, 2020.

[Hanzely et al., 2020]
Hanzely, F., Hanzely, S., Horváth, S., and Richtárik, P.
Lower bounds and optimal algorithms for personalized federated learning.
Advances in Neural Information Processing Systems, 33:0 2304–2315, 2020.

[He et al., 2016]
He, K., Zhang, X., Ren, S., and Sun, J.
Deep residual learning for image recognition.
In Proceedings of the IEEE conference on computer vision and pattern recognition, pp.  770–778, 2016.

[Hong et al., 2020]
Hong, M., Wai, H.-T., Wang, Z., and Yang, Z.
A two-timescale framework for bilevel optimization: Complexity analysis and application to actor-critic.
arXiv preprint arXiv:2007.05170, 2020.

[Horváth & Richtárik, 2020]
Horváth, S. and Richtárik, P.
A better alternative to error feedback for communication-efficient distributed learning.
arXiv preprint arXiv:2006.11077, 2020.

[Huang et al., 2022]
Huang, Y., Lin, Q., Street, N., and Baek, S.
Federated learning on adaptively weighted nodes by bilevel optimization.
arXiv preprint arXiv:2207.10751, 2022.

[Ji et al., 2021]
Ji, K., Yang, J., and Liang, Y.
Bilevel optimization: Convergence analysis and enhanced design.
In International conference on machine learning, pp.  4882–4892. PMLR, 2021.

[Ji et al., 2022]
Ji, K., Liu, M., Liang, Y., and Ying, L.
Will bilevel optimizers benefit from loops.
arXiv preprint arXiv:2205.14224, 2022.

[Juditsky et al., 2011]
Juditsky, A., Nemirovski, A., and Tauvel, C.
Solving variational inequalities with stochastic mirror-prox algorithm.
Stochastic Systems, 10 (1):0 17–58, 2011.

[Kairouz et al., 2021]
Kairouz, P., McMahan, H. B., Avent, B., Bellet, A., Bennis, M., Bhagoji, A. N., Bonawitz, K., Charles, Z., Cormode, G., Cummings, R., et al.
Advances and open problems in federated learning.
Foundations and Trends® in Machine Learning, 140 (1–2):0 1–210, 2021.

[Karimi et al., 2016]
Karimi, H., Nutini, J., and Schmidt, M.
Linear convergence of gradient and proximal-gradient methods under the polyak-łojasiewicz condition.
In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2016, Riva del Garda, Italy, September 19-23, 2016, Proceedings, Part I 16, pp.  795–811. Springer, 2016.

[Karimireddy et al., 2020]
Karimireddy, S. P., Kale, S., Mohri, M., Reddi, S., Stich, S., and Suresh, A. T.
Scaffold: Stochastic controlled averaging for federated learning.
In International conference on machine learning, pp.  5132–5143. PMLR, 2020.

[Khaled & Richtárik, 2022]
Khaled, A. and Richtárik, P.
Better theory for sgd in the nonconvex world.
Transactions on Machine Learning Research, 2022.

[Khaled et al., 2020]
Khaled, A., Mishchenko, K., and Richtárik, P.
Tighter theory for local sgd on identical and heterogeneous data.
In International Conference on Artificial Intelligence and Statistics, pp.  4519–4529. PMLR, 2020.
URL <http://proceedings.mlr.press/v108/bayoumi20a/bayoumi20a-supp.pdf>.

[Konečný et al., 2016]
Konečný, J., McMahan, H. B., Ramage, D., and Richtárik, P.
Federated optimization: Distributed machine learning for on-device intelligence.
arXiv preprint arXiv:1610.02527, 2016.

[Krizhevsky et al., 2009]
Krizhevsky, A., Hinton, G., et al.
Learning multiple layers of features from tiny images.

[Kulkarni et al., 2020]
Kulkarni, V., Kulkarni, M., and Pant, A.
Survey of personalization techniques for federated learning.
In 2020 Fourth World Conference on Smart Trends in Systems, Security and Sustainability (WorldS4), pp.  794–797. IEEE, 2020.

[Lai et al., 2021]
Lai, F., Zhu, X., Madhyastha, H. V., and Chowdhury, M.
Oort: Efficient federated learning via guided participant selection.
In 15th $\{$USENIX$\}$ Symposium on Operating Systems Design and Implementation ($\{$OSDI$\}$ 21), pp.  19–35, 2021.

[Lamport et al., 1982]
Lamport, L., Shostak, R., and Pease, M.
The byzantine generals problem.
ACM Transactions on Programming Languages and Systems, 40 (3):0 382–401, 1982.

[Li et al., 2022]
Li, J., Gu, B., and Huang, H.
A fully single loop algorithm for bilevel optimization without hessian inverse.
In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp.  7426–7434, 2022.

[Li et al., 2020]
Li, T., Sahu, A. K., Talwalkar, A., and Smith, V.
Federated learning: Challenges, methods, and future directions.
IEEE signal processing magazine, 370 (3):0 50–60, 2020.

[Lian et al., 2017]
Lian, X., Zhang, C., Zhang, H., Hsieh, C.-J., Zhang, W., and Liu, J.
Can decentralized algorithms outperform centralized algorithms? a case study for decentralized parallel stochastic gradient descent.
Advances in Neural Information Processing Systems, 30, 2017.
URL <https://arxiv.org/pdf/1705.09056.pdf>.

[Liu et al., 2021]
Liu, R., Gao, J., Zhang, J., Meng, D., and Lin, Z.
Investigating bi-level optimization for learning and vision from a unified perspective: A survey and beyond.
IEEE Transactions on Pattern Analysis and Machine Intelligence, 440 (12):0 10045–10067, 2021.

[Lojasiewicz, 1963]
Lojasiewicz, S.
A topological property of real analytic subsets.
Coll. du CNRS, Les équations aux dérivées partielles, 1170 (87-89):0 2, 1963.

[Luo et al., 2022]
Luo, B., Xiao, W., Wang, S., Huang, J., and Tassiulas, L.
Tackling system and statistical heterogeneity for federated learning with adaptive client sampling.
In IEEE INFOCOM 2022-IEEE Conference on Computer Communications, pp.  1739–1748. IEEE, 2022.

[Lyu et al., 2020]
Lyu, L., Yu, H., Ma, X., Sun, L., Zhao, J., Yang, Q., and Yu, P. S.
Privacy and robustness in federated learning: Attacks and defenses.
arXiv preprint arXiv:2012.06337, 2020.

[Ma et al., 2022]
Ma, J., Long, G., Zhou, T., Jiang, J., and Zhang, C.
On the convergence of clustered federated learning.
arXiv preprint arXiv:2202.06187, 2022.

[Marfoq et al., 2021]
Marfoq, O., Neglia, G., Bellet, A., Kameni, L., and Vidal, R.
Federated multi-task learning under a mixture of distributions.
Advances in Neural Information Processing Systems, 34:0 15434–15447, 2021.

[McMahan et al., 2017]
McMahan, B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. A.
Communication-efficient learning of deep networks from decentralized data.
In Artificial intelligence and statistics, pp.  1273–1282. PMLR, 2017.

[Mishchenko et al., 2019]
Mishchenko, K., Gorbunov, E., Takáč, M., and Richtárik, P.
Distributed learning with compressed gradient differences.
arXiv preprint arXiv:1901.09269, 2019.

[Mishchenko et al., 2022]
Mishchenko, K., Malinovsky, G., Stich, S., and Richtárik, P.
Proxskip: Yes! local gradient steps provably lead to communication acceleration! finally!
In International Conference on Machine Learning, pp.  15750–15769. PMLR, 2022.

[Necoara et al., 2019]
Necoara, I., Nesterov, Y., and Glineur, F.
Linear convergence of first order methods for non-strongly convex optimization.
Mathematical Programming, 175:0 69–107, 2019.

[Nemirovski et al., 2009]
Nemirovski, A., Juditsky, A., Lan, G., and Shapiro, A.
Robust stochastic approximation approach to stochastic programming.
SIAM Journal on optimization, 190 (4):0 1574–1609, 2009.

[Nguyen et al., 2020]
Nguyen, H. T., Sehwag, V., Hosseinalipour, S., Brinton, C. G., Chiang, M., and Poor, H. V.
Fast-convergent federated learning.
IEEE Journal on Selected Areas in Communications, 390 (1):0 201–218, 2020.

[Polyak, 1963]
Polyak, B. T.
Gradient methods for the minimisation of functionals.
USSR Computational Mathematics and Mathematical Physics, 30 (4):0 864–878, 1963.

[Ribero & Vikalo, 2020]
Ribero, M. and Vikalo, H.
Communication-efficient federated learning via optimal client sampling.
arXiv preprint arXiv:2007.15197, 2020.

[Sadiev et al., 2022]
Sadiev, A., Borodich, E., Beznosikov, A., Dvinskikh, D., Chezhegov, S., Tappenden, R., Takáč, M., and Gasnikov, A.
Decentralized personalized federated learning: Lower bounds and optimal algorithm for all personalization modes.
EURO Journal on Computational Optimization, 10:0 100041, 2022.

[Shalev-Shwartz et al., 2009]
Shalev-Shwartz, S., Shamir, O., Srebro, N., and Sridharan, K.
Stochastic convex optimization.
In COLT, volume 2, pp.  5, 2009.

[Shamir, 2017]
Shamir, O.
An optimal algorithm for bandit and zero-order convex optimization with two-point feedback.
The Journal of Machine Learning Research, 180 (1):0 1703–1713, 2017.

[Shokri & Shmatikov, 2015]
Shokri, R. and Shmatikov, V.
Privacy-preserving deep learning.
In Proceedings of the 22nd ACM SIGSAC conference on computer and communications security, pp.  1310–1321, 2015.

[Song et al., 2022]
Song, Z., Li, W., Jin, K., Shi, L., Yan, M., Yin, W., and Yuan, K.
Communication-efficient topologies for decentralized learning with $ o (1) $ consensus rate.
Advances in Neural Information Processing Systems, 35:0 1073–1085, 2022.

[Stich, 2018]
Stich, S. U.
Local sgd converges fast and communicates little.
In International Conference on Learning Representations, 2018.

[Stich et al., 2018]
Stich, S. U., Cordonnier, J.-B., and Jaggi, M.
Sparsified sgd with memory.
Advances in Neural Information Processing Systems, 31, 2018.

[Tang et al., 2021]
Tang, X., Guo, S., and Guo, J.
Personalized federated learning with contextualized generalization.
arXiv preprint arXiv:2106.13044, 2021.

[Wang et al., 2021]
Wang, J., Charles, Z., Xu, Z., Joshi, G., McMahan, H. B., Al-Shedivat, M., Andrew, G., Avestimehr, S., Daly, K., Data, D., et al.
A field guide to federated optimization.
arXiv preprint arXiv:2107.06917, 2021.

[Wu & Wang, 2021]
Wu, H. and Wang, P.
Fast-convergent federated learning with adaptive weighting.
IEEE Transactions on Cognitive Communications and Networking, 70 (4):0 1078–1088, 2021.

[Zhang et al., 2021]
Zhang, C., Xie, Y., Bai, H., Yu, B., Li, W., and Gao, Y.
A survey on federated learning.
Knowledge-Based Systems, 216:0 106775, 2021.

[Zhang et al., 2023]
Zhang, Y., Khanduri, P., Tsaknakis, I., Yao, Y., Hong, M., and Liu, S.
An introduction to bi-level optimization: Foundations and applications in signal processing and machine learning.
arXiv preprint arXiv:2308.00788, 2023.

§ PROOF OF THEOREM <REF>

Let Assumptions <ref> and <ref> hold. Then after $T$ iterations of with $\gamma \leq \frac{1}{2L}$ outputs $\x^{i}$, $i=0,\cdots, T-1$ such that

\begin{equation}
\frac{1}{T}\sum_{t=0}^{T-1}\E \norm*{\nabla\f\rbr*{\x^{t}}}^{2}
\le \frac{2\rbr*{\f\rbr*{\x^{0}} - \f\rbr*{\x^{*}}}}{T\gamma} + \frac{2\sigma^2 \gamma L}{\gn} + \frac{2 \delta}{\gamma}, \label{eq:non_cvx_result_appendix}
\end{equation}

where $\delta$ is the accuracy of solving the problem in Line <ref> and $G = |\cG|$.
Moreover if Assumption <ref> additionally holds, then after $T$ iterations of with $\gamma \leq \frac{1}{2L}$ outputs $\x^{T}$ such that

\begin{equation}
\E\f\rbr*{\x^{T}} - \f^* \le \rbr*{1 - \gamma\mu}^T \rbr*{\f\rbr*{\x^{0}} - \f^*}  + \frac{\sigma^2 \gamma L}{\mu\gn} + \frac{\delta T}{\gamma\mu}. \label{eq:PL_result_appendix}
\end{equation}

We write $\g_i^t$ or simply $\g_i$ instead of $\g_i(\x^t, \xiv_i^t)$ when there is no ambiguity.
Then, the update rule in can be written as

\begin{equation*}
\x^{t+1} = \x^t - \gamma \sumin \w^{t+1}_i  \g_{i}(\x^{t}),
\end{equation*}

where $\w^{t+1}$ is an approximate solution of

\begin{eqnarray*}
\min_{\w \Delta_1^n} \f\rbr*{\x^t - \gamma \sumin \w_i \g_{i}(\x^{t})}
\end{eqnarray*}

that satisfies

\begin{eqnarray*}
\E\sbr*{\f\rbr*{\x^{t+1}} | \x^t, \xiv^t} - \min_{\w} \f\rbr*{\x^t - \gamma \sumin \w_i \g_{i}(\x^{t})} \le \delta.
\end{eqnarray*}

By definition of the minimum, we have

\begin{eqnarray*}
\lefteqn{\min_{\w\in \Delta^\n_1} \f\rbr*{\x^t - \gamma \sumin \w_i  \g_{i}(\x^{t})} \le \f\rbr*{\x^t - \frac{\gamma}{\gn} \sum_{i \in \cG}  \g_{i}(\x^{t})}}
\\        &\oset{\eqref{eq:lipschitzness}}{\le}& \f\rbr*{\x^{t}} - \frac{\gamma}{\gn} \inp*{\nabla\f\rbr*{\x^{t}}}{\sum_{i \in \cG}  \g_{i}(\x^{t})} + \frac{L\gamma^2}{2}\norm*{\frac{1}{\gn}\sum_{i \in \cG}  \g_{i}(\x^{t})}^{2}
\\
&\le& \f\rbr*{\x^{t}} - \frac{\gamma}{\gn} \inp*{\nabla\f\rbr*{\x^{t}}}{\sum_{i \in \cG}  \g_{i}(\x^{t})} + \gamma^2 L\norm*{\nabla\f\rbr*{\x^{t}} - \frac{1}{\gn}\sum_{i \in \cG}  \g_{i}(\x^{t})}^{2} + \gamma^2 L\norm*{\nabla\f\rbr*{\x^{t}}}^{2}.
\end{eqnarray*}

The last two inequalities imply

\begin{eqnarray}
\lefteqn{\E\sbr*{\f\rbr*{\x^{t+1}} | \x^t, \xiv^t} } \notag
\\
&\le& \f\rbr*{\x^{t}}
- \frac{\gamma}{\gn}
\inp*{\nabla\f\rbr*{\x^{t}}}{\sum_{i \in \cG}  \g_{i}(\x^{t})}
+ \gamma^2 L
\norm*{
\nabla\f\rbr*{\x^{t}} - \frac{\sum_{i \in \cG}  \g_{i}(\x^{t})}{\gn}
}^{2} + \gamma^2 L
\norm*{\nabla\f\rbr*{\x^{t}}}^{2} + \delta.
\nonumber
\end{eqnarray}

Taking the full expectation we get

\begin{eqnarray}
\E[f(x^{t+1})] &\leq& \E[f(x^t)] - \gamma(1 - \gamma L)\E\left[\|\nabla f(x^t)\|^2\right] + \gamma^2L\E\left[\left\|\nabla\f\rbr*{\x^{t}} - \frac{\sum_{i \in \cG}  \g_{i}(\x^{t})}{\gn}\right\|^2\right] + \delta \notag\\
&\overset{\gamma \leq \frac{1}{2L}}{\leq}&\E[f(x^t)] - \frac{\gamma}{2}\E\left[\|\nabla f(x^t)\|^2\right] + \frac{\gamma^2L}{G^2}\sum\limits_{i\in \cG}\E\left[\left\|\nabla\f\rbr*{\x^{t}} - \g_{i}(\x^{t})\right\|^2\right] + \delta \notag\\
&\overset{\eqref{eq:xi-var}}{\leq}& \E[f(x^t)] - \frac{\gamma}{2}\E\left[\|\nabla f(x^t)\|^2\right] + \frac{\gamma^2 L\sigma^2}{G} + \delta. \label{eq:suff_decr}
\end{eqnarray}

The above is equivalent to

\begin{eqnarray*}
\frac{\gamma}{2}\E \norm*{\nabla\f\rbr*{\x^{t}}}^{2}
\le \E\f\rbr*{\x^{t}} -  \E\f\rbr*{\x^{t+1}} + \frac{\sigma^2 \gamma^2 L}{\gn} + \delta,
\end{eqnarray*}

which concludes the first part of the proof.

Next, summing the inequality for $t \in \{0,1,\dots,T-1\}$
leads to

\begin{eqnarray*}
\frac{1}{T}\sum_{t=0}^{T-1}\E \norm*{\nabla\f\rbr*{\x^{t}}}^{2}
% \\
&\le& \frac{2\rbr*{\f\rbr*{\x^{0}} -  \E\f\rbr*{\x^{T}}}}{T\gamma} + \frac{2\sigma^2 \gamma L }{\gn} + \frac{2\delta }{\gamma}
\\
&\le& 2\rbr*{\f\rbr*{\x^{0}} - \f\rbr*{\x^{*}}} + \frac{2\sigma^2 \gamma L }{\gn} + \frac{2\delta }{\gamma}.
\end{eqnarray*}

Combining (<ref>) with (<ref>) gives

\begin{eqnarray*}
\E[f(x^{t+1}) - f^*] &\leq& (1 - \gamma\mu)\E[f(x^t) - f^*] + \frac{\gamma^2 L\sigma^2}{G} + \delta.
\end{eqnarray*}

Unrolling the above recurrence, we obtain (<ref>).

§ MISSING PLOTS

In this section we provide plots of evolution of weights. The following plots partially duplicate plots from the main parts(except plots with evolution) and run with the same setup as corresponding plots in the main part.

CIFAR10: $\alpha = 0.5$

CIFAR10: $\alpha = 0.7$

CIFAR10: $\alpha = 0.9$

CIFAR10: $\alpha = 0.99$

GoEmotions: $\alpha = 0.5$

GoEmotions: $\alpha = 0.7$

GoEmotions: $\alpha = 0.9$

GoEmotions: $\alpha = 0.99$

GoEmotions (extra validation): $\alpha = 0.5$

GoEmotions (extra validation): $\alpha = 0.7$

GoEmotions (extra validation): $\alpha = 0.9$

GoEmotions (extra validation): $\alpha = 0.99$