FedExP: Speeding Up Federated Averaging via Extrapolation
Authors: Divyansh Jhunjhunwala
Shiqiang Wang
Gauri Joshi
Venue: ICLR 2023 notable top 25% (2023)
Keywords: Federated Learning
Optimization
Step Size

Abstract: Federated Averaging (FedAvg) remains the most popular algorithm for Federated Learning (FL) optimization due to its simple implementation, stateless nature, and privacy guarantees combined with secure aggregation. Recent work has sought to generalize the vanilla averaging in FedAvg to a generalized gradient descent step by treating client updates as pseudo-gradients and using a server step size. While the use of a server step size has been shown to provide performance improvement theoretically, the practical benefit of the server step size has not been seen in most existing works. In this work, we present FedExP, a method to adaptively determine the server step size in FL based on dynamically varying pseudo-gradients throughout the FL process. We begin by considering the overparameterized convex regime, where we reveal an interesting similarity between FedAvg and the Projection Onto Convex Sets (POCS) algorithm. We then show how FedExP can be motivated as a novel extension to the extrapolation mechanism that is used to speed up POCS. Our theoretical analysis later also discusses the implications of FedExP in underparameterized and non-convex settings. Experimental results show that FedExP consistently converges faster than FedAvg and competing baselines on a range of realistic FL datasets. 

Published as a conference paper at ICLR 2023
FEDEXP: SPEEDING UP FEDERATED AVERAGING VIA EXTRAPOLATION
Divyansh Jhunjhunwala' Shiqiang Wang', Gauri Joshi' CCamegie Mellon University, 'IBM Research (djhunjhu, g Bandrew. emu, edu, wangshiq@us. ibm com
ABSTRACT
Federated Averaging (FedAvg) remains the most popular algorithm for Federated Learning (FI.) optimization due to its simple implementation, stateless nature, and privacy guarantees combined with secure aggregation. Recent work has sought generalize the vanilla averaging in FedAvg to generalized gradient descent step by treating client updates as gauuudogrrddents and using a server step size. While the use of a server step size has been shown to provide performance improvement theoretically, the practical benefit of the server step size has not been seen most existing works. In this work, present FedExP, method to adaptively deter- mine the server step size in FL. based on dynamically varying pseudo-grat -graddeess throughout the FL process. We begin by considering the overparameterized convex regime, where we reveal an interesting similarity between F FedAvg and the Pro- jection Onto Convex Sets (POCS) algorithm. We then show how FedExP can be motivated novel extension to the extrapolation mechanism that is used to speed up POCS. Our theoretical analysis later also discusses the implications of FedExp in underparameterizeed and non-convex settings. Experimental results show that FedExP consistently converges faster than FedAvg and competing baselines on range of realistic FL datasets.
INTRODUCTION
Federated Learning (FL) has emerged as a key distributed learning paradigm in which a central server orchestrates the training of a machine learning model across a network of devices. FL is based on the fundamental premise that data never leaves clients device, clients only communicate model updates with the server. Federated Averaging or FedAvg, first introduced by McMMahn et al. (2017). remains the most popular algorithm in this setting due to the simplicity of its implementation, stateless nature (i.e., clients do not maintain local parameters during training) and the ability to incorporate privacy-preserving protocols such as secure aggregation (Bonawitz al. 2016 Kadhe et al. 2020).
Slowdown Due to Heterogeneity. One the most persistent problems in FedAvg is the slowdown in model convergence due to data heterogeneity across clients. Clients usually perform multiple steps of gradient descent on their heterogeneous objectives before communicating with the server in FedAvg, which leads to what is colloquially known as client drift error (Karimireddy 2019). The effect of heterogeneity is further exacerbated by the constraint that only a fraction of the total number of clients may be available for training every round (Kairouz al. 2021). Various techniques have been proposed to combat this slowdown, among the most popular being variance reduction techniques such Karimireddy al. (2019), Mishchenko al. (2022); Mitra et al.)
al..
(2021), but they either lead to clients becoming stateful, add extra computation or communication requirements have privacy limitations.
Server Step Size. Recent work has sought to deal with this slowdown by using two separate step sizes in FedAvg client step size used by the clients minimize their local objectives and a server step size used by the server to update the global model by treating client updates as pseudo-gradients (Karimireddy et al.) 2019 Reddi al. 202T). To achieve the fastest convergence rate, these works propose keeping the client step size (1/rVTT) and the server step size as ((VM)). where T is the number of communication rounds, T is the number of local steps and M' is the number of clients. Using small client step size mitigates client drift, and large server

---

Published as a conference paper at ICLR 2023
step size prevents global slowdown. While this idea may be asymptotically optimal. it is not always effective in practical non-asymptotic and communication-ling limited settings (Charles & Konečný 2020). In practice, a small client step size severely slows down convergence in the initial rounds and cannot be fully compensated for by a large server step size (see Figure T). Also, if local objectives differ signif- icantly. then it may be beneficial to use smaller values of the server step size (Malinovsky al. 2022).
6.8
8.2
42
55
9:1
Server Step Size (ng)
9.3
16
46
58
71
56
46
60
73
77
8.2
Therefore, we seek to answer the following question: For moderate client step size, can we adapt the server step size according to the local progress made by the clients and the heterogeneity of their ob- jectives? In general, is challenging to answer this question because it is difficult to obtain knowledge of the heterogeneity between the local objectives and appropriately use it to adapt the server step size.
60
50
-0
6.4
6.1
0.01
0.03
0.1 0.3
Clent Step Size (nt)
Figure 1: Test accuracy (%) achieved by different server and client step sizes on EMNIST dataset (Cohen et al., 2017) after 50 rounds (details of experimental setup are in Section 6 and Appendix D).
Our Contributions. In this paper, we take a novel approach to address the question posed above. We begin by considering the case where the models are overparameterized, i.c., the number of model parameters larger than the total number of data points across all clients. This is often true for modern deep neural network models (Zhang et al., 2017) Jacot al.) 2018) and the small datasets collected by edge clients in the FL setting. In this overparameterized regime, the global minimizer becomes common minimizer for all local objectives, even though they may be arbitrarily heterogeneous. Using this fact, we obtain a novel connection between FedAvg and the Projection Onto Convex Sets (POCS) algorithm, which is used to find a point in the intersection of some convex sets.
Based on this connection, we find an interesting analogy between the server step size and the extrapolation parameter that is used to speed up POCS (Pierra 1984). We propose new extensions to the extrapolated POCS algorithm to support inexact and noisy projections as F FedAv. In particular, we derive time-varying bound on the progress made by clients towards the global minimum and show how this bound can be used to adaptively estimate a good server step size at each round. The result is proposed algorithm FedExP, which is method adaptively determine the server step in that round.
size in each round of FL based on the pseudo-grading
Although motivated by the overparameterized regime, our proposed FedExP algorithm performs well (both theoretically and empirically) in the general case, where the model can be either overpa- rameterized or underparameterized. For this general case, we derive the convergence upper bounds for both convex and non-convex objectives. Some highlights of our work are as follows.
We reveal a novel connection between FedAvg and the POCS algorithm for finding point in the intersection of convex sets.
The proposed F FedExP algorithm is simple to implement with virtually no additional communica- tion, computation, storage required at clients the server. It is well suited for both cross-device and cross-silo FL., and is compatible with partial client participation.
Experimental results show that FedExP converges 1.4-2x x faster than FedAvg and most compet- ing baselines on standard FL tasks.
Related Work. Popular algorithms for adaptively tuning the step size when training neural
networks include Adagrad (Duchi et al.] (01T)) and its variants RMSProp (Tieleman et al. 2012) and Adadelta (Zeieer (2012). These algorithms consider the notion of coordinate-wis adaptivity and adapt the step size separately for each dimension of the parameter vector based on the magnitude of the accumulated gradients. While these algorithms can be extended to the federated setting using the concept of meeudo-gradient as done by RRedii et al. (2021) these extensions are agnostic to inherent data heterogeneity across clients, which is central to FL. On the contrary, FedExP is explicitly designed for FL settings and uses client-centric notion of adaptivity that utilizes the heterogeneity of client updates in each round. The work closest to us is Johnson et al. (020)) which proposes method to adapt the step size for large-batch training by estimating the gradient diversity YYnn et al. 2018) of a minibatch. This result has been improved in a recent work by Horváth et al.] (2022) However, both Johnson et al. (000)) Horváth et al. (0022) focus on the centralized setting. In
2

---

Published as a conference paper at ICLR 2023
FedExP, similar concept, but within a federated environment which comes with stronger theoretical motivation, since client data are inherently diverse in this case. We defer a more detailed discussion of other adaptive step size methods and related work to Appendix A
2 PROBLEM FORMULATION AND PRELIMINARIES
As in most standard federated learning frameworks, we consider the problem of optimizing the model parameters w E IRd to minimize the global objective function F(w) defined as follows:
M
min F(w): WERd
= M
Σ F((w),
where Fi(w) := [DI] Es.D,, ((w, 6.) is the empirical risk objective computed on the local data D; the the i-th client. Here, E(:.-) is loss function and 6; represents data sample from the empirical local data distribution D, The total number of clients in the FL system is denoted by M. Without loss of generality, we assume that all the M client objectives are given equal weight in the global objective function defined in (I). Our algorithm and analysis can be directly extended to the case where client objectives are unequally weighted, e.g.. proportional to local dataset sizes [Da].
FedAvg. We focus on solving T using FedAvg (McMahan et al. 2017 Kairouz al. 2021]. At round of FedAvg, the server sends the current global model w()) to all clients. Upon receiving the global model, clients perform r steps of local stochastic gradient descent (SGD) to compute their updates {A (4) Ji=1 for round follows.
Perform Local SGD:
(t,k+1)
wi (t.N)
V.( (f.R) ((tk)) Ah: E (0, ................................ (2)
Compute Local Difference:
(()) w (1,7)
where w (((10) (( for all i [M], 7 is the client step size and (w ,
(f.kk) ((tik) represents
stochastic gradient computed on the minibatch e(t.k) sampled randomly from D,
Server Optimization in FedAvg. In vanilla FedAvg (McMahan et al. 2017 the global model would simply be updated as the average of the client local models, that is, w (++1) M >M w (t,7)
To improve over this, recent work (Reddi et al. 2021 Hsu et al. 2019) has focused optimizing the and multiplying
server aggregation process by treating the client updates (!) as
by server step size when aggregating them follows.
Generalized FedAvg Global Update:
w ((+1)) (t) 17g
(t)
(4)
where A = M Σ Note that setting fl recovers the vanilla FedAvg update.
1=1 A (+) is the aggregated client update in round t and fy acts as server step size.
While the importance of the server step size has been theoretically well established in these works, we find that its practical relevance has not been explored. In this work, we take a step towards bridging this gap between theory and practice by adaptively tuning the value of My that we use in every round.
3 PROPOSED ALGORITHM: FEDEXP
Before discussing our proposed algorithm, we first highlight useful and novel connection between FedAvg and the POCS algorithm used to find point the intersection of some convex sets.
3.1 MOTIVATION FOR EXTRAPOLATION
Connection Between FedAvg and POCS in the Overparameterized Convex Regime. Consider the case where the local objectives of the clients (F,(W)) are convex. In this case, we know that the set of minimizers of F, (w) given by S [w E arg min (.(w) is also convex set for all i E [M]. Now let us assume that we are in the overparameterized regime where is sufficiently larger than the total number of data points across clients. In this regime, the model can
3

---

Published as a conference paper at ICLR 2023
fit all the training data at clients simultaneously and hence be minimizer for all local objectives. Thus we assume that the global minimum satisfies w E S; Vi E [M]. Our original problem in (I) can then be reformulated as trying to find point in the intersection of convex sets {S;]at since w St Vi E [M]] One of the most popular algorithms to do so is the Projection Onto Convex Sets (POCS) algorithm (Gurin et al. 1967). In POCS, at every iteration the current model is updated as follows
Generalized POCS update:
POCS
((++)
W
POCS
(2)
A ( ΣM P. P:
Woocs
(1)
(5)
where (w POCS is a projection of Weucs on the set S and À is known as the relaxation coefficient (Combettes 1997).
(t)
Extrapolation in POCS. Combettes (197) notes that POCS has primarily been used with À 1,
with studies failing to demonstrate a systematic benefit of X < 1 or À > 1 (Mandell 1984). This prompts Combettes (197)) to study adaptive method of setting X, first introduced by Pierra (9844)
as follows:
E Pi(w((() -w(]] M EA Pi(w(()- w(ol
1
Pierra 1984) refer to the POCS algorithm with this adaptive ( as Extrapolated Parallel Projection Method (EPPM). This is referred to as extrapolation since always have X 2 by Jensen's inequality. The intuition behind EPPM lies in showing that the update with the proposed ((()
always satisfies ||weoos ((+1) -w"l"
Weecs - w=12 thereby achieving asymptotic convergence. Experimental results in Pierra (1984) and Combettes (1997) show that EPPM can give order-wiss wise
speedup over POCS, motivating us to study this algorithm in the FL context.
3.2 INCORPORATING EXTRAPOLATION IN FI.
Note that to implement POCS do not need to explicitly know the sets {S only need to know how to compute a projection on these sets. From this point of view, we see that FedAvg proceeds similarly to POCS. In each round, clients receive w(() from the server and run multiple SGD steps to compute an "approximate projection" w (f,T) of (() on their solution sets S, These approximate projections are then aggregated at the server to update the global model. In this case, the
relaxation coefficient À plays exactly the same role as the server step size 1]g in FedAvg.
Inspired by this observation and the idea of extrapolation in POCS, we seek to understand if a similar idea can be applied to tune the server step size fl in FedAvg. Note that the EPPM algorithm makes use of exact projections to prove convergence which is not available to us in FL settings. This is further complicated by the fact that the client updates are noisy due to the stochasticity in sampling minibatches. We find that in order to use an EPPM-like step size the use of exact projections can be relaxed to the following condition, which bounds the distance of the local models from the global minimum as follows.
Approximate projection condition in FL: / D w (3.7)
w-11
< [|w(1) w111
(6)
where (())) and {w (M Timi are the global and local client models, respectively, at round and w" is a global minimum. Intuitively, this condition suggests that after the local updates, the local models are closer to the optimum w' on average as compared to model w(() at the beginning of that round. We first show that this condition (6) holds in the overparameterized convex regime under some conditions.
The full proofs for lemmas and theorems in this paper are included in Appendix C
Lemma 1. Let F((w) be convex and L-smooth for all i E [M] and let w* be a common minimizer of all F. (w). Assuming clients run full- gradient descent minimize their local objectives with n S I/L, then (6) holds for all and T 2 1.
In the case with stochastic gradient noise or when the model is underparameterized, although (6) may not hold in general, we expect it to satisfied at least during the initial phase of training when (t( w=1 is large and clients make common progress towards minimum.
We refer here to parallel implementation of POCS. This is also known Parallel Projection Method (PPM) and Simultaneous Iterative Reconstruction Technique RT) in some literature (Combettes) 1997).
4

---

Published as a conference paper at ICLR 2023
Algorithm Proposed Algorithm: FedExP
1: Input: w(()) number rounds T, local iteration steps т. parameters 7)t
2: For t = 0, ..... T communication rounds do:
3: Global server does:
4: Send w(t) to all clients
5: Clients i E [M]] in parallel do:
6: Set (((1.0) w(t,0)
7: For k =
local iterations do:
Update w; (1,6+1) f (t.k) Nn (w, (t.b) (t.k)
8:
9: Send A < w(!)
to the server
Global server does:
(2) Compute Á (4) 3/4 EE A and Tg max 1, E M5
|| 1//M (1300]2+ +
12: Update global model with w((+)) + (()) 1/9 A
(4)
Given that (6) holds, we now consider the generalized FedAvg update with server step size 170
(t)
in
round Our goal is to find the value of Tlg that minimizes the distance of w(++1) w
(t)
[w((++)) -w w"| |[w(()
w + (ma ) | 2nh (1)
(w(() w (01)
(7)
Setting the derivative of the RHS of to zero have,
(ng)) Jopt (w(t)) w A (
ΣN ((1)) w" 1
JM i==1
(4)
(8)
|
M A 10
2M
A
where the last inequality follows from (a. b) = |[l|all + ||b||
||a - bll definition of A
(9) in
and 66. Note that depending on the values of (4, y we may have (ng (2) Jopt Thus, we see
that 66 acts as suitable replacement for projection to justify the use of extrapolation in FL settings.
3.3 PROPOSED ALGORITIM
Motivated by our findings above, we propose the following server step size for the generalized FedAvg update at each round:
M 1
(ng (t) )eddDE = max
(9)
2M 10 +e)
We term our algorithm Federated Extrapolated Averaging FedExP, in reference to the original EPPM algorithm which inspired this work. Note that our proposed step size satisfies the property
that
(m
upt (nl
(e) Dredezel I (nn >apt 1| when ( holds, which can be seen by comparing
8 and Since ( depends quadratically on 7y we can show that in this case
|[w((+1))
(1) w1 < [(+ w'll implying we are at least as close to the optimum as the FedAvg update. In the rest the paper, we denote (ng²)
(1g FedExP
(+) )eocre as flg (t) when the context clear.
Importance Adding Small Constant to Denominator. In the case where (6) does not hold, using the lower bound established in (8) can cause the proposed step size to blow up. This is especially true
towards the end of training where we can have |
0 but A
* Thus we propose
add small positive constant to the denominator in 9 to prevent this blow-up. For large enough e our algorithm reduces to FedAvg and therefore tuning can be a useful tool to interpolate between vanilla averaging and extrapolation. Similar techniques exist in adaptive algorithms such Adam (Kingma & Ba 2015) and Adagrad (Duchi et 2011) to improve stability.
Compatibility with Partial Client Participation and Secure Aggregation. Note that FedExP
can be easily extended to support partial participation of clients by calculating (1) using only the updates of participating clients, i.c., the averaging and division in ( will be only over the clients that participate in the round. Furthermore, since the server only needs to estimate the average of
(0) pseudo-gradient norms, 1/8 be computed with secure aggregation, similar to computing (1)
5

---

Published as a conference paper at ICLR 2023
Connection with Gradient Diversity. We that our lower bound (1g (9)
Jopt naturally depends
on the similarity of the client updates with each other. In the case where T and clients run [[
full-batch gradient descent,
lower bound (8) reduces >i=1
/2M (
which is used as a measure of data-heterogeneng in many FL works WWang et al. 2020 Haddadpour & Mahdavi 2019). Our lower bound suggests using larger step-sizes as this gradient diversity increases. which can be useful tool to speed up training in heterogeneous settings. This is an orthogonal approach to existing optimization methods to tackle heterogeneity such as Karimireddy et al. (2020b); Li et al. (2020); Acarr et al. (202T), which propose additional regularization terms adding control variates to the local client objectives to limit the impact of heterogeneity.
4 CONVERGENCE ANALYSIS
Our analysis so far has focused on the overparameterized convex regime to motivate our algorithm. In this section we discuss the convergence of our algorithm in the presence of underparameterizzaton and non-convexity. We would like to emphasize that (6) not needed to show convergence of FedExP: it is only needed to motivate why FedExP might be beneficial. To show general convergence, we only require that t7 be sufficiently small and the standard assumptions stated below.
Challenge in incorporating stochastic noise and partial participation. Our current analysis fo- cuses the case where clients are computing full-batch gradients in every step with full participation. This is primarily due to the difficulty in decoupling the effect of stochastic and sampling noise on % and the at round t, then Eges ( (4)
{A M
{
==1 To be more specific, if use g(t) to denote the randomness
Egce
(mg Eg(s)
[+]] which significantly complicates the
proof. This purely a theoretical limitation. Empirically, our results in Section 6 show that FedExP
performs well with both SGD and partial client participation.
Assumption 1. (L-smoothness) Local objective (((w) is differentiable and L-smooth for all i E [M], ie., [|VF; w) VF:(w')| S L|[w - w'll Vw, w' E Rd.
Assumption 2. (Bounded data heterogenenity optimum) The norm of the client gradients the global optima w is bounded as follows: 1 Ei (F((w*) S <o?
Theorem 1. (F; convex) Under Assumptions
22 and assuming clients compute full-batch gradients with full participation and 17 S OFT. the iterates (w(t)) generated by FedExP satisfy.
(0) - ww'l
)
((w((1)) - F* 00
+0 10 (vē-(r- 1)LO") +
-
(nITO.) Ty:=noise optimum
(10)
ThT Ets Th
Ty:=elient drift everr
Ty=:inttllizaionn evror
where 1l (1) is the FedExP server step size round and (T) Ei 1/ (!))
LiED No
For the non-convex case, we need the data heterogeneity be bounded everywhere as follows. Assumption 3. (Bounded global gradient variance) There exists a constant o² > O such that the global gradient variance is bounded as follows. M ΣM [VF(() (F)))) < o
Vw E Rd 3 and assuming clients compute full-batch Σ 1 the iterates (w(()) generated by FedExP satisfy, brL:
Theorem 2. (F; are non-convex) Under Assumptions)
gradients with full participation and Tp
F F(w(()) Fo
min V F( 5o
+ (n + O (mLLoo)
(1I)
THTT ETD 1l
To: drift evror
Ta:= global variance
Ty:minitaaiiationn ervor
(0 is the FedExP server step size at round t. where Bi
Discussion. In the convex case, the O (l) wI -w*ll//mrTt /nTT) + +0 T I)Lo (Khaled et al. (020) and in the non-convex case by O ((F(w9) F*)/nTT) + O (n 1.2 1)0z) (Wang al. 2020). A careful inspection
of FedAvg can be bounded by

---

Published as a conference paper at ICLR 2023
Trajectory Iterates
Global Objective
to Optimum
Festeg
FesEuP
244
int losdon
Clobul Meammm
10"
(we W 1P
2.3-
Fiwe
2.2
20
10"
Tacliog Jast bantts)
Tecifig 3.art lestl)
FecE.P (Lst boeks)
TecEuP (Leet tamm))
20
FecExP (log ledd too laotess
FecEUP (Wa las two testtos)
Training rounds
Training rounds
Figure 2: Training characteristics of FedAvg and FedExP for the 2-D toy problem Section
The last iterate of FedExF has an oscillating behavior in F(w) but monotonically decreases ww(t) w*|;; the average of the last two iterates lies in lower loss region than the last iterate.
reveals that the impact of Ti on convergence of FedExP is different from FedAvg (effect of Ta is the same). We see that since Σ "i (4) 2 T. FedExP reduces T₁ faster than F FedAvg. However this comes at the price of an increased error floor due to Tg. Thus, the larger step-sizes in FedExP help us reach the vicinity of optimum faster, but ultimately end up saturating at a higher error floor due to noise around the optimum. Note that the impact of the error floor can be controlled by setting the client step size N7 appropriately. Moreover, the overparameterized convex regime where o² O, the effect of T2 and T3 vanishes and thus FedExP clearly outperforms FedAvg. This aligns well with our initial motivation of using extrapolation in the overparameterized regime.
5 FURTHER INSIGHTS INTO FEDEXP
In this section, we discuss some further insights into the training of FedExP and how we leverage these insights to improve the performance of FedExP.
FedExP monotonically decreases |w(t) w*| but not necessarily F(w(O) - F(w"). Recall that our original motivation for the F FedExP step size was aimed at trying to minimize the distance to the optimum give by [ w ||2 when (6) holds. Doing so satisfies [| w+|
[w(())
w* but does not necessarily satisfy ((w((+1)) S F(w(").
To better illustrate this phenomenon, we consider the following toy example in R2. We consider a setup with two clients, where the objective at each client is given as follows:
Fi (w) (3w1 + w2 - -3)3; F2(w) = (wi + W2 - -3)2
(12)
We denote the set of minimizers of Fi (w) and (2(w) by S {w 3w + w2 = 3) and
Sj = [w Wi + w2 = 3) respectively. Note that Sj and S intersect at the point w- =000]) making it global minimum. To minimize their local objectives, assume clients descent with r + oo in every round" Figure 2 shows the trajectory of the iterates generated by FedExP and FedAvg. We see that while |/ - w| decreases monotonically for FedExP, F(w(()) does not do so and in fact has an oscillating nature as we discuss below.
gradient
Understanding oscillations in F(w(t). We see that the oscillations in F(w(+) are caused by FedExP iterates trying to minimize their distance from the solution sets Sj and Sž simultaneously. The initialization point (()) is closer to St than S2, which causes the FedExP iterate at round to move towards Sg. then back towards Si and so on. To understand why this happens, consider
(+) the case where (1) 0, A * 0. In this case, we have nh = 2 and therefore (t+1))
(()) 2 w (t.r) which indicates that FedExP is now trying to minimize | This
gives us the intuition that the FedExP update in round / is trying to minimize the objectives of the clients that have
|
0. While this leads to temporary increase in global loss F(w(t)) in
TTh local models will be an exact projection of the global model on the solution sets {S; Him. In this case, the lower bound in (8) can be improved by factor of 2 and therefore use 1// for this experiment (see Appendix C.4 aan Appendix C.4. I for proof).
(21)) 0012
7

---

Published as a conference paper at ICLR 2023
FedExP (ours)
FedAvg
SCAFFOLD
FedAdagrad
Synthetic
EMNIST
CIFAR-10
CIFAR-100
CINIC-10
65
10
Mesn Stuared Entr
Accusnyy (h1)
(%)
Tort Accurny Phi
10"
60
40-
8 15
10"
18
10'19
50
-0"
40-
100
250 Taskingg rounds
250 Trsining munds
500
260
500
260
500
Tishiing rounds
Tranning rounds
Training rounds
Figure 3: Experimental results on synthetic linear regression experiments and range of realistic FL tasks. FedExP consistently gives faster convergence compared to baselines while adding extra computation, communication or storage at clients or server.
some rounds as shown in Figure distance to the global optimum w
2 it is beneficial in the long run as it leads to a faster decrease in
Averaging last two iterates in FedExP. Given the oscillating behavior of the iterates of FedExP, we find that measuring progress on F(w) using the last iterate can be misleading. Motivated by this finding, we propose to set the final model as the average of the last two iterates of FedExP. While the last iterate oscillates between regions that minimize the losses Pi (w) and F22 (w) respectively, the behavior of the average of the last two iterates is more stable and proceeds along a globally low loss region. Interestingly, we find that the benefits of averaging the iterates of F FedExP also extend to training neural networks with multiple clients in practical FL scenarios (see Appendix [DII. In practice, the number of iterates average over could also be hyperparameter for FedExP, but find that averaging the last two iterates works well, and we use this for our other experiments.
EXPERIMENTS
We evaluate the performance of FedExP synthetic and real FL tasks. For our synthetic experiment, we consider a distributed overparameterized linear regression problem. This experiment aligns most closely with our theory and allows us to carefully examine the performance of FedEx? when (6)
holds. For realistic FL tasks, we consider image classification on the following datasets i) EMNIST
(Cohen al. 2017), ii) CIFAR-10 (Krizhevsky et al.) 2009), iii) CIFAR- 100 (Krizhevsky et al. 2009).
iv) CINIC-10 (Darlow et al.) 2018). In all experiments, we compare against the following baselines
i) FedAvg, ii) SCAFFOLD Karimireddy et al. 2020b), and iii) FedAdagrad (Reddi et al. 2021) which is a federated version of the popular Adagrad algorithm. To the best of our knowledge, we are not aware of any other baselines that adaptively tune the server step size in FL.
Experimental Setup. For the synthetic experiment, we consider a setup with 20 clients, 30 samples at each client, and model size to be 1000, making this overparameterrize problem. The data at each client is generated following a similar procedure as the synthetic dataset in Li et al. (2020). We use the federated version of EMNIST available at Caldas al. (2019), which is naturally partitioned into 3400 clients. For 1TFAR-10/100 we artifically partition the data into 100 clients, and for CINIC-10 partition the data into 200 clients. In both cases, follow a Dirichlet distribution with a 0.3 for the partitioning to model heterogeneity among client data (Hsu et al. 2019). For EMNIST we use the same CNN architecture used in Reddi al. (202I). For CIFARIO. CIFARIO0 and CINIC-10 we use a ResNet-18 18 model (He et al.) 2016). For our baselines, we find the best performing Ny and n by grid-searce tuning. For FecExP we optimize for c and ny by grid search. We fix the number of participating clients to 20, minibatch size to 50 and number of local updates to 20 for all experiments. In Appendix D we provide additional details and results, including the best performing hyperparameters, comparison with FedProx (Li et al. 2020), and results for more rounds.
FedExP comprehensively outperforms FedAvg and baselines. Our experimental results in Figure demonstrate that FedExP clearly outperforms FedAvg and competing baselines that use the best performing 74 and np found by grid search. Moreover, FedExP does not require additional communication or storage at clients or server unlike SCAFFOLD and FedAdagrad. The order- wise improvement in the case of the convex linear regression experiment confirms
theoretical motivation for FedExP outlined in Section In this case, since (6) is satisfied, we know that the FedExP iterates are always moving towards the optimum. For realistic FL tasks, we see consistent
8

---

Published as a conference paper at ICLR 2023
Table 1: Table showing the average number of rounds reach desired accuracy for FedExP and baselines. FedExP provides a consistent speedup over all baselines.
Dataset
Target Acc.
FedExF
FedAvg
SCAFFOLD
FedAdagrad
EMNIST
84%
186
328 (1.76x)
232 (1.24x)
277 (1.48x)
CIFAR-I0
72%
267
434 (1.62x)
429 (1.61x)
419 (1.56x
CIFAR-100
40%
242
500 (2.06x)
> 500 5000(22066x
494 (2.04x)
CINIC-10 58%
318
450 (1.42x)
470
444 (1.00)
speedup of over 1.4 - 2x over FedAvg. This verifies that FedExP also provides performance
improvement more general settings with realistic datasets and models. Plots showing Tig can be found in Appendix D.5 The key takeaway from our experiments is that adapting the server step size allows FedExp to take much larger steps some (but not all) rounds compared to the constant optimum step size taken by our baselines, leading to large speedup.
Comparison with FedAdagrad. As discussed in Section I FedAdagzad and FedExP use different notions of adaptivity; FedAdagrad uses coordinate-win adaptivity, while FedExP uses client-based adaptivity. We believe that the latter is more meaningful for FI. settings as seen in our experiments. In many experiments, especially image classification tasks like CIFAR, the gradients produced are dense with relatively little variance in coordinate-win gradient magnitudes (Reddi et al., 2021 Zhang et al. 2020). In such cases, FedAdagrad is unable to leverage any coordinate-level information and gives almost the same performance as FedAvg.
Comparison with SCAFFOLD. We that FedExP outperforms SCAFFOLD in all experiments, showing that adaptively tuning the server step size is sufficient to achieve speedup in FL settings. Furthermore, SCAFFOLD even fails to outperform FedAvg for the more difficult CIFAR and CINIC datasets. Several other papers have reported similar findings, including Reddi et al. (202)): Karimireddy et al. (2020a); Yu et al. (2022). Several reasons have been postulated for this behavior, including the staleness of control variates (Reddi et al., 2021) and the difficulty in characterizing client drift in non-convex scenarios (Yu et al. 2022). Thus, while theoretically attractive, simply using variance reduction techniques such as SCAFFOLE may not provide any speedup in practice.
CIFAR-10
Adding extrapolation to SCAFFOLD. We note that SCAFFOLD only modifies the Local SGD procedure at clients and keeps the global aggregation at the server unchanged. Therefore, it is easy to modify the SCAFFOLD algorithm to use extrapolation when updating the global model at the server (algorithm details in Appendix Db. Figure 4 shows the result our proposed extrapolated SCAF FOLD on the CIFAR-10 dataset. Interestingly, we observe that while SCAFFOLD alone fails to outperform F FedAvg, the extrapolated version SCAFFOLD achieves the best performance among all algorithms. This result highlights the importance of carefully tuning the server step size to achieve the best performance for variance-reduction algorithms. is also possible to add extrapolation to algorithms with server momentum (Appendix F
Aocu (4]
FedEUP SO
Fediig
SCAFFOLD
100
200
30C
400
200
Training rounds
Figure 4: Adding extrap- olation to SCAFFOLD for greater speedup.
7 CONCLUSION
In this paper, we have proposed FedExP, a novel extension of FedAvg that adaptively determines the server step size used in every round of global aggregation in FL. Our algorithm is based on the key observation that FedAvg can be seen as approximate variant of the POCS algorithm, especially for overparameterized convex objectives. This has inspired us to leverage the idea of extrapolation that is used to speed up POCS in a federated setting, resulting in FedExP. We have also discussed several thcoretical and empirical perspectives of F FecExP. In particular, we have explained some design choices in FedExP and how it can be used in practical scenarios with partial client participation and secure aggregation. We have also shown the convergence of FedExP for possibly underparameterized models and non-convex objectives. Our experimental results have shown that FedExP consistently outperforms baseline algorithms with virtually no additional computation or communication at clients or server. We have also shown that the idea of extrapolation can be combined with other techniques, such as the rrrducio-redction method in SCAFPOLD, for greater speedup. Future work will study the convergence analysis of F FedExP with stochastic gradient noise and the incorporation of extrapolation into wider range of algorithms used in FL..

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
This work was supported in part by NSF grants CCF 2045694, CNS-2112471, ONR N00014-23--1 2149, and the CMU David Barakat and Vernn Owen- wren-Barakat Fellowship.
REFERENCES
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. In International Conference on Learning Representations, 2021.
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances in Neural Information Processing Systems, 32, 2019.
Larry Armijo. Minimization of functions having Lipschitz continuous first partial derivatives. Pacific Journal of mathematics, (6(1):1-3, 1966.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. Advances in Neural Information Processing Systems, 32, 2019.
Jonathan Barzilai and Jonathan M Borwein. Two-point step size gradient methods. IMA journal of numerical analysis, 8(1):141-14)) 1988.
K. A. Bonawitz, Vladimir Ivanov, Ben Kreuter, Antonio Marcedone, H. Brendan McMahan, Sarvar Patel, Daniel Ramage, Aaron Segal, and Karn Seth. Practical secure aggregation for federated learning user-held data. In NeurIPS Workshop on Private Multi-Party Machine Learning, 2016.
Stephen Boyd and Jon Dattarro. Alternating projections, 2003. httts:://vet stanford.eduu
class/ ee 3 /920/
pdf
Oleg Burdakov, Yu-Hong Dai, and Na Huang. Stabilized Barzilai-Boree Borwein method. Journal of Computational Mathematics, 37(6):916-936, 2019.
Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMa- han, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. In Workshop on Federated Learning for Data Privacy and Confidentiality, 2019.
Zachary Charles and Jakub Konečný. On the outsized importance of learning rates in local update methods. arXiv preprint arXiv:2007.00878. 2020.
Gregory Cohen, Saced Afshar, Jonathan Tapson, and Andre Van Schaik. Emnist: Extending mnist to handwritten letters. In 20/7 International Joint Conference on Neural Networks (IJCNN), pp. 2921-2926. IEEE, 2017.
Patrick L Combettes. Convex set theoretic image recovery by extrapolated iterations of parallel subgradient projections. IEEE Transactions on Image Processing, 6(4):493-506. 1997.
Luke N Darlow, Elliot J Crowley, Antreas Antoniou, and Amos J Storkey. CINIC-10 is not Imagenet or CIFAR-10. arXiv preprint arXiv: 100.00v000000000 2018
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Local SGD optimizes overparam- eterized neural networks polynomial time. In International Conference on Artificial Intelligence and Statistics, pp. 6840-6861 PMLR. 2022.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(7), 2011.
AA Goldstein. Optimization Lipschitz continuous functions. Mathematical Programming, 13(1): 14-22, 1977.
10

---

Published as a conference paper at ICLR 2023
Leonid Georgievich Gurin, Boris Teodorovich Polyak, and È V Raik. The method of projections for finding the common point of convex sets. Zhurnal Vychislitel noi Matematiki Matematicheskoi Fiziki, 7(6):1211-1228, 1967.
Farzin Haddadpour and Mehrdad Mahdavi. On the convergence of local descent methods in federated learning. arXiv preprint arXiv:1910. 144.2 2019.
Elad Hazan and Sham Kakade. Revisiting the Polyak step size. arXiv preprint arXiv: 1905.00313, 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp. 770-778, 2016.
Samuel Horváth, Konstantin Mishchenko, and Peter Richtárik, Adaptive learning rates for faster stochastic gradient methods. arXiv preprint arXiv:2208. 05287, 2022.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint
2019.
Baihe Huang, Xiaoxiao Li, Zhao Song, and Xin Yang. FL-NTK: A neural tangent kernel-based framework for federated learning analysis. International Conference on Machine Learning, pp- 4423-4436 PMLR, 2021.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in Neural Information Processing Systems, 31, 2018.
Tyler Johnson, Pulkit Agrawal, Haijie Gu, and Carlos Guestrin. Adascale SGD: user-friendly algorithm for distributed training. In International Conference Machine Learning, pp. 491 1- 4920. PMLR, 2020.
Swanand Kadhe, Nived Rajaraman, O Ozan Koyluoglu, and Kannan Ramchandran. FastSecAgg; Scal- able secure aggregation for privacy-pree preeeering federated learning. arXiv preprint arXiv:2009. 1244, 2020.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Ad- vances and open problems in federated learning. Foundations and Trends@ in Machine Learning, 14(1-2):1-216 2021.
Sai Praneeth Karimireddy, Quentin Rebjock, Sebastian Stich, and Martin Jaggi. Error feedback fixes SignSGD and other gradient compression schemes. In Proceedings of the 36th International Conference on Machine Learning, volume 97, pp. 3252-3261. PMLR, 2019.
Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning. arXiv preprint NXXv:2008.03600 2020a.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference Machine Learning, pp. 5132-5143, PMLR, 2020b.
Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. Tighter theory for local sgd on identical and heterogeneous data. In International Conference on Artificial Intelligence and Statistics, pp. 4519 4529. PMLR, 2020.
Diederík Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poster), 2015. URL.tttp:/arxxi orq/ abs het.//111..6.00
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2:429 2020.
11

---

Published as a conference paper at ICLR 2023
Nicolas Loizou, Sharan Vaswani, Issam Hadj Laradji, and Simon Lacoste- Julien. Stochastic polyak step-size for sgd: An adaptive learning rate for fast convergence. In International Conference on Artificial Intelligence and Statistics. pp. 1506-1314 PMLR, 2021.
Grigory Malinovsky, Konstantin Mishchenko, and Peter Richtárik. Server-side stepsizes and sampling without replacement provably help in federated optimization. arXiv preprint arXiv:2201. 1066,
2022.
Yura Malitsky and Konstantin Mishchenko. Adaptive gradient descent without descent. In Pro- ceedings of the 37th International Conference Machine Learning, volume 119 of PMLR, pp- 6702-671 67 12, 2020.
Jan Mandel. Convergence of the cyclical relaxation method for linear inequalities. Mathematical programming, 30(2):218-225 1984.
Brendan McMMaan,, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication-effin learning of deep networks from decentralized data. Artificial Intelli- gence and Statistics, pp. 1272- PMLR, 2017.
Konstantin Mishchenko, Grigory Malinovsky, Sebastian Stich, and Peter Richtarik. ProxSkip: Yes! Local gradient steps provably lead to communication acceleration! Finally! Proceedings of the 39th International Conference on Machine Learning, volume 162, pp. 15750-15769, PMLR, 2022.
Aritra Mitra, Rayana Jaafar, George Pappas, and Hamed Hassani. Linear convergence in federated learning: Tackling client heterogeneity and sparse gradients. Advances in Neural Information Processing Systems, 34: 14666-1461,, 2021.
Guy Pierra. Decomposition through formalization in a product space. Mathematical Programming, 28(1):96-115, 1984.
Boris Teodorovich Polyak. Minimization of unsmooth functionals. USSR Computational Mathematics and Mathematical Physics, 9(3):14-29, 1969. ISSN 0041-5553. doi: https://doi.orr/100 1016/
0041-5553(69)90001--5
Marcos Raydan. On the barzilai and borwein choice of steplength for the gradient method. IMA Journal of Numerical Analysis, 13(3):321-322 1993.
Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Representations, 2021.
Tijmen Ticleman, Geoffrey Hinton, et al. Lecture 6.5 6.5-RMSProp: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning. 4(2):26-31, 2012.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and Vincent Poor. inconsistency problem in heterogeneous federated optimization. Advances in Neural Information Processing Systems, 33:7611 1-7233, 2020.
Tackling the objective
Dong Yin, Ashwin Pananjady, Max Lam, Dimitris Papailiopoulos, Kannan Ramchandran, and Peter Bartlett. Gradient diversity: key ingredient for scalable distributed learning. In International Conference on Artificial Intelligence and Statistics, pp. 1998-2007 PMLR, 2018
Yaodong Yu, Alexander Wei, Sai Praneeth Karimireddy, Yi Ma, and Michael I Jordan. TCT: Convexifying federated learning using bootstrapped neural tangent kernels. arXiv preprint arXiv:2207.06343, 2022.
Kai Yue, Richeng Jin, Ryan Pilgrim, Chau-Wai Wong, Dror Baron, and Huaiyu Dai. Neural tangent kernel empowered federated learning. In International Conference on Machine Learning, pp. 25783-25803. PMLR, 2022.
Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprint 2012.
12

---

Published as a conference paper at ICLR 2023
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understand- ing deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiy Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33:15383-15533 2020.
13

---

Published as a conference paper at ICLR 2023
APPENDIX
A Additional Related Work
15
Table of Notation and Schematic
16
B. Table of Notation
16
B.2 Schematic of Client-Server communication FedExP
16
Proofs
17
Proof of Lemma
17
C.2
Convergence Analysis for Convex Objectives
18
Convergence Analysis for Non-Convex Objectives
4
Exact Projection with Gradient Descent for Linear Regression
23
D Additional Experiments and Setup Details
26
D.I
Impact of Averaging Iterates for Neural Networks
26
D.2 Dataset Details
26
D33
Hyperparameter Details
................................
27
D.4
Sensitivity of FedExP
27
D.5
Additional Results
28
Combining Extrapolation with SCAFFOLD
32
Combining Extrapolation with Server Momentum
33
14

---

Published as a conference paper at ICLR 2023
A ADDITIONAL RELATED WORK
In this section, we provide further discussion on some additional related work that complements our discussion in Section I
Adaptive Step Size in Gradient Descent. Here we briefly discuss methods for tuning the step size in gradient descent and the challenges in applying them to the FL setting. Early methods to tune the step size in gradient descent were based on line search (or backtracking) strategies (Armijo 1966; Goldstein 1977). However, these strategies need to repeatedly compute the function value or gradient within an iteration, making them computationally expensive. Another popular class of adaptive step sizes is based on the Polyak step size (Polyak) 1969 Hazan & Kakade 2019 Loizou eet al.) 2021). Similar to FedExP, the Polyak step size is derived from trying to mininize the distance the optimum for convex functions. However it is not clear how this can be extended to the federated setting where we only have access to pseudo-gradient, Also, the Polyak step size requires knowledge of the function value at the optimum which hard estimate. Another related class of step sizes the Barzilai-Borwe these are known to provably work only for quadratic functions (Raydan 1993 Burdakov al. 2019) only. A recent work (Malitsky & Mishchenko 2020) alleviates some of the concerns associated with these classical methods by setting the step size as an approximation of the inverse local Lipschitz constant; however it is again not clear how this intuition can be applied to the federated setting. An orthogonal line of work has focused on methods that adapt to the geometry of the data using gradient information in previous iterations, the most popular among them being Adagrad (Duchi al.) 2011) and its extensions RMSP rop (Tieleman et al.) 2012) and Adadel (Zeiler 2012). There exist federated counterparts of these algorithms, namely FedAdagrad however, as show in experiments these methods can fail to even outperform FedAvg in standard FL. tasks.
stepsize (Barzilai & Borwein 1988). However, to the best of our knowledge,
Overparameterization in FL. Inspired by the success of analyzing deep neural networks in the neural tangent kernel (NTK) regime (Jacot et al. 2018 Arora et al., 2019) Allen-Zhu al. 2019), recent work has looked at studying the convergence of overparameterized neural networks in the FL setting. Huang et al. 2021) and Deng et al. (022) show that for a sufficiently wide neural network and proper step size conditions, FedAvg will converge presence of data heterogeneity. We note that these works are primarily concerned with convergence analysis, whereas focus is developing practical algorithm that is inspired by characteristics in the overparameterized regime for speeding up FL training. Another recent line of work has looked at utilizing NTK style Jacobian features for learning FL model in just few rounds of communication Yu et al. 2022; Yue et al. 2022). While interesting, these approaches are orthogonal to our current work.
globally optimal solution even in the
15

---

Published as a conference paper at ICLR 2023
B TABLE OF NOTATION AND SCHEMATIC
B.1 TABLE OF NOTATION
Table 2: Summary of notation used in paper
Symbol
Description
L2 norm
Number of clients
((..)
Loss function
D
Dataset at i-th client
(.(w)
Local objective at i-th client
F(w)
Global objective at server
q
Client step size
fg
Server step size
w(O
Global model at round t
fg
(
FedExP
step size round t
w
Local model at i-th client at t-th round and k-th iteration Number of local SGD steps
A
Update of i-th client at round
(
Average of client updates at round f Set of minimizers of F) (w) Number of communication rounds Small constant added to denominator of FedExP step size Global minimum
ST
w
F
Minimum value of global objective L-smoothness constant used in Assumption
a2
Upper bound on variance of client gradients at optimum (see Assumption 2
Upper bound variance of client gradients (see Assumption 3
B.2 SCHEMATIC OF CLIENT-SERVER COMMUNICATION IN FEDEXP
At each round 4, the server first sends global model w(() to all clients. Clients perform local optimization (!( to compute their local models w! (f,r) and send back their update A
w (1.r and norm update | to the server. This procedure is illustrated in Figure 5
O | |III
1) Server sends global will clients
3) Clients send back 1/5 w() -w and !14!"11
SG
Client
Client
Client M
W
2) Clienta perform local tesining
Figure 5: Schematic of client-server communication in FedExP.
16

---

Published as a conference paper at ICLR 2023
C
PROOFS
We first state some preliminary lemmas that will used throughout the proofs.
Lemma 2. (Jensen's 's inequality) For any E R i E (1, M):
M
2
M
1
1
a
|
(13)
M
M
i=1
M
2
Σ
a
< M
|a.1|
(14)
i=1
2=1
We also note the following known result related to the Bregman divergence. Lemma 3.
(Khaled et al. 2020 If F is smooth and convex, then
IVF(w)- VF(w")] s (I(F(w) - F(w'=((===)) (VF(w") w w-w)).
(15)
Lemma 4. (Co-coercivityy of convex smooth function) If F is L-smooth and convex then,
(VF(w)- - VF(W'), ww'w)),-ww> 2 T [VF(w) VF(W)]]
(16)
A direct consequence of this lemma is,
WF((),.-wW - 2 L [F(2))?
(17)
where w' is minimizer of F(w).
C.1 PROOF OF LEMMA 1
Let (w) be the local objective at a client and w be the global minimum. From the overparameteri- zation assumption, we know that w is also a minimizer for F; (w). We have,
w (1,R)
w'
(c,k-1) m V F (w.
(t.k ()) w=
(18)
(1,k-1) w 1| 2m (WW, (1 (4. w;
(1,6-1)
w + +m
VF(w
(1-
|
(19)
w (1,k-1) w'
2m
VF(W (t)-)))
2
+ +nī ((
(20)
L
(t,k-1)
w
m
VF(W
())))
(21)
L
where 20
follows from 17 and 21 to T we have,
follows from m7 S A. Summing the above inequality from
w w-1²² 5 (1) w
111 T
V
(22)
k=0
Thus we have,
M
1
(tr) w
w
(() w ||
W MT
Σ
VF
(23)
M
(()) w' ||²
(24)
This completes the proof of this lemma.
17

---

Published as a conference paper at ICLR 2023
C.2 CONVERGENCE ANALYSIS FOR CONVEX OBJECTIVES
Our proof technique is inspired by Khaled al. (2020) with some key differences. The biggest difference is the incorporation of the adaptive FedExP server step sizes which Khaled et al. (2020) does account for. Another difference is that provide convergence guarantees in terms of number of rounds T while Khaled et al. (2020) focus on number of iterations T = TT. We highlight the specific steps where we made adjustments the analysis Khaled al. (2020) below.
We begin by modifying Khaled et al. (2020) Lemma 11 and Lemma 13) to bound client drift in every round instead of every iteration.
Lemma 5. (Bounding client aggregate gradients)
3L M
V W
//
ΣΣ
wi
(1,k)
- F((w)) + 3TO
(25)
Proof of Lemma 5
M
VF w /
M
VW (1)))
() - VFi(w(0) VF(W(0)- VF +V(W")+A((WW)) (26)
VF((W)
M
M
3
-
FF((W(u)))
3
Σ
())( - vFF(w)) (27)
M
M
M
[V](2")|1
M
3L-
M
W(t (t.)) - w//0|² -
-)) - -F*) + 3TO2
(28)
M
The first term in 288 follows from T-smoothness of (w(w,, the second term follows from Lemma and the third term follows from bounded noise at optimum.
Lemma 6. (Bounding client drift)
M ΣΣ
w. w.(bbi]? &
(()
(r - F(w")) + 7 1)0 (29)
M
18

---

Published as a conference paper at ICLR 2023
Proof of Lemma 6
w-10)
w (2
(w: (w: (111))
` M
(30)
M
$ के M
(31)
S TTT 1) M
V (w /
(32)
i= K=U
M
S (T8 1) 1
w(10 -w w/(ibb/??
+ (1))((-1)(((((())) - F(w"))
(33)
M
k=0
+ +3177 - 1)0
1 M
ΣΣ
((t)
w! f
+ - -F(w"))
(34)
2M
1=1 k=0
+ +
1)0²
where (33 uses Lemma 5 and 34 uses 1 S 1
GrL
Therefore we have,
w(t) w wcbbl11 12 (2)(<(())))(((
1)T((F(w(()) - F(ww)) 7 1)03 (35)
M
Proof of Theorem 1:
We define the following auxiliary variables that will used in the proof.
Aggregate Client Gradient: h Σ ((
(36)
We also define h())
1
ΣM 1i=1 h
Recall that the update of the global model can be written as w(++)) = ((() -
-1/g (1) (0))
We have
w-(1 (++1) w | (()
( m))()) - -w*/ w11
(37)
(1) w" ||² -2 (2) m ww' w", ((()))))))))))))))))
+ +(((g))) (1) nì
(1)
(38)
< ((()) w || 2m m w w", h())
+ 11
(t)
m
M
h
(4)
(39)
M
A
where (39) follows from 19g
M h Incquality (39) key step in our proof and the differentaaiing factor in our approach from Khaled et al. (2020). Following similar technique as (f)
Khaled et al. (2020) bound (17g ) [ [0 will end up requiring the condition vi S 1//T.mg (t) (t)
which cannot be satisfied in our setup due to the adaptive choice of 1/g Therefore we first upper
19

---

Published as a conference paper at ICLR 2023
and focus further bounding this quantity in the rest of the proof, which does not require the aforementioned condition. Note that this comes at the expense of the additional T3 error seen in our final convergence bound in Theorem T
bound (ng )np
|| by 1/0nn?
M Ei=1 h (1)
M
Therefore.
| - w' - w((0) -w"
w-12 2n m (w' - w", (0)
M
(8)
(40)
M
Bounding T2
have,
M T2 =
(41)
M
i=1
M
(w ((ak)
(42)
M
M
V ( ((.R)) 2
(43)
M
k==
$
L
ΣΣ
(c.A) wi
- +3
F²)
(44)
M
where (33 follows from Jensen's inequality and and 444 follows from Lemma 5
Bounding Ti
1 M
T
w (1.))
(45)
M
M
w(()) w VFF(W!
(46)
M 1=1 A=0
We have,
((()) -w') -.', (w(((()))))) = w(()) -(w())
w(( (w(((())))) + (w((k)) -w", WR,
(47)
From T-smoothness of F; we have,
(w()) w. (F((((()))) 2 Fe(w(()) (((w!())) - 2 w(0) w.
(C.K)
(48)
From convexity of F; we have,
(w (t.k) w V..(W ((.)( Pi(ww(.At) F.(w)).
(49)
Therefore, adding the above inequalities we have,
(w()) (F((((()))) 2 Fi(w(9) F.(w") L /w()) w.
(1,))
(50)
Substituting 50 in 46 we have,
Ti 2 (0) -F(w")
w(() ww(
(51)
2M
Am
k=
20

---

Published as a conference paper at ICLR 2023
Here would like to note that the bound for Tj is our contribution and is needed in our proof due to the relaxation in (39). The bound for T2 follows a similar technique as Khaled et al. (2020 Lemma 12).
Substituting the bounds for Ti and T2 in 440 we have,
W (++1)
w 112 2ng" MT(L - 3mm 3 -F(w") + mi==
w
w
M
(3ml") n TL² L2 + 11 m L)
w. (f.8) -w
w//0/²
M
km0
- (((() w'
M1 MET
F F(w((0) - F(w")) + 3ng (1) 12
(52)
M
1
w((ck) w(lll
2mg m T.
M
k=0
< w w
-1 WI (F(w(0) F(w")) + 1//
(53)
24ng (1) m/ (T (F) (FW-))))) + 12mg (1) 1// 1/² -
L)
< (t) w² m11²
1l9
(1)
T
(P( P(w"))))(("))))---)))))))))>>))) k 02
(54)
3
12
(t) n (T 1)Lo
where both 52 and 55 use S 5.F and 553 uses Lemma
Rearranging terms and averaging over all rounds we have,
11
(1)
((w( - F(w")
3
/w(O) 11
(1) UITTo 36/ -1)Lo (55)
ΣE 199
(1)
2 1g T
This implies,
[w(0) - w 11²
F(w(T))- - -F(""= S
+ (n/ 11)_0) +0 (nIO²?
(56)
THT Ei 7l
(2)
T- 2100 ( (t)) Ei 1/
where (T)
This completes the proof.
C.3 CONVERGENCE ANALYSIS FOR NON-CONVEX OBJECTIVES
Our proof technique is inspired by Wang al. (020)) and we use one of their intermediate results to bound client drift in non-convex settings as we describe below. We highlight the specific steps where we made adjustments to the analysis of Wang et al.) (2020) below.
We begin by defining the following auxiliary variables that will used in the proof.
Normalized Gradient:
h
(((((()))
(57)
We also define (1()
(t) ΣM
h
Lemma 7. (Bounding client drift in Non-Conver Setting)
M
1
v (w(( -n 201l2
VF W ( (0/l
+5m 110
(58)
M
21

---

Published as a conference paper at ICLR 2023
Proof of Lemma 7
Let D LT( - 1). We have the following bound from equation (87) in
Wang et al. (2020)
1
D D
VF VF(W(00)
+ 1-D
Do²
(59)
M
From 12 S 6.1. we have D S which implies --D S
9 and 8
I-D 1 8.
Therefore have,
M - nno0 - (F(
9D
(
+
8 o
(60)
M
+5 $ VF(
1))
(61)
Proof Theorem 2:
The update of the global model can be written as follows,
(++1) (t) n/ hc)(
(62)
Now using the Lipschizzssmoothness assumption we have,
22
(0)
F(w(++1)) F(w((t) S -n niT D (VF(w(")), (1)
(11
(+)
(63)
2
(t)
M
s n Ti (VF(w()), h²t)
T1
ni
(64)
2M
ΣM M 6000
h
where (64) uses 179
As in the convex case, inequality 664 is key step in our proof
and the differentiating factor in our approach from Wang et al. (2020). Following a similar technique as Wang et al. (020) to bound (71)) ) [1] /2 will need the condition yp S 1/2LTT/ (1)
which cannot be satisfied in our setup due to the adaptive choice of 1/ Therefore first upper /22 and focus on further bounding this quantity in the rest of the proof, which does not require the aforementioned condition. Note that this comes at the expense of the additional T3 error seen in our final convergence bound in Theorem
bound (1g (1) ) || 1? by nh 2 Σ h 1b/"///////
Therefore have,
(t)
2 2
M
F(w((+1))- F(w(()) S -n (FF(w()), h()))
11g
(65)
2M
24
Bounding Ti
We have,
M
Ti = ()(w
Σ
h
(8)
(66)
M
i=0
M
M
2/
1
VF (0))
||@
+
1
h
VF(w(()
h
(4)
(67)
M
2
M
1=1
1=1
M
1
VF
1
V
(w(((()) -1(")
(68)
2
2M
22

---

Published as a conference paper at ICLR 2023
where (67) uses (a. = |||| [b|] definition of the global objective function F.
2/2 ||a -bll and 68 uses Jensen's inequality and the
Bounding T2
We have,
T2
h
(0)
(69)
M
M
Σ
|n.
(w((w(()) VFF(((()) - VF(w(()) + [F((
(70)
M
M
(11/0- [))((( ()
Σ
VF(((w()) - WFFWW))) / +||>r(/(>>])))
(71)
M
1=1
3
Σ
bb.0 -
VF
+ +o3 +
VF
(72)
M
where 771) uses Jensen's inequality.
72 uses bounded data heterogeneity assumption.
Here would like to note that the bound for T2 is our contribution and is needed in our proof due to the relaxation in 39 The bound for T₁ follows a similar technique as in Wang et al. 2020).
Substituting the Tj and T3 bounds into 65 we have,
M
- - -1/0
F(w((+1))- F(w (1) S -n T
(73)
2
2M
i=1
T
30 و + 3
VF 1|²
3
I1.00
M
M
-n1²0 n1.01² + 3m Lo?
Σ
(l:(w(()) -
+
M
(74)
+ + Lo + 5n L - 1)02)
(75)
where 77 uses 1/t
S
75 uses Lemma
Tr
Thus rearranging terms and averaging over all rounds we have,
EE 1l
(#)
(1)2
8(F(w(0)) - FF) + LO 1)0g + 2Am
(76)
Σ 110
Er- 119 INT
This implies,
min (F((((0)) 0//p 5
(F(w(0)--F)) ( F²))
+ (1 + 0 (mLro?).
(77)
ΣT- Mli n1T
(t)
This completes the proof.
C.4 EXACT PROJECTION WITH GRADIENT DESCENT FOR LINEAR REGRESSION
Let F(w) A|ww bil² where A is (n x d) matrix and b is a dimensional vector. We assume that d 2 n here and A has rank n. The singular value decomposition (SVD) of A can be written as,
A UEVT E
O
V
(78)
V
23

---

Published as a conference paper at ICLR 2023
where U is an (n n) orthogonal matrix, Σ is (n x n) diagonal matrix, V1 is (d x n) matrix with orthogonal columns and V2 is a (d x (d (d-n)) n)) matrix with orthogonal columns. Here V1 is basis for the row space of A, while V2 is a basis for the null space of A. We first prove the following lemmas about the set of minimizers
F(w) and the projection on this set.
Lemma 8. The set of minimizers of F F(w) is given by,
S' =V2 VV blw E Rd).
(79)
Proof. Let w V2 x V VI U for some x E Rd. We have,
Aw = UE, V (V2V x V
(80)
= b
(81)
where the last line uses V2 = 0, V 11 = 1,UT = I. This implies | - -bll = Thus
any in S is minimizer of F(w).
Now let w be a minimizer of F(w), implying Aw = -E2 V w b. We have,
w = V2V2 w V.Vi w
(82)
= =V2V w + ViE, U
(83)
where (22 uses ViVI + V.V. and 83 uses UE, V w = b. Thus any minimizer of F(w)
must lie in S"
Combining the above statements we have.
w is a minimizer of F(w) + w S".
(84)
which completes the proof.
Lemma 9. The projection of any E IRd on S is given by.
Ps. (w) arg J[w w'l2 V.V. ViE V U w'Es-
(85)
Proof. When w E S", it is easy to see that this holds. Therefore we consider the case where w & S
Let V2V, w + V, E TTT and Ps- (w-(w) V2V1 wo + V where wo o w. We have,
w V2V2 Wo V,Σ UT bll
(86)
V2 (w- - wo) V, V,
bl2 (ViV + V2V. 1I
(87)
V (w wo)l VV
(88)
V//22V V (w- 2 +|||2+ x||
(89)
||w - x|2
(90)
leading to contradiction. The cross term in 88 is zero since V V2 0. Equation (89 the definition of x.
follows by
We now show that rumning gradient descent on F(w) starting from w with sufficiently small step size converges to Ps-(w),
Lemma 10. Let ("),, w...... the iterates generated by running gradient descent on F(w) with (0) w and learning rate Th S Amaxo where Amax is the largest eigen value of A A. Then limy- (T) Ps- (w).
Proof. By the gradient descent update we have,
w ((+1) = =w (t)
m(A T Aw (1) - A ATb)
(91)
= A A A (1)) FMATT
(92)
24

---

Published as a conference paper at ICLR 2023
Therefore,
((7) -I (I- A AATA)) w'() (0) +
Σ
R A
(93)
T
V
R
V(I
(94)
T
= -V -
V2 w + V
(1-
E,U'b.
(95)
In the limit T and with rp S Amax we have,
T-1
lim (I n =EtT"* and lim T-000
Σ [I (I-
Σ
(96)
T-000
1/
Thus,
lim w(T) V2V2 (0 (.) + VV T-400
(97)
= =Ps- (w))
(98)
Ps- (w).
(99)
C.4.1 IMPROVING LOWER BOUND IN (8) IN THE CASE OF EXACT PROJECTIONS
(t,7) Let S be convex and let w E S, for all E [M]. We assume that w w! = Psy (w())) Vi [M].
i.e. the local models are exact projection of (()) on their respective solution sets. From 8 we have
ΣM (wit) w w",
(w(t) W w., AA))
(ng (t) Dopt
(100)
|
M |
(wi()) w', 1 follows,
We can lower bound
(w(1) (100-w"AA -,,')) (t))) w ((rT) + w! (())) w", w()) -w! w
(101)
(() w + (m((.)) -w", w(()) w (
(102)
2 ((()) w m/
(103)
[0
(104)
(105)
where 103) uses the fact that w (t,r) w", ((t) w(( ((T))
2 0 following the properties of projec-
tion (Boyd Dattarro 2003).
Thus we have,
DM 1012
ops
M
2
(106)
Note here the improvement by a factor of 2 in the lower bound compared to 8
25

---

Published as a conference paper at ICLR 2023
D ADDITIONAL EXPERIMENTS AND SETUP DETAILS
Our code is available the following link htttarr/gginuub.ccm//rbaaaaaaas
D.1 IMPACT OF AVERAGING ITERATES FOR NEURAL NETWORKS
As discussed in Section we find that setting the final FedExP model as the average of the last two iterates also improves performance when training neural networks in practical FL scenarios. To demonstrate this, we consider an experiment on the CIFAR- 10 dataset with 10 clients, where the data at each client is distributed using Dirichlet distribution with 0.3. We set the number of local steps to be T 20 and train a CNN model having the same architecture as outlined in McMahan et al.) (2017) with full client participation. Figure shows the training accuracy as a function of the last iterate and the average of last two iterates for FedAvg and FedExP. We see that the last iterate of FedExP has oscillating behavior that can hide improvements in training accuracy. On the other hand, the average of the last two iterates of FedExP produces a more stable training curve and shows considerable improvement in the final accuracy. Note however that this improvement only shows for F ;edExP; averaging iterates does not make significant difference for FedAvg.
FedExP (Aug. of lest two iterates)
FedExP (Last iterane)
FedExP
FedAvg (Avg. last two (terates)
FedAug (Last iterale)
80
Tii Acturayy (%)
Server Step Size ngt
3.0
wwwm
40-
20
80
20
100
Training rounds
Training rounds
Figure Benefit of averaging the last two iterates for FedExP in training CNN model on CIFAR-10. Note that averaging does not make significant difference for FedAvg.
D.2 DATASET DETAILS
Here provide more details about the datasets used in Section
Synthetic Linear Regression. In this case we assume that the local objective of each client is given by (.w) = ||A,w w bill where A, E (30x 1000), b, E R030 and w 1000 We set the number of clients to be M = 20. Note that since 2 Eis n, this is an overparameterized convex problem. To generate A; and b;, follow similar process
et al. (2020). We have (A4)j: -N N(m,, La) and (b:); = w (A.)j: where m; ~ N(u,, 1), w. N N(y; 1), u, (O, 0.1), W 0.1).
EMNIST. EMNIST an image classification task consisting of handwritten characters associated with 62 labels. The federated EMNIST dataset available at at(Caldas et al. (2019) is naturally partitioned into 3400 clients based on the identities of the character authors. The number of training and test samples is 671,585 and 77,483 respectively.
CIFAR-10/100 CIFAR-10 is natural image dataset consisting of 60,000 32x32 images divided into 10 classes. CIFAR- 100 uses a finer labeling of the CIFAR images to divide them into 100 classes making it harder dataset for image classification. In both cases the number of training examples and test examples is 50,000 and 10,000 respectively. To simulate a federated setting. we artificially partition the training data into 100 clients following the procedure outlined in Hsu et al. (2019).
CINIC-10. CINIC- 10 is a natural image dataset that can be used as a direct replacement of CIFAR for machine learning tasks. It is intended to act harder dataset than CIFAR- 10 while being easier than CIFAR- 100. The number of training and test examples is both 90,000, We partition the training data into 200 clients in this case, following similar procedure for CIFAR.
26

---

Published as a conference paper at ICLR 2023
D.3 HYPERPARAMETER DETAILS
For our baselines, we find the best performing "/ and vR by grid-search tuning. For FedExP we search for and !n. This is done by running algorithms for 50 rounds and finding the parameters that achieve the highest training accuracy averaged over the last 10 rounds. We provide details of the grid used below for each experiment below.
Grid for Synthetic.
For FedAvg and SCAFFOLD, the grid for f is {10', 100.5 10 10' 100). For FedAdagrad,
the grid for fly is (10-1, 10-00-.. 0.5 10-0 100.5 10-} 3. For FedExP we keep = 0 in this experiment
as (6) is satisfied in this case. The grid for nn is {10-2 10-15, 10-1,, 10-1, 10-0.0.. 1100) for all algorithms.
Grid for Neural Network Experiments.
For FedAvg and SCAFFOLD the grid for 72 {10-1, 1000,1101000000 1,10-.3 0.5 100 100.5 100). For FedAdagrad, the grid for 7g is {10-2 10-5... 100 10 100). For FedExP the grid for is (10-3 10-2.3 2.3 10-2 10-10- 10-1). The grid for v is {10-2, 10-1.5, 10-1, 10-0.5, 100) for all algorithms.
We use lower values of fl the grid for FedAdagrad based on observations from Reddi al. (202T) which show that FedAdagrad performs better with smaller values of the server step size. We provide details of the best performing hyperparameters below.
Table Base- 10 logarithm of the best combination of and rp for FedExP and combination of and fl for baselines. For the synthetic dataset we keep = 0 for FedExP.
FedExP
FedAvg
SCAFFOLD
FedAdagrad
Dataset
Vlg
Tl
11
flg
1
Synthetic
-1
-1
-1
-1
-1
EMNIST
-1
-0.5
00.5
-0.5
-0.5
-0.5
CIFAR-10
-3
-2
0
-2
0
-2
-1
-2
CIFAR-100
-3
-2
-2
-2
-1
-2
CINIC-100
-3 -2
0
-2
0
-2
-1
-2
Other hyperparameters are kept the same for all algorithms. In particular, we apply a weight decay of 0.0001 for all algorithms and decay n by factor of 0.998 in every round. We also use gradient clipping to improve stability the algorithms as done in previous works (Acar et al. 2021). In all experiments fix the number of participating clients be 20, minibatch size be 50 (for the synthetic dataset this reduces to full-batch gradient descent) and number local updates r to be 20.
D.4 SENSITIVITY OF FEDEXP TO
To evaluate the sensitivity of FedExP to e, we compute the training accuracy of FedExP after 500 rounds for varying and on different tasks. For each task, fix np to be the value used in our experiments in Section and only vary c. The results are summarized below.
Table 4: Training accuracy obtained by FedExP with different choices of after 500 rounds of training on various tasks. Value of y is fixed for each task (10 0.5 for EMNIST and 10 10-2 for others). Results averaged over last 10 rounds.
Dataset
= 10
=
10 2.5
< 10
10
10
EMNIST
85.40
86. 86.26
85.73
85.49
84.90
CIFAR-10
84.79
77.82
77.63
77.66
77.64
CIFAR-100
59.01
44.76
44.21
44.37
44.40
CINIC-10
66.31
60.93
61.05
60.47
60.96
We see that the sensitivity of e similar to that of the 7 parameter which is added to the denominator of FedAdam and FedAdagrad (Reddi al. 2021) prevent the step size from blowing up.
27

---

Published as a conference paper at ICLR 2023
Keeping too large reduces the adaptivity of the method and makes the behavior similar FedAvg. At the same time, keeping e too small may not also be beneficial always as seen in the case of EMNIST. In practice, find that grid search for in the range {10 10 2.5 10 10 10-1-1, 1.5 10 usually suffices to yield a good value c. general rule thumb would be to start with = 10-3 and increase till the performance drops.
D.5 ADDITIONAL RESULTS
In this section, we provide additional results obtained from our experiments.
Synthetic Linear Regression. Note that for the synthetic linear regression experiments there is no test data. Also note that there is no randomness in this experiment since clients compute full-batch gradients with full participation. We provide the plot of "/g for FedExP in Figure 7 see that FedExP takes much larger steps in some (but not all) rounds compared the constant optimum step size taken by our baselines, leading to a large speedup. Recall that we also let
in this
(t) experiment (since it aligns with our theory) which also explains the larger values of flg taken by FedExP in this case.
FedExP (ours)
FedAvg
SCAFFOLD
FedAdagrad
Server Slep Size 1i
150
100
50
0
25
50
75
100
125
150
175
200
Training rounds
Figure Global learning rates for synthetic data with linear regression. Results from single instance of experiment.
EMNIST. For EMNIST we observe that SCAFFOLD gives slightly better training loss than FedExP towards the end of training. As described in Section extrapolation can be combined with the variance-reduction in SCAFFOLD (the resulting algorithm is referred to as SCAFFOLD ExP) to further improve performance. This gives the best result in this case shown in Figure g
FedExP (ours)
SCAFFOLD-ExP (ours)
FedAvg
SCAFFOLD
FedAdagrad
0.70
86
0.65
85
Server Step Size n
Training Loss
0.60
Training Apcuracy
0.55
0.50
0.45
0.40
80
250
500
0
250
500
250
500
Training rounds
Training rounds
Training rounds
Figure 8: Additional results for EMNIST dataset. Mean and standard deviation from experiments with 20 different random seeds. The shaded areas show the standard deviation.
28

---

Published as a conference paper at ICLR 2023
CIFAR-10, CIFAR-100 and CINIC-10. From Figure 3 and Figures 9
that FedExP
comprehensively outperforms baselines in these cases, achieving almost 10% 20% higher accuracy than the closest baseline by the end of training. The margin of improvement is most in CIFAR- 100, which can be considered as the toughest dataset our experiments. This points to the practical utility of FedExP in challenging FL scenarios.
FedExP (ours)
FedAvg
SCAFFOLD
FedAdagrad
1.4
80
2.0
Step Size n/d
Training Loss
1.2
Training Accuracy
70
1.5
60
1.0
50
0.5
40
0.4
30
0.0
250
500
250
500
250
500
Training rounds
Training rounds
Training rounds
Figure 9: Additional results for CIFAR-10 dataset. Mean and standard deviation from experiments with different random seeds. The shaded areas show the standard deviation.
FedExP (ours)
FedAvg
SCAFFOLD
FedAdagrad
60
Server Step Size m
50
Training Loss
Training Accuracy
3-
40
30
2
250
500
0
250
500
250
500
Training rounds
Training rounds
Training rounds
Figure 10: Additional results for C0F dataset. Mean and standard deviation from experiments with different random seeds. The shaded areas show the standard deviation.
FedExP (ours)
FedAvg
SCAFFOLD
FedAdagrad
1.8
2.0
65
Training Accuracy (%)
Server Step nlg
Training Loss
60
1.5
55
1.0
1.2
50
0.5
1.0
45
40
0
250
500
250
500
250
500
Training rounds
Training rounds
Training rounds
Figure 11: Additional results for CINIC-10 dataset. Mean and standard deviation from experiments with 5 different random seeds. The shaded areas show the standard deviation.
29

---

Published as a conference paper at ICLR 2023
Long-Term Behavior of Algorithms and Comparison with FedProx. To evaluate the long-term behavior of different algorithms, we ram the experiments for 2000 rounds. Here, we also consider an additional algorithm, namely FedP rox, for comparison. For fair comparison, parameter of FedP rox for each dataset, by doing a grid search over the range {10-3 10-2 10-1, (0-1.1) .1] done in the original FedProx paper (Li et al., 2020). The results of EMNIST, CIFAR-10, CIFAR- 100, and CINIC-10 in Figures [2H [
have tuned the
Table 5 are from experiments with 3 different random seeds. Except for the synthetic dataset, the plots show mean and standard deviation values across all the random seeds and also over moving average window of size 20.
FedExP (ours)
FedAvg
SCAFFOLD
FedAdagrad
FedProx
Synthetic
EMNIST
CIFAR-10
CIFAR-100
CINIC-10
066
10"1
4.U
2.4
Mean Squane Ener
10
Tradura Lona
Trahing Loos
10"79
SE
2.0
Trahing
S4
1000
2000
1000
2000
1000
2000
1000
2020
1000
2000
Training rounds
Trakiing munds
Training rounds
rounds
Traniig iounds
Thaining munds
Figure 12: Training loss results of FedExP, FedAvg, SCAFFOLD, FedAdagrad and FedP rox on the Synthetic, EMNIST, CIFAR- MA,R-I,CIIAA--II and CINIC-10 datasets for 2000 rounds.
FedExP (ours)
FedAvg
SCAFFOLD
FedAdagrad
FedProx
EMNIST
CIFAR-10
CIFAR-100
CINIC-10
Trartrg Kcainti (5)
Traniigg Acaanttyyy (S))
Tra Fouitic (56)
75
Transing (%g
20-
88
60
80
85
40
34
40
1500
1500
500
1000
1500
Figure 13: Training accuracy results of FedExP, FedAvg, SCAFFOLD, FedAdagrad and
FedPros rox on the EMNIST, CIFAR-10, CIFAR-100 and CINIC-10 datasets for 2000 rounds.
FedExP (ours)
FedAvg
SCAFFOLD
FedAdagrad
FedProx
CIFAR-100
IB50
85.
Ter: Accunoyy (X0
Tost Accumay (3)
60
75.0
80
12.8
500
1000
1600
2000
500
1000
1500
2000
503
1200
1500
2000
500
1030
1600
2000
Trining munds
Training mundss
Training mondds
Trainigg minds
Figure 14: Test accuracy results of FedExP, FedAvg, SCAFFOLD, FedAdagrad and FedP rox on the EMNIST, CIFAR-10, CIFAR-100 and CINIC-10 datasets for 2000 rounds.
30

---

Published as a conference paper at ICLR 2023
Table Test accuracy obtained by FedExP and baselines after 2000 rounds of training on various tasks. Results are averaged across 3 random seeds and last 20 rounds.
Dataset
FedExP
FedAvg
SCAFFOLD
FedAdagrad
FedProx
EMNIST
86.96 + 0.58
85.78 + 0.35
86.22 + 0.35
85.53 + 1.04
85.77 + 0.39
CIFAR-10
82.94 + 0.42
80.10 + 0.56
82.02 + 0.30
80.21 + 0.60
80.16 + 0.59
CIFAR-100
54.65 + 0.49
49.63 + 0.37
49.40 = 0.38
49.64 0.39
49.47 + 0.31
CINIC-10
66.45 1.28
64.87 + 0.44
64.61 = 0.49
64.87 + 0.44
64.52 + 0.45
We see that FedExP continues to outperform baselines including FedP rox in the long-term behavior well.
31

---

Published as a conference paper at ICLR 2023
E COMBINING EXTRAPOLATION WITH SCAFFOLD
As described Section 6 the extrapolation step can added to the SCAFFOLD algorithm in a similar way FedExP. The detailed steps of this SCAFFOLD ExP algorithm shown in Algorithm
Algorithm 2 SCAFFOLD-EXP
Input: w(0), control variate c00), C; (0)
Vi [M], number of rounds T, local iteration steps т,
parameters T):
2: For 0, ......... 1 communication rounds do:
3: Global server do:
4: Send w('), c( to all clients
5: Clients i [M]] in parallel do:
6: Set w; (t,0) f- (1.0)
7:
For l == - local iterations do:
8:
Update w! (L.k22) + w (l,k) - [V; (w (t.k) c((.k)
+ c(0)
;5
9:
Compute A. + w(t) w;
((())) and !!()) + c(t)
L
A
Send A and ! to the server
10:
Update local control variate c, (++1) c0 v!!0
12:
Global server do:
Compute A (1) +- -i=1 A (1) and (2) mmx 1, 1. M || 11/2 2M (14(0)]++)) +e)}
13:
14:
Update global model with w (r++) w(t) 77/ (1)
Compute (()) N7 M (1(1)
15:
16:
Update global control variate with (1++1) + c()) (1(1)
32

---

Published as a conference paper at ICLR 2023
F COMBINING EXTRAPOLATION WITH SERVER MOMENTUM
We begin by recalling some notation from our work. The vector w(() is the global model at round t and (1) is the average of client updates at round t. The server momentum update at round can
be written as v(t) (1) + (y((-1) (let v = 0) and the global model update can be written as
((++1) w((") 1/
(t) (t). Our goal is now to find T]g that minimizes
(t)
||w(++)) w"ll We have,
w -
w ||
w(() w²
+ (ng )²
wool
2
2n²² (w/(1) - -w", (107)
Setting the derivative of the RHS of (107) to zero we have,
(w((t) w w(,v())
m Jpp
(108)
|/
Our goal now is to find a lower bound on (w(t)) w,, v()). We have the following lemma.
Lemma 11. Assume that (w(()) w", A (1) 2 m(t)
MM
(4)
/M (see Appendix C.4.1 for
all L 2 0 and v (m(() + EK=D (3/2) //2 ||v(0)]
for all < L 1. Then,
t-1
(())) w", v(00) 2 (t)
+ Σ
(B/2)==*))) m (8)
(109)
k=
which implies,
(! + Et-1 (3/2 mm
(ng opp
(110)
2 1(00/
Proof. We proceed via a proof by induction. (w(()) - w", (')) (w(")
The statement clearly holds at = since
wx, A (()) 2 m))
Now assuming the lemma holds at L - 1 we have,
w(0 w", v())
(0) w. " (1) - 3 w(()) - ww, w",r(1-1))
(111)
(w(()) w (() ) + B (w(t-1)- (
1//6--1 (t-11) (1-1))((-11 - w", v(t-1))
(112)
2 (t) + 3 ( (wllt-)) - w", ,(t-1)) (((-11) /
(113)
m (1) Σ (B
(114)
where the last line follows from the fact that (w(t-1) 1) - w", v((--()) (t-11) 1) (8/2
10 (t-1) < (m( + (322 1- Thus propose to keep the following server step size when using server momentum,
and
(t) (9 + (B/2
(115)
v(1lv(() +
where (t) EM
/M. Note that we also add small constant to the denominator to prevent the step size from blowing up done for FedExP. We call server momentum with this step size as FedExP-M.
We compare the performance of FedExP-te M with FedAdam and FedAvg-M (FedAvg with server momentum) on the CIFAR- 10 and CIFAR- datasets as shown in Figures [7] where the mean and standard deviation values computed over 3 random seeds and moving average window of
33

---

Published as a conference paper at ICLR 2023
size 20. The experimental setup is the same described in Section 6 The hyperparameters m, for FedExP-KX and 7) 7, for F F'dAdam and F FedAvg-M were tuned following a similar process as described in Appendix D.3] and their resulting values are in Table
Table Base-10 10 logarithm of the best combination of c and d for FedExP-M and combination of and fly for FedAdam and FedAvg-M.
FedExP
FedAdam
FedAvgm-M
Dataset
1R
Vig
CIFAR-10
-3
-2
-2
-2
-2
CIFAR-100
-3
-2
-2
-2
-2
FedExP-M (ours)
FecAAg-M
FedAdam
CIFAR-10
CIFAR-100
10"
10"
10"
300
1000
1000
2000
000
1000
1000
2000
Training rounds
Training rounds
Figure 15: Training loss results of FedExP MM, FedAdam and FedAvg M on the CIFARIO and CIFAR 100 datasets.
FodExP-M
Fod/ug-M
FodAdam
CIFAR-10
CIFAR-100
Thahiig Thahhhg Acouracy (%)
80-
90
500
1000
1600
2000
500
1000
1800
2000
Triiiin
Training rounds
Figure 16: Training accuracy results of FedExPPPB MM, FedAdam and FedAvg- M on the CIFARIO and CIFARION 100 datasets.
FedExP-M (ours)
FedAug-M
FedAdam
CIFAR-10
CIFAR-100
90
60
Text Acturssyy (5)
88
48
Test
80
1000
1000
1500
2000
Traihing rounds
Trainig rouncs
Figure 17: Test accuracy results of FedExP M, FedAdam and FedAvg M on the CIFARIO and CIFARIO0 100 datasets.
Our result shows that server momentum can be successfully combined with extrapolation for the best speed-up among all baselines. The behavior of FedAdam and F FedAvg-M are quite similar in these experiments which be attributed to the dense nature of the gradients in image classification as
34

---

Published as a conference paper at ICLR 2023
We note that this is only preliminary result and future work will look to study the effect of combining server momentum and extrapolation more rigorously.
discussed in Section 6
35


Reviews:
[{"review_id": "aRi8b12D9fz", "reviewer": "Reviewer_LWDX", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "- In federated learning settings, aggregating the learning of all participating clients has always been a challenging issue in cases of data heterogeneity across clients which is very probable in real world scenarios. To tackle this, recent works have optimized the server aggregation process of FedAvg algorithm by treating client updates as \"pseudo gradients\" with fixed server step size. However, using a fixed step size is not ideal for varying data distribution across clients.\n- To tackle this issue, this paper proposes a novel federated learning algorithm FedExP which adaptively determines the server step size in each round of federated learning such that it can handle varying degree of similarity in data distribution across clients. In doing so, the authors first establish a novel connection between FedAvg and POCS algorithms for overparameterized convex objectives and further draw the inspiration from the EPPM extrapolation algorithm (used to speed up POCS algorithm) resulting in adaptively determining server step size in FedExP for better and fast convergence of global model.\n- Key feature of FedExP is that there is virtually no requirement for additional communication, computation, or storage at clients or the server.\n- Experimental evaluation demonstrates that FedExP algorithm converges faster and consistently outperforms baselines over 5 datasets.\n- novelty is simple and limited yet effectively applied to federated learning paradigm and have shown improved performance in terms of speed-up/efficiency and client-server data communication\n\n", "summary_of_the_review": "The dynamic tuning of server step size results in efficient (faster) convergence and performance boost across multiple datasets while efficiently tackling the issue of knowledge interference during server aggregation due to heterogeneous data distribution across clients.\n\n- The novelty is limited and simple yet effective in Federated learning paradigm \n- the proposed method is effective both in terms on efficiency and prediction performance \n- addresses the sustainability concerns and may be applicable and scale when a number of clients participate in Federated averaging, especially in IOT scenarios and cloud-edge continuum", "strength_and_weaknesses": "Strengths:\n- novelty is simple yet effective application in federated learning paradigm\n- The motivation is well founded and the claims are sound.\n- Paper is clearly presented and easy to follow.\n- Mathematical formulation is very detailed and explanatory.\n- Proposed FedExP model is effective and efficient.\n- Proposed FedExP model consistently outperforms two strong federated learning baseline methods on global model convergence time and test accuracy over multiple datasets.\n\nWeaknesses:\n- inspired by related works, especially the application of extrapolation in federated learning\n- novelty is limited (yet effective)\n- Missing comparison with the relevant FedProx method.\n\tLi, Tian, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. \"Federated optimization in heterogeneous networks.\" Proceedings of Machine Learning and Systems 2 (2020): 429-450.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ZamqbB-lWLD", "reviewer": "Reviewer_AfRg", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies adaptive server step sizes for federated learning (FL). The authors propose a method called FedExP, which uses a server step size akin to the POCS algorithm from the literature on finding feasible points inside the intersections of sets. The authors use a similar expression to estimate how much one can extrapolate in the direction proposed by FedAvg.\n\nThe authors prove convergence of the method for convex and nonconvex problems. They also run experiments on convex problems that clearly illustrate the theory, as well as on FL benchmarks.\n\n-------------\n## update after reading authors' feedback\n\nThe authors addressed my main concern about the short runs. I can see the figures in the appendix that include longer runs.  \nThe authors also added the missing assumptions to Theorem 2.  \nI believe it is very clear that the paper should be accepted.", "summary_of_the_review": "The paper shows a good mixture of theoretical and empirical results, both of which seem of sufficient quality. There are a couple of ways the paper can be improved, but I believe that it meets the bar for acceptance.", "strength_and_weaknesses": "## Strength\n1. The work is written in a clear and intuitive way.  \n2. The considered problem is quite meaningful as server side stepsizes have been shown to give an improvement in practice and building better adaptive estimates is a promising direction.  \n3. The theory is rigorous and considers both convex and nonconvex settings.  \n4. There are experiments to support the theory as well as to show that the method will be useful in practice. \n\n## Weaknesses:  \n1. In the experiments, it appears that most algorithms were terminated before convergence. Therefore, we do not see the final values of the test accuracy. In practice, these values are very important, so I would like to know if the authors can present plots with longer training.  \n2. Theorem 2 is analyzed under the assumption that the gradients are deterministic and I do not see if this stated anywhere in the text.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "9wJe3l6x6M", "reviewer": "Reviewer_Qfnz", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the problem of federated optimization. It shows a connection between FedAvg and Projection Onto Convex Sets algorithm (under overparametrized convex setting) and based on this proposes a method (FedExP) with adaptive server step-size strategy by using the gradient diversity measure. The authors perform a theoretical analysis of FedExP and show convergence under standard reasonable assumptions in convex and non-convex settings. In the experimental section proposed method is studied in both illustrative and realistic FL settings and is shown to outperform selected baselines.", "summary_of_the_review": "This work is insightful from theoretical perspective and presents useful practical results. I have some comments and questions (in the previous section) on parts of the paper, which will be hopefully addressed. Overall, I think that submission is worth publishing after the minor revision.", "strength_and_weaknesses": "## Strengths\n\n1. Novel and insightful view of an important Federated Averaging algorithm, which has the potential for fruitful theoretical and practical advancements in the area of federated optimization. \n\n3. The suggested step-size strategy can be applicable to a wider set of local gradient methods, which was illustrated for SCAFFOLD in the Appendix.\n\n2. Thorough empirical study of the proposed method, which shows its insightful properties on toy problems and promising results for synthetic and realistic benchmarks.\n\n## Weaknesses\n\n1. Convergence analysis focuses on the local full-batch and full client participation. which is quite far from current practical federated learning settings. Although I do not consider it a big issue, as almost all local optimization methods do not provide any theoretical benefits over simplest distributed (mini-batch stochastic) gradient descent baseline for heterogenous problems.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%