The Power of Extrapolation in Federated Learning
Authors: Hanmin Li
Kirill Acharya
Peter Richtárik
Venue: NeurIPS 2024 poster (2024)
Keywords: Federated Learning
Optimization
Primary area: optimization
TL;DR: We extend the FedProx algorithm to its extrapolated counterpart with better theoretical guarantees.

Abstract: We propose and study several server-extrapolation strategies for enhancing the theoretical and empirical convergence properties of the popular federated learning optimizer FedProx [Li et al., 2020]. While it has long been known that some form of extrapolation can help in the practice of FL, only a handful of works provide any theoretical guarantees. The phenomenon seems elusive, and our current theoretical understanding remains severely incomplete. In our work, we focus on smooth convex or strongly convex problems in the interpolation regime. In particular, we propose Extrapolated FedProx (FedExProx), and study three extrapolation strategies: a constant strategy (depending on various smoothness parameters and the number of participating devices), and two smoothness-adaptive strategies; one based on the notion of gradient diversity (FedExProx-GraDS), and the other one based on the stochastic Polyak stepsize (FedExProx-StoPS). Our theory is corroborated with carefully constructed numerical experiments.

The Power of Extrapolation in Federated Learning
Hanmin Li GenAl Center of Excellence KAUST, Saudi Arabia
Kirill Acharya GenAl Center Excellence KAUST, Saudi Arabia acharya. www.hysthh.kkhyhtthh.hhh edu
hanmin, di@kaust.edae edu.
Peter Richtárik
GenAl Center of Excellence KAUST, Saudi Arabia
peter. peter..richtarikkkkuust edu.
Abstract
We propose and study several server-extrapolation strategies for enhancing the theoretical and empirical convergence properties of the popular federated learning optimizer FedProx [Li al., 2020]. While has long been known that some form of extrapolation can help in the practice of FL, only handful of works provide any theoretical guarantees. The phenomenon seems elusive, and our current theoretical understanding remains severely incomplete. In our work, we focus smooth convex or strongly convex problems in the interpolation regime. particular, we propose Extrapolated FedProx (FedExProx), and study three extrapolation strategies: a constant strategy (depending on various smoothness parameters and the number of participating devices), and two smoothhess-adaptis strategies; one based on the notion of gradient diversity (FedExProx-GraDSS and the other one based the stochastic Polyak stepsize (FedExProx-Stan -SoPPS.. Our theory is corroborated with carefully constructed numerical experiments.
Introduction
Federated learning (FL) is distributed training approach for machine learning models, where multiple clients collaborate under the guidance of a central server to optimize a loss function [Konečný et al., 2016, McMahan al., 2017] 7]. This method allows clients contribute to model training while keeping their data private, as it avoids the need for direct data sharing. Often, federated optimization formulated the minimization of finite-sum objective function,
f(x) = 72 -
min ZER4
where each fi Rd + R is the empirical risk of model associated with the i-th client. The federated averaging method (FedAvg) is among the most favored strategies for addressing federated learning problems, proposed by McMahan et al. [2017], Mangasarian and Solodov [1993]. In FedAvg, the server initiates an iteration by selecting a subset of clients for participation in a given round. Each chosen client then proceeds with local training, employing gradient-based techniques like gradient descent (GD) or stochastic gradient descent (SCD) with random reshuffling, as discussed by Bubeck al. [2015), Gower al. [2019], Moulines and Bach [201] 1], Sadiev et al. [2022].
Kirill was student at MIPT during his internship at KAUST.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Li et al. [2020] proposed replacing the local training of each client via SGD in FedAvg with the computation of proximal term, resulting the FedProx algorithm.
1 prox, (2%)
Fk+1
(2)
71
where 7 > 0 is the step size, and the proximal operator is defined
proxyf: (I) == arg min zERd
f (2) + 27 |= ==1?
Contrary to gradient-basee based methods like GD and SGD, algorithms based on proximal operation, such as proximal point method (PPM) [Rockafellar, 1976, Parikh et al., 2014] and stochastic proximal point methods (SPPH) [Asi and Duchi, 2019, Bertsekas, 2011, Khaled and Jin, 2022, Patrascu and Necoara, 2018, Richtárik and Takáe, 2020] benefit from stability against inaccuracies in learning rate specification [Ryu and Boyd, 2014]. Indeed, for GD and SGD, step size that is excessively large can result in divergence of the algorithm, whereas step size that is too small can significantly deteriorate the convergence rate of the algorithm. PPM formally introduced and popularized by the seminal paper of Rockafellar [1976] to solve the variational inequality problems. In practice, the stochastic variant SPPM is more frequently used.
is known that the proximal operator applied to proper, closed and convex function can be viewed as the projection to some level set of the same function depending the value of y. In particular, if we let each fü be the indicator function of a nonempty closed convex set Xr then prox,,. (-) becomes the projection IIx (.) onto the set Xt In this case, FedProx in (2) becomes the parallel projection method for convex feasibility problem [Censor et al., 2001, 2012, Combettes, 1997a, Necoara et al., 2019], we additionally assume
X'
X + 0. 0
well known fact about the parallel projection method is that its empirical efficiency can often be improved by extrapolation [Combettes, 1997a, Necoara al., 2019]. This involves moving further along the line that connects the last iterate Iı and the average projection point, resulting in the iteration
Tk++1 Tk al Σ Tx (Ir) 2t
where as 1 defines extrapolation level. Despite the various heuristic rules proposed over the years for setting as [Bauschke et al., 2006, Censor et al., 2001, Combettes, 1997b], which have demonstrated satisfactory practical performance, it was only recently that the theoretical foundation explaining the success of extrapolation techniques for solving convex feasibility problems was unveiled by Necoara et al. [2019], where the authors considered randomized version of (3) named Randomized Projection Method (RPM). The practical success extrapolation has spurred numerous extensions of existing algorithms. Notably, Jhunjhunwala et al. [2023] combined FedAvg with extrapolation, resulting in FedExP, leveraging insights from an effective heuristic rule [Combettes, 1997b] for setting og follows:
Ee |aa II (xx))l
04k
DE (IT 1xx, (II)
However, the authors did not consider the case of a constant extrapolation parameter, nor did they disclose the relationship between the extrapolation parameter and the stepsize of SGD. The extrapolation parameter can be viewed as server side stepsize in the context federated learning, discussed by Malinovsky et al. [2023].
its effectiveness
In the field of fixed point methods, extrapolation is also known as over-relaxation [Rechardson, 1911]. It is technique used to effectively accelerate the convergence of fixed point methods, including gradient algorithms and proximal splitting algorithms [lutzeler and Hendrickx, 2019, Condat et al., 2023].

---

1.1 Contributions
Our paper contributes in the following ways; for the notations used please refer Appendix A.
Based the insights gained from the convex feasibility problem, we extend FedProx to its extrapolated counterpart FedExProx for both convex and strongly' convex interpolation problems (See Table 1). By optimally setting the constant extrapolation parameter, we obtain
iteration complexity O (1((tt////)) )² in the convex case and O (=(1+77mmax) log (t)
1
in the strongly convex case, when all the clients participate in the training (full participation). We reveal the dependence of the optimal extrapolation parameter on smoothness, indicating that simply averaging the iterates from local training on the server is suboptimal. Instead, extrapolation should be applied to achieve faster convergence. Specifically, compared to FedProx with the same step size "; our method is always at least + 7Lmax + +lmax times better in terms of iteration complexity, see Remark 5.
Our method, FedExProx, improves upon the worst-case iteration complexity O ( ("rman) of (See Table 2). The improvement could lead to acceleration up to a factor of n. see Remark 6. Furthermore, we extend FedExProx to client partial participation setting, showing the dependence of optimal extrap- olation parameter on 7 which is the number of clients participating in the training and the benefits of a larger т. In particular, show that compared the single client setting, with complexity O kmmu)) the full participation version enjoys speed-up up to factor of see Remark 7.
FedExP [Jhunjhunwala et al., 2023] to O (1+72mmax)
Our theory uncovers the relationship between the extrapolation parameter and the step size in typical gradient-type methods, leveraging the power of the Moreau envelope. We also recover RPM of Necoara et al. [2019] as a special case in our analysis (see Remark 12), and show that the heuristic outlined in (4), is in fact step size based gradient diversity [Horváth et al., 2022, Yin et al., 2018] for the Moreau envelopes of client functions.
Building on the insights from Horváth et al. [2022], we propose two adaptive rules for determining the extrapolation parameter: based on gradient diversity (FedExProx-Gra) and the stochastic Polyak step size (FedExProx-StoP))) [Horváth et al., 2022, Loizou et al., 2021]. The proposed methods eliminate reliance the unknown smoothness constant and exhibit "semi- adaitiiiaptiiity", mcaning the algorithm converges with any local step size y and by selecting sufficiently large y, ensure that we lose most factor of in iteration complexity.
We validate our theory with numerical experiments. Numerical evidence suggests that FedExProx achieves 2x or higher speed-up in terms of iteration complexity compared to FedProx and improved performance compared to FedExP. The framework and the plots are included in the Appendix.
1.2 Related work
Stochastic gradient descent. SGD [Robbins and Monro, 1951, Ghadimi and Lan, 2013, Gower et al., 2019, Gorbunov et al., 2020] stands as a cornerstone algorithm utilized across the fields of machine learning. In its simplest form, the algorithm is written as zx+1 2g 11--17a 9xxx), where 17 0 is a scalar step size, g (αxx) represents a stochastic estimator of the true gradient (**x). We recover GD when g(*x)= = V f(xx) The evolution of SGD has been marked by significant advancements since its introduction by Robbins and Monro [1951], leading to various adaptations like stochastic batch gradient descent [Nemirovski al., 2009] and compressed gradient descent [Alistarh et al., 2017, Khirirat et al., 2018]. Gower et al. [2019] presented a framework for analyzing SGD with arbitrary sampling strategies in the convex setting based on expected smoothness, which was later extended by Gorbunov et al. [2020] to the case of local SCD. While many methods have been crafted leverage the stochastic nature of g( substantial research efforts are also dedicated to finding a
'Strongly convex: f is p-strongly convex. PAs we later see in Theorem 1, here Lmax = maxieln" Le. where each L; is smoothness of St. Ly is the smoothness constant of M 1 >. M7,.
3

---

Table I: General comparison of FedExP, RPM and FedExProx in terms of conditions and convergence. Each entry indicates whether the method has the corresponding feature (v/) or not (X). We use the sign where a feature is not applicable to the corresponding method.
Features
FedExP
RPM
FedExProx
Does not require interpolation regime
x
Does not require convexity"
Acceleration in the strongly convex setting"
X
/
Does not require smoothness'
X
/
Allows for partial participation of clients"
X
Works with constant extrapolation parameter
X
Smoothness and partial participation influence extrapolation
X
Semi-adee 'daptivia
X
RPM refers to the randomized projection method of Necoara et al. [2019]. Our method includes it as a special case, see Remark 12
Convexity:
local objective fa is convex, which is the indicator function of the convex set X, in RPM.
The strong convexity pertains f. and for RPM, it indicates that the lincar regularity condition is satisfied.
Smoothness: fi Li-smooth. Our algorithm also applies in the non mon-s case; see Appendix F.2. Jhunjhunwala et al. [2023] provides no convergence guarantee client partial participation setting. The concept of "semi-adaptivity"" is explained in Remark 9.
Table 2: Comparison of convergence of FedExP, FedProx, FedExProx, FedExProx-Gralle and FedExProx-Stoon StoPS. The local step size of FedExP set to be the largest possible value 1/ezr. in the full batch case, where f is the number of local iterations of GD performed. We assume the assumptions of Theorem 1 also hold here. The notations are introduced in Theorem and Theorem 2. The convergence for our methods are described for arbitrary >> O. We use K to denote the total number of iterations. For FedExProx, optimal constant extrapolation is used. The O (-) notation hidden for all complexities in this table.
Full Participation
Method
General Case
Best Case
Worst Case
FedExP
Olomax/ / aup
6Lemax II oue
Glemax/ K
FedProx
Lmax/ /max/ nr (2-1(2-1 Lmax/nk
FedExProx (New)
Lyy 2y(144ymmaz//xx
Leaxx/
FedExProx-GraDl (New) (L+ ExE0 an
(14 ax.
FedExProx-StoPS (New) (1+7lmsx)/7 IK-- ox.s
(1+> // ex.s
(L+>
The O&,P here is determined according the theory of Jhunjhunwala et al. [2023].
Notice that we always have LL, < 1. so the complexity of FedProx is strictly worse than FedExProx. We have Ly (1 + Lmmax) < Immax, see Remark 6.
d We leave out factor of (1+7Lmax)) ((tmxx//((2+mmm))))) which is constant between (1) 1)-
See Remark II for lower bound of ax,s, using which we can rewrite the rate L1 (1mmx)))
better stepsize. An illustration of this is the coordinate-wist adaptive step size Adagrad [Duchi et al., 2011]. 1]. Another approach involves employing matrix step size, as demonstrated by Safaryan et al. [2021], Li et al. [2023, 2024]. Our analysis builds on the theory of SGD mainly adapted from Gower et al. [2019] with additional consideration on the upper bound the step size.
Stochastic proximal point method. PPM was first introduced by Rockafellar [1976] to address the problems variational inequalities its inception. Its transition to stochastic case, motivated by need to efficiently solve large scale optimization problems, results in SPPM. It is often assumed that the proximity operator can be computed efficiently for the algorithm be practical. Over the years, SPPM has been the subject of extensive research, documented by Bertsekas [2011], Bianchi [2016],

---

Patrascu and Necoara [2018]. Unlike traditional gradient- based methods, SPPM is more robust to inaccuracies in learning rate specifications, as demonstrated by Ryu and Boyd [2014]. Asi and Duchi [2019] studied APROX, which includes SPPM as the special case using the full proximal model; APROX was later extended into minibatch case by Asi et al. [2020]. However, this extension was based on model averaging rather than iterate averaging. The convergence rate of SPPM has been analyzed in various contexts by Khaled and Jin [2022], Ryu and Boyd [2014], Yuan and Li [2022], revealing that its performance does not surpass that of SGD in non-convex regimes.
Projection onto convex sets. The projection method originated from efforts solve systems of linear equations or linear inequalities [Kaczmarz, 1937, Von Neumann, 1949, Motzkin and Schoenberg, 1954]. Subsequently. it was generalized address the convex feasibility problem [Combettes, 1997b]. Typically, the method involves projecting onto set X,, where is determined through sampling or other strategies. A particularly relevant method to our paper is the parallel projection method, in which individual projections onto the sets are performed in parallel, and their results are averaged in order to produce the next iterate. It is well-established experimentally that the parallel projection method can be accelerated through extrapolation, with numerous successful heuristics having been proposed adaptively set the extrapolation parameter [Bauschke et al., 2006, Pierra, 1984]. However, only recently a theory was proposed by Necoara et al. [2019] to explain this phenomenon. Necoara et al. [2019] introduced stochastic reformulations of the convex feasibility problem and revealed how the optimal extrapolation parameter depends on the smoothness of the setting and the size of the minibatch. A better result under a linear regularity condition, which is connected to strong convexity, was also obtained. However, the explanation provided by Necoara et al. [2019] was not satisfactory, as it failed to clarify why adaptive rules based on gradient diversity are effective.
Moreau envelope. The concept of the Moreau envelope, also known Moreau-Yosida regulariza- tion, was first introduced by Moreau [1965] as a mathematical tool for handling non-smooth functions. particularly relevant property of the Moreau envelope that executing proximal minimization algorithms on the original objective is equivalent to applying gradient methods to its Moreau envelope [Ryu and Boyd, 2014]. Based on this observation, Davis and Drusvyatskiy [2019] conducted an anal- ysis of several methods, including SPPM for weakly convex and Lipschitz functions. The properties of the Moreau envelope and its applications have been thoroughly investigated in many works including Jourani et al. [2014], Planiden and Wang [2016, 2019]. Beyond its role in proximal minimization algorithms, the Moreau envelope has been utilized in the contexts of personalized federated learning [T Dinh et al., 2020] and meta-learning [Mishchenko et al., 2023].
Adaptive step size. One of the most crucial hyperparameters in training machine learning models with gradient-based methods is the step size. For GD and SGD, determining the step size often depends on the smoothness parameter, which is typically unknown, posing challenges in practical step size selection. There has been a growing interest in adaptive step sizes, leading to the development of numerous adaptive methods that enable real-time computation of the step size. Examples include Adagrad [Duchi et al., 2011], 1]. RMSProp [Hinton et a..], and ADAM [Kingma and Ba, 2015]. Recently, several studies have attempted extend the Polyak step size beyond deterministic settings, leading to the development of the stochastic Polyak step size [Richtárik and Takác, 2020, Horváth et al., 2022, Loizou et al., 2021, Orvieto et al., 2022]. Gradient diversity, first introduced by Yin et al. [2018], was subsequently analyzed theoretically by Horváth et al. [2022].
2 Preliminaries
We now introduce the several definitions and assumptions that are used throughout the paper. Definition (Proximity operator). The proximity operator of an extended-real-valuee function
6: Rd + R U [+oo} with step size > 0 is defined as
1 1|: 2
prox (x) = arg min =ER
(x))
+
x 2
is known that for a proper, closed and convex function Φ. the minimizer of ((2) + 1 ||=- x||
2- exists and unique. Throughout this paper, we assume that the proximal operators are evaluated exactly, with no approximation inexactness.
5

---

Algorithm Extrapolated SPPM (FedExProx) with partial client participation
1: Parameters: extrapolation parameter og 0, step size for the proximity operator 0, starting point Zo E RR number of clients n, total number of iterations K, number of clients participate in the training 7, for simplicity,
7-nice sampling example
2: for k = 0, K do
3: The server samples Sk (1,2,...,n) n) uniformly from all subsets of cardinality
4: The server computes
3k+1 = 2k + Ok
prox,ft (Ix)) IKk
(7)
VES&
5: end for
Definition 2 (Moreau envelope). The Moreau envelope of an extenneed-eaaa-vllleee function Φ Rd + R (+00) with step size 7 >> O is defined as
1 1| z11? +
M (x)
min =ER
((z)
27
The following assumptions are used in our analysis. We use the notation [n]] for the set nn..
Assumption (Differeeniaabliity) The function fi in (1) is differentiable for all i E [n]]
Assumption (Interpolation regime). There exists I. Rdd such that V fi( for all E [m].
Note that Assumption 2 indicates that each and f are lower bounded. In this paper, focus cases where the interpolation regime holds. This assumption often holds in modern deep learning which overparameterized where the number of parameters greatly exceeds the number of data points, as justified by Arora et al. [2019], Montanari and Zhong [2022]. Our motivation for this assumption partly arises from the convex feasibility problem [Combettes, 1997a, Necoara al., 2019], wherein the intersection X is presumed nonempty. This is equivalent to assuming that the interpolation regime holds when fi is the indicator function of the nonempty closed convex set X, Further motivations derived from the proof for this assumption will be discussed later.
Assumption 3 (Convexity). each Js
The function Ji Rd + R is convex for all E [m]. This means that for
0 0 f((z) f((v) (V)).,, v) ER.
(5)
Assumption 4 (Smoothness). Function Rd is Li-smooth, La > for all i [n]. This means that for each fü.
- filv) (V)l))((ff(l)) (Vf((), S La || v11 Ve.yER".
(6)
We will use Lmax to denote max,e[m] Li-
is important to note that the smoothness assumption here is not necessary to obtain a convergence result, see Appendix F.2 for the detail. We introduce this assumption highlight how the optimal extrapolation parameter depends on smoothness it is present. The following strong convexity assumption is introduced that, if adopted, enables to achieve better results.
Assumption (Strong convexity). The function f is p-strongly convex, μ > 0. That is
f(x) f(v) (Vf(), y) > 4
vll Ve,y E Rd.
We first present our algorithm FedExProx as Algorithm 1. In the subsequent sections, we first present the theory in the stochastic setting for FedExProx with fixed extrapolation parameter in Section 3. Then we proceed to adaptive versions of our algorithm which eliminates the dependence on the unknown smoothness constant in Section 4.

---

3 Constant extrapolation
In order demonstrate the convergence result of our algorithm in the stochastic setting, we T-nice sampling as the way of selecting clients for partial participation. This refers to that in each uniformly at random from all subsets of size iteration, the server samples set Sy C ,I, n
We want to emphasize that the sampling strategy here is merely an example, it is possible to use other client sampling strategies.
Theorem 1. Suppose Assumption I (Diifferentiaiility) Assumption 2 (Interpolation regime), Assump- tion 3 (Convexity) and Assumption 4 (Smoothness) hold. If we use a fixed extrapolation parameter and any step size 0 7 < +00, then the average iterate of Algorithm I ag = =a (0: (im..
satisfies
|00-1.11
[[ - inf S C' (7,
K
where K is the number of iteration, 3x is sampled uniformly random from the first K iterates (2o.T1, (1,............................. Ik- C (K- T,) is defined as
1 1 TImax C(v,T,a) == acy x7,,))
and Lyg : 11 T- L Lmmax
+ n( - Lyy
T(n - 1) + y L
7(n 11
where Lmax= - max; Ly Ly is the smoothness constant of M (x) :- 1. E Mj, (x). Ie fix
and r the optimal constant extrapolation parameter is given by ax, == 7bs, > 1, which results in the following convergence guarantee:
E [f - inf S C(y, (, 120 x.ll
Loysr (1 + yLmaa) | 12 K
The proof of this theorem relies on the reformulation of the update rule in (7), using the identity VM,, j. (ar) = (ar ,r.x.,, (x)) given in Lemma 2, which holds for any x E Rdd, into the following
form:
rk+1
Ork Y
VM7, (Ix).
8
LE
We can then apply our modified theory for SGD given in Theorem 3, which is adapted from Gower et al. [2019], to obtain function value suboptimality in terms of M (z). The results are then translated back to function value suboptimality in terms of f. Note that (8) unveils the connection between the step size of gradient type methods and extrapolation parameter in our case.
Remark 1. Theorem 1 provides convergence guarantee for Algorithu 1 in the convex case. If in addition, we assume Assumption 5 (Strong convexity) holds, the rate can be improved and we obtain linear convergence. See Corollary 1 for the details.
Remark 2. Theorem indicates convergence for any O < < +30. Indeed, as it is proved by Lemma 7, we have C (7, T. x.,7) = L (1 + Lmmax) I Lmax holds for any 0 < < +oc In cases where there exists at least one L; < L we have (ÿ, a, Lmaax.
One may question the necessity of the interpolation regime assumption. This assumption crucial our analysis. Besides allowing us to revisit the convex feasibility problem setting, also guarantees that M (z) has the same set of minimizers as f(x) as illustrated by Lemma 8. It also allows us to improve the upper bound on the step size by a factor of the SGD theory, which is demonstrated in Theorem 3 in the Appendix.
Remark 3.
Remark 4. From the reformulation presented in (8), we see the best extrapolation parameter is obtained when ogyy y is the best step size for SGD running on global objective M7 (x), Since the best step size is affected by the smoothness and the minibatch size, so is the best extrapolation parameter.
We can also compare our algorithm with FedProx in the convex overparameterizee regime.
Remark 5. Our algorithm includes FedProz as special case when a 1. To recover its result, Ly
we simply plug in 1, the resulting condition number is C(y, T, = -(2-qImn.)) ). Compared to FedProz, Algorithum 1 with the same y > demonstratte superior performance, with the acceleration factor being quantified by
C(y,7,1) 22 2+ L + y T'max 2 4.
C(y,T,ay;
7

---

See Lemma 14 for the proof. This suggests that the approach of the server averaging all iterates following local computation suboptimal.
In the following paragraphs, study some special cases,
Full participation case
For the full participation case (T = -i), using definition from Theorem
1 - >, 71oy
Ly,n = Lep C' (1, Ly (1 + XLmax) Imax-
yn
In this case, we can compare our method with FedExP in the convex overparameterized setting. Remark 6. Assume the conditions in Theorem 1 hold, the worst case iteration complexity of FedEaP is given by O
( (Lm while for Algorithm it is O (C(),)a),)) As suggested by Lemma 7, Algorithm has better iteration complexity (C (7, n, a, Lmax) whenever there exists Li + Lmax for some i E [m], and the acceleration could reach up to factor of as suggested by Example In general, the speed-up up in the worst case quantified by
Limax
-
Li
Linax
Limax
Li
S
S 12. 1
1 + +Lmmaxxx
1 7b:
Single client case
For the single client case (T = 1), using definition from Theorem
1 Immax
Liaxx C (1,
02 =1
>1,
Ty,,1
Tmmax-
1 TLmax
Remark 7. Compared with full and partial client participation, the following relations hold for any E [n].
(1,,n, C (T, C (7,1, and On111 S Onyz S 00yyas Vr E [m]
Since the iteration complexity of FedEmProw given by O (Chi)aan))) ) the above inequalities tell us a larger client minibatch size leads to larger extrapolation and better iteration complexity.
Specifically, Lemma suggests the improvement over the single client case could be as much as factor of n (C (,, 1 (r, ())1) as suggested by Example 1.
4 Adaptive extrapolation
Observe that Theorem 1, in order to determine the optimal extrapolation, we require the knowledge of Lyyr, which is typically unknown. Although theoretically it suggests that simply averaging the iterates may result in suboptimal performance, in practice, this implication is less significant. To address this issue, introduced
variants of FedExProx, based on gradient diversity and stochastic Polyak step size, given their relation to the extrapolation parameter in our cases.
Theorem Suppose Assumption I (Differentiabilitty) Assumption 2 (Interpolation regime), Assump- tion 3 (Convexity) and Assumption (Smoothness) hold.
(i) (FedEnPron-Gradl If we are using Oik- = O'kGG, where
1 - I - prox,f; (xx))l 1,
O.k.
(10)
En (Ix prox (xx)) 11
then the iterates of Algorithm with satisfy
7
( Lrmax
I120 - x.
E [[((k)) inf
2 + TTmaxx
where Fx chosen randomly from the first K iterates (2o, 2r, IK- with probabilities Pk: o//K/ an.a
8

---

(ii) (FedEaProz-StoPS) If we are using ay = Ok,S, where.
1 - Mj, (#x) - inf MT.
2 1
(11)
| Ei= VMG. (ow)
2
then the iterates of Algorithm with T = n satisfy
[[[(xx)] - inf (7 Trmax)
|1:0 - x.
(12)
Ek OKK.5
where FKK chosen randomly from the first K iterates {ro, T1, IK- 1 with probabilities / 2 a..8-
Theorem 2 describes the convergence in the full participation setting, However, we can also extend it the stochastic setting by implementing stochastic version of these adaptive step size rules for gradient-based methods [Horváth et al., 2022, Loizou et al., 3. See Theorem 5 in the Appendix for the details.
Remark In fact, the adaptive rule based on gradient diversity can be improved by using T4 instead of 2 as the maximum of local smoothness constant of Moreau envelops, resulting in the extrapolation,
Lemax
ag =de= == 1+
1 A 2x fr (Ix)
(13)
Lmmax
E.= (Ix prox, (Ix))
One can obtain a slightly better convergence guarantee than the FedEaProw GraDS case in Theo- rem see Corollary 2 in the Appendix. However, the requires the knowledge of Liax in order to compute 14ylmax
Teemax
GraDS benefits from semiiadaptiiitty.. This refers to the fact that the algorithm converges for any choice of > Although a smaller hinders convergence, setting it to at least Temax convergence factor 2.
Remark 9. Note that, compared to classical gradient-based methods,
limits the worsening of the
Remark 10. Compared to FedExProz with the optimal constant extrapolation parameter, we gain "semi-adappiiity" here by using the gradient diversity based extrapolation. However, this results in losing the favorable dependence of convergence on T. and instead establishes a dependence on Imax-
Remark 11. For FedEaPron-StoPP as it is suggested by Lemma 20, the convergence depends on the favorable smoothness constant Loy, rather than on Lmmax- However, this comes at the price of having to know the minimum of each individual Moreau envelope.
For a detailed discussion of the adaptive variants of FedExProx, we refer the readers to Appendix F.5. Since one of our starting points the RPM by Necoara al. [2019] solve the convex feasibility problem with non-smooth local objectives, we have also adapted our method to non-smooth cases, as detailed in Theorem 4 in the Appendix. We also provided a discussion of our method in the non-interpolated setting and in the non-convex setting in Appendix F.
Finally, we support our findings with experiments, see Figure for a simple experiment confirming that FedExProx indeed has better iteration complexity than FedProx. For more details on the experiments, refer the readers Appendix in the Appendix. Notice that practice, each local proximity operator can be solved using different oracles. Clients may use GD or SGD to solve the local problem to a certain accuracy. The complexity this subroutine depends on the local stepsize. If y is large, the local problem becomes harder to solve because we aim to minimize the local objective itself. Conversely, if small, the problem easier since we do not stray from the current iterate. As the choice of subroutine affects local computation complexity, comparing it directly with FedExP becomes complicated. Therefore, we compare the iteration complexity (number of communication rounds) of the two algorithms, assuming efficient local computations carried out by the clients.

---

2 de 2,2933
T- 01..0. 2.0099
1 10, 2.3157
Flad 1062
Figure 1: Comparison of FedExProx and FedProx in terms of iteration complexity in the full participation setting. The notation y here denotes the local step size of the proximity operator and an,n is the corresponding optimal extrapolation parameter computed in (9) in the full participation case. In all cases, our proposed algorithm outperforms FedProx, suggesting that the practice of simply averaging the iterates is suboptimal.
Conclusion
5.1
Limitations
Our analysis of FedExProx serves as an initial step in adding extrapolation to FedProx, which currently relies on the suboptimal practice of the server merely averaging the iterates. While we discuss the behavior of our algorithm in non-interpolated and non-convex scenarios, our analysis only validates the effectiveness of extrapolation under the interpolation regime assumption.
5.2 Future Work
As we have just mentioned, extending our method and analysis beyond interpolation and convex regime is intriguing. In this case, new techniques may be needed for variance reduction. It is also interesting to investigate whether extrapolation can be applied together with client-specific personalization.
Acknowledgement
The research reported in this publication was supported by funding from King Abdullah University of Science and Technology (KAUST): i) KAUST Baseline Research Scheme, ii) Center of Excellence for Generative AI, under award number 5940, iii) SDAIA-KAUST Center of Excellence in Artificial Intelligence and Data Science. The work was done during Kirill Acharya's internship KAUST. Kirill Acharya was also affiliated with MIPT, ISP RAS, Russia. The work of Kirill Acharya was also supported by Grant App. No. 2 to Agreement No. 075-03-2024-214.
References
D. Alistarh, D. Grubic, J. Li, R. Tomioka, and M. Vojnovic. QSGD: Communication-cfficinn SGD via gradient quantization and encoding. Advances in Neural Information Processing Systems, 30, 2017.
S. Arora, S. Du, W. Hu, Z. Li, and R. Wang. Fine-grained analysis of optimization and generaliza- tion for overparameterized two-layer neural networks. In International Conference on Machine Learning, pages 322 332. PMLR, 2019.
H. Asi and J. C. Duchi. Stochastic (approximate) proximal point methods: Convergence, optimality, and adaptivity. SIAM Journal on Optimization, 29(3):2257-229) 2019.
H. Asi, Chadha, G. Cheng, and C. Duchi. Minibatch stochastic approximate proximal point methods. Advances in Neural Information Processing Systems,
2020.
10

---

H. H. Bauschke, P. L. Combettes, and S. G. Kruk. Extrapolation algorithm for affine-convex feasibility problems. Numerical Algorithms, 41:239 274, 2006.
A. Beck. First-order methods in optimization. SIAM, 2017.
D. P. Bertsekas. Incremental proximal methods for large scale convex optimization. Mathematical Programming, 129(2):163-99))) 2011.
Bianchi. Ergodic convergence of a stochastic proximal point algorithm. SIAM Journal on Optimization, (6(4))2235-22)) 2260, 2016.
A. Böhm and S. J. Wright. Variable smoothing for weakly convex composite functions. Journal of Optimization Theory and Applications, 188:628-649 2021.
S. Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends@ in Machine Learning, 8(3-4):231-357,22 2015.
Y. Censor, T. Elfving, and G. Herman. Averaging strings of sequential iterations for convex feasibility problems. In Studies in Computational Mathematics, volume 8, pages 101-113. Elsevier, 2001.
Y. Censor, W. Chen, P. L. Combettes, R. Davidi, and G. T. Herman. On the effectiveness of projec- tion methods for convex feasibility problems with linear inequality constraints. Computational Optimization and Applications, 51:1065-1088 1088, 2012.
P. Combettes. Convex set theoretic image recovery by extrapolated iterations of parallel subgradient projections. IEEE Transactions on Image Processing, 6(4):493-506 1997a.
P. L. Combettes. Hilbertian convex feasibility problem: Convergence of projection methods. Applied Mathematics and Optimization, 35(3):311-333 1997b.
L. Condat, D. Kitahara, A. Contreras, and A. Hirabayashi. Proximal splitting algorithms for convex optimization: A tour of recent advances, with new twists. SIAM Review, 65(2):375 435, 2023.
D. Davis and D. Drusvyatskiy. Stochastic basel minimization of weakly convex functions. SIAM Journal on Optimization, 29(1):207-239, 2019..
J. Duchi, E. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(7), 2011.
S. Ghadimi and G. Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic program- ming. SIAM Journal on Optimizaaiion, 23(4):2341-2366 2013.
E. Gorbunov, F. Hanzely, and P. Richtárik. unified theory of SGD: Variance reduction, sampling, quantization and coordinate descent. In International Conference on Artificial Intelligence and Statistics, pages 680-690. PMLR, 2020.
R. M. Gower, N. Loizou, X. Qian, A. Sailanbayev, E. Shulgin, and P. Richtárik. SGD: General analysis and improved rates. In International Conference on Machine Learning, pages 5200-5209. PMLR, 2019.
G. Hinton, N. Srivastava, and K. Swersky. Neural networks for machine learning lecture overview of bani- gradient descent.
S. Horváth, K. Mishchenko, and P. Richtárik. Adaptive learning rates for faster stochastic gradient methods. arXiv preprint arXiv:2208. 05287, 2022.
F. lutzeler and J. M. Hendrickx. A generic online acceleration scheme for optimization algorithms via relaxation and inertia. Optimization Methods and Software, 34(2):383 405, 2019.
D. Jhunjhunwala, Wang, and G. Joshi. FedExP: Speeding up federated averaging via extrapolation. In International Conference on Learning Representations, 2023.
A. Jourani, Thibault, and D. Zagrodny. Differential properties of the moreau envelope. Journal of Functional Analysis, 166(3):1185-12)))) 1238772--,,,,,11 2014.

---

S. Kaczmarz. Approximate solution of systems of linear equations. International Journal of Control, 1766):1269-1271, 1937.
A. Khaled and C. Jin. Faster federated optimization under second-order similarity. In The Eleventh International Conference on Learning Representations, 2022.
A. Khaled and P. Richtárik. Better theory for SGD in the nonconvex world. Transactions on Machine Learning Research, 2023.
S. Khirirat, H. R. Feyzmahdavian, and M. Johansson. Distributed learning with compressed gradients. arXiv preprint arXiv: 1806.0677, 2018.
D. Kingma and J. Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, San Diego, CA, USA, 2015.
J. Konečný, H. B. McMahan, F. X. Yu, P. Richtárik, A. T. Suresh, and D. Bacon. Federated learning: Strategics for improving communication efficiency. arXiv preprint arX1v:1610:05492, 8, 2016.
H. Li, A. Karagulyan, and Richtárik. Marina meets matrix stepsizes: Variance reduced distributed non-convex optimization. arXiv preprint arXiv:2310.0461 2023.
H. Li, A. Karagulyan, and P. Richtárik. Det-CGD: Compressed gradient descent with matrix stepsizes for non-convex optimization, International Conference on Learning Representations, 2024.
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 2020.
N. Loizou, S. Vaswani, I. H. Laradji, and S. Lacoste-Julien. Stochastic Polyak step-size for SGD: An adaptive learning rate for fast convergence. In International Conference on Artificial Intelligence and Statistics, pages 1306-1314. PMLR, 2021.
G. Malinovsky, K. Mishchenko, and P. Richtárik. Server-side stepsizes and sampling without replacement provably help in federated optimization. In Proceedings of the 4th International Workshop on Distributed Machine Learning. DistributedML. '23, page 85-104, 2023.
O. L. Mangasarian and M. V. Solodov. Backpropagation convergence via deterministic nommonotone perturbed minimization. Advances in Neural Information Processing Systems, 6, 1993.
B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. Arcas. Communication-effcccnn learning of deep networks from decentralized data. In Artificial Intelligence and Statistics, pages 1273-1282. PMLR, 2017.
K. Mishchenko, S. Hanzely, and P. Richtárik. Convergence of first-order algorithms for meta-learning with Moreau envelopes. arXiv preprint arXiv:2301. 066806, 2023.
A. Montanari and Y. Zhong. The interpolation phase transition in neural networks: Memorization and generalization under lazy training. The Annals of Statistics, 50(5):8116 2--28116677777 2022.
J.JJ. Moreau, Proximité et dualité dans un espace Hilbertien. Bulletin de la Société Mathématique de France, 93:273-299, 1965.
S. Motzkin and I. J. Schoenberg. The relaxation method for linear inequalities. Canadian Journal of Mathematics, 6:393-404, 1954.
E. Moulines and F. Bach. Non-asymptotic analysis of stochastic approximation algorithms for machine learning. Advances in Neural Information Processing Systems, 24, 2011.
I. Necoara, P. Richtárik, and A. Patrascu. Randomized projection methods for convex feasibility: Conditioning and convergence rates. SIAM Journal on Optimization, 29(4):2814-288 2019.
A. Nemirovski, A. Juditsky, G. Lan, and A. Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal on Optimization, 19(4):1574-1609, 2009.
A. Orvieto, S. Lacoste-Julien, and N. Loizou. Dynamics of SGD with stochastic Polyak stepsizes: Truly adaptive variants and convergence exact solution. Advances in Neural Information Processing Systems, 35-26943-26944, 2022.
12

---

N. Parikh, S. Boyd, et al. Proximal algorithms. Foundations and Trends@ in Optimization, 1(3): 127-239, 2014.
A. Patrascu and Necoara. Nonasymptotic convergence of stochastic proximal point methods for constrained convex optimization. Journal of Machine Learning Research, 18(198):1-422 2018.
G. Pierra. Decomposition through formalization in product space. Mathematical Programming. 28: 96-115, 1984.
C. Planiden and X. Wang. Strongly convex functions, Moreau envelopes, and the generic nature of convex functions with strong minimizers. SIAM Journal on Optimization, 26(2):1341-1364 2016.
C. Planiden and X. Wang. Proximal mappings and Moreau envelopes of single- convex piecewise cubic functions and multivariable gauge functions. Nonsmooth Optimization and Its Applications, pages 89-130, 2019.
L. Rechardson. The approximate arithmetical solution by finite difference of physical problems involving differential equations, with an application to the stresses in masonary dam. R. Soc. London Phil. Trans. A, 210:307-357, 1911.
Richtárik and M. Takác. Stochastic reformulations linear systems: algorithms and convergence theory. SIAM Journal on Matrix Analysis and Applications, 41(2):487-524, 2020.
H. Robbins and S. Monro. A stochastic approximation method. The Annals of Mathematical Statistics, pages 400-407, 1951.
R. T. Rockafellar. Monotone operators and the proximal point algorithm. SIAM Journal on Control and Optimization, 14(5)88778898, 1976.
E. Ryu and S. Boyd. Stochastic proximal iteration: non-asymptotic improvement upon stochastic gradient descent. Author website, early draft, 2014.
A. Sadiev, G. Malinovsky, E. Gorbunov, I. Sokolov, A. Khaled, K. Burlachenko, and P. Richtárik. Federated optimization algorithms with random reshuffling and gradient compression. arXiv preprint ar arXiv:2206.0700 2022.
M. Safaryan, Hanzely, and P. Richtárik. Smoothness matrices beat smoothness constants: Bet- ter communication compression techniques for distributed optimization. Advances in Neural Information Processing Systems,
2021.
C. T Dinh, N. Tran, Jnd J. Nguyen. Personalized federated learning with Moreau envelopes. Advances in Neural Information Processing Systems, 33:21394-21480 2020.
J. Von Neumann. On rings of operators. reduction theory. Annals of Mathematics, pages 401-483 1949.
D. Yin, A. Pananjady, M. Lam, D. Papailiopoulos, K. Ramchandran, and P. Bartlett Gradient diversity: key ingredient for scalable distributed learning. In International Conference Artificial Intelligence and Statistics, pages 1998-2007. PMIR, 2018.
Y. Yu, X. Zheng, M. Marchetti Marchetti-Bowick, and E. Xing. Minimizing nonconvex non-separable functions. In Artificial Intelligence and Statisties, pages 1107-1115. PMLR, 2015.
X. Yuan and P. Li. On convergence FedProx: Local dissimilarity invariant bounds, non-smoothness and beyond. Advances in Neural Information Processing Systems, 35:10752-10765, 2022.
13

---

Contents
Notations
15
B Basic Facts
15
Properties of Moreau envelope
16
D Technical lemmas
18
Theory of SGD
F Additional analysis on FedExProx
19
F.1
FedExProx in the strongly convex case
19
F2
FedExProx in the non-smooth case
20
F3
Discussion on the non-interpolation case
21
F.4
Discussion on the non-convex case
21
F55 Additional notes on adaptive variants
23
F.6
Extension of adaptive variants into client partial participation (PP) setting
24
Missing proofs of theorems and corollaries
25
G.I Proof of Theorem
25
G.2 Proof of Theorem 2
27
G.3 Proof of Theorem
29
G.4 Proof of Theorem
31
G.5 Proof of Theorem 5
32
G.6 Proof of Corollary
35
G.7 Proof of Corollary 2
35
H Missing proofs of lemmas
36
H.1 Proof of Lemma
36
H.2 Proof of Lemma
37
H.3 Proof of Lemma
37
................................
H.4 Proof of Lemma
37
H.5 Proof of Lemma
37
H.6 Proof of Lemma 6
37
H.7 Proof of Lemma
38
H.8 Proof of Lemma
39
H.9 Proof of Lemma 9
39
H.10 Proof of Lemma 10
39
H.11 Proof of Lemma
H.12 Proof of Lemma
H.13 Proof of Lemma 13
41
14

---

H.14 Proof of Lemma 14
42
H.15 Proof of Lemma 15
43
H.16 Proof of Lemma
H.17 Proof of Lemma 17
43
H.18 Proof of Lemma 18
43
H.19 Proof of Lemma
43
H.20 Proof of Lemma 20
44
I
Experiments
44
1.1
Experiment settings
1.2
Large dimension regime
45
1.2.1
Comparison of FedExProx and FedProx
45
I,2.2
Comparison of FedExProx with different local step size
47
I,2.3
Comparison of FedExProx and its adaptive variants
48
A Notations
Throughout the paper, we use the notation |||| to denote the standard Euclidean norm defined on Rdd and () denote the standard Euclidean inner product. Given a differentiable function f _ + R. its gradient is denoted as fz)).. For convex function f _R ~ R. use (f to denote its subdifferential at x. We use the notation D (æ, y) to denote the Bregman divergence associated with a function f I2 + between x and y. The notation inf f is used to denote the minimum of function Rd + R. We use prox,, (z) to denote the proximity operator of function (4) to denote the corresponding Moreau Envelope. The proper functions. We denote the average of the Moreau envelope of each local objective f. by the notation M R + R. Specifically, we define M (x) n Σ M7 (z). Note that A17 (I) has an implicit dependence on y, smoothness constant is denoted by Ly. We say an extended real-valued function f Rd + IR U {+o)) is proper if there exists x IRd such that f(x) < +oc. We say an extended real-valued function f: Idd + R U {+oo} closed its epigraph is closed set. The following Table 3 summarizes the commonly used notations and quantities appeared in this paper.
Rd R with y notation D is used for the infimal convolution of
x E Rdd and M7
B
Basic Facts
Fact (First prox theorem). [Beck, 2017, Theorem 6.3] Let f IRd + R be proper, closed and fonnett Then prox, (x:) is a singleton for any x E Rdd
Fact 2 (Second prox theorem). [Beck, 2017, Theorem 6.39] Let : dd (+oo) be proper, closed and convex function. Then for any E IRd, the following three claims are equivalent:
(i) u = prox (x)
(ii) x u 0f()).
(ii) (x (x- u. u. f(y) fu) for any y E Rd.
Fact 3 (Bregman divergence). The Bregman divergence associated with function f between x, Rdd is defined as,
Dj (x.1) f(x) f(v) (V)))
(14)
f is convex, then for any v, y Rdd
Dy (x, 2 0.
(15)
15

---

Table 3: Summary of frequently used notations and quantities in this paper.
Notations
Explanation
n
The total number of clients. The dimension of the model. The model which belongs to R The total number of iterations. The model at &-th iteration.
d
x
K
mx
an
The extrapolation parameter iteration k. Each local objective function.
f((x)
The step size in the proximity operator.
f(x)
The global objective f.
prox, (x)
The proximity operator associated with and y > 0 at point R
M7, (x)
The Moreau envelope associated with fi and 0 point RR
M (x)
The average of M7, (æ).
Li
The smoothness constant of fi. The smoothness constant of The strong convexity constant of f- The smoothness constant of M7 f;
u
Liax
The maximum of all I4, for É [n]]
/max// (mma))
The maximum of smoothness constant of each M7,, E [n].
by
The smoothness constant of M
Lyg
The interpolation between the Ly and Amax/(-->hmaa) induced 7-nice sampling. The optimal extrapolation parameter FedExProx under 7-nice sampling.
The convergence rate of FedExProx with r-nice sampling in the convex case.
Ok.G
The gradient diversity extrapolation in the k-th itcration defined in Theorem 2.
an.s
The stochastic Polyak extrapolation in the k-th iteration defined in Theorem 2.
The improved gradient diversity based extrapolation used Corollary 2.
Dy (x,
The Bregman divergence associated with f between two points x, IRd The set of indices server sampled in the k-th iteration.
Sk
Or.kk.G
The gradient diversity based extrapolation the k-th iteration for SelExProx-GraDS k-th iteration RedExProx-StoPS-
0r,4,88
The stochastic Polyak based extrapolation
If f is convex, L-smooth and differentiable, the following inequalities hold for any x. y E Rd
Σ [Vf(x) Vf(w)l (z. y) + Dy (y. x) - L || v/1
1
[V v < 2D, (a. v) S || y11
(16)
Fact (Increasing function), Let f(x T where ? 0. Then fx) is monotone increasing when x > O.
C Properties of Moreau envelope
In this section, explore the properties of the Moreau envelope of individual functions fi, and the global objective M Ei M7, Before that, we present the definition of infimal convolution
Definition 3 (Infimal convolution). The infimal convolution of two proper functions f.g Rdd R U { is defined via the following formula
(fDg) (z) min {f(2) + -
ZERR
One key observation is that M7 can be viewed as the infimal convolution of the proper, closed and convex function f and the real-valued convex function 1
27 11-1122 This observation enables us to infer
the convexity and smoothness of the Moreau envelope from the properties of the original function. First, present two lemmas about basic properties of Moreau envelope.
16

---

Lemma 1 -eeal--alued Let f IR' [U(+))) be proper. closed and convex function. Then its Moreau envelope Mi for any y is real-valued function. particular, the following identity holds for x E IR according to the definition of Moreau envelope,
1 - prox,, - (z)]l²
Mj (a) = f (prox,, (x))
27
Lemma 2 (Differentiability of Moreau envelope). [Beck, 2017, Theorem 6.60] Let f R H R RU (+ be a proper, closed and convex function. Then its Moreau envelope M2 for any > is
3 -smooth, and for any E Idd, have
VMJ (1) ( (p-ooxx, (x))
We then focus on the relation between individual ft and f:" The following lemma suggests that the convexity of individual fi guarantees the convexity of M
Lemma (Convexity of Moreau envelope). [Beck, 2017, Theorem 6.55] Let f :Rd - U (+0o}
be a proper and convex function. Then M7 is convex function.
is also true that the smoothness of individual fi indicates the smoothness of MT;
Lemma 4 (Smoothness of Moreau envelope). Let f B + R be convex and L-smooth function. Then MJ
L L -sooohh.
I+71
One notable fact is that f; and Mf, have the same set of minimizers.
Lemma 5 (Minimizer equivalence). Let f Rd + R U {+00} } be a proper, closed and convex function. Then for any y > 0, and MJ has the same set of minimizers.
In addition, Mÿ is global lower bound of f-
Lemma 6 (Individual lower bound). Let f ; R2 + U be proper, closed and convex function. Then the Moreau envelope MJ satisfies M7 (z) S f(x) for all x E Rd.
Next, we focus the global objective M (z). The following lemma bounds its smoothness constant from both above and below.
Lemma (Global convexity and smoothness). Let each fo be proper, closed convex and L--smooth. Then M is convex and Ly, -moooth with
La 21
L 1
SL
n
n
As result of the above inequalities, we have the following inequality the condition number defined in Theorem which holds for any E [n].
Ly (1 <Imaax) C (1, C (7,
C (7,,1,0,,1) Temax-
When there exists at least one Li < Lmax. we have C (y,n, az < C (T,,,, < LLmax C (r, 1, a.,1). Even L; Limax holds for all [n], there are cases (See Example I in the proof.) that C (7, (,,)) 1C (7,1, - Lmax.
n
A key observation in this case is the generalization of Lemma 5 into the finite-sum setting under the interpolation regime.
Lemma 8 (Minimizer equivalence). If we let every f. Rd R UU(+))) be proper, closed and convex, then f(x 1 -it f((x) has the same set of minimizers and minimum as
1
M (z)
M7. (x),
n
we in the interpolation regime and 0 < co.
17

---

The following lemma generalizes Lemma 6 into the finite-sum sum setting.
Lemma (Global lower bound). Let each f: Rd R {+oo} be proper, closed and convex. Then the following inequality holds for any Rd and 0.
M (x) M7 (o) S f(x).
In addition, we assume we are in the interpolation regime, then M 7 and f have the same set of minimizers, for any I, in this set of minimizers, the following identity holds,
2 (I.) ) = Mj (I.) f(x.).
Besides the global lower bound provided above, there is also a relation between the function value suboptimality of M and f.
Lemma 10 (Suboptimality bound). Suppose Assumption I (Differentiabiliyy, 2 (Interpolation Regime), (Convexity) and 4 (Smoothness) hold, for any minimizer 2, of M (2), all x E Rd, the following inequality holds for each client objective,
Mf. (x) - M7. (x.) 2 1+ (Ji(z)- -f(xx.)
(17)
Furthermore, this suggests
M (x) M (Ix) 1 7Tmax
- Si(x)))
(18)
above function suboptimality bound the star strong convexity of M7 from the strong convexity of f.
direct consequence
Lemma 11. (Star strong convexity) Assume Assumption 1 (Differentiabilit)) Assumption 2 (Interpo- lation Regime), Assumption (Convexity), Assumptioo 4 (Smoothness) and Assumption 5 (Strong convexity) hold, then the convex function M (x) satisfies the following inequality,
M (z) M (xx)) 2 1+ L
P
xz 1=1
2
for any x E TRd and a minimizer v, of M (x)
The star strong convexity property of M allows us to improve the sublinear convergence in the convex regime linear convergence.
D Technical lemmas
Lemma 12. Let f: Rd + I be proper, closed and convex function. Then x is minimizer of f if and only x prox,, (x).
Lemma 13. Assume we are working with the finite-sum sum problem f , fo where each fi is convex and L4-smooth, f is convex and L-smooth. Then the smoothness of L satisfies
Σ LLSLS Σ L.
2
1=1
where both bounds are attainable.
Lemma 14. Assume that all the conditions mentioned Theorem hold, then the condition number T(),T,1) of FedProw and the condition number C (y, T, (,,,)) of the optimal constant extrapolation parameter a. Tbbr satisfy the following inequality,
22+ 0(7,5,1))222--
1(
Lmmax 4 VT E [n].
C T, 01y,7)
TTmax
Lemma 15. Assume that all the conditions mentioned in Theorem hold, then the following inequalities hold.
C (7 & (m, a,,)) S (7,1,04,), Vr e [n],
and
0,,1 a,,5 2 ay,, Vr E [n].
18

---

E
Theory of SGD
In order to prove main theorem, partly rely the theory of SGD. The following theorem the convergence of SGD with 7-nice sampling is adapted from Gower et al. [2019]. We introduce modifications the proof technique and tailor the theorem specifically the interpolation regime. In this context, the upper bound on the step size is increased by a factor of 2. We first formulate the algorithm as follows for completeness.
Algorithm 2 SGD with 7-nice sampling
1: Parameters: learning rate 1j 0, starting point E Rd minibbtch size E (1,2.....) ,n)
2: for == 0,
2,...... do
3: The server samples S& S (1,2.... ..... n) uniformly from all subsets of cardinality т
4: The server performs one gradient step
1 Σ (xx)) 1)
Fk+1 Tk
S.ESs
5: end for
Theorem 3. Assume Assumption (Differentiability), 2 (Interpolation regime), 3 (Convexity), 4 (Smoothness) hold. Additionally, assume J is L-smooth where L S Ei= L 3 If we are running SGD with T-nice sampling using step size that satisfies O < <n A, where
2
7 - Lmax + n(T-1)
n2 T
A= ==
and
Lmax = max Lis
7(n -
then the iterates of Algorithm 2 satisfy
1 Fx 11
E [[(xx)) inf
nA,)
K
where K is the total number of iterations, F is chosen uniformly at random from the first K iterates {zo, xK-)). If. additionally, assume the following property (which will refer to as star strong convexity" holds, then the iterates of Algorithmm satisfy
x.112] |0EK mock=x11]] & (1 (1-n(2 712 - ( (1) K
||200 x=.1?
F Additional analysis FedExProx
In this section, we provide some additional details on the analysis of FedExProx and its adaptive
variants.
F.1
FedExProx in the strongly convex case
The following corollary summarizes the convergence guarantee in the strongly convex case.
Corollary 1. Suppose the assumptions in Theorem hold, and assume in addition that Assumption 5 (Strong Convexity) holds, then we achieve linear convergence for the final iterate of Algorithm 1, which satisfies
R -1IK- x1. (1-07(2-or1)))
1|:0 x.112.
2 + LImax)
where the definition of L,,r is given in Theorem 1. Fixing the choice of " and the optimal extrapolation parameter that minimizes the convergence rate is given by ay,7 JLa, 1. which results in the following convergence in the strongly convex case:
[ 2=11 < 3(1
K
|1z0 xz x2.12
2L (1 + 7mmaa)
'This is justified by Lemma 13.
19

---

As one can observe, by additionally assuming I strong convexity of the original function f. we improve the sublinear convergence in the convex case into linear convergence.
F.2 FedExProx in the non-smooth case
Our analysis also adapts to the non-smooth cases. This is based on the observation that even if we only assume Assumption (differentiability) Assumption 2 (Interpolation Regime) and Assumption (Convexity) hold and do not have additional assumptions on smoothness, still each M7 is --smoohh because of Lemma 2. Thus, the theory of SGD in the convex smooth case still applies. However, there are some differences from the smooth case. For the sake of simplicity, we will mainly focus on the convex non-smooth case with a constant extrapolation parameter, the results in the strongly convex regime and with adaptive extrapolation can be obtained similarly as in the proof of Theorem 1 and Theorem 2.
Theorem 4. Assume Assumption (Differentiabilityy) 2 (Interpolation Regime) and 3 (Convexity)
hold. If we choose constant extrapolation parameter ag = satisfying
2
0
"Ln,
where L, is the smoothness constant of M (x)
Σ M7, (I), Lyp is given by
1 n(r - 1)
Ly T(n 1) 7
T(n ([n-1) Lyy
Then the iterates of Algorithm 1 ssatisfy
|1z0 - 1.1
MM (ox) - inf I a (2 - ,
K
where k is chosen uniformly from the first K iterates {oo, 21.... K--1), -1 is easy to see that the best a is given by
1 21,
0, 7
where the corresponding convergence is given by
7M (mx) inf 7M <
I[n 7(n (n- 1 12,)))
n1 T-
n(T 1)
| 1.1
K
Remark 12. Notice that in this case we recover the convergence result of RPM presented in Necoara et al. [2019] in the convex case. Indeed, if each fs(x) = Ix, (z), then we have
prox, (x) = 1Tx, (x), Vo E RRd,
and
MM (x)
- Ix (1) and 7 (x) 2 Σ 1 IT, (x)]
Since we are in the interpolation regime, inf yM 0, and the convergence result becomes
n- n(T - 1) ((n - 1) T(n - 1) 72.) 1 1.1 ||KK- 1xx (IK) -
| 1211
2
K
Notice that here Yh.y S 1 is the smoothness constant associated with each distance function || (x)]l. The difference in the coefficients on the left-hand side from the original results presented in Necoara et al. [2019] results from different sampling strategies employed.
key difference in the non-smooth setting is that extrapolation in some cases may not be beneficiary, as illustrated by the following remark.
Remark 13. In the non-smone non-smooth case, is possible that L 1, where the optimal ax, = 1, in this cuse, extrupolation will not generate any benefis. However, as it is mentioned by Necoara et al. [2019], there are many examples where T.o < 1 and extrapolation indeed accelerates the algorithm. This is different from the smooth case, where extrapolation always helps.
20

---

Remark 14. Since we do not assume smoothness. Lemma 10 no longer applies. Therefore. the convergence result is stated in terms of the function value suboptimality of Moreau envelope instead of the original objective J which is used in the smooth case.
Using similar approach, it is also possible to obtain convergence guarantee for FedExProx in the strongly convex non-smooth regime, assuming addition that M7 (4r) sy--trongly convex, where we recover the convergence result of RPM in Necoara et al. [2019] in cases where the smooth and linear regularity conditions are both satisfied. The following Table 4 confirms that our analysis of FedExProx recovers the theory of RPM special case.
Table 4: Comparison of iteration complexity of RPM from Necoara et al. [2019] obtained using our theory and the original theory. In both cases, the optimal extrapolation parameter is used. The notation (++) hidden. the error level reached by function value suboptimality for convex case, squared distance to the solution for strongly convex case.
Sctting
Original Theory
Our Theory
Convex smooth case
Poe -1 1 ?hi
11 1
Strongly convex smooth case
Lygz log [x00 P4
Lyiz Pry
log
[z0-2.
20
The smoothness here does not refer each being L--mmotth, but M being -my--mooth This corresponds the smooth regularity condition presented Necoara et al. [2019].
(2) Here the strongly convex setting meaning that the linear regularity condition in Necoara et al. [2019] is satisfied. In our theory, it refers to M (ar) being sty-slroggly convex with joy < Ly.
F33 Discussion on the non-interpolation case
For the non-interpolation regime cases, we assume that Assumption (Differentiability) Assump- tion 3 (Convexity) and Assumption (Smoothness) hold. The differences are listed as follows
(i) Although ft and M7, have the same set of minimizers, I and M does not necessarily have the same set of minimizers. This will lead to the convergence of FedExProx to the minimizer z² of M (z) instead of v, of f. As a result, we will only converge to neighborhood of the x, depending on the specific setting.
(ii) Since we are not in the interpolation regime, the upper bound on the step size of SGD with sampling is reduced by a factor of 2. Thus, the optimal extrapolation parameter af, in the non-interpolated cases is also halved, a' a.. As a result, it is possible that a', < 1. The same phenomenon also observed FedExP of Jhunjhunwala et al. [2023], where their heuristic in determining the extrapolation parameter adaptively is also reduced by factor of 2 in non overparameterized cases.
Observe that all of the above results in both smooth/non-smool interpolatednoon-inerpppo cases suggests that the practice of server simply averaging the iterates it obtained from local training is suboptimal.
F.4 Discussion on the non-convex case
In the non-convex case, we assume Assumption 1 (Differentiability) holds, and we need the following additional assumptions on f R2 + IR and f; 12d ::
Assumption (Lower boundedness). Function fo is lower bounded by inf fa.
Assumption 7 (Weak convexity). Function fi is p weakly convex. this means that f. + 2 I1-111 convex.
We have the following lemma under the above assumptions:
Lemma 16. [Böhm and Wright, 2021, Lemma 3.1] Let f be a proper, closed, p-weakly convex function and let y <, 1. Then the Moreau envelope M2 is continuously differentiable on Rdd with
MM (x) = (x prox,, (z))
21

---

- In addition, the Moreau envelope is max {2. 77 1 -mooth. We will thereby denote the smoothness constant Leg.p-
Indeed, if the stepsize y in this case is chosen properly such that 1/4 p. then it is straight forward to
the function within the proximity operator prox,,; given by f. + } |1-11 is strongly convex. Thus the proximity operator still results in singleton. Lemma 16 allows to again reformulate the original algorithm using the gradient of Moreau envelope. The only difference from the convex regime is that the Moreau envelope Mj, is not necessarily convex. The following lemmas illustrate the connection between Mf, and fi:
Lemma 17. [Yu et al., 2015, Proposition 7] Let " > O, f be a closed, proper function that lower bounded. Then M7 S S. inf M3 inf f, arg ming M7 (x) arg ming f(x) S
{x: x E prox,s {-)
Lemma 18. Let R + R be p-weakly convex with p > 0 and differentiable. If take 0 < 7 < p then M3, has the same set of stationary points as fi.
For the sake of simplicity, we will consider only the full participation case with a constant extrapola- tion parameter ag a. The following lemma describes the convergence of GD in the non-convex case, which is adapted from the theory of Khaled and Richtárik [2023].
Lemma 19. Assume function f is L-smooth and lower bounded. If we are running GD with a constant stepsize ii satisfying 0 1j < +. Then for any K 1, the iterates vy of GD satisfy
E [[VJ(CE))P] ((((z)) - -inf) f)
min oskSK-1
nn
Now we directly apply Lemma 19 our case,
1. Since each M7, is smooth, M is L,, -mooth with Ly S Lygo which result in the following bound on the extrapolation parameter
0 < a
Notice that in this case we have the following estimation of 7 y
1 - min
{ 1 1-
PP
This suggests that extrapolation may not be much beneficiary in the non-convex case.
2. The following convergence guarantee can be obtained.
= 2 2M (xo) inf MY)
min
E [IVM
OSkSK
ayK
Notice that by Lemma 17, know that M7. (20) S Si (xo). We also have inf M >
1 E inf M7 Jt = E inf fo since inf M7, = inf f; true for each client by
Lemma 17. Thus, we have
1
M (zo) inf < f(2o) inf f inf f
inf fc.
n
We can relax the above convergence guarantee and obtain
[VM < 2(f(xo) - inf f) + K
2 (inf
Ei= inf fi)
min OSKSK
E
K
The above convergence guarantee indicates that the algorithm converges to some stationary points of 7 (z) in the non-convex case.
22

---

3. In the non non-convex case, we did not assume anything similar to the interpolation regime in the convex case. As result, we did not know the relation between the set of stationary points of M (x)) and [(r), denoted )" and J/ respectively. However, if we assume, in addition, that each stationary point y E y" of M also a stationary point of each M then y/ also
a stationary point of f/ according to Lemma 18. Thus, Vf (v') A E f (v') = 0, which indicates y' E y. As a result, we have y C VV. This means that under this additional assumption, the algorithm converges a stationary point of f.
E5
Additional notes adaptive variants
Notes on gradient diversity variant. In general, the gradient diversity step size 77s used in SGD to solve the finite sum minimization problem
1 Σ f(())
min *ER"
f(x)
=
71 i=1
can be written
1 Ei= (f(x)) n
1J :
Lmax
| E 1i=1 [m()))] 2
where Tmmax is the maximum of local smoothness constants, In our case, since each local Moreau envelope 1492; -mmooth and 7 -mmoothh we can use both Lmax (here in Corollary if we know
L,
1+7Lmax
L and 1 (in original Theorem 2, if we do not know Lmaxx as the maximum of local smoothness. We present the convergence result of Algorithm with the following rule given in (13),
1. E | proxyf, (Ix))l 2
1 +
Gg
7Lmax
ER (Ik prox, (rx))]
Corollary 2. Suppose all the assumptions mentioned in Theorem 2 hold, if we are using (13) to determine d'sc each iteration for Algorithm with ni, then the iterates satisfy
E [[(x)]] - -iint < +
120 - 2
ΣRJ o
where FK is chosen randomly from the first K iterates {Io, I1, ... IK-1) with probabilities pk a/s o/../EKK K-1 o
Notice that compared to the case of FedExProx GraaDS in Theorem 2, the convergence rate given in Corollary is indeed better. This can be seen by comparing them directly, for FedExProx-GraDS. we have
LLmax
(+ Lmax
| - 111
E [| - inf f - 7mmax
Σ k= Ok.
and for Algorithm with a'eaG given in (13), we have
E [f(xx]] - finl < ( Linax
1 - s.a1
Σk= O'l.
(-+ Lemaxx
[r0 - x.1l
1 + t
Σ Σ 06.
Since
Lmax
1 7Lmax
Vy > 0,
1+7 1+7 <bmax
2 + TLmax
the convergence of Algorithm 1 in the full participation case with (13) given in Corollary 2 is indeed better than FedExProx GraDS. However, this adaptive rule is only practical when we have the knowledge of local smoothness.
"Note that It7l, for any 7 0.
23

---

Notes stochastic Polyak variant. In this paragraph, we further elaborate on the convergence of FedExProx- StooS.. We start by providing lower bound on the adaptive extrapolation parameter.
Lemma 20. Suppose that all assumptions mentioned in Theorem 2 hold, then the following inequali- ties hold for any x E E= and x, that is a minimizer of f.
1 DE (M], (x) Mj, (I.))
1
E VMG. (a))
2
Using the above lower bound, we can farther write the convergence of FedEaProz-StoPSee
|:0 11
B [f( - inf S 22L, (1 2
K
Observe that recover the favorable dependence of convergence the smoothness of M However, this comes at the price of having to know each M7, (x.) or, equivalently in the interpolation regime, knowing M (z..)
F.6 Extension of adaptive variants into client partial participation (PP) setting
In this subsection, we extend the adaptive variants of FedExProx into the stochastic setting. We will refer to them as SadExProx-GraDS- we consider that the server chooses the client using the r-nice sampling strategy we have intro- duced before in Algorithm 1. The following theorem summarizes the convergence guarantee of FedExProx GraDsSPP -PP and FedExProx -oooxx-StPPPPPPP the convex case. Its extension to the strongly convex case where we additionally assume Assumption (Strong convexity) is straight forward.
and, PedExProx-StoPS-- respectively. Specifically,
Theorem 5. Suppose Assumption I (Differentiaaiility) Assumption 2 (Interpolation regime), Assump- tion (Convexity) and Assumption (Smoothness) hold. Assume we are running FedEaProw with T-nice client sampling.
(i) TredEoProo-GraDS-ee If we are using 0₁ = =Kr,k, Sk), where
1 Eies, 22 -prox,,, (Ix²
ar,,,((%, Sk) 11 EiEs, (xx - prox, (xx))
(19)
Then the iterates of Algorithm satisfy
Lemax
( Temax
1 xz
[[[(a))] - inf
(20)
2+ 22 L
inf OT,K, K'
where K the total number of iteration, ūK samples uniformly at random from the first K iterates {zo, TK,1), inf a,, is defined as
inf inf O,, ZER SEM)
Or, (I,
satisfying
ar. Sa) 2 inf a, 2 1.
(i) (FedEnProz-StoPS-PP If we are using as = aT.,,s(T St), where
Or,,,s(k + E MM. (xx) - inf M7))
(21)
/ E. VMG. (aw)
Then the iterates of Algorithm 1 satisfy
Ir x. 11
[[f(*x)] inf f $
+ Tmmax
(22)
inf a
24

---

where K is the total number of iteration, b is sampled uniformly at random from the first iterates {og, 2I,...,IK-11 inf a, defined as
inf kr, = xER SE kr (x,
inf
satisfying
(T S&) 2 inf ar.ks 2
TTmmax
Remark 15. For FedEmProz GraDS-PP, different from the full participation setting, the denomina- for of the sublinear term on the right-hand side of (20) is replaced by K .inf K.G.GGG
(i) In the single client case (T 1), have
011,, inf in
(ii) In the partial participation case (1 < n)) is possible that inf a,,K,G > 1,
resulting in acceleration compared to single client case.
(ii) For the full participation case (T = n)) we have Ol Onn,k.GG
and
O&,G 2 K inf An.K.G
thus the convergence guarantee here is relaxed version of that presented in Theorem 2.
similar discussion also applies FedExProx-StooS-- in the client partial participation setting. Remark 16. For FedEmProw StoPS-PP, different from the full participation setting, the denomina- tor of the sublinear term on the right-hand side of (22) is replaced by K inf k,
(i) In the single client case (T = = we huve
1 (( 01,ks inf 01, 2
LLmax
(ii) In the partial participation case (1 < n), is possible that
inf Ot,,k, >
2 1 L
resulting in acceleration compared to single client case.
(iii) For the full participation case (T = n1) we have
0g.5 Kn,
and
a 2 inf an.k.S:
k=0
thus the convergence guarantee here is relaxed version of that presented Theorem 2.
The following Table summarizes the convergence of new algorithms and their variants appeared in paper.
G Missing proofs of theorems and corollaries
G.1 Proof Theorem
The proof of this theorem can be divided into three parts.
25

---

Table 5: Summary of convergence of new algorithms appeared in our paper in the convex setting. The O (-) notation is hidden for all complexities in this table. For convergence in the full client participation case, results of Theorem I and Theorem 2 are used where the relevant notations are defined. For convergence in the partial participation, the results of Theorem are used.
Method
Full Participation
Partial Participation
Single Client
FedExProx
Ly(ttymmxx/K
Lmax/K
FedExProx-GraDS
(147hamaxl/y EK-I ung
(147hmax// (1
PcdExPox-SStoS fmnkkmee/y:tttt ok,8
(1+ (onnkmax)//(tte (> (1 ( (x
Step 1: Reformulate the algorithm using Moreau envelope. any
We know from Lemma that for
Rd
VMG. (x) - (x-- prox,s, (x))
Using the above identity, we can rewrite the update rule given in (7) in the following form,
1 Σ T iES%
Fk+1 26 0x7
VMJ. (rt).
(23)
The above reformulation suggests that running FedExProx with 7-nice sampling strategy is equivalent to running SGD with 7-nice sampling to the global objective M (a:) n Σ M (ar) with step size au Now, it seems natural apply the theory of SGD adapted in Theorem 3. However, before proceeding, we list the properties we know about the global objective M and each local objective M7.-
1. Each M7, (x) is convex. This is consequence of a direct application of Lemma to each fi. Since M is the average of convex functions Mf. we conclude that M7 (x) is also convex.
2. Each M7, (1)) HH,,, smooth, where Lg is the smoothness constant fa. This proved by applying Lemma 4 to each f.. Drawing on Lemma 13 for justification, it is reasonable to assume M7 (x) is smooth with L, S b 1t7., -smooth.
3. Each M7, (x) has the same set of minimizers and minimum as fi. This result arises from applying Lemma each function fi.
4. Furthermore, Assumption 2 (Interpolation Regime) holds, M (x) and f(z) have the same set of minimizers and minimum. This is demonstrated in Lemma 8.
L. Step 2: Applying the theory of gradient type methods. Notice that here M7, (x) is 147L. - -mooth and convex, M (x) is convex and smooth. Furthermore, due to the assumption of interpolation regime, MM (x) and fzx have the same set of minimizers. Applying the theory of SGD with - sampling in this case, where
Li 71
n(T - 1)
A, Lygt 7(n
max sE(m]
+
Lay
1)
7(n
Notice that using Fact 4, know that
Li y+
Fact
Linxx
max
KE [x]
1 1 T
thus Ly can be simplified and written
n T
I Txnax
n(( Lag
Lg
(n-1)
1 - 7L T
7(n 1)
26

---

where Lmax max; Lg. setting,
We obtain the following result given that 0 < ay <
__ in the convex Lyp
Theurem 3
[o x. x.11
B [M (Ix) M (Ix))
(r 122-0yly.
where ž is sampled uniformly at random from the first K iterates fro, 21,...; HK KK-1). However, the convergence mentioned pertains to M2 (z)). Given our objective is to solve (1), it is necessary to reinterpret this outcome in terms of f.
Step 3: Translate the result into function values of f. This step is only needed in the convex setting. We use the lower bound in Lemma 10,
(18)
M (Tx) - M (x+) 2
(J(īk) - ((x))
+ ILmax
to obtain the following result
1 1 I
|100 x211?
E [f -
ar (2 ag
K
Observe that we have
C,, 1 Tmmax
ary (2 - L,,
and its numerator does not depend on a. we fix the choice of y and т. then the denominator is maximized when ory,,, 1. This yields the optimal constant extrapolation parameter an,r
_1 and the following convergence corresponding to it
-00 x.
B [[((k)]] f(x.) Laysr (1 + mmmaa)
K
Finally, notice that
Lemma 1 74e
yhy S
< 1,
7
1=1
for any y > O. This suggests that,
y yTa
-
7mma
n(T - 1)
T(n 1) 7) TTr
7(n 1) 1 L
n - 7
n(( 1)
+ 7(n 1) =1
7(n 1)
which in turn tells us a,,7 T,,, > This concludes the proof.
G.2 Proof of Theorem
We start with the following decomposition,
|
2=11 = Ia 0. VM (xx)
12,1 x. 11 2047 (V (æx) x I -x, - x.) + aft? | (x)]
(24)
Case 1: FedExProx-Gree For gradient diversity based axs we have
Ei= 7V (rw))l # E VMJ. (fx)
VM (Ix)]l
VMM (xx)
For the last term of (24),
VMG. (Ix) ||
o o2² ||VM
0 (x=)]]
n
2. =0.o7
1
V.7. (Ix) - VMJ. (ax))l
n
2
Σ
L
(DM1, (xx, + Dar, (z).
1+7L,
i=l
27

---

where the last inequality follows from the Ityb: obtain using Fact 4 that
L.
-moothhess of MJ, given in Lemma 4. We further
Fac o VV (xx))l
Liax
S ax,07²
(DMr (I), Dar (I., (x)) IR))
1+7
7m (DMr (IK, 2x) + Dar (In)
= -R. 1 7/omax
(25)
For the second term of (24), we have
Y (VM (zx) 2h - xx) 20 (VM (Rx) VM (ax)) , - In)
-0 (Dar (I), + DM (z+;
(26)
Plugging (26) and (25) into (24), we have
[2 a 2211 ax.07 (2 1 1 //maxx
Immax
(D (IR, + Darr (Ix, Ix))
Notice that know that
(15) DM (I), (14) M (Ix) - -M (x.), DMr (I,,
2 0.
As result, we obtain
Ik x. ag.07 (2 ILmax | 2.11² <
(M7 (xx) M (x.)).
1 1 1
Summing up the above recursion for k = 0, 1, ..... K - 1, we notice that many of them will telescope and M (x.) inf M due to interpolation regime as it is proved by Lemma 8. Thus, we obtain
7 (2
7 Temax
ani (M" (xx) - inf M s ||000 x=11²
1 ybmax
0=
Denote Pk ax.a/Ek-1 an for 0, 0,1, K probabilities pı from the first K iterates {2o, 21, ,....., IK then we can further write the above recursion as
If we pick ãK randomly according to
E M (x*)] inf M S + 7/max
1 | x.11 ?
+ LLmax
2k=o On.
Utilizing Lemma 10, we further obtain,
E [ff(E inf f 1 + 7Lmax
(++ Lmax
11a0 111
LLmax
Σ
K-
k= Q3.G
The above inequality indicates convergence. Indeed, by convexity of standard Euclidean norm, we have
E (It prox,f; (Ix
Gk 2
=.
E (IK- prox, (xx)) ||
This tells that
OKK. 2 K.
0=
Case 2: RedExProx-Stort StoPS For stochastic Polyak step size based ag,s, since we are in the interpolation regime, by Lemma 9, we have
1 n inf Mn.
M (x+) inf M"
Σ n
As a result,
1 Σ" [M7 (Ix) - inf
M (xx) M (x.) 71 VM (2x)
0k.S
En VMG. (mx)
28

---

We have for the last term of (24),
s |M (rx))l² 54 (Ix) M2 (I.)).
(27)
For the second term of (24), we have
SY (VM (zx) 2h - n.) sa (VM (st), 2x - -2a)
(5)
s (M7 (I.) - -M (xx))
s7 (M" (xx) - M (x.)).
(28)
where the inequality due to the convexity of M Plugging (28) and (27) into (24), we obtain
| - 2=11 & 100 J. 2 OKk.s"Y (M (Ix)) - M (x+)).
Summing up the above recursion for 0, 0......, K 1, we notice that many of them will telescope. Thus, we obtain
K
ou,s (M (xx) inf MT) S | - -x.1
Denote Pk a ox.s for k = 0, 1.... K - 1. If we sample I² randomly according
to probabilities ps from the first K iterates ( 21, JK IK-1), J. we can further write the above recursion
E M (**) inf M - ||:0 - xx
Ek=0 04k5
Utilizing the local bound in Lemma 10, we further obtain,
[S(**)] inf 22) [ Lmmx
||a0 x=ll
(29)
k =k=o 0x,8
Notice that the above inequality indeed indicates convergence, since
K
K-1
M (Ix) - M (x.).
1
OK.S
2
VV (xx))l
24L4
k=0
k=
where the inequality follows from Lemma 20. The above upper bounds allow us to further write the convergence (29) as
- - x. 11?
E
[[(**)] inf S 2L, (1 + 2 Lmaxx
K
This concludes the proof.
G.3 Proof of Theorem
We start from the decomposition,
+ {
| 11 | = 2n 24
v((())
Σ
vi(ax))
F.,
iES.
iES.
where Sp is the set sampled at iteration k: Taking expectation conditioned on Tke we have
Es. ||akt1 - =
|
| - x. 112 2n (ox x., (f(z)) + Es,
Σ
Vilcon)
IESt
We can write the second inner product term
(Ix F., vf(zx) (f(x,)) Dy (xp, + Dy (x., 2K))
(30)
29

---

where Dy (I), 2x) denotes the Bregman divergence associated with f between zy and x,. For the last squared norm term, we first define the indicator random variable X as
when E Sk. {1,
XA.i
0,
when i $ Sk.
Since we are in the interpolation regime, we have
27
2
Ess
= V)((z))
Esp.
Xa. (Vf((a))) v.f((z)))
iESt
Denote 0R
Vii((rk) v((())))
Es.
Xkii (fr(())) V(((o+))
Es
X RR
1=1
1 Est
Σ X ||ax,1|? +
Σ
Xa.iii (ax, akd)
72
1si4isn
n
Bis [xH.]] ||ax
Σ
Esg (xa,Xaa (aree ans)
ISijisn
| 11
T
Σ ak!i
||
nT -1
1=1
i=1
1
+ n(r n(T-1) 1) 7(n
(31)
T (nn
1)
n
For the first term above in (31), due the smoothness and convexity of each f., we have
n
(Dz, (x)) + Ds. (MRR
n
S Lmmax
Σ
(D). (I., + Dj. (Ix, x..)
Lmax (Dy (x., xk) + Dy (Ix)
where the first incquality is obtained as a result of Fact 3. For the second term, we have due to the smoothness and convexity of f,
1
1
(v(((x)) v.((z.))
n
S LD, (I,, + Dy (xI,,
where the inequality is obtained using Fact 3. Combining the above two inequalities and plugging them into (31), we obtain
7
Lmax + nl[ - 1) 7(n - 1)
EsN.
V((TT))
1. (Ds (I,, (wr, 2k) Dy (xx, xx)))
T(n 117-1 1)
VESt
(32)
30

---

Notice that we already defined A, as
12 -
max + n(T n(r-1) -)
A,
T(n 1] L.
7(n -
Combining (30) and (32), we have
Es. [ - 1=112] If we require 0 <) A." have 17(2 nd,) O.
I 11 - 22,11
- 1/1211177 (Dy (x) + Dy (x),
Convex regime. remains notice that D, (z)) + Dy (Ir,, 2R) 2 Dy (10k) = [(ak) f(x.) 2 0, and we have
Es, [1 - x=11?] s || x=11" 1122 - n),) (f(zz) f(xx))
Taking expectation again and using tower property, we get
E | - =-1l" s S ||a - 2=11² n)(2 ( (B inf f).
Unrolling this recurrence, we get
1 x2
E [f((x)]] inf
1/2 -1,))
K
where K the total number of iterations, Rx is selected uniformly random from the first K iterates {2o, Il,.., I
Star strongly convex regime. Due star strong convexity of f. further lower bound the Bregman divergence
Dj (Ix. = f(xx) - f(xx) 2 11
2 |100 12112.
and we have
Esp | x. x111²] (1 1/2 -n,,) -1d)))- ()) |1: x211?
Taking expectation again, using tower property get
E || 2.113 s (1 (1- 11- nAr)
14
E || 22113)
2
Unrolling the recurrence, get
E | x0113] S [11==(212 - (-(=,n,)) 5) K
||z0 x1.
This concludes the proof.
G.4 Proof of Theorem
Since each fa is proper, closed and convex, by Lemma we know that each M7 is -smooth. Therefore, it is reasonable to assume that M = 1 Σn i== Mý, is T.,-shooth, with L. < 7' Theorem in this case, we obtain,
Applying
Thearem S
|1000 1.11
M (ōx) inf MY
any (2 - Lg,r
where ã x is chosen uniformly at random from the first iterates {xo, LT, IK - and
n-T
n(T - 1)
L
7 7 7(n - Lay
31

---

Multiplying both sides by 7, we obtain
200 - v.
YM (ēk) inf yM 2
(2 - (7
K
It is easy to that the coefficient the right-hand side is minimized when "/, convergence is given by
and the
+ 1a(T= -1) 1) =(n -11 1) -3l,) | x.11 11
n-T
PM (fx) inf y
7(n-1)
K
Notice that Ty 3. As result,
1 - 1. a, 701y
G.S Proof of Theorem
Case of FedExProx-GraDS-PP We start with the following identity
| - x. = |100 - x2112 -20,,
=
VMJ
(xx) I - 2x
7
iES
+ a N. 72
Σ
VMJ. (ars)
(33)
1ES
For the last term, we have
a. 72
2
Σ Σ VMG (xx)
- 72
Σ
VM7 f (rx) |
SeSu
seSa
= 92
Σ
VMJ (xx) V.G. (z.) (a.)
where the last step is due to the assumption that we are the interpolation regime. Using Fact 3, we can further upper bound the above expression,
1
VMG (Ir))
iESt
Li
(DMr, (ox, + Dasy, (ser
1 + TTa
SESt
Lumax
(Dary, (IK) + Me, (k.,
I _a, . 1 7Lmax
(34)
FES.
where the last inequality due to Fact 4. Now we look at the second term in Equation (33).
1
-20 .7
VM7
(IR),
--20r,k,
Σ
VMM (xx) M7, (xx)) DK x+
ieSt
=-20y,k,6
Σ
(DMr, (IR, + DM (I (2x;
(35)
VES.
32

---

Plugging (34) and (35) into (33), we obtain,
[6++ - x.1
$ | - =x11=² OT,,K, 77
ILmax
Σ
(DMy, (w)s + Dary (Ix, Fk.)
+ + ITuaax
SESx
2 2 7/maxx
$ | x+11= - Or,KK. '9
My (xx) My, (oa))
(36)
1 + Lmmxx
where the last inequality is due to
(15) Dar (Tk, v.) (14) M7 (Ix) - MM (z.), and DM7, (I.,
0.
Now we want to lower bound k,, notice that it can be viewed as function of the iterate r and the sampled set S. Therefore, use the notation
inf aT,
inf .ERV,SC (S Or,ko S).
As result, we have
Kr,k,G inf ar,kG 2
where the second inequality comes from the convexity of standard Euclidean norm. Plugging this lower bound into (36), we obtain
| x=112
S | x=1l= inf 0,, 7
2 ( 7
1 Σ
(M7 M7, (x+)
1 1 Iumax
SES
Taking expectation conditioned on 2t, we have
Es. || xx11²]]
2 + 7 TTmax
1 Σ M1.. (axx) - M7. (xx+)
< |, x2112 inf A. ?
1 Lmmxx
i=1
2 7
|2 xx, inf a
(M (xx) inf M)
I >L
where the last identity is due to the fact that are in the interpolation regime. Using Lemma 10, have
Es, || ==.1
2 + 7 Lmax
| x2112 inf a. 7
(f(xx) inf f).
1 Lunax
1 Linax
Taking expectation again and using tower property, obtain
|| - 13
LLmx < R ||08 x. 1111 inf o. y + y ILmax
[f((xx) inf f]
1 TLmax
Following the same step as Theorem 1, we can unroll the above recurrence and obtain
E6 [f((k)]] inf (1 11+7 7 ) (++ +
Lmsax)
+ Lmmax
inf a. K
where K the total number of iterations, Fx is sampled uniformly at random from the first K -itterttes {2o, Zi..... I-
33

---

Case of FedExProx-StooS-PP We start with the following identity
|| x.4 x=11²= 1 - x=11²
2. 8 "9
VMG. f. (Ix) - Ix
1ESt
+ a T, 72.
Σ VES
VMG, (zx)
(37)
For the last term of Equation (37), we have
H. . 72.
1
Σ
VMG (Fx)
= KT 7 :
My. (ITK) - inf MT7.
T
7 rESt
LES#
= 0T,KK8 '7-
Σ
DMr, (ta))
(38)
7
rESt
While for the second term we have
-2.
VMG. (Iz), 2 I.
iES
-2, -
DM7 (I), + Dary. (k.,
iES
(15)
S -20,,, -
Σ
DM7, (x),
(39)
sESA
Plugging (38) and (39) into (37), we obtain
1 Σ (M. (M, (xx) inf (nn)
| 1=1 || x=1l kr 7
(40)
VESk
Now want to lower bound &,S. notice that it can be viewed function of the iterate and the sampled set S. Therefore, we use the notation
inf dr,KKS
inf (CRM
ar,ss Sx
As result, we have
ar,k, inf 0r,k,S-
L, Notice that since each M7, is 1471, -smooth, we conclude that the function 1 Eies. Mj, least Lemax -smooth Using the smoothness of the mentioned function and Fact lower bound on T4ry/mmax
inf Kr, is obvious,
1
inf Ofk,r.S 2.
1+
Timas It7
2
7: Lemax)
This means that we have
1
(1
Kr,KS 2 inf kT
2
2
7Lmax
Using the above lower bound in (40), we have
| x, 11 < I I.11² inf infar ar -
1
Σ (M.. (zx) (zx) -inf inf M7))
iinfMM))
ie
Taking expectation conditioned Tkk and noticing that we are in the interpolation regime, we obtain
Es, [|0 - ...1 < b 2=11² inf kr r (M7 (xx) - inf M).
Same M (æ), its smoothness constant can be much better.
34

---

Using Lemma 10, we have
Est ||a x.1l²] Lemma 10
|| x.1l² inf kr,k
(f(xx) inf f).
1 7
Now, following the exact same steps in the previous case of PedExProx-Graaee we result in
[[[(xx)] inf f () T.max
-0 1.
inf k, K"
where K is the total number of iterations, Tx is sampled uniformly at random from the first K-iterates
(To.T1, ,
G.6 Proof Corollary
If additionally we assume is p-strongly convex, then from Lemma 11, know it indicates the following star strong convexity M7 holds,
M (x) M (x.)
P
||
x. 11
1 7m
Thus, we apply Theorem 3 with r-nice sampling in the star strong convexity case, and obtain the following result:
E | xx1l²) Theorem (1 (1- 11-r7(-0r,,m,))) b,,,,,)) (1 + y
1 12112.
Since the convergence here is stated in terms of squared distance to the minimizer, do not need further transformation. Notice that the convergence rate in this case,
1 1 a2 agy 2 ?Lmax)
1. In case of a = Cly,, the convergence is given by is also minimized when a = axy = 7K,,t"
< (1- 2L, (1 + LLmxx)
K || xx11?
E
| 2211?]
This concludes the proof.
G.7 Proof of Corollary
Similar to the proof of Theorem 2, we start with the following identity
[xx+1 - .. 1 Ix V (#x) x.
||
2,1l" a (VM7 (xx) In xx)) + (ah.a)² ( (z)l². (41)
The extrapolation parameter can be rewritten as
VM7 (Ix) 2
1 + 7Lmax
Iomax
VM (xx)11
We have for the last term of (41),
(ak,,) ? / V (Ix)
n
= = . (7+ Trnax
VM (ra)
n
= (x+ Lunax
VM7% (xx) - VMyffr
T=1
< <oho0 (x+ Liax
Σ i=1
+L, (Dasz, (xR, + DMg, (x).
L
1
35

---

where the last inequality follows from the Itys for x. >> O, we further obtain
4, -moothhess of My. Utilizing the monotonicity of 1474,
(ok.a) IIM (rr))l1
Fast
S ak.07 (1+
Lunax
1
Σ
DMr, (x) + Dary. (x))
Trmaxx
+ Tmmax
= (7+ Lmaxx
Lunax
(Das (I), + DMr (kx,
1+ 7Lmmxx
= =d (DM7 (I)) + Dair (I1,,
(42)
For the second term of (41), we have
-20' (V (rr), Ih - v.a 20 (VM (Ix) Tt -Ir)
22 (VM" (ax) VM (Ir, , -2.--11 In)
-20 (Dar (x)) + DM- (R+;0k))
(43)
Plugging (43) and (42) into (41), we obtain
|a+++1 - 1=2 S ||:x x=11² o2 (D x.) + Dar- xx),
Notice that know that
(14) M (Ix) M2 (x.), DMr (xx, ok) 2 0.
(15)
DAN (xx,
As a result, we have
| - xx,1| S || 12 7 (M (xx) M (z.)).
Summing up the above recursion for k = 0.1..... ..... K - 1, we notice that many of them telescope, we obtain
K-1
Σ
o (M" (xx) - inf M S 11000 x=112
Denote Pk o/ a for k 0, 1, K If sample 3x randomly according
to probabilities ps from the first K iterates (2o, 01, IK x we can further write the above recursion
1 [r0 - x=112
E [ (Ιx)] inf M S 7 Σ or
Utilizing the local bound in Lemma 10, we further obtain,
[f(=*)) pint < (1) + Lmax
120 x211? Σ Ekー1 -1 kg
-
This concludes the proof.
H
Missing proofs of lemmas
H.1 Proof of Lemma
Notice that since f is proper, closed and convex, by Fact 1, prox,, (x) a singleton. We use the notation z(x) = prox.., (x). Using the definition of prox,, (I)), we see that
Mf (w) f((xx)
||=(x) x11
27
= f (prox,, (x)) + 27
proxx, (2) - -xl|2
Now, assume Mf (x) +00. We have for any = Rdd
1
(x:) (x) ((t)))
+oo = M
|=(x) x1 [(=)+ || x1|2
2-
27
which means that is also optimal, which contradicts the uniqueness =(x) = prox.s (2). This indicates that M7 (x) +00, thus, it is real-valued, which concludes the proof.
36

---

H.2 Proof of Lemma 2
Let fr be the convex conjugate f, using Corollary 6.56 in the book by Beck [2017], we have (M7)* f II-I1* We know that the convex conjugate of proper, closed and convex function is also proper closed and convex. As result, f + 2 11-11" is y-strongly convex. This indicates that (M?) is y-strongly convex, which implies M7 is smooth.. Notice that we have
5 (I) arg min ZER4
f(z)
1| x=| J
2/
by the definition of proximity operator. Using Theorem 5.30 from Beck [2017], we have
VMj (o) = (r - (x))
This completes the proof.
H.3 Proof of Lemma! 3
To prove this lemma, we use Theorem 2.19 in the book by Beck [2017]. From the key observation that My is the infimal convolution of the proper, convex function f and the real-valued convex function
11-11,
we deduce that M7 is convex. This completes the proof.
2y
H.4 Proof of Lemma 4
Let f* be the convex conjugate of f. From Corollary 6.56 book by Beck [2017], it holds that 11-11". Since f is L-smooth, we deduce that f is --strongly convex, and thus M7) is f. ---ronly convex. This suggests that (M7) is 1+yb -strongy convex, which in turn implies M7 is 147 1 smooth. This completes the proof.
Mf) = - +
H.S Proof Lemma
Notice that since M3 is convex and differentiable, the condition VM (x) 0 gives its set of minimizers. This optimality condition
be written exactly as prox, f (ar) according to Lemma 2. Using Lemma 12, we know this condition also gives the set minimizers of f. which suggests that and M7 have the same set of minimizers. Pick any 2. Rd that is minimizer of using Lemma 1, we have
inf M7 = MM7 (x.)
= s {prox,, (Ix)) + 24
11. prox,s (xx))
=xxx) inf f.
This completes the proof.
H.6 Proof of Lemma 6
For any x e Rd, we have
M3 (x) min
f(2)
||2 2=1
ZERd
2
< f(x) + 27
a| - 10
f(x).
This completes the proof.
37

---

H.7 Proof of Lemma! 7
From Lemma 3 and Lemma 4, we immediately obtain that each M7, is convex and L,
smooth.
1+y1,
This immediately suggests that M = - My. is convex and -moothh with
L
by s
14
i=1
Then by Lemma 13, we have
Li yاا
Lemma
Ly.
i=1
Combing the above two inequalities, have
L 1 SI.
L
n
1=
We then look the condition number defined in Theorem 1. It is easy verify that C n, ay,0) = Ly (1 + +Imax) and, C' (7, 1, 1,,11 = -maxx
As result,
C (,,,,, ay,a) = T (1 + 7mmax)
1 n
+ TTmax
n Σ L
1+
Tmmax C m),
Notice that we can write C therefore
as an interpolation between C (77 n, ar.m) and C (7, 1, (xy,))
L (1 + mmmx) I C Cr,, n, an,) S C C7, a7,7) s C7,, (,,1,,,,, 1,,1,) = Lmaxx
In cases where there exists least one Li < Lmmax we have
n
L 1 1
Σ
Lmaxx
71 1=1
1
which is true for all < < +00. Thus, C(y, < (T,,,, < Lmnax C , ay,1).
Now we give an example that when all L; Lmaxx still C (7, n, a.y.)) 10 (1, n Lmax.
1
Example 1. Consider the setting where f: RR + IR is defined f((x) 2x, for some 0 0. 0
Here D; denotes the i-th coordinate of the vector x E Rd f Rd + IR is given by f(x) 2n | It is easy show that for each fi is convex, f-smooth function and the smoothness constant cannot be improved since
0
0
0
||
24
2
2
2
j
For fxx 2n Moreau envelope of fe(x))
0
||= apparently, it smooohh and convex. We have the following formula for the
Mf. (x) = 1 1+ 1+00 x
As expected, each one of them is convex and I+70 smooth. For MY (x)) it is given by n
M (x)
n Mf. (x) 2 n(1 + 70)
||0|
thus, we know it is convex and Iny m[1+50j smooth. In this case
T.max
0
L, (1+ yLmx))
(1 + 78) n,
m[1+70]
which is
1 C7,, 1, 1y,1))
Ly (1 + Limxx) C (7, ayon
-maxx
n
n
38

---

H.8 Proof of Lemma 8
By Lemma 5, we know that fa and M7, have the same set of minimizers and minimum. Denote the set of minimizers X2 since we in the interpolation regime, know that the set of minimizers of f is given by,
X
Xs + 0.
Now we prove that every x: in X is a minimizer of M = - MT. This true since x E X minimizes each Si. thus Mf, at the same time. The minimum is given by
1 Σ inf Mj - - Σ inf J 10-=in inf
inf M
We then prove that every x & X' not a minimizer of f. If X, there exists at least one set X; such that X,, Thus M f; (x) > inf MJ, This indicates that M (x) > inf M, which means, X is not minimizer of M.
H.9
Proof of Lemma 9
From Lemma 6, is clear that Mf is global lower bound of satisfying M7 (x) S f(x) for any E R2d and > 0. Notice that the definition of M indicates that
M (x)
M7. (x)
fi(z;) + || x0l2
min RRR
27
1
S min sER
ft(z) +
1| x1
24
i=1
min SER
[1(2) +
|| x012
27
M3 (z),
holds for any E Rd and 0. Combining the above result, we have M7 (1) S Mf S f(x) for any x E Rd and y > Notice that in Lemma we have already shown that M and f have the same set of minimizers and minimum in the interpolation regime. A direct application of Lemma 5 indicates that M7 and f have the same set of minimizers and minimum. Therefore, combining the above statement, we know that M M7 and ff have the same set of minimizers and minimum. Thus, for any x, belongs to the set of minimizers, we have
M (xx) = Mf (x.)) = f(x)).
This completes the proof.
H.10 Proof of Lemma 10
We start from noticing that according to Lemma 1, the following identity is true for Moreau envelope,
M7. (x) = fprox,, (x)) + || prox,s, (z)]]?
1
(44)
27
For the second squared norm term, we have the following inequality due to the smoothness of each f. and the fact that f.f (prax,,. (x)) = + xx prox, (z)).
x prox, (z)) = (I - prox,, (x), x - prox, fi (x))
= ( (x)), - -rox,, (z))
L1
2 (f((z) - fo (prox,,, (x))) 2
|| prox (z))
39

---

which leads to the following lower bound:
- prox,, (x)]|
(fi(z) Si (prox,s, (x))).
I 2
Plug in the above inequality into (44) and notice that inf M 1. n. E inf M2, 1, n Ei= inf fs. we obtain the following lower bound on M, (r),
M2 (x) inf M7. 2 fi (prox, f (x)) + 2+ y (Fi(x) - Se (prox,,i (x))) - -inf ft
2 + zIi (fi(z) - -if fi) +
2 + TT
(Fi(prox,,, (2)) - inf fi)
(45)
Now let us look at the term Si (prox,ys, (x)) inf f. Using again L--smoothness of Si, we have
Li
f((x) fi(prox, f. (1:) - (prox, v. (x)), x prox s, (x)) S 2
|| -roxx (x)]
Notice that x - pr f. (x) VViipr 7/. (r)) As a result, we have,
fi(x) - (proxx (z)) 2 l
V (prox (=)) S (prox (x)),
2
which is
fa()) - inf f: (7+ 24
fi (z))]] S fi (prox,Ji (x)) inf ft.
2
Using the interpolation regime assumption, we have
V Si (prox,,. (x))]| = [.. (prox,,, (x)) - [1(
< 2I, (prox,, (c),
2T, (Jo(prxxy,, (x)) - -inf in)ff.).
where the inequality is obtained using Fact 3. As a result, we obtain the following bound,
1 fs (prox, s. (x)) inf f. 2 1 L. + t (f((z) inf fi)
(1+71,)2 (Ji(z) - inf ft).
Plug the above lower bound into (45), we obtain
M7. (I:) inf Mj, 2 1 1 Lb (fi(z) - inf ft),
(46)
Notice that we have M (z) 1 Em= Mf (x). Since we are in the interpolation regime, from Lemma 9, we know that
Mf. (x.) = 1
inf M = M (x+) =
Σ
inf M7
i=1
1=1
and
f f= f(x.)
1
f((x.) Σ
m
inf f.
n
n1
We average (46) for each i E [r]] and obtain
MT (x:) inf M
(fi(x) inf fi)
L
1 Σ (f((z) - inf f.)
LLaxx
i=l
(f(x) - inf f).
+ Tymax
This concludes the proof.
40

---

H.11 Proof of Lemma 11
We start with picking any point 2 E Rd, since we are in the interpolation regime, according to Lemma 9, we have M (x.) fxx)). Applying Lemma 10, we get
1
M (x) M (x++) 2 ++ +ymmux (f(z) f(x+))
(47)
We know that from the μ-strong convexity of f, we have for any a E Rdd,
f(x) f(x.) (Vf(z.), x - x.) 2 1
| x.1l 11
2
Notice that since Vf((x)) 0, we have
4 1 - 1.12.
- f(x 2
(48)
2
Combining the above two inequalities (47) and (48), have
M (z) M (x.)
1
||
x=11?
mmma
2
This concludes the proof.
H.12 Proof of Lemma 12
Notice that Rd is minimizer of f and only 0 0f(z). This inclusion holds if and only
0 E 0 (z)), which can be rewritten as x E 0 (1r). By the equivalence of (i) and (ii) in Fact 2, the above condition is the same as x = prox,, (x).
H.13 Proof of Lemma 13
Since each ft is L--mmooth the following function convex for every E [n].
Li
|z||
fa (z).
2
Thus,
E Li
n n
||0||/
f((x),
also a convex function, which indicates f(x is also 1 n
Σi= smooth. This means that
L S
L 1=1
(49)
Now notice that the L-smoothness of f is equivalent to the following function being convex
nL
|
Σ fl(z).
Pick any j E [n], we have
nL 2
|
Σ
f(()) +
Σ M(x)
nL ||0|| - (z)
2
1=1
Since all functions are convex and the sum of convex functions is convex,
nL.
||0|| fz)),
2
convex, which indicates that fz( also nL-smooth. As result, for every [n]. have L 2 Lj. Summing up the inequality for every j E [n]]. we have
1
Σ4
(50)
41

---

Combining (49) and (50), we have
n2
n L.
Σ
L
Σ
Lx
n i=1
72
In order to demonstrate that both bounds are tight in the above inequality, consider cases where they are identities.
Consider the case that each function f((x) } L, ||~ is easy to see that f(z)
1 ( E Li) ||==1]² In this case L -i=1 Li, the upper bound is an identity.
(ii): Consider the case that each function fi()) -0. where 0 >> 0 is constant, 2; is the i-th coordinate of x E Rd. In this case f(x) = 1 1 It is easy to verify that this
2
case L, = 6, L = A 0. Thus n² Ei=1 Li = L, the lower bound is an identity.
This concludes the proof.
H.14 Proof of Lemma 14
From the definition C(y,T, T, 1) and C' (7/ T, 07,7), we know that
C(),T,1)
C C (7,T, T, a,,)) 7Lao (2- 7L
Now let = p we have the following bound on according to the definition of L,,r given in Theorem 1.
= 7 Tog
n1 7 Tmmx
+ n(T - 1) T(n -
I
T(n - 1)
+ Lmax
Notice that in Lemma have shown that
Lemma
L;
Ly
+ 7
and due Fact 4, we have
Ti
Tmmax
+ 7
1 + Lmaxx
As result,
11 -
Xmmax
+ n((T - 1) Lmaa
ALmax + xlumxx
1 5
< 1.
(nn 101] 1) 1 1 Tmmax
T(n - 1) 1 ++yImax
It is easy to show that g(t)
is monotone decreasing when t E [0, 1}, thus 1(2-1)
1
C(
C Cy, T
Lomax ItyLmax
(1
7 Lemax I+7 Lomis
2 7Lmax
+ Lymxx
AM.MM
where the last inequality is due to the AM-GM inequality. This concludes the proof.
42

---

H.15 Proof of Lemma 15
and we know that As suggested by Lemma 7, we have
C (7,n0,y,,)) S C (
S (1,
Vr e
Notice that 0,,r is given by
0, "Loe
71 T
Imax
n( -11 1)
Lag
+
Ly.
7(1 n)
1 L
7(n - 1)
From Lemma and Fact 4, we know that
Lemaaa
L yb
Fact
Liaxx
71
1 +mxxaax
Consequently, L,,r decreases as T increases. Therefore, Ogyr increases with the increase of T, as illustrated by the following inequality
0,,1 S 0x S 0y,,, Vr E [m].
This concludes the proof.
H.16 Proof of Lemma 16
refer the readers to the proof of Lemma 3.1 of Böhm and Wright [2021].
H.17
Proof of Lemma 17
We refer the readers to the proof of Proposition 7 of Yu al. [2015].
H.18 Proof of Lemma 18
Observe that since 0 < y < 1, we do have f 1/2 11-112 being strongly convex. This indicates that prox, is always singleton and therefore M7 is differentiable, suggested by Lemma 16. Notice that æ is stationary point of M7 and only if VMJ (x) = This equivalent to (x prox,, (x:) = 0, which is pr (x). In addition, prox,, (x) is equivalent
Vs(+) 1 (a-
which is vx(x)= = 0. Combining the above statements, we have vf(x) = if and only
VM (x) 0. This suggests that the two functions have the same set of stationary points.
H.19 Proof of Lemma 19
Apply Theorem 1 of Khaled and Richtárik [2023], notice that in this GD satisfy the expected smoothness assumption given in Assumption Khaled and Richtárik [2023] with = 0, = and C 0, we obtain that when the step size y satisfies
<<
LB L'
where L is the smoothness constant of f. the iterates of GD satisfy
2 (f(()) inf f)
min E 0
<
This completes the proof.
43

---

H.20 Proof of Lemma 20
Notice that we are in the interpolation regime, by Lemma 8, we know that of and M have the same set of minimizers and minimum. As result,
M (x.)
M7. (Ix) Lemmna f(x.).
(51)
From the above inequality, we obtain that
Ei=1 M7 (x) M7. (I.))
(51) M (x) - ' (x.))
7. VM (x:)]2
7 1. E. VM, (x)
Then by the smoothness of M and Fact 3, we have
VV (x) VM (x.))²
M (x) M (x.)
Fact 2h
7 VM (x:)11
7. VM (x)]
1
2
Thus, by combining the above inequalities, have
MM. (x) M7 (sx))
1
| VM7. (a)
2
Notice that from the definition of ak,s for FedExProx-StoPP we have
1 E 7 Ei VM7. (xx)
M7. (Ix) M7, (I+)
1
2
Therefore, using the above lower bound, it is straight forward to further relax (12) to
||00
1.1
E [[{(**)] - inf f 2L, (1 2 Lyxx)
K
This concludes the proof.
I Experiments
In this section, we describe the settings and results numerical experiments to demonstrate the effectiveness of our method.
I.1 Experiment settings
We consider the overparameterizee linear regression problem in the finite sum setting
-
f(x) - - f()))
min E
where d is the dimension of the problem, n is the total number of clients, each function fi has the following form
Mxx) A| 10:x-b,l" 1,11
where A; E xd b, E R": Here is the number of samples on each client. It is easy to see that for each function fi, we have
ff((z) A, A A,b. and V2f((x) A, A, 2 Oa.
44

---

Thus, it follows that
[ A, A. and v²f(2) 1 7
vf(x)= Σ 1=1
The problem therefore convex. Notice that one implicit assumption for the class of proximal point methods in practice is that the proximity operator can be computed efficiently. In the setting of linear regression, we have the following closed form formula for the proximity operator ,rox, which holds for any a E IRd
A. A. A. + +1)
prox, A (x)
(52)
Observe that in the linear regression problem, since we know the closed form expression of each fü and f. we know the corresponding smoothness constant
L = Amaa (A, A.).
Notice that from Lemma we have
1 || prox, (fi) (z)]l² M3. (x) (prox, (ft) + 24
Since we know prox, (fi) in closed form using (52), we also know each local Moreau envelope in closed form, and thus the same for M E Mf, As a result, we can deduce L,, for M7 In our experiments, we pick
Ex= n, so that we are in the interpolation regime. Each A, is generated randomly from uniform distribution between [0. 1). and the corresponding vector b, is also generated from the same uniform distribution. In order to find a minimizer T., we run gradient descent for sufficient amount of iterations. All the codes for the experiments are written in Python 3.11 with NumPy and SciPy package. The code was run on a machine with AMD Ryzen 9 5900HX Radeon Graphics @ 3.3 GHz and 8 cores 16 threads. For experiment in the small dimension regime, each algorithm considers here only takes seconds to finish. For larger experiments, depending on the specific implementation. the algorithms typically take a few minutes to half an hour to finish. For FedProx, FedExP and our method FedExProx in the full participation case, the algorithm for specific dataset is deterministic, while in case where client sampling is taken into account, the randomness of the algorithms comes from the specific sampling strategy used. Our code is publicly available at the following link: https: /tap::://anooyynn /aaonymoouuuu..o 4open. science/r Ftteee///eeeeeeeoeeeeeeeeeeeeeee F2622
1.2
Large dimension regime
In this section we provide the numerical experiments in the large dimension regime, where n, = 20 for each E [n). n 30, d = 900.
I.2.1 Comparison of FedExProx and FedProx
In this section, we compare the performance of FedProx with our method FedExProx in the full participation case and in the client partial participation case, demonstrating that the extrapolated counterpart outperforms FedProx in iteration complexity. Notice that here we are only concerned with iteration complexity, since the amount of computations is almost the same for the two algorithms. The only difference is that for FedExProx, instead of simply averaging the iterates obtained from each client, the server performs extrapolation. From Figure 2, it is easy to see that our proposed algorithm FedExProx outperforms FedProx, which provides numerical evidence for our thcoretical findings. Notably, in order to achieve the small level of function value sub-optimality, FedExProx typically requires only half the number of iterations needed by FedProx, which indicates a factor of 2 speed up in terms of iteration complexity. Another observation is that, a,,, is decreasing y increases, which suggests that when local step sizes are small, the practice of simply averaging the iterates is far from optimal.
We also compare the performance of the two algorithms in the client partial participation setting. As one can observe from Figure 3, FedExProx still outperforms FedProx in the client partial participation setting, which further corroborates our theoretical findings. Observe that at,, here increases becomes larger, which coincides with predictions in Remark 7.
45

---

G.D0R 422309
U.B. 223105
0.2, 20055
F(6) 104)
= Oo 20199
7.00 3 2.0550
7 10,0 2000770
2008
Pricttons
Paittoon
Figure 2: complexity in the full participation setting. For this experiment ? is picked from the set {0.0001, 0.001, 0.01, 0.1, I, 10). the (ty, indicates the optimal constant extrapolation parameter as defined Theorem 1. For each choice y, the two algorithms are run for 10000 iterations, respectively.
Comparison of convergence of FedExProx and FedProx in terms of iteration
0,000L 3,2211,
0,200L 3,22880,
0,2001 5.29328,
Flad 1162
Fba 1662
Bld 1662
2000
1000
2:08
600
17007
480
7 0.000 Tepe
0.03 To
(Tea Fiee
(162 Fle2
(Te2 Fie
Eacation
Earation
Baratioo
Figure 3: Comparison of convergence of FedExProx and FedProx in terms of iteration complexity the client partial participation setting. For this experiment y picked from the set (0.0001, 0.001). }. the client minibatch size т is chosen from {10, 15. 20} and the ar,,n indicates the optimal constant extrapolation parameter as defined in Theorem For each choice of y and r, the two algorithm are run for K 10000 iterations, respectively.
46

---

Figure 4: Comparison in terms of iteration complexity for FedExProx with different step sizes y chosen from {0.0001, 0.001, 0.01, 1, 10, 100} in the full participation setting. In the figure, we use FedExP with different iterations of local training sub-figures. The local step size for FedExP set to be the largest possible value Gthmax where Lmax max; Li.
{1, 5, 10} benchmark in the three
1:0 21005
1o: 21507
Figure 5: Comparison in terms of iteration complexity for FedExProx with different step sizes ? chosen from (0.0001, 0.0005, 0.01, 10) in the client partial participation case. Different client minibatch sizes are used, the minibatch size is chosen from (5, 10, 20).
I.2.2 Comparison of FedExProx with different local step size
In this section, we compare the performance in terms iteration complexity for FedExProx with different local step sizes. We also include FedExP as a reference. The local step size of FedExP is chosen be a where is the number of gradient descent iterations performed by each client for local training, Lmax = max; L,, where L, is the smoothness constant of fs-
As one can observe from Figure 4, for our proposed method FedExProx, the larger y is, the faster it will converge. However, 7 becomes larger, the improvement in iteration complexity becomes trivial at some point. Note that for different 7, the complexities required to compute the proximity operator locally varies and often larger 7 requires more computation than smaller y. Compared to FedExP with the best local step size 6tLmase FedExProx with a large enough ? is better in terms of iteration complexity. In the case where the computation proximity operator efficient, our method has better computation complexity as well. Notice that small y leads to slow down of our method, and we do not claim that the iteration complexity of FedExProx is always better than FedExP. However, it is provable that FedExProx indeed has better worst case iteration complexity. We want to emphasize a key difference between FedExP and our method is that we do not have any constraints on the local step size y, and our method converges for arbitrary local step size ? > while for FedExP, a misspecified step size could lead to divergence.
We also compare FedExProx with different step sizes in the client sampling case, see Figure 5. However, since there is no explicit convergence guarantee for FedExP in this case, we did not include FedExP in the plot.
In the client partial participation case, the same behavior of how our proposed algorithm FedExProx changes according to different local step sizes 7 is observed. A small leads to slow convergence of the algorithm, while for large 7, the convergence is improved. However, at some point, the improvement becomes trivial.
47

---

00025
T- 0.005
- 0.05
1344
200
R66
15 1162
(150
12
1: 22005
No 21005
Figure 6: Comparison of FedExProx, FedExProx-GraDS and FedExProx-Steee in terms of itera- tion complexity with different step sizes chosen from {0.0005, 0.0005, 0.05, 0.5, 1, 5) in the full participation setting.
1.2.3 Comparison of FedExProx and its adaptive variants
In this section, compare FedExProx and its two adaptive variants FedExProx-GraDS and TedExProx-StoPS. We first focus on the full participation case. Note that in this case, the all the algorithms deterministic. For FedExProx-GraDs: extrapolation parameter is given by
it is suggested by Theorem 2, the
1 E.= M3 - prox (I+)
O% OR,G
1. > im (xt prox, (xx))
The server use the local iterates it received from each client to compute ag,G directly. If, in addition, we know Lmaxs we can implement a version that has a better theoretical guarantee,
1 + yLeaax
n E
||< - prox, (I+))
ak. =
Lmax
n 1 Liz1 i=
(x* - prox, (Ix))
For FedExProx-StoPS, we have
1 DE=
[M.. (Æx) inf M7.)
0g 0k,5
E.= VM7, (ox)
In order to implement CK,s. the server requires each client to send the function value of its Moreau envelope at the current iterate to it, and we need to know each inf M7, which, according to Lemma 5, the same as inf f..
From Figure 6, we can observe that in all cases when is sufficiently large, FedExProx-StoPP is the best among the three algorithms considered, and FedExProx ox GraDD outperforms FedExProx, this provides numerical evidence for the effectiveness of our proposed algorithms. In the cases when 7 is small, the convergence of FedExProx- GraDS seems to better than other two algorithms.
We also plot the difference of extrapolation parameter used by the algorithms in each iteration. From Figure observe that when 7 small, 0g,G often much larger than as.s, resulting in better convergence of FedExProx-Gee GraDD observed in the first two plots of Figure 6. When
48

---

00025
0,305
0.05
8
8
haratioss
Io 2t00s
Figure 7: Comparison of the extrapolation parameter ag used by FedExProx, FedExProx-GraDS and FedExProx-StoP5 in each iteration with different step {0.0005, 0.0005, 0.05, 0.5, 1, 5} in the full participation setting.
sizes Y chosen from
becomes larger, O0.G and ar,s become comparable, and their performance is also comparable, with FedExProx-StoPS slightly better than FedExProx-GraDS
We also conduct the experiment where we take client partial participation into account. We can observe from Figure that in all cases, the two adaptive variants FedExProx- GraDS-PP and FedExProx-StoPP5
outperform FedExProx in iteration complexity, and between the two adap- tive variants, FedExProx-Geee GGaDDS better one almost the time. However, FedExProx-GraDS seems to be more stable than FedExProx-Ster StoPS, especially when is small.
49

---

7I 0.1.1 20
05,7 20
10,1 20
Fla1 F062
Flo F06
0.001.F
Fln) R16
(ln) fl62
(16)
12-
1:0 21005
110 21505
BBS
100
(1e)
2008
M00
0.00
2008
100
100
M0D
Figure 8: Comparison of FedExProx, FedExProx-Graie and FedExProx-StoPP in terms of iteration complexity with different step sizes y in the client partial participation (PP) setting. The client minibatch size is chosen from (5, 10, 20), for each minibatch size, a step size y (0.001, 0.005, 0.1, 0.5, 1,5, 10, 50, 100, 500} is randomly selected.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction section accurately reflect the contributions made in this paper, which mainly presented in Section 3, Section 4 and some parts of the Appendix.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long is clear that these goals not attained by the paper.
50

---

2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitations of the work are discussed Section 5.1.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e-g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: Full set of assumptions and a complete and correct proof are described for every fact, lemma, theorem and corollary appeared in this paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
51

---

Answer:
[Yes]
Justification: The details of the experiments are included in the experiment section in Appendix I. The code is also provided in the corresponding link.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, it may be that access to the model is limited in some way (e.g.. registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The details of the experiments are described in detail in Appendix I, and the code is given in the corresponding anonymous link.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.ce)
lublic/guides Soaiddes/CodSuumissseees cotmiiisssoooiicyyyc for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: //nips. wttpccccc//pcc/cccccccccc....llt httpc/ggiidee/ooeuuuuueeeeeeeeee http://oeeeummissiooooooi.oooo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
52

---

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All the details of the experiments and link to anonymized repository are provided which is enough understand the experiment.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report enror bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: The details are depicted in the experiment section, and notice that for the full participation case of our proposed methods, it is deterministic for specific dataset. No errors are needed in this case.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
53

---

Answer: [Yes]
Justification: The computation resources needed for the experiments are described in the experiment section.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: / /tt///nuuu...//. wwwprrps.ccppuulll. cctppblic/eececcccccccccccc c/tthiieieieeeeeoeeee cechicssuudelinees
Answer: [Yes]
Justification: The research conducted in this paper conform with the NeurIPS Code of Ethics in every aspect.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
No potential social impact is expected by the authors.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (c.g., disinformation. generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out, For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
54

---

11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: The paper contain no such risks in the authors' expectation.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: The paper does not use existing assets.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The experiment and code for this paper well documented. The details of the dataset used is described detail in the experiment section of the paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
55

---

The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper dose not involve crowdsoureing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjeets
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
56


Reviews:
[{"review_id": "LwMhIubi0S", "reviewer": "Reviewer_xABN", "summary": "In this paper, the authors present an enhanced version of the FedProx algorithm for federated learning, named FedExProx. Unlike the FedProx, this algorithm incorporates an extrapolation step on the server following the computation of the proximal operator on each client. The authors investigate both constant and adaptive extrapolation step-sizes and provide the corresponding convergence results. Numerical experiments demonstrate that this method surpasses FedProx in terms of convergence rate.", "questions": "In Table 2, the authors may need to include the convergence rate of FedProx for comparison, as the proposed algorithm is more similar to an enhanced version of FedProx rather than FedExP. Another suggestion is to compare the convergence rates of these algorithms in the partial participation setting, since as shown in Table 5, the proposed algorithm still performs well in this scenario.\\\nIn line 220, Remark 2, the authors claim that data heterogeneity is successfully managed. However, as mentioned in the weaknesses section, Assumptions 2 and 3 ensure no data heterogeneity. Therefore, this claim may be incorrect.", "limitations": "The authors have addressed the main limitation of this paper in Section 5.1.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This paper introduces the global step on the server side to the existing FedProx algorithm for the first time, resulting in a new algorithm called FedExProx. The analysis also includes an adaptive global step size, inspired by FedExP. The paper provides clear convergence results, and the theoretical proofs for the main theorems are also presented with great clarity.", "weaknesses": "This paper has two main weaknesses. \\\nFirst, the benefits gained from the extrapolation step are not evident. As shown in Table 2 of this manuscript, the only noticeable difference in the convergence rate compared to the existing FedExP appears to be a difference in the constant factor. The order of parameters, such as $\\tau$ (the number of participating clients or the total number of clients in full participation case) and $T$ (the total number of iterations), remains the same as in existing works. Given that this paper primarily focuses on theoretical contributions, this result is not sufficiently strong.\\\nSecond, the assumptions appear to be too strong. Assumptions 2 (interpolation regime) and 3 (convexity) imply that all functions $f_i$ share the same optimal point. This is a very strong assumption, as it essentially indicates that there is no data heterogeneity, unlike existing works such as FedExP and FedProx, which allow for bounded data heterogeneity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedbacks on our paper. Here is a detailed response to the weaknesses and questions the reviewer mentioned.\n- Weakness 1: `First, the benefits gained from the extrapolation step are not evident. As shown in Table 2 of this manuscript, the only noticeable difference in the convergence rate compared to the existing FedExP appears to be a difference in the constant factor. Given that this paper primarily focuses on theoretical contributions, this result is not sufficiently strong.`\n\n\nWe respectfully disagree with the reviewer. While Table 2 shows that the worst-case convergence of FedExP and FedExProx differs only by a constant factor, key differences remain. FedExP uses an adaptive step size based on gradient diversity, whereas FedExProx uses a constant step size. The adaptive version of FedExProx offers a better worst-case convergence guarantee, as it depends on $L_{\\gamma}(1+\\gamma L_{\\max})$, which can be up to $n$ times better than $L_{\\max}$ according to Lemma 7. Thus, our method has a stronger theoretical guarantee. Additionally, no convergence guarantee was provided for FedExP with a constant extrapolation parameter in the original paper and it was not clear whether or not a constant extrapolation parameter would be effective or not, so a theoretical comparison with FedExProx under these conditions was not possible. In general, due to the differences in the adaptive rules for selecting the extrapolation parameter, it is challenging to compare the convergence directly. However, our experiments in Figure 4 confirm that the iteration complexity of our method, with a properly chosen step size $\\gamma$ and a constant extrapolation $\\alpha$, outperforms FedExP with an adaptive extrapolation parameter.\n\n---\n- Weakness 2: `Second, the assumptions appear to be too strong. Assumptions 2 (interpolation regime) and 3 (convexity) imply that all functions share the same optimal point. This is a very strong assumption, as it essentially indicates that there is no data heterogeneity, unlike existing works such as FedExP and FedProx, which allow for bounded data heterogeneity.`\n\nIndeed, the interpolation regime is a strong assumption. However, in deep learning scenarios, we are often in the overparameterized regime, which is even stronger. As noted in Appendix F.3, without assuming the interpolation regime, the method converges to a neighborhood and the optimal extrapolation constant is reduced by a factor of $2$. This paper focuses on the interpolation regime, as we are inspired by the extrapolation used in the projection methods of convex feasibility problems and the similarity between proximal operations and projections. We also provide discussions of our proposed algorithm in the non-smooth setting, non-convex setting and strongly convex setting in Appendix F. \n\n---\n- Question 1: `In Table 2, the authors may need to include the convergence rate of FedProx for comparison, as the proposed algorithm is more similar to an enhanced version of FedProx rather than FedExP. Another suggestion is to compare the convergence rates of these algorithms in the partial participation setting, since as shown in Table 5, the proposed algorithm still performs well in this scenario.`\n\nThanks for the suggestion. We will add a comparison of FedProx and FedExProx and those algorithms in the partial participation setting in the next version of the paper.\n\n---\n- Question 2: `In line 220, Remark 2, the authors claim that data heterogeneity is successfully managed. However, as mentioned in the weaknesses section, Assumptions 2 and 3 ensure no data heterogeneity. Therefore, this claim may be incorrect.`\n\nThanks for pointing this out, this is indeed a typo, and we will delete the last sentence in Remark 2 in the next version of the paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3va7COjV0N", "reviewer": "Reviewer_eQ2L", "summary": "This paper proposes and analyzes several server-side extrapolation strategies to enhance the theoretical and empirical convergence properties of FedProx.\nThe authors present the convergence properties of the proposed methods for smooth convex and strongly convex problems in the interpolation regime. \nTheoretical results demonstrate that the proposed methods have better dependence on the smoothness coefficient $L$ than FedExP in the general case. \nSpecifically, they achieve an iteration complexity of $O(L_{\\gamma}(1+\\gamma L_{\\max})/\\epsilon)$ compared to FedExP's $\\mathcal{O}(L_\\max/\\epsilon)$. Experimental results validate the theoretical analysis, demonstrating improved performance over both FedProx and FedExP.", "questions": "1. The main concern is the suitability of Assumption 2. Are there any reasons for this assumption?\n2. There are some typo:\n- Line 212, Eq.(8), $\\frac{1}{n}\\sum_{i\\in \\mathcal{S}_k}$, Should this be $\\frac{1}{\\vert \\mathcal{S}_k\\vert}$?\n- Does Eq.(23) assume full participation? It does not seem to align with the partial participation described in Eq.(7) and and the setting in Theorem 1.\n- Line 868, the second equation, Should it be $\\chi^2$?\n\nThe author should review the rest of the paper to address these typo.\n\nI will raise my score if the authors address my concerns.", "limitations": "NA.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- A clear and solid proof is presented in the paper. I have reviewed it and believe the result is correct.\n\n- The authors also talk about the setting where the coefficient $L_{\\gamma, \\tau}$ is unknow. I think these results are useful in practical applications.", "weaknesses": "1. Assumption 2, the interpolation regime, seems too strong. Are there any other published papers that use this assumption? If $\\nabla f_i(x_*)=0$ for all $i\\in [n]$, it suggests that all clients in the system have identical local datasets, which is impractical.\n2. The proposed method is only designed for proximal-based algorithms, which limits its application.\n3. The experiments are conducted on a small dataset with a logistic regression problem. I encourage the author to conduct experiments on more complex models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedbacks on our paper. Here is a detailed response to the weaknesses and questions the reviewer mentioned.\n- Weakness 1 & Question 1: `Assumption 2, the interpolation regime, seems too strong. Are there any other published papers that use this assumption? If for all , it suggests that all clients in the system have identical local datasets, which is impractical.`\n\nIndeed, the interpolation regime is a strong assumption. However, in deep learning scenarios, we are often in the overparameterized regime, which is stronger. We emphasize that the interpolation regime does not imply that all local clients have the same datasets, but rather that they share a common minimizer $x_\\star$​. An example is the convex feasibility problem where the convex sets $\\mathcal{X}_i$​ intersect. Published papers, such as [1] and [2], use the interpolation regime assumption, with theoretical justifications provided in [3].\n\nWe discuss our method's performance in the non-interpolation, non-smooth, and non-convex cases in Appendix F. Specifically, without the interpolation regime assumption, our method converges to a neighborhood of the solution, and the optimal constant extrapolation parameter is reduced by a factor of 2. This paper focuses on the interpolation regime, as we were inspired by the extrapolation technique used in projection methods for convex feasibility problems and the similarity between projections and proximal operations. Prior to our method, it was unclear whether a constant extrapolation parameter would be effective.\n\n[1] A. M. Subramaniam, A. Magesh and V. V. Veeravalli, \"Adaptive Step-Size Methods for Compressed SGD With Memory Feedback,\" _IEEE Transactions on Signal Processing_, 2024.\n\n[2] N. Ion, P. Richtárik and A. Patrascu. \"Randomized projection methods for convex feasibility: Conditioning and convergence rates.\" _SIAM Journal on Optimization_, 2019.\n\n[3] S. Arora, S. Du, W. Hu, Z. Li and R. Wang, \"Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks\" _PMLR_, 2019.\n\n---\n\n- Weakness 2: `The proposed method is only designed for proximal-based algorithms, which limits its application.`\n\nThe benefits of extrapolation also apply to gradient-based algorithms, allowing one to choose the appropriate algorithm for the specific setting. Extrapolation with gradient-based algorithms in federated settings has already been considered in [1]. In our experiments, we use the proposed algorithm FedExP as a benchmark. Our method achieves a better worst-case convergence guarantee than FedExP. Additionally, compared to gradient-based algorithms, proximal algorithms often enjoy enhanced stability.\n\nOur intuition behind this paper stems from the extrapolation technique used in the projection method to solve the convex feasibility problem and the similarity between projections and proximal operations. This is why we consider proximal algorithms in the first place.\n\n[1] J. Divyansh, S. Wang and G. Joshi, \"FedExp: Speeding Up Federated Averaging via Extrapolation.\" _ICLR_ 2023.\n\n---\n\n- Weakness 3: `The experiments are conducted on a small dataset with a logistic regression problem.`\n\nWe thank the reviewer for the feedback, more experiments will be included in the next version of the paper.\n\n---\n- Question 2: `The author should review the rest of the paper to address these typo.`\n\nAll typos mentioned will be corrected in the next version of the paper. We will carefully check the rest of the paper.\nSpecifically, \n- In line 212, Eq (8), it should be $\\frac{1}{| S_k |}$ instead of $\\frac{1}{n}$.\n- There is indeed a typo in Eq (23), we do not assume full participation here.\n- In line 868, yes, there is a square missing here.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mjIQo7tXva", "reviewer": "Reviewer_HgaS", "summary": "The paper presents a new method FedExProx, a federated learning method based on proximal splitting using extrapolation. The method combines the proximal splitting approach from FedProx with the extrapolation from FedExp. FedExProx is shown to improve upon both previous methods in multiple ways, among others faster convergence rates. Then, two adaptive methods are introduced that allow to choose the extrapolation parameter without prior knowledge of the smoothness constants of the individual Moreau envelopes and $L_max$ for both the full and partial participation scenario.", "questions": "-", "limitations": "-", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper introduces a new method that improves on both FedProx and FedExP in terms of iteration complexity (assuming that the functions a proxable). Further, it extends FedExP to the partial participation case and the adaptive variants do not require a local step size as opposed to FedExP.", "weaknesses": "**Comparison to FedExP:** FedExProx assumes that the prox problems can be solved exactly but the analysis of FedExp takes into account the number of local iterations needed. Hence I find the comparison between the convergence rates to be unfair. For smooth problems the prox problems can be approximated very cheaply and the analysis of prox problems can typically be extended to allow for some error in the prox solution. For general convex and L-smooth functions, if one chooses $1/L$ as the prox parameter then prox problem becomes strongly convex and smooth with condition number 2 hence computing an epsilon approximation takes only a logarithmic amount of steps. In your example you argue that there exists a closed form solution for the prox. But for that same example there also exists a closed-form solution for the original problem, so I am not convinced by that argument. I think that the paper should emphasize the assumption that the prox of the function is cheaply computable. \n\n**Extrapolation:** Assuming that the prox can be solved exactly, then based on Eq. 8, I don't see the advantage of considering $\\gamma$ and $\\alpha_k$ separately? Couldn't one just have one step size and compute that using one of the proposed adaptive methods directly? What is gained from decoupling the prox factor from the extrapolation parameter? In FedExp it seems that this is useful because one can decouple the local from the global step size, but since you assume that the prox can computed in closed form, there is not need to adapt the prox parameter right? The only reason I see for this being useful is if one wants to choose the $\\gamma$ in order to be able to solve the proxes efficiently.\n\nI am willing to raise my score if the authors clarify these points.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedbacks on our paper. Here is a detailed response to the weaknesses and questions the reviewer mentioned.\n- Weakness 1: `FedExProx assumes that the prox problems can be solved exactly but the analysis of FedExp takes into account the number of local iterations needed. Hence I find the comparison between the convergence rates to be unfair. For smooth problems the prox problems can be approximated very cheaply and the analysis of prox problems can typically be extended to allow for some error in the prox solution. For general convex and L-smooth functions, if one chooses as the prox parameter then prox problem becomes strongly convex and smooth with condition number 2 hence computing an epsilon approximation takes only a logarithmic amount of steps. In your example you argue that there exists a closed form solution for the prox. But for that same example there also exists a closed-form solution for the original problem, so I am not convinced by that argument. I think that the paper should emphasize the assumption that the prox of the function is cheaply computable.`\n\nWe partly agree with the reviwer on that the paper should emphasize that the proximity operator is computable. This will appear in the next version of the paper. In fact, almost every proximal algorithm needs to assume that the proximity operator is computable in some sense. \n\nIn this paper, we assume that the proximity operator is solved exactly for simplicity, as our goal is to demonstrate the effectiveness of extrapolation combined with proximal algorithms in the federated learning setting. Considering inexactness and obtaining a convergence guarantee to a neighborhood is certainly possible but beyond the scope of this paper.\n\nFor FedExP, the number of local training rounds plays a role in determining the local stepsize and if we take the largest local stepsize in the interpolation regime, then the convergence is independent of the local training round $\\tau$. For FedExProx, however, the amount of local computation needed is hidded in $\\gamma$ and often the larger $\\gamma$ is, the harder it is to compute the proximity operator. However, this does not prevent us from comparing their iteration complexity, as the idea of local training aims to reduce the total number of training rounds and, consequently, the communication complexity. It is neither feasible to directly compare the total number of computations for the two algorithms nor meaningful from our perspective.\n\n---\n- Weakness 2: `Assuming that the prox can be solved exactly, then based on Eq. 8, I don't see the advantage of considering $\\gamma$ and $\\alpha_k$ separately? Couldn't one just have one step size and compute that using one of the proposed adaptive methods directly? What is gained from decoupling the prox factor from the extrapolation parameter? In FedExP it seems that this is useful because one can decouple the local from the global step size, but since you assume that the prox can computed in closed form, there is not need to adapt the prox parameter $\\gamma$ right? The only reason I see for this being useful is if one wants to choose the in order to be able to solve the proxes efficiently.`\n\nWe did not manually separate $\\alpha_k$​ and $\\gamma$ to consider them individually; this separation is inherent to the algorithm. This can be seen from the original formulation in Eq (7) of Algorithm 1. The parameter $\\gamma$ is the local step size associated with each client and determines the effort needed to solve the proximity operator. Often, the larger $\\gamma$ is, the more challenging the local problem becomes. Note that we do not assume the proximity operator can be solved in closed form. The parameter $\\alpha_k$​ is used for extrapolation. It just so happens that after the reformulation in Eq (8), their product becomes the step size for a gradient-based algorithm to minimize the average of Moreau envelopes. Note also that different $\\gamma$ values correspond to different Moreau envelopes $M^{\\gamma}_{f_i}$ as local objectives. As a result, $\\gamma$ influences the local problems we try to solve and cannot be considered in combination with $\\alpha_k$​ directly.\n\nIt is fair to ask which $\\gamma$ is the optimal local step size since we can find the optimal constant extrapolation parameter $\\alpha_k$​ for each $\\gamma$. However, this requires more information about the smoothness of the average of the Moreau envelope $L_{\\gamma}$, which is usually unavailable.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for their time and effort.

The reviewers highlighted several key strengths of our paper. Notably, we introduced an extrapolation parameter to the FedProx algorithm for the first time, developed adaptive versions that eliminate dependence on the unknown smoothness parameter, extended the algorithm to handle partial participation, and provided clear and robust proofs for our claims.

The reviewers also had questions and concerns, which we addressed in the individual responses.

Meta-review: This paper introduces Extrapolated FedProx (FedExProx) to enhance the convergence of FedProx in federated learning. The authors propose three extrapolation strategies: a constant strategy and two adaptive strategies based on gradient diversity (FedExProx-GraDS) and the stochastic Polyak stepsize (FedExProx-StoPS). Theoretical insights are supported by robust numerical experiments, making this a valuable contribution to the community. All reviewers have reached a consensus that this paper should be accepted.
Decision: Accept (poster)