The Power of Extrapolation in Federated Learning

Authors: hanmin  li
kirill  acharya
peter  richtárik
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-10 22:11:55

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.

Main paper: 
 introduction: 
 Federated learning (FL) is a 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 et al., 2017] . This method allows clients to contribute to model training while keeping their data private, as it avoids the need for direct data sharing. Often, federated optimization is formulated as the minimization of a finite-sum objective function,EQUATIONwhere each f i : R d → R is the empirical risk of model x associated with the i-th client. The federated averaging method (FedAvg) is among the most favored strategies for addressing federated learning problems, as 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 (SGD) with random reshuffling, as discussed by Bubeck et al. [2015] , Gower et al. [2019] , Moulines and Bach [2011] , Sadiev et al. [2022] . ‡ Kirill was a student at MIPT during his internship at KAUST. Li et al. [2020] proposed replacing the local training of each client via SGD in FedAvg with the computation of a proximal term, resulting in the FedProx algorithm.EQUATIONwhere γ > 0 is the step size, and the proximal operator is defined as prox γfi (x) := arg minz∈R d f i (z) + 1 2γ ∥z -x∥ 2 .Contrary to gradient-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 (SPPM) [Asi and Duchi, 2019 , Bertsekas, 2011 , Khaled and Jin, 2022 , Patrascu and Necoara, 2018 , Richtárik and Takác, 2020] benefit from stability against inaccuracies in learning rate specification [Ryu and Boyd, 2014] . Indeed, for GD and SGD, a step size that is excessively large can result in divergence of the algorithm, whereas a step size that is too small can significantly deteriorate the convergence rate of the algorithm. PPM was 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.It is known that the proximal operator applied to a proper, closed and convex function can be viewed as the projection to some level set of the same function depending on the value of γ. In particular, if we let each f i be the indicator function of a nonempty closed convex set X i , then prox γfi (•) becomes the projection Π Xi (•) onto the set X i . 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] , if we additionally assumeX := n i=1 X i ̸ = ∅.A well known fact about the parallel projection method is that its empirical efficiency can often be improved by extrapolation [Combettes, 1997a , Necoara et al., 2019] . This involves moving further along the line that connects the last iterate x k and the average projection point, resulting in the iterationEQUATIONwhere α k ≥ 1 defines extrapolation level. Despite the various heuristic rules proposed over the years for setting α k [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 of 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 α k as follows:EQUATIONHowever, 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 a server side stepsize in the context of federated learning, its effectiveness was discussed by Malinovsky et al. [2023] .In the field of fixed point methods, extrapolation is also known as over-relaxation [Rechardson, 1911] .It is a technique used to effectively accelerate the convergence of fixed point methods, including gradient algorithms and proximal splitting algorithms [Iutzeler and Hendrickx, 2019, Condat et al., 2023] .Our paper contributes in the following ways; for the notations used please refer to Appendix A.• Based on the insights gained from the convex feasibility problem, we extend FedProx to its extrapolated counterpart FedExProx for both convex and strongly 1 convex interpolation problems (See Table 1 ). By optimally setting the constant extrapolation parameter, we obtain iteration complexity O Lγ (1+γLmax) ϵ 2 in the convex case and O Lγ (1+γLmax) µ log 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 2 + 1 γLmax + γL max times better in terms of iteration complexity, see Remark 5.• Our method, FedExProx, improves upon the worst-case iteration complexity O Lmax ϵ of FedExP [Jhunjhunwala et al., 2023] to O Lγ (1+γLmax) ϵ(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 extrapolation parameter on τ which is the number of clients participating in the training and the benefits of a larger τ . In particular, we show that compared to the single client setting, with complexity O Lmax ϵ , the full participation version enjoys a speed-up up to a factor of n, see Remark 7.• 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 a step size based on 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-GraDS), and the stochastic Polyak step size (FedExProx-StoPS) [Horváth et al., 2022 , Loizou et al., 2021] . The proposed methods eliminate reliance on the unknown smoothness constant and exhibit "semi-adaptivity", meaning the algorithm converges with any local step size γ and by selecting a sufficiently large γ, we ensure that we lose at most a factor of 2 in iteration complexity.• We validate our theory with numerical experiments. Numerical evidence suggests that FedExProx achieves a 2× 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.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 x k+1 = x k -η •g(x k ), where η > 0 is a scalar step size, g(x k ) represents a stochastic estimator of the true gradient ∇f (x k ). We recover GD when g(x k ) = ∇f (x k ). 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 et 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 SGD. While many methods have been crafted to leverage the stochastic nature of g(x k ), substantial research efforts are also dedicated to finding a 1 Strongly convex: f is µ-strongly convex. 2 As we later see in Theorem 1, here Lmax = max i∈[n] Li, where each Li is the smoothness of fi, Lγ is the smoothness constant ofM γ = 1 n n i=1 M γ f i .Table 1 : General comparison of FedExP, RPM and FedExProx in terms of conditions and convergence. Each entry indicates whether the method has the corresponding feature (✓) or not (✗). We use the sign "-" where a feature is not applicable to the corresponding method. a RPM refers to the randomized projection method of Necoara et al. [2019] . Our method includes it as a special case, see Remark 12 b Convexity: local objective fi is convex, which is the indicator function of the convex set Xi in RPM. c The strong convexity pertains to f , and for RPM, it indicates that the linear regularity condition is satisfied. d Smoothness: fi is Li-smooth. Our algorithm also applies in the non-smooth case; see Appendix F.2. e Jhunjhunwala et al. [2023] provides no convergence guarantee for client partial participation setting. f The concept of "semi-adaptivity" is explained in Remark 9.Table 2: Comparison of convergence of FedExP, FedProx, FedExProx, FedExProx-GraDS and FedExProx-StoPS. The local step size of FedExP is set to be the largest possible value 1 /6tL in the full batch case, where t 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 1 and Theorem 2. The convergence for our methods are described for arbitrary γ > 0. We use K to denote the total number of iterations. For FedExProx, optimal constant extrapolation is used. The O (•) notation is hidden for all complexities in this table.General Case Best Case Worst Case 1+γLmax) /γK a The α k,P here is determined according to the theory of Jhunjhunwala et al. [2023] .EQUATIONb Notice that we always have γLγ < 1, so the complexity of FedProx is strictly worse than FedExProx. c We have Lγ (1 + γLmax) ≤ Lmax, see Remark 6. d We leave out a factor of (1+γLmax) /(2+γLmax) which is a constant between ( 1 2 , 1). e See Remark 11 for a lower bound of α k,S , using which we can rewrite the rate asLγ (1+γLmax) K .better stepsize. An illustration of this is the coordinate-wise adaptive step size Adagrad [Duchi et al., 2011] . Another approach involves employing matrix step size, as demonstrated by Safaryan et al. [2021] , Li et al. [2023 Li et al. [ , 2024]] . Our analysis builds on the theory of SGD mainly adapted from Gower et al. [2019] with additional consideration on the upper bound of the step size.Stochastic proximal point method. PPM was first introduced by Rockafellar [1976] to address the problems of variational inequalities at its inception. Its transition to stochastic case, motivated by the 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 to be practical. Over the years, SPPM has been the subject of extensive research, as 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 to solve systems of linear equations or linear inequalities [Kaczmarz, 1937 , Von Neumann, 1949 , Motzkin and Schoenberg, 1954] . Subsequently, it was generalized to address the convex feasibility problem [Combettes, 1997b] . Typically, the method involves projecting onto a set X i , where i 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 to 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 as Moreau-Yosida regularization, was first introduced by Moreau [1965] as a mathematical tool for handling non-smooth functions.A particularly relevant property of the Moreau envelope is 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 analysis 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] , RMSProp [Hinton et al.] , and ADAM [Kingma and Ba, 2015] . Recently, several studies have attempted to 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] . 
preliminaries: 
 We now introduce the several definitions and assumptions that are used throughout the paper. Definition 1 (Proximity operator). The proximity operator of an extended-real-valued function ϕ : R d → R ∪ {+∞} with step size γ > 0 is defined asprox γϕ (x) := arg min z∈R d ϕ(z) + 1 2γ ∥z -x∥ 2 .It is known that for a proper, closed and convex function ϕ, the minimizer of ϕ(z) + 1 2γ ∥z -x∥ 2 exists and is unique. Throughout this paper, we assume that the proximal operators are evaluated exactly, with no approximation or inexactness.Algorithm 1 Extrapolated SPPM (FedExProx) with partial client participation 1: Parameters: extrapolation parameter α k > 0, step size for the proximity operator γ > 0, starting point x 0 ∈ R d , number of clients n, total number of iterations K, number of clients participate in the training τ , for simplicity, we use τ -nice sampling as an example 2: for k = 0, 1, 2 . . . K -1 do 3:The server samples S k ⊆ {1, 2, . . . , n} uniformly from all subsets of cardinality τ 4:The server computesEQUATION5: end for Definition 2 (Moreau envelope). The Moreau envelope of an extended-real-valued function ϕ : R d → R ∪ {+∞} with step size γ > 0 is defined asM γ ϕ (x) := min z∈R d ϕ(z) + 1 2γ ∥z -x∥ 2 .The following assumptions are used in our analysis. We use the notation [n] for the set {1, . . . , n}.Assumption 1 (Differentiability). The functionf i in (1) is differentiable for all i ∈ [n].Assumption 2 (Interpolation regime). There existsx ⋆ ∈ R d such that ∇f i (x ⋆ ) = 0 for all i ∈ [n].Note that Assumption 2 indicates that each f i and f are lower bounded. In this paper, we focus on cases where the interpolation regime holds. This assumption often holds in modern deep learning which are 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 et al., 2019] , wherein the intersection X is presumed nonempty. This is equivalent to assuming that the interpolation regime holds when f i is the indicator function of the nonempty closed convex set X i . Further motivations derived from the proof for this assumption will be discussed later.Assumption 3 (Convexity). The functionEQUATIONAssumption 4 (Smoothness). Functionf i : R d → R is L i -smooth, L i > 0 for all i ∈ [n]. This means that for each f i ,EQUATION)We will useL max to denote max i∈[n] L i .It 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 to highlight how the optimal extrapolation parameter depends on smoothness if it is present. The following strong convexity assumption is introduced that, if adopted, enables us to achieve better results.Assumption 5 (Strong convexity). The function f is µ-strongly convex, µ > 0. That isf (x) -f (y) -⟨∇f (y), x -y⟩ ≥ µ 2 ∥x -y∥ 2 , ∀x, y ∈ R d .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 a 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. 
constant extrapolation: 
 In order to demonstrate the convergence result of our algorithm in the stochastic setting, we use τ -nice sampling as the way of selecting clients for partial participation. This refers to that in each iteration, the server samples a set S k ⊆ {1, 2, . . . , n} uniformly at random from all subsets of size τ . 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 1 (Differentiability), Assumption 2 (Interpolation regime), Assumption 3 (Convexity) and Assumption 4 (Smoothness) hold. If we use a fixed extrapolation parameterα k = α ∈ 0, 2 γLγ,τand any step size 0 < γ < +∞, then the average iterate of Algorithm 1 satisfiesE [f (x K )] -inf f ≤ C (γ, τ, α) • ∥x 0 -x ⋆ ∥ 2 K ,where K is the number of iteration, xK is sampled uniformly at random from the first K iterates {x 0 , x 1 , . . . , x K-1 }, C (γ, τ, α) is defined asC (γ, τ, α) := 1 + γL max αγ (2 -αγL γ,τ )and L γ,τ := n -τ τ (n -1)L max 1 + γL max + n(τ -1) τ (n -1) L γ ,whereL max = max i L i , L γ is the smoothness constant of M γ (x) := 1 n n i=1 M γ fi (x).If we fix γ and τ the optimal constant extrapolation parameter is given by α γ,τ := 1 γLγ,τ > 1, which results in the following convergence guarantee:E [f (x K )] -inf f ≤ C(γ, τ, α γ,τ ) • ∥x 0 -x ⋆ ∥ 2 K = L γ,τ (1 + γL max ) • ∥x 0 -x ⋆ ∥ 2 K .The proof of this theorem relies on the reformulation of the update rule in (7), using the identity ∇M γ fi (x) = 1 γ x -prox γfi (x) given in Lemma 2, which holds for any x ∈ R d , into the following form:EQUATIONWe 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 γ (x). 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 Algorithm 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 1 indicates convergence for any 0 < γ < +∞. Indeed, as it is proved by Lemma 7, we have C (γ, τ, α γ,τ ) = L γ,τ (1 + γL max ) ≤ L max holds for any 0 < γ < +∞. In cases where there exists at least one L i < L max , we have C (γ, τ, α γ,τ ) < L max . Remark 3. One may question the necessity of the interpolation regime assumption. This assumption is crucial to our analysis. Besides allowing us to revisit the convex feasibility problem setting, it also guarantees that M γ (x) 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 2 in the SGD theory, which is demonstrated in Theorem 3 in the Appendix. Remark 4. From the reformulation presented in (8), we see the best extrapolation parameter is obtained when α k γ is the best step size for SGD running on global objective M γ (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 overparameterized regime. Remark 5. Our algorithm includes FedProx as a special case when α = 1. To recover its result, we simply plug in α = 1, the resulting condition number is C(γ, τ, 1) = 1+γLmax γ(2-γLγ,τ ) . Compared to FedProx, Algorithm 1 with the same γ > 0 demonstrates superior performance, with the acceleration factor being quantified by C(γ, τ, 1)C (γ, τ, α γ,τ ) ≥ 2 + 1 γL max + γL max ≥ 4.See Lemma 14 for the proof. This suggests that the approach of the server averaging all iterates following local computation is suboptimal.In the following paragraphs, we study some special cases, Full participation case For the full participation case (τ = n), using definition from Theorem 1EQUATIONIn 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 ofFedExP is given by O Lmax ϵ , while for Algorithm 1, it is O C(γ,n,αγ,n) ϵ. As suggested by Lemma 7, Algorithm 1 has a better iteration complexity (C (γ, n, α γ,n ) < L max ) whenever there exists L i ̸ = L max for some i ∈ [n], and the acceleration could reach up to a factor of n as suggested by Example 1. In general, the speed-up in the worst case is quantified byL max 1 + γL max • 1 n n i=1 L i 1 + γL i -1 ≤ L max C (γ, n, α γ,n ) ≤ n • L max 1 + γL max • 1 n n i=1 L i 1 + γL i -1.Single client case For the single client case (τ = 1), using definition from Theorem 1α γ,1 = 1 + 1 γL max > 1, L γ,1 = L max 1 + γL max , C (γ, 1, α γ,1 ) = L max .Remark 7. Compared with full and partial client participation, the following relations hold for anyτ ∈ [n], C (γ, n, α γ,n ) ≤ C (γ, τ, α γ,τ ) ≤ C (γ, 1, α γ,1 ) and α γ,1 ≤ α γ,τ ≤ α γ,n , ∀τ ∈ [n].Since the iteration complexity of FedExProx is given by O C(γ,τ,αγ,τ ) ϵ, the above inequalities tell us a larger client minibatch size τ leads to a larger extrapolation and a better iteration complexity. Specifically, Lemma 7 suggests the improvement over the single client case could be as much as a factor of n (C (γ, n, α γ,n ) = 1 n C (γ, 1, α γ,1 )) as suggested by Example 1. 
adaptive extrapolation: 
 Observe that in Theorem 1, in order to determine the optimal extrapolation, we require the knowledge of L γ,τ , 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, we introduced two variants of FedExProx, based on gradient diversity and stochastic Polyak step size, given their relation to the extrapolation parameter in our cases. Theorem 2. Suppose Assumption 1 (Differentiability), Assumption 2 (Interpolation regime), Assumption 3 (Convexity) and Assumption 4 (Smoothness) hold.(i) (FedExProx-GraDS): If we are using α k = α k,G , whereEQUATIONthen the iterates of Algorithm 1 with τ = n satisfyE [f (x K )] -inf f ≤ 1 + γL max 2 + γL max • 1 γ + L max • ∥x 0 -x ⋆ ∥ 2 K-1 k=0 α k,G, where xK is chosen randomly from the first K iterates {x 0 , x 1 , ..., x K-1 } with probabilitiesp k = α k,G/ K-1 k=0 α k,G .(ii) (FedExProx-StoPS): If we are using α k = α k,S , where,EQUATIONthen the iterates of Algorithm 1 with τ = n satisfyEQUATION)where xK is chosen randomly from the first K iterates {x 0 , x 1 , ..., x K-1 } with probabilitiesp k = α k,S/ K-1 k=0 α k,S .Theorem 2 describes the convergence in the full participation setting. However, we can also extend it to the stochastic setting by implementing a stochastic version of these adaptive step size rules for gradient-based methods [Horváth et al., 2022 , Loizou et al., 2021] . See Theorem 5 in the Appendix for the details.Remark 8. In fact, the adaptive rule based on gradient diversity can be improved by using Lmax 1+γLmax instead of 1 γ as the maximum of local smoothness constant of Moreau envelops, resulting in the extrapolation,EQUATION)One can obtain a slightly better convergence guarantee than the FedExProx-GraDS case in Theorem 2, see Corollary 2 in the Appendix. However, the requires the knowledge of L max in order to compute 1+γLmax γLmax . Remark 9. Note that, compared to classical gradient-based methods, FedExProx-GraDS benefits from "semi-adaptivity". This refers to the fact that the algorithm converges for any choice of γ > 0. Although a smaller γ hinders convergence, setting it to at least 1 Lmax limits the worsening of the convergence to a factor of 2.Remark 10. Compared to FedExProx with the optimal constant extrapolation parameter, we gain "semi-adaptivity" here by using the gradient diversity based extrapolation. However, this results in losing the favorable dependence of convergence on L γ and instead establishes a dependence on L max .Remark 11. For FedExProx-StoPS, as it is suggested by Lemma 20, the convergence depends on the favorable smoothness constant L γ , rather than on L max . 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 is the RPM by Necoara et al. [2019] to 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 1 for a simple experiment confirming that FedExProx indeed has a better iteration complexity than FedProx. For more details on the experiments, we refer the readers to Appendix I in the Appendix. Notice that in 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 of this subroutine depends on the local stepsize. If γ is large, the local problem becomes harder to solve because we aim to minimize the local objective itself. Conversely, if it is small, the problem is easier since we do not stray far 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 are carried out by the clients. Iterations10 -4 10 -3 10 -2 10 -1 10 0 f (xk) -f (x ) γ = 10, αγ,n = 2.01576 FedProx FedExProxFigure 1: Comparison of FedExProx and FedProx in terms of iteration complexity in the full participation setting. The notation γ here denotes the local step size of the proximity operator and α γ,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. 
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.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. Throughout the paper, we use the notation ∥•∥ to denote the standard Euclidean norm defined on R d and ⟨•, •⟩ to denote the standard Euclidean inner product. Given a differentiable function f : R d → R, its gradient is denoted as ∇f (x). For a convex function f : R d → R, we use ∂f (x) to denote its subdifferential at x. We use the notation D f (x, y) to denote the Bregman divergence associated with a function f : R d → R between x and y. The notation inf f is used to denote the minimum of a function f : R d → R. We use prox γf (x) to denote the proximity operator of function f : R d → R with γ > 0 at x ∈ R d , and M γ f (x) to denote the corresponding Moreau Envelope. The notation □ is used for the infimal convolution of two proper functions. We denote the average of the Moreau envelope of each local objective f i by the notation M γ : R d → R. Specifically, we defineM γ (x) = 1 n n i=1 M γ f (x). Note that M γ (x) has an implicit dependence on γ, its smoothness constant is denoted by L γ . We say an extended real-valued function f : R d → R ∪ {+∞} is proper if there exists x ∈ R d such that f (x) < +∞. We say an extended real-valued function f : R d → R ∪ {+∞} is closed if its epigraph is a closed set. The following Table 3 summarizes the commonly used notations and quantities appeared in this paper.Fact 1 (First prox theorem). [Beck, 2017, Theorem 6.3 ] Let f : R d → R be a proper, closed and convex function. Then prox f (x) is a singleton for any x ∈ R d .Fact 2 (Second prox theorem). [Beck, 2017, Theorem 6.39 ] Let f : R d → R ∪ {+∞} be a proper, closed and convex function. Then for any x, u ∈ R d , the following three claims are equivalent:EQUATIONEQUATION)Table 3 : Summary of frequently used notations and quantities in this paper.The total number of clients. dThe dimension of the model. xThe model which belongs to R d .The total number of iterations.x kThe model at k-th iteration. α kThe extrapolation parameter at iteration k. fi(x)Each local objective function. γThe step size in the proximity operator.The global objective f .prox γf i (x)The proximity operator associated with fi and γ > 0at point x ∈ R d . M γ f i (x)The Moreau envelope associated with fi and γ > 0at point x ∈ R d . M γ (x)The average of M γ f i (x).The smoothness constant of fi. LThe smoothness constant of f . µThe strong convexity constant of f .L i/(1+γL i )The smoothness constant of M γ f iThe maximum of all Li, for i ∈[n]. Lmax /(1+γLmax) The maximum of the smoothness constant of each M γ f i for i ∈ [n].The smoothness constant of M γ .The interpolation between the Lγ and Lmax /(1+γLmax) induced by τ -nice sampling.The optimal extrapolation parameter of FedExProx under τ -nice sampling.C (γ, τ, α)The convergence rate of FedExProx with τ -nice sampling in the convex case.α k,GThe gradient diversity extrapolation in the k-th iteration defined in Theorem 2. α k,SThe stochastic Polyak extrapolation in the k-th iteration defined in Theorem 2.α ′ k,GThe improved gradient diversity based extrapolation used in Corollary 2.D f (x, y)The Bregman divergence associated with f between two points x, y ∈ R d .The set of indices server sampled in the k-th iteration.α τ,k,GThe gradient diversity based extrapolation in the k-th iteration for FedExProx-GraDS-PP. α τ,k,SThe stochastic Polyak based extrapolation in the k-th iteration for FedExProx-StoPS-PP.If f is convex, L-smooth and differentiable, the following inequalities hold for any x, y ∈ R d ,EQUATION)Fact 4 (Increasing function). Let f (x) = x 1+γx , where γ > 0. Then f (x) is monotone increasing when x > 0.In this section, we explore the properties of the Moreau envelope of individual functions f i , and the global objectiveM γ = 1 n n i=1 M γ fi .Before that, we present the definition of infimal convolution Definition 3 (Infimal convolution). The infimal convolution of two proper functions f, g : R d → R ∪ {+∞} is defined via the following formula(f □g) (x) = min z∈R d {f (z) + g(x -z)} .One key observation is that M γ f can be viewed as the infimal convolution of the proper, closed and convex function f and the real-valued convex function 1 2γ ∥•∥ 2 . This observation enables us to infer the convexity and smoothness of the Moreau envelope from the properties of the original function.First, we present two lemmas about basic properties of Moreau envelope.Lemma 1 (Real-valuedness). Let f : R d → R ∪ {+∞} be a proper, closed and convex function.Then its Moreau envelope M γ f for any γ > 0 is a real-valued function. In particular, the following identity holds for x ∈ R d according to the definition of Moreau envelope,M γ f (x) = f prox γf (x) + 1 2γ x -prox γf (x) 2 .Lemma 2 (Differentiability of Moreau envelope). [Beck, 2017, Theorem 6 .60] Let f : R d → R ∪ {+∞} be a proper, closed and convex function. Then its Moreau envelope M γ f for any γ > 0 is 1 γ -smooth, and for any x ∈ R d , we have∇M γ f (x) = 1 γ x -prox γf (x) .We then focus on the relation between individual f i and M γ fi . The following lemma suggests that the convexity of individual f i guarantees the convexity of M γ fi . Lemma 3 (Convexity of Moreau envelope). [Beck, 2017, Theorem 6 .55] Let f : R d → R ∪ {+∞} be a proper and convex function. Then M γ f is a convex function.It is also true that the smoothness of individual f i indicates the smoothness of M γ fi . Lemma 4 (Smoothness of Moreau envelope). Let f : R d → R be a convex and L-smooth function.Then M γ f is L 1+γL -smooth.One notable fact is that f i and M γ fi have the same set of minimizers. Lemma 5 (Minimizer equivalence). Let f : R d → R ∪ {+∞} be a proper, closed and convex function. Then for any γ > 0, f and M γ f has the same set of minimizers.In addition, M γ f is a global lower bound of f . Lemma 6 (Individual lower bound). Let f : R d → R ∪ {+∞} be a proper, closed and convex function. Then the Moreau envelopeM γ f satisfies M γ f (x) ≤ f (x) for all x ∈ R d .Next, we focus on the global objective M γ (x). The following lemma bounds its smoothness constant from both above and below.Lemma 7 (Global convexity and smoothness). Let each f i be proper, closed convex and L i -smooth. Then M is convex and L γ -smooth with1 n 2 n i=1 L i 1 + γL i ≤ L γ ≤ 1 n n i=1 L i 1 + γL i .As a result of the above inequalities, we have the following inequality on the condition number defined in Theorem 1 which holds for any τ ∈ [n],L γ (1 + γL max ) = C (γ, n, α γ,n ) ≤ C (γ, τ, α γ,τ ) ≤ C (γ, 1, α γ,1 ) = L max .When there exists at least oneL i < L max , we have C (γ, n, α γ,n ) < C (γ, τ, α γ,τ ) < L max = C (γ, 1, α γ,1 ). Even L i = L max holds for all i ∈ [n], there are cases (See Example 1 in the proof.) that C (γ, n, α γ,n ) = 1 n C (γ, 1, α γ,1 ) = 1 n L max .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 i : R d → R ∪ {+∞} be proper, closed and convex, thenf (x) = 1 n n i=1 f i (x)has the same set of minimizers and minimum asM γ (x) = 1 n n i=1 M γ fi (x) ,if we are in the interpolation regime and 0 < γ < ∞.The following lemma generalizes Lemma 6 into the finite-sum setting. Lemma 9 (Global lower bound). Let each f i : R d → R ∪ {+∞} be proper, closed and convex. Then the following inequality holds for any x ∈ R d and γ > 0,M γ (x) ≤ M γ f (x) ≤ f (x).In addition, if we assume we are in the interpolation regime, then M γ , M γ f and f have the same set of minimizers, for any x ⋆ in this set of minimizers, the following identity holds,M γ (x ⋆ ) = M γ f (x ⋆ ) = 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 1 (Differentiability), 2 (Interpolation Regime), 3 (Convexity) and 4 (Smoothness) hold, for any minimizer x ⋆ of M γ (x), all x ∈ R d , the following inequality holds for each client objective,EQUATION)Furthermore, this suggestsEQUATIONA direct consequence of the above function suboptimality bound is the star strong convexity of M γ from the strong convexity of f . Lemma 11. (Star strong convexity) Assume Assumption 1 (Differentiability), Assumption 2 (Interpolation Regime), Assumption 3 (Convexity), Assumption 4 (Smoothness) and Assumption 5 (Strong convexity) hold, then the convex function M γ (x) satisfies the following inequality,M γ (x) -M γ (x ⋆ ) ≥ µ 1 + γL max • 1 2 ∥x -x ⋆ ∥ 2 ,for any x ∈ R d and a minimizer x ⋆ of M γ (x).The star strong convexity property of M γ allows us to improve the sublinear convergence in the convex regime into linear convergence.Lemma 12. Let f : R d → R be a proper, closed and convex function. Then x is a minimizer of f if and only if x = prox γf (x).Lemma 13. Assume we are working with the finite-sum problem f = 1 n n i=1 f i , where each f i is convex and L i -smooth, f is convex and L-smooth. Then the smoothness of L satisfies1 n 2 n i=1 L i ≤ L ≤ 1 n n i=1 L i ,where both bounds are attainable. Lemma 14. Assume that all the conditions mentioned in Theorem 1 hold, then the condition number C(γ, τ, 1) of FedProx and the condition number C (γ, τ, α γ,τ ) of the optimal constant extrapolation parameter α ⋆ = 1 γLγ,τ satisfy the following inequality,C(γ, τ, 1) C (γ, τ, α γ,τ ) ≥ 2 + 1 γL max + γL max ≥ 4 ∀τ ∈ [n].Lemma 15. Assume that all the conditions mentioned in Theorem 1 hold, then the following inequalities hold,C (γ, n, α γ,n ) ≤ C (γ, τ, α γ,τ ) ≤ C (γ, 1, α γ,1 ) , ∀τ ∈ [n],andα γ,1 ≥ α γ,τ ≥ α γ,n , ∀τ ∈ [n].In order to prove our main theorem, we partly rely on the theory of SGD. The following theorem on the convergence of SGD with τ -nice sampling is adapted from Gower et al. [2019] . We introduce modifications to the proof technique and tailor the theorem specifically to 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 τ -nice sampling 1: Parameters: learning rate η > 0, starting point x 0 ∈ R d , minibatch size τ ∈ {1, 2, . . . , n} 2: for k = 0, 1, 2, . . . doThe server samples S k ⊆ {1, 2, . . . , n} uniformly from all subsets of cardinality τ 4:The server performs one gradient stepx k+1 = x k -η • 1 τ ξi∈S k ∇f ξi (x k ).5: end for Theorem 3. Assume Assumption 1 (Differentiability), 2 (Interpolation regime), 3 (Convexity), 4 (Smoothness) hold. Additionally, assume f is L-smooth where L ≤ 1 n n i=1 L i .3 If we are running SGD with τ -nice sampling using step size η that satisfies 0 < η < 2 Aτ , whereEQUATIONL max := max i L i ,then the iterates of Algorithm 2 satisfyE [f (x K )] -inf f ≤ 1 η(2 -ηA τ ) • ∥x 0 -x ⋆ ∥ 2 K ,where K is the total number of iterations, xK is chosen uniformly at random from the first K iterates {x 0 , x 1 , . . . , x K-1 }. If, additionally, we assume the following property (which we will refer to as "star strong convexity") holds, then the iterates of Algorithm 2 satisfyE ∥x K -x ⋆ ∥ 2 ≤ 1 -η(2 -ηA τ ) • µ 2 K ∥x 0 -x ⋆ ∥ 2 .In this section, we provide some additional details on the analysis of FedExProx and its adaptive variants.The following corollary summarizes the convergence guarantee in the strongly convex case. Corollary 1. Suppose the assumptions in Theorem 1 hold, and assume in addition that Assumption 5 (Strong Convexity) holds, then we achieve linear convergence for the final iterate of Algorithm 1, which satisfiesE ∥x K -x ⋆ ∥ 2 ≤ 1 -αγ(2 -αγL γ,τ ) • µ 2 (1 + γL max ) K ∥x 0 -x ⋆ ∥ 2 ,where the definition of L γ,τ is given in Theorem 1. Fixing the choice of γ and τ , the optimal extrapolation parameter that minimizes the convergence rate is given by α γ,τ = 1 γLγ,τ > 1, which results in the following convergence in the strongly convex case:E ∥x K -x ⋆ ∥ 2 ≤ 1 - µ 2L γ,τ (1 + γL max ) K ∥x 0 -x ⋆ ∥ 2 .As one can observe, by additionally assuming µ strong convexity of the original function f , we improve the sublinear convergence in the convex case into linear convergence.Our analysis also adapts to the non-smooth cases. This is based on the observation that even if we only assume Assumption 1 (differentiability), Assumption 2 (Interpolation Regime) and Assumption 3 (Convexity) hold and do not have additional assumptions on smoothness, still each M γ fi is 1 γ -smooth 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 1 (Differentiability), 2 (Interpolation Regime) and 3 (Convexity) hold. If we choose a constant extrapolation parameter α k = α satisfying0 < α < 2 γL γ,τ,where L γ is the smoothness constant of M γ (x) = 1 n n i=1 M γ fi (x), L γ,τ is given by L γ,τ = n -τ τ (n -1) • 1 γ + n(τ -1) τ (n -1) • L γ .Then the iterates of Algorithm 1 satisfyγM γ (x K ) -inf γM γ ≤ 1 α (2 -αγL γ,τ ) • ∥x 0 -x ⋆ ∥ 2 K ,where xK is chosen uniformly from the first K iterates {x 0 , x 1 , . . . , x K-1 }. It is easy to see that the best α is given byα ⋆ = 1 γL γ,τ ≥ 1,where the corresponding convergence is given byγM γ (x K ) -inf γM γ ≤ n -τ τ (n -1) + n(τ -1) τ (n -1) γL γ • ∥x 0 -x ⋆ ∥ 2 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 f i (x) = I Xi (x), then we haveprox γfi (x) = Π Xi (x) , ∀x ∈ R d ,andγM γ fi (x) = 1 2 ∥x -Π Xi (x)∥ 2 , and γM γ (x) = 1 2 • 1 n n i=1 ∥x -Π Xi (x)∥ 2 .Since we are in the interpolation regime, inf γM γ = 0, and the convergence result becomes1 2 • 1 n n i=1 ∥x K -Π Xi (x K )∥ 2 ≤ n -τ τ (n -1) + n(τ -1) τ (n -1) γL γ • ∥x 0 -x ⋆ ∥ 2 K .Notice that here γL γ ≤ 1 is the smoothness constant associated with each distance function1 2 ∥x -Π Xi (x)∥ 2 .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.A 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-smooth case, it is possible that γL γ = 1, where the optimal α ⋆ = 1, in this case, extrapolation will not generate any benefits. However, as it is mentioned by Necoara et al. [2019] , there are many examples where γL γ < 1 and extrapolation indeed accelerates the algorithm. This is different from the smooth case, where extrapolation always helps.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 f which is used in the smooth case.Using a similar approach, it is also possible to obtain a convergence guarantee for FedExProx in the strongly convex non-smooth regime, assuming in addition that M γ (x) is µ γ -strongly 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 as a special case. Necoara et al. [2019] obtained using our theory and the original theory. In both cases, the optimal extrapolation parameter is used. The notation O(•) is hidden. ε is the error level reached by function value suboptimality for convex case, squared distance to the solution for strongly convex case.Original Theory Our TheoryConvex + smooth case (1) γLγ,τ• ∥x 0 -x⋆∥ 2 ε γLγ,τ • ∥x 0 -x⋆∥ 2 ε Strongly convex + smooth case (2) Lγ,τ µγ • log ∥x 0 -x⋆∥ 2 ε Lγ,τ µγ • log ∥x 0 -x⋆∥ 2 ε (1)The smoothness here does not refer to each fi being Li-smooth, but γM γ being γLγ-smooth. This corresponds to the smooth regularity condition presented in 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 γ (x) being µγ-strongly convex with µγ < Lγ.For the non-interpolation regime cases, we assume that Assumption 1 (Differentiability), Assumption 3 (Convexity) and Assumption 4 (Smoothness) hold. The differences are listed as follows (i) Although f i and M γ fi have the same set of minimizers, f and M γ does not necessarily have the same set of minimizers. This will lead to the convergence of FedExProx to the minimizer x ′ ⋆,γ of M γ (x) instead of x ⋆ of f . As a result, we will only converge to a 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 α ′ ⋆ in the non-interpolated cases is also halved, α ′ ⋆ = 1 2 α ⋆ . As a result, it is possible that α ′ ⋆ ≤ 1. The same phenomenon is also observed in FedExP of Jhunjhunwala et al. [2023] , where their heuristic in determining the extrapolation parameter adaptively is also reduced by a factor of 2 in non overparameterized cases.Observe that all of the above results in both smooth/non-smooth, interpolated/non-interpolated cases suggests that the practice of server simply averaging the iterates it obtained from local training is suboptimal.In the non-convex case, we assume Assumption 1 (Differentiability) holds, and we need the following additional assumptions on f : R d → R and f i : R d → R: Assumption 6 (Lower boundedness). Function f i is lower bounded by inf f i . Assumption 7 (Weak convexity). Function f i is ρ > 0 weakly convex, this means thatf i + ρ 2 ∥•∥ 2 is 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, ρ-weakly convex function and let γ < 1 ρ . Then the Moreau envelopeM γ f is continuously differentiable on R d with ∇M γ f (x) = 1 γ x -prox γf (x) .In addition, the Moreau envelope is max 1 γ , ρ 1-γρ -smooth. We will thereby denote the smoothness constant as L γ,ρ .Indeed, if the stepsize γ in this case is chosen properly such that 1 γ > ρ, then it is straight forward to see the function within the proximity operator prox γfi given byf i + 1 2 • 1 γ ∥•∥ 2 is strongly convex.Thus the proximity operator still results in a singleton. Lemma 16 allows us to again reformulate the original algorithm using the gradient of Moreau envelope. The only difference from the convex regime is that the Moreau envelope M γ fi is not necessarily convex. The following lemmas illustrate the connection between M γ fi and f i : Lemma 17. [Yu et al., 2015, Proposition 7] Let γ > 0, f be a closed, proper function that is lower bounded. ThenM γ f ≤ f , inf M γ f = inf f , arg min x M γ f (x) = arg min x f (x) ⊆ x : x ∈ prox γf (x) .Lemma 18. Let f : R d → R be ρ-weakly convex with ρ > 0 and differentiable. If we take 0 < γ < 1 ρ , then M γ fi has the same set of stationary points as f i .For the sake of simplicity, we will consider only the full participation case with a constant extrapolation parameter α k = α. 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 η satisfying 0 < η < 1 L . Then for any K ≥ 1, the iterates x k of GD satisfymin 0≤k≤K-1 E ∥∇f (x k )∥ 2 ≤ 2 (f (x 0 ) -inf f ) ηK .Now we directly apply Lemma 19 in our case,1. Since each M γ fi is L γ,ρ -smooth, M γ is L γ -smooth with L γ ≤ L γ,ρ, which result in the following bound on the extrapolation parameter0 < α < 1 γL γ .Notice that in this case we have the following estimation of γL γ ,1 γL γ ≥ 1 γL γ,ρ = min 1, 1 -γρ γρ .This suggests that extrapolation may not be much beneficiary in the non-convex case.2. The following convergence guarantee can be obtained.min 0≤k≤K-1 E ∥∇M γ (x k )∥ 2 ≤ 2 (M γ (x 0 ) -inf M γ ) αγK .Notice that by Lemma 17, we know thatM γ fi (x 0 ) ≤ f i (x 0 ). We also have inf M γ ≥ 1 n n i=1 inf M γ fi = 1 n n i=1 inf f i since inf M γ fi = inf f i istrue for each client by Lemma 17. Thus, we haveM γ (x 0 ) -inf M γ ≤ f (x 0 ) -inf f + inf f - 1 n n i=1 inf f i .We can relax the above convergence guarantee and obtainmin 0≤k≤K-1 E ∥∇M γ (x k )∥ 2 ≤ 2 (f (x 0 ) -inf f ) αγK + 2 inf f -1 n n i=1 inf f i αγK .The above convergence guarantee indicates that the algorithm converges to some stationary points of M γ (x) in the non-convex case.3. In the non-convex case, we did not assume anything similar to the interpolation regime in the convex case. As a result, we did not know the relation between the set of stationary points of M γ (x) and f (x), denoted as Y ′ and Y, respectively. However, if we assume, in addition, that each stationary point y ′ ∈ Y ′ of M γ is also a stationary point of each M γ fi , then y ′ is also a stationary point of f i according to Lemma 18. Thus, ∇f (y ′ ) = 1 n n i=1 ∇f i (y ′ ) = 0, which indicates y ′ ∈ Y. As a result, we have Y ′ ⊆ Y. This means that under this additional assumption, the algorithm converges to a stationary point of f .Notes on gradient diversity variant. In general, the gradient diversity step size η k used in SGD to solve the finite sum minimization problemmin x∈R d f (x) = 1 n n i=1 f i (x) ,can be written asη k := 1 L max • 1 n n i=1 ∥∇f i (x k )∥ 2 1 n n i=1 ∇f i (x k ) 2 ,where L max is the maximum of local smoothness constants. In our case, since each local Moreau envelope is Li 1+γLi -smooth and 1 γ -smooth4 , we can use both Lmax 1+γLmax (here in Corollary 2, if we know L max ) and 1 γ (in original Theorem 2, if we do not know L max ) as the maximum of local smoothness. We present the convergence result of Algorithm 1 with the following rule given in (13),α ′ k,G = 1 + γL max γL max • 1 n n i=1 x k -prox γfi (x k ) 2 1 n n i=1 x k -prox γfi (x k ) 2 .Corollary 2. Suppose all the assumptions mentioned in Theorem 2 hold, if we are using (13) to determine α ′ k,G in each iteration for Algorithm 1 with τ = n, then the iterates satisfyE [f (x K )] -f inf ≤ 1 γ + L max • ∥x 0 -x ⋆ ∥ 2 K-1 k=0 α ′ k,G.where xK is chosen randomly from the first K iterates {x 0 , x 1 , ..., x K-1 } with probabilitiesp k = α ′ k,G/ K-1 k=0 α ′ k,G .Notice that compared to the case of FedExProx-GraDS in Theorem 2, the convergence rate given in Corollary 2 is indeed better. This can be seen by comparing them directly, for FedExProx-GraDS, we haveE [f (x K )] -inf f ≤ 1 + γL max 2 + γL max • 1 γ + L max • ∥x 0 -x ⋆ ∥ 2 K-1 k=0 α k,G, and for Algorithm 1 with α ′ k,G given in (13), we haveE [f (x K )] -f inf ≤ 1 γ + L max • ∥x 0 -x ⋆ ∥ 2 K-1 k=0 α ′ k,G = γL max 1 + γL max • 1 γ + L max • ∥x 0 -x ⋆ ∥ 2 K-1 k=0 α k,G . Since γL max 1 + γL max ≤ 1 + γL max 2 + γL max , ∀γ > 0,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.Notes on stochastic Polyak variant. In this paragraph, we further elaborate on the convergence of FedExProx-StoPS. We start by providing a lower bound on the adaptive extrapolation parameter. Lemma 20. Suppose that all assumptions mentioned in Theorem 2 hold, then the following inequalities hold for any x ∈ R d and x ⋆ that is a minimizer of f ,1 n n i=1 M γ fi (x) -M γ fi (x ⋆ ) γ • 1 n n i=1 ∇M γ fi (x) 2 ≥ 1 2γL γ .Using the above lower bound, we can further write the convergence of FedExProx-StoPS asE f (x K ) -inf f ≤ 2L γ (1 + 2γL max ) • ∥x 0 -x ⋆ ∥ 2 K .Observe that we recover the favorable dependence of convergence on the smoothness of M γ . However, this comes at the price of having to know each M γ fi (x ⋆ ) or, equivalently in the interpolation regime, knowing M γ (x ⋆ ).In this subsection, we extend the adaptive variants of FedExProx into the stochastic setting. We will refer to them as FedExProx-GraDS-PP and, FedExProx-StoPS-PP respectively. Specifically, we consider that the server chooses the client using the τ -nice sampling strategy we have introduced before in Algorithm 1. The following theorem summarizes the convergence guarantee of FedExProx-GraDS-PP and FedExProx-StoPS-PP in the convex case. Its extension to the strongly convex case where we additionally assume Assumption 5 (Strong convexity) is straight forward. Theorem 5. Suppose Assumption 1 (Differentiability), Assumption 2 (Interpolation regime), Assumption 3 (Convexity) and Assumption 4 (Smoothness) hold. Assume we are running FedExProx with τ -nice client sampling.(i) (FedExProx-GraDS-PP): If we are using α k = α τ,k,G (x k , S k ), whereEQUATION)Then the iterates of Algorithm 1 satisfyEQUATION)where K is the total number of iteration, xK is samples uniformly at random from the firstK iterates {x 0 , x 1 , . . . , x K-1 }, inf α τ,k,G is defined as inf α τ,k,G := inf x∈R d ,S⊆[n],|S|=τ α τ,k,G (x, S) . satisfying α τ,k,G (x k , S k ) ≥ inf α τ,k,G ≥ 1.(ii) (FedExProx-StoPS-PP): If we are using α k = α τ,k,S (x k , S k ), whereEQUATION)Then the iterates of Algorithm 1 satisfyEQUATION)where K is the total number of iteration, xK is sampled uniformly at random from the first K iterates {x 0 , x 1 , . . . , x K-1 }, inf α τ,k,G is defined asinf α τ,k,S := inf x∈R d ,S⊆[n],|S|=τ α τ,k,S (x, S) . satisfying α τ,k,S (x k , S k ) ≥ inf α τ,k,S ≥ 1 2 1 + 1 γL max .Remark 15. For FedExProx-GraDS-PP, different from the full participation setting, the denominator of the sublinear term on the right-hand side of (20) is replaced by K • inf α τ,k,G .(i) In the single client case (τ = 1), we haveα 1,k,G = inf α 1,k,G = 1.(ii) In the partial participation case (1< τ < n), it is possible that inf α τ,k,G > 1,resulting in acceleration compared to single client case.(iii) For the full participation case (τ = n), we haveEQUATIONK-1 k=0 α k,G ≥ K • inf α n,k,G ,thus the convergence guarantee here is a relaxed version of that presented in Theorem 2.A similar discussion also applies to FedExProx-StoPS-PP in the client partial participation setting.Remark 16. For FedExProx-StoPS-PP, different from the full participation setting, the denominator of the sublinear term on the right-hand side of (22) is replaced by K • inf α τ,k,S .(i) In the single client case (τ = 1), we haveα 1,k,S ≥ inf α 1,k,G = 1 2 1 + 1 γL max .(ii) In the partial participation case (1< τ < n), it is possible that inf α τ,k,S > 1 2 1 + 1 γL max ,resulting in acceleration compared to single client case.(iii) For the full participation case (τ = n), we haveEQUATIONK-1 k=0 α k,S ≥ K • inf α n,k,S ,thus the convergence guarantee here is a relaxed version of that presented in Theorem 2.The following Table 5 summarizes the convergence of new algorithms and their variants appeared in our paper.The proof of this theorem can be divided into three parts. FedExProx Lγ (1+γLmax) /K Lγ,τ (1+γLmax) /K Lmax /K FedExProx-GraDS (1+γLmax) /γ• K-1 k=0 α k,G (1+γLmax) /(γK•inf α τ,k,G) (1+γLmax) /(γK) FedExProx-StoPS (1+γLmax) /γ• K-1 k=0 α k,S (1+γLmax) /(γK•inf α τ,k,S ) (1+γLmax) /(γK•inf α 1,k,S )Step 1: Reformulate the algorithm using Moreau envelope. We know from Lemma 2 that for any x ∈ R d .∇M γ fi (x) = 1 γ x -prox γfi (x) .Using the above identity, we can rewrite the update rule given in (7) in the following form,EQUATION)The above reformulation suggests that running FedExProx with τ -nice sampling strategy is equivalent to running SGD with τ -nice sampling to the global objective M γ (x) = 1 n n i=1 M γ fi (x) with step size α k γ. Now, it seems natural to 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 M γ fi .1. Each M γ fi (x) is convex. This is a consequence of a direct application of Lemma 3 to each f i . Since M γ is the average of convex functions M γ fi , we conclude that M γ (x) is also convex.2. Each M γ fi (x) is Li 1+γLi -smooth, where L i is the smoothness constant of f i . This is proved by applying Lemma 4 to each f i . Drawing on Lemma 13 for justification, it is reasonable to assumeM γ (x) is L γ -smooth with L γ ≤ 1 n n i=1 Li 1+γLi -smooth. 3. Each M γ fi (x)has the same set of minimizers and minimum as f i . This result arises from applying Lemma 5 to each function f i . 4. Furthermore, if Assumption 2 (Interpolation Regime) holds, M γ (x) and f (x) have the same set of minimizers and minimum. This is demonstrated in Lemma 8.Step 2: Applying the theory of gradient type methods. Notice that here M γ fi (x) is Li 1+γLi -smooth and convex, M γ (x) is convex and L γ -smooth. Furthermore, due to the assumption of interpolation regime, M γ (x) and f (x) have the same set of minimizers. Applying the theory of SGD with τ -nice sampling in this case, whereA τ = L γ,τ = n -τ τ (n -1) • max i∈[n] L i 1 + γL i + n(τ -1) τ (n -1) L γ .Notice that using Fact 4, we know thatmax i∈[n] L i 1 + γL i Fact 4 = L max 1 + γL max ,thus L γ can be simplified and written asL γ,τ = n -τ τ (n -1) • L max 1 + γL max + n(τ -1) τ (n -1) L γ ,where L max = max i L i . We obtain the following result given that 0 < αγ < 2 Lγ,τ in the convex setting,E [M γ (x K )] -M γ (x ⋆ ) Theorem 3 ≤ 1 αγ(2 -αγL γ,τ ) • ∥x 0 -x ⋆ ∥ 2 K ,where xK is sampled uniformly at random from the first K iterates {x 0 , x 1 , . . . , x K-1 }. However, the convergence mentioned pertains to M γ (x). 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,M γ (x K ) -M γ (x ⋆ ) (18) ≥ 1 1 + γL max (f (x K ) -f (x ⋆ )) ,to obtain the following resultE [f (x K )] -f (x ⋆ ) ≤ 1 + γL max αγ (2 -αγL γ,τ ) • ∥x 0 -x ⋆ ∥ 2 K .Observe that we haveC (γ, τ, α) = 1 + γL max αγ (2 -αγL γ,τ ) ,and its numerator does not depend on α. If we fix the choice of γ and τ , then the denominator is maximized when αγL γ,τ = 1. This yields the optimal constant extrapolation parameter α γ,τ = 1 γLγ,τ and the following convergence corresponding to itE [f (x K )] -f (x ⋆ ) ≤ L γ,τ (1 + γL max ) • ∥x 0 -x ⋆ ∥ 2 K .Finally, notice thatγL γ Lemma 13 ≤ 1 n n i=1 γL i 1 + γL i < 1,for any γ > 0. This suggests that,γL γ,τ = n -τ τ (n -1) • γL max 1 + γL max + n(τ -1) τ (n -1) γL γ < n -τ τ (n -1) + n(τ -1) τ (n -1) = 1,which in turn tells us α γ,τ = 1 γLγ,τ > 1. This concludes the proof.We start with the following decomposition,EQUATION)Case 1: FedExProx-GraDS For gradient diversity based α k , we haveα k = α k,G = 1 n n i=1 γ∇M γ fi (x k ) 2 ∥γ∇M γ (x k )∥ 2 = 1 n n i=1 ∇M γ fi (x k ) 2 ∥∇M γ (x k )∥ 2 .For the last term of (24),α 2 k,G γ 2 ∥∇M γ (x k )∥ 2 = α k,G γ 2 • 1 n n i=1 ∇M γ fi (x k ) 2 = α k,G γ 2 • 1 n n i=1 ∇M γ fi (x k ) -∇M γ fi (x ⋆ ) 2 (16) ≤ α k,G γ 2 • 1 n n i=1 L i 1 + γL i D M γ f i (x k , x ⋆ ) + D M γ f i (x ⋆ , x k ) ,where the last inequality follows from the Li 1+γLi -smoothness of M γ fi given in Lemma 4. We further obtain using Fact 4 thatEQUATION)For the second term of ( 24), we haveEQUATION)Plugging ( 26) and ( 25) into ( 24), we have∥x k+1 -x ⋆ ∥ 2 ≤ ∥x k -x ⋆ ∥ 2 -α k,G γ 2 - γL max 1 + γL max (D M γ (x k , x ⋆ ) + D M γ (x ⋆ , x k )) .Notice that we know thatD M γ (x k , x ⋆ ) (14) = M γ (x k ) -M γ (x ⋆ ) , D M γ (x ⋆ , x k ) (15) ≥ 0.As a result, we obtain∥x k+1 -x ⋆ ∥ 2 ≤ ∥x k -x ⋆ ∥ 2 -α k,G γ 2 - γL max 1 + γL max (M γ (x k ) -M γ (x ⋆ )) .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γ 2 - γL max 1 + γL max K-1 k=0 α k,G (M γ (x k ) -inf M γ ) ≤ ∥x 0 -x ⋆ ∥ 2 . Denote p k = α k,G/ K-1 k=0 α k,G for k = 0, 1, ..., K -1.If we pick xK randomly according to probabilities p k from the first K iterates {x 0 , x 1 , . . . , x K-1 }, then we can further write the above recursion asE M γ xK -inf M γ ≤ 1 + γL max 2 + γL max • 1 γ • ∥x 0 -x ⋆ ∥ 2 K-1 k=0 α k,G.Utilizing Lemma 10, we further obtain,E [f (x K )] -inf f ≤ 1 + γL max 2 + γL max • 1 γ + L max • ∥x 0 -x ⋆ ∥ 2 K-1 k=0 α k,G.The above inequality indicates convergence. Indeed, by convexity of standard Euclidean norm, we haveα k,G ≥ 1 n n i=1 x k -prox γfi (x k ) 2 1 n n i=1 x k -prox γfi (x k ) 2 = 1.This tells us thatK-1 k=0 α k,G ≥ K.Case 2: FedExProx-StoPS For stochastic Polyak step size based α k,S , since we are in the interpolation regime, by Lemma 9, we haveM γ (x ⋆ ) = inf M γ = 1 n n i=1 inf M γ fi .As a result,α k = α k,S = 1 n n i=1 M γ fi (x k ) -inf M γ fi γ 1 n n i=1 ∇M γ fi (x k ) 2 = M γ (x k ) -M γ (x ⋆ ) γ ∥∇M γ (x k )∥ 2 .We have for the last term of ( 24),EQUATION)For the second term of ( 24), we haveEQUATION)where the inequality is due to the convexity of M γ . Plugging ( 28) and ( 27) into ( 24), we obtain∥x k+1 -x ⋆ ∥ 2 ≤ ∥x k -x ⋆ ∥ 2 -α k,S γ (M γ (x k ) -M γ (x ⋆ )) .Summing up the above recursion for k = 0, 1, ..., K -1, we notice that many of them will telescope. Thus, we obtainγ K-1 k=0 α k,S (M γ (x k ) -inf M γ ) ≤ ∥x 0 -x ⋆ ∥ 2 .Denote p k = α k,S/ K-1 k=0 α k,S for k = 0, 1, ..., K -1. If we sample xK randomly according to probabilities p k from the first K iterates {x 0 , x 1 , . . . , x K-1 }, we can further write the above recursion asE M γ xK -inf M γ ≤ 1 γ • ∥x 0 -x ⋆ ∥ 2 K-1 k=0 α k,S.Utilizing the local bound in Lemma 10, we further obtain,EQUATIONNotice that the above inequality indeed indicates convergence, sinceK-1 k=0 α k,S = K-1 k=0 M γ (x k ) -M γ (x ⋆ ) γ ∥∇M γ (x k )∥ 2 ≥ 1 2γL γ ,where the inequality follows from Lemma 20. The above upper bounds allow us to further write the convergence in ( 29) asE f (x K ) -inf f ≤ 2L γ (1 + 2γL max ) • ∥x 0 -x ⋆ ∥ 2 K .This concludes the proof.We start from the decomposition,∥x k+1 -x ⋆ ∥ 2 = ∥x k -x ⋆ ∥ 2 -2η x k -x ⋆ , 1 τ i∈S k ∇f i (x k ) + η 2 1 τ i∈S k ∇f i (x k ) 2 ,where S k is the set sampled at iteration k. Taking expectation conditioned on x k , we haveE S k ∥x k+1 -x ⋆ ∥ 2 = ∥x k -x ⋆ ∥ 2 -2η ⟨x k -x ⋆ , ∇f (x k ) -∇f (x ⋆ )⟩ + η 2 E S k   1 τ i∈S k ∇f i (x k ) 2   .We can write the second inner product term asEQUATIONwhere D f (x k , x ⋆ ) denotes the Bregman divergence associated with f between x k and x ⋆ . For the last squared norm term, we first define the indicator random variable χ k,i asχ k,i = 1, when i ∈ S k , 0, when i / ∈ S k .Since we are in the interpolation regime, we haveEQUATION)For the first term above in (31), due to the smoothness and convexity of each f i , we have1 n n i=1 ∥a k,i ∥ 2 = 1 n n i=1 ∥∇f i (x k ) -∇f i (x ⋆ )∥ 2 ≤ 1 n n i=1 L i (D fi (x ⋆ , x k ) + D fi (x k , x ⋆ )) ≤ L max 1 n n i=1 (D fi (x ⋆ , x k ) + D fi (x k , x ⋆ )) = L max (D f (x ⋆ , x k ) + D f (x k , x ⋆ )), where the first inequality is obtained as a result of Fact 3. For the second term, we have due to the smoothness and convexity of f ,1 n n i=1 a k,i 2 = 1 n n i=1 (∇f i (x k ) -∇f i (x ⋆ )) 2 = ∥∇f (x k ) -∇f (x ⋆ )∥ 2 ≤ L (D f (x ⋆ , x k ) + D f (x k , x ⋆ )), where the inequality is obtained using Fact 3. Combining the above two inequalities and plugging them into (31), we obtainEQUATION)Notice that we already defined A τ asA τ = n -τ τ (n -1) • L max + n(τ -1) τ (n -1) • L.Combining ( 30) and ( 32), we haveE S k ∥x k+1 -x ⋆ ∥ 2 ≤ ∥x k -x ⋆ ∥ 2 -η(2 -ηA τ ) (D f (x ⋆ , x k ) + D f (x k , x ⋆ )) .If we require 0 < η < 2 Aτ , we have η(2 -ηA τ ) ≥ 0.Convex regime. It remains to notice thatD f (x k , x ⋆ ) + D f (x ⋆ , x k ) ≥ D f (x k , x ⋆ ) = f (x k ) - f (x ⋆ )≥ 0, and we haveE S k ∥x k+1 -x ⋆ ∥ 2 ≤ ∥x k -x ⋆ ∥ 2 -η(2 -ηA τ ) (f (x k ) -f (x ⋆ )) .Taking expectation again and using tower property, we getE ∥x k+1 -x ⋆ ∥ 2 ≤ E ∥x k -x ⋆ ∥ 2 -η(2 -ηA τ ) (E [f (x k )] -inf f ) .Unrolling this recurrence, we getE [f (x K )] -inf f ≤ 1 η(2 -ηA τ ) • ∥x 0 -x ⋆ ∥ 2 K ,where K is the total number of iterations, xK is selected uniformly at random from the first K iterates {x 0 , x 1 , . . . , x K-1 }.Star strongly convex regime. Due to star strong convexity of f , we further lower bound the Bregman divergenceD f (x k , x ⋆ ) = f (x k ) -f (x ⋆ ) ≥ µ 2 ∥x k -x ⋆ ∥ 2 .and we haveE S k ∥x k+1 -x ⋆ ∥ 2 ≤ 1 -η(2 -ηA τ ) • µ 2 ∥x k -x ⋆ ∥ 2 .Taking expectation again, using tower property we getE ∥x k+1 -x ⋆ ∥ 2 ≤ 1 -η(2 -ηA τ ) • µ 2 E ∥x k -x ⋆ ∥ 2 .Unrolling the recurrence, we getE ∥x K -x ⋆ ∥ 2 ≤ 1 -η(2 -ηA τ ) • µ 2 K ∥x 0 -x ⋆ ∥ 2 .This concludes the proof.Since each f i is proper, closed and convex, by Lemma 2, we know that eachM γ fi is 1 γ -smooth. Therefore, it is reasonable to assume that M γ = 1 n n i=1 M γ fi is L γ -smooth, with L γ ≤ 1 γ .Applying Theorem 3 in this case, we obtain,M γ (x K ) -inf M γ Theorem 3 ≤ 1 αγ (2 -αγL γ,τ ) • ∥x 0 -x ⋆ ∥ 2 K ,where xK is chosen uniformly at random from the first K iterates {x 0 , x 1 , . . . , x K-1 }, andL γ,τ = n -τ τ (n -1) • 1 γ + n(τ -1) τ (n -1) • L γ .Multiplying both sides by γ, we obtainγM γ (x K ) -inf γM γ ≤ 1 α (2 -αγL γ,τ ) • ∥x 0 -x ⋆ ∥ 2 K .It is easy to see that the coefficient on the right-hand side is minimized when α = 1 γLγ,τ , and the convergence is given byγM γ (x K ) -inf γM γ ≤ n -τ τ (n -1) + n(τ -1) τ (n -1) • γL γ • ∥x 0 -x ⋆ ∥ 2 K .Notice that L γ ≤ 1 γ . As a result,α ⋆ = 1 γL γ ≥ 1.Case of FedExProx-GraDS-PP. We start with the following identityEQUATION)For the last term, we haveα 2 τ,k,G • γ 2 • 1 τ i∈S k ∇M γ fi (x k ) 2 = α τ,k,G • γ 2 • 1 τ i∈S k ∇M γ fi (x k ) 2 = α τ,k,G • γ 2 • 1 τ i∈S k ∇M γ fi (x k ) -∇M γ fi (x ⋆ ) 2 ,where the last step is due to the assumption that we are in the interpolation regime. Using Fact 3, we can further upper bound the above expression,EQUATION)where the last inequality is due to Fact 4. Now we look at the second term in Equation ( 33).-EQUATION)Plugging ( 34) and ( 35) into (33), we obtain,EQUATIONwhere the last inequality is due toEQUATION≥ 0.Now we want to lower bound α τ,k,G , notice that it can be viewed as a function of the iterate x and the sampled set S. Therefore, we use the notationinf α τ,k,G = inf x∈R d ,S⊆[n],|S|=τ α τ,k,G (x, S) .As a result, we haveα τ,k,G ≥ inf α τ,k,G ≥ 1,where the second inequality comes from the convexity of standard Euclidean norm. Plugging this lower bound into (36), we obtain∥x k+1 -x ⋆ ∥ 2 ≤ ∥x k -x ⋆ ∥ 2 -inf α τ,k,G • γ 2 + γL max 1 + γL max • 1 τ i∈S k M γ fi (x k ) -M γ fi (x ⋆ ) .Taking expectation conditioned on x k , we haveE S k ∥x k+1 -x ⋆ ∥ 2 ≤ ∥x k -x ⋆ ∥ 2 -inf α τ,k,G • γ 2 + γL max 1 + γL max • 1 n n i=1 M γ fi (x k ) -M γ fi (x ⋆ ) = ∥x k -x ⋆ ∥ 2 -inf α τ,k,G • γ 2 + γL max 1 + γL max • (M γ (x k ) -inf M ) ,where the last identity is due to the fact that we are in the interpolation regime. Using Lemma 10, we haveE S k ∥x k+1 -x ⋆ ∥ 2 ≤ ∥x k -x ⋆ ∥ 2 -inf α τ,k,G • γ 2 + γL max 1 + γL max • 1 1 + γL max (f (x k ) -inf f ) .Taking expectation again and using tower property, we obtainE ∥x k+1 -x ⋆ ∥ 2 ≤ E ∥x k -x ⋆ ∥ 2 -inf α τ,k,G • γ 2 + γL max 1 + γL max • 1 1 + γL max E [f (x k ) -inf f ] .Following the same step as Theorem 1, we can unroll the above recurrence and obtainE [f (x K )] -inf f ≤ 1 + γL max 2 + γL max • 1 γ + L max • ∥x 0 -x ⋆ ∥ 2 inf α τ,k,G • K ,where K is the total number of iterations, xK is sampled uniformly at random from the first K-iterates {x 0 , x 1 , . . . , x K-1 }.Case of FedExProx-StoPS-PP. We start with the following identityEQUATION)For the last term of Equation ( 37), we haveEQUATIONWhile for the second term we haveEQUATIONPlugging ( 38) and ( 39) into (37), we obtainEQUATIONNow we want to lower bound α τ,k,S , notice that it can be viewed as a function of the iterate x and the sampled set S. Therefore, we use the notationinf α τ,k,S = inf x∈R d ,S⊆[n],|S|=τ α τ,k,S (x, S) .As a result, we have α τ,k,S ≥ inf α τ,k,S .Notice that since each M γ fi is Li 1+γLi -smooth, we conclude that the function 1τ i∈S k M γ fi is at least Lmax 1+γLmax -smooth 5 .Using the smoothness of the mentioned function and Fact 3, a lower bound on inf α τ,k,S is obvious,inf α k,τ,S ≥ 1 2 • Lmax 1+γLmax γ = 1 2 1 + 1 γL max .This means that we haveα τ,k,S ≥ inf α τ,k,S ≥ 1 2 1 + 1 γL max .Using the above lower bound in (40), we have∥x k+1 -x ⋆ ∥ 2 ≤ ∥x k -x ⋆ ∥ 2 -inf α τ,k,S • γ • 1 τ i∈S k M γ fi (x k ) -inf M γ fi .Taking expectation conditioned on x k , and noticing that we are in the interpolation regime, we obtainE S k ∥x k+1 -x ⋆ ∥ 2 ≤ ∥x k -x ⋆ ∥ 2 -inf α τ,k,S • γ • (M γ (x k ) -inf M ) .Using Lemma 10, we haveE S k ∥x k+1 -x ⋆ ∥ 2 Lemma 10 ≤ ∥x k -x ⋆ ∥ 2 -inf α τ,k,S • γ 1 + γL max • (f (x k ) -inf f ) .Now, following the exact same steps as in the previous case of FedExProx-GraDS, we result inE [f (x K )] -inf f ≤ 1 γ + L max • ∥x 0 -x ⋆ ∥ 2 inf α τ,k,S • K ,where K is the total number of iterations, xK is sampled uniformly at random from the first K-iterates {x 0 , x 1 , . . . , x K-1 }.If additionally we assume f is µ-strongly convex, then from Lemma 11, we know it indicates the following star strong convexity of M γ holds,M γ (x) -M γ (x ⋆ ) ≥ µ 1 + γL max • 1 2 ∥x -x ⋆ ∥ 2 .Thus, we apply Theorem 3 with τ -nice sampling in the star strong convexity case, and obtain the following result:E ∥x K -x ⋆ ∥ 2 Theorem 3 ≤ 1 -αγ(2 -αγL γ,τ ) • µ 2 (1 + γL max ) K ∥x 0 -x ⋆ ∥ 2 .Since the convergence here is stated in terms of squared distance to the minimizer, we do not need further transformation. Notice that the convergence rate in this case,1 -αγ(2 -αγL γ,τ ) • µ 2 (1 + γL max ) ,is also minimized when α = α γ,τ = 1 γLγ,τ . In case of α = α γ,τ , the convergence is given byE ∥x K -x ⋆ ∥ 2 ≤ 1 - µ 2L γ,τ (1 + γL max ) K ∥x 0 -x ⋆ ∥ 2 .This concludes the proof.Similar to the proof of Theorem 2, we start with the following identityEQUATION)The extrapolation parameter can be rewritten asα ′ k,G = 1 + γL max γL max • 1 n n i=1 ∇M γ fi (x k ) 2 ∥∇M γ (x k )∥ 2 .We have for the last term of (41),α ′ k,G 2 γ 2 ∥∇M γ (x k )∥ 2 = α ′ k,G γ • γ + 1 L max 1 n n i=1 ∇M γ fi (x k ) 2 = α ′ k,G γ • γ + 1 L max • 1 n n i=1 ∇M γ fi (x k ) -∇M γf i x ⋆ 2 ≤ α ′ k,G γ • γ + 1 L max • 1 n n i=1 L i 1 + γL i D M γ f i (x k , x ⋆ ) + D M γ f i (x ⋆ , x k ) ,where the last inequality follows from the Li 1+γLi -smoothness of M γ fi . Utilizing the monotonicity of x 1+γx , for x > 0, we further obtainα ′ k,G 2 γ 2 ∥∇M γ (x k )∥ 2 Fact 4 ≤ α ′ k,G γ • γ + 1 L max • L max 1 + γL max • 1 n n i=1 D M γ f i (x k , x ⋆ ) + D M γ f i (x ⋆ , x k ) = α ′ k,G γ • γ + 1 L max • L max 1 + γL max • (D M γ (x k , x ⋆ ) + D M γ (x ⋆ , x k )) = α ′ k,G γ (D M γ (x k , x ⋆ ) + D M γ (x ⋆ , x k )) .(42) For the second term of (41), we have-2α ′ k,G γ ⟨∇M γ (x k ) , x k -x ⋆ ⟩ = 2α ′ k,G γ ⟨∇M γ (x k ) , x ⋆ -x k ⟩ = 2α ′ k,G γ ⟨∇M γ (x k ) -∇M γ (x ⋆ ) , x ⋆ -x k ⟩ = -2α ′ k,G γ (D M γ (x k , x ⋆ ) + D M γ (x ⋆ , x k )) .(43) Plugging ( 43) and ( 42) into (41), we obtain∥x k+1 -x ⋆ ∥ 2 ≤ ∥x k -x ⋆ ∥ 2 -α ′ k,G γ (D M γ (x k , x ⋆ ) + D M γ (x ⋆ , x k )). Notice that we know thatD M γ (x k , x ⋆ ) (14) = M γ (x k ) -M γ (x ⋆ ) , D M γ (x ⋆ , x k ) (15) ≥ 0. As a result, we have ∥x k+1 -x ⋆ ∥ 2 ≤ ∥x k -x ⋆ ∥ 2 -α ′ k,G γ (M γ (x k ) -M γ (x ⋆ )). Summing up the above recursion for k = 0, 1, ..., K -1, we notice that many of them telescope, we obtainγ K-1 k=0 α ′ k,G (M γ (x k ) -inf M γ ) ≤ ∥x 0 -x ⋆ ∥ 2 . Denote p k = α ′ k,G/ K-1 k=0 α ′k,G for k = 0, 1, ..., K -1. If we sample xK randomly according to probabilities p k from the first K iterates {x 0 , x 1 , . . . , x K-1 }, we can further write the above recursion asE [M γ (x K )] -inf M γ ≤ 1 γ • ∥x 0 -x ⋆ ∥ 2 K-1 k=0 α ′ k,G.Utilizing the local bound in Lemma 10, we further obtain,E f (x K ) -f inf ≤ 1 γ + L max • ∥x 0 -x ⋆ ∥ 2 K-1 k=0 α ′ k,G.This concludes the proof.Notice that since f is proper, closed and convex, by Fact 1, prox γf (x) is a singleton. We use the notation z(x) = prox γf (x). Using the definition of prox γf (x), we see thatM γ f (x) = f (z(x)) + 1 2γ ∥z(x) -x∥ 2 = f prox γf (x) + 1 2γ prox γf (x) -x 2 . Now, assume M γ f (x) = +∞. We have for any z ∈ R d , +∞ = M γ f (x) = f (z(x)) + 1 2γ ∥z(x) -x∥ 2 ≤ f (z) + 1 2γ ∥z -x∥ 2 ,which means that z is also optimal, which contradicts the uniqueness z(x) = prox γf (x). This indicates that M γ f (x) < +∞, thus, it is real-valued, which concludes the proof.Let f ⋆ be the convex conjugate of f , using Corollary 6.56 in the book by Beck [2017] , we haveM γ f ⋆ = f ⋆ + γ 2 ∥•∥ 2 .We know that the convex conjugate of a proper, closed and convex function is also proper closed and convex. As a result,f ⋆ + γ 2 ∥•∥ 2 is γ-strongly convex. This indicates that M γ f ⋆ is γ-strongly convex, which implies M γ f is 1 γ -smooth. Notice that we have prox γf (x) = arg min z∈R d f (z) + 1 2γ ∥z -x∥ 2 ,by the definition of proximity operator. Using Theorem 5.30 from Beck [2017] , we have∇M γ f (x) = 1 γ x -prox γf (x) .This completes the proof.To prove this lemma, we use Theorem 2.19 in the book by Beck [2017] . From the key observation that M γ f is the infimal convolution of the proper, convex function f and the real-valued convex function1 2γ ∥•∥ 2 ,we deduce that M γ f is convex. This completes the proof.Let f ⋆ be the convex conjugate of f . From Corollary 6.56 in the book by Beck [2017] , it holds thatM γ f ⋆ = f ⋆ + γ 2 ∥•∥ 2 . Since f is L-smooth, we deduce that f ⋆ is 1 L -strongly convex, and thus M γ f ⋆ is 1 L + γ-strongly convex. This suggests that M γ f ⋆ is 1+γL L -strongly convex, which in turn implies M γ f is L 1+γL -smooth.This completes the proof.Notice that since M γ f is convex and differentiable, the condition ∇M γ f (x) = 0 gives its set of minimizers. This optimality condition can be written exactly as x = prox γf (x) according to Lemma 2. Using Lemma 12, we know this condition also gives the set of minimizers of f , which suggests that f and M γ f have the same set of minimizers. Pick any x ⋆ ∈ R d that is a minimizer of f , using Lemma 1, we haveinf M γ f = M γ f (x ⋆ ) = f prox γf (x ⋆ ) + 1 2γ x ⋆ -prox γf (x ⋆ ) 2 = f (x ⋆ ) = inf f.This completes the proof.For any x ∈ R d , we haveM γ f (x) = min z∈R d f (z) + 1 2γ ∥z -x∥ 2 ≤ f (x) + 1 2γ ∥x -x∥ 2 = f (x).This completes the proof.From Lemma 3 and Lemma 4, we immediately obtain that each M γ fi is convex and Li 1+γLi -smooth. This immediately suggests that M = 1 n n i=1 M γ fi is convex and L γ -smooth withL γ ≤ 1 n n i=1 L i 1 + γL i .Then by Lemma 13, we have1 n 2 n i=1 L i 1 + γL i Lemma 13 ≤ L γ .Combing the above two inequalities, we have1 n 2 n i=1 L i 1 + γL i ≤ L γ ≤ 1 n n i=1 L i 1 + γL i .We then look at the condition number defined in Theorem 1. It is easy to verify thatC (γ, n, α γ,n ) = L γ (1 + γL max ) and, C (γ, 1, α γ,1 ) = L max . As a result, C (γ, n, α γ,n ) = L γ (1 + γL max ) ≤ 1 n n i=1 L i • 1 + γL max 1 + γL i ≤ L max = C (γ, n, 1) , Notice that we can write C (γ, τ, α γ,τ ) as an interpolation between C (γ, n, α γ,n ) and C (γ, 1, α γ,1 ), thereforeL γ (1 + γL max ) ≤ C (γ, n, α γ,n ) ≤ C (γ, τ, α γ,τ ) ≤ C (γ, 1, α γ,1 ) = L max .In cases where there exists at least one L i < L max , we have1 n n i=1 L i • 1 + γL max 1 + γL i < L max .which is true for all 0 < γ < +∞. Thus, C (γ, n, α γ,n ) < C (γ, τ, α γ,τ ) < L max = C (γ, 1, α γ,1 ). Now we give an example that when allL i = L max , still C (γ, n, α γ,n ) = 1 n C (γ, 1, α γ,1 ) = 1 n L max .Example 1. Consider the setting where f i : R d → R is defined as f i (x) = θ 2 x 2 i for some θ > 0. Here x i denotes the i-th coordinate of the vector x ∈ R d , f : R d → R is given by f (x) = θ 2n ∥x∥ 2 . It is easy to show that for each f i is a convex, θ-smooth function and the smoothness constant θ cannot be improved sinceθ 2 ∥x∥ 2 - θ 2 x 2 i = θ 2 j̸ =i x 2 j .For f (x) = θ 2n ∥x∥ 2 , apparently, it is θ n -smooth and convex. We have the following formula for the Moreau envelope of f i (x),M γ fi (x) = 1 2 • θ 1 + γθ • x 2 i .As expected, each one of them is convex and θ 1+γθ -smooth. For M γ (x), it is given byM γ (x) = 1 n n i=1 M γ fi (x) = 1 2 • θ n(1 + γθ) • ∥x∥ 2 ,thus, we know it is convex and L γ = θ n(1+γθ) -smooth. In this caseL max L γ (1 + γL max ) = θ θ n(1+γθ) • (1 + γθ) = n, which is L γ (1 + γL max ) = C (γ, n, α γ,n ) = 1 n C (γ, 1, α γ,1 ) = 1 n L max .H.8 Proof of Lemma 8By Lemma 5, we know that f i and M γ fi have the same set of minimizers and minimum. Denote the set of minimizers as X i , since we are in the interpolation regime, we know that the set of minimizers of f is given by,X = n i=1 X i ̸ = ∅.Now we prove that every x in X is a minimizer of M = 1 n n i=1 M γ fi . This is true since x ∈ X minimizes each f i , thus M γ fi at the same time. The minimum is given byinf M = 1 n n i=1 inf M γ fi = 1 n n i=1 inf f i = inf f.We then prove that every x / ∈ X is not a minimizer of f . If x / ∈ X , there exists at least one set X j such that x / ∈ X j . Thus M γ fj (x) > inf M γ fj . This indicates that M γ (x) > inf M , which means, x / ∈ X is not a minimizer of M .H.9 Proof of Lemma 9From Lemma 6, it is clear that M γ f is a global lower bound of f satisfying M γ f (x) ≤ f (x) for any x ∈ R d and γ > 0. Notice that the definition of M γ indicates thatM γ (x) = 1 n n i=1 M γ fi (x) = 1 n n i=1 min zi∈R d f i (z i ) + 1 2γ ∥z i -x∥ 2 ≤ min z∈R d 1 n n i=1 f i (z) + 1 2γ ∥z -x∥ 2 = min z∈R d 1 n n i=1 f i (z) + 1 2γ ∥z -x∥ 2 = M γ f (x), holds for any x ∈ R d and γ > 0. Combining the above result, we have M γ (x) ≤ M γ f (x) ≤ f (x) for any x ∈ R d and γ > 0. Notice that in Lemma 8, 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 M γ f and f have the same set of minimizers and minimum. Therefore, combining the above statement, we know that M γ , M γ f and f have the same set of minimizers and minimum. Thus, for any x ⋆ belongs to the set of minimizers, we haveM γ (x ⋆ ) = M γ f (x ⋆ ) = f (x ⋆). This completes the proof.We start from noticing that according to Lemma 1, the following identity is true for Moreau envelope,EQUATION)For the second squared norm term, we have the following inequality due to the smoothness of each f i and the fact that∇f i prox γfi (x) = 1 γ x -prox γfi (x) , x -prox γfi (x) 2 = x -prox γfi (x) , x -prox γfi (x) = γ ∇f i (prox γfi (x)), x -prox γfi (x) ≥ γ f i (x) -f i prox γfi (x) - γL i 2 x -prox γfi (x) 2 ,which leads to the following lower bound:x -prox γfi (x) 2 ≥ 1 1 γ + Li 2 f i (x) -f i prox γfi (x) .Plug in the above inequality into (44) and notice that inf M = 1 n n i=1 inf M γ fi = 1 n n i=1 inf f i , we obtain the following lower bound on M γ fi (x),EQUATIONNow let us look at the term f i prox γfi (x) -inf f . Using again L i -smoothness of f i , we havef i (x) -f i (prox γfi (x)) -∇f i (prox γfi (x)), x -prox γfi (x) ≤ L i 2 x -prox γfi (x) 2 .Notice that x -prox γfi (x) = γ∇f i (prox γfi (x)). As a result, we have,f i (x) -γ ∇f i (prox γfi (x)) 2 - L i γ 2 2 ∇f i (prox γfi (x)) 2 ≤ f i (prox γfi (x)),which isf i (x) -inf f i -γ + γ 2 L i 2 ∇f i (prox γfi (x)) 2 ≤ f i prox γfi (x) -inf f i .Using the interpolation regime assumption, we have∇f i prox γfi (x) 2 = ∇f i prox γfi (x) -∇f i (x ⋆ ) 2 ≤ 2L i D fi prox γfi (x) , x ⋆ = 2L i f i (prox γfi (x)) -inf f i, where the inequality is obtained using Fact 3. As a result, we obtain the following bound,f i prox γfi (x) -inf f i ≥ 1 1 + γL i (2 + γL i ) (f i (x) -inf f i ) = 1 (1 + γL i ) 2 (f i (x) -inf f i ) .Plug the above lower bound into (45), we obtainEQUATIONNotice that we have M γ (x) = 1 n n i=1 M γ fi (x). Since we are in the interpolation regime, from Lemma 9, we know thatinf M γ = M γ (x ⋆ ) = 1 n n i=1 M γ fi (x ⋆ ) = 1 n n i=1 inf M γ fi ,andinf f = f (x ⋆ ) = 1 n n i=1 f i (x ⋆ ) = 1 n n i=1 inf f i .We average (46) for each i ∈ [n] and obtainM γ (x) -inf M γ ≥ 1 n n i=1 1 1 + γL i (f i (x) -inf f i ) ≥ 1 1 + γL max • 1 n n i=1 (f i (x) -inf f i ) = 1 1 + γL max (f (x) -inf f ) .This concludes the proof.Combining ( 49) and ( 50), we have1 n 2 n i=1 L i ≤ L ≤ 1 n n i=1 L i .In order to demonstrate that both bounds are tight in the above inequality, we consider cases where they are identities.(i): Consider the case that each functionf i (x) = 1 2 • L i • ∥x∥ 2 , it is easy to see that f (x) = 1 2 • 1 n n i=1 L i • ∥x∥ 2 .In this case L = 1 n n i=1 L i , the upper bound is an identity.(ii): Consider the case that each functionf i (x) = 1 2 • θ • x 2 i , where θ > 0 is a constant, x i is the i-th coordinate of x ∈ R d . In this case f (x) = 1 2 • θ n ∥x∥ 2 .It is easy to verify that in thiscase L i = θ, L = 1 n θ. Thus 1 n 2 n i=1 L i = L, the lower bound is an identity.This concludes the proof.From the definition of C(γ, τ, 1) and C (γ, τ, α γ,τ ), we know thatC(γ, τ, 1) C (γ, τ, α γ,τ ) = 1 γL γ,τ (2 -γL γ,τ ).Now let t = γL γ,τ , we have the following bound on t according to the definition of L γ,τ given in Theorem 1.t = γL γ,τ = n -τ τ (n -1) • γL max 1 + γL max + n(τ -1) τ (n -1) • γL γ .Notice that in Lemma 7, we have shown thatL γ Lemma 7 ≤ 1 n n i=1 L i 1 + γL i ,and due to Fact 4, we have1 n n i=1 L i 1 + γL i Fact 4 ≤ L max 1 + γL max .As a result,t ≤ n -τ τ (n -1) • γL max 1 + γL max + n(τ -1) τ (n -1) • γL max 1 + γL max = γL max 1 + γL max < 1.It is easy to show that g(t) = 1 t(2-t) is monotone decreasing when t ∈ [0, 1], thusC(γ, τ, 1) C (γ, τ, α γ,τ ) ≥ 1 γLmax 1+γLmax 1 -γLmax 1+γLmax = 2 + 1 γL max + γL max AM-GM ≥ 4,where the last inequality is due to the AM-GM inequality. This concludes the proof.As suggested by Lemma 7, we haveC (γ, n, α γ,n ) ≤ C (γ, τ, α γ,τ ) ≤ C (γ, 1, α γ,1 ) , ∀τ ∈ [n].Notice that α γ,τ is given byα γ,τ = 1 γL γ,τ ,and we know thatL γ,τ = n -τ τ (1 -n) • L max 1 + γL max + n(τ -1) τ (n -1) • L γ .From Lemma 7 and Fact 4, we know thatL γ Lemma 7 ≤ 1 n n i=1 L i 1 + γL i Fact 4 ≤ L max 1 + γL max .Consequently, L γ,τ decreases as τ increases. Therefore, α γ,τ increases with the increase of τ , as illustrated by the following inequalityα γ,1 ≤ α γ,τ ≤ α γ,n , ∀τ ∈ [n].This concludes the proof.We refer the readers to the proof of Lemma 3.1 of Böhm and Wright [2021] .H.17 Proof of Lemma 17We refer the readers to the proof of Proposition 7 of Yu et al. [2015] .Observe that since 0 < γ < 1 ρ , we do have f + 1 2 • 1 γ ∥•∥ 2 being strongly convex. This indicates that prox γf is always a singleton and therefore M γ f is differentiable, as suggested by Lemma 16. Notice that x is stationary point of M γ f if and only if ∇M γ f (x) = 0. This is equivalent to1 γ x -prox γf (x) = 0, which is x = prox γf (x). In addition, x = prox γf (x) is equivalent to ∇f (x) + 1 γ (x -x) = 0,which is ∇f (x) = 0. Combining the above statements, we have ∇f (x) = 0 if and only if ∇M γ f (x) = 0. This suggests that the two functions have the same set of stationary points.Apply Theorem 1 of Khaled and Richtárik [2023] , notice that in this case GD satisfy the expected smoothness assumption given in Assumption 2 of Khaled and Richtárik [2023] with A = 0, B = 1 and C = 0, we obtain that when the step size η satisfies0 < η < 1 LB = 1 L ,where L is the smoothness constant of f , the iterates of GD satisfymin 0≤k≤K-1 E ∥∇f (x k )∥ 2 ≤ 2 (f (x 0 ) -inf f ) ηK .This completes the proof.Notice that we are in the interpolation regime, by Lemma 8, we know that f and M γ have the same set of minimizers and minimum. As a result,EQUATIONFrom the above inequality, we obtain that1 n n i=1 M γ fi (x) -M γ fi (x ⋆ ) γ • 1 n n i=1 ∇M γ fi (x) 2 (51) = M γ (x) -M γ (x ⋆ ) γ • ∥∇M γ (x)∥ 2 .Then by the smoothness of M γ and Fact 3, we haveM γ (x) -M γ (x ⋆ ) γ • ∥∇M γ (x)∥ 2 Fact 3 ≥ 1 2Lγ ∥∇M γ (x) -∇M γ (x ⋆ )∥ 2 γ • ∥∇M γ (x)∥ 2 = 1 2γL γ .Thus, by combining the above inequalities, we have1 n n i=1 M γ fi (x) -M γ fi (x ⋆ ) γ • 1 n n i=1 ∇M γ fi (x) 2 ≥ 1 2γL γ .Notice that from the definition of α k,S for FedExProx-StoPS, we haveα k,S = 1 n n i=1 M γ fi (x k ) -M γ fi (x ⋆ ) γ • 1 n n i=1 ∇M γ fi (x k ) 2 ≥ 1 2γL γ .Therefore, using the above lower bound, it is straight forward to further relax ( 12) toE f (x K ) -inf f ≤ 2L γ (1 + 2γL max ) • ∥x 0 -x ⋆ ∥ 2 K .This concludes the proof.In this section, we describe the settings and results of numerical experiments to demonstrate the effectiveness of our method.We consider the overparameterized linear regression problem in the finite sum settingmin x∈R d f (x) = 1 n n i=1 f i (x) ,where d is the dimension of the problem, n is the total number of clients, each function f i has the following formf i (x) = 1 2 ∥A i x -b i ∥ 2 ,whereA i ∈ R ni×d , b i ∈ R ni .Here n i is the number of samples on each client. It is easy to see that for each function f i , we have∇f i (x) = A ⊤ i A i x -A ⊤ i b i , and ∇ 2 f i (x) = A ⊤ i A i ⪰ O d .Thus, it follows that∇f (x) = 1 n n i=1 A ⊤ i A i x -A ⊤ i b i , and ∇ 2 f (x) = 1 n n i=1 A ⊤ i A i ⪰ O d .The problem is 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 prox γfi , which holds for any x ∈ R d ,EQUATIONObserve that in the linear regression problem, since we know the closed form expression of each f i and f , we know the corresponding smoothness constantL i = λ max A ⊤ i A i . Notice that from Lemma 1, we have M γ fi (x) = f i prox γ (f i ) + 1 2γ x -prox γ (f i ) (x) 2 .Since we know prox γ (f i ) in closed form using (52), we also know each local Moreau envelope in closed form, and thus the same forM γ = 1 n n i=1 M γ fi .As a result, we can deduce L γ for M γ . In our experiments, we pick d ≥ n i=1 n i so that we are in the interpolation regime. Each A i is generated randomly from a uniform distribution between [0, 1), and the corresponding vector b i is also generated from the same uniform distribution. In order to find a minimizer x ⋆ , 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 a 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://anonymous.4open.science/r/FedExProx-F262/In this section we provide the numerical experiments in the large dimension regime, where n i = 20 for each i ∈ [n], n = 30, d = 900.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 theoretical 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, α γ,n is decreasing as γ 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 α γ,τ here increases as τ becomes larger, which coincides with our predictions in Remark 7. In this section, we compare the performance in terms of iteration complexity for FedExProx with different local step sizes. We also include FedExP as a reference. The local step size of FedExP is chosen to be 1 6tLmax , where t is the number of gradient descent iterations performed by each client for local training, L max = max i L i , where L i is the smoothness constant of f i .As one can observe from Figure 4 , for our proposed method FedExProx, the larger γ is, the faster it will converge. However, as γ becomes larger, the improvement in iteration complexity becomes trivial at some point. Note that for different γ, the complexities required to compute the proximity operator locally varies and often larger γ requires more computation than smaller γ. Compared to FedExP with the best local step size 1 6tLmax , FedExProx with a large enough γ is better in terms of iteration complexity. In the case where the computation of proximity operator is efficient, our method has a better computation complexity as well. Notice that small γ 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 a 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 γ, and our method converges for arbitrary local step size γ > 0, 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 γ is observed. A small γ leads to slow convergence of the algorithm, while for large γ, the convergence is improved. However, at some point, the improvement becomes trivial. In this section, we compare FedExProx and its two adaptive variants FedExProx-GraDS and FedExProx-StoPS. We first focus on the full participation case. Note that in this case, the all the algorithms are deterministic. For FedExProx-GraDS, as it is suggested by Theorem 2, the extrapolation parameter is given byα k = α k,G := 1 n n i=1 x k -prox γfi (x k ) 2 1 n n i=1 x k -prox γfi (x k ) 2 .The server can use the local iterates it received from each client to compute α k,G directly. If, in addition, we know L max , we can implement a version that has a better theoretical guarantee,α k,G := 1 + γL max γL max • 1 n n i=1 x k -prox γfi (x k ) 2 1 n n i=1 x k -prox γfi (x k ) 2 .For FedExProx-StoPS, we haveα k = α k,S = 1 n n i=1 M γ fi (x k ) -inf M γ fi γ 1 n n i=1 ∇M γ fi (x k ) 2 .In order to implement α k,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 M γ fi which, according to Lemma 5, is the same as inf f i .From Figure 6 , we can observe that in all cases when γ is sufficiently large, FedExProx-StoPS is the best among the three algorithms considered, and FedExProx-GraDS outperforms FedExProx, this provides numerical evidence for the effectiveness of our proposed algorithms. In the cases when γ is small, the convergence of FedExProx-GraDS seems to be better than the other two algorithms. We also plot the difference of extrapolation parameter used by the algorithms in each iteration. From Figure 7 , observe that when γ is small, α k,G is often much larger than α k,S , resulting in better convergence of FedExProx-GraDS as observed in the first two plots of Figure 6 . When γ becomes larger, α k,G and α k,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 8 that in all cases, the two adaptive variants FedExProx-GraDS-PP and FedExProx-StoPS-PP outperform FedExProx in iteration complexity, and between the two adaptive variants, FedExProx-GraDS is the better one almost all the time. However, FedExProx-GraDS seems to be more stable than FedExProx-StoPS, especially when γ is small. The client minibatch size is chosen from {5, 10, 20}, for each minibatch size, a step size γ ∈ {0.001, 0.005, 0.1, 0.5, 1, 5, 10, 50, 100, 500} is randomly selected.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 are 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 in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in 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. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.This is justified by Lemma 13.Note that L i 1+γL i < 1 γ for any γ > 0.Same as M γ (x), its smoothness constant can be much better. 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: This paper has two main weaknesses. \
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. 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.\
Second, 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.
questions: 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.\
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.
limitations: limitations: The authors have addressed the main limitation of this paper in Section 5.1.
paper_summary: paper_summary: This paper proposes and analyzes several server-side extrapolation strategies to enhance the theoretical and empirical convergence properties of FedProx.
The authors present the convergence properties of the proposed methods for smooth convex and strongly convex problems in the interpolation regime. 
Theoretical results demonstrate that the proposed methods have better dependence on the smoothness coefficient $L$ than FedExP in the general case. 
Specifically, 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.
strength_weakness: strength_weakness: 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.
2. The proposed method is only designed for proximal-based algorithms, which limits its application.
3. The experiments are conducted on a small dataset with a logistic regression problem. I encourage the author to conduct experiments on more complex models.
questions: questions: 1. The main concern is the suitability of Assumption 2. Are there any reasons for this assumption?
2. There are some typo:
- Line 212, Eq.(8), $\frac{1}{n}\sum_{i\in \mathcal{S}_k}$, Should this be $\frac{1}{\vert \mathcal{S}_k\vert}$?
- 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.
- Line 868, the second equation, Should it be $\chi^2$?

The author should review the rest of the paper to address these typo.

I will raise my score if the authors address my concerns.
limitations: limitations: NA.
paper_summary: paper_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.
strength_weakness: strength_weakness: **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. 

**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.

I am willing to raise my score if the authors clarify these points.
questions: questions: -
limitations: limitations: -

Comments: title: Thanks!!
Thank you!
Did our response address your concerns?
comment: Thanks for the information!
There is a version of FedAvg that employs both a local learning rate and a global learning rate. For more details, you can refer to the paper "Achieving Linear Speedup with Partial Worker Participation in Non-IID Federated Learning."
Thank you for your response.

Regarding your concern, FedAvg indeed uses a constant learning rate $\eta_l$ for solving each local optimization problem. However, it does not employ extrapolation ($\eta_g = 1$). In its extrapolated version, FedExP, the constraint $\eta_l \leq \frac{1}{6\tau L}$​ applies to the local stepsize $\eta_l$​. The extrapolation parameter $\eta_g$​, on the other hand, is determined adaptively in each round based on gradient diversity. The original paper does not provide a theory for using a constant $\eta_g$​ in FedExP.

Hope this clarifies your concern. We appreciate your time and efforts on reviewing our paper.
We will respond to this question shortly.

authors
Thank you for addressing my concerns. Regarding the replies about the first weakness, I do not fully agree with your explanation. The authors claim that this paper uses a constant step size compared to existing works such as FedExp. However, to the best of my knowledge, in cases where there is no gradient noise and no data heterogeneity, FedAvg can also use a constant learning rate. Moreover, in FedExp, the global step size \(\eta\) is required to satisfy \(\eta \le \frac{1}{\tau L}\), which still constitutes a constant step size, doesn't it?

Based on the responses provided, I will raise my score to 5.
Thank you!
Dear Reviewer xABN,

Thanks again for your review. Did our rebuttal address your concerns? Please note that the other two reviewers increased their scores from 5 to 6 after rebuttal. Our average score is 5.33 now, which may still be considered borderline (the official email from NeurIPS said that scores between 4.5-5.5 are borderline). Your score may therefore have a large influence on the fate of our work. 

Please let us know whether the concerns were addressed. If that is the case, we would of course be happy if this could be reflected in the score. If not, please let us know what remains to be explained -- today is the last day we can do so.

Thanks again!

Authors
Thank you for addressing the majority of my concerns. Given that Assumption 2 is relevant in deep learning scenarios, I recommend that the authors validate the proposed method within neural network settings in the next version of the paper. 

Based on the responses provided, I am inclined to raise my score to 6.
Thank you for your timely response. We appreciate your effort in reviewing our work.

- Weakness 1: We agree with the reviewer and will include a discussion on our focus on communication complexity in the paper.

- Weakness 2: We now better understand the reviewer's concern. Indeed, we did not include such explanations in the paper. In the next version, we will provide a discussion on solving the local problems and the role of the two parameters in this case.

We will add the following discussion to the paper in its next version.
> Each local proximity operator can be solved using different oracles. In practice, clients may use gradient descent or stochastic gradient descent to solve the local problem to a certain accuracy. The complexity of this subroutine depends on the local stepsize $\gamma$. If $\gamma$ is large, the local problem becomes harder to solve because we aim to minimize the local objective itself. Conversely, if $\gamma$ is small, the problem is easier since we do not stray far 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 of the two algorithms, assuming efficient local computations are carried out by the clients.

Please let us know if anything is unclear.
> Weakness 1

I don't understand why it is neither feasible nor meaningful to compare the complexity of two algorithms that are designed for precisely the same setting, i.e. smooth, convex function in the interpolation regime. I think it is fair that you want to focus on the communication complexity rather than the local iteration complexity, but in this case this should be clearly discussed in the paper.

> Weakness 2

What I meant with my comment was precisely that the fact that the differentiation between $\gamma$ and $\alpha_k$ only becomes relevant when taking into account the interaction between the communication complexity and the amount of local steps required. So one the one hand you simply assume that you can solve these local problem or at least you do not discuss how to solve them, but then you say that having the ability to set $\gamma$ and $\alpha_k$ separately is useful to adapt to the local problems.

I've raised my score to 6.