Title: Federated Learning with Client Subsampling, Data Heterogeneity, and Unbounded Smoothness: A New Algorithm and Lower Bounds
Venue: NeurIPS 2023
Accepted: yes
Keywords: federated learning

client subsampling

nonconvex optimization

relaxed smoothness

data heterogeneity

lower bound
Abstract: We study the problem of Federated Learning (FL) under client subsampling and data heterogeneity with an objective function that has potentially unbounded smoothness. This problem is motivated by empirical evidence that the class of relaxed smooth functions, where the Lipschitz constant of the gradient scales linearly with the gradient norm, closely resembles the loss functions of certain neural networks such as recurrent neural networks (RNNs) with possibly exploding gradient. We introduce EPISODE++, the first algorithm to solve this problem. It maintains historical statistics for each client to construct control variates and decide clipping behavior for sampled clients in the current round. We prove that EPISODE++ achieves linear speedup in the number of participating clients, reduced communication rounds, and resilience to data heterogeneity. Our upper bound proof relies on novel techniques of recursively bounding the client updates under unbounded smoothness and client subsampling, together with a refined high probability analysis. In addition, we prove a lower bound showing that the convergence rate of a special case of clipped minibatch SGD (without randomness in the stochastic gradient and with randomness in client subsampling) suffers from an explicit dependence on the maximum gradient norm of the objective in a sublevel set, which may be large. This effectively demonstrates that applying gradient clipping to minibatch SGD in our setting does not eliminate the problem of exploding gradients.  Our lower bound is based on new constructions of hard instances tailored to client subsampling and a novel analysis of the trajectory of the algorithm in the presence of clipping. Lastly, we provide an experimental evaluation of EPISODE++ when training RNNs on federated text classification tasks, demonstrating that EPISODE++ outperforms strong baselines in FL. The code is available at https://github.com/MingruiLiu-ML-Lab/episode_plusplus.

Federated Learning with Client Subsampling, Data Heterogeneity, and Unbounded Smoothness: A New Algorithm and Lower Bounds

Abstract

Westudy the problem of Federated Learning (FL) under client subsampling and data heterogeneity with an objective function that has potentially unbounded smoothness. This problem is motivated by empirical evidence that the class of relaxed smooth functions, where the Lipschitz constant of the gradient scales linearly with the gradient norm, closely resembles the loss functions of certain neural networks such as recurrent neural networks (RNNs) with possibly exploding gradient. We introduce EPISODE++, the first algorithm to solve this problem. It maintains historical statistics for each client to construct control variates and decide clipping behavior for sampled clients in the current round. We prove that EPISODE++ achieves linear speedup in the number of participating clients, reduced communication rounds, and resilience to data heterogeneity. Our upper bound proof relies on novel techniques of recursively bounding the client updates under unbounded smoothness and client subsampling, together with a refined high probability analysis. In addition, we prove a lower bound showing that the convergence rate of a special case of clipped minibatch SGD (without randomness in the stochastic gradient and with randomness in client subsampling) suffers from an explicit dependence on the maximum gradient norm of the objective in a sublevel set, which may be large. This effectively demonstrates that applying gradient clipping to minibatch SGD in our setting does not eliminate the problem of exploding gradients. Our lower bound is based on new constructions of hard instances tailored to client subsampling and a novel analysis of the trajectory of the algorithm in the presence of clipping. Lastly, we provide an experimental evaluation of EPISODE++ when training RNNs on federated text classification tasks, demonstrating that EPISODE++ outperforms strong baselines in FL. The code is available at https://github.com/MingruiLiu-ML-Lab/episode_plusplus .

1 Introduction

Federated Learning (FL) [33, 23] is a distributed learning paradigm in which many clients collaboratively train a machine learning model while communicating over a network, which preserves privacy and leverages parallelism across many clients. Minimizing communication cost, accounting for data heterogeneity across clients, and allowing for partial client participation are core principles of FL. Interest in FL has grown in recent years, especially with user-facing applications such as next word prediction on smartphones [19].

Optimization is a central part of FL algorithms, and most work on non-convex optimization assumes a smooth objective in both the single machine setting [15, 16, 1] and the FL setting [41, 48, 24, 26]. However, recent work [52, 8] has shown empirical evidence that certain neural networks (LSTMs [21], and Transformers [42]) do not satisfy this assumption, but do satisfy a weaker condition known as relaxed smoothness [52]. Under relaxed smoothness, techniques such as gradient clipping [37] are essential for avoiding exploding gradients. To avoid the negative effects of exploding gradients, GD without gradient clipping requires a step size inversely proportional to the maximum gradient norm of the objective in a sublevel set (denoted as M ), resulting in very slow convergence. The usefulness of gradient clipping under relaxed smoothness matches observations of training these neural networks in practice, for example on natural language tasks, which are common in FL [23].

However, little work yet exists for FL in the relaxed smoothness setting. Liu et al. [32] introduced a communication-efficient gradient clipping algorithm for FL under relaxed smoothness, with the additional assumption of homogeneous client data and a distributional assumption on the noise of stochastic gradients. The EPISODE algorithm [9] was subsequently introduced to handle heterogeneous client data in this setting, but requires full client participation, that is, that every client participates in every communication round. This significantly decreases the practical applicability of EPISODE, since full client participation is rarely achievable with large-scale FL in practice [23].

In this work, we introduce EPISODE++, the first algorithm for FL under relaxed smoothness, client heterogeneity, and client subsampling. EPISODE++ maintains statistics of the history of gradients for each client, and uses these statistics to (1) correct each local update step to approximate an update on the global loss and (2) determine at which steps the clipping operation should be performed. We prove that EPISODE++ achieves linear speedup in the number of participating clients, has reduced communication cost, and enjoys convergence rate independent of client heterogeneity.

A previous line of work in the non-convex smooth stochastic setting [45, 47, 27] compares FL algorithms against a classical baseline: minibatch SGD [40]. Despite the impressive results of newer

algorithms in practice, only a few works have proven that some FL algorithms have a theoretical advantage over minibatch SGD [24, 27]. It is therefore natural to ask what is the analogue of minibatch SGD in the relaxed smooth setting, and whether it is possible to improve upon this analogue. In this work we consider clipped minibatch SGD, which is obtained by applying gradient clipping to each update of minibatch SGD. We demonstrate a surprising negative result for clipped minibatch SGD: under client subsampling, relaxed smoothness, and client heterogeneity, the convergence rate of clipped minibatch SGD still depends on M . This implies that gradient clipping does not help minibatch SGD in this setting.

Our contributions can be summarized as follows:

We introduce EPISODE++, the first algorithm for FL under relaxed smoothness, client heterogeneity, and client subsampling. We prove that EPISODE++ achieves linear speedup in the number of participating clients, reduced communication cost, and has convergence rate independent of client heterogeneity. Table 1 shows a detailed comparison of the complexity of various algorithms. To achieve this result, we introduce novel techniques of recursively bounding the client updates in the presence of unbounded smoothness and client data heterogeneity, together with a refined high probability analysis.

We demonstrate a lower bound for clipped minibatch SGD in which the convergence rate depends on M (the maximum gradient norm of the objective function in a sublevel set). This shows that, in our setting, clipped minibatch SGD is susceptible to exploding gradients, and avoiding them requires a very small learning rate, which slows down convergence. Our lower bound is based on new constructions of hard instances tailored to client subsampling and a novel analysis of the trajectory of the algorithm in the presence of clipping.

We empirically evaluate EPISODE++ against strong baselines when training RNNs on federated text classification tasks. The results show that EPISODE++ consistently outperforms baselines across various client participation ratios and is resilient to heterogeneous client data, which is consistent with our theory.

2 Related Work

Federated Learning FL was proposed by [33], where the authors designed the FedAvg algorithm, which is also referred to as local SGD in the literature [41, 31, 48]. The local SGD algorithm has been analyzed in various settings, including convex smooth setting [41, 11, 28, 24, 45, 47, 26, 49, 17, 25], convex composite setting [50, 2, 35], and nonconvex smooth setting [22, 43, 31, 18, 48, 28, 24, 38, 53, 26]. There is a line of work which specially considered client partial participation in FL [6, 13, 24, 29, 44, 7] under convex or nonconvex smooth settings. Recently, Liu et al. [32] and Crawshaw et al. [9] considered FL with nonconvex and relaxed smooth functions for homogeneous and heterogeneous data respectively. However, they assume full client participation and neither of them are applicable to the case of client subsampling.

Relaxed Smoothness Relaxed smoothness was proposed by [52] as a relaxation of the standard smoothness condition, which is used to model the exploding gradient problem in training deep neural networks such as recurrent neural networks [36, 37] and long-short term memory networks [52], language models [14, 34] and transformers [8]. Zhang et al. [52] proved that gradient clipping converges faster than any fixed step size gradient descent for relaxed smooth functions. The complexity bound in [52] was further improved by [51]. Recently, there is a line of work which considered different algorithms and various analyses under relaxed smoothness [8, 39, 12]. However, all of them focused on single machine setting and may not be applicable to FL setting.

Lower bounds in Federated Learning There are several lower bound results for FL algorithms. Woodworth et al. [45, 47] compared minibatch SGD and local SGD in the regime of federated stochastic convex optimization setting for homogeneous and heterogeneous data and established lower bounds for local SGD. Woodworth et al. [46] proved a min-max complexity of distributed stochastic convex optimization for any intermittent communication algorithm. Glasgow et al. [17] established improved lower bounds of local SGD in convex optimization setting for both homogeneous and heterogeneous data. However, all of these lower bounds are not applicable to our setting where the problem instance is relaxed smooth with heterogeneous data and client subsampling.

3 Problem Setup

We consider federated learning with heterogeneous and stochastic objectives, where the goal is to minimize the average loss function across N clients. For i ∈ [ N ] , let f i ( x ) = E ξ ∼D i [ F i ( x ; ξ )] be the objective of the i -th client, where D i is the underlying data distribution of the i -th client. Then the global objective is

Since each f i is not necessarily convex, we consider the problem of finding an ϵ -stationary point, that is, a point x ∈ R d such that ∥∇ f ( x ) ∥ ≤ ϵ .

Most works on non-convex optimization [15, 16, 1] consider the case where each f i is L -smooth, so that ∥∇ 2 f i ( x ) ∥ ≤ L for every x ∈ R d . However, several works [52, 8] have shown empirical evidence that objective functions corresponding to some types of neural networks (such as LSTMs [21] and Transformers [42]) do not satisfy this condition, but do satisfy a strictly weaker condition known as ( L 0 , L 1 ) -smoothness, or relaxed smoothness. A second-order differentiable function g : R d → R is called ( L 0 , L 1 ) -smooth if ∥∇ 2 g ( x ) ∥ ≤ L 0 + L 1 ∥∇ g ( x ) ∥ for all x ∈ R d . Notice that any L -smooth function is ( L, 0) -smooth.

In this work we consider the problem described in (1) under the following assumptions.

Assumption 1. (i) Denoting by x 0 the initial iterate, there exists some ∆ > 0 such that f ( x 0 ) -min x ∈ R d f ( x ) ≤ ∆ . (ii) Each f i and f is ( L 0 , L 1 ) -smooth. (iii) There exist κ ≥ 0 , ρ ≥ 1 such that ∥∇ f i ( x ) ∥ ≤ κ + ρ ∥∇ f ( x ) ∥ for all x ∈ R d . (iv) There exists σ ≥ 0 such that E ξ ∼D i [ ∇ F i ( x ; ξ )] = ∇ f i ( x ) and ∥∇ F i ( x ; ξ ) -∇ f i ( x ) ∥ ≤ σ almost surely for ξ ∼ D i .

Assumption ( i ) is standard in non-convex optimization [15, 16]. Assumption ( ii ) is typically used in the FL literature [32, 9]. Assumption ( iii ) is used in heterogeneous federated learning [24] and describes the heterogeneity between client objectives: if κ = 0 and ρ = 1 , then all client objectives f i are equal. Assumption ( iv ) is common in the relaxed smoothness setting [52, 51, 32, 8, 9].

In addition, we consider the case of partial client participation in federated learning, also known as client subsampling. With partial participation, only S out of N clients will participate in each communication round, which exacerbates the issue of client heterogeneity.

4 Algorithm and Convergence Analysis

4.1 Main Challenges and Algorithm Design

We first illustrate why existing algorithms such as SCAFFOLD [24] and EPISODE [9] are not able to handle heterogeneous data, relaxed smoothness and client subsampling simultaneously. The analysis of SCAFFOLD crucially requires the function to be L -smooth to recursively bound (1) the lag error from client subsampling and (2) client drift from local updates, but this argument is not applicable for relaxed smooth functions whose gradient information changes quickly. EPISODE has convergence guarantees for relaxed smooth functions and heterogeneous data, but only with full client participation. A naive variant of EPISODE in the client subsampling case does not work: the indicator of gradient clipping is based only on information from clients participating in the current round and ignores information from unsampled clients, which introduces non-negligible bias from client heterogeneity.

To address these challenges, we design a new algorithm named EPISODE++, which is presented in Algorithm 1. Similar to EPISODE [9], our algorithm utilizes episodic gradient clipping , which determines whether a clipping operation will be performed depending on the size of the average control variate G r . This means that during each round, either (1) all clients perform a normalized update for all steps, or (2) all clients will perform an unnormalized update for all steps. However, different from EPISODE, our algorithm corrects local updates with control variates G i r computed as the averaged stochastic gradient over the previous round in which client i participated, as in SCAFFOLD. As we will show in our proof, the episodic gradient clipping together with the update correction strategy allows the algorithm to progress in a stable manner: and it will sufficiently decrease the objective value and also avoid the negative effect of possibly exploding gradients.

Algorithm 1 EPISODE++

4.2 Convergence Result

The following result proves that EPISODE++ converges to an ϵ -stationary point with high probability.

The above result holds for a wide range of η and I , and for any noise level σ . The corollary below summarizes the best possible iteration complexity and communication complexity implied by Theorem 1. The full proofs of Theorem 1 and Corollary 1 can be found in Appendix A.5 and A.6.

4.3 Proof Sketch

In this section we provide a sketch for the proof of Theorem 1. We wish to establish the descent inequality for the global objective function in each round by applying Lemma A.3 in [51]:

where A = 1 + e C -e C -1 C , B = e C -1 C and C > 0 is an absolute constant. However, using this descent inequality requires ∥ ¯ x r +1 -¯ x r ∥ ≤ C/L 1 . Achieving a universal bound on the distance over two consecutive rounds is nontrivial in the presence of client subsampling, data heterogeneity and relaxed smoothness, which requires a new analysis.

Compared with SCAFFOLD [24] and EPISODE [9], the main difficulty of analyzing EPISODE++ lies in controlling the distance between local weights x i r,k and the synchronization point ¯ x r under relaxed smoothness and client subsampling. In EPISODE, the magnitude of g i r,k can be bounded in terms of ∥∇ f i ( x i r,k ) -∇ f i (¯ x r ) ∥ and ∥ G r ∥ , since G i r is evaluated at the synchronization point ¯ x r . Then the distance can be bounded recursively when clipping does not happen (i.e., ∥ G r ∥ ≤ γ/η ). However, EPISODE++ utilizes historical gradients to construct the indicator G r , which means the increment also depends on the lag error ∥ G i r -∇ f i (¯ x r ) ∥ . Due to relaxed smoothness, we cannot bound the lag error as in SCAFFOLD [24] because we do not know the distance between ¯ x r and where G i r is evaluated. In the extreme case that client i has never been sampled before round r , we have G i r = ∇ F i (¯ x 0 ; ˜ ξ i ) , so the lag error can be unbounded when r is large.

To address the issue mentioned above, we analyze the convergence of EPISODE++ in a highprobability framework. Let y i r,k denote the local model of client i at step k during the most recent round in which client i participated (before round r ), and let q i r denote the index of this round. Since client subsampling independent across rounds, P ( q i r -r ≥ K ) = (1 -S/N ) K . Therefore, we introduce the event E := { max 0 ≤ r ≤ R -1 , 1 ≤ i ≤ N r -q i r ≤ K } , where K is a logarithmic factor such that P ( E ) ≥ 1 -δ . Under E , we obtain the following lemma to bound the distance of local updates.

Lemma 1. Suppose that (9) holds. Then for any r ≥ 0 we have

The proof of Lemma 1 is deferred to Appendix A.3, which relies on a jump start analysis. If ∥ G r ∥ > γ/η , then (4) trivially holds due to the clipping operation. When ∥ G r ∥ ≤ γ/η , a recursive argument shows that the discrepancy ∥ x i r,k -¯ x r ∥ depends on the magnitude of the increment at the starting point, that is the lag error ∥ E [ G i r ] -∇ f i (¯ x r ) ∥ and ∥ G r ∥ after removing noise. According to the construction of G i r , we know E [ G i r ] = 1 I ∑ I k =1 ∇ f i ( y i r,k ) = 1 I ∑ I k =1 ∇ f i ( x i q i r ,k ) . Therefore, under the event E , the lag error at the r -th round can be bounded by the discrepancies of previous K +1 rounds at most, which is the bound (5). In summary, the discrepancy at any round can be controlled recursively if the discrepancy at the initial round is small. This insight motivates the initialization G i 0 = ∇ F i (¯ x 0 , ˜ ξ i ) in Algorithm 1, which enjoys zero initial lag error.

Finally, Lemma 1 shows that the condition of (3) can be satisfied by choosing η, γ, I , which establishes descent of the global objective from ¯ x r to ¯ x r +1 . Summing from r = 0 , . . . , R -1 and applying concentration bounds over martingale difference sequences to yield high probability bounds for error terms coming from stochastic gradient noise and client subsampling yields

where A r = {∥ G r ∥ ≤ γ/η } denotes the clipping indicator, and U ( x ) , V ( x ) are defined in the proof of Theorem 1. According to the choice of η, γ , the dominant term of U (¯ x r ) is -γI ∥∇ f (¯ x r ) ∥ and that of V (¯ x r ) is -ηI ∥∇ f (¯ x r ) ∥ 2 . Plugging into (6) and rearranging proves Theorem 1.

5 Lower Bound for Clipped Minibatch SGD

Clipped minibatch SGD is a natural extension of minibatch SGD [40, 45] to the relaxed smooth setting (see pseudocode in Algorithm 2, Appendix B). Clipped minibatch SGD is nearly identical to minibatch SGD: with the addition of gradient clipping to each round's update. In the similar spirit of [45], we are interested in this algorithm because it has the same computation and communication structure as EPISODE++ and it is important to understand whether EPISODE++ has any advantage over clipped minibatch SGD. In fact, we will show that clipped minibatch SGD is significantly hindered by the combination of relaxed smoothness, client heterogeneity, and client subsampling.

Assumption 2. There exists M > 0 such that ∥∇ f ( x ) ∥ ≤ M for all x with f ( x ) ≤ f ( x 0 ) .

Aline of work on relaxed smoothness in the single-machine setting [52, 8] has shown that the number of iterations required to find an ϵ -stationary point by gradient descent (GD) under relaxed smoothness is ˜ Ω ( ∆ L 1 M ϵ 2 ) , while that of GD with gradient clipping is O ( ∆ L 0 ϵ 2 ) . In this way, gradient clipping can remove the dependence on M , which can be large, and significantly speed up optimization.

Theorem 2 shows a surprising result: under some conditions on the participation ratio, number of clients, and heterogeneity parameter ρ , clipped minibatch SGD requires ˜ Ω ( ∆ L 1 M ϵ 2 ) rounds, showing that applying gradient clipping to minibatch SGD does not eliminate dependence on M .

Theorem 2. Fix ϵ > 0 , 0 < δ < 1 , L 0 > 0 , L 1 > 0 , κ > 0 , ρ > 2+log(2 -δ ) log(2 -δ ) , M > max( L 0 L 1 , ϵ ) , N ≥ ( ρ +1)(1+log(2 -δ )) ( ρ -1) log(2 -δ ) -2 . Define Q = ⌊ κ +( ρ -1) M κ +( ρ +1) M N ⌋ . Let F ( L 0 , L 1 , M, κ, ρ, N ) denote the set of problem instances { f i } N i =1 satisfying Assumptions 1(ii)-(iv) and 2 with σ = 0 . For any fixed choice of parameters γ, η based on the knowledge of above constants, there exists { f i } N i =1 ∈ F ( L 0 , L 1 , M, κ, ρ, N ) such that clipped minibatch SGD initialized at x 0 with 1 ≤ S ≤ log(2 -δ )( Q +1) N -Q +log(2 -δ ) will satisfy P ( ∥∇ f ( x r ) ∥ < ϵ for some 0 ≤ r ≤ R -1) > 1 -δ only if

The proof is included in Appendix B. This result shows that clipped minibatch SGD in our setting suffers the same problem as GD in the single-machine setting: divergence can only be avoided with a very small step size, leading to slow convergence. In contrast, the convergence rate of EPISODE++ in the same setting is independent of M .

The proof of Theorem 2 analyzes clipped minibatch SGD for three different problem instances. The first contains linear local objectives with high heterogeneity: if the clipping threshold is sufficiently small ( γ η ≤ M ), then clipped minibatch SGD will never converge with probability δ . The second instance contains homogeneous, exponential local objectives: the learning rate must be sufficiently small η < O ( 1 L 1 M ) to avoid divergence due to the exponentially increasing gradient magnitude. However, with a large clipping threshold and small learning rate, the convergence of clipped minibatch SGD will depend on M for the third problem instance, which has homogeneous linear objectives. Note that our lower bound is different from previous lower bounds in [52, 51] which are in the settings of single machine [52] or almost sure bounded noise [51], since our lower bound is considering noise from client subsampling and client data heterogeneity which is not almost surely bounded.

6 Experiments

To validate our theory, we evaluate EPISODE++ and baselines in the training of RNNs for two text classification tasks. We compare EPISODE++ to CELGC [32], clipped minibatch SGD, and NaiveParallelClip [32], which is a naive parallel implementation of SGD with gradient clipping that requires communication at every iteration. As an ablation study, we also evaluate two algorithms closely related to EPISODE++. The first is a naive extension of EPISODE [9] for client subsampling, where each participating client's control variate G i r is resampled at the beginning of each round, and G r = 1 S ∑ i ∈S r G i r . The second is an algorithm which we refer to as SCAFFOLDClip [9], which applies gradient clipping to each local step of SCAFFOLD [24]. We evaluate these six

SNLI

2

4

6

8

Participating Clients

0.3

0.4

0.5

0.6

0.7

Train Loss

2

4

6

8

Participating Clients

0.70

0.75

0.80

Test Accuracy

50%

30%

10%

Data similarity

0.3

0.4

0.5

0.6

0.7

Train Loss

50%

30%

10%

Data similarity

0.70

0.75

0.80

Test Accuracy

EPISODE++

CELGC

NaiveParallelClip

Clipped Minibatch SGD

SCAFFOLDClip

EPISODE

(a) Training loss and testing accuracy for SNLI dataset.

2

4

6

8

Participating Clients

0.3

0.4

0.5

0.6

Train Loss

2

4

6

8

Participating Clients

0.74

0.76

0.78

Test Accuracy

20%

10%

0%

Data similarity

0.3

0.4

0.5

0.6

Train Loss

20%

10%

0%

Data similarity

0.74

0.76

0.78

Test Accuracy

Sentiment140

EPISODE++

CELGC

NaiveParallelClip

Clipped Minibatch SGD

SCAFFOLDClip

EPISODE

algorithms on natural language inference with the SNLI dataset [4] and sentiment classification with the Sentiment140 dataset [5]. More experimental results can be found in Appendix C.

6.1 Setup

All experiments use uniform client sampling, a batch size of 64 (on each client) and the multi-class hinge loss. See Appendix C.1 for full details on hyperparameters. All experiments were implemented in PyTorch and ran on eight NVIDIA V100 GPUs.

SNLI SNLI [4] is a 3-way text classification task, in which the logical relationship of a pair of sentences must be classified as either entailment, neutral, or contradictory. The dataset contains 570k pairs of sentences. Because SNLI is a centralized dataset, we follow the heterogeneity protocol in [24] to divide the dataset into clients according to a similarity parameter s between 0% and 100% . According to this protocol, s % of each client's local dataset is allocated from a randomly shuffled set of examples, while the remaining (100 -s )% is allocated from a set of examples which is sorted by label. In this way, the similarity of the label distributions of client datasets grows with s . The network consists of a one-layer bidirectional RNN encoder followed by a three-layer fully connected classifier. We train for R = 5375 communication rounds with I = 4 for all algorithms except NaiveParallelClip, which uses R = 21500 and I = 1 , so that every algorithm runs the same number of training steps.

Sentiment140 Sentiment140 [5] is a sentiment prediction task designed for FL. The dataset is comprised of tweets, each labeled as either positive or negative. We follow the data processing steps of [28] to discard users with small datasets and to split into training and testing sets. In order to control the data heterogeneity between clients, we follow a similar protocol as described for SNLI to form client datasets by combining the datasets of original users in the Sentiment140 dataset. The process is nearly identical to that of SNLI, but here we allocate users to each local dataset instead of examples . s % of each client's local dataset is allocated from a randomly shuffled set of users , while the remaining (100 -s )% is allocated from a set of users sorted by the proportion of positive samples. We train for R = 2000 communication rounds with I = 4 for all algorithms except NaiveParallelClip, which uses R = 8000 and I = 1 . We use the same network architecture as SNLI.

0

2000

4000

Rounds

0.4

0.6

0.8

1.0

Train Loss

SNLI

0

2000

4000

Rounds

0.4

0.5

0.6

0.7

0.8

Test Accuracy

0

1000

2000

Rounds

0.4

0.6

0.8

Train Loss

Sentiment140

0

1000

2000

Rounds

0.65

0.70

0.75

Test Accuracy

EPISODE++

CELGC

NaiveParallelClip

Clipped Minibatch SGD

SCAFFOLDClip

EPISODE

We study the effect of client subsampling and data heterogeneity by varying each of these values in a controlled way. We first fix the heterogeneity s = 30% and vary S ∈ { 2 , 4 , 6 , 8 } , then we fix the number of participating clients S = 4 and vary s ∈ { 10% , 30% , 50% } . It should be noted that in the experiments with varying S , we always train for a fixed number of iterations RI . This means that separate training runs with different S will have the same per-client computation cost (number of gradient computations), but the total computation cost of a training run scales with S . For this setting, we use N = 8 clients.

To simulate large-scale federated learning, we also include results with a larger number of clients N = 128 . For this setting, we use S = 16 participating clients in each round, and data heterogeneity of s = 30% (SNLI) and s = 10% (Sentiment140).

6.2 Results

Figure 1 contains the final training loss and testing accuracy for the variety of settings of client participation and data heterogeneity. For a single setting of participation and heterogeneity, Figure 2 and Table 2 show learning curves and average results over three trials, respectively. More learning curves are given in Appendix C.2. Learning curves for large-scale experiments are shown in Figure 3. EPISODE++ achieves the minimum training loss and maximum testing accuracy of all algorithms in nearly every setting. Only with full participation S = 8 does NaiveParallelClip achieve a lower training loss than EPISODE++, but EPISODE++ maintains a higher testing accuracy. Also, NaiveParallelClip requires a much larger communication cost than EPISODE++ to perform the same number of training iterations: when the number of communication rounds is fixed, EPISODE++ significantly outperforms NaiveParallelClip.

Effect of Subsampling The first two plots of Figures 1(a) and 1(b) show the performance of each algorithm as the number of participating clients S varies over { 2 , 4 , 6 , 8 } with fixed data similarity s = 30% . Clipped minibatch SGD, NaiveParallelClip, and EPISODE all exhibit degraded performance as S decreases, whereas EPISODE++, SCAFFOLDClip, and CELGC maintain relatively constant performance as S decreases. Despite the constant performance of CELGC and SCAFFOLDClip under client sampling, both algorithms are significantly outperformed by EPISODE++.

0

2000

4000

Rounds

0.2

0.4

0.6

0.8

1.0

Train Loss

SNLI

0

2000

4000

Rounds

0.4

0.5

0.6

0.7

0.8

Test Accuracy

0

1000

2000

Rounds

0.0

0.2

0.4

0.6

0.8

Train Loss

Sentiment140

0

1000

2000

Rounds

0.65

0.70

0.75

Test Accuracy

EPISODE++

CELGC

NaiveParallelClip

Clipped Minibatch SGD

SCAFFOLDClip

EPISODE

Effect of Heterogeneity The last two plots of each row in Figure 1 show each algorithm's performance as the data similarity s varies over { 50% , 30% , 10% } for SNLI and { 20% , 10% , 0% } for Sentiment140, with fixed S = 4 . For both datasets, decreasing data similarity negatively impacts the performance of clipped minibatch SGD, CELGC, and NaiveParallelClip. EPISODE++, SCAFFOLDClip, and EPISODE are able to maintain performance as data similarity decreases, though EPISODE++ maintains a significantly better performance than SCAFFOLDClip and EPISODE.

Large-Scale Experiments With a larger number of clients ( N = 128 , S = 16 ), the relative performance of each algorithm is similar to the N = 8 setting, as shown in Figure 3. Here, the proportion of participating clients S/N = 1 / 8 is smaller than that of the N = 8 experiments, where S/N ≥ 1 / 4 . This suggests that the effect of partial client participation may be stronger in the large-scale experiments. EPISODE++ still outperforms all other algorithms in the large-scale setting, while maintaining about the same test accuracy as the N = 8 setting, demonstrating the effectiveness of EPISODE++ for large-scale federated learning.

Comparison with Ablations By comparing EPISODE++ against the closely related EPISODE and SCAFFOLDClip, we can see that the use of information from previous rounds and episodic gradient clipping are both critical for the superior performance of EPISODE++. EPISODE only utilizes information from the currently participating clients, ignoring information from clients that participated in previous rounds. As a result, the performance of EPISODE degrades as S decreases. On the other hand, SCAFFOLDClip determines whether to perform clipping individually for each local step, as opposed to the episodic gradient clipping of EPISODE++ and EPISODE. Although SCAFFOLDClip maintains performance under changes in the participation ratio and data similarity, the level it maintains is significantly lower than that of EPISODE++.

Communication Cost NaiveParallelClip suffers a large communication cost for the same number of training iterations compared with other algorithms, due to the cost of synchronizing clients at every iteration. As shown in Figure 2 and Figure 3, EPISODE++ outperforms all other algorithms by a wide margin when the number of communication rounds is fixed. Also, EPISODE requires twice the number of communication operations per training round, which doubles the time required for communication per round compared to all other algorithms.

7 Conclusion

We have presented EPISODE++, the first algorithm for FL with heterogeneous data and client subsampling under relaxed smoothness. We proved that EPISODE++ finds an ϵ -stationary point with high probability, and its convergence rate satisfies linear speedup and resilience to heterogeneity while enjoying reduced communication. We also presented a lower bound showing that the convergence rate of a special case of clipped minibatch SGD in our setting suffers a dependence on M (the maximum gradient norm of the objective in a sublevel set), implying that applying gradient clipping to minibatch SGD does not alleviate the problem of exploding gradients. Our experimental results for RNN training on text classification tasks demonstrate the superior performance of EPISODE++ compared to baselines. One limitation of our current work is that our lower bound assumes σ = 0 , and we plan to get a better lower bound in the future for σ > 0 .

Acknowledgments and Disclosure of Funding

We would like to thank the anonymous reviewers for their helpful comments. Michael Crawshaw is supported by the Institute for Digital Innovation fellowship from George Mason University. Michael Crawshaw and Mingrui Liu are both supported by a grant from George Mason University. Computations were run on ARGO, a research computing cluster provided by the Office of Research Computing at George Mason University (URL: https://orc.gmu.edu). The work of Yajie Bao was done when he was virtually visiting Mingrui Liu's research group in the Department of Computer Science at George Mason University.

References

Yossi Arjevani, Yair Carmon, John C Duchi, Dylan J Foster, Nathan Srebro, and Blake Woodworth. Lower bounds for non-convex stochastic optimization. Mathematical Programming , 199 (1-2):165-214, 2023.

Yajie Bao, Michael Crawshaw, Shan Luo, and Mingrui Liu. Fast composite optimization and statistical recovery in federated learning. In International Conference on Machine Learning , pages 1508-1536. PMLR, 2022.

Rémi Bardenet and Odalric-Ambrym Maillard. Concentration inequalities for sampling without replacement. Bernoulli , 21(3):1361 - 1385, 2015. doi: 10.3150/14-BEJ605. URL https: //doi.org/10.3150/14-BEJ605 .

Samuel R. Bowman, Gabor Angeli, Christopher Potts, and Christopher D. Manning. A large annotated corpus for learning natural language inference. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing (EMNLP) . Association for Computational Linguistics, 2015.

Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Koneˇ cn` y, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv:1812.01097 , 2018.

Wenlin Chen, Samuel Horvath, and Peter Richtarik. Optimal client sampling for federated learning. arXiv preprint arXiv:2010.13723 , 2020.

Yae Jee Cho, Pranay Sharma, Gauri Joshi, Zheng Xu, Satyen Kale, and Tong Zhang. On the convergence of federated averaging with cyclic client participation. arXiv preprint arXiv:2302.03109 , 2023.

Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang. Robustness to unbounded smoothness of generalized signsgd. Advances in neural information processing systems , 2022.

Michael Crawshaw, Yajie Bao, and Mingrui Liu. Episode: Episodic gradient clipping with periodic resampled corrections for federated learning with heterogeneous data. In The Eleventh International Conference on Learning Representations , 2023.

Ashok Cutkosky and Harsh Mehta. High-probability bounds for non-convex stochastic optimization with heavy tails. Advances in Neural Information Processing Systems , 34:4883-4895, 2021.

Aymeric Dieuleveut and Kumar Kshitij Patel. Communication trade-offs for local-sgd with large step size. Advances in Neural Information Processing Systems , 32:13601-13612, 2019.

Matthew Faw, Litu Rout, Constantine Caramanis, and Sanjay Shakkottai. Beyond uniform smoothness: A stopped analysis of adaptive sgd. arXiv preprint arXiv:2302.06570 , 2023.

Yann Fraboni, Richard Vidal, Laetitia Kameni, and Marco Lorenzi. Clustered sampling: Low-variance and improved representativity for clients selection in federated learning. In International Conference on Machine Learning , pages 3407-3416. PMLR, 2021.

Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N Dauphin. Convolutional sequence to sequence learning. In International Conference on Machine Learning , pages 12431252. PMLR, 2017.

Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM Journal on Optimization , 23(4):2341-2368, 2013.

Saeed Ghadimi and Guanghui Lan. Accelerated gradient methods for nonconvex nonlinear and stochastic programming. Mathematical Programming , 156(1-2):59-99, 2016.

Margalit R Glasgow, Honglin Yuan, and Tengyu Ma. Sharp bounds for federated averaging (local sgd) and continuous perspective. In International Conference on Artificial Intelligence and Statistics , pages 9050-9090. PMLR, 2022.

Farzin Haddadpour, Mohammad Mahdi Kamani, Mehrdad Mahdavi, and Viveck Cadambe. Local sgd with periodic averaging: Tighter analysis and adaptive synchronization. In Advances in Neural Information Processing Systems , pages 11080-11092, 2019.

Andrew Hard, Kanishka Rao, Rajiv Mathews, Swaroop Ramaswamy, Françoise Beaufays, Sean Augenstein, Hubert Eichner, Chloé Kiddon, and Daniel Ramage. Federated learning for mobile keyboard prediction. arXiv preprint arXiv:1811.03604 , 2018.

Elad Hazan, Kfir Y. Levy, and Shai Shalev-Shwartz. Beyond convexity: Stochastic quasi-convex optimization. In Proceedings of the 28th International Conference on Neural Information Processing Systems - Volume 1 , NIPS'15, page 1594-1602, Cambridge, MA, USA, 2015. MIT Press.

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation , 9(8): 1735-1780, 1997.

Peng Jiang and Gagan Agrawal. A linear speedup analysis of distributed deep learning with sparse and quantized communication. In Advances in Neural Information Processing Systems , pages 2525-2536, 2018.

Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. Foundations and Trends® in Machine Learning , 14(1-2):1-210, 2021.

Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning , pages 5132-5143. PMLR, 2020.

Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. Tighter theory for local sgd on identical and heterogeneous data. In International Conference on Artificial Intelligence and Statistics , pages 4519-4529. PMLR, 2020.

Anastasia Koloskova, Nicolas Loizou, Sadra Boreiri, Martin Jaggi, and Sebastian Stich. A unified theory of decentralized sgd with changing topology and local updates. In International Conference on Machine Learning , pages 5381-5393. PMLR, 2020.

Kfir Y Levy. Slowcal-sgd: Slow query points improve local-sgd for stochastic convex optimization. arXiv preprint arXiv:2304.04169 , 2023.

Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems , 2:429-450, 2020.

Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of fedavg on non-iid data. arXiv preprint arXiv:1907.02189 , 2019.

Xiaoyu Li and Francesco Orabona. A high probability analysis of adaptive sgd with momentum. arXiv preprint arXiv:2007.14294 , 2020.

Tao Lin, Sebastian U Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't use large mini-batches, use local sgd. arXiv preprint arXiv:1808.07217 , 2018.

Mingrui Liu, Zhenxun Zhuang, Yunwen Lei, and Chunyang Liao. A communication-efficient distributed gradient clipping algorithm for training deep neural networks. Advances in Neural Information Processing Systems , 35:26204-26217, 2022.

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics , pages 1273-1282. PMLR, 2017.

Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and optimizing LSTM language models. In International Conference on Learning Representations , 2018.

Konstantin Mishchenko, Ahmed Khaled, and Peter Richtárik. Proximal and federated random reshuffling. In International Conference on Machine Learning , pages 15718-15749. PMLR, 2022.

Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. Understanding the exploding gradient problem. corr abs/1211.5063 (2012). arXiv preprint arXiv:1211.5063 , 2012.

Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning , pages 1310-1318. Pmlr, 2013.

Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konecny, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. ICLR , 2021.

Amirhossein Reisizadeh, Haochuan Li, Subhro Das, and Ali Jadbabaie. Variance-reduced clipping for non-convex optimization. arXiv preprint arXiv:2303.00883 , 2023.

Ohad Shamir and Nathan Srebro. Distributed stochastic optimization and learning. In 2014 52nd Annual Allerton Conference on Communication, Control, and Computing (Allerton) , pages 850-857, 2014. doi: 10.1109/ALLERTON.2014.7028543.

Sebastian U Stich. Local sgd converges fast and communicates little. arXiv preprint arXiv:1805.09767 , 2018.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Jianyu Wang and Gauri Joshi. Cooperative sgd: A unified framework for the design and analysis of communication-efficient sgd algorithms. arXiv preprint arXiv:1808.07576 , 2018.

Shiqiang Wang and Mingyue Ji. A unified analysis of federated learning with arbitrary client participation. In Annual Conference on Neural Information Processing Systems , 2022.

Blake Woodworth, Kumar Kshitij Patel, Sebastian Stich, Zhen Dai, Brian Bullins, Brendan Mcmahan, Ohad Shamir, and Nathan Srebro. Is local sgd better than minibatch sgd? In International Conference on Machine Learning , pages 10334-10343. PMLR, 2020.

Blake Woodworth, Brian Bullins, Ohad Shamir, and Nathan Srebro. The min-max complexity of distributed stochastic convex optimization with intermittent communication. arXiv preprint arXiv:2102.01583 , 2021.

Blake E Woodworth, Kumar Kshitij Patel, and Nati Srebro. Minibatch vs local sgd for heterogeneous distributed learning. Advances in Neural Information Processing Systems , 33:6281-6292, 2020.

Hao Yu, Rong Jin, and Sen Yang. On the linear speedup analysis of communication efficient momentum sgd for distributed non-convex optimization. arXiv preprint arXiv:1905.03817 , 2019.

Honglin Yuan and Tengyu Ma. Federated accelerated stochastic gradient descent. Advances in Neural Information Processing Systems , 33:5332-5344, 2020.

Honglin Yuan, Manzil Zaheer, and Sashank Reddi. Federated composite optimization. In International Conference on Machine Learning , pages 12253-12266. PMLR, 2021.

Bohang Zhang, Jikai Jin, Cong Fang, and Liwei Wang. Improved analysis of clipping algorithms for non-convex optimization. Advances in Neural Information Processing Systems , 33:1551115521, 2020.

Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. In International Conference on Learning Representations , 2020.

Xinwei Zhang, Mingyi Hong, Sairaj Dhople, Wotao Yin, and Yang Liu. Fedpd: A federated learning framework with optimal rates and adaptivity to non-iid data. arXiv preprint arXiv:2005.11418 , 2020.

A Deferred Proofs of the Upper Bound

A.1 Notation and Preliminaries

Let q i 0 = 0 , y i 0 ,k = ¯ x 0 , and define

and

The intermediate variable q i r is the most recent round before (not including) round r in which client i participated. Similarly, y i r,k is the local model at step k during the most recent round before (not including) round r in which client i participated. Then we know that y i r,k = x i q i r ,k . Therefore G i r = 1 ∑ I -1 ∇ F i ( y i ; ξ i i ) . Also, we will initialize G i 0 = ∇ F i (¯ x i 0 ; ˜ ξ i ) for all clients.

Denote by A r = {∥ G r ∥ ≤ γ η } . Our analysis is simplified by taking the following conditions throughout the proof:

where C > 0 is a constant, A and B are defined in terms of C (see Lemma 2), and K is defined in Lemma 8. Usage of these conditions will be explicitly stated when used.

Through out this section, we denote the filtration

We use E r [ · ] = E [ · | F r ] and P r ( · ) to denote the conditional expectation and probability given the filtration F r . Notice that, given F r , the global weight ¯ x r is fixed but the subsampling set S r is still random and is independent of F r .

A.2 Auxiliary Lemmas

Lemma 2 (Corollary A.4 in [51]) . Let f : R d → R be ( L 0 , L 1 ) -smooth and C > 0 . For any x , y ∈ R d with ∥ x -y ∥ ≤ C L 1 ,

where A = 1 + e C -e C -1 C and B = e C -1 C .

Lemma 3 (Lemma A.3 in [51]) . Let f : R d → R be ( L 0 , L 1 ) -smooth and C > 0 . For any x , y ∈ R d with ∥ x -y ∥ ≤ C L 1 ,

where A = 1 + e C -e C -1 C and B = e C -1 C .

Lemma 4 (Lemma B.1 in [51]) . For any µ ≥ 0 and x , y ∈ R d ,

Lemma 5 (Lemma 1 in [30]) . Assume that Z 1 , . . . , Z T is a martingale difference sequence with respect to the filtration F t and E t [exp( Z 2 t /σ 2 t )] ≤ exp(1) for all t , where σ 1 , . . . , σ T is a sequence of random variables such that σ t ∈ F t . Then for any fixed λ > 0 and δ ∈ (0 , 1) , with probability at least 1 -δ ,

Lemma 6 (Improved Serfling's inequality, Proposition 2.3 in [3]) . Let X = ( x 1 , . . . , x N ) be a finite population of N real numbers and X 1 , . . . , X n denote a random sample without replacement from X . Let µ X = 1 N ∑ N i =1 x i , σ 2 X = 1 N ∑ N i =1 ( x i -µ ) 2 , and ˆ µ n = 1 n ∑ n i =1 X i . If max i ∈ [ N ] | x i | ≤ b , then for any λ > 0 ,

Lemma 7 (Lemma 12 in [10]) . Suppose X 1 , . . . , X T is a martingale difference sequence in a Hilbert space such that ∥ X t ∥ ≤ b almost surely for some constant b . Further, assume E t -1 [ ∥ X t ∥ 2 ] ≤ σ 2 t with probability 1 for some constants σ t . Then with probability at least 1 -3 δ , for all k :

A.3 Proof of Lemma 1

Define

Ξ r is the average 'lag" (over clients) of the correction G i r behind ∇ f (¯ x r ) due to client sampling, since each G i r is set according to stochastic gradients at y i r,k .

For any δ ∈ (0 , 1) , define and

The event E occurs when every round in which a client is sampled occurs after no more than K rounds following the previous round in which that client was sampled. This means that under E , the correction G i r was computed within K rounds of round r . The next lemma shows that E occurs with probability 1 -δ .

Lemma 8. For any δ with 0 < δ < 1 , P ( E ) ≥ 1 -δ .

Proof. If S = N , then r -q i r = 1 , and we are done. Otherwise, since each round's participating clients are sampled uniformly without replacement and independently at each round,

Also by the choice of K : K ≥ log ( RN δ ) log ( N N -S ) , we can show that ( 1 -S N ) K ≤ δ RN . So P ( r -q i r ≤ K ) ≥ 1 -δ RN . Therefore

Lemma9. Let r ≥ 0 . If ∥ ¯ x r -y i r,k ∥ ≤ C L 1 for all i ∈ [ N ] and k ∈ { 0 , . . . , I -1 } , and BL 1 ρ Ξ r < 1 , then

Proof. Due to the assumption ∥ ¯ x r -y i r,k ∥ ≤ C L 1 , we have

where we used Lemma 2 in ( i ) and Assumption 1(iv) in ( ii ) . Rearranging to isolate ∥∇ f (¯ x r ) ∥ gives the result.

Lemma 10. Suppose that Equation 9 holds, then for any i ∈ [ N ] and k ∈ [ I ] ,

Proof. Notice that, due to the clipping operation,

holds for any k ≤ I . Next we verify Equation 11 under the event A 0 . To see this, we first consider the case k = 1 ,

where we used the initialization G i 0 = ∇ F i (¯ x 0 ; ˜ ξ i ) and ∥ G 0 ∥ ≤ γ η under the event A 0 . This indicates that Equation 11 holds with k = 1 .

Now suppose that Equation 11 holds for some k with 1 ≤ k < I . Then we have

where ( i ) holds due to the initialization G 0 = ∇ F i (¯ x 0 ; ˜ ξ i ) ; ( ii ) follows from ∥ G 0 ∥ ≤ γ η ; and ( iii ) follows from Lemma 2. Note that the conditions of Lemma 2 are satisfied here by the inductive hypothesis together with Equation 9.

Notice that Ξ 0 = 0 , so by Lemma 9, we have

Therefore

where ( i ) holds due to Equation 9 and ( ii ) holds because of the inductive hypothesis that Equation 11 holds for k . This completes the induction over k under the event A 0 . This fact together with Equation 12 proves that Equation 11 holds for all k = 1 , . . . , I .

Lemma 11. Suppose that Equation 9 holds. If max k ∈ [ I ] ∥ x i q,k -¯ x q ∥ ≤ 2 I (2 ση + γ ) holds for any q ≤ r -1 , then we have

Proof. Under ¯ A r , Equation 13 trivially holds due to clipping. Next we verify Equation 13 under A r . First consider the case k = 1 .

Recall that y i r,k = x i q i r ,k from the definitions Equation 7 and Equation 8, so

where ( i ) holds due to inductive hypothesis for rounds q ∈ { 0 , . . . , r -1 } ; ( ii ) holds due to the event E such that r -q i r ≤ K ; and ( iii ) follows from Equation 9. The above bound on ∥ ¯ x r -y i r,k ∥ verifies Equation 14 and shows that the condition of Lemma 2 is satisfied. Then applying it to Equation 15 yields

where we used Equation 16 and Assumption 1 (iii) and (iv). Invoking Equation 9 and Equation 16 gives

Therefore the conditions of Lemma 9 are satisfied, and we can bound ∥∇ f (¯ x r ) ∥ as follows:

where we used Equation 9. Finally, plugging back into Equation 17 yields

where we used Equation 9 again. This proves Equation 13 for the base case k = 1 . Now suppose that Equation 13 holds for some k with 1 ≤ k < I . Then

where we used ∥ G r ∥ ≤ γ η under the event A r and Assumption 1 (iii). We can individually bound each of the two terms 1 A r ∥∇ f i ( x i r,k ) -∇ f i (¯ x r ) ∥ and 1 A r ∥∇ f i (¯ x r ) -∇ f i ( y i r,j ) ∥ . By the inductive hypothesis over k , together with Equation 9: ∥ x i r,k -¯ x r ∥ ≤ 2 I (2 ση + γ ) ≤ C L 1 . So we can apply Lemma 2 and obtain:

where we used Assumption 1 (iii) and (iv), Equation 18, and Equation 9. Similarly, by the inductive hypothesis over r ,

So we can apply Lemma 2:

where the last inequality follows from the assumption Equation 9.

Finally, plugging Equation 20 and Equation 21 into Equation 19 gives

where we used the inductive hypothesis over k and the inductive hypothesis over r . This fact together with 1 ¯ A r ∥ x i r,k +1 -¯ x r ∥ ≤ 2( k +1)(2 ση + γ ) (which trivially holds due to the clipping operator), we prove that Equation 13 holds for k +1 , completing the induction over k . Therefore Equation 13 holds for all k ∈ [ I ] .

Lemma 1 restated. Suppose that Equation 9 holds. Then for any r ≥ 0 we have

Proof. We will use induction to prove Equation 22 by a induction over r . And Equation 23 holds as a consequence of the inductive assumption due to Lemma 11. Invoking Lemma 10, we can verify the base case of the induction over r . Now we assume max k ∈ [ I ] ,i ∈ [ N ] ∥ x i q,k -¯ x q ∥ ≤ 2( K + 1) I (2 ση + γ ) holds for any q ≤ r -1 . Then according to Lemma 11, we can prove Equation 22 and Equation 23 immediately. It also completes the induction of Equation 22 over r .

Corollary 2. Suppose that Equation 9 holds. Then for any r ≥ 0 and i ∈ [ N ] ,

Proof. All proofs of this corollary are under the event E . Equation 24 follows directly from Equation 14 and the definition Ξ r = 1 NI ∑ N i =1 ∑ I k =1 ∥ y i r,k -¯ x r ∥ . Equation 25 is exactly Equation 18.

Due to Lemma 1, we can apply Lemma 2 such that

So, using Equation 24 and Equation 25,

where we used Equation 9. Equation 27 follows immediately from Equation 26 by

A.4 Proof of descent inequality

Lemma 12. Suppose Equation 9 holds. Then with probability 1 -15 δ , Algorithm 1 satisfies

where

Γ

1

=

AL

0

+

BL

1

κ

+4

BL

1

ρ

2

σ

+

γ

η

.

Proof. By Lemma 8, P ( E ) ≥ 1 -δ . The remainder of the proof will suppose the event E happens, and the results will hold with probability 1 -δ . By Lemma 1 and Equation 9, ∥ ¯ x r +1 -¯ x r ∥ ≤ C L 1 . Denote A r = {∥ G r ∥ ≤ γ/η } . Therefore we may apply Lemma 3 to obtain

(

)

Summing over r = 0 , . . . , R -1 yields

We proceed by bounding each of the four terms in Equation 28.

Bounding A 1 . Denote

Then we can write

It follows that

Since ∥ x i r,k -¯ x r ∥ ≤ C/L 1 due to Equation 13, we can apply Lemma 2 such that

where ( i ) holds due to Equation 13; ( ii ) follows from Assumption 1(iv); and ( iii ) follows from Equation 9. Plugging into Equation 29,

Denote each of the three last terms (those involving ϵ i r,k , ϵ (1) r , and ϵ (2) r , respectively) as B 1 , B 2 , B 3 . B 1 can be bounded using the martingale concentration bound in Lemma 5.

Denoting G r,k = σ ( F r , S r , { ξ i r,ℓ : i ∈ S r , ℓ ≤ k } ) , we have

so that { 1 A r 〈 ∇ f (¯ x r ) , 1 S ∑ i ∈S r ϵ i r,k 〉} r,k is a martingale difference sequence with respect to G r,k . Since 1 A r ∈ F r ∈ G r,k for any 0 ≤ k ≤ I , we have

where we used Cauchy-Schwarz, Jensen's inequality, and ∥ ϵ i r,k ∥ ≤ σ . Therefore we can apply Lemma 5 to obtain

with probability 1 -δ for any λ 1 > 0 .

The concentration bounds for B 2 and B 3 will rely on the concentration bound for sampling without replacement in Lemma 6. We start with B 2 . Notice that

so we must upper bound 1 A r |⟨∇ f (¯ x r ) , G i r ⟩| to apply Lemma 6. Under A r ,

where ( i ) holds due to Equation 25 and Equation 26; and ( ii ) holds since ρ ≥ 1 . Therefore

and the condition of Lemma 6 is satisfied. Notice that 1 A r ∈ F r because G r is constructed by the history gradients before round r . Therefore, applying Lemma 6 to Equation 32 gives 4

Let

Then we know E r [ Y r ] ≤ Y r -1 by Equation 34, which implies

Therefore for any λ 2 > 0 , we have

where the last line uses Markov's inequality. Repeating this argument for -1 A r ⟨∇ f (¯ x r ) , ϵ (1) r ⟩ , we can obtain

which means that

holds with probability

The same argument can be applied to B 3 . Notice that

so we must upper bound 1 A r |⟨∇ f (¯ x r ) , ∇ f i (¯ x r ) ⟩| to apply Lemma 6. Using Equation 25 in Corollary 2, we have

This matches the corresponding upper bound in Equation 33 from the bound of B 2 . We may therefore apply an identical argument as in the case of B 2 and obtain

with probability 1 -2 δ .

Combining Equation 31, Equation 35, and Equation 36 into Equation 30, yields that, with probability 1 -5 δ ,

where we chose

Bounding A 2 . From Lemma 4,

and

Therefore

and

where we used Equation 26. Finally

where we used the choice µ = 3 4 .

Bounding A 3 . To bound A 3 , we begin by bounding 1 A r ∥ ¯ x r +1 -¯ x r ∥ 2 .

where we denoted ˜ ϵ i r,k = ϵ i q i r ,k = ∇ F i ( y i r,k ; ξ i q i r ,k ) -∇ f i ( y i r,k ) . The first three terms on the RHS of Equation 38 can be bounded with similar arguments as those used previously:

where ( i ) holds due to Equation 13 and Assumption 1(iv); and ( ii ) follows from the condition on η in Equation 9. Due to Lemma 1, we can use the same argument to get

and

Plugging back into Equation 38,

We can bound the remaining terms using the concentration inequality from Lemma 7. We write the index of S r as { i 1 , ..., i S } For k = 0 , ..., I -1 and for j ∈ [ S ] , let X kS + j = ε i j r,k and H kS + j = σ ( F r , S r , { ξ i r,ℓ } 0 ≤ ℓ ≤ k -1 ,i ∈S r , { ξ i m r,k } m ≤ j ) for k ≥ 1 , where { ξ i r,ℓ } 0 ≤ ℓ ≤-1 ,i ∈S r = ∅ . For any t = kS + j , it holds that

where the second inequality holds due to the independence between different clients given F r . Then { X t } t ≤ IS is a martingale difference sequence with respect to {H t } t ≤ IS . Also, ∥ X t ∥ ≤ σ almost surely for all t ∈ [ IS ] and E t -1 [ ∥ X t ∥ ] ≤ σ 2 . Therefore, by Lemma 7, with probability 5 at least 1 -δ ,

Without loss of generality, we assume log(1 /δ ) ≥ 1 . We can apply the exact same argument to bound the remaining noise terms:

and

each with probability 1 -3 δ . Plugging back into Equation 39,

1 A r ∥ ¯ x r +1 -¯ x r ∥ 2 ≤ 1 A r ( 5 η 2 I 2 ∥∇ f (¯ x r ) ∥ 2 +192Γ 2 1 η 4 I 4 ( 2 σ + γ η ) 2 + 432 σ 2 η 2 I log 1 δ S ) , with probability 1 -9 δ . Also, using Equation 25, 1 A r ( AL 0 + BL 1 ∥∇ f (¯ x r ) ∥ ) ≤ 1 A r ( AL 0 +4 BL 1 ( 2 σ + γ η )) ≤ 1 A r Γ 1 .

Finally,

with probability 1 -9 δ .

Bounding A 4 : Due to normalization of the update g i r,k under ¯ A r , we have

Combining the respective bounds for A 1 , A 2 , A 3 , A 4 into Equation 28,

which is the desired result. Note that the bound on A 1 holds with probability 1 -5 δ and the bound on A 2 holds with probability 1 -9 δ , and we initially supposed the event E , which holds with probability 1 -δ . So the overall result holds with probability at least 1 -15 δ .

A.5 Proof of Theorem 1

and γ = ( 72 σ + AL 0 BL 1 ρ ) η , then Algorithm 1 satisfies 1 R ∑ R -1 r =0 ∥∇ f (¯ x r ) ∥ ≤ 35 ϵ with probability at least 1 -15 δ , as long as R ≥ 8∆ ϵ 2 ηI .

Proof. First, under our choice of η and γ ,

and

where the last line follows from B,ρ ≥ 1 . Therefore Equation 9 holds under our choice of η and γ , so the condition of Lemma 12 is satisfied. Denoting

Lemma 12 gives us

We will bound each of 1 ¯ A r U (¯ x r ) and 1 A r V (¯ x r ) under our choices of γ and η . To bound U (¯ x r ) , notice

and

So

where the last inequality holds due to the fact ϵ ≤ AL 0 16 BL 1 ρ ≤ 72 σ + AL 0 BL 1 ρ = γ η . To bound V (¯ x r ) , notice

since η ≤ 32Γ 1 I ( 74 σ + AL 0 BL 1 ρ ) = 32Γ 1 I ( 2 σ + γ η ) , and

where the last line uses ϵ ≤ AL 0 16 BL 1 ρ ≤ AL 0 16 BL 1 . Lastly,

So

where the last line came from the inequality x 2 ≥ 2 ax + a 2 with x = ∥∇ f (¯ x r ) ∥ and a = ϵ . Combining the bounds of U (¯ x r ) and V (¯ x r ) ,

Plugging this into Equation 40 yields

and by rearranging we have

where the last line comes from R ≥ 8∆ ϵ 2 ηI . Finally, η ≤ ∆ log 1 δ min { 1 12 σ 2 , 1 Γ 2 I } implies

Plugging into Equation 41 gives the result.

A.6 Proof of Corollary 1

Corollary 1 restated. If, under the setting of Theorem 1, we additionally choose

then with probability 1 -14 δ , Algorithm 1 will reach an ϵ -stationary point with iteration complexity

and so

Proof. From the condition on ϵ ,

From the condition on I ,

and

Therefore η = Sϵ 2 216 AL 0 σ 2 log 1 δ . Choosing R = 8∆ ϵ 2 ηI ,

B Deferred Proofs of the Lower Bound

Algorithm 2 Clipped Minibatch SGD

1: Initialize x 0

2: for r = 0 , 1 , . . . , R -1 do

3: Sample S r ⊂ [ N ] uniformly at random such that |S r | = S

4: g r = 1 SI ∑ i ∈S r ∑ I -1 k =0 ∇ F i ( x r , ξ i r,k )

5: Update x r +1 ← x r -min ( η, γ ∥ g r ∥ ) g r

6: end for

Our proof of Theorem 2 analyzes three separate functions satisfying Assumption 1, and concludes that, in order for clipped minibatch SGD to avoid divergence on the first two functions, one must choose parameters γ, η such that convergence on the third function is slow. These three functions are separately analyzed in Lemmas 13, 15, and 16. Each lemma assumes a particular value of the initial point x 0 : this assumption can be made without loss of generality, since for the initialization x 0 we may always translate each function to simulate the assumed initialization.

and γ η ≤ M . For any 0 < ϵ < M , there exists a problem instance { f i } N i =1 ∈ F ( L 0 , L 1 , M, κ, ρ, N ) such that, with probability at least δ , clipped minibatch SGD with parameters γ, η, S will generate iterates { x r } R r =0 with ∥∇ f ( x r ) ∥ > ϵ for all r .

The proof of Lemma 13 will require the following lemma.

Lemma 14 (Lemma 5.3 in [20]) . Let { X t } ∞ t =1 be a Markov chain over states { i } ∞ i =0 , such that 0 is an absorbing state, and the transition distribution elsewhere is as follows:

Define the absorb probabilities α i := P ( ∃ t > 0 : X t = 0 | X 0 = i ) , then:

Proof of Lemma 13. Define and

Also, define

and with global objective

Note that b > M . For this problem instance, we define the stochastic objective as F i ( x ; ξ ) = f i ( x ) , so that the true gradient of each local objective is returned by each gradient query of the algorithm.

By construction, each local objective f i and the global objective f is twice continuously differentiable, with gradient bounded by M . Also, each objective is ( L 0 , L 1 ) -smooth (in fact, they are L 0 -smooth), since for | x | ≤ v ,

and for | x | ≥ v , | f ′′ i ( x ) | = 0 . Finally, the collection of objectives satisfies the heterogeneity condition. For i ≤ Q ,

and for i ≥ Q +1 ,

where we used

together with the fact ϕ ≤ κ +( ρ -1) M κ +( ρ +1) M . Therefore { f i } N i =1 ∈ F ( L 0 , L 1 , M, κ, ρ, N ) .

Now, we analyze the behavior of clipped minibatch SGD on the instance { f i } N i =1 with initialization x 0 = v +1 . Let r 0 be the index of the first round (if one exists) in which x r < v . For each round r , define the event B r = {S r ⊂ [ Q ] } . To compute P ( B r ) , first notice

Therefore,

where we used the inequality ( 1 + a x ) x ≤ e a .

Consider the problem instance defined by

Now define the auxiliary sequence { y r } R r =0 as follows:

We claim that y r ≤ x r for all r ≤ r 0 . For any r < r 0 , if B r occurs then y r +1 -y r = x r +1 -x r , since in this case g r = -M and x r +1 -x r = γ , due to the fact that γ η ≤ M . If B r does not occur, then x r +1 -x r ≥ -γ , due to the clipping operation, so that y r +1 -y r ≤ x r +1 -x r . Therefore for all r ≤ r 0 ,

which proves our claim. Denoting D = {| f ′ ( x r ) | ≥ M for all r ≥ 1 } ,

where ( i ) comes from applying Lemma 14 to the sequence y r and ( ii ) uses ⌈ x 0 -v γ ⌉ ≥ 1 and 1 -P ( B r ) P ( B r ) ≤ 1 -δ .

The following Lemma uses the same function as in the first half of the proof of Theorem 2 in [8]. However, the sequence of iterates computed by clipped minibatch SGD (our setting) evolves differently than that of GD (their setting).

Lemma 15. Suppose γ η > M , M > L 0 L 1 , and η ≥ 2 L 1 M ( 1 + log L 1 L M 0 ) . For any 0 < ϵ < M , there exists a problem instance { f i } N i =1 ∈ F ( L 0 , L 1 , M, κ, ρ, N ) such that clipped minibatch SGD with parameters γ, η , and any S will generate iterates { x r } R r =0 with ∥∇ f ( x r ) ∥ > ϵ for all r .

Proof. Consider the following function,

and the problem instance f i = f for all i ∈ [ N ] , with the initialization x 0 = 1 L 1 ( 1 + log ML 1 L 0 ) . We define the stochastic objective for this problem instance as F i ( x ; ξ ) = f i ( x ) , so that the true gradient of each local objective is returned by each gradient query of the algorithm.

Note that f is bounded from below and ( L 0 , L 1 ) -smooth. Also, since all clients have the same objective and gradients are computed deterministically, this problem instance satisfies Assumption 1. Also, our setting of x 0 is consistent with the definition of M , since

First, define w := 1 L 1 ( 1 + log ( γ η L 1 L 0 )) . By the condition γ η > M , we know w > x 0 , and by construction f ′ ( w ) = γ η . Since f ′ is increasing on [0 , ∞ ) , and decreasing on ( -∞ , 0] , we know | f ′ ( w ) | ≥ γ η if and only if | x | ≥ w . This means that the clipping operation will be performed when | x | ≥ w , and will not be performed for | x | < w .

Now we analyze the behavior of clipped minibatch SGD on this problem instance. Suppose that | x 0 | ≤ | x r | ≤ w for some r ≥ 0 . Then

where the first inequality comes from the definition of x 0 and the condition η ≥ 2 L 1 M ( 1 + log L 1 L M 0 ) , and the second inequality comes from the fact that ϕ ( x ) := 2 L 1 x L 0 exp( L 1 x -1) is decreasing on [ 1 L 1 , ∞ ) . Therefore

̸

Since sign ( f ′ ( x r )) = sign ( x r ) , this implies | x r +1 | = | x r -ηf ′ ( x r ) | ≥ 2 | x r | - | x r | = | x r | . This shows that the sequence of iterates { x r } is non-decreasing in absolute value until (if ever) the absolute value exceeds w . If | x r | ≤ w for all r , then we are done, since in this case we have | x r | ≥ | x 0 | and therefore |∇ f ( x r ) | ≥ |∇ f ( x 0 ) | = M > ϵ for all r .

Otherwise, let ¯ r be the first index r for which | x r | > w . Without loss of generality, assume x ¯ r > 0 (for x ¯ r < 0 , the same argument applies with signs reversed). Since | x ¯ r -1 | ≤ w ,

So

Therefore | f ′ ( x ¯ r +1 ) | ≥ | f ′ ( -w ) | = | f ′ ( w ) | = γ η , and so x ¯ r +2 = x ¯ r +1 + γ = x ¯ r . We can then show by induction that x ¯ r +2 n = x ¯ r and x ¯ r +2 n +1 = x ¯ r +1 for all n ≥ 0 . Since | f ′ ( x ¯ r ) | , | f ′ ( x ¯ r +1 ) | ≥ | f ′ ( w ) | > | f ′ ( x 0 ) | > ϵ , we have | f ′ ( x r ) | > ϵ for all r .

The following lemma is nearly identical to the second half of the proof of Theorem 2 in [8]. We include it here for the sake of completeness.

Lemma 16. Suppose γ η > M and η < 2 L 1 M ( 1 + log L 1 L M 0 ) . For any 0 < ϵ < M , there exists a problem instance { f i } N i =1 ∈ F ( L 0 , L 1 , M, κ, ρ, N ) such that clipped minibatch SGD with parameters γ, η , and any S requires at least

rounds in order to find an ϵ -stationary point.

Proof. Consider the following function,

and the problem instance f i = f for all i ∈ [ N ] , with the initialization x 0 = 3 ϵ 2 L 0 + d . We define the stochastic objective for this problem instance as F i ( x, ξ ) = f i ( x ) , so that the true gradient of each local objective is returned by each gradient query of the algorithm.

Note that f is bounded from below and ( L 0 , L 1 ) -smooth. Also, since all clients have the same objective and gradients are computed deterministically, this problem instance satisfies Assumption 1. Also, our setting of x 0 is consistent with the definition of M , since | f ′ ( x 0 ) | = ϵ < M .

Now we analyze the behavior of clipped minibatch SGD on this problem instance. Since this problem instance is homogeneous with deterministic gradients, we have ∥ g r ∥ ≤ ϵ < M < γ η , so clipped minibatch SGD will always perform unnormalized updates. Therefore, for x r ≥ 3 ϵ 2 L 0 ,

Therefore

as long as 2 ϵr L 1 M ( 1 + log L 1 L M 0 ) ≤ d , or r ≥ dL 1 M 2 ϵ ( 1+log L 1 L M 0 ) . Notice that

so

Plugging into the above bound on r tells us that x r > x 0 as long as r ≥ L 1 M ( f ( x 0 ) -f ∗ -15 ϵ 2 16 L 0 ) 2 ϵ 2 ( 1+log L 1 L M 0 ) . The result follows from the fact that | f ′ ( x r ) | = ϵ for all x r > x 0 .

We can now combine Lemmas 13, 15, and 16 to prove Theorem 2.

Proof of Theorem 2. Let γ, η > 0 be given. We will prove the requirement on R for three cases of γ, η . In the case γ η ≤ M , then Lemma 13 demonstrates that clipped minibatch SGD cannot guarantee to find an ϵ -stationary point with probability greater than 1 -δ for any R . Second, if γ η > M and η ≥ 2 L 1 M ( 1 + log L 1 L M 0 ) , then Lemma 15 demonstrates that clipped minibatch SGD can fail to find an ϵ -stationary point for any R . In the remaining case, i.e. γ η > M and η < 2 L 1 M ( 1 + log L 1 L M 0 ) , Lemma 16 demonstrates that for finding an ϵ -stationary point, the number of iterations clipped minibatch SGD requires is at least

C Additional Experimental Results

C.1 Hyperparameter Information

Learning Rate and Clipping Threshold For each dataset, we tune the hyperparameters γ and η with separate grid searches. Specifically, we first tune the clipping threshold γ η with grid search, then we tune the learning rate η with grid search. During the grid search for η , γ is chosen so that the clipping threshold γ η is equal to the tuned value from the search for γ η . Due to computational constraints, we do not perform this search for every algorithm in every setting. Instead, we follow the tuning process of [9]: we use the above procedure to tune η and γ for CELGC [32] separately for each S ∈ { 2 , 4 , 6 , 8 } , while fixing s = 30% (SNLI) or s = 10% (Sent140). We then reuse the tuned values of γ and η for all other algorithms and settings sharing the same dataset and S . Our theory suggests that the best learning rate η depends on the number of participating clients S , and we follow this guidance by tuning the parameters separately for each value of S . The search ranges and tuned values for each dataset are shown in Table 3.

Network Architecture For both datasets, the network is composed of a one-layer bidirectional RNN encoder followed by a three-layer fully connected classifier. The encoder has hidden size 2048 and max pooling, and the decoder has hidden size 512 with tanh activations. Input sentences are encoded as sequences of GloVe vectors. For SNLI, where each input is a pair of sentences, each sentence is passed through the encoder separately, and the resulting representations are concatenated and used as input for the encoder.

C.2 Learning Curves

Figures 4 and 5 contain learning curves (by communication rounds) of training loss and testing accuracy for all settings, for SNLI and Sentiment140 respectively. The experiments show that our proposed algorithm EPISODE++ significantly outperform all other baselines in various settings on these two datasets, including different client participation ratio and different heterogeneity level.

C.3 Training with Homogeneous Data

Here we include learning curves for an additional experiment that uses homogeneous data across clients, i.e. s = 0% . In this setting, we use N = 8 and S = 4 , and all other settings are the same as described in Section 6. We compare the six algorithms described in Section 6. Results are shown in Figure 6. The results are consistent with the experiments that use heterogeneous data: EPISODE++ outperforms all other algorithms in terms of training loss and testing accuracy.

0

2000

4000

Rounds

0.4

0.6

0.8

1.0

1.2

Train Loss

SNLI (S=2, s=30%)

0

2000

4000

Rounds

0.4

0.5

0.6

0.7

0.8

Test Accuracy

0

2000

4000

Rounds

0.4

0.6

0.8

1.0

Train Loss

SNLI (S=4, s=30%)

0

2000

4000

Rounds

0.4

0.5

0.6

0.7

0.8

Test Accuracy

0

2000

4000

Rounds

0.4

0.6

0.8

1.0

Train Loss

SNLI (S=6, s=30%)

0

2000

4000

Rounds

0.4

0.5

0.6

0.7

0.8

Test Accuracy

0

2000

4000

Rounds

0.4

0.6

0.8

1.0

Train Loss

SNLI (S=8, s=30%)

0

2000

4000

Rounds

0.4

0.5

0.6

0.7

0.8

Test Accuracy

0

2000

4000

Rounds

0.4

0.6

0.8

1.0

Train Loss

SNLI (S=4, s=50%)

0

2000

4000

Rounds

0.4

0.5

0.6

0.7

0.8

Test Accuracy

0

2000

4000

Rounds

0.4

0.6

0.8

1.0

1.2

Train Loss

SNLI (S=4, s=10%)

0

2000

4000

Rounds

0.4

0.5

0.6

0.7

0.8

Test Accuracy

EPISODE++

CELGC

NaiveParallelClip

Clipped Minibatch SGD

SCAFFOLDClip

EPISODE

0

1000

2000

Rounds

0.4

0.6

0.8

Train Loss

Sentiment140 (S=2, s=30%)

0

1000

2000

Rounds

0.65

0.70

0.75

Test Accuracy

0

1000

2000

Rounds

0.4

0.6

0.8

Train Loss

Sentiment140 (S=4, s=30%)

0

1000

2000

Rounds

0.65

0.70

0.75

Test Accuracy

0

1000

2000

Rounds

0.4

0.6

0.8

Train Loss

Sentiment140 (S=6, s=30%)

0

1000

2000

Rounds

0.65

0.70

0.75

Test Accuracy

0

1000

2000

Rounds

0.3

0.4

0.5

0.6

0.7

0.8

Train Loss

Sentiment140 (S=8, s=30%)

0

1000

2000

Rounds

0.65

0.70

0.75

Test Accuracy

0

1000

2000

Rounds

0.3

0.4

0.5

0.6

0.7

0.8

Train Loss

Sentiment140 (S=4, s=20%)

0

1000

2000

Rounds

0.65

0.70

0.75

Test Accuracy

0

1000

2000

Rounds

0.3

0.4

0.5

0.6

0.7

0.8

Train Loss

Sentiment140 (S=4, s=0%)

0

1000

2000

Rounds

0.65

0.70

0.75

Test Accuracy

EPISODE++

CELGC

NaiveParallelClip

Clipped Minibatch SGD

SCAFFOLDClip

EPISODE

0

2000

4000

Rounds

0.4

0.6

0.8

1.0

Train Loss

SNLI

0

2000

4000

Rounds

0.4

0.6

0.8

Test Accuracy

0

1000

2000

Rounds

0.4

0.6

Train Loss

Sentiment140

0

1000

2000

Rounds

0.65

0.70

0.75

Test Accuracy

EPISODE++

CELGC

NaiveParallelClip

Clipped Minibatch SGD

SCAFFOLDClip

EPISODE

Peer reviews:
comments: SUMMARY:
This manuscript introduces a new federated learning algorithm, EPISODE++, to accelerate the convergence speed of federated learning and provides theoretical proofs for the upper bound of convergence speed.

STRENGTHS:
The method proposed in this manuscript improves the EPISODE algorithm by addressing the issue of client heterogeneity while allowing partial client participation. The manuscript also provides theoretical proofs for the convergence speed.

WEAKNESSES:
The experimental section does not explain how the hyperparameters for each baseline method were selected, and it does not discuss how the performance of the method is affected when the number of clients reaches hundreds.

QUESTIONS:
1. In Figure 1 (a) and (b), the left two figures show that NaiveParallelClip has lower training loss than other methods when the number of clients is 8, but its testing accuracy is lower than that of EPISODE++. How can this phenomenon be explained?

2. How does the performance of the method change when there are more clients (resulting in larger variances for $G_{r+1}$ and $G_{r+1}^i$)? Does it outperform classical methods like fedprox?

3. How were the hyperparameters (learning rate, training steps) chosen for each method?

4. How does the proposed method perform when the client data are i.i.d.?

LIMITATIONS:
The method proposed in this manuscript only slightly modifies the EPISODE method, and the experimental section is not comprehensive.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-24

comments: SUMMARY:
The authors propose a federated learning algorithm that can work under (L0,L1)-smooth function. Different from the previous work, the authors consider the partial-participant setting, modifying the previous algorithm, showing the convergence of the new algorithm, and giving a lower bound of the communication iterations under this setting. The experiments show that the proposed algorithm performs well.

STRENGTHS:
1. The authors propose a new algorithm that overcomes the bias introduced by client heterogeneity and client sampling.

2. The authors prove the convergence of the proposed algorithm and a lower bound of communication iterations under this setting.

3. The experimental results show the the proposed algorithm performs much better than previous work.

WEAKNESSES:
1. The key motivation is not clear to me. For me, it is hard to justify that with a uniform sampling strategy on clients, why the bias will occur? And it is hard to see how the proposed algorithm can fix the introduced bias.

QUESTIONS:
See weakness.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a novel algorithm for non-convex federated learning that addresses the challenges of relaxed smoothness, client heterogeneity, and client subsampling. The authors begin by discussing the limitations of existing algorithms such as SCAFFOLD and EPISODE in handling these challenges simultaneously. They then introduce their algorithm, EPISODE++, which is designed to overcome these limitations. The algorithm is initialized with a set of parameters and then iteratively updated through a series of communication rounds. In each round, a subset of clients is selected randomly, and each client performs local updates based on its own data. The algorithm also includes a gradient clipping step to control the magnitude of the gradient updates.\n\nThe authors provide a detailed theoretical analysis of the convergence properties of EPISODE++. They show that the algorithm can achieve a linear speedup over the standard federated averaging algorithm (FedAvg) under certain conditions. They also demonstrate that EPISODE++ can significantly outperform clipped minibatch SGD, another popular algorithm for non-convex optimization, in terms of the number of iterations required to find an epsilon-stationary point. The paper concludes with an extensive set of experiments that validate the theoretical findings. The authors show that EPISODE++ outperforms other state-of-the-art algorithms on a variety of benchmark datasets and neural network architectures, including LSTMs and Transformers.\n\nSTRENGTHS:\n1. The authors provide a detailed theoretical analysis of the convergence properties of EPISODE++. They show that under certain conditions, the algorithm can achieve a linear speedup over the standard federated averaging algorithm (FedAvg), and significantly outperform clipped minibatch SGD in terms of the number of iterations required to find an epsilon-stationary point. This rigorous analysis strengthens the credibility of the proposed algorithm.\n\n2. The paper includes an extensive set of experiments that validate the theoretical findings. The authors demonstrate that EPISODE++ outperforms other state-of-the-art algorithms on a variety of benchmark datasets and neural network architectures, including LSTMs and Transformers. There are many other FL papers that only conduct very toy experiments that are not convincing enough. This empirical evidence provides strong support for the effectiveness of the proposed algorithm.\n\nWEAKNESSES:\nWhile the paper presents a novel algorithm and provides a detailed theoretical analysis, the scope of the study appears to be limited to non-convex federated learning. It would be interesting to see how EPISODE++ performs in other contexts or problem domains.\n\nQUESTIONS:\nThe paper could provide more information about the implementation details of the EPISODE++ algorithm. For instance, how are the parameters of the algorithm chosen in practice? Are there any specific strategies or heuristics for setting these parameters? Providing such information could help other researchers to replicate the results and apply the algorithm to their own problems.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

comments: SUMMARY:
The paper investigates Federated Learning (FL) under client subsampling and data heterogeneity, focusing on functions with potentially unbounded smoothness, and introduces the proposed algorithm to address the problem, EPISODE++. EPISODE++ has demonstrated benefits including linear speedup with client numbers, reduced communication rounds, and resilience to data heterogeneity. The authors provide theoretical convergence analysis and experimental results validate the effectiveness of their method.

STRENGTHS:
1. The paper is well-written and easy to follow.
2. Provides novel techniques when proofing both upper and lower bounds.
3. Theoretically demonstrates the benefit of proposed methods EPISODE++, and also shows the shortage of existing clipped minibatch SGD.
4. The proposed algorithm achieves significant improvement over existing methods.

WEAKNESSES:
1. The experiments with N=8 may not fully reflect the actual performance with large-scale FL, for example, FedAvg uses 100 clients in their experiments, where more heterogeneity and client shift might result in different behavior. I would be happy to see if there's more benefit of EPOSIDE++
2. Missing error bar and not indicates the number of repetition runs in experiments.

QUESTIONS:
See weakness

LIMITATIONS:
The authors have mentioned limitations of their work.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-11