Unbounded Gradients in Federated Learning with Buffered Asynchronous Aggregation
Abstract
Synchronous updates may compromise the efficiency of cross-device federated learning once the number of active clients increases.
The FedBuff algorithm (Nguyen et al.
[16]) alleviates this problem by allowing asynchronous updates (staleness), which enhances the scalability of training while preserving privacy via secure aggregation.
We revisit the FedBuff algorithm for asynchronous federated learning and extend the existing analysis by removing the boundedness assumptions from the gradient norm.
This paper presents a theoretical analysis of the convergence rate of this algorithm when heterogeneity in data, batch size, and delay are considered.
Introduction
Federated learning (FL) is an approach in machine learning theory and practice that allows training models on distributed data sources [12, 14].
The distributed structure of FL has numerous benefits over traditional centralized methods, including parallel computing, efficient storage and improvements in data privacy.
However, this framework also presents communication efficiency, data heterogeneity, and scalability challenges.
Several works have been proposed to improve the performance of FL [1, 5, 7].
Existing works usually address a subset of these challenges while imposing additional constraints or limitations in other aspects.
For example, the work in [8] shows a trade-off between privacy, communication efficiency, and accuracy gains for the distributed discrete Gaussian mechanism for
FL with secure aggregation.
One of the most important advantages of FL is scalability.
Training models on centralized data stored on a single server can be problematic when dealing with large amounts of data.
Servers may be unable to handle the load or clients might refuse to share their data with a third party.
In FL, the data is distributed across many devices, potentially improving data privacy and computation scalability.
However, this also presents some challenges.
First, keeping the update mechanism synchronized across all devices may be very difficult when the number of clients is large [23].
Second, even if feasible, imposing synchronization results in huge (unnecessary) delays in the learning procedure [1].
Finally, each client often might have different data distributions, which can impact the convergence of algorithms [9, 21].
In synchronous FL, e.g., FedAvg [12, 14], the server first sends a copy of the current model to each client.
The clients then train the model locally on their private data and send the model updates back to the server.
The server then aggregates the client updates to produce a new shared model.
The process is repeated for many rounds until the shared model converges to the desired accuracy.
However, the existence of delays, message losses, and stragglers hinders the performance of distributed learning.
Several works have been proposed to improve the scalability of federated/distributed learning via enabling asynchronous communications [1, 6, 11, 13, 15, 18, 20, 23].
In the majority of these results, each client immediately communicates the parameters to the server after applying a series of local
[Table]
updates.
The server updates the global parameter once it receives any client update.
This has the benefit of reducing the training time and better scalability in practice and theory [1, 6, 15, 17] since the server can start aggregating the client updates as soon as they are available.
The setup, known as “vanilla” asynchronous FL, has several challenges that must be addressed.
First, due to the nature of asynchronous updates, the clients are supposed to deal with staleness, where the client updates are not up-to-date with the current model on the server [16].
Moreover, the asynchronous setup may imply potential risks for privacy due to the lack of secure aggregation, i [Table] e., the immediate communication of every single client to the server [2, 3].
In [16], the authors proposed an algorithm called federated learning with buffered asynchronous aggregation ( FedBuff ), which modifies pure asynchronous FL by enabling secure aggregation while clients perform asynchronous updates.
This novel method is considered a variant of asynchronous FL while serving as an intermediate approach between synchronous and asynchronous FL.
FedBuff [16] is shown to converge for the class of smooth and non-convex objective functions under the boundedness of the gradient norm.
By removing this assumption, we provide a new analysis for
FedBuff and improve the existing theory by extending it to a broader class of functions.
We derive our bounds based on stochastic and heterogeneous variance and the maximum delay between downloads and uploads across all the clients.
Table 1 summarizes the properties and rate of our analysis for
FedBuff algorithm alongside and provides a comparison with existing analyses for FedAsync [23] and FedAvg [12, 14].
The rates reflect the complexity of the number of updates performed by the central server.
The speed of asynchronous algorithms is faster since the constraint for synchronized updates is removed in asynchronous variations.
To our knowledge, this is the first analysis for (a variant of) asynchronous federated learning with no boundedness assumption on the gradient norm.
Following is an outline of the remainder of this paper.
The problem setup and FedBuff algorithm are presented in Section 2.
Moreover, our convergence result and its corresponding assumptions are provided in Section 2.
We state detailed proof of our result in section 3.
Finally, we conclude remarks and prospects for future research in Section 4.
Problem Setup & Algorithm
In this section, we first state the problem setup and then explain the FedBuff algorithm [16].
We consider a set of n clients and one server, where each client i ∈ [ n ] owns a private function f i : R d → R and the goal is to jointly minimize the average local cost functions via finding a d -dimensional parameter w ∈ R d that [Equation]
[Equation]
with f i ( w ) :=
E ξ i ∼ p
i
[ ℓ i ( w, ξ i )] , [Equation]
where ℓ i : R d × S i → R is a cost function that determines the prediction error of w over a single data point ξ i ∈ S i on user i , and p i represents user
i ’s data distribution over S i , for i ∈ [ n ] .
In the above definition, f i ( · ) is the local cost function of client i , and f ( · ) denotes the global (average) cost function which the clients try to collaboratively minimize.
Now, let D i be a data batch sampled from p i .
Similar to (1), we denote the stochastic cost function ˜ f
i ( w, D i ) as follows:
[Equation]
Minimization of (1) by having access to an oracle of samples and its variants are extensively studied for many different frameworks [7].
Now, we are ready to explain the FedBuff algorithm.
2.1 Federated Learning with Buffered Asynchronous Aggregation (FedBuff):
Let w 0 be the initialization parameter at the server.
The ultimate goal is to minimize the cost function in (1), using an algorithm via access to the stochastic gradients.
All clients can communicate with the server, and each client i ∈ [ n ] communicates when its connection to the server is stable.
First, let us explain the FedBuff algorithm from the client and server perspectives.
• Client Algorithm : Each client i requests to read the server’s parameter w ∈ R d 1 once the connection is stable and the server is ready to send the parameter.
There is often some delay in this step which we call the download delay.
This may be originated from factors such as unstable connection, bandwidth limit, or communication failure.
For example, maybe the server seeks to reduce the simultaneously active users by setting client i on hold.
The download delay can simply model all these factors.
Once the parameter is received (downloaded) from the server, client i performs Q steps of local stochastic gradient descent starting from the downloaded model w for its cost function f i ( · ) .
In words, agent i runs a Q -step algorithm (loop of size Q ), where at each local round q ∈
Q − 1 } , client i samples a data batch D i,q with respect to distribution p i and performs one step of gradient descent with local stepsize η > 0 .
Finally, agent i returns the updates (the difference between the initial and final parameters) to the server.
We refer to the time required to broadcast parameters to the server as the upload delay, which could have similar factors as the download delay.
Agent repeats all this procedure until the server sends a termination message.
Algorithm 1 summarizes the pseudo-code of operations at client i ∈ [ n ] , where Steps 4-8 show the local updates performed at the agent.
Moreover, ∆ i in Step 9 denotes the difference communicated to the server.
do 6:
8: end • Server Algorithm : The server considers
an initialization for parameter w 0 ∈ R d .
Then, starting from timestep t = the server repeats an iterative procedure in addition to sending its parameters to the clients upon their request.
Algorithm 2 describes the server operations in FedBuff.
In a nutshell, the algorithm consists of two parts, (i) secure aggregation of client updates in a buffer with size K ≥ 1 2 , and (ii) updating the parameters using the aggregated updates.
In other words, let k, t respectively denote the indices associated with buffer and server updates.
The server starting from t = 0 , receives updates broadcast by the agents asynchronously depending on their upload & download delays as well as the time required for Q local updates.
A secure buffered aggregates these updates, up to K separate updates received by the clients in ∆ 0 , initially set to zero.
By indexing k , we keep track of uploaded updates on the server.
When the buffer saturates of K different updates, server uses the aggregator parameter ∆ 0 and updates its parameter waccording to line 9 of Algorithm 2.
Then, the server increases its update counter t and removes all updates from the buffer, i.e., k = 0 .
In this algorithm, we denote the agent which sends the k -th update at round t by index i t,k ∈ [ n ] .
Basically, server repeats Steps 5-14 until some convergence criteria be satisfied.
After the convergence, the server sends a termination message to all the clients.
As we described above, the crucial novelty of this algorithm is on the server side, where the server operations, with the help of a secure buffered aggregation, control the staleness and prevent unnecessary access to individual updates.
Note that for K = 1 , the presented algorithm reduces to vanilla asynchronous federated learning with no buffer aggregation.
Figure 1 illustrates the update schedule for FedBuff and provides a comparison with the asynchronous updates in FedAvg [14].
As shown on the left of Figure 1, the vertical lines with light blue color are associated with uploaded updates.
Note that the buffer size is K = 2 in this example.
These vertical lines are of two types, (i) solid or (ii) hatched.
The solid lines reflect the time the buffer is full, and hence the server performs an update.
Contrary to FedBuff, under the synchronous updates (as shown in the right figure), the server should halt the training procedure until all clients selected within one round receive the updates.
Convergence Result
This section presents our main result alongside a few standard assumptions.
First, to be coherent with the proof in [16], let us denote τ ti to be the timestep of the last downloaded parameter on client i ∈ [ n ] up to the t -th update at the server.
Now, we are ready to introduce the assumptions in our analysis for FedBuff, i.e., Algorithms 1 & 2.
Assumption 1 (Bounded Staleness) .
For all clients, i ∈ [ n ] and server steps t ≥ 0 , the staleness or effective delay between the download and upload steps is bounded by some constant τ , i.e.,
[Equation]
) and the server receives updates uniformly, i.e., i t,k ∼ Uniform([ n ]) .
Note that τ ti is the timestep of the last parameter downloaded via agent i up to timestep t at the server.
Therefore, if agent i contributes in the ( t +1) -th update, i.e., i t,k
= i , for some k ∈ { 0 , 1 , . .
, K − 1 } , the difference between the download and upload rounds is bounded.
This is a standard assumption in the analysis of asynchronous algorithms with heterogeneous data on the clients.
Assumption 2 (Smoothness) .
For all clients i ∈ [ n ] , function f i : R d → R is bounded below, differentiable, and L -smooth, i.e., for all w, u ∈ R d ,
[Equation]
This assumption guarantees the necessary conditions for analyzing smooth & non-convex functions.
Note that boundedness from below can be relaxed only to the global cost function f , i.e., it is sufficient to only assume that f ⋆ := min w ∈ R d f ( w ) >
−∞ in our analysis instead of (5) for all i ∈ [ n ] .
Now, we introduce the assumptions on bounded stochasticity and heterogeneity.
Assumption 3 (Bounded Variance) .
For all clients i ∈ [ n ] , the variance of a stochastic gradient ∇ ℓ i ( w, ξ i ) on a single data point ξ i ∈ S i is bounded, i.e., for all w ∈ R d [Equation]
[Equation]
This assumption is conventional in the analysis of stochastic optimization algorithms and has been used in many relevant works [1, 9–11, 16, 19, 21].
Note that as we defined the stochastic loss in (2) and used the stochastic gradients in Step 7, we also need to show the stochastic variance for the gradients of the sampled batches.
For simplicity, let us assume that all batch sizes are of size at least b , therefore according to (6), we have:
[Equation]
Assumption 4 (Bounded Population Diversity) .
For all w ∈ R d , the gradients of local functions f i ( w ) and the global function f ( w ) satisfy the following property:
[Equation]
In our analysis, we work with heterogeneous cost functions.
Therefore, it is a reasonable and conventional assumption to assume that the boundedness of the population diversity [4, 5, 16, 20].
The inequality in 9 measures the variance of local full gradients from the average full gradient, which resembles to the expressions in (6) & (7).
Fallah et al.
[5] discusses the connection of this bound to the similarity of local data distributions p i , for all i ∈ [ n ] .
We now move to present our convergence result under the above assumptions.
Theorem 1.
Let Assumptions 1-4 hold, β = 1 K , and η = 1 Q √ LT .
Then, the following property holds for the joint iterates of Algorithms 1 and 2: for any timestep T ≥ 160 L ( Q +7)
( τ +1) 3 at the server
[Equation]
We present the proof for Theorem 1 in Appendix A.
The above theorem states the convergence of FedBuff algorithm to a first-order stationary point.
This result states a convergence rate of O (cid:16) 1 √ T (cid:17) + O (cid:16) τ 2 T (cid:17) , where the term affected by the maximum delay (second term) decays faster, hence the same convergence complexity as the synchronized counterpart.
Note that this rate states the number of updates occurring on the server (iteration complexity), which in the case of asynchronous updates, practically converges much faster ( 3 . [-
 ×according-]3 × according to [16]) than synchronized updates.
Remark 1.
The choice of β in Theorem 1 is an arbitrary option that implies the rate in the theorem statement.
The convergence proof may hold for any choices of β , such that βK
= O (1) .
Remark 2.
In our analysis for Theorem 1, we considered bounded population diversity in Assumption 4.
One can see that by relaxing this assumption to a stronger variant
[Equation]
i.e., uniformly bounded heterogeneity 5 , nγ 2 can be replaced with γ 2 in the third term of the rate.
Conclusion
This paper studied the convergence properties of asynchronous federated learning via secure buffered aggregation.
By removing the boundedness assumption on the gradient norms, we presented a novel analysis of the convergence of FedBuff algorithm, where we showed a sublinear convergence rate of O ( ϵ 2 )
+ O ( τ 2 ϵ ) to an ϵ -first-order stationary solution.
We also discussed the dependence of this rate on the batch size, stochasticity variance, data heterogeneity, and maximum delays.
We leave the privacy analysis of Fed-Buff with gradient clipping and noise addition to future studies.
Also, the communication complexity of this method and the extensions to decentralized setups remain for future work.
Acknowledgments and Disclosure of Funding
This work was partially funded by ARPA-H Strategic Initiative Seed Fund #916012.
Part of this material is based upon work supported by the National Science Foundation under Grants # and #2213568.
A Proof of Theorem
Proof of Theorem 1.
Before proceeding with the proof let us state some inequalities.
For any set of m vectors { w i } mi =1 such that w i ∈ R d and a constant α > the following properties hold: for all i, j ∈ [ m ] :
[Equation]
10c
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
For simplicity, let us denote ˜ ∇ f i ( w ) = ∇ ˜
f i ( w, D i ) .
Therefore, at round t , the server updates its parameter by receiving ∆ t , as follows:
[Equation]
[Equation]
Due to Assumption 2, we can infer that f is L -smooth, thus
[Equation]
[Equation]
First, we provide a lower bound on term S 1 in (14).
Prior to show the bound, let us denote ˜ g ti =
[Equation]
g t =
[Equation]
n n
[Equation]
i =1 g ti .
Therefore,
[Equation]
Moreover, the following holds for S 2 in (14):
[Equation]
Now, according to (14), (21)
[Equation]
and (22), we have:
[Equation]
where we bound S
[Equation]
S 4 as follows:
[Equation]
i = (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
Q − 1 (cid:88) q = ∇ f i (cid:16) w τ ti i,q (cid:17) − Q ∇ f i ( w t ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) (28) =
n n (cid:88) i =
(cid:13)(cid:13) (cid:13)(cid:13)(cid:13)
Q − 1 (cid:88) q
= (cid:104) ∇ f i (cid:16) w τ ti i,q (cid:17) − ∇ f i ( w t ) (cid:105) (cid:13)(cid:13) (cid:13)(cid:13)(cid:13) (29) (10d)
≤ Q n n (cid:88) i =
Q − 1 (cid:88) q
= (cid:13)(cid:13)(cid:13) ∇ f i (cid:16) w τ ti i,q (cid:17) − ∇ f i ( w t ) (cid:13)(cid:13)(cid:13) 2 , (30)
[Equation]
[Equation] [Equation]
[Equation]
[Equation] [Equation]
[Equation]
=
[Equation]
[Equation]
∇
[Equation]
[Equation]
Therefore, due to (25)-(30)
[Equation]
and (35)-(40)
[Equation]
we have
[Equation]
Hence, it is sufficient to bound
[Equation]
in
[Equation] [Equation]
[Equation]
Now, we show a bound on the evolution of local updates at an arbitrary round s ≥ [Equation]
i
[Equation]
e., the distance between w si,q and w s [Equation] which we will use to provide a bound on S
[Equation]
E (cid:13)(cid:13)(cid:13) w si,q − 1 − w s (cid:13)(cid:13)(cid:13)(54) + 4(1+2 Q ) η 2 E (cid:34)(cid:13)(cid:13)(cid:13)
˜ ∇ f i (cid:0) w si,q − 1 (cid:1) − ∇ f i (cid:0) w si,q − 1 (cid:1) (cid:13)(cid:13)(cid:13)+ (cid:13)(cid:13)(cid:13) ∇ f i (cid:0) w si,q − 1 (cid:1) − ∇ f i ( w s ) (cid:13)(cid:13)(cid:13)+ (cid:13)(cid:13) (cid:13) ∇ f i ( w s ) − ∇ f ( w s ) (cid:13)(cid:13) (cid:13)+ (cid:13)(cid:13)(cid:13) ∇ f ( w s ) (cid:13)(cid:13)(cid:13)(cid:35) (55) (4) , (7) ≤ (cid:18)
+2 Q (cid:19) E (cid:13)(cid:13)(cid:13) w si,q − 1 − w s (cid:13)(cid:13)(cid:13)(56)
+ 4(1+2 Q ) η(cid:34)
ˆ σ 2 + L 2 E (cid:13)(cid:13)(cid:13) w si,q − 1 − w s (cid:13)(cid:13)(cid:13)+ E (cid:13)(cid:13)(cid:13) ∇ f i ( w s ) − ∇ f ( w s ) (cid:13)(cid:13)(cid:13) 2 + E (cid:13)(cid:13)(cid:13) ∇ f ( w s ) (cid:13)(cid:13)(cid:13)
[Equation]
Note that we can select stepsize η ≤ 1 4 L ( Q +1) such that
[Equation]
therefore, due to (52)-(57) and (58), we have:
[Equation]
for all q ∈ [ Q ] .Again, note that according to Algorithm
[Equation]
we have:
[Equation]
w
[Equation]
68) =
[Equation]
n [Equation] τ
[Equation] [Equation] [Equation] [Equation]
(cid:13) (cid:13)
[Equation] [Equation]
n
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
(cid:16)
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
Let ϕ = 8 η 2 Q 2 (1+2 Q )(1+ β 2 K 2 ) , then according to (47)-(72)
[Equation]
we have
[Equation]
Note that according to (3), we know that: τ ti ∈ { t − τ . . .
, t } , therefore:
[Equation]
and similarly, for any s ∈ { t − τ . . . , t } and j ∈ [ n ] ,
[Equation]
Moreover, we have:
[Equation]
Therefore, due to (73)-(79), we have: nϕ
[Equation]
[Equation]
By combining (41)-(46) and (80)-(82)
[Equation]
we have the following inequality:
[Equation]
{+[Equation] ∇ f ( w s ) (cid:13)(cid:13)(cid:13) 
 + τ s (cid:88)+} u
= s − τ
E (cid:13)(cid:13)(cid:13) ∇ f ( w u ) (cid:13)(cid:13)(cid:13)(cid:35) (85) + ηβKQL 2 (1+ τ 2 ) (1+4 ηβKL ) ϕ
(cid:2)
ˆ σ 2 + nγ 2 (cid:3) (86)
≤ E f ( w t )
+ 2 η 2 Lβ 2 K 2 Q 2 (cid:2)
ˆ σ 2 + γ 2 (cid:3) (87) − ηβKQ (cid:2) 1 − 4 ηLβKQ − QL 2 (1+ τ 2 ) (1+4 ηβKL ) ϕ
(cid:3) E (cid:13)(cid:13) ∇ f ( w t ) (cid:13)(cid:13)(88) + ηβKQL 2 (1+4 ηβKL )
ϕ t − 1 (cid:88) s
= t − τ (cid:34)
E (cid:13)(cid:13)(cid:13) ∇ f ( w s ) (cid:13)(cid:13)(cid:13)
+ τ s (cid:88) u
= s − τ
E (cid:13)(cid:13)(cid:13) ∇ f ( w u ) (cid:13)(cid:13)(cid:13)(cid:35) (89) + ηβKQL 2 (1+ τ 2 ) (1+4 ηβKL ) ϕ
(cid:2)
ˆ σ 2 + nγ 2 (cid:3) .
Now, we can obtain the following inequality by rearranging the terms in (83)-(90):
[Equation]
whereby mixing the terms in (92)
[Equation]
we obtain:
[Equation]
Finally
[Equation]
we add (96)-(100), for t = 0 , [Equation]
T −
[Equation]
and divide by T to show that:
[Equation]
Let us fix β = 1 K and η = 1 Q √ LT .
Thus [Equation] we know that the following inequality holds
[Equation]
for T ≥ 160 L ( Q +7)( τ +1) 3 .
Note that under this choices for η and β , we also have η ≤ 1 4 L ( Q +1) , which we used in (58).
Therefore, we can conclude the final result in Theorem 1 as follows:
[Equation]