F ED E X P: S PEEDING UP F EDERATED A VERAGING VIA E XTRAPOLATION
A BSTRACT
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 Pro jection 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.
I NTRODUCTION
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 a clients device, as clients only communicate model updates with the server.
Federated Averaging or FedAvg , first introduced by McMahan 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 et al., 2016; Kadhe et al., 2020).
Slowdown Due to Heterogeneity.
One of 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 et al., 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 in every round (Kairouz et al., 2021).
Various techniques have been proposed to combat this slowdown, among the most popular being variance reduction techniques such as Karimireddy et al.
; Mishchenko et al.
; Mitra et al.
, but they either lead to clients becoming stateful, add extra computation or communication requirements or have privacy limitations.
Server Step Size.
Recent work has sought to deal with this slowdown by using two separate step sizes in FedAvg – a client step size used by the clients to 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 et al., 2021).
To achieve the fastest convergence rate, these works propose keeping the client step size as O (cid:0) 1 /τ √ T (cid:1) and the server step size as O (cid:0) √ τM (cid:1) , where T is the number of communication rounds, τ is the number of local steps and M is the number of clients.
Using a small client step size mitigates client drift, and a large server step size prevents global slowdown.
While this idea may be asymptotically optimal, it is not always effective in practical non-asymptotic and communication-limited settings (Charles & Koneˇcn`y, 2020).
[-e ( η g ) 5.8 8.2 42 
-]
[Figure]
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 1).
Also, if local objectives differ significantly, then it may be beneficial to use smaller values of the server step size (Malinovsky et al., 2022).
Therefore, we seek to answer the following question: For a 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 objectives?
In general, it 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.
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.e., the number of model parameters is 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 et al., 2018) and the small datasets collected by edge clients in the FL setting.
In this overparameterized regime, the global minimizer becomes a 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 in FedAvg .
In particular, we derive a 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 our proposed algorithm FedExP , which is a method to adaptively determine the server step size in each round of FL based on the pseudo-gradients in that round.
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 overparameterized 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 a point in the intersection of convex sets.
• The proposed FedExP algorithm is simple to implement with virtually no additional communication, computation, or storage required at clients or 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 – 2 × faster than FedAvg and most competing baselines on standard FL tasks.
Related Work.
Popular algorithms for adaptively tuning the step size when training neural networks include Adagrad (Duchi et al., 2011) and its variants RMSProp (Tieleman et al., 2012) and Adadelta (Zeiler, 2012).
These algorithms consider the notion of coordinate-wise 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 pseudo-gradients as done by Reddi 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 a client-centric notion of adaptivity that utilizes the heterogeneity of client updates in each round.
The work closest to us is Johnson et al.
(2020), which proposes a method to adapt the step size for large-batch training by estimating the gradient diversity (Yin et al., 2018) of a minibatch.
This result has been improved in a recent work by Horváth et al.
However, both Johnson et al.
; Horváth et al.
(2022) focus on the centralized setting.
In
FedExP , we use a similar concept, but within a federated environment which comes with a 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.
P ROBLEM F ORMULATION AND P RELIMINARIES
As in most standard federated learning frameworks, we consider the problem of optimizing the model parameters w ∈ R d to minimize the global objective function F ( w ) defined as follows:
[Equation]
where F i ( w ) := [Equation] i ∈D i ℓ ( w , δ i ) is the empirical risk objective computed on the local data set D i at the the i -th client.
Here, ℓ ( · , · ) is a loss function and δ i represents a data sample from the empirical local data distribution D i .
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 (1).
Our algorithm and analysis can be directly extended to the case where client objectives are unequally weighted, e.g., proportional to local dataset sizes |D i | .
FedAvg .
We focus on solving (1) using FedAvg (McMahan et al., 2017; Kairouz et al., 2021).
At round t of FedAvg , the server sends the current global model w ( t ) to all clients.
Upon receiving the global model, clients perform τ steps of local stochastic gradient descent (SGD) to compute their updates { ∆ ( t ) i } Mi =1 for round t as follows.
Perform Local
[Equation]
Compute Local Difference
: ∆ (
[Equation]
where w ( t, 0) i = w ( t ) for all i ∈
[ M ] , η l is the client step size and ∇ F i ( w ( t,k ) i , ξ ( t,k ) ) represents a stochastic gradient computed on the minibatch ξ ( t,k ) i sampled randomly from D i .
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 ( t +1) =
1 M (cid:80) Mi
=1 w ( t,τ ) i .
To improve over this, recent work (Reddi et al., 2021; Hsu et al., 2019) has focused on optimizing the server aggregation process by treating the client updates ∆ ( t ) i as “pseudo-gradients” and multiplying by a server step size when aggregating them as follows.
Generalized
[Equation]
where ¯∆ ( t ) = 1 M (cid:80) Mi =1 ∆ ( t ) i is the aggregated client update in round t and η g acts as server step size .
Note that setting η g = 1 recovers the vanilla FedAvg update.
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 η g that we use in every round.
P ROPOSED A LGORITHM : F ED E X P
Before discussing our proposed algorithm, we first highlight a useful and novel connection between
FedAvg and the POCS algorithm used to find a point in the intersection of some convex sets
. 3. M OTIVATION FOR E XTRAPOLATION
Connection Between FedAvg and POCS in the Overparameterized Convex Regime.
Consider the case where the local objectives of the clients { F i ( w ) } Mi =1 are convex.
In this case, we know that the set of minimizers of F i ( w ) given by S ∗ i = { w : w ∈ arg min F i ( w ) } is also a convex set for all i ∈
[ M ] .
Now let us assume that we are in the overparameterized regime where d is sufficiently larger than the total number of data points across clients.
In this regime, the model can fit all the training data at clients simultaneously and hence be a minimizer for all local objectives.
Thus we assume that the global minimum satisfies w ∗
∈
S ∗ i , ∀ i ∈
[ M ] .
Our original problem in (1) can then be reformulated as trying to find a point in the intersection of convex sets {S ∗ i } Mi =1 since w ∗ ∈ S ∗ i , ∀ i ∈
[ 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 1 .
Generalized POCS update:
[Equation]
where P i ( w ( t ) POCS ) is a projection of w ( t ) POCS on the set S ∗
i
and λ is known as the relaxation coefficient (Combettes, 1997).
Extrapolation in POCS .
Combettes (1997) notes that POCS has primarily been used with λ = 1 , with studies failing to demonstrate a systematic benefit of λ < 1 or λ > 1 (Mandel, 1984).
This prompts Combettes (1997) to study an adaptive method of setting λ , first introduced by Pierra (1984) as follows:
[Equation]
Pierra (1984) refer to the POCS algorithm with this adaptive λ ( t ) as Extrapolated Parallel Projection
Method ( EPPM ).
This is referred to as extrapolation since we always have λ ( t ) ≥ 1 by Jensen’s inequality.
The intuition behind EPPM lies in showing that the update with the proposed λ ( t ) always satisfies (cid:13)(cid:13) w ( t +1) [Equation] − w
[Equation]
(cid:13)(cid:13) 2 < (cid:13)(cid:13) w [Equation] t ) POCS − w ∗ (cid:13)(cid:13) 2 , thereby achieving asymptotic convergence.
Experimental results in Pierra (1984) and Combettes (1997) show that EPPM can give an order-wise speedup over POCS , motivating us to study this algorithm in the FL context.
3. I NCORPORATING E XTRAPOLATION IN FL
Note that to implement POCS we do not need to explicitly know the sets {S ∗ i } Mi =1 ; we 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 ( t ) from the server and run multiple
SGD steps to compute an “approximate projection” w ( t,τ ) i of w ( t ) on their solution sets S ∗ i .
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 η 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 η g 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:
[Equation]
where w ( t ) and { w ( t,τ ) i } Mi =1 are the global and local client models, respectively, at round t 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 ( t ) 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 i ( w ) be convex and L -smooth for all i ∈
[ M ] and let w ∗ be a common minimizer of all F i ( w ) .
Assuming clients run full-batch gradient descent to minimize their local objectives with η l ≤ 1 /L , then (6) holds for all t and τ ≥ 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 be satisfied at least during the initial phase of training when (cid:13)(cid:13) w ( t ) − w ∗ (cid:13)(cid:13) 2 is large and clients make common progress towards a minimum.
1: Input: w (0) , number of rounds T , local iteration steps τ , parameters η l , ϵ 2: For t = 0 , . .
, T − 1 communication rounds do : 3:
Global server does: 4:
Send w ( t ) to all clients 5:
Clients i ∈
[ M ] in parallel do: 6:
Set w ( t, 0) i ← w ( t, 0) 7:
For k = 0 , . . .
, τ − 1 local iterations do: 8:
Update w ( t,k +1)
i ← w ( t,k ) i − η l ∇ F i ( w ( t,k ) i , ξ ( t,k ) i ) 9:
Send ∆ ( t ) i ← w ( t ) − w ( t,τ ) i to the server 10:
Global server does: 11:
Compute ¯∆ ( t ) ← 1 M (cid:80) Mi =1 ∆ ( t ) i and η ( t ) g ← max (cid:110) 1 , (cid:80) Mi
=1 (cid:13)(cid:13)
∆ ( t ) i (cid:13)(cid:13)
(cid:46) 2 M (cid:16)(cid:13)(cid:13)
¯∆ ( t ) (cid:13)(cid:13) 2 + ϵ (cid:17)(cid:111) 12:
Update global model with w ( t +1) ← w ( t ) − η ( t )
g ¯∆ ( t )
Given that (6) holds, we now consider the generalized FedAvg update with a server step size η ( t ) g in round t .
Our goal is to find the value of η ( t ) g that minimizes the distance of w ( t +1) to w ∗ :
[Equation]
Setting the derivative of the RHS of (7) to zero we have,
[Equation]
where the last inequality follows from ⟨ a , b ⟩ = 12 [ ∥ a ∥ 2 + ∥ b ∥ 2 − ∥ a − b ∥ 2 ] , definition of ∆ ( t ) i in (3) and (6).
Note that depending on the values of { ∆ ( t ) i } Mi =
[Equation]
we may have ( η ( t ) g ) opt ≫ 1 .
Thus [Equation] we see that (6) acts as a suitable replacement for projection to justify the use of extrapolation in FL settings
[Equation]
3. P ROPOSED A LGORITHM
Motivated by our findings above
[Equation]
we propose the following server step size for the generalized
FedAvg update at each round:
[Equation]
We term our algorithm Federated Extrapolated Averaging or FedExP , in reference to the original
EPPM algorithm which inspired this work.
Note that our proposed step size satisfies the property that (cid:12)(cid:12) ( η ( t ) g ) opt − ( η ( t ) g ) FedExP (cid:12)(cid:12) ≤ (cid:12)(cid:12) ( η ( t ) g ) opt − 1 (cid:12)(cid:12) when (6) holds, which can be seen by comparing (8) and (9).
Since (7) depends quadratically on η ( t ) g , we can show that in this case (cid:13)(cid:13) w ( t +1) − ( η ( t ) g )
FedExP
¯∆ ( t ) − w ∗ (cid:13)(cid:13) 2 ≤ (cid:13)(cid:13) w ( t +1) − w ∗ (cid:13)(cid:13) 2 , implying we are at least as close to the optimum as the FedAvg update.
In the rest of the paper, we denote ( η ( t ) g ) FedExP as η ( t ) g when the context is clear.
Importance of 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 (cid:13)(cid:13)
¯∆ ( t ) (cid:13)(cid:13) 2 ≈ 0 but (cid:13)(cid:13)(cid:13) ∆ ( t ) i (cid:13)(cid:13)(cid:13)
2 ̸ = 0 .
Thus we propose to add a small positive constant ϵ to the denominator in (9) to prevent this blow-up.
For a large enough ϵ 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 as Adam (Kingma & Ba, 2015) and Adagrad (Duchi et al., 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 η ( t ) g using only the updates of participating clients, i.e., the averaging and division in (9) will be only over the clients that participate in the round.
Furthermore, since the server only needs to estimate the average of pseudo-gradient norms, η ( t ) g can be computed with secure aggregation, similar to computing ¯ ∆ ( t ) .
Connection with Gradient Diversity.
We see that our lower bound on ( η ( t ) g ) opt naturally depends on the similarity of the client updates with each other.
In the case where τ = 1 and clients run full-batch gradient descent, our lower bound (8) reduces to (cid:80) Mi
=1 (cid:13)(cid:13) ∇ F i ( w ( t ) )
(cid:13)(cid:13) 2 (cid:14) 2 M (cid:13)(cid:13) ∇ F ( w ( t ) )
(cid:13)(cid:13)which is used as a measure of data-heterogeneity in many FL works (Wang et al., 2020; Haddadpour &
Mahdavi, 2019).
Our lower bound suggests using larger step-sizes as this gradient diversity increases, which can be a 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.
Acar et al.
(2021), which propose additional regularization terms or adding control variates to the local client objectives to limit the impact of heterogeneity.
C ONVERGENCE A NALYSIS
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 underparameterization and non-convexity.
We would like to emphasize that (6) is 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 η l be sufficiently small and the standard assumptions stated below.
Challenge in incorporating stochastic noise and partial participation.
Our current analysis focuses on 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 η ( t ) g and the pseudo-gradients { ∆ ( t ) i } Mi =1 .
To be more specific, if we use ξ ( t ) to denote the randomness at round t , then E ξ ( t ) (cid:104) ( η ( t ) g ) ¯∆ ( t ) (cid:105) ̸
= E ξ ( t ) (cid:104) ( η ( t ) g ) (cid:105) E ξ ( t ) (cid:2)
¯∆ ( t ) (cid:3) which significantly complicates the proof.
This is 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 F i ( w ) is differentiable and L -smooth for all i ∈
[ M ] ,
[Equation] [Equation]
[Equation]
Assumption 2.
(Bounded data heterogenenity at optimum)
The norm of the client gradients at the global optima w ∗ is bounded as follows:
[Equation]
Theorem 1.
( F i are convex) Under Assumptions 1,2 and assuming clients compute full-batch gradients with full participation and η l ≤ 16 τL , the iterates { w ( t ) } generated by FedExP satisfy,
[Equation]
= noise at optimum
[Equation]
where η ( t ) g is the FedExP server step size at round t and ¯ w ( T ) = (cid:80) [Equation] [Equation]
0 η
[Equation]
g w
[Equation]
[Equation]
For the non-convex case, we need the data heterogeneity to be bounded everywhere as follows.
Assumption 3.
(Bounded global gradient variance)
There exists a constant σ 2 g > 0 such that the global gradient variance is bounded as follows.
[Equation]
Theorem 2.
( F i are non-convex)
Under Assumptions 1, 3 and assuming clients compute full-batch gradients with full participation and η l ≤ 16 τL , the iterates { w ( t ) } generated by FedExP satisfy,
[Equation]
where η ( t ) g is the FedExP server step size at round t .
Discussion.
In the convex case, the error of
FedAvg can be bounded by
O [Equation] ∥ w (0) − w ∗ ∥ 2 /η l τT (cid:1) + O (cid:0) η 2 l τ ( τ − 1)
Lσ 2 ∗ (cid:1) (Khaled et al., 2020) and in the non-convex case by O (cid:0) ( F ( w 0 ) − F ∗ ) /η
l τT (cid:1) + O (cid:0) η 2 l L 2 τ ( τ − 1) σ 2 g (cid:1) (Wang et al., 2020).
A careful inspection
[Figure]
reveals that the impact of T 1 on convergence of FedExP is different from FedAvg (effect of T 2 is the same).
We see that since (cid:80) T − 1 t =0 η ( t ) g ≥ T , FedExP reduces T 1 faster than FedAvg .
However this comes at the price of an increased error floor due to T 3 .
Thus, the larger step-sizes in FedExP help us reach the vicinity of an optimum faster, but can 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 η l appropriately.
Moreover, in the overparameterized convex regime where σ 2 ∗ = 0 , the effect of T 2 and T 3 vanishes and thus FedExP clearly outperforms FedAvg .
This aligns well with our initial motivation of using extrapolation in the overparameterized regime.
F URTHER I NSIGHTS INTO F ED E X P
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 (cid:13)(cid:13) w ( t ) − w ∗ (cid:13)(cid:13)
but not necessarily F ( w ( t ) )
− F ( w ∗ ) .
Recall that our original motivation for the FedExP step size was aimed at trying to minimize the distance to the optimum give by (cid:13)(cid:13) w ( t +1) − w ∗ (cid:13)(cid:13) 2 , when (6) holds.
Doing so satisfies (cid:13)(cid:13) w ( t +1) − w ∗ (cid:13)(cid:13) 2 ≤ (cid:13)(cid:13) w ( t ) − w ∗ (cid:13)(cid:13) 2 but does not necessarily satisfy F ( w ( t +1) ) ≤ F ( w ( t ) ) .
To better illustrate this phenomenon, we consider the following toy example in R 2 .
We consider a setup with two clients, where the objective at each client is given as follows:
[Equation]
We denote the set of minimizers of F 1 ( w ) and F 2 ( w ) by S ∗ 1 = { w : 3 w 1 + w 2 = 3 } and
S ∗ 2 = { w : w 1 + w 2 = 3 } respectively.
Note that S ∗ 1 and S ∗ 2 intersect at the point w ∗ =
[0 , 3] , making it a global minimum.
To minimize their local objectives, we assume clients run gradient descent with τ → ∞ in every round 2 .
Figure 2 shows the trajectory of the iterates generated by
FedExP and FedAvg .
We see that while (cid:13)(cid:13) w ( t ) − w ∗ (cid:13)(cid:13) 2 decreases monotonically for FedExP ,
F ( w ( t ) ) does not do so and in fact has an oscillating nature as we discuss below.
Understanding oscillations in F ( w ( t ) ) .
We see that the oscillations in F ( w ( t ) ) are caused by
FedExP iterates trying to minimize their distance from the solution sets S ∗ 1 and S ∗ 2 simultaneously.
The initialization point w (0) is closer to S ∗ 1 than S ∗ 2 , which causes the FedExP iterate at roundto move towards S ∗ 2 , then back towards S ∗ 1 and so on.
To understand why this happens, consider the case where ∆ ( t )= 0 , ∆ ( t )̸
= 0 .
In this case, we have η ( t ) g = 2 and therefore w ( t +1) =
w ( t ) − 2 ¯∆ ( t ) =
w ( t,τ ), which indicates that FedExP is now trying to minimize (cid:13)(cid:13)(cid:13) ∆ ( t +1)(cid:13)(cid:13)(cid:13) 2 .
This gives us the intuition that the FedExP update in round t is trying to minimize the objectives of the clients that have (cid:13)(cid:13)(cid:13) ∆ ( t ) i (cid:13)(cid:13)(cid:13) 2 ≫ 0 .
While this leads to a temporary increase in global loss F ( w ( t ) ) in
[Figure]
some rounds as shown in Figure 2, it is beneficial in the long run as it leads to a faster decrease in distance to the global optimum w ∗ .
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 F 1 ( w ) and F 2 ( 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 FedExP also extend to training neural networks with multiple clients in practical FL scenarios (see Appendix D.1).
In practice, the number of iterates to average over could also be a hyperparameter for FedExP , but we find that averaging the last two iterates works well, and we use this for our other experiments.
E XPERIMENTS
We evaluate the performance of FedExP on 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 FedExP when (6) holds.
For realistic FL tasks, we consider image classification on the following datasets i) EMNIST
(Cohen et 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 an overparameterized problem.
The data at each client is generated following a similar procedure as the synthetic dataset in Li et al.
We use the federated version of EMNIST available at Caldas et al.
(2019), which is naturally partitioned into 3400 clients.
For CIFAR-10/100 we artifically partition the data into 100 clients, and for CINICwe partition the data into 200 clients.
In both cases, we follow a Dirichlet distribution with α = 0 .for
the partitioning to model heterogeneity among client data (Hsu et al., 2019).
For EMNIST we use the same CNN architecture used in Reddi et al.
For CIFAR10, CIFAR100 and CINICwe use a ResNet-18 model (He et al., 2016).
For our baselines, we find the best performing η g and η l by grid-search tuning.
For FedExP we optimize for ϵ and η l 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 3 demonstrate that FedExP clearly outperforms FedAvg and competing baselines that use the best performing η g and η l found by grid search.
Moreover, FedExP does not require additional communication or storage at clients or server unlike SCAFFOLD and FedAdagrad .
The orderwise improvement in the case of the convex linear regression experiment confirms our theoretical motivation for FedExP outlined in Section 3.2.
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 a consistent
[Table]
speedup of over 1 . 4 − 2 × over FedAvg .
This verifies that FedExP also provides performance improvement in more general settings with realistic datasets and models.
Plots showing η ( t ) g 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 in some (but not all) rounds compared to the constant optimum step size taken by our baselines, leading to a large speedup.
Comparison with FedAdagrad .
As discussed in Section 1, FedAdagrad and FedExP use different notions of adaptivity; FedAdagrad uses coordinate-wise adaptivity, while FedExP uses client-based adaptivity.
We believe that the latter is more meaningful for FL 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-wise 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 see 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.
Karimireddy et al.
(2020a);
Yu et al.
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 SCAFFOLD may not provide any speedup in practice.
[Figure]
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 E).
Figure 4 shows the result of our proposed extrapolated SCAFFOLD on the CIFARdataset.
Interestingly, we observe that while SCAFFOLD alone fails to outperform FedAvg , the extrapolated version of 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.
It is also possible to add extrapolation to algorithms with server momentum (Appendix F).
C ONCLUSION
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 an 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 theoretical and empirical perspectives of FedExP .
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 variance-reduction method in SCAFFOLD , for greater speedup.
Future work will study the convergence analysis of FedExP with stochastic gradient noise and the incorporation of extrapolation into a wider range of algorithms used in FL A CKNOWLEDGMENTS
This work was supported in part by NSF grants CCF 2045694, CNS-2112471, ONR N00014-23-1- and the CMU David Barakat and LaVerne Owen-Barakat Fellowship.
A PPENDIX A Additional Related Work
B
C Proofs
C.1 Proof of Lemma 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.
Objectives . . . . . . . . . . . . . . . . . . . .
C.
C.4 Exact Projection with Gradient Descent for Linear Regression . . . . . . . . . . .
D Additional Experiments and Setup Details
D.1 Impact of Averaging Iterates for Neural Networks . . . . . . . . . . . . . . . . . .
D.2 Dataset Details . . .
D.3 Hyperparameter Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.4 Sensitivity of FedExP to ϵ . . . .
D.
Additional Results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E Combining Extrapolation with SCAFFOLD F Combining Extrapolation with Server Momentum
A A DDITIONAL R ELATED W ORK
In this section, we provide further discussion on some additional related work that complements our discussion in Section 1.
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 et al., 2021).
Similar to FedExP , the Polyak step size is derived from trying to minimize the distance to 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-gradients.
Also, the Polyak step size requires knowledge of the function value at the optimum which is hard to estimate.
Another related class of step sizes is the Barzilai-Borwein stepsize (Barzilai & Borwein, 1988).
However, to the best of our knowledge, these are known to provably work only for quadratic functions (Raydan, 1993; Burdakov et 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 et al., 2011) and its extensions RMSProp (Tieleman et al., 2012) and Adadelta (Zeiler, 2012).
There exist federated counterparts of these algorithms, namely FedAdagrad ; however, as we show in our experiments these methods can fail to even outperform FedAvg in standard FL tasks.
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 et 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.
(2022) show that for a sufficiently wide neural network and proper step size conditions, FedAvg will converge to a globally optimal solution even in the presence of data heterogeneity.
We note that these works are primarily concerned with convergence analysis, whereas our focus is on developing a 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 a FL model in just a few rounds of communication (Yu et al., 2022;
Yue et al., 2022).
While interesting, these approaches are orthogonal to our current work.
B T ABLE OF N OTATION AND S CHEMATIC B. T ABLE OF N OTATION
[Table]
B. S CHEMATIC OF C LIENT -S ERVER COMMUNICATION IN F ED E X P
At each round t
[Table]
the server first sends global model w ( t ) to all clients.
Clients perform local optimization on w ( t ) to compute their local models w ( t,τ ) i and send back their update ∆ ( t ) i = w ( t ) i − w ( t,τ ) i and norm of update (cid:13)(cid:13)(cid:13) ∆ ( t ) i (cid:13)(cid:13)(cid:13) 2 to the server.
This procedure is illustrated in Figure
[Table]
+VZYf0zxJfsvz9qyJ2b56lXLjy7m7kF+b/cqKJsZ1wrU1aERq4aZZXmVPDFifhUWZSk5x6AtMrPyuUZWJDkDxn5JSQ3Ld8GR4N+st1Pvg66e3G7jnX2kr1imyxh79ke+8wO2ZDJYCPYCT4EH8MX4W54EH5aScOg/fOcXYvwyx/NnMjS</latexit> 3) Clients send back  ( t ) i =
w ( t )  w ( t,⌧ ) i and k  ( t )
i k<latexit sha1_base64="Gv0cxlC+5B5gZAK6uC+GYh7cdq8=">AAACAHicbVDLSsNAFJ3UV62vVBcu3ASLUEFKIqIuC25cVrAPaGKYTCft0MmDmRulhGz8ERduXCji1s9w586ln+Gk7UJbD1w4nHMv997jxZxJMM1PrbCwuLS8Ulwtra1vbG7p5e2WjBJBaJNEPBIdD0vKWUibwIDTTiwoDjxO297wIvfbt1RIFoXXMIqpE+B+yHxGMCjJ1XftAMPA89O7zLVu0ioc2YCTw8zVK2bNHMOYJ9aUVOpl7+vhe3DScPUPuxeRJKAhEI6l7FpmDE6KBTDCaVayE0ljTIa4T7uKhjig0knHD2TGgVJ6hh8JVSEYY/X3RIoDKUeBpzrzc+Wsl4v/ed0E/HMnZWGcAA3JZJGfcAMiI0/D6DFBCfCRIpgIpm41yAALTEBlVlIhWLMvz5PWcc06rVlXKg0TTVBEe2gfVZGFzlAdXaIGaiKCMvSIntGLdq89aa/a26S1oE1ndtAfaO8/bR6Z8Q==</latexit> w ( t,⌧ )<latexit sha1_base64="fXxOmsKy1kGcjU02Z7mKlRvjWHQ
[Table]
[Figure]
[Table]
+Tcl/+G6WPaWdeThDM7hEny4hTLcQxVqIGAMz/AKb07mvDjvzsdiNOcsM6fwB87nD/Nlkyg=</latexit
[Figure]
<latexit sha1_base64="vC017/Bqmjrtiuc75oiguazEKfc=">AAAB/XicbVDJSgNBFOxxjXEbl5uXxiB4CjMi6jGQi8eIZoFkCD09L0mTnoXuN2Icgr/ixYMiXv0Pb/6NnWQOmljQUFTV470uP5FCo+N8W0vLK6tr64WN4ubW9s6uvbff0HGqONR5LGPV8pkGKSKoo0AJrUQBC30JTX9YnfjNe1BaxNEdjhLwQtaPRE9whkbq2ocdhAfMaFXGaUBvQZnwuGuXnLIzBV0kbk5KJEeta391gpinIUTIJdO67ToJehlTKLiEcbGTakgYH7I+tA2NWAjay6bXj+mJUQLai5V5EdKp+nsiY6HWo9A3yZDhQM97E/E/r51i78rLRJSkCBGfLeqlkmJMJ1XQQCjgKEeGMK6EuZXyAVOMoymsaEpw57+8SBpnZfei7N6clypOXkeBHJFjckpcckkq5JrUSJ1w8kieySt5s56sF+vd+phFl6x85oD8gfX5A3bZlS
[Figure]
C P ROOFS
We first state some preliminary lemmas that will used throughout the proofs.
Lemma 2.
(Jensen’s inequality)
For any a i ∈ R d , i ∈ { 1 , 2 , . . .
, M } : (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
M
M (cid:88) i = a i (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ≤
M
M (cid:88) i =1 ∥ a i ∥ 2 , (13) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) M (cid:88)
i = a i (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ≤ M
M (cid:88) i =1 ∥ a i ∥ 2 .
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
[Equation]
Lemma 4.
(Co-coercivity of convex smooth function)
If F is L -smooth and convex then,
[Equation]
A direct consequence of this lemma is,
[Equation]
17) where w ∗ is a minimizer of F ( w ) .
C. P ROOF OF L EMMA
Let F i ( w ) be the local objective at a client and w ∗ be the global minimum.
From the overparameterization assumption, we know that w ∗ is also a minimizer for F i ( w ) .
We have,
[Equation]
[Equation]
[Equation]
where (20) follows from (17) and (21) follows from η l ≤ 1 L .
Summing the above inequality from k = 0 to τ − 1 we have,
[Equation]
Thus we have,
[Equation]
This completes the proof of this lemma.
C. C ONVERGENCE A NALYSIS FOR C ONVEX O BJECTIVES
Our proof technique is inspired by Khaled et 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 not account for.
Another difference is that we provide convergence guarantees in terms of number of rounds T while Khaled et al.
(2020) focus on number of iterations T ′ =
Tτ .
We highlight the specific steps where we made adjustments to the analysis of Khaled et 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
[Equation]
Lemma
[Equation]
(Bounding client aggregate gradients)
[Equation]
Proof of Lemma 5:
[Equation]
The first term in (28) follows from L -smoothness of F i ( w ) , the second term follows from Lemmaand the third term follows from bounded noise at optimum
[Equation]
Lemma 6.
(Bounding client drift)
[Equation] [Equation]
where (33) uses Lemma 5 and (34) uses η l ≤ 16 τL .
Therefore we have,
[Equation]
Proof of Theorem
We define the following auxiliary variables that will used in the proof
[Equation]
Aggregate Client
[Equation]
We also define ¯ h ( t )
= 1 M (cid:80) Mi
=1 h ( t ) i .
Recall that the update of the global model can be written as w ( t +1)
=
w ( t ) − η ( t ) g
η l ¯ h ( t ) .
We have
[Equation]
where (39) follows from η ( t ) g ≤
[Equation]
) ∥ 2 .
Inequality (39) is a key step in our proof and the differentiating factor in our approach from Khaled et al.
Following a similar technique as
Khaled et al.
(2020) to bound ( η ( t ) g ) 2 η 2 l (cid:13)(cid:13)
¯ h ( t ) (cid:13)(cid:13) 2 will end up requiring the condition η l ≤ 1 / 8 Lη ( t ) g , which cannot be satisfied in our setup due to the adaptive choice of η ( t ) g .
Therefore we first upper bound
[Equation]
2 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 T 3 error seen in our final convergence bound in Theorem 1.
Therefore,
[Equation]
Bounding T
We have,
[Equation]
where (43) follows from Jensen’s inequality and and (44) follows from Lemma
[Equation]
Bounding T
[Equation]
We have,
[Equation]
From L -smoothness of F i we have,
[Equation]
From convexity of F i we have,
[Equation]
Therefore
[Equation]
adding the above inequalities we have,
[Equation]
Substituting (50) in (46) we have,
[Equation]
Here we would like to note that the bound for T 1 is our contribution and is needed in our proof due to the relaxation in (39).
The bound for T 2 follows a similar technique as Khaled et al.
Lemma 12).
Substituting the bounds for T 1 and T 2 in (40) we have,
[Equation]
where both (52) and (55) use η l ≤ 16 τL , and (53) uses Lemma 6.
Rearranging terms and averaging over all rounds we have,
[Equation]
This implies,
[Equation]
where ¯ w ( T ) =
(cid:80) T − 1 t =
η ( t ) g w ( t ) (cid:80) T −t =
η ( t ) g .
This completes the proof.
C. C ONVERGENCE A NALYSIS FOR N ON -C ONVEX O BJECTIVES
Our proof technique is inspired by Wang et al.
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.
[Equation]
We also define ¯ h ( t )
= 1 M (cid:80) M i =1 h ( t ) i .
Lemma 7.
(Bounding client drift in Non-Convex Setting)
[Equation]
Proof of Lemma 7: Let D = 4 η 2 l L 2 τ ( τ − 1) .
We have the following bound from equation (87) in
Wang et al.
[Equation]
From η l ≤ 16 τL we have D ≤ 19 which implies
[Equation]
≤ 98 and D
[Equation]
D ≤ 18 .
Therefore we have,
[Equation]
Proof of Theorem 2:
The update of the global model can be written as follows,
[Equation]
Now using the Lipschitz-smoothness assumption we have,
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
[Equation]
(64) where (64) uses η ( t ) g ≤ (cid:80) Mi
=1 (cid:13)(cid:13)(cid:13) h ( t ) i (cid:13)(cid:13)(cid:13)
M ∥ ¯ h ( t ) ∥ 2 .
As in the convex case, inequality (64) is a key step in our proof and the differentiating factor in our approach from Wang et al.
Following a similar technique as Wang et al.
(2020) to bound ( η ( t ) g )
η 2 l τ 2 L (cid:13)(cid:13) ¯ h ( t ) (cid:13)(cid:13) 2 / 2 will need the condition η l ≤ 1 / 2 Lτη ( t ) g , which cannot be satisfied in our setup due to the adaptive choice of η ( t ) g .
Therefore we first upper bound ( η ( t ) g ) 2 η 2 l τ 2 L 2 (cid:13)(cid:13)
¯ h ( t ) (cid:13)(cid:13) 2 by η ( t ) g
η 2 l τ 2 L 1 M (cid:80) Mi =1 (cid:13)(cid:13)(cid:13) h ( t ) i (cid:13)(cid:13)(cid:13) 2 (cid:46) 2 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 T 3 error seen in our final convergence bound in Theorem 2.
Therefore we have,
[Equation] [Equation]
Bounding T
We have,
[Equation] [Equation]
( t
[Equation]
[Equation]
i
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
[Equation] ( t ) )
[Equation]
i = h ( t
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) [Equation] 67) ≥
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
( w ( t ) )
(cid:13)(cid:13)(cid:13) 2 −
[Equation]
M
M (cid:88) i = (cid:13)(cid:13)(cid:13)
[Equation]
where (67) uses ⟨ a , b ⟩ = 12 ∥ a ∥ 2 + 12 ∥ b ∥ 2 − 12 ∥ a − b ∥ 2 and (68) uses Jensen’s inequality and the definition of the global objective function F .
Bounding T
We have,
[Equation]
where (71) uses Jensen’s inequality, (72) uses bounded data heterogeneity assumption.
Here we would like to note that the bound for T 2 is our contribution and is needed in our proof due to the relaxation in (39).
The bound for T 1 follows a similar technique as in Wang et al.
Substituting the T 1 and T 2 bounds into (65) [Equation] we have,
[Equation] [Equation]
where (74) uses η l ≤ 16 τL , (75) uses Lemma 7.
Thus rearranging terms and averaging over all rounds we have,
[Equation]
This implies,
[Equation]
This completes the proof.
C. E XACT P ROJECTION WITH G RADIENT D ESCENT FOR L INEAR R EGRESSION
Let F ( w ) =
∥
Aw − b ∥ 2 where A is a ( n × d ) matrix and b is a n dimensional vector.
We assume that d ≥ n here and A has rank n .
The singular value decomposition (SVD) of A can be written as,
[Equation]
where U is an ( n × n ) orthogonal matrix, Σ is an ( n × n ) diagonal matrix, V 1 is a ( d × n ) matrix with orthogonal columns and V 2 is a ( d × ( d − n )) matrix with orthogonal columns.
Here V 1 is a basis for the row space of A , while V 2 is a basis for the null space of A .
We first prove the following lemmas about the set of minimizers of F ( w ) and the projection on this set.
Lemma 8.
The set of minimizers of F ( w ) is given by, [Equation] Proof. Let w = V 2 V
⊤ 2 x + V 1 Σ − 1 1 U ⊤ b for some x ∈ R d .
We have,
[Equation]
{+where the last line uses V ⊤ 1 V 2 = 0 , V ⊤ 1 V 
+}
=
I , UU ⊤ =
I .
This implies ∥
Aw − b ∥ 2 = 0 .
Thus any w in S ∗ is a minimizer of F ( w ) .
Now let w ∗ be a minimizer of F ( w ) , implying Aw ∗ = UΣ 1 V ⊤ 1 w ∗ =
b .
We have,
[Equation]
{+where (82) uses V 1 V ⊤ 1 + V 2 V ⊤ 
+}
= I and (83) uses UΣ 1 V ⊤ 1 w ∗ = b . Thus any minimizer of F ( w ) must lie in S ∗ . Combining the above statements we have,
[Equation]
which completes the proof.
Lemma 9.
The projection of any w ∈ R d on S ∗ is given by,
[Equation]
Proof.
When w ∈ S ∗ , it is easy to see that this holds.
Therefore we consider the case where w / ∈ S ∗ .
Let x = V 2 V ⊤ 2 w + V 1 Σ − 1 1 U ⊤ b and P S ∗ ( w ) =
V 2 V ⊤ 2 w 0 + V 1 Σ − 1 1 U ⊤ b where w 0 ̸ = w .
We have,
[Equation]
leading to a contradiction.
The cross term in (88) is zero since V ⊤ 1 V 2 = 0 .
Equation (89) follows by the definition of x .
We now show that running gradient descent on F ( w ) starting from w with a sufficiently small step size converges to P S ∗ ( w ) .
Lemma 10. Let w (0) , w (1) , . . be the iterates generated by running gradient descent on F ( w ) with w (0)
=
w and learning rate η l ≤ λ max , where λ max is the largest eigen value of A ⊤ A . Then lim T →∞ w ( T ) = P S ∗ ( w ) . Proof.
By the gradient descent update we have,
[Equation]
Therefore,
[Equation]
In the limit T → ∞ and with η l ≤ λ max
[Equation]
we have, [Equation] Thus,
[Equation]
C.4. I MPROVING L OWER B OUND IN (8) IN THE C ASE OF E XACT P ROJECTIONS
Let S ∗
i be convex and let w ∗ ∈ S
i for all i ∈
[ M ] .
We assume that w ( t,τ ) i
[Equation]
P S ∗ i ( w ( t ) ) ∀ i ∈
[ M ] , i
[Equation]
e., the local models are an exact projection of w ( t ) on their respective solution sets.
From (8) we have,
[Equation]
[Equation]
We can lower bound (cid:68) w ( t ) − w ∗ , ∆ ( t ) i (cid:69) as follows,
[Equation]
where (103) uses the fact that (cid:68) w ( t,τ )
i
[Equation]
0 following the properties of projection (Boyd & Dattarro, 2003).
Thus we have, [Equation] [Table] D A DDITIONAL E XPERIMENTS AND S ETUP D ETAILS
Our code is available at the following link https://github.com/Divyansh03/FedExP .
D. I MPACT OF A VERAGING I TERATES FOR N EURAL N ETWORKS
As discussed in Section 5, 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 a Dirichlet distribution with α = 0 .
We set the number of local steps to be τ = 20 and train a CNN model having the same architecture as outlined in McMahan et al.
(2017) with full client participation.
Figure 6 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 an 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 a considerable improvement in the final accuracy.
Note however that this improvement only shows for FedExP ; averaging iterates does not make significant difference for FedAvg .
[Figure]
[Figure]
D ATASET D ETAILS
Here we provide more details about the datasets used in Section 6.
Synthetic Linear Regression.
In this case we assume that the local objective of each client is given by F i ( w ) =
∥
A i w − b i ∥ 2 where A i ∈ R (30 × 1000) , b i ∈ R 30 and w ∈ R 1000 .
We set the number of clients to be M = 20 .
Note that since d ≥ (cid:80) Mi =1 n i , this is an overparameterized convex problem.
To generate A i and b i , we follow a similar process as Li et al.
We have ( A i ) j : ∼ N ( m i , I d ) and ( b i ) j = w ⊤ i
( A i ) j : where m i ∼ N ( u i , 1) , w i ∼ N ( y i , 1) , u i ∼ N (0 , 0 . 1) ,
y i ∼ N (0 , 0 . 1) .
EMNIST.
EMNIST is an image classification task consisting of handwritten characters associated with 62 labels.
The federated EMNIST dataset available 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 a 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 a 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.
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 as a 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 a similar procedure as for CIFAR.
D. H YPERPARAMETER D ETAILS
For our baselines, we find the best performing η g and η l by grid-search tuning.
For FedExP we search for ϵ and η l .
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 η g is { 10 0 , 10 0 . 5 , 10 0 . 5 , 10 1 , 10 2 } .
For FedAdagrad , the grid for η g is { 10 − 1 , 10 − 0 .
5 , 10 − 0 , 10 0 .
For FedExP we keep ϵ = 0 in this experiment as (6) is satisfied in this case.
The grid for η l is { 10 − 2 , 10 − 1 .
5 , 10 − 1 , 10 − 0 .
5 , 10 0 } for all algorithms.
Grid for Neural Network Experiments.
For FedAvg and SCAFFOLD the grid for η g is { 10 − 1 , 10 − 0 .
For FedAdagrad , the grid for η g is { 10 − 2 , 10 − 1 .
5 , 10 − 1 , 10 − 0 .
For FedExP the grid for ϵ is { 10 − 3 , 10 − 2 .
5 , 10 − 2 , 10 − 1 .
5 , 10 − 1 } .
The grid for η l is { 10 − 2 , 10 − 1 .
5 , 10 − 1 , 10 − 0 .
5 , 10 0 } for all algorithms.
We use lower values of η g in the grid for FedAdagrad based on observations from Reddi et al.
(2021) which show that FedAdagrad performs better with smaller values of the server step size.
We provide details of the best performing hyperparameters below.
[Table]
SCAFFOLD
FedAdagrad ϵ
[Equation]
Synthetic *
[Equation]
CIFAR-
[Equation]
CIFAR-
[Equation] [Equation]
Other hyperparameters are kept the same for all algorithms.
In particular, we apply a weight decay of 0.0001 for all algorithms and decay η l by a factor of 0.998 in every round.
We also use gradient clipping to improve stability of the algorithms as done in previous works (Acar et al., 2021).
In all experiments we fix the number of participating clients to be 20, minibatch size to be 50 (for the synthetic dataset this reduces to full-batch gradient descent) and number of local updates τ to be 20.
D. S ENSITIVITY OF F ED E X P TO ϵ
To evaluate the sensitivity of FedExP to ϵ , we compute the training accuracy of FedExP afterrounds for varying ϵ and on different tasks.
For each task, we fix η l to be the value used in our experiments in Section 6 and only vary ϵ .
The results are summarized below.
[Table]
We see that the sensitivity of ϵ is similar to that of the τ parameter which is added to the denominator of FedAdam and FedAdagrad (Reddi et al., 2021) to prevent the step size from blowing up.
Keeping ϵ too large reduces the adaptivity of the method and makes the behavior similar to FedAvg .
At the same time, keeping ϵ too small may not also be beneficial always as seen in the case of EMNIST.
In practice, we find that a grid search for ϵ in the range { 10 − 3 , 10 − 2 .
5 , 10 − 2 , 10 − 1 .
5 , 10 − 1 } usually suffices to yield a good value of ϵ .
A general rule of thumb would be to start with ϵ = 10 − 3 and increase ϵ till the performance drops.
D. A DDITIONAL R ESULTS
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 η ( t ) g for FedExP in Figure 7.
We see that
FedExP takes much larger steps in some (but not all) rounds compared to the constant optimum step size taken by our baselines, leading to a large speedup.
Recall that we also let ϵ = 0 in this experiment (since it aligns with our theory) which also explains the larger values of η ( t ) g taken by
FedExP in this case.
[Figure]
EMNIST.
For EMNIST we observe that SCAFFOLD gives slightly better training loss than
FedExP towards the end of training.
As described in Section 6, 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 as shown in Figure 8.
[Figure]
CIFAR-10, CIFAR-100 and CINIC-10.
From Figure 3 and Figures 9–11, we see 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 in our experiments.
This points to the practical utility of FedExP even in challenging FL scenarios
[Figure]
[Figure]
[Figure]
Long-Term Behavior of Algorithms and Comparison with FedProx .
To evaluate the long-term behavior of different algorithms, we ran the experiments for 2000 rounds.
Here, we also consider an additional algorithm, namely FedProx , for comparison.
For fair comparison, we have tuned the µ parameter of FedProx for each dataset, by doing a grid search over the range { 10 − 3 , 10 − 2 , 10 − 1 , 1 } as done in the original FedProx paper (Li et al., 2020).
The results of EMNIST, CIFAR-10, CIFAR100, and CINIC-10 in Figures 12–14 and 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 a moving average window of size 20 .
[Figure]
[Figure] [Figure]
[Table]
E C OMBINING E XTRAPOLATION WITH SCAFFOLD
As described in Section 6, the extrapolation step can be added to the SCAFFOLD algorithm in a similar way as FedExP . The detailed steps of this SCAFFOLD-ExP algorithm are shown in Algorithm 2. Algorithm 2 SCAFFOLD-ExP 1: Input: w (0) , control variate c (0) , c (0) i , ∀ i ∈
[ M ] , number of rounds T , local iteration steps τ , parameters η l , ϵ 2:
For t = 0 , . . .
, T − 1 communication rounds do : 3:
Global server do: 4:
Send w ( t ) , c ( t ) to all clients 5:
Clients i ∈
[ M ] in parallel do: 6:
Set w ( t, 0) i ← w ( t, 0) 7:
For k = 0 , . . .
, τ − 1 local iterations do: 8:
Update w ( t,k +1)
i ← w ( t,k ) i − η l (cid:16) ∇ F i ( w ( t,k ) i , ξ ( t,k ) i ) − c ( t ) i + c ( t ) (cid:17) 9:
Compute ∆ ( t ) i ← w ( t ) − w ( t,τ ) i and Ψ ( t )
i ← c ( t ) − 1 τη l ∆ ( t ) i 10:
Send ∆ ( t ) i and Ψ ( t ) i to the server 11:
Update local control variate c ( t +1)
i ← c ( t ) i − Ψ
( t ) i 12:
Global server do: 13:
Compute
[Table]
2 + ϵ (cid:17)(cid:111) 14:
Update global model with w ( t +1) ← w ( t ) − η ( t )
g ¯∆ ( t ) 15:
Compute ¯Ψ ( t ) ← 1 M (cid:80)
Mi
=1 Ψ ( t ) i 16:
Update global control variate with c ( t +1) ← c ( t ) − ¯Ψ ( t )
F C OMBINING E XTRAPOLATION WITH S ERVER M OMENTUM
We begin by recalling some notation from our work.
The vector w ( t ) is the global model at round t and ¯∆ ( t ) is the average of client updates at round t .
The server momentum update at round t can be written as v ( t ) =
¯∆ ( t ) + β v ( t − 1) (let v − 1 = 0 ) and the global model update can be written as w ( t +1) =
w ( t ) − η ( t ) g v ( t ) .
Our goal is now to find η ( t ) g that minimizes (cid:13)(cid:13) w ( t +1) − w ∗
[Equation]
We have,
[Equation]
Setting the derivative of the RHS of (107) to zero we have,
[Equation]
Our goal now is to find a lower bound on ⟨ w ( t ) − w ∗ , v ( t ) ⟩ .
We have the following lemma.
Lemma 11.
Assume that (cid:10)
[Equation]
[Equation]
∆ ( t ) i (cid:13) (cid:13)(cid:13) 2 /M
(see Appendix C.4.1) for all t ≥ 0 and η ( r ) g ≤ ( m ( r ) + (cid:80)
r − 1 k =
( β/ 2) r − k m ( k ) ) /
2 (cid:13)(cid:13) v ( r ) (cid:13)(cid:13) 2 for all r < t − 1 .
Then,
[Equation]
t
[Equation]
∗ ,
[Equation]
t
[Equation]
which implies,
[Equation]
Proof.
We proceed via a proof by induction.
The statement clearly holds at t = 0 since (cid:10)
[Equation]
Now assuming the lemma holds at t − 1 we have,
[Equation]
where the last line follows from the fact that (cid:10) w ( t − 1) − w ∗ , v ( t − 1) (cid:11)
[Equation]
m ( t − 1) + (cid:80) t − 2 k =
( β/ 2) t − 1 − k m ( k ) and η ( t − 1) g
[Equation]
Thus we propose to keep the following server step size when using server momentum,
[Equation]
where m ( t ) =
(cid:80) M i =1 (cid:13) (cid:13) (cid:13) ∆ ( t ) i (cid:13) (cid:13) (cid:13) 2 /M .
Note that we also add a small constant ϵ to the denominator to prevent the step size from blowing up as done for FedExP .
We call server momentum with this step size as FedExP-M .
We compare the performance of FedExP-M with FedAdam and FedAvg-M ( FedAvg with server momentum) on the CIFAR-10 and CIFAR-100 datasets as shown in Figures 15–17, where the mean and standard deviation values are computed over 3 random seeds and a moving average window of size 20 .
The experimental setup is the same as described in Section 6.
The hyperparameters η l , ϵ for FedExP-M and η l , η g for FedAdam and FedAvg-M were tuned following a similar process as described in Appendix D.3, and their resulting values are in Table 6.
[Figure]
[Figure] [Figure]
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 FedAvg-M are quite similar in these experiments which can be attributed to the dense nature of the gradients in image classification as discussed in Section 6.
We note that this is only a preliminary result and future work will look to study the effect of combining server momentum and extrapolation more rigorously.