# Explicit Personalization and Local Training: Double Communication
Acceleration in Federated Learning

Kai Yi Laurent Condat Peter Richtárik
xx
King Abdullah University of Science and Technology (KAUST)
Thuwal, Kingdom of Saudi Arabia

(May 18, 2023)

###### Abstract

Federated Learning is an evolving machine learning paradigm, in which multiple
clients perform computations based on their individual private data,
interspersed by communication with a remote server. A common strategy to
curtail communication costs is _Local Training_ , which consists in performing
multiple local stochastic gradient descent steps between successive
communication rounds. However, the conventional approach to local training
overlooks the practical necessity for client-specific _personalization_ , a
technique to tailor local models to individual needs. We introduce Scafflix, a
novel algorithm that efficiently integrates explicit personalization with
local training. This innovative approach benefits from these two techniques,
thereby achieving doubly accelerated communication, as we demonstrate both in
theory and practice.

###### Contents

  1. 1 Introduction
  2. 2 Related work
    1. 2.1 Local Training (LT) methods in Federated Learning (FL)
    2. 2.2 Personalization in FL
  3. 3 Proposed algorithm Scafflix and convergence analysis
  4. 4 Experiments
    1. 4.1 Prelude: Convex Logistic Regression
    2. 4.2 Neural Network Training: Datasets and Baselines for Evaluation
    3. 4.3 Analysis of Generalization with Limited Communication Rounds
    4. 4.4 Key Ablation Studies
  5. 5 Conclusion
  6. A Proposed i-Scaffnew algorithm
  7. B From i-Scaffnew to Scafflix
  8. C Proof of Corollary 2
  9. D Additional Experimental Results

##  1 Introduction

Due to privacy concerns and limited computing resources on edge devices,
centralized training with all data first gathered in a datacenter is often
impossible in many real-world applications of data science and artificial
intelligence. As a result, Federated Learning (FL) has gained increasing
interest as a framework that enables multiple clients to do local
computations, based on their personal data kept private, and to communicate
back and forth with a server. FL is classically formulated as an empirical
risk minimization problem of the form

$\min_{x\in\mathbb{R}^{d}}\left[f(x)\coloneqq\frac{1}{n}\sum_{i=1}^{n}f_{i}(x)\right],$
(ERM)

where $f_{i}$ is the local objective on client $i$, $n$ is the total number of
clients, $x$ is the global model.

Thus, the usual approach is to solve (ERM) and then to deploy the obtained
globally optimal model
$x^{\star}\coloneqq{\arg\min_{x\in\mathbb{R}^{d}}}\,f(x)$ to all clients. To
reduce communication costs between the server and the clients, the practice of
updating the local parameters multiple times before aggregation, known as
_Local Training_ (LT) (Povey et al., 2015; Moritz et al., 2016; McMahan et
al., 2017; Li et al., 2020b; Haddadpour and Mahdavi, 2019; Khaled et al.,
2019, 2020; Karimireddy et al., 2020; Gorbunov et al., 2020a; Mitra et al.,
2021), is widely used in FL. LT, in its most modern form, is a communication-
acceleration mechanism, as we detail in Section 2.1.

Meanwhile, there is a growing interest in providing _personalization_ to the
clients, by providing them more-or-less customized models tailored to their
individual needs and heterogeneous data, instead of the one-size-fits-all
model $x^{\star}$. We review existing approaches to personalization in Section
2.2. If personalization is pushed to the extreme, every client just uses its
private data to learn its own locally-optimal model

$x^{\star}_{i}\coloneqq{\arg\min_{x\in\mathbb{R}^{d}}}\,f_{i}(x)$

and no communication at all is needed. Thus, intuitively, more personalization
means less communication needed to reach a given accuracy. In other words,
personalization is a communication-acceleration mechanism, like LT.

Therefore, we raise the following question:

> Is it possible to achieve double communication acceleration in FL by jointly
> leveraging the acceleration potential of personalization and local training?

For this purpose, we first have to formulate personalized FL as an
optimization problem. A compelling interpretation of LT (Hanzely and
Richtárik, 2020) is that it amounts to solve an implicit personalization
objective of the form:

$\min_{x_{1},\ldots,x_{n}\in\mathbb{R}^{d}}\frac{1}{n}\sum_{i=1}^{n}f_{i}(x_{i})+\frac{\lambda}{2n}\sum_{i=1}^{n}{\left\lVert\bar{x}-x_{i}\right\rVert}^{2},$
(1)

where $x_{i}\in\mathbb{R}^{d}$ denotes the local model at client
$i\in[n]\coloneqq\\{1,\ldots,n\\}$,
$\bar{x}\coloneqq\frac{1}{n}\sum_{i=1}^{n}x_{i}$ is the average of these local
models, and $\lambda\geq 0$ is the implicit personalization parameter that
controls the amount of personalization. When $\lambda$ is small, the local
models tend to be trained locally. On the other hand, a larger $\lambda$ puts
more penalty on making the local models $x_{i}$ close to their mean $\bar{x}$,
or equivalently in making all models close to each other, by pushing towards
averaging over all clients. Thus, LT is not only compatible with
personalization, but can be actually used to implement it, though implicitly:
there is a unique parameter $\lambda$ in (1) and it is difficult evaluate the
amount of personalization for a given value of $\lambda$.

The more accurate model FLIX for personalized FL was proposed by Gasanov et
al. (2022). It consists for every client $i$ to first compute locally its
personally-optimal model $x^{\star}_{i}$, and then to solve the problem

$\min_{x\in\mathbb{R}^{d}}\,\tilde{f}(x)\coloneqq\frac{1}{n}\sum_{i=1}^{n}f_{i}\big{(}\alpha_{i}x+(1-\alpha_{i})x^{\star}_{i}\big{)},$
(FLIX)

where $\alpha_{i}\in[0,1]$ is the explicit and individual personalization
factor for client $i$. At the end, the personalized model used by client $i$
is the explicit mixture

$\tilde{x}_{i}^{\star}\coloneqq\alpha_{i}x^{\star}+(1-\alpha_{i})x^{\star}_{i},$

where $x^{\star}$ is the solution to (FLIX). A smaller value of $\alpha_{i}$
gives more weight to $x^{\star}_{i}$, which means more personalization. On the
other hand, if $\alpha_{i}=1$, the client $i$ uses the global model
$x^{\star}$ without personalization. Thus, if all $\alpha_{i}$ are equal to 1,
there is no personalization at all and (FLIX) reverts to (ERM). So, (FLIX) is
a more general formulation of FL than (ERM). The functions in (FLIX) inherit
smoothness and strong convexity from the $f_{i}$, so every algorithm
appropriate for (ERM) can also be applied to solve (FLIX). Gasanov et al.
(2022) proposed an algorithm also called FLIX to solve (FLIX), which is simply
vanilla distributed gradient descent (GD) applied to (FLIX).

In this paper, we first redesign and generalize the recently-proposed Scaffnew
algorithm (Mishchenko et al., 2022), which features LT and has an accelerated
communication complexity, and propose Individualized-Scaffnew (i-Scaffnew),
wherein the clients can have different properties. We then apply and tune
i-Scaffnew for the problem (FLIX) and propose our new algorithm for
personalized FL, which we call Scafflix. We answer positively to the question
above and prove that Scafflix enjoys a doubly accelerated communication
complexity, by jointly harnessing the acceleration potential of LT and
personalization. That is, its communication complexity depends on the square
root of the condition number of the functions $f_{i}$ and on the $\alpha_{i}$.
In addition to establishing the new state of the art for personalized FL with
our theoretical guarantees, we show by extensive experiments that Scafflix is
efficient in real-world learning setups and outperforms existing algorithms.

##  2 Related work

###  2.1 Local Training (LT) methods in Federated Learning (FL)

Theoretical evolutions of LT in FL have been long-lasting, spanning five
generations from empirical results to accelerated communication complexity.
The celebrated FedAvg algorithm proposed by McMahan et al. (2017) showed the
feasibility of communication-efficient learning from decentralized data. It
belongs to the first generation of LT methods, where the focus was on
empirical results and practical validations (Povey et al., 2015; Moritz et
al., 2016; McMahan et al., 2017).

The second generation of studies on LT for solving (ERM) was based on
homogeneity assumptions, such as bounded gradients ($\exists
c<+\infty,{\left\|\nabla f_{i}(x)\right\|}\leq c,x\in\mathbb{R}^{d}$,
$i\in[n]$) (Li et al., 2020b) and bounded gradient diversity
($\frac{1}{n}\sum_{i=1}^{n}{\left\lVert\nabla f_{i}(x)\right\rVert}^{2}\leq
c{\left\lVert\nabla f(x)\right\rVert}^{2}$) (Haddadpour and Mahdavi, 2019).
However, these assumptions are too restrictive and do not hold in practical FL
settings (Kairouz et al., 2019; Wang et al., 2021).

The third generation of approaches, under generic assumptions on the convexity
and smoothness of the functions, exhibited sublinear convergence (Khaled et
al., 2019, 2020) or linear convergence to a neighborhood (Malinovsky et al.,
2020).

Recently, popular algorithms have emerged, such as Scaffold (Karimireddy et
al., 2020), S-Local-GD (Gorbunov et al., 2020a), and FedLin (Mitra et al.,
2021), successfully correcting for the client drift and enjoying linear
convergence to an exact solution under standard assumptions. However, their
communication complexity remains the same as with GD, namely
$\mathcal{O}(\kappa\log\epsilon^{-1})$, where $\kappa\coloneqq L/\mu$ is the
condition number.

Finally, Scaffnew was proposed by Mishchenko et al. (2022), with accelerated
communication complexity $\mathcal{O}(\sqrt{\kappa}\log\epsilon^{-1})$. This
is a major achievement, which proves for the first time that LT is a
communication acceleration mechanism. Thus, Scaffnew is the first algorithm in
what can be considered the fifth generation of LT-based methods with
accelerated convergence. Subsequent works have further extended Scaffnew with
features such as variance-reduced stochastic gradients (Malinovsky et al.,
2022), compression (Condat et al., 2022), partial client participation (Condat
et al., 2023), asynchronous communication of different clients (Maranjyan et
al., 2022), and to a general primal–dual framework (Condat and Richtárik,
2023). The fifth generation of LT-based methods also includes the 5GCS
algorithm (Grudzień et al., 2023), based on a different approach: the local
steps correspond to an inner loop to compute a proximity operator inexactly.
Our proposed algorithm Scafflix generalizes Scaffnew and enjoys even better
accelerated communication complexity, thanks to a better dependence on the
possibly different condition numbers of the functions $f_{i}$.

###  2.2 Personalization in FL

We can distinguish three main approaches to achieve personalization:

a) One-stage training of a single global model using personalization
algorithms. One common scheme is to design a suitable regularizer to balance
between current and past local models (Li et al., 2021) or between global and
local models (Li et al., 2020a; Hanzely and Richtárik, 2020). The FLIX model
(Gasanov et al., 2022) achieves explicit personalization by balancing the
local and global model using interpolation. Meta-learning is also popular in
this thread, as evidenced by T Dinh et al. (2020), which proposes a federated
meta-learning framework that utilizes Moreau envelopes and a regularizer to
balance personalization and generalization.

b) Training a global model and fine-tuning every local client or knowledge
transfer/distillation. This approach allows knowledge transfer from a source
domain trained in the FL manner to target domains (Li and Wang, 2019), which
is especially useful for personalization in healthcare domains (Chen et al.,
2020; Yang et al., 2020).

c) Collaborative training between the global model and local models. The basic
idea behind this approach is that each local client trains some personalized
parts of a large model, such as the last few layers of a neural network.
Parameter decoupling enables learning of task-specific representations for
better personalization (Arivazhagan et al., 2019; Bui et al., 2019), while
channel sparsity encourages each local client to train the neural network with
sparsity based on their limited computation resources (Horvath et al., 2021;
Alam et al., 2022; Mei et al., 2022).

Despite the significant progress made in FL personalization, many approaches
only present empirical results. Our approach benefits from the simplicity and
efficiency of the FLIX framework and enjoys accelerated convergence.

##  3 Proposed algorithm Scafflix and convergence analysis

We generalize Scaffnew (Mishchenko et al., 2022) and propose Individualized-
Scaffnew (i-Scaffnew), shown as Algorithm 2 in the Appendix. Its novelty with
respect to Scaffnew is to make use of different stepsizes $\gamma_{i}$ for the
local SGD steps, in order to exploit the possibly different values of $L_{i}$
and $\mu_{i}$, as well as the different properties $A_{i}$ and $C_{i}$ of the
stochastic gradients. This change is not straightforward and requires to
rederive the whole proof with a different Lyapunov function and to formally
endow $\mathbb{R}^{d}$ with a different inner product at every client.

We then apply and tune i-Scaffnew for the problem (FLIX) and propose our new
algorithm for personalized FL, which we call Scafflix, shown as Algorithm 1.

We analyze Scafflix in the strongly convex case, because the analysis of
linear convergence rates in this setting gives clear insights and allows us to
deepen our theoretical understanding of LT and personalization. And to the
best of our knowledge, there is no analysis of Scaffnew in the nonconvex
setting. But we conduct several nonconvex deep learning experiments to show
that our theoretical findings also hold in practice.

Algorithm 1 Scafflix for (FLIX)

1: input: stepsizes $\gamma_{1}>0,\ldots,\gamma_{n}>0$; probability
$p\in(0,1]$; initial estimates $x_{1}^{0},\ldots,x_{n}^{0}\in\mathbb{R}^{d}$
and ${\color[rgb]{0.75,0,0}h_{1}^{0},\ldots,h_{n}^{0}}\in\mathbb{R}^{d}$ such
that $\sum_{i=1}^{n}{\color[rgb]{0.75,0,0}h_{i}^{0}}=0$, personalization
weights $\alpha_{1},\ldots,\alpha_{n}$

2: at the server,
$\gamma\coloneqq\left(\frac{1}{n}\sum_{i=1}^{n}\alpha_{i}^{2}\gamma_{i}^{-1}\right)^{-1}$
$\diamond$ $\gamma$ is used by the server at Step 11

3: at clients in parallel, $x_{i}^{\star}\coloneqq\arg\min f_{i}$$\diamond$
not needed if $\alpha_{i}=1$

4: for $t=0,1,\ldots$ do

5: flip a coin $\theta^{t}\coloneqq\\{1$ with probability $p$, 0
otherwise$\\}$

6: for $i=1,\ldots,n$, at clients in parallel, do

7: $\tilde{x}_{i}^{t}\coloneqq\alpha_{i}x_{i}^{t}+(1-\alpha_{i})x_{i}^{\star}$
$\diamond$ estimate of the personalized model $\tilde{x}_{i}^{\star}$

8: compute an estimate $g_{i}^{t}$ of $\nabla f_{i}(\tilde{x}_{i}^{t})$

9: $\hat{x}_{i}^{t}\coloneqq
x_{i}^{t}-\frac{\gamma_{i}}{\alpha_{i}}\big{(}g_{i}^{t}-{\color[rgb]{0.75,0,0}h_{i}^{t}}\big{)}$
$\diamond$ local SGD step

10: if $\theta^{t}=1$ then

11: send $\frac{\alpha_{i}^{2}}{\gamma_{i}}\hat{x}_{i}^{t}$ to the server,
which aggregates
$\bar{x}^{t}\coloneqq\frac{\gamma}{n}\sum_{j=1}^{n}\frac{\alpha_{i}^{2}}{\gamma_{i}}\hat{x}_{j}^{t}$
and broadcasts it to all clients $\diamond$ communication, but only with small
probability $p$

12: $x_{i}^{t+1}\coloneqq\bar{x}^{t}$

13:
${\color[rgb]{0.75,0,0}h_{i}^{t+1}}\coloneqq{\color[rgb]{0.75,0,0}h_{i}^{t}}+\frac{p\alpha_{i}}{\gamma_{i}}\big{(}\bar{x}^{t}-\hat{x}_{i}^{t}\big{)}$$\diamond$
update of the local control variate $\color[rgb]{0.75,0,0}h_{i}^{t}$

14: else

15: $x_{i}^{t+1}\coloneqq\hat{x}_{i}^{t}$

16:
${\color[rgb]{0.75,0,0}h_{i}^{t+1}}\coloneqq{\color[rgb]{0.75,0,0}h_{i}^{t}}$

17: end if

18: end for

19: end for

######  Assumption 1 (Smoothness and strong convexity).

In the problem (FLIX) (and (ERM) as the particular case $\alpha_{i}\equiv 1$),
we assume that for every $i\in[n]$, the function $f_{i}$ is $L_{i}$-smooth and
$\mu_{i}$-strongly convex,111A function
$f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ is said to be $L$-smooth if it is
differentiable and its gradient is Lipschitz continuous with constant $L$;
that is, for every $x\in\mathbb{R}^{d}$ and $y\in\mathbb{R}^{d}$, $\|\nabla
f(x)-\nabla f(y)\|\leq L\|x-y\|$, where, here and throughout the paper, the
norm is the Euclidean norm. $f$ is said to be $\mu$-strongly convex if
$f-\frac{\mu}{2}\|\cdot\|^{2}$ is convex. We refer to Bauschke and Combettes
(2017) for such standard notions of convex analysis. for some
$L_{i}\geq\mu_{i}>0$. This implies that the problem is strongly convex, so
that its solution $x^{\star}$ exists and is unique.

We also make the two following assumptions on the stochastic gradients
$g_{i}^{t}$ used in Scafflix (and i-Scaffnew as a particular case with
$\alpha_{i}\equiv 1)$.

######  Assumption 2 (Unbiasedness).

We assume that for every $t\geq 0$ and $i\in[n]$, $g_{i}^{t}$ is an unbiased
estimate of $\nabla f_{i}(\tilde{x}_{i}^{t})$; that is,

$\mathbb{E}\\!\left[g_{i}^{t}\;|\;\tilde{x}_{i}^{t}\right]=\nabla
f_{i}(\tilde{x}_{i}^{t}).$

To characterize unbiased stochastic gradient estimates, the modern notion of
_expected smoothness_ is well suited (Gower et al., 2019; Gorbunov et al.,
2020b):

######  Assumption 3 (Expected smoothness).

We assume that, for every $i\in[n]$, there exist constants $A_{i}\geq L_{i}$
222We can suppose $A_{i}\geq L_{i}$. Indeed, we have the bias-variance
decomposition $\mathbb{E}\\!\left[{\left\lVert g_{i}^{t}-\nabla
f_{i}(\tilde{x}_{i}^{\star})\right\rVert}^{2}\;|\;\tilde{x}_{i}^{t}\right]={\left\lVert\nabla
f_{i}(\tilde{x}_{i}^{t})-\nabla
f_{i}(\tilde{x}_{i}^{\star})\right\rVert}^{2}+\mathbb{E}\\!\left[{\left\lVert
g_{i}^{t}-\nabla
f_{i}(\tilde{x}_{i}^{t})\right\rVert}^{2}\;|\;\tilde{x}_{i}^{t}\right]\geq{\left\lVert\nabla
f_{i}(\tilde{x}_{i}^{t})-\nabla
f_{i}(\tilde{x}_{i}^{\star})\right\rVert}^{2}$. Assuming that $L_{i}$ is the
best known smoothness constant of $f_{i}$, we cannot improve the constant
$L_{i}$ such that for every $x\in\mathbb{R}^{d}$, ${\left\lVert\nabla
f_{i}(x)-\nabla f_{i}(\tilde{x}_{i}^{\star})\right\rVert}^{2}\leq
2L_{i}D_{f_{i}}(x,\tilde{x}_{i}^{\star})$. Therefore, $A_{i}$ in (2) has to be
$\geq L_{i}$.  and $C_{i}\geq 0$ such that, for every $t\geq 0$,

$\mathbb{E}\\!\left[{\left\lVert g_{i}^{t}-\nabla
f_{i}(\tilde{x}_{i}^{\star})\right\rVert}^{2}\;|\;\tilde{x}_{i}^{t}\right]\leq
2A_{i}D_{f_{i}}(\tilde{x}_{i}^{t},\tilde{x}_{i}^{\star})+C_{i},$ (2)

where $D_{\varphi}(x,x^{\prime})\coloneqq f(x)-f(x^{\prime})-\langle\nabla
f(x^{\prime}),x-x^{\prime}\rangle\geq 0$ denotes the Bregman divergence of a
function $\varphi$ at points $x,x^{\prime}\in\mathbb{R}^{d}$.

Thus, unlike the analysis in Mishchenko et al. (2022)[Assumption 4.1], where
the same constants are assumed for all clients, since we consider
personalization, we individualize the analysis: we consider that each client
can be different and use stochastic gradients characterized by its own
constants $A_{i}$ and $C_{i}$. This is more representative of practical
settings. Assumption 3 is general and covers in particular the following two
important cases (Gower et al., 2019):

  1. 1.

(bounded variance) If $g_{i}^{t}$ is equal to $\nabla
f_{i}(\tilde{x}_{i}^{t})$ plus a zero-mean random error of variance
$\sigma_{i}^{2}$ (this covers the case of the exact gradient $g_{i}^{t}=\nabla
f_{i}(\tilde{x}_{i}^{t})$ with $\sigma_{i}=0$), then Assumption 3 is satisfied
with $A_{i}=L_{i}$ and $C_{i}=\sigma_{i}^{2}$.

  2. 2.

(sampling) If $f_{i}=\frac{1}{n_{i}}\sum_{j=1}^{n_{i}}f_{i,j}$ for some
$L_{i}$-smooth functions $f_{i,j}$ and $g_{i}^{t}=\nabla
f_{i,j^{t}}(\tilde{x}_{i}^{t})$ for some $j^{t}$ chosen uniformly at random in
$[n_{i}]$, then Assumption 3 is satisfied with $A_{i}=2L_{i}$ and
$C_{i}=\big{(}\frac{2}{n_{i}}\sum_{j=1}^{n_{i}}\left\|\nabla
f_{i,j}(\tilde{x}_{i}^{\star})\right\|^{2}\big{)}-2\left\|\nabla
f_{i}(\tilde{x}_{i}^{\star})\right\|^{2}$ (this can be extended to minibatch
and nonuniform sampling).

We now present our main convergence result:

######  Theorem 1 (fast linear convergence).

In (FLIX) and Scafflix, suppose that Assumptions 1, 2, 3 hold and that for
every $i\in[n]$, $0<\gamma_{i}\leq\frac{1}{A_{i}}$. For every $t\geq 0$,
define the Lyapunov function

$\Psi^{t}\coloneqq\frac{1}{n}\sum_{i=1}^{n}\frac{\gamma_{\min}}{\gamma_{i}}\left\|\tilde{x}_{i}^{t}-\tilde{x}_{i}^{\star}\right\|^{2}+\frac{\gamma_{\min}}{p^{2}}\frac{1}{n}\sum_{i=1}^{n}\gamma_{i}\left\|h_{i}^{t}-\nabla
f_{i}(\tilde{x}_{i}^{\star})\right\|^{2},$ (3)

where $\gamma_{\min}\coloneqq\min_{i\in[n]}\gamma_{i}$. Then Scafflix
converges linearly: for every $t\geq 0$,

$\mathbb{E}\\!\left[\Psi^{t}\right]\leq(1-\zeta)^{t}\Psi^{0}+\frac{\gamma_{\min}}{\zeta}\frac{1}{n}\sum_{i=1}^{n}\gamma_{i}C_{i},$
(4)

where

$\zeta=\min\left(\min_{i\in[n]}\gamma_{i}\mu_{i},p^{2}\right).$ (5)

It is important to note that the range of the stepsizes $\gamma_{i}$, the
Lyapunov function $\Psi^{t}$ and the convergence rate in (4)–(5) do not depend
on the personalization weights $\alpha_{i}$; they only play a role in the
definition of the personalized models $\tilde{x}_{i}^{t}$ and
$\tilde{x}_{i}^{\star}$. Indeed, the convergence speed essentially depends on
the conditioning of the functions $x\mapsto
f_{i}\big{(}\alpha_{i}x+(1-\alpha_{i})x_{i}^{\star}\big{)}$, which are
independent from the $\alpha_{i}$. More precisely, let us define, for every
$i\in[n]$,

$\kappa_{i}\coloneqq\frac{L_{i}}{\mu_{i}}\geq
1\quad\mbox{and}\quad\kappa_{\max}=\max_{i\in[n]}\kappa_{i},$

and let us study the complexity of of Scafflix to reach $\epsilon$-accuracy,
i.e. $\mathbb{E}\\!\left[\Psi^{t}\right]\leq\epsilon$. If, for every
$i\in[n]$, $C_{i}=0$, $A_{i}=\Theta(L_{i})$, and
$\gamma_{i}=\Theta(\frac{1}{A_{i}})=\Theta(\frac{1}{L_{i}})$, the iteration
complexity of Scafflix is

$\mathcal{O}\left(\left(\kappa_{\max}+\frac{1}{p^{2}}\right)\log(\Psi^{0}\epsilon^{-1})\right).$
(6)

And since communication occurs with probability $p$, the communication
complexity of Scafflix is

$\mathcal{O}\left(\left(p\kappa_{\max}+\frac{1}{p}\right)\log(\Psi^{0}\epsilon^{-1})\right).$
(7)

Note that $\kappa_{\max}$ can be much smaller than
$\kappa_{\mathrm{global}}\coloneqq\frac{\max_{i}L_{i}}{\min_{i}\mu_{i}}$,
which is the condition number that appears in the rate of Scaffnew with
$\gamma=\frac{1}{\max_{i}A_{i}}$. Thus, Scafflix is much more versatile and
adapted to FL with heterogeneous data than Scaffnew.

######  Corollary 1 (case $C_{i}\equiv 0$).

In the conditions of Theorem 1, if
$p=\Theta\big{(}\frac{1}{\sqrt{\kappa_{\max}}}\big{)}$ and, for every
$i\in[n]$, $C_{i}=0$, $A_{i}=\Theta(L_{i})$, and
$\gamma_{i}=\Theta(\frac{1}{A_{i}})=\Theta(\frac{1}{L_{i}})$, the
communication complexity of Scafflix is

$\mathcal{O}\left(\sqrt{\kappa_{\max}}\log(\Psi^{0}\epsilon^{-1})\right).$ (8)

######  Corollary 2 (general stochastic gradients).

In the conditions of Theorem 1, if $p=\sqrt{\min_{i\in[n]}\gamma_{i}\mu_{i}}$
and, for every $i\in[n]$,

$\gamma_{i}=\min\left(\frac{1}{A_{i}},\frac{\epsilon\mu_{\min}}{2C_{i}}\right)$
(9)

(or $\gamma_{i}\coloneqq\frac{1}{A_{i}}$ if $C_{i}=0$), where
$\mu_{\min}\coloneqq\min_{j\in[n]}\mu_{j}$, the iteration complexity of
Scafflix is

$\mathcal{O}\\!\left(\\!\left(\max_{i\in[n]}\max\left(\frac{A_{i}}{\mu_{i}},\frac{C_{i}}{\epsilon\mu_{\min}\mu_{i}}\right)\right)\log(\Psi^{0}\epsilon^{-1})\\!\right)\\!=\\!\mathcal{O}\\!\left(\\!\max\left(\max_{i\in[n]}\frac{A_{i}}{\mu_{i}},\max_{i\in[n]}\frac{C_{i}}{\epsilon\mu_{\min}\mu_{i}}\right)\log(\Psi^{0}\epsilon^{-1})\\!\right)$
(10)

and its communication complexity is

$\mathcal{O}\left(\\!\max\left(\max_{i\in[n]}\sqrt{\frac{A_{i}}{\mu_{i}}},\max_{i\in[n]}\sqrt{\frac{C_{i}}{\epsilon\mu_{\min}\mu_{i}}}\right)\log(\Psi^{0}\epsilon^{-1})\right).$
(11)

If $A_{i}=\Theta(L_{i})$ uniformly, we have
$\max_{i\in[n]}\sqrt{\frac{A_{i}}{\mu_{i}}}=\Theta(\sqrt{\kappa_{\max}})$.
Thus, we see that thanks to LT, the communication complexity of Scafflix is
accelerated, as it depends on $\sqrt{\kappa_{\max}}$ and
$\frac{1}{\sqrt{\epsilon}}$.

In the expressions above, the acceleration effect of personalization is not
visible: it is “hidden” in $\Psi^{0}$, because every client computes
$x_{i}^{t}$ but what matters is its personalized model $\tilde{x}_{i}^{t}$,
and
$\left\|\tilde{x}_{i}^{t}-\tilde{x}_{i}^{\star}\right\|^{2}=\alpha_{i}^{2}\left\|x_{i}^{t}-x^{\star}\right\|^{2}$.
In particular, assuming that $x_{1}^{0}=\cdots=x_{n}^{0}=x^{0}$ and
$h_{i}^{0}=\nabla f_{i}(\tilde{x}_{i}^{0})$, we have

$\Psi^{0}\leq\frac{\gamma_{\min}}{n}\left\|x^{0}-x^{\star}\right\|^{2}\sum_{i=1}^{n}\alpha_{i}^{2}\left(\frac{1}{\gamma_{i}}+\frac{\gamma_{i}L_{i}^{2}}{p^{2}}\right)\leq\big{(}\max_{i}\alpha_{i}^{2}\big{)}\frac{\gamma_{\min}}{n}\left\|x^{0}-x^{\star}\right\|^{2}\sum_{i=1}^{n}\left(\frac{1}{\gamma_{i}}+\frac{\gamma_{i}L_{i}^{2}}{p^{2}}\right),$

and we see that the contribution of every client to the initial gap $\Psi^{0}$
is weighted by $\alpha_{i}^{2}$. Thus, the smaller the $\alpha_{i}$, the
smaller $\Psi^{0}$ and the faster the convergence. This is why personalization
is an acceleration mechanism in our setting.

##  4 Experiments

Figure 1: The objective gap $f(x^{k})-f^{\star}$ and the squared gradient norm
${\left\lVert\nabla f(x^{k})\right\rVert}^{2}$ against the number $k$ of
communication rounds for Scafflix and GD on the problem (FLIX). We set all
$\alpha_{i}$ to the same value for simplicity. The dashed line represents GD,
while the solid line represents Scafflix. We observe the double communication
acceleration achieved through explicit personalization and local training.
Specifically, (a) for a given algorithm, smaller $\alpha_{i}$s (i.e. more
personalized models) lead to faster convergence; (b) comparing the two
algorithms, Scafflix is faster than GD, thanks to its local training
mechanism.

We first consider a convex logistic regression problem to show that the
empirical behavior of Scafflix is in accordance with the theoretical
convergence guarantees available in the convex case. Then, we make extensive
experiments of training neural networks on large-scale distributed
datasets.333Code is available at https://github.com/WilliamYi96/Scafflix.

###  4.1 Prelude: Convex Logistic Regression

We begin our evaluation by considering the standard convex logistic regression
problem with an $l_{2}$ regularizer. This benchmark problem is takes the form
(ERM) with

${f}_{i}{(x)}\coloneqq\frac{1}{n_{i}}\sum_{j=1}^{n_{i}}\log\left(1+\exp(-b_{i,j}x^{T}a_{i,j})\right)+\frac{\mu}{2}{\left\lVert
x\right\rVert}^{2},$ (12)

where $\mu$ represents the regularization parameter, $n_{i}$ is the total
number of data points present at client $i$; $a_{i,j}$ are the training
vectors and the $b_{i,j}\in\\{-1,1\\}$ are the corresponding labels. Every
function $f_{i}$ is $\mu$-strongly convex and $L_{i}$-smooth with
$L_{i}=\frac{1}{4n_{i}}\sum_{j=1}^{n_{i}}{\left\lVert
a_{i,j}\right\rVert}^{2}+\mu$. We set $\mu$ to $0.1$ for this experiment. We
employ the mushrooms, a6a, and w6a datasets from the LibSVM library (Chang and
Lin, 2011) to conduct these tests. The data is distributed evenly across all
clients, and the $\alpha_{i}$ are set to the same value. The results are shown
in Fig. 1. We can observe the double acceleration effect of our approach,
which combines explicit personalization and accelerated local training. Lower
$\alpha_{i}$ values, i.e. more personalization, yield faster convergence for
both GD and Scafflix. Moreover, Scafflix is much faster than GD, thanks to its
specialized local training mechanism.

###  4.2 Neural Network Training: Datasets and Baselines for Evaluation

To assess the generalization capabilities of Scafflix, we undertake a
comprehensive evaluation involving the training of neural networks using two
widely-recognized large-scale FL datasets.

#### Datasets.

Our selection comprises two notable large-scale FL datasets: Federated
Extended MNIST (FEMNIST) (Caldas et al., 2018), and Shakespeare (McMahan et
al., 2017). FEMNIST is a character recognition dataset consisting of 671,585
samples. In accordance with the methodology outlined in FedJax (Ro et al.,
2021), we distribute these samples randomly across 3,400 devices. For all
algorithms, we employ a Convolutional Neural Network (CNN) model, featuring
two convolutional layers and one fully connected layer. The Shakespeare
dataset, used for next character prediction tasks, contains a total of 16,068
samples, which we distribute randomly across 1,129 devices. For all algorithms
applied to this dataset, we use a Recurrent Neural Network (RNN) model,
comprising two Long Short-Term Memory (LSTM) layers and one fully connected
layer.

#### Baselines.

The performance of our proposed Scafflix algorithm is benchmarked against
prominent baseline algorithms, specifically FLIX (Gasanov et al., 2022) and
FedAvg (McMahan et al., 2016). The FLIX algorithm optimizes the FLIX objective
utilizing the SGD method, while FedAvg is designed to optimize the ERM
objective. We employ the official implementations for these benchmark
algorithms. Comprehensive hyperparameter tuning is carried out for all
algorithms, including Scafflix, to ensure optimal results. For both FLIX and
Scafflix, local training is required to achieve the local minima for each
client. By default, we set the local training batch size at $100$ and employ
SGD with a learning rate selected from the set
$C_{s}\coloneqq\\{10^{-5},10^{-4},\cdots,1\\}$. Upon obtaining the local
optimum, we execute each algorithm with a batch size of $20$ for 1000
communication rounds. The model’s learning rate is also selected from the set
$C_{s}$.

###  4.3 Analysis of Generalization with Limited Communication Rounds

In this section, we perform an in-depth examination of the generalization
performance of Scafflix, particularly in scenarios with a limited number of
training epochs. This investigation is motivated by our theoretical evidence
of the double acceleration property of Scafflix. To that aim, we conduct
experiments on both FEMNIST and Shakespeare. These two datasets offer a varied
landscape of complexity, allowing for a comprehensive evaluation of our
algorithm. In order to ensure a fair comparison with other baseline
algorithms, we conducted an extensive search of the optimal hyperparameters
for each algorithm. The performance assessment of the generalization
capabilities was then carried out on a separate, held-out validation dataset.
The hyperparameters that gave the best results in these assessments were
selected as the most optimal set.

In order to examine the impact of personalization, we assume that all clients
have same $\alpha_{i}\equiv\alpha$ and we select $\alpha$ in
$\\{0.1,0.3,0.5,0.7,0.9\\}$. We present the results corresponding to
$\alpha=0.1$ in Fig. 2. Additional comparative analyses with other values of
$\alpha$ are available in the Appendix. As shown in Fig. 2, it is clear that
Scafflix outperforms the other algorithms in terms of generalization on both
the FEMNIST and Shakespeare datasets. Interestingly, the Shakespeare dataset
(next-word prediction) poses a greater challenge compared to the FEMNIST
dataset (digit recognition). Despite the increased complexity of the task,
Scafflix not only delivers significantly better results but also achieves this
faster. Thus, Scafflix is superior both in speed and accuracy.

Figure 2: Comparative generalization analysis with baselines. We set the
communication probability to $p=0.2$. The left figure corresponds to the
FEMNIST dataset with $\alpha=0.1$, while the right figure corresponds to the
Shakespeare dataset with $\alpha=0.3$.

###  4.4 Key Ablation Studies

In this section, we conduct several critical ablation studies to verify the
efficacy of our proposed Scafflix method. These studies investigate the
optimal personalization factor for Scafflix, assess the impact of the number
of clients per communication round, and examine the influence of the
communication probability $p$ in Scafflix.

#### Optimal Personalization Factor.

In this experiment, we explore the effect of varying personalization factors
on the FEMNIST dataset. The results are presented in Fig. 3(a). We set the
batch size to 128 and determine the most suitable learning rate through a
hyperparameter search. We consider linearly increasing personalization factors
within the set $\\{0.1,0.3,0.5,0.7,0.9\\}$. An exponential scale for $\alpha$
is also considered in the Appendix, but the conclusion remains the same.

We note that the optimal personalization factor for the FEMNIST dataset is
$0.3$. Interestingly, personalization factors that yield higher accuracy also
display a slightly larger variance. However, the overall average performance
remains superior. This is consistent with expectations as effective
personalization may emphasize the representation of local data, and thus,
could be impacted by minor biases in the model parameters received from the
server.

(a) $\alpha$s

(b) $\tau$s

(c) $p$s

Figure 3: Key ablation studies: (a) evaluate the influence of difference
personalization factor $\alpha$, (b) examinate the effect of different numbers
of clients participating to communication, (c) compare different values of the
communication probability $p$.

#### Number of Clients Communicating per Round.

In this ablation study, we examine the impact of varying the number of
participating clients in each communication round within the Scafflix
framework. By default, we set this number to 10. Here, we conduct extensive
experiments with different client numbers per round, choosing $\tau$ from
$\\{1,5,10,20\\}$. The results are presented in Fig. 3(b). We can observe that
Scafflix shows minimal sensitivity to changes in the batch size for local
training. However, upon closer examination, we find that larger batch sizes,
specifically $\tau=10$ and $20$, demonstrate slightly improved generalization
performance.

#### Selection of Communication Probability $p$.

In this ablation study, we explore the effects of varying the communication
probability $p$ in Scafflix. We select $p$ from $\\{0.1,0.2,0.5\\}$, and the
corresponding results are shown in Fig. 3(c). We can clearly see that a
smaller value of $p$, indicating reduced communication, facilitates faster
convergence and superior generalization performance. This highlights the
benefits of LT, which not only makes FL faster and more communication-
efficient, but also improves the learning quality.

##  5 Conclusion

In the contemporary era of artificial intelligence, improving federated
learning to achieve faster convergence and reduce communication costs is
crucial to enhance the quality of models trained on huge and heterogeneous
datasets. To address this challenge, we introduced Scafflix, a novel algorithm
that achieves double communication acceleration by redesigning the objective
to support explicit personalization for individual clients, while leveraging a
state-of-the-art local training mechanism. We provided complexity guarantees
in the convex setting, and also validated the effectiveness of our approach in
the nonconvex setting through extensive experiments and ablation studies. We
believe that our work is a significant contribution to the important topic of
communication-efficient federated learning and offers valuable insights for
further investigation in the future.

## References

  * Alam et al. [2022] S. Alam, L. Liu, M. Yan, and M. Zhang.  Fedrolex: Model-heterogeneous federated learning with rolling sub-model extraction.  _arXiv preprint arXiv:2212.01548_ , 2022.
  * Arivazhagan et al. [2019] M. G. Arivazhagan, V. Aggarwal, A. K. Singh, and S. Choudhary.  Federated learning with personalization layers.  _arXiv preprint arXiv:1912.00818_ , 2019.
  * Bauschke and Combettes [2017] H. H. Bauschke and P. L. Combettes.  _Convex Analysis and Monotone Operator Theory in Hilbert Spaces_.  Springer, New York, 2nd edition, 2017.
  * Bui et al. [2019] D. Bui, K. Malik, J. Goetz, H. Liu, S. Moon, A. Kumar, and K. G. Shin.  Federated user representation learning.  _arXiv preprint arXiv:1909.12535_ , 2019.
  * Caldas et al. [2018] S. Caldas, P. Wu, T. Li, J. Konečný, H. B. McMahan, V. Smith, and A. Talwalkar.  LEAF: A benchmark for federated settings.  _arXiv preprint arXiv:1812.01097_ , 2018.
  * Chang and Lin [2011] C.-C. Chang and C.-J. Lin.  LibSVM: A library for support vector machines.  _ACM Transactions on Intelligent Systems and Technology (TIST)_ , 2(3):27, 2011.
  * Chen et al. [2020] Y. Chen, X. Qin, J. Wang, C. Yu, and W. Gao.  Fedhealth: A federated transfer learning framework for wearable healthcare.  _IEEE Intelligent Systems_ , 35(4):83–93, 2020\.
  * Condat and Richtárik [2023] L. Condat and P. Richtárik.  RandProx: Primal-dual optimization algorithms with randomized proximal updates.  In _Proc. of Int. Conf. Learning Representations (ICLR)_ , 2023\.
  * Condat et al. [2022] L. Condat, I. Agarsky, and P. Richtárik.  Provably doubly accelerated federated learning: The first theoretically successful combination of local training and compressed communication.  _arXiv preprint arXiv:2210.13277_ , 2022.
  * Condat et al. [2023] L. Condat, G. Malinovsky, and P. Richtárik.  TAMUNA: Accelerated federated learning with local training and partial participation.  preprint arXiv:2302.09832, 2023.
  * Gasanov et al. [2022] E. Gasanov, A. Khaled, S. Horváth, and P. Richtárik.  Flix: A simple and communication-efficient alternative to local methods in federated learning.  In _24th International Conference on Artificial Intelligence and Statistics (AISTATS 2022)_ , 2022.
  * Gorbunov et al. [2020a] E. Gorbunov, F. Hanzely, and P. Richtárik.  Local SGD: unified theory and new efficient methods.  In _NeurIPS_ , 2020a.
  * Gorbunov et al. [2020b] E. Gorbunov, F. Hanzely, and P. Richtárik.  A unified theory of SGD: Variance reduction, sampling, quantization and coordinate descent.  In _Proc. of 23rd Int. Conf. Artificial Intelligence and Statistics (AISTATS), PMLR 108_ , 2020b.
  * Gower et al. [2019] R. M. Gower, N. Loizou, X. Qian, A. Sailanbayev, E. Shulgin, and P. Richtárik.  SGD: General analysis and improved rates.  In _Proc. of 36th Int. Conf. Machine Learning (ICML)_ , volume PMLR 97, pages 5200–5209, 2019.
  * Grudzień et al. [2023] M. Grudzień, G. Malinovsky, and P. Richtárik.  Can 5th Generation Local Training Methods Support Client Sampling? Yes!  In _Proc. of Int. Conf. Artificial Intelligence and Statistics (AISTATS)_ , Apr. 2023.
  * Haddadpour and Mahdavi [2019] F. Haddadpour and M. Mahdavi.  On the convergence of local descent methods infederated learning.  _arXiv preprint arXiv:1910.14425_ , 2019.
  * Hanzely and Richtárik [2020] F. Hanzely and P. Richtárik.  Federated learning of a mixture of global and local models.  _arXiv preprint arXiv:2002.05516_ , 2020.
  * Horvath et al. [2021] S. Horvath, S. Laskaridis, M. Almeida, I. Leontiadis, S. Venieris, and N. Lane.  Fjord: Fair and accurate federated learning under heterogeneous targets with ordered dropout.  _Advances in Neural Information Processing Systems_ , 34:12876–12889, 2021.
  * Kairouz et al. [2019] P. Kairouz, H. B. McMahan, B. Avent, A. Bellet, M. Bennis, A. N. Bhagoji, K. Bonawitz, Z. Charles, G. Cormode, R. Cummings, R. G. D’Oliveira, H. Eichner, S. E. Rouayheb, D. Evans, J. Gardner, Z. Garrett, A. Gascón, B. Ghazi, P. B. Gibbons, M. Gruteser, Z. Harchaoui, C. He, L. He, Z. Huo, B. Hutchinson, J. Hsu, M. Jaggi, T. Javidi, G. Joshi, M. Khodak, J. Konečný, A. Korolova, F. Koushanfar, S. Koyejo, T. Lepoint, Y. Liu, P. Mittal, M. Mohri, R. Nock, A. Özgür, R. Pagh, M. Raykova, H. Qi, D. Ramage, R. Raskar, D. Song, W. Song, S. U. Stich, Z. Sun, A. T. Suresh, F. Tramèr, P. Vepakomma, J. Wang, L. Xiong, Z. Xu, Q. Yang, F. X. Yu, H. Yu, and S. Zhao.  Advances and open problems in federated learning.  _Foundations and Trends®in Machine Learning_ , 14(1–2):1–210, 2019.
  * Karimireddy et al. [2020] S. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. Suresh.  SCAFFOLD: Stochastic controlled averaging for on-device federated learning.  In _ICML_ , 2020.
  * Khaled et al. [2019] A. Khaled, K. Mishchenko, and P. Richtárik.  First analysis of local GD on heterogeneous data.  In _NeurIPS Workshop on Federated Learning for Data Privacy and Confidentiality_ , pages 1–11, 2019.
  * Khaled et al. [2020] A. Khaled, K. Mishchenko, and P. Richtárik.  Tighter theory for local SGD on identical and heterogeneous data.  In _The 23rd International Conference on Artificial Intelligence and Statistics (AISTATS 2020)_ , 2020.
  * Li and Wang [2019] D. Li and J. Wang.  Fedmd: Heterogenous federated learning via model distillation.  _arXiv preprint arXiv:1910.03581_ , 2019.
  * Li et al. [2021] Q. Li, B. He, and D. Song.  Model-contrastive federated learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 10713–10722, 2021.
  * Li et al. [2020a] T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith.  Federated optimization in heterogeneous networks.  _Proceedings of Machine learning and systems_ , 2:429–450, 2020a.
  * Li et al. [2020b] X. Li, K. Huang, W. Yang, S. Wang, and Z. Zhang.  On the conver-gence of FedAvg on non-IID data.  In _International Conference on Learning Representations_ , 2020b.
  * Malinovsky et al. [2020] G. Malinovsky, D. Kovalev, E. Gasanov, L. Condat, and P. Richtárik.  From local SGD to local fixed point methods for federated learning.  In _Proc. of 37th Int. Conf. Machine Learning (ICML)_ , 2020.
  * Malinovsky et al. [2022] G. Malinovsky, K. Yi, and P. Richtárik.  Variance reduced Proxskip: Algorithm, theory and application to federated learning.  _arXiv preprint arXiv:2207.04338_ , 2022.
  * Maranjyan et al. [2022] A. Maranjyan, M. Safaryan, and P. Richtárik.  Gradskip: Communication-accelerated local gradient methods with better computational complexity.  _arXiv preprint arXiv:2210.16402_ , 2022.
  * McMahan et al. [2016] B. McMahan, E. Moore, D. Ramage, and B. Agüera y Arcas.  Federated learning of deep networks using model averaging.  _arXiv preprint arXiv:1602.05629_ , 2016.
  * McMahan et al. [2017] H. B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. Agüera y Arcas.  Communication-efficient learning of deep networks from decentralized data.  In _Proceedings of the 20th International Conference on Artificial Intelligence and Statistics (AISTATS)_ , 2017.
  * Mei et al. [2022] Y. Mei, P. Guo, M. Zhou, and V. Patel.  Resource-adaptive federated learning with all-in-one neural composition.  In _Advances in Neural Information Processing Systems_ , 2022.
  * Mishchenko et al. [2022] K. Mishchenko, G. Malinovsky, S. Stich, and P. Richtárik.  ProxSkip: Yes! Local gradient steps provably lead to communication acceleration! Finally!  In _39th International Conference on Machine Learning (ICML 2022)_ , 2022.
  * Mitra et al. [2021] A. Mitra, R. Jaafar, G. Pappas, and H. Hassani.  Linear convergence in federated learning: Tackling client heterogeneity and sparse gradients.  In _Advances in Neural Information Processing Systems 34_ , 2021.
  * Moritz et al. [2016] P. Moritz, R. Nishihara, I. Stoica, and M. I. Jordan.  SparkNet: Training deep networks in Spark.  In _International Conference on Learning Representations (ICLR)_ , 2016.
  * Povey et al. [2015] D. Povey, X. Zhang, and S. Khudanpur.  Parallel training of DNNs with natural gradient and parameter averaging.  In _ICLR Workshop_ , 2015.
  * Ro et al. [2021] J. H. Ro, A. T. Suresh, and K. Wu.  Fedjax: Federated learning simulation with jax.  _arXiv preprint arXiv:2108.02117_ , 2021.
  * T Dinh et al. [2020] C. T Dinh, N. Tran, and J. Nguyen.  Personalized federated learning with moreau envelopes.  _Advances in Neural Information Processing Systems_ , 33:21394–21405, 2020.
  * Wang et al. [2021] J. Wang, Z. Charles, Z. Xu, G. Joshi, H. B. McMahan, B. A. y Arcas, M. Al-Shedivat, G. Andrew, S. Avestimehr, K. Daly, D. Data, S. Diggavi, H. Eichner, A. Gadhikar, Z. Garrett, A. M. Girgis, F. Hanzely, A. Hard, C. He, S. Horvath, Z. Huo, A. Ingerman, M. Jaggi, T. Javidi, P. Kairouz, S. Kale, S. P. Karimireddy, J. Konecny, S. Koyejo, T. Li, L. Liu, M. Mohri, H. Qi, S. J. Reddi, P. Richtarik, K. Singhal, V. Smith, M. Soltanolkotabi, W. Song, A. T. Suresh, S. U. Stich, A. Talwalkar, H. Wang, B. worth, S. Wu, F. X. Yu, H. Yuan, M. Zaheer, M. Zhang, T. Zhang, C. Zheng, C. Zhu, and W. Zhu.  A field guide to federated optimization.  _arXiv preprint arXiv:2107.06917_ , 2021.
  * Yang et al. [2020] H. Yang, H. He, W. Zhang, and X. Cao.  Fedsteg: A federated transfer learning framework for secure image steganalysis.  _IEEE Transactions on Network Science and Engineering_ , 8(2):1084–1094, 2020.

##  Appendix A Proposed i-Scaffnew algorithm

We consider solving (ERM) with the proposed i-Scaffnew algorithm, shown as
Algorithm 2 (applying i-Scaffnew to (FLIX) yields Scafflix, as we discuss
subsequently in Section B).

Algorithm 2 i-Scaffnew for (ERM)

1: input: stepsizes $\gamma_{1}>0,\ldots,\gamma_{n}>0$; probability
$p\in(0,1]$; initial estimates $x_{1}^{0},\ldots,x_{n}^{0}\in\mathbb{R}^{d}$
and ${\color[rgb]{0.75,0,0}h_{1}^{0},\ldots,h_{n}^{0}}\in\mathbb{R}^{d}$ such
that $\sum_{i=1}^{n}{\color[rgb]{0.75,0,0}h_{i}^{0}}=0$.

2: at the server,
$\gamma\coloneqq\left(\frac{1}{n}\sum_{i=1}^{n}\gamma_{i}^{-1}\right)^{-1}$
$\diamond$ $\gamma$ is used by the server for Step 9

3: for $t=0,1,\ldots$ do

4: flip a coin $\theta^{t}\coloneqq\\{1$ with probability $p$, 0
otherwise$\\}$

5: for $i=1,\ldots,n$, at clients in parallel, do

6: compute an estimate $g_{i}^{t}$ of $\nabla f_{i}(x_{i}^{t})$

7: $\hat{x}_{i}^{t}\coloneqq
x_{i}^{t}-\gamma_{i}\big{(}g_{i}^{t}-{\color[rgb]{0.75,0,0}h_{i}^{t}}\big{)}$
$\diamond$ local SGD step

8: if $\theta^{t}=1$ then

9: send $\frac{1}{\gamma_{i}}\hat{x}_{i}^{t}$ to the server, which aggregates
$\bar{x}^{t}\coloneqq\frac{\gamma}{n}\sum_{j=1}^{n}\frac{1}{\gamma_{i}}\hat{x}_{j}^{t}$
and broadcasts it to all clients $\diamond$ communication, but only with small
probability $p$

10: $x_{i}^{t+1}\coloneqq\bar{x}^{t}$

11:
${\color[rgb]{0.75,0,0}h_{i}^{t+1}}\coloneqq{\color[rgb]{0.75,0,0}h_{i}^{t}}+\frac{p}{\gamma_{i}}\big{(}\bar{x}^{t}-\hat{x}_{i}^{t}\big{)}$$\diamond$
update of the local control variate $\color[rgb]{0.75,0,0}h_{i}^{t}$

12: else

13: $x_{i}^{t+1}\coloneqq\hat{x}_{i}^{t}$

14:
${\color[rgb]{0.75,0,0}h_{i}^{t+1}}\coloneqq{\color[rgb]{0.75,0,0}h_{i}^{t}}$

15: end if

16: end for

17: end for

######  Theorem 2 (fast linear convergence).

In (ERM) and i-Scaffnew, suppose that Assumptions 1, 2, 3 hold and that for
every $i\in[n]$, $0<\gamma_{i}\leq\frac{1}{A_{i}}$. For every $t\geq 0$,
define the Lyapunov function

$\Psi^{t}\coloneqq\sum_{i=1}^{n}\frac{1}{\gamma_{i}}\left\|x_{i}^{t}-x^{\star}\right\|^{2}+\frac{1}{p^{2}}\sum_{i=1}^{n}\gamma_{i}\left\|h_{i}^{t}-\nabla
f_{i}(x^{\star})\right\|^{2}.$ (13)

Then i-Scaffnew converges linearly: for every $t\geq 0$,

$\mathbb{E}\\!\left[\Psi^{t}\right]\leq(1-\zeta)^{t}\Psi^{0}+\frac{1}{\zeta}\sum_{i=1}^{n}\gamma_{i}C_{i},$
(14)

where

$\zeta=\min\left(\min_{i\in[n]}\gamma_{i}\mu_{i},p^{2}\right).$ (15)

###### Proof.

To simplify the analysis of i-Scaffnew, we introduce vector notations: the
problem (ERM) can be written as

$\mathrm{find}\
\mathbf{x}^{\star}=\operatorname*{arg\,min}_{\mathbf{x}\in\mathcal{X}}\
\mathbf{f}(\mathbf{x})\quad\mbox{s.t.}\quad W\mathbf{x}=0,$ (16)

where $\mathcal{X}\coloneqq\mathbb{R}^{d\times n}$, an element
$\mathbf{x}=(x_{i})_{i=1}^{n}\in\mathcal{X}$ is a collection of vectors
$x_{i}\in\mathbb{R}^{d}$,
$\mathbf{f}:\mathbf{x}\in\mathcal{X}\mapsto\sum_{i=1}^{n}f_{i}(x_{i})$, the
linear operator $W:\mathcal{X}\rightarrow\mathcal{X}$ maps
$\mathbf{x}=(x_{i})_{i=1}^{n}$ to
$(x_{i}-\frac{1}{n}\sum_{j=1}^{n}\frac{\gamma}{\gamma_{j}}x_{j})_{i=1}^{n}$,
for given values $\gamma_{1}>0,\ldots,\gamma_{n}>0$ and their harmonic mean
$\gamma=\left(\frac{1}{n}\sum_{i=1}^{n}\gamma_{i}^{-1}\right)^{-1}$. The
constraint $W\mathbf{x}=0$ means that $\mathbf{x}$ minus its weighted average
is zero; that is, $\mathbf{x}$ has identical components $x_{1}=\cdots=x_{n}$.
Thus, (16) is indeed equivalent to (ERM).
$\mathbf{x}^{\star}\coloneqq(x^{\star})_{i=1}^{n}\in\mathcal{X}$ is the unique
solution to (16), where $x^{\star}$ is the unique solution to (ERM).

Moreover, we introduce the weighted inner product in $\mathcal{X}$:
$(\mathbf{x},\mathbf{y})\mapsto\langle\mathbf{x},\mathbf{y}\rangle_{\boldsymbol{\gamma}}\coloneqq\sum_{i=1}^{n}\frac{1}{\gamma_{i}}\langle
x_{i},y_{i}\rangle$. Then, the orthogonal projector $P$ onto the hyperspace
$\\{\mathbf{y}\in\mathcal{X}\ :\ y_{1}=\cdots=y_{n}\\}$, with respect to this
weighted inner product, is
$P:\mathbf{x}\in\mathcal{X}\mapsto\mathbf{\bar{x}}=(\bar{x})_{i=1}^{n}$ with
$\bar{x}=\frac{\gamma}{n}\sum_{i=1}^{n}\frac{1}{\gamma_{i}}x_{i}$ (because
$\bar{x}$ minimizes
$\left\|\mathbf{\bar{x}}-\mathbf{x}\right\|^{2}_{\boldsymbol{\gamma}}$, so
that $\frac{1}{n}\sum_{i=1}^{n}\frac{1}{\gamma_{i}}(\bar{x}-x_{i})=0$). Thus,
$P$, as well as $W=\mathrm{Id}-P$, where $\mathrm{Id}$ denotes the identity,
are self-adjoint and positive linear operators with respect to the weighted
inner product. Moreover, for every $\mathbf{x}\in\mathcal{X}$,

$\left\|\mathbf{x}\right\|^{2}_{\boldsymbol{\gamma}}=\left\|P\mathbf{x}\right\|^{2}_{\boldsymbol{\gamma}}+\left\|W\mathbf{x}\right\|^{2}_{\boldsymbol{\gamma}}=\left\|\mathbf{\bar{x}}\right\|^{2}_{\boldsymbol{\gamma}}+\left\|W\mathbf{x}\right\|^{2}_{\boldsymbol{\gamma}}=\frac{n}{\gamma}\left\|\bar{x}\right\|^{2}+\left\|W\mathbf{x}\right\|^{2}_{\boldsymbol{\gamma}},$

where $\mathbf{\bar{x}}=(\bar{x})_{i=1}^{n}$ and
$\bar{x}=\frac{\gamma}{n}\sum_{i=1}^{n}\frac{1}{\gamma_{i}}x_{i}$.

Let us introduce further vector notations for the variables of i-Scaffnew: for
every $t\geq 0$, we define the _scaled_ concatenated control variate
$\mathbf{h}^{t}\coloneqq(\gamma_{i}h_{i}^{t})_{i=1}^{n}$,
$\mathbf{h}^{\star}\coloneqq(\gamma_{i}h_{i}^{\star})_{i=1}^{n}$, with
$h_{i}^{\star}\coloneqq\nabla f_{i}(x^{\star})$,
$\mathbf{\bar{x}}^{t}\coloneqq(\bar{x}^{t})_{i=1}^{n}$,
$\mathbf{w}^{t}\coloneqq(w_{i}^{t})_{i=1}^{n}$, with $w_{i}^{t}\coloneqq
x_{i}^{t}-\gamma_{i}g_{i}^{t}$,
$\mathbf{w}^{\star}\coloneqq(w_{i}^{\star})_{i=1}^{n}$, with
$w_{i}^{\star}\coloneqq x_{i}^{\star}-\gamma_{i}\nabla f_{i}(x_{i}^{\star})$,
$\mathbf{\hat{h}}^{t}\coloneqq\mathbf{h}^{t}-pW\mathbf{\hat{x}}^{t}$. Finally,
we denote by $\mathcal{F}_{0}^{t}$ the $\sigma$-algebra generated by the
collection of $\mathcal{X}$-valued random variables
$\mathbf{x}^{0},\mathbf{h}^{0},\ldots,\mathbf{x}^{t},\mathbf{h}^{t}$ and by
$\mathcal{F}^{t}$ the $\sigma$-algebra generated by these variables, as well
as the stochastic gradients $g_{i}^{t}$.

We can then rewrite the iteration of i-Scaffnew as:

$\mathbf{\hat{x}}^{t}\coloneqq\mathbf{w}^{t}+\mathbf{h}^{t}$

if $\theta^{t}=1$ then

$\mathbf{x}^{t+1}\coloneqq\mathbf{\bar{x}}^{t}$

$\mathbf{h}^{t+1}\coloneqq\mathbf{h}^{t}-pW\mathbf{\hat{x}}^{t}$

else

$\mathbf{x}^{t+1}\coloneqq\mathbf{\hat{x}}^{t}$

$\mathbf{h}^{t+1}\coloneqq\mathbf{h}^{t}$

end if

We suppose that $\sum_{i=1}^{n}h^{0}_{i}=0$. Then, it follows from the
definition of $\bar{x}^{t}$ that
$\frac{\gamma}{n}\sum_{j=1}^{n}\frac{1}{\gamma_{i}}(\bar{x}^{t}-\hat{x}_{j}^{t})=0$,
so that for every $t\geq 0$, $\sum_{i=1}^{n}h^{t}_{i}=0$; that is,
$W\mathbf{h}^{t}=\mathbf{h}^{t}$.

Let $t\geq 0$. We have

$\displaystyle\mathbb{E}\\!\left[\left\|\mathbf{x}^{t+1}-\mathbf{x}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}\;|\;\mathcal{F}^{t}\right]$
$\displaystyle=p\left\|\mathbf{\bar{x}}^{t}-\mathbf{x}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+(1-p)\left\|\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star}\right\|^{2}_{\boldsymbol{\gamma}},$

with

$\left\|\mathbf{\bar{x}}^{t}-\mathbf{x}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}=\left\|\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}-\left\|W\mathbf{\hat{x}}^{t}\right\|^{2}_{\boldsymbol{\gamma}}.$

Moreover,

$\displaystyle\left\|\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}$
$\displaystyle=\left\|\mathbf{w}^{t}-\mathbf{w}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+2\langle\mathbf{w}^{t}-\mathbf{w}^{\star},\mathbf{h}^{t}-\mathbf{h}^{\star}\rangle_{\boldsymbol{\gamma}}$
$\displaystyle=\left\|\mathbf{w}^{t}-\mathbf{w}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}-\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+2\langle\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star},\mathbf{h}^{t}-\mathbf{h}^{\star}\rangle_{\boldsymbol{\gamma}}$
$\displaystyle=\left\|\mathbf{w}^{t}-\mathbf{w}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}-\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+2\langle\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star},\mathbf{\hat{h}}^{t}-\mathbf{h}^{\star}\rangle_{\boldsymbol{\gamma}}-2\langle\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star},\mathbf{\hat{h}}^{t}-\mathbf{h}^{t}\rangle_{\boldsymbol{\gamma}}$
$\displaystyle=\left\|\mathbf{w}^{t}-\mathbf{w}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}-\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+2\langle\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star},\mathbf{\hat{h}}^{t}-\mathbf{h}^{\star}\rangle_{\boldsymbol{\gamma}}+2p\langle\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star},W\mathbf{\hat{x}}^{t}\rangle_{\boldsymbol{\gamma}}$
$\displaystyle=\left\|\mathbf{w}^{t}-\mathbf{w}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}-\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+2\langle\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star},\mathbf{\hat{h}}^{t}-\mathbf{h}^{\star}\rangle_{\boldsymbol{\gamma}}+2p\left\|W\mathbf{\hat{x}}^{t}\right\|^{2}_{\boldsymbol{\gamma}}.$

Hence,

$\displaystyle\mathbb{E}\\!\left[\left\|\mathbf{x}^{t+1}-\mathbf{x}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}\;|\;\mathcal{F}^{t}\right]$
$\displaystyle=\left\|\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}-p\left\|W\mathbf{\hat{x}}^{t}\right\|^{2}_{\boldsymbol{\gamma}}$
$\displaystyle=\left\|\mathbf{w}^{t}-\mathbf{w}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}-\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+2\langle\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star},\mathbf{\hat{h}}^{t}-\mathbf{h}^{\star}\rangle_{\boldsymbol{\gamma}}+p\left\|W\mathbf{\hat{x}}^{t}\right\|^{2}_{\boldsymbol{\gamma}}.$

On the other hand, we have

$\displaystyle\mathbb{E}\\!\left[\left\|\mathbf{h}^{t+1}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}\;|\;\mathcal{F}^{t}\right]$
$\displaystyle=p\left\|\mathbf{\hat{h}}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+(1-p)\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}$

and

$\displaystyle\left\|\mathbf{\hat{h}}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}$
$\displaystyle=\left\|(\mathbf{h}^{t}-\mathbf{h}^{\star})+(\mathbf{\hat{h}}^{t}-\mathbf{h}^{t})\right\|^{2}_{\boldsymbol{\gamma}}$
$\displaystyle=\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+\left\|\mathbf{\hat{h}}^{t}-\mathbf{h}^{t}\right\|^{2}_{\boldsymbol{\gamma}}+2\langle\mathbf{h}^{t}-\mathbf{h}^{\star},\mathbf{\hat{h}}^{t}-\mathbf{h}^{t}\rangle_{\boldsymbol{\gamma}}$
$\displaystyle=\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}-\left\|\mathbf{\hat{h}}^{t}-\mathbf{h}^{t}\right\|^{2}_{\boldsymbol{\gamma}}+2\langle\mathbf{\hat{h}}^{t}-\mathbf{h}^{\star},\mathbf{\hat{h}}^{t}-\mathbf{h}^{t}\rangle_{\boldsymbol{\gamma}}$
$\displaystyle=\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}-\left\|\mathbf{\hat{h}}^{t}-\mathbf{h}^{t}\right\|^{2}_{\boldsymbol{\gamma}}-2p\langle\mathbf{\hat{h}}^{t}-\mathbf{h}^{\star},W(\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star})\rangle_{\boldsymbol{\gamma}}$
$\displaystyle=\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}-p^{2}\left\|W\mathbf{\hat{x}}^{t}\right\|^{2}_{\boldsymbol{\gamma}}-2p\langle
W(\mathbf{\hat{h}}^{t}-\mathbf{h}^{\star}),\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star}\rangle_{\boldsymbol{\gamma}}$
$\displaystyle=\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}-p^{2}\left\|W\mathbf{\hat{x}}^{t}\right\|^{2}_{\boldsymbol{\gamma}}-2p\langle\mathbf{\hat{h}}^{t}-\mathbf{h}^{\star},\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star}\rangle_{\boldsymbol{\gamma}}.$

Hence,

$\displaystyle\mathbb{E}\\!\left[\left\|\mathbf{x}^{t+1}-\mathbf{x}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}\;|\;\mathcal{F}^{t}\right]$
$\displaystyle+\frac{1}{p^{2}}\mathbb{E}\\!\left[\left\|\mathbf{h}^{t+1}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}\;|\;\mathcal{F}^{t}\right]$
$\displaystyle=\left\|\mathbf{w}^{t}-\mathbf{w}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}-\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+2\langle\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star},\mathbf{\hat{h}}^{t}-\mathbf{h}^{\star}\rangle_{\boldsymbol{\gamma}}+p\left\|W\mathbf{\hat{x}}^{t}\right\|^{2}_{\boldsymbol{\gamma}}$
$\displaystyle\quad+\frac{1}{p^{2}}\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}-p\left\|W\mathbf{\hat{x}}^{t}\right\|^{2}_{\boldsymbol{\gamma}}-2\langle\mathbf{\hat{h}}^{t}-\mathbf{h}^{\star},\mathbf{\hat{x}}^{t}-\mathbf{x}^{\star}\rangle_{\boldsymbol{\gamma}}$
$\displaystyle=\left\|\mathbf{w}^{t}-\mathbf{w}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+\frac{1}{p^{2}}\left(1-p^{2}\right)\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}.$
(17)

Moreover, for every $i\in[n]$,

$\displaystyle\left\|w_{i}^{t}-w_{i}^{\star}\right\|^{2}$
$\displaystyle=\left\|x_{i}^{t}-x^{\star}-\gamma_{i}\big{(}g_{i}^{t}-\nabla
f_{i}(x^{\star})\big{)}\right\|^{2}$
$\displaystyle=\left\|x_{i}^{t}-x^{\star}\right\|^{2}-2\gamma_{i}\langle
x_{i}^{t}-x^{\star},g_{i}^{t}-\nabla
f_{i}(x^{\star})\rangle+\gamma_{i}^{2}\left\|g_{i}^{t}-\nabla
f_{i}(x^{\star})\right\|^{2},$

and, by unbiasedness of $g_{i}^{t}$ and Assumption 2,

$\displaystyle\mathbb{E}\\!\left[\left\|w_{i}^{t}-w_{i}^{\star}\right\|^{2}\;|\;\mathcal{F}_{0}^{t}\right]$
$\displaystyle=\left\|x_{i}^{t}-x^{\star}\right\|^{2}-2\gamma_{i}\langle
x_{i}^{t}-x^{\star},\nabla f_{i}(x_{i}^{t})-\nabla f_{i}(x^{\star})\rangle$
$\displaystyle\quad+\gamma_{i}^{2}\mathbb{E}\\!\left[\left\|g_{i}^{t}-\nabla
f_{i}(x^{\star})\right\|^{2}\;|\;\mathcal{F}^{t}\right]$
$\displaystyle\leq\left\|x_{i}^{t}-x^{\star}\right\|^{2}-2\gamma_{i}\langle
x_{i}^{t}-x^{\star},\nabla f_{i}(x_{i}^{t})-\nabla
f_{i}(x^{\star})\rangle+2\gamma_{i}^{2}A_{i}D_{f_{i}}(x_{i}^{t},x^{\star})$
$\displaystyle\quad+\gamma_{i}^{2}C_{i}.$

It is easy to see that $\langle x_{i}^{t}-x^{\star},\nabla
f_{i}(x_{i}^{t})-\nabla
f_{i}(x^{\star})\rangle=D_{f_{i}}(x_{i}^{t},x^{\star})+D_{f_{i}}(x^{\star},x_{i}^{t})$.
This yields

$\displaystyle\mathbb{E}\\!\left[\left\|w_{i}^{t}-w_{i}^{\star}\right\|^{2}\;|\;\mathcal{F}_{0}^{t}\right]$
$\displaystyle\leq\left\|x_{i}^{t}-x^{\star}\right\|^{2}-2\gamma_{i}D_{f_{i}}(x^{\star},x_{i}^{t})-2\gamma_{i}D_{f_{i}}(x_{i}^{t},x^{\star})+2\gamma_{i}^{2}A_{i}D_{f_{i}}(x_{i}^{t},x^{\star})$
$\displaystyle\quad+\gamma_{i}^{2}C_{i}.$

In addition, the strong convexity of $f_{i}$ implies that
$D_{f_{i}}(x^{\star},x_{i}^{t})\geq\frac{\mu_{i}}{2}\left\|x_{i}^{t}-x^{\star}\right\|^{2}$,
so that

$\displaystyle\mathbb{E}\\!\left[\left\|w_{i}^{t}-w_{i}^{\star}\right\|^{2}\;|\;\mathcal{F}_{0}^{t}\right]$
$\displaystyle\leq(1-\gamma_{i}\mu_{i})\left\|x_{i}^{t}-x^{\star}\right\|^{2}-2\gamma_{i}(1-\gamma_{i}A_{i})D_{f_{i}}(x_{i}^{t},x^{\star})+\gamma_{i}^{2}C_{i},$

and since we have supposed $\gamma_{i}\leq\frac{1}{A_{i}}$,

$\displaystyle\mathbb{E}\\!\left[\left\|w_{i}^{t}-w_{i}^{\star}\right\|^{2}\;|\;\mathcal{F}_{0}^{t}\right]$
$\displaystyle\leq(1-\gamma_{i}\mu_{i})\left\|x_{i}^{t}-x^{\star}\right\|^{2}+\gamma_{i}^{2}C_{i}.$

Therefore,

$\mathbb{E}\\!\left[\left\|\mathbf{w}^{t}-\mathbf{w}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}\;|\;\mathcal{F}_{0}^{t}\right]\leq\max_{i\in[n]}(1-\gamma_{i}\mu_{i})\left\|\mathbf{x}^{t}-\mathbf{x}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+\sum_{i=1}^{n}\gamma_{i}C_{i}$

and

$\displaystyle\mathbb{E}\\!\left[\Psi^{t+1}\;|\;\mathcal{F}_{0}^{t}\right]$
$\displaystyle=\mathbb{E}\\!\left[\left\|\mathbf{x}^{t+1}-\mathbf{x}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}\;|\;\mathcal{F}^{t}_{0}\right]+\frac{1}{p^{2}}\mathbb{E}\\!\left[\left\|\mathbf{h}^{t+1}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}\;|\;\mathcal{F}^{t}_{0}\right]$
$\displaystyle\leq\max_{i\in[n]}(1-\gamma_{i}\mu_{i})\left\|\mathbf{x}^{t}-\mathbf{x}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+\frac{1}{p^{2}}\left(1-p^{2}\right)\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+\sum_{i=1}^{n}\gamma_{i}C_{i}$
$\displaystyle\leq(1-\zeta)\left(\left\|\mathbf{x}^{t}-\mathbf{x}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}+\frac{1}{p^{2}}\left\|\mathbf{h}^{t}-\mathbf{h}^{\star}\right\|^{2}_{\boldsymbol{\gamma}}\right)+\sum_{i=1}^{n}\gamma_{i}C_{i}$
$\displaystyle=(1-\zeta)\Psi^{t}+\sum_{i=1}^{n}\gamma_{i}C_{i},$ (18)

where

$\zeta=\min\left(\min_{i\in[n]}\gamma_{i}\mu_{i},p^{2}\right).$

Using the tower rule, we can unroll the recursion in (18) to obtain the
unconditional expectation of $\Psi^{t+1}$. ∎

##  Appendix B From i-Scaffnew to Scafflix

We suppose that Assumptions 1, 2, 3 hold. We define for every $i\in[n]$ the
function $\tilde{f}_{i}:x\in\mathbb{R}^{d}\mapsto
f_{i}\big{(}\alpha_{i}x+(1-\alpha_{i})x^{\star}_{i}\big{)}$. Thus, (FLIX)
takes the form of (ERM) with $f_{i}$ replaced by $\tilde{f}_{i}$.

We want to derive Scafflix from i-Scaffnew applied to (ERM) with $f_{i}$
replaced by $\tilde{f}_{i}$. For this, we first observe that for every
$i\in[n]$, $\tilde{f}_{i}$ is $\alpha_{i}^{2}L_{i}$-smooth and
$\alpha_{i}^{2}\mu_{i}$-strongly convex. This follows easily from the fact
that $\nabla\tilde{f}_{i}(x)=\alpha_{i}\nabla
f_{i}\big{(}\alpha_{i}x+(1-\alpha_{i})x^{\star}_{i}\big{)}$.

Second, for every $t\geq 0$ and $i\in[n]$, $g_{i}^{t}$ is an unbiased estimate
of $\nabla
f_{i}(\tilde{x}_{i}^{t})=\alpha_{i}^{-1}\nabla\tilde{f}_{i}(x_{i}^{t})$.
Therefore, $\alpha_{i}g_{i}^{t}$ is an unbiased estimate of
$\nabla\tilde{f}_{i}(x_{i}^{t})$ satisfying

$\mathbb{E}\\!\left[{\left\lVert\alpha_{i}g_{i}^{t}-\nabla\tilde{f}_{i}(x^{\star})\right\rVert}^{2}\;|\;x_{i}^{t}\right]=\alpha_{i}^{2}\mathbb{E}\\!\left[{\left\lVert
g_{i}^{t}-\nabla
f_{i}(\tilde{x}_{i}^{\star})\right\rVert}^{2}\;|\;\tilde{x}_{i}^{t}\right]\leq
2\alpha_{i}^{2}A_{i}D_{f_{i}}(\tilde{x}_{i}^{t},\tilde{x}_{i}^{\star})+\alpha_{i}^{2}C_{i}.$

Moreover,

$\displaystyle D_{f_{i}}(\tilde{x}_{i}^{t},\tilde{x}_{i}^{\star})$
$\displaystyle=f_{i}(\tilde{x}_{i}^{t})-f_{i}(\tilde{x}_{i}^{\star})-\langle\nabla
f_{i}(\tilde{x}_{i}^{\star}),\tilde{x}_{i}^{t}-\tilde{x}_{i}^{\star}\rangle$
$\displaystyle=\tilde{f}_{i}(x_{i}^{t})-\tilde{f}_{i}(x^{\star})-\langle\alpha_{i}^{-1}\nabla\tilde{f}_{i}(x^{\star}),\alpha_{i}(x_{i}^{t}-x^{\star})\rangle$
$\displaystyle=\tilde{f}_{i}(x_{i}^{t})-\tilde{f}_{i}(x^{\star})-\langle\nabla\tilde{f}_{i}(x^{\star}),x_{i}^{t}-x^{\star}\rangle$
$\displaystyle=D_{\tilde{f}_{i}}(x_{i}^{t},x^{\star}).$

Thus, we obtain Scafflix by applying i-Scaffnew to solve (FLIX), viewed as
(ERM) with $f_{i}$ replaced by $\tilde{f}_{i}$, and further making the
following substitutions in the algorithm: $g_{i}^{t}$ is replaced by
$\alpha_{i}g_{i}^{t}$, $h_{i}^{t}$ is replaced by $\alpha_{i}h_{i}^{t}$ (so
that $h_{i}^{t}$ in Scafflix converges to $\nabla
f_{i}(\tilde{x}_{i}^{\star})$ instead of
$\nabla\tilde{f}_{i}(x^{\star})=\alpha_{i}\nabla
f_{i}(\tilde{x}_{i}^{\star})$), $\gamma_{i}$ is replaced by
$\alpha_{i}^{-2}\gamma_{i}$ (so that the $\alpha_{i}$ disappear in the
theorem).

Accordingly, Theorem 1 follows from Theorem 2, with the same substitutions and
with $A_{i}$, $C_{i}$ and $\mu_{i}$ replaced by $\alpha_{i}^{2}A_{i}$,
$\alpha_{i}^{2}C_{i}$ and $\alpha_{i}^{2}\mu_{i}$, respectively. Finally, the
Lyapunov function is multiplied by $\gamma_{\min}/n$ to make it independent
from $\epsilon$ when scaling the $\gamma_{i}$ by $\epsilon$ in Corollary 2.

We note that i-Scaffnew is recovered as a particular case of Scafflix if
$\alpha_{i}\equiv 1$, so that Scafflix is indeed more general.

##  Appendix C Proof of Corollary 2

We place ourselves in the conditions of Theorem 1. Let $\epsilon>0$. We want
to choose the $\gamma_{i}$ and the number of iterations $T\geq 0$ such that
$\mathbb{E}\\!\left[\Psi^{T}\right]\leq\epsilon$. For this, we bound the two
terms $(1-\zeta)^{T}\Psi^{0}$ and $\frac{\gamma_{\min}}{\zeta
n}\sum_{i=1}^{n}\gamma_{i}C_{i}$ in (4) by $\epsilon/2$.

We set $p=\sqrt{\min_{i\in[n]}\gamma_{i}\mu_{i}}$, so that
$\zeta=\min_{i\in[n]}\gamma_{i}\mu_{i}$. We have

$T\geq\frac{1}{\zeta}\log(2\Psi^{0}\epsilon^{-1})\Rightarrow(1-\zeta)^{T}\Psi^{0}\leq\frac{\epsilon}{2}.$
(19)

Moreover,

$(\forall i\in[n]\mbox{ s.t.\ }C_{i}>0)\
\gamma_{i}\leq\frac{\epsilon\mu_{\min}}{2C_{i}}\Rightarrow\frac{\gamma_{\min}}{\zeta
n}\sum_{i=1}^{n}\gamma_{i}C_{i}\leq\frac{\epsilon}{2}\frac{\left(\min_{j\in[n]}\gamma_{j}\right)\left(\min_{j\in[n]}\mu_{j}\right)}{\min_{j\in[n]}\gamma_{j}\mu_{j}}\leq\frac{\epsilon}{2}.$

Therefore, we set for every $i\in[n]$

$\gamma_{i}\coloneqq\min\left(\frac{1}{A_{i}},\frac{\epsilon\mu_{\min}}{2C_{i}}\right)$

(or $\gamma_{i}\coloneqq\frac{1}{A_{i}}$ if $C_{i}=0$), and we get from (19)
that $\mathbb{E}\\!\left[\Psi^{T}\right]\leq\epsilon$ after

$\mathcal{O}\left(\left(\max_{i\in[n]}\max\left(\frac{A_{i}}{\mu_{i}},\frac{C_{i}}{\epsilon\mu_{\min}\mu_{i}}\right)\right)\log(\Psi^{0}\epsilon^{-1})\right)$

iterations.

##  Appendix D Additional Experimental Results

Figure 4: As part of our experimentation on the FEMNIST dataset, we performed
complementary ablations by incorporating various personalization factors,
represented as $\alpha$. In the main section, we present the results obtained
specifically with $\alpha=0.1$. Furthermore, we extend our analysis by
highlighting the outcomes achieved with $\alpha$ values spanning from 0.3 to
0.9, inclusively.

Figure 5: In our investigation of the Shakespeare dataset, we carried out
complementary ablations, considering a range of personalization factors
denoted as $\alpha$. The selection strategy for determining the appropriate
$\alpha$ values remains consistent with the methodology described in the above
figure.

Figure 6: Ablation studies with different values of the personalization factor
$\alpha$. The left figure is the complementary experiment of linearly
increasing $\alpha$ with full batch size; the right is the figure with
exponentially increasing $\alpha$ with default batch size 20.