EPISODE: E PISODIC G RADIENT C LIPPING WITH P ERIODIC R ESAMPLED C ORRECTIONS FOR F EDERATED L EARNING WITH H ETEROGENEOUS D ATA
Anonymous authors
Paper under double-blind review
A BSTRACT
Gradient clipping is an important technique for deep neural networks with exploding gradients, such as recurrent neural networks.
Recent studies have shown that the loss functions of these networks do not satisfy the conventional smoothness condition, but instead satisfy a relaxed smoothness condition, i.e., the Lipschitz constant of the gradient scales linearly in terms of the gradient norm.
Due to this observation, several gradient clipping algorithms have been developed for nonconvex and relaxed-smooth functions.
However, the existing algorithms only apply to the single-machine or multiple-machine setting with homogeneous data across machines.
It remains unclear how to design provably efficient gradient clipping algorithms in the general Federated Learning (FL) setting with heterogeneous data and limited communication rounds.
In this paper, we design EPISODE, the very first algorithm to solve FL problems with heterogeneous data in the nonconvex and relaxed smoothness setting.
The key ingredients of the algorithm are two new techniques called episodic gradient clipping and periodic resampled corrections .
At the beginning of each round, EPISODE resamples stochastic gradients from each client and obtains the global averaged gradient, which is used to (1) determine whether to apply gradient clipping for the entire round and (2) construct local gradient corrections for each client.
Notably, our algorithm and analysis provide a unified framework for both homogeneous and heterogeneous data under any noise level of the stochastic gradient, and it achieves state-of-the-art complexity results.
In particular, we prove that EPISODE can achieve linear speedup in the number of machines, and it requires significantly fewer communication rounds.
Experiments on several heterogeneous datasets, including text classification and image classification, show the superior performance of EPISODE over several strong baselines in FL.
I NTRODUCTION
Gradient clipping (Pascanu et al., 2012; 2013) is a well-known strategy to improve the training of deep neural networks with the exploding gradient issue such as Recurrent Neural Networks (RNN)
(Rumelhart et al., 1986; Elman, 1990; Werbos, 1988) and Long Short-Term Memory (LSTM) (Hochreiter &
Schmidhuber, 1997).
Although it is a widely-used strategy, formally analyzing gradient clipping in deep neural networks under the framework of nonconvex optimization only happened recently (Zhang et al., 2019a; 2020a; Cutkosky & Mehta, 2021;
Liu et al., 2022).
In particular, Zhang et al. (2019a) showed empirically that the gradient Lipschitz constant scales linearly in terms of the gradient norm when training certain neural networks such as AWD-LSTM (Merity et al., 2018), introduced the relaxed smoothness condition (i.e., ( L 0 , L 1 ) -smoothness 1 ), and proved that clipped gradient descent converges faster than any fixed step size gradient descent.
Later on, Zhang et al.
(2020a) provided tighter complexity bounds of the gradient clipping algorithm.
Federated Learning (FL) (McMahan et al., 2017a) is an important distributed learning paradigm in which a single model is trained collaboratively under the coordination of a central server without revealing client data 2 .
FL has two critical features: heterogeneous data and limited communication.
Although there is a vast literature on FL (see (Kairouz et al., 2019) and references therein), the theoretical and algorithmic understanding of gradient clipping algorithms for training deep neural networks in the FL setting remains nascent.
To the best of our knowledge, Liu et al.
(2022) is the
[Table]
+
[Table]
only work that has considered a communication-efficient distributed gradient clipping algorithm under the nonconvex and relaxed smoothness conditions in the FL setting.
In particular, Liu et al.
(2022) proved that their algorithm achieves linear speedup in terms of the number of clients and reduced communication rounds.
Nevertheless, their algorithm and analysis are only applicable to the case of homogeneous data.
In addition, the analyses of the stochastic gradient clipping algorithms in both single machine (Zhang et al., 2020a) and multiple-machine setting (Liu et al., 2022) require strong distributional assumptions on the stochastic gradient noise 3 , which may not hold in practice.
In this work, we introduce a provably computation and communication efficient gradient clipping algorithm for nonconvex and relaxed-smooth functions in the general FL setting (i.e., heterogeneous data, limited communication) and without any distributional assumptions on the stochastic gradient noise .
Compared with previous work on gradient clipping (Zhang et al., 2019a; 2020a;
Cutkosky & Mehta, 2020; Liu et al., 2022) and FL with heterogeneous data (Li et al., 2020a;
Karimireddy et al., 2020), our algorithm design relies on two novel techniques: episodic gradient clipping and periodic resampled corrections .
In a nutshell, at the beginning of each communication round, the algorithm resamples each client’s stochastic gradient; this information is used to decide whether to apply clipping in the current round (i.e., episodic gradient clipping ), and to perform local corrections to each client’s update (i.e., periodic resampled corrections ).
These techniques are very different compared with previous work on gradient clipping.
Specifically, (1) In traditional gradient clipping (Pascanu et al., 2012;
Zhang et al., 2019a; 2020a; Liu et al., 2022), whether or not to apply the clipping operation is determined only by the norm of the client’s current stochastic gradient.
Instead, we use the norm of the global objective’s stochastic gradient (resampled at the beginning of the round) to determine whether or not clipping will be applied throughout the entire communication round.
(2) Different from Karimireddy et al.
(2020) which uses historical gradient information from the previous round to perform corrections, our algorithm utilizes the resampled gradient to correct each client’s local update towards the global gradient, which mitigates the effect of data heterogeneity.
Notice that, under the relaxed smoothness setting, the gradient may change quickly around a point at which the gradient norm is large.
Therefore, our algorithm treats a small gradient as more “reliable” and confidently applies the unclipped corrected local updates; on the contrary, the algorithm regards a large gradient as less “reliable” and in this case clips the corrected local updates.
Our major contributions are summarized as follows.
• We introduce EPISODE, the very first algorithm for optimizing nonconvex
and ( L 0 , L 1 ) -smooth functions in the general FL setting with heterogeneous data and limited communication.
The algorithm design introduces novel techniques, including episodic gradient clipping and periodic
resampled corrections.
To the best of our knowledge, these techniques are first introduced by us and crucial for algorithm design.
• Under the nonconvex and relaxed smoothness condition, we prove that the EPISODE algorithm can achieve linear speedup in the number of clients and reduced communication rounds in the heterogeneous data setting, without any distributional assumptions on the stochastic gradient noise.
In addition, we show that the degenerate case of EPISODE matches state-of-the-art complexity results under weaker assumptions 4 .
Detailed complexity results and a comparison with other relevant algorithms are shown in Table 1.
• We conduct experiments on several heterogeneous medium and large scale datasets with different deep neural network architectures, including a synthetic objective, text classification, and image classification.
We show that the performance of the EPISODE algorithm is consistent with our theory, and it consistently outperforms several strong baselines in FL.
R ELATED W ORK
Gradient Clipping Gradient clipping is a standard technique in the optimization literature for solving convex/quasiconvex problems (Ermoliev, 1988; Nesterov, 1984; Shor, 2012; Hazan et al., 2015; Mai & Johansson, 2021; Gorbunov et al., 2020), nonconvex smooth problems (Levy, 2016;
Cutkosky & Mehta, 2021), and nonconvex distributionally robust optimization (Jin et al., 2021).
Menon et al.
(2019) showed that gradient clipping can help mitigate label noise.
Gradient clipping is a well-known strategy to achieve differential privacy (Abadi et al., 2016; McMahan et al., 2017b;
Andrew et al., 2021; Zhang et al., 2021).
In the deep learning literature, gradient clipping is employed to avoid exploding gradient issue when training certain deep neural networks such as recurrent neural networks or long-short term memory networks (Pascanu et al., 2012; 2013) and language models (Gehring et al., 2017; Peters et al., 2018; Merity et al., 2018).
Zhang et al. (2019a) initiated the study of gradient clipping for nonconvex and relaxed smooth functions.
Zhang et al.
(2020a) provided an improved analysis over Zhang et al. (2019a).
However, none of these works apply to the general FL setting with nonconvex and relaxed smooth functions.
Federated Learning FL was proposed by McMahan et al.
(2017a),
to enable large-scale distributed learning while keep client data decentralized to protect user privacy.
McMahan et al.
(2017a) designed the FedAvg algorithm, which allows multiple steps of gradient updates before communication.
This algorithm is also known as local SGD (Stich, 2018; Lin et al., 2018; Wang & Joshi, 2018; Yu et al., 2019c).
The local SGD algorithm and their variants have been analyzed in the convex setting (Stich, 2018;
Stich et al., 2018; Dieuleveut & Patel, 2019; Khaled et al., 2020; Li et al., 2020a; Karimireddy et al., 2020; Woodworth et al., 2020a;b; Koloskova et al., 2020; Yuan et al., 2021) and nonconvex smooth setting (Jiang & Agrawal, 2018; Wang & Joshi, 2018;
Lin et al., 2018;
Basu et al., 2019;
Haddadpour et al., 2019; Yu et al., 2019c;b; Li et al., 2020a; Karimireddy et al., 2020; Reddi et al., 2021; Zhang et al., 2020b; Koloskova et al., 2020).
Recently, in the stochastic convex optimization setting, several works compared local SGD and minibatch SGD in the homogeneous (Woodworth et al., 2020b) and heterogeneous data setting (Woodworth et al., 2020a), as well as the fundamental limit (Woodworth et al., 2021).
For a comprehensive survey, we refer the readers to Kairouz et al.
; Li et al.
(2020a)
and references therein.
The most relevant work to ours is Liu et al.
(2022), which introduced a communication-efficient distributed gradient clipping algorithm for nonconvex and relaxed smooth functions.
However, their algorithm and analysis does not apply in the case of heterogeneous data as considered in this paper.
P ROBLEM S ETUP AND P RELIMINARIES
Notations In this paper, we use ⟨· , ·⟩ and ∥ · ∥ to denote the inner product and Euclidean norm in space R d .
We use 1 ( · ) to denote the indicator function.
We let I r be the set of iterations at the r -th round, that is I r
= { t r , ..., t r +1 − 1 } .
The filtration generated by the random variables before step t r is denoted by F r .
We also use E r [ · ] to denote the conditional expectation E [ ·|F r ] .
The number of clients is denoted by N and the length of the communication interval is denoted by I , i.e., |I r
|
=
I for r = 0 , 1 , ..., R .
Let f i ( x ) :=
E ξ i ∼D i [ F i ( x ; ξ i )] be the loss function in i -th client for i ∈
[ N ] , where the local distribution D i is unknown and may be different across i ∈
[ N ] .
In the FL setting, we aim to minimize the following overall averaged loss function:
[Equation]
We focus on the case that each f i is non-convex, in which it is NP-hard to find the global minimum of f .
Instead we consider finding an ϵ -stationary point (Ghadimi & Lan, 2013; Zhang et al., 2020a).
Definition 3.1.
For a function h : R d → R , a point x ∈ R d is called ϵ -stationary if ∥∇ h ( x ) ∥ ≤ ϵ .
Most existing works in the non-convex FL literature (Yu et al., 2019a; Karimireddy et al., 2020) assume each f i is L -smooth, i.e., ∥∇ f i ( x ) −
∇ f i ( y ) ∥
≤ L ∥ x − y ∥ for any x , y ∈ R d .
However it is shown in Zhang et al.
(2019a)
that L -smoothness may not hold for certain neural networks such as
RNNs and LSTMs.
( L 0 , L 1 ) -smoothness in Definition 3.2 was proposed by Zhang et al.
(2019b)
and is strictly weaker than L -smoothness.
Under this condition, the local smoothness of the objective can grow with the gradient scale.
For AWD-LSTM (Merity et al., 2018), empirical evidence of ( L 0 , L 1 ) -smoothness was observed in Zhang et al.
(2019b).
Definition 3.2.
A second order differentiable function h : R d → R is ( L 0 , L 1 ) -smooth if ∥∇ 2 h ( x ) ∥ ≤ L 0 + L 1 ∥∇ h ( x ) ∥ holds for any x ∈ R d .
Suppose we only have access to the stochastic gradient
∇ F i ( x ; ξ ) for ξ ∼ D i in each client.
Next we make the following assumptions on objectives and stochastic gradients.
Assumption 1.
Assume f
i for i ∈
[ N ] and f defined in (1) satisfy:
(i) f
i is second order differentiable and ( L 0 , L 1 ) -smooth.
(ii) Let x
∗ be the global minimum of f and x 0 be the initial point.
There exists some ∆ > 0 such that f ( x 0 ) − f ( x ∗ ) ≤ ∆ .
(iii)
For all x ∈ R d , there exists some σ ≥ 0 such that E ξ i ∼D i [ ∇ F i ( x ; ξ i )]
= ∇ f i ( x ) and ∥∇ F i ( x ; ξ i ) − ∇ f i ( x ) ∥ ≤ σ almost surely.
(iv) There exists some κ ≥ 0 and ρ ≥ 1 such that ∥∇ f i ( x ) ∥ ≤ κ + ρ ∥∇ f ( x ) ∥ for any x ∈ R d .
Remark : (i) and (ii) are standard in the non-convex optimization literature (Ghadimi & Lan, 2013), and (iii) is a standard assumption in the ( L 0 , L 1 ) -smoothness setting (Zhang et al., 2019b; 2020a; Liu et al., 2022).
(iv) is used to bound the difference between the gradient of each client’s local loss and the gradient of the overall loss, which is commonly assumed in the FL literature with heterogeneous data (Karimireddy et al., 2020).
When κ = 0 and ρ = 1 , (iv) corresponds to the homogeneous setting.
A LGORITHM AND A NALYSIS 4. M AIN C HALLENGES AND A LGORITHM D ESIGN
Main Challenges We first illustrate why the prior local gradient clipping algorithm (Liu et al., 2022) would not work in the heterogeneous data setting.
Liu et al.
(2022) proposed the first communicationefficient local gradient clipping algorithm (CELGC) in a homogeneous setting for relaxed smooth functions, which can be interpreted as the clipping version of FedAvg.
Let us consider a simple heterogeneous example with two clients in which CELGC fails.
Denote f 1 ( x ) = 12 x 2 + a 1 x and f 2 ( x ) = 12 x 2 + a 2 x with a 1 = − γ − 1 , a 2 = γ + 2 , and γ > 1 .
We know that the optimal solution for f = f 1 + f 2is x ∗ = − a 1 + a 2= − 12 , and both f 1 and f 2 are ( L 0 , L 1 ) -smooth with L 0 = 1 and
L 1 = 0 .
Consider CELGC with communication interval I = 1 (i.e., communication happens at every iteration), starting point x 0 = 0 , η = 1 /L
= 1 , clipping threshold γ , and σ = 0 .
In this setting, after the first iteration, the model parameters on the two clients become γ and − γ respectively, so that the averaged model parameter after communication returns to 0 .
This means that the model parameter of
CELGC remains stuck at 0 indefinitely, demonstrating that CELGC cannot handle data heterogeneity.
We then explain why the stochastic controlled averaging method (SCAFFOLD) (Karimireddy et al., 2020) for heterogeneous data does not work in the relaxed smoothness setting.
SCAFFOLD utilizes the client gradients from the previous round to constructing correction terms which are added to each client’s local update.
Crucially, SCAFFOLD requires that the gradient is Lipschitz so that gradients from the previous round are good approximations of gradients in the current round with controllable errors.
This technique is not applicable in the relaxed smoothness setting: the gradient may change dramatically, so historical gradients from the previous round are not good approximations of the current gradients anymore due to potential unbounded errors.
Algorithm Design To address the challenges brought by heterogeneity and relaxed smoothness, our idea is to clip the local updates similarly as we would clip the global gradient (if we could access it).
The detailed description of EPISODE is stated in Algorithm 1.
Specifically, we introduce two novel techniques: (1) Episodic gradient clipping.
At the r -th round, EPISODE constructs a global indicator( ∥ G r ∥ ≤ γ/η ) to determine whether to perform clipping in every local update during the round for all clients (line 6).
(2) Periodic resampled corrections.
EPISODE resamples fresh gradients with constant batch size at the beginning of each round (line 3-5).
In particular, at the beginning of the r -th round, EPISODE samples stochastic gradients evaluated at the current averaged global weight ¯ x r in all clients to construct the control variate G r , which has two roles.
The first is to construct the global clipping indicator according to ∥ G r ∥ (line 10).
The second one is to correct the bias between local gradient and global gradient through the variate g it in local updates (line 10).
4. M AIN R ESULTS
Theorem 4.1.
Suppose Assumption 1 holds.
For any tolerance ϵ ≤ 3 AL 0 5 BL 1 ρ , we choose the hyper parameters as η ≤ min (cid:26) 216Γ I , ϵ 180Γ Iσ ,
Nϵ16 AL 0 σ(cid:27) and γ = (cid:18)
σ + AL
BL 1 ρ (cid:19) η, (2) where
Γ =
AL
+ BL 1 κ + BL 1 ρ (cid:16) σ + γη (cid:17) .
Then EPISODE satisfies 1 R +1 (cid:80)
Rr =
E [ ∥∇ f
(¯ x r ) ∥ ] ≤ 3 ϵ as long as the number of communication rounds satisfies R ≥ 4∆ ϵ 2 ηI .
Remark 1: The result in Theorem 4.1 holds for arbitrary noise level, while the complexity bounds in the stochastic case of Zhang et al.
(2020a)
; Liu et al.
both require σ ≥ 1 .
In addition, this theorem can automatically recover the complexity results in Liu et al.
, but does not require their symmetric and unimodal noise assumption.
The improvement upon previous work comes from a better algorithm design, as well as a more careful analysis in the smoothness and individual discrepancy in the non-clipped case (see Lemma 4.2 and 4.3).
Remark 2:
In Theorem 4.1, when we choose η = min (cid:110) 1 216Γ I ,
[Equation]
, the total communication complexity to find an ϵ -stationary point is no more than
[Equation]
= O
[Equation]
Next we present some implications of the complexity (3): 1.
When I ≲
L 0 σ ( L 0 + L 1 ( κ + σ ))
Nϵ and σ ≳ ϵ , EPISODE has communication complexity O ( ∆ L 0 σ 2 NIϵ 4 ) .
In this case, EPISODE enjoys a better communication complexity than the naive parallel version of the algorithm in Zhang et al.
(2020a)
, that is O
( ∆ L 0 σ 2 Nϵ 4 ) .
Moreover, the iteration complexity of EPISODE is T = RI = O ( ∆ L 0 σ 2 Nϵ 4 ) , which achieves the linear speedup w.r.t.
the number of clients N .
This matches the result of Liu et al.
(2022) in the homogeneous data setting.
When I ≳
L 0 σ ( L 0 + L 1 ( κ + σ ))
Nϵ and σ ≳ ϵ , the communication complexity of EPISODE is
O ( ∆( L 0 + L 1 ( κ + σ )) σ ϵ) .
This term does not appear in Theorem III of Karimireddy et al.
but it appears here due to the difference in the construction of the control variates.
In fact, the communication complexity of EPISODE is still lower than the naive parallel version of Zhang et al.
(2020a)
if the number of clients satisfies N ≲
[Equation]
)) ϵ .
When 0 < σ ≲ ϵ , EPISODE has communication complexity O ( ∆( L 0 + L 1 ( κ + σ )) ϵ) .
Under this particular noise level, the algorithms in Zhang et al.
(2020a)
; Liu et al.
(2022) do not guarantee convergence because their analyses crucially rely on the fact that σ ≳ ϵ .
When σ = 0 , EPISODE has communication complexity O ( ∆( L 0 + L 1 κ ) ϵ) .
This bound includes an additional constant L 1 κ compared with the complexity results in the deterministic case (Zhang et al., 2020a), which comes from data heterogeneity and infrequent communication.
4. P ROOF S KETCH OF T HEOREM 4.
Despite recent work on gradient clipping in the homogeneous setting (Liu et al., 2022), the analysis of
Theorem 4.1 is highly nontrivial since we need to cope with ( L 0 , L 1 ) -smoothness and heterogeneity simultaneously.
In addition, we do not require a lower bound of σ and allow for arbitrary σ ≥ 0 .
The first step is to establish the descent inequality of the global loss function.
According to the ( L 0 , L 1 ) -smoothness condition, if ∥ ¯ x r +1 − ¯ x r ∥ ≤ C/L 1 , then
[Equation]
where A r := {∥ G r ∥ ≤ γ/η } , ¯ A r is the complement of A r , and A, B, C are constants defined in
Lemma A.1.
To utilize the inequality (4), we need to verify that the distance between ¯ x r +
and ¯ x r is small almost surely.
In the algorithm of Liu et al.
(2022), clipping is performed in each iteration based on the magnitude of the current stochastic gradient, and hence the increment of each local weight is bounded by the clipping threshold γ .
For each client in EPISODE, whether to perform clipping is decided by the magnitude of G r at the beginning of each round.
Therefore, the techniques in Liu et al.
(2022) to bound the individual discrepancy cannot be applied to EPISODE.
To address this issue, we introduce Lemma 4.1, which guarantees that we can apply the properties of relaxed smoothness (Lemma A.1 and A.2) to all iterations in one round, in either case of clipping or non-clipping.
Lemma 4.1.
Suppose 2 ηI ( AL 0 + BL 1 κ + BL 1 ρ ( σ + γ/η )) ≤ 1 and max { 2 ηI (2 σ + γ/η ) , γI } ≤
CL 1 .
Then for any i ∈ [ N ] and t − 1 ∈
I r , it almost surely holds that
[Equation]
Equipped with Lemma 4.1, the condition ∥ ¯ x r +1 − ¯ x r ∥ ≤ 1 N (cid:80) Ni
=1 ∥ x it r +1 − ¯ x r ∥ ≤ C/L 1 can hold almost surely with a proper choice of η .
Then it suffices to bound the terms from (4) in expectation under the events A r and ¯ A r respectively.
To deal with the discrepancy term E [ ∥ x it − ¯ x r ∥ 2 ] for t − 1 ∈ I r , Liu et al.
(2022) directly uses the almost sure bound for both cases of clipping and non-clipping.
Here we aim to obtain a more delicate bound in expectation for the non-clipping case.
The following lemma, which is critical to obtain the unified bound from Theorem 4.1 under any noise level, gives an upper bound for the local smoothness of f i at x .
Lemma 4.2.
Under the conditions of Lemma 4.1, for all x ∈ R d such that ∥ x − ¯ x r ∥ ≤ 2 ηI (2 σ + γ/η ) , the following inequality almost surely holds:
[Equation]
From (5), we can see that all iterations in the r -th round satisfy the condition of Lemma 4.2 almost surely.
Hence we are guaranteed that each local loss f i is L -smooth over the iterations in this round under the event A r , where L = L 0 + L 1 ( κ + ( ρ + 1)( γ/η + 2 σ )) .
In light of this, the following lemma gives a bound in expectation of the individual discrepancy.
We denote p r =
E r
[ 1 ( A r )] .
Lemma 4.3.
Under the conditions of Lemma 4.1, for any i ∈ [ N ] and t − 1 ∈
I r , we have
[Equation]
It is worthwhile noting that the bound in (6) involves a quadratic term of ∥∇ f
(¯ x r ) ∥ , whereas it is linear in (7).
The role of the linear bound is to deal with 1 ( A r ) ∥∇ f (¯ x r )
∥∥ ¯ x r +1 − ¯ x r ∥ 2 from the descent inequality (4), since directly substituting (6) would result in a cubic term which is hard to analyze.
With Lemma 4.1, 4.2 and 4.3, we obtain the following descent inequality.
Lemma 4.4.
Under the conditions of Lemma 4.1, let Γ =
AL
+ BL 1 ( κ + ρ ( γ/η + σ )) .
Then it holds that for each 0 ≤ r ≤ R − 1 ,
[Equation]
where the definitions of V (¯ x r ) and U (¯ x r ) are given in Appendix C.1.
The detailed proof of Lemma 4.4 is deferred in Appendix C.1.
With this Lemma, the descent inequality is divided into V (¯ x r ) (objective value decrease during the non-clipping rounds) and U (¯ x r ) (objective value decrease during the clipping rounds).
Plugging in the choices of η and γ yields
[Equation]
The conclusion of Theorem 4.1 can then be obtained by substituting (9) into (8) and summing over r .
E XPERIMENTS
In this section, we present an empirical evaluation of EPISODE to validate our theory.
We present results in the heterogeneous FL setting on three diverse tasks: a synthetic optimization problem satisfying ( L 0 , L 1 ) -smoothness, natural language inference on the SNLI dataset (Bowman et al., 2015), and ImageNet classification (Deng et al., 2009).
We compare EPISODE against FedAvg (McMahan et al., 2017a), SCAFFOLD (Karimireddy et al., 2020), CELGC (Liu et al., 2022), and a naive distributed algorithm which we refer to as Naive Parallel Clip 5 We include additional experiments on the CIFAR-10 dataset (Krizhevsky et al., 2009) in Appendix E.4.
We have also added running time results in Appendix F, ablation study in Appendix G and new experiments on federated learning benchmark datasets in Appendix H.
5. S ETUP
All non-synthetic experiments were implemented with PyTorch (Paszke et al., 2019) and run on a cluster with eight NVIDIA Tesla V100 GPUs.
Since SNLI , CIFAR-10, and ImageNet are centralized datasets, we follow the non-i.i.d.
partitioning protocol in (Karimireddy et al., 2020) to split each dataset into heterogeneous client datasets with varying label distributions.
Specifically, for a similarity parameter s ∈ [0 , 100] , each client’s local dataset is composed of two parts.
The first s % is comprised of i.i.d.
samples from the complete dataset, and the remaining (100 − s )% of data is sorted by label.
Synthetic To demonstrate the behavior of EPISODE and baselines under ( L 0 , L 1 ) -smoothness, we consider a simple minimization problem in a single variable.
Here we have N = 2 clients with:
[Equation]
where the parameter H controls the heterogeneity between the two clients.
Notice that f 1 and fsatisfy ( L 0 , L 1 ) -smoothness but not traditional L -smoothness.
Proposition 5.1.
For any x ∈ R and i = 1 , 2 , it holds that ∥∇ f i ( x ) ∥ ≤ 2 ∥∇ f ( x ) ∥ + κ ( H ) , where κ ( H ) < ∞ and is a positive increasing function of H for H ≥ 1 .
According to Proposition 5.1, Assumption 1(iv) will be satisfied with ρ = 2 and κ = κ ( H ) , where κ ( H ) is an increasing function of H .
The proof of this proposition is deferred to Appendix E.1.
[Figure]
SNLI Following Conneau et al.
(2017), we train a BiRNN network for 25 epochs using the multiclass hinge loss and a batch size of 64 on each worker.
The network is composed of a one layer
BiRNN encoder with hidden size 2048 and max pooling, and a three-layer fully connected classifier with hidden size 512.
The BiRNN encodes a sentence (represented as a sequence of GloVe vectors (Pennington et al., 2014)), and the classifier predicts the relationship of two encoded sentences as either entailment, neutral, or contradiction.
For more hyperparameter information, see Appendix E.3.
To determine the effects of infrequent communication and data heterogeneity on the performance of each algorithm, we vary I ∈ { 2 , 4 , 8 , 16 } and s ∈ { 10% , 30% , 50% } .
We compare EPISODE,
CELGC, and the Naive Parallel Clip.
Note that the training process diverged when using SCAFFOLD, likely due to a gradient explosion issue, since SCAFFOLD does not use gradient clipping.
ImageNet Following Goyal et al.
(2017), we train a ResNet-
(He et al., 2016) for 90 epochs using the cross-entropy loss, a batch size of 32 for each worker, clipping parameter γ = 1 . 0 , momentum with coefficient 0 .
9 , and weight decay with coefficient 5 ×
10 − 5 .
We initially set the learning rate η = 0 .
1 and decay by a factor of 0.1 at epochs 30 , 60 , and 80 .
To analyze the effect of data heterogeneity in this setting, we fix I = 64 and vary s ∈ { 50% , 60% , 70% } .
Similarly, to analyze the effect of infrequent communication, we fix s = 60% and vary I ∈ { 64 , 128 } .
We compare the performance of FedAvg, CELGC, EPISODE, and SCAFFOLD.
5. R ESULTS
Synthetic Figure 3 in Appendix E.2 shows the objective value throughout training, where the heterogeneity parameter H varies over { 1 , 2 , 4 , 8 } .
CELGC exhibits very slow optimization due to the heterogeneity across clients: as H increases, the optimization progress becomes slower and slower.
In contrast, EPISODE maintains fast convergence as H varies.
We can also see that EPISODE converges to the minimum of global loss, while CELGC fails to do so when H is larger.
SNLI Results for the SNLI dataset are shown in Figure 1.
To demonstrate the effect of infrequent communication, Figures 1(a) and 1(b) show results for EPISODE, CELGC, and Naive Parallel Clip as the communication interval I varies (with fixed s = 30% ).
After 25 epochs, the test accuracy
Interval
Similarity
Algorithm
Train loss
Test acc.
FedAvg 1.
CELGC 1.
SCAFFOLD 1.
EPISODE 0.
FedAvg 0.
CELGC 0.
SCAFFOLD 0.
EPISODE 0.
FedAvg 0.
CELGC 0.
SCAFFOLD 0.
EPISODE 0.
FedAvg 1.
CELGC 1.
SCAFFOLD 1.
EPISODE 1.
[Figure]
of EPISODE nearly matches that of Naive Parallel Clip for all I ≤ 8 , while CELGC lags 2-3% behind Naive Parallel Clip for all values of I .
Also, EPISODE nearly matches the test accuracy of
Naive Parallel Clip with as little as 8 times fewer communication rounds.
Lastly, EPISODE requires significantly less communication rounds to reach the same training loss as CELGC.
For example,
EPISODE with I = 4 , s = 30% takes less than 5000 rounds to reach a training loss of 0 .
4 , while
CELGC does not reach 0 .
4 during the entirety of training with any I .
To demonstrate the effect of client data heterogeneity, Figure 1(c) shows results for varying values of s (with fixed I = 4 ).
Here we can see that EPISODE is resilient against data heterogeneity: even with client similarity as low as s = 10% , the performance of EPISODE is the same as s = 50% .
Also, the testing accuracy of EPISODE with s = 10% is nearly identical to that of the Naive Parallel
Clip.
On the other hand, the performance of CELGC drastically worsens with more heterogeneity: even with s = 50% , the training loss of CELGC is significantly worse than EPISODE with s = 10% .
ImageNet Figure 2 shows the performance of each algorithm at the end of training for all settings (left) and during training for the setting I = 64 and s =
(right).
Training curves for the rest of the settings are given in Appendix E.5.
EPISODE outperforms all baselines in every experimental setting, especially in the case of high data heterogeneity.
EPISODE is particularly dominant over other methods in terms of the training loss during the whole training process, which is consistent with our theory.
Also, EPISODE exhibits more resilience to data heterogeneity than CELGC and
SCAFFOLD: as the client data similarity deceases from 70% to 50% , the test accuracies of CELGC and SCAFFOLD decrease by 0.8% and 0.7%, respectively, while the test accuracy of EPISODE decreases by 0.4%.
Lastly, as communication becomes more infrequent (i.e., the communication interval I increases from 64 to 128 ), the performance of EPISODE remains superior to the baselines.
C ONCLUSION
We have presented EPISODE, a new communication-efficient distributed gradient clipping algorithm for federated learning with heterogeneous data in the nonconvex and relaxed smoothness setting.
We have proved convergence results under any noise level of the stochastic gradient.
In particular, we have established linear speedup results as well as reduced communication complexity.
Further, our experiments on both synthetic and real-world data show demonstrate the superior performance of EPISODE compared to competitive baselines in FL.
Our algorithm is suitable for the cross-silo federated learning setting such as in healthcare and financial domains (Kairouz et al., 2019), and we plan to consider cross-device setting in the future.
A P RELIMINARIES
We use F r to denote the filtration generated by
[Equation]
It means that given F r , the global solution ¯ x r is fixed, but the randomness of A r , G ir and G r still exists.
In addition, for t ∈ I r , we use H t to denote the filtration generated by
[Equation]
Recall the definitions of G ir and G r ,
[Equation]
Hence we have
[Equation]
and ∥ G r − ∇ f
[Equation]
¯
[Equation]
∥ ≤
[Equation]
, hold almost surely due to Assumption 1(iii).
Also, the local update rule of EPISODE is
[Equation]
where g it = ∇ F i ( x it ; ξ it ) − G ir + G r , A r = {∥ G r ∥ ≤ γ/η } and ¯ A r
= {∥ G r ∥ > γ/η } .
A. A UXILIARY L EMMAS
Lemma A.1 (Lemma A.2 in Zhang et al.
(2020a)) .
Let f be ( L 0 , L 1 ) -smooth, and C > 0 be a constant.
For any x , x ′ ∈ R d such that ∥ x − x ′ ∥ ≤ C/L 1 , we have
[Equation]
where A = 1 + e C − e C − 1 C and B =
e C − 1 C .
Lemma A.2 (Lemma A.3 in Zhang et al.
(2020a)) .
Let f be ( L 0 , L 1 ) -smooth, and C > 0 be a constant.
For any x , x ′ ∈ R d such that ∥ x − x ′ ∥ ≤ C/L 1 , we have ∥∇ f ( x ′ )
− ∇ f ( x ) ∥ ≤ ( AL
+ BL 1 ∥∇ f ( x ) ∥ )
∥ x ′ − x ∥ , where A = 1 + e C − e C − 1 C and B = e C − 1 C .
Here we choose C ≥ 1 such that A ≥ 1 and B ≥ 1 .
Lemma A.3 (Lemma B.1 in Zhang et al.
(2020a)) .
Let µ > 0 and u , v ∈ R d .
Then
[Equation]
B P ROOF OF L EMMAS IN S ECTION 4. B. P ROOF OF L EMMA 4.
Lemma 4.1 restated.
Suppose 2 ηI ( AL 0 + BL 1 κ + BL 1 ρ ( σ + γη )) ≤ 1 and max (cid:110) 2 ηI (2 σ + γη ) , γI (cid:111) ≤ CL 1 , where the relation between A , B and C is stated in Lemma
A.1 and A.2.
Then for any i ∈ [ N ] and t − 1 ∈
I r , it almost surely holds that
[Equation]
(10) and
[Equation]
Proof of Lemma 4.1.
To show (10) holds, it suffices to show that under the event A r ,
[Equation]
holds for any t r + 1 ≤ t ≤ t r +
and i ∈
[ N ] .
We will show it by induction.
In particular, to show that this fact holds for t = t r + 1 , notice
[Equation]
where we used the fact that ∥ G r ∥ ≤ γη under A r , and ∥∇ F i (¯ x r ; ξ it r ) − ∇ F i (¯ x r ) ∥ ≤ σ ,
∥ G ir − ∇ F i (¯ x r ) ∥
≤ σ hold almost surely.
Now, denote Λ = 2 (cid:16) 2 σ + γη (cid:17) and suppose that
[Equation]
Then we have
[Equation]
Using our assumption η Λ I ≤ C/L 1 together with the inductive assumption (12), we can apply
Lemma A.2 to obtain
[Equation]
where ( i ) comes from the heterogeneity assumption ∥∇ f i ( x ) ∥ ≤ κ + ρ ∥∇ f ( x ) ∥ for all x and ( ii ) from the assumption 2 ηI ( AL 0 + BL 1 κ + BL 1 ρ ( σ + γη )) ≤ 1 .
Substituting this into Equation (13) yields
[Equation]
which completes the induction and the proof of Equation (10)
[Equation]
Next, to show Equation (11)
[Equation]
notice that under the event ¯ A r we have
[Equation]
B. P ROOF OF L EMMA 4.
Lemma 4.2 restated.
Suppose 2 ηI ( AL 0 + BL 1 κ + BL 1 ρ ( σ + γη )) ≤ 1 and max (cid:110) 2 ηI (cid:16) 2 σ + γη (cid:17) , γI (cid:111) ≤ CL 1 .
Then for all x ∈ R d such that ∥ x − ¯ x r ∥ ≤ 2 ηI (cid:16) 2 σ + γη (cid:17) , we have the following inequality almost surely holds:
[Equation]
Proof of Lemma 4.2.
Under the event A r = {∥ G r ∥ ≤ γ/η } .
From the definition of ( L 0 , L 1 ) - smoothness we have
[Equation]
where we used the heterogeneity assumption ∥∇ f i ( x ) ∥ ≤ κ + ρ ∥∇ f
(¯ x r ) ∥ for all x to obtain ( i ) and the fact ∥∇ f (¯ x r ) ∥ ≤ ∥∇ f (¯ x r ) − G r
∥
+ ∥ G r ∥ to obtain ( ii ) .
Now, for all x such that ∥ x − ¯ x r ∥ ≤ 2 ηI (2 σ + γη ) , according to our assumptions, we have ∥ x − ¯ x r ∥ ≤ 2 ηI (2 σ + γη ) ≤ CL 1 .
Hence we can apply Lemma A.2 to x and ¯ x r , which yields
[Equation]
where ( i ) comes from the assumption 2 ηI ( AL 0 + BL 1 κ + BL 1 ρ ( σ + γη )) ≤ 1 .
Substituting this result into Equation (15) yields
[Equation]
B. P ROOF OF L EMMA 4.
Lemma 4.3 restated.
Suppose 2 ηI ( AL 0 + BL 1 κ + BL 1 ρ ( σ + γη )) ≤ 1 and max (cid:110) 2 ηI (2 σ + γη ) , γI (cid:111) ≤ CL 1 , we have both
[Equation]
hold for any t − 1 ∈
I r .
Proof of Lemma 4.3.
Under the event A r , the local update rule is given by
[Equation]
where g
[Equation]
= ∇ F
[Equation]
it
[Equation]
G
[Equation]
+ G r .
Using the basic inequality ( a + b ) 2 ≤ (1 + 1 /λ )
a 2 + ( λ + 1) b 2 for any λ > 0 , we have
[Equation]
The equality ( i ) and ( ii ) hold since F r ⊆
H t for t ≥ t r such that
[Equation]
and
[Equation]
Let L =
L 0 + L 1 ( κ + ( ρ + 1)( γη + 2 σ )) .
Applying the upper bound for Hessian matrix in Lemma 4.2 and the premise in 4.1, we have
[Equation]
[Equation]
where the second inequality follows from ∥ G ir − ∇ f i (¯ x r ) ∥ ≤ σ almost surely.
Plugging the final bound of (19) into (18) yields
[Equation]
By recursively invoking (20)
[Equation]
we are guaranteed that
[Equation]
The inequality ( i ) comes from
[Equation]
which is true because 2 ηI ( AL 0 + BL 1 κ + BL 1 ρ ( σ + γη )) ≤ 1 and A, B ≥ 1 .
The inequality ( ii ) comes from ( 2 I +1)
I ( I +1) ≤ e 2 I for any I ≥ 1 .
The inequality ( iii ) holds since ∥ G r −∇ f ( ¯ x r ) ∥ ≤ σ almost surely.
Therefore, we have proved (16).
In addition, for (17), we notice that
[Equation]
The inequality ( iv ) holds since ∥ G r ∥ ≤ γ/η holds under the event A r and ∥ G r − ∇ f
(¯ x r ) ∥ ≤ σ almost surely.
C P ROOF OF M AIN R ESULTS C. P ROOF OF L EMMA 4.
Lemma 4.4 restated.
Under the conditions of Lemma 4.
[Equation]
let p r
= P ( A r |F r ) , Γ =
AL
+
BL 1 ( κ + ρ ( γη + σ )) .
Then it holds that for each 1 ≤ r ≤ R − 1 ,
[Equation]
where
[Equation]
and
[Equation]
Proof.
We begin by applying Lemma A.1 to obtain a bound on f (¯ x r +1 )
− f (¯ x r )
[Equation]
but first we must show that the conditions of Lemma A
[Equation]
1 hold here.
Note that
[Equation]
where the last step is due to the conditions of Lemma 4.1.
This shows that we can apply Lemma A. to obtain
[Equation]
Let p r
=
P
A r |F r )
, then 1 − p r
=
P
¯ A r |F r ) .
Notice that p r is a function of ¯ x r .
The last term in
Equation (21) can be bounded as follows:
[Equation]
where ( i ) comes from an application of Lemma 4.1 with t = t r +
[Equation]
and ( ii ) comes from an application of (17) in Lemma 4.3.
Substituting (22) into (21) gives
[Equation]
We introduce three claims to bound the first three terms in (23), whose proofs are deferred to Section
D.
Claim C.1.
Under the conditions of Lemma 4.4, we have
[Equation]
Claim C.2.
Under the conditions of Lemma 4.4, we have − η E r (cid:34)
N
N (cid:88)
[Equation]
(cid:88) t
[Equation]
(cid:35) ≤ p r (cid:20)(cid:18)
− ηI
[Equation]
36 I
[Equation]
I
[Equation]
η
[Equation]
σ
[Equation]
Γ
[Equation]
(cid:21) −
[Equation]
E r   1 ( A r ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 1 N
N (cid:88) i = (cid:88) t ∈I r ∇ f ( x it ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 2   , where
Γ =
AL
+ BL 1 (cid:16) κ
+ ρ (cid:16) σ + γη (cid:17)(cid:17) .
Claim C.3.
Under the conditions of Lemma 4.4, we have
[Equation]
i =
[Equation]
(cid:88) t
[Equation]
r
[Equation]
f
[Equation]
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
  .
[Table]
[Equation]
where the last inequality holds since η/ (2 I ) ≥ 4 η 2 due to the assumption 4 AL 0 ηI ≤ 1 .
Then we can finish the proof of Lemma 4.4 by noticing that p r = E r
[ 1 ( A r )] and 1 − p r = E r
[ 1 ( ¯ A r )] .
C. P ROOF OF T HEOREM 4.
Theorem 4.1 restated.
Suppose Assumption 1 hold.
For any ϵ ≤ 3 AL 0 5 BL 1 ρ , we choose η ≤ min (cid:26) 856Γ I ,
[Equation] [Equation]
where Γ =
AL
+ BL 1 κ + BL 1 ρ (cid:16) σ + γη (cid:17) .
The output of EPISODE satisfies
[Equation]
as long as R ≥ 4∆ ϵ 2 ηI .
Proof.
In order to apply Lemma 4.4, we must verify the conditions of Lemma 4.1 under our choice of hyperparameters.
From our choices of η and γ , we have
[Equation]
Also
[Equation]
where ( i ) comes from the condition η ≤ 1 / (856Γ I ) in (24), ( ii ) is true due to the fact that B, C ≥and ρ ≥ 1 .
Lastly, it also holds that
[Equation]
Therefore the conditions of Lemma 4.1 are satisfied, and we can apply Lemma 4.4.
Denoting
U
[Equation]
and
[Equation]
Lemma
[Equation]
4 tells us that
[Equation]
We will proceed by bounding each U ( x ) and V ( x ) by the same linear function of ∥∇ f ( x ) ∥ .
To bound U ( x ) , notice
[Equation]
where ( i ) comes from ρ ≥ 1 and ( ii ) comes from 856Γ ηI ≤ 1 and ( iii ) holds since γ/η = 11 σ + AL
BL 1 ρ and ( iv ) comes from ϵ ≤ 3 AL 0
BL 1 ρ .
Also, we have
[Equation]
Plugging Equations (28) and (29) into Equation (25) yields
[Equation]
Now to bound V ( x ) , we have
[Equation]
(31) where ( i ) comes from η ≤ 1 856Γ I and Γ > BL 1 σ + AL 0 /ρ for ρ > 1 .
Using the assumption η ≤ ϵ 180 I Γ σ , it holds that
[Equation]
where ( ii ) comes from 16 BL 1 σ + AL 0 <
5Γ .
Lastly, we have
[Equation]
where ( i ) comes from η ≤ min (cid:110)
[Equation]
σ
[Equation]
Plugging Equations (31), (32), and (33) into (26) then yields
[Equation]
We can then use the inequality x 2 ≥ 2 ax − a 2 with x = ∥∇ f ( x ) ∥ and
a =
ϵ to obtain
[Equation] [Equation]
Having bounded U ( x ) and V ( x ) , we can return to (27).
Using (30), we can see
[Equation]
so the RHS of (34) is an upper bound of both U ( x ) and V ( x ) .
Plugging this bound into (27) and taking total expectation then gives
[Equation]
Finally, denoting ∆ =
f (¯ x
[Equation]
− f ∗ , we can unroll the above recurrence to obtain
[Equation]
where the last inequality comes from our choice of R ≥ 4∆ ϵ 2 ηI .
D D EFERRED P ROOFS OF S ECTION C D. P ROOF OF C LAIM C.
Proof.
Starting from Lemma A.3 with u = ∇ f
(¯ x r ) and v = g it , we have
[Equation]
Under ¯ A r = {∥ G r ∥ > γη } , note that g it =
∇ F i ( x it ; ξ it ) − G ir + G r , and we have ∥ g it ∥ ≥ ∥
[Equation]
r
[Equation]
x it
[Equation]
ξ
[Equation]
and
[Equation]
Plugging these two inequalities into (35) yields
[Equation]
Under ¯ A r
[Equation]
we know ∥ x it − ¯ x r ∥ ≤ γI
[Equation]
and γI ≤ CL 1 by assumption.
Therefore we can apply Lemma
A.2 to obtain
[Equation]
This implies that
[Equation]
Combining this with the choice µ = 2 /
[Equation]
we have the final bound:
[Equation]
where we used the heterogeneity assumption ∥∇ f i (¯ x r ) ∥ ≤ κ + ρ ∥∇ f
(¯ x r ) ∥ .
D. P ROOF OF C LAIM C.
Proof.
Recall the event A r = {∥ G r ∥ ≤ γ/η } , we have
I E r (cid:34)
N
N (cid:88) i = (cid:88) t ∈I r 1 ( A r ) ⟨∇ f
(¯ x r ) , g it ⟩ (cid:35) =
E r (cid:34) 1 ( A r ) (cid:42)
I ∇ f (¯ x r ) , (cid:88) t ∈I r
N
N (cid:88) i
= g it (cid:43)(cid:35) ( i ) =
E r (cid:34) 1 ( A r ) (cid:42)
I ∇ f (¯ x r ) , (cid:88) t ∈I r
N
N (cid:88) i = ∇ F i ( x i t ; ξ i t ) (cid:43)(cid:35) ( ii ) =
E r (cid:34) 1 ( A r ) (cid:42)
I ∇ f (¯ x r ) , (cid:88) t ∈I r
N
N (cid:88) i = ∇ f i ( x it ) (cid:43)(cid:35) ( iii ) =
[Equation]
 
[Equation]
( A
[Equation]
The equality ( i ) is obtained from the fact that 1 N (cid:80) Ni =1 g it = 1 N (cid:80)
Ni =1 ∇ F i ( x it , ξ it ) − G ir + G r = 1 N (cid:80) Ni =1 ∇ F i ( x it , ξ it ) .
The equality ( ii ) holds due to the tower property such that for t > t r
E (cid:2) 1 ( A r ) ∇ F i ( x i t , ξ i t ) (cid:12)(cid:12) F r (cid:3)
=
E (cid:2) 1 ( A r ) E (cid:2) ∇ F i ( x i t , ξ i t ) (cid:12)(cid:12)
H t (cid:3) (cid:12)(cid:12) F r (cid:3) =
E (cid:2) 1 ( A r ) ∇ f i ( x i t ) (cid:12)(cid:12) F r (cid:3) ; for t = t r
[Equation]
which is true since G r = 1 N (cid:80) Ni =1 ∇ F i (¯ x r ; (cid:101) ξ ir ) is independent of ∇ F i (¯ x r , ξ it r ) given F r , and ( iii ) holds because 2 ⟨ a, b ⟩ = ∥ a ∥ 2 + ∥ b ∥ 2 − ∥ a − b ∥ 2 .
Let Γ =
AL
+ BL 1 (cid:16) κ
+ ρ (cid:16) σ + γη (cid:17)(cid:17) .
Notice that we can apply the relaxed smoothness in Lemma
A.2 to obtain
[Equation]
The inequality ( i ) holds since ∥∇ f
(¯ x r ) ∥ ≤ ∥∇ f (¯ x r ) − G r
∥
+ ∥ G r ∥ ≤ σ + γ/η almost surely under the event A r .
The inequality ( ii ) follows from the bound (16) in Lemma 4.3.
Therefore, we are guaranteed that
[Equation]
Multiplying both sides of (36) by − η/I and substituting (37) then yields − η E r
[Equation]
N
[Equation]
t ∈I r ∇ f ( x it ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 2   + p r η
I E r  
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:88) t ∈I r (cid:32)
N
N (cid:88) i = ∇ f i ( x it )
− ∇ f
(¯ x r ) (cid:33)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 2  
[Table]
p
[Table]
− ηI 2 + 36 I
[Equation]
η
[Equation]
Γ
[Equation]
(cid:19) ∥∇ f (¯ x r ) ∥
[Equation]
+ 126 I 3 η 3 σ
[Equation]
Γ
[Equation]
(cid:21) − η 2 I E r   1 ( A r ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 1 N
N (cid:88) i = (cid:88) t ∈I r ∇ f ( x it ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 2   .
D. P ROOF OF C LAIM C.
Proof.
From the definition of ¯ x r +1 , we have
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
N
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
where ( i ) is obtained by noticing that 1 N (cid:80) Ni =
g it = 1 N (cid:80) Ni =1 ∇ F i ( x it , ξ it ) , and ( ii ) holds by the fact that each client’s stochastic gradients ∇ F i ( x it , ξ it ) are sampled independently from one another.
Similarly, let s ∈
I r with s > t .
, we can see that
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
=
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
= 0 .
Therefore, we have
[Equation]
x it
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
 
[Equation]
A
[Equation]
x it
[Equation]
i ( x it ) (cid:13)(cid:13)
[Equation]
(cid:105) =
[Equation]
N
[Equation]
N
[Equation]
r
[Equation]
( A r ) E
[Equation]
(cid:104)(cid:13)(cid:13) ∇ F i ( x it
[Equation]
ξ it ) − ∇ f i ( x it ) (cid:13)(cid:13) 2 (cid:12)(cid:12) H t (cid:105)(cid:105) ≤ p r Iσ
N .
And the desired result is obtained by plugging (39) into (38).
E A DDITIONAL E XPERIMENTAL R ESULTS E. P ROOF OF P ROPOSITION 5.
Proof.
Recall the definition of f 1 ( x ) and f 2 ( x ) ,
[Equation]
which means
[Equation]
and
[Equation]
It follows that
[Equation]
Let g ( x ) =
10 H | x | − (cid:12)(cid:12) 4 x 3 − 9 x 2 (cid:12)(cid:12) , next we will characterize g ( x ) in different region.
(i) When x ∈ ( −∞ , 0) , g ( x ) = 4 x 3 − 9 x 2 − 10 Hx .
The root for the derivative of g ( x ) in this region is
[Equation]
It follows that
[Equation]
where the last inequality follows from x 1 ≤ 0 .
(ii) When x ∈ (0 , 94 ) , g ( x ) = 4 x 3 − 9 x 2 + 10 Hx .
The derivative of g ( x ) is greater than 0 in this case since 18 2 − 480 H ≤ 0 for H ≥ 1 .
Then we have
[Equation] [Equation]
(iii) When x ∈ ( 94 , + ∞ ) , g ( x ) = − 4 x 3 + 9 x 2 + 10 Hx .
The root for the derivative of g ( x ) is − 12 x 2 + 18 x + 10 H =
=
⇒ x
= x
[Equation]
= −
[Equation]
+
[Equation] [Equation]
[Equation]
Then we have
[Equation]
Combining (41), (42) and (43), we are guaranteed that
[Equation]
Substituting this bound into (40), we get
[Equation]
And κ ( H ) < ∞ is an increasing function of H .
E. S YNTHETIC TASK
For two algorithms
[Figure]
we inject uniform noise over [ − 1 , 1] into the gradient at each step, and tune γ/η ∈ { 5 , 10 , 15 } and tune η ∈ { 0 . 1 , 0 . 01 , 0 . 001 } .
We run each algorithm for 500 communication rounds and the length of each communication round is I = 8 .
The results are showed in Figure
[Figure]
CELGC and EPISODE on synthetic task.
[Figure]
E. SNLI
The learning rate η and the clipping parameter γ are tuned with search in the following way: we vary γ ∈ { 0 . 01 , 0 . 03 , 0 . 1 } and for each γ we vary η so that the clipping threshold γ/η varies over { 0 . 1 , 0 . 333 , 1 . 0 , 3 . 333 , 10 . 0 } , leading to 15 pairs ( η, γ ) .
We decay both η and γ by a factor of 0 .at
epochs 15 and 20 .
We choose the best pair ( η, γ ) according to the performance on a validation set, and the corresponding model is evaluated on a held-out test set.
Note that we do not tune ( γ, η ) separately for each algorithm.
Instead, due to computational constraints, we tune the hyperparameters for the baseline CELGC under the setting I = 4 , s = 50% and re-use the tuned values for the rest of the settings.
E. CIFARE.4. S ETUP
We train a ResNet-50 (He et al., 2016) for 150 epochs using the cross-entropy loss and a batch size of 64 for each worker.
Starting from an initial learning rate η 0 = 1 . 0 and clipping parameter γ = 0 .
5 , we decay the learning rate by a factor of 0.5 at epochs 80 and 120.
In this setting, we decay the clipping parameter γ with the learning rate η , so that the clipping threshold γη remains constant during training.
We present results for I = 8 and s ∈ { 50% , 70% } .
We include the same baselines as the experiments of the main text, comparing EPISODE to FedAvg, SCAFFOLD, and CELGC.
E.4. R ESULTS
Training loss and testing accuracy during training are shown below in Figure 4.
In both settings,
EPISODE is superior in terms of testing accuracy and nearly the best in terms of training loss.
E. I MAGE N ET
The training curves (training and testing loss) for each ImageNet setting are shown below in Figure 5.
[Figure]
[Table]
[Figure]
[Table]
F R UNNING T IME R ESULTS
To demonstrate the utility of EPISODE for federated learning in practical settings
[Table]
we also provide a comparison of the running time of each algorithm on the SNLI dataset.
Our experiments were run on eight NVIDIA Tesla V100 GPUs distributed on two machines.
The training loss and testing accuracy of each algorithm (under the settings described above) are plotted against running time below.
Note that these are the same results as shown in Figure 1, plotted against time instead of epochs or communication rounds.
On the SNLI dataset, EPISODE reaches a lower training loss and higher testing accuracy with respect to time, compared with CELGC and NaiveParallelClip.
Table 2 shows that, when I ≤ 8 , EPISODE requires significantly less running time to reach high testing accuracy compared with both CELGC and NaiveParallelClip.
When I = 16 , CELGC and NaiveParallelClip nearly match, indicating that I = 16 may be close to the theoretical upper bound on I for which fast convergence can be guaranteed.
Also, as the client data similarity decreases, the running time requirement of EPISODE to reach high test accuracy stays nearly constant (e.g., when I = 4 ), while the running time required by CELGC steadily increases.
This demonstrates the resilience of EPISODE’s convergence speed to heterogeneity.
Training curves for the same experiment are shown in Figure 6.
G A BLATION S TUDY
In this section, we introduce an ablation study which disentangles the role of the two components of
EPISODE’s algorithm design: periodic resampled corrections and episodic clipping.
Using the SNLI dataset, we have evaluated several variants of the EPISODE algorithm constructed by removing one algorithmic component at a time, and we compare the performance against EPISODE along with variants of the baselines mentioned in the paper.
Our ablation study shows that both components of
EPISODE’s algorithm design (periodically resampled corrections and episodic clipping) contribute to the improved performance over previous work.
Our ablation experiments follow the same setting as the SNLI experiments in the main text.
The network architecture, hyperparameters, and dataset are all identical to the SNLI experiments described in the main text.
In this ablation study, we additionally evaluate multiple variants of EPISODE and baselines, which are described below:
Minutes 0.
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
T r a i n l o ss
Minutes 0.
[Figure] [Figure]
T e s t a cc u r a c y
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
• SCAFFOLD (clipped): The SCAFFOLD algorithm (Karimireddy et al., 2020) with gradient clipping applied at each iteration.
This algorithm, as a variant of CELGC, determines the gradient clipping operation based on the corrected gradient at every iteration on each machine.
• EPISODE (unclipped): The EPISODE algorithm with clipping operation removed.
• FedAvg: The FedAvg algorithm (McMahan et al., 2017a).
We include this to show that clipping in some form is crucial for optimization in the relaxed smoothness setting.
• SCAFFOLD: The SCAFFOLD algorithm (Karimireddy et al., 2020).
We
include
this to show that SCAFFOLD-style corrections are not sufficient for optimization in the relaxed smoothness setting.
We compare these four algorithm variations against the algorithms discussed in the main text, which include EPISODE, CELGC, and NaiveParallelClip.
Following the protocol outlined in the main text, we train each one of these algorithms while varying the communication interval I and the client data similarity parameter s .
Specifically, we evaluate six settings formed by first fixing s = 30% and varying I ∈ { 2 , 4 , 8 , 16 } , then fixing I = 4 and varying s ∈ { 10% , 30% , 50% } .
Note that the results of NaiveParallelClip are unaffected by I and s , since
NaiveParallelClip communicates at every iteration.
For each of these six settings, we provide the training loss and testing accuracy reached by each algorithm at the end of training.
Final results for all settings are given in Table 3, and training curves for the setting I = 4 , s = 30% are shown in
Figure 7.
[Table]
From these results, we can conclude that both components of EPISODE (periodic resampled corrections and episodic clipping) contribute to EPISODE’s improved performance.
• Replacing periodic resampled corrections with SCAFFOLD-style corrections yields the variant SCAFFOLD (clipped).
In every setting, SCAFFOLD (clipped) performs slightly better than CELGC, but still worse than EPISODE.
This corroborates the intuition that SCAFFOLDstyle corrections use slightly outdated information compared to that of EPISODE, and this information lag caused worse performance in this ablation study.
• On the other hand, clipping is essential for EPISODE to avoid divergence.
By removing clipping from EPISODE, we obtain the variant EPISODE (unclipped), which fails to learn entirely.
EPISODE (unclipped) never reached a test accuracy higher than 35% , which is barely higher than random guessing, since SNLI is a 3-way classification problem.
In summary, both
[Figure]
periodic resampled corrections and episodic clipping contribute to the improved performance of
EPISODE over baselines.
In addition, FedAvg and SCAFFOLD show similar divergence behavior as EPISODE (unclipped).
None of these three algorithms employ any clipping or normalization in updates, and consequently none of these algorithms are able to surpass random performance on SNLI.
Finally, although NaiveParallelClip appears to be the best performing algorithm from this table, it requires more wall-clock time than any other algorithms due to its frequent communication.
For a comparison of the running time results, see Table 2 in Appendix F.
H N EW E XPERIMENTS ON F EDERATED L EARNING B ENCHMARK : S ENTIMENT 140 D ATASET
To evaluate EPISODE on a real-world federated dataset, we provide additional experiments on the Sentiment140 benchmark from the LEAF benchmark (Caldas et al., 2018).
Sentiment140 is a sentiment classification problem on a dataset of tweets, where each tweet is labeled as positive or negative.
For this setting, we follow the experimental setup of Li et al.
(2020b): training a 2-layer
LSTM network with 256 hidden units on the cross-entropy classification loss.
We also follow their data preprocessing steps to eliminate users with a small number of data points and split into training and testing sets.
We perform an additional step to simulate the cross-silo federated environment (Kairouz et al., 2019) by partitioning the original Sentiment140 users into eight groups (i.e., eight machines).
To simulate heterogeneity between silos, we partition the users based on a non-i.i.d.
sampling scheme similar to that of our SNLI experiments.
Specifically, given a silo similarity parameter s , each silo is allocated s % of its users by uniform sampling, and (100 − s )% of its users from a pool of users which are sorted by the proportion of positive tweets in their local dataset.
This way, when s is small, different silos will have a very different proportion of positive/negative samples in their respective datasets.
We evaluate NaiveParallelClip, CELGC, and EPISODE in this cross-silo environment with I = 4 and s ∈ { 0 , 10 , 20 } .
We tuned the learning rate η , and the clipping parameter γ with grid search over the values η ∈ { 0 . 01 , 0 .
3 , 1 . 0 } and γ ∈ { 0 . 01 , 0 .
Results are plotted in Figures 8 and 9.
Overall, EPISODE is able to nearly match the training loss and testing accuracy of NaiveParallelClip while requiring significantly less running time , and the performance of EPISODE does not degrade as the client data similarity s decreases.
Figure 8 shows that, with respect to the number of training steps, EPISODE remains competitive with NaiveParallelClip and outperforms CELGC.
In particular, the gap between EPISODE and CELGC grows as the client data similarity decreases, showing that
EPISODE can adapt to data heterogeneity.
On the other hand, Figure 9 shows that, with a fixed time budget, EPISODE is able to reach lower training loss and higher testing accuracy than both CELGC and NaiveParallelClip in all settings.
This demonstrates the superior performance of EPISODE in practical scenarios.
[Figure] [Figure]
(a) I = 4 , s = 20%
[Figure]