Reducing Communication in Federated Learning with a Novel Single-Loop Variance Reduction Method
Abstract
In Federated Learning (FL), inter-client heterogeneity causes two types of errors: (i) client drift error which is induced by multiple local updates, (ii) client sampling error due to partial participation of clients at each communication.
While several solutions have been offered to the former one, there is still much room of improvement on the latter one.
We provide a fundamental solution to this client sampling error.
The key is a novel single-loop variance reduction algorithm, SLEDGE (Single-Loop mEthoD for Gradient Estimator), which does not require periodic computation of full gradient but achieves near-optimal gradient complexity in the nonconvex finite-sum setting.
While sampling a small number of clients at each communication round, the proposed FL algorithm, FLEDGE, requires provably fewer or at least equivalent communication rounds compared to any existing method, for finding first and even second-order stationary points in the general nonconvex setting, and under the PL condition.
Moreover, under less Hessianheterogeneity between clients, the required number of communication rounds approaches to ˜Θ(1) .
1. Introduction
Federated learning (FL) is a paradigm of distributed learning, where each local client has access to local dataset to train a local model, and periodically local model parameters are exchanged to update a global model in the server [22, 31, 43].
Avoiding share of local data itself, FL provides privacy protection and encourages the usage of distributed Big Data [17, 24].
Since clients such as smartphones and organizations are physically separated, the main bottlenecks in FL are synchronization and communication between clients and the server.
To reduce the number of communication rounds, local update has been adopted [25, 29, 31].
Local update means that, between communication rounds, the parameters of a local model are updated several times inside each client using only its local data, and then aggregated in the server at the next communication round.
Local update causes client drift error, but if that is carefully corrected, local update can provably reduce communication rounds, especially when clients have less heterogeneity [20, 32], which had experimentally been observed [19, 31, 49].
On the other hand, client sampling, meaning that not all but only a part of clients are sampled to participate in each communication, is also widely used to reduce the communication complexity (the total number of parameters communicated).
Due to huge parameter size, a large number of clients, low bandwidth communication, and increase of communication cost for secure computation, which are characteristics of FL, client sampling has become indispensable in designing FL algorithms [17].
However, while a mass of literature has established treatments for client drift error [19, 20, 32], client sampling error has not effectively been controlled until recently.
In fact, the convergence rate of FedAvg [31], one of the most famous FL algorithms, can be dominated by client sampling error rather than client drift error [14], and we can see that this is also true for other FL algorithms.
Especially, in order to reduce communication rounds by taking advantage of less client heterogeneity with local update, client sampling is not allowed [32] or increasingly larger sampling size is required as heterogeneity gets smaller [20], due to the client sampling error.
Recently, FedVarp [14] tackled client sampling error applying a variance reduced method of
SAGA [9, 40] to FedAvg [31].
Variance reduction is a technique to construct a gradient estimator with a smaller variance than vanilla SGD by recursively utilizing minibatch gradients at previously obtained anchor points [9, 16, 41].
Although we agree that applying variance reduction should be a right direction, a large part of the problem still remains.
In fact, their algorithm requires relatively large client sample size of O ( P 23 ) to the total number of clients P , and cannot take advantage of heterogeneity, resulting in sub-optimal communication rounds and complexity compared to the state-of-the-art FL methods [20, 32].
We considered that this is because existing variance reduction methods are not suitable for applying FL methods.
In general, variance reduction methods are measured only in terms of gradient complexity.
However, when applying them to FL methods, we additionally want single-loop structure (that is, not to require periodic full or large minibatch gradient since this leads to impractical full client participation in FL), and fewer gradient complexity under less heterogeneity.
From these perspectives, existing methods are not satisfactory.
Indeed, SAGA satisfies single-loop structure, but not the others.
We consider the limitations of the aforementioned work of Jhunjhunwala et al.
[14] came from this point.
As for other variance reduction methods, SARAH [35], SPIDER [10], and NestedSVRG [50] requires periodic full gradient.
STORM [7] is sub-optimal in the nonconvex finite-sum setting, and its application to FL, MimeMVR [20], requires larger sampling size as heterogeneity gets smaller as mentioned; ZeroSARAH [28] cannot benefit from less heterogeneity, and its distributed version does not considers local update.
See also Appendix A.1.
Thus, we must design a novel variance reduction algorithm for FL.
1.1. Contributions
We consider FL problems with finite clients, and so variance reduction in the finite-sum setting.
First, we developed a novel single-loop variance reduction method called SLEDGE (SingleLoop mEthoD for Gradient Estimator) for the nonconvex finite-sum problems.
SLEDGE does not require periodic computation of full gradients, and satisfies the followings: (i) nearly optimal gradient complexity of ˜ O ( √ nε 2 ) for finding ε -first-order stationary points with data size n , (ii) second-order optimality as the first such single-loop algorithm, (iii) exponential convergence under the PolyakŁojasiewicz (PL) condition, and (iv) fewer complexity under the less heterogeneity assumption.
Next, we combined SLEDGE with local updates into an efficient federated algorithm, FLEDGE.
FLEDGE appropriately controls client sampling error, and achieves the followings, with the interclient heterogeneity ζ and the number of local updates K : (i) For first-order stationary points, the number of required communication rounds is ˜ O
( 1 Kε 2 + ζε 2 ) , when the number of sampled clients at each step p is larger than √ P .
(ii)
Adding small perturbation, FLEDGE also can find SOSPs.
(iii)
Under the µ -PL condition, FLEDGE exhibits exponential convergence, and the number of communication rounds depends on µ only through LKµ + ζµ .
For finding first and second order stationary points, our rates are smaller than or equivalent to those of all existing FL methods across all range of the inter-client heterogeneity ζ .
As ζ → 0 , these rates approach to ˜ O
(1) under taking p ≃ √ P , while previous methods require full client participation [32, 33] or increasingly larger client sample size [20].
For PL (and even strongly convex) functions, FLEDGE breaks the dependency on the condition number when ζ ≪ L and K
≫ 1 , for the first time.
For detailed comparison of existing
FL methods and discussion on required local budget, see Appendix A.2.
2. SLEDGE: Single-Loop Method for Gradient Estimator
We formalize the finite-sum problems as follows:
[Equation]
Our goal is to find a solution x that is an ε -first-order stationary point (i.e., ∥∇ f ( x ) ∥ ≤ ε ), and an ( ε, δ ) -second-order stationary point (SOSP; i.e., ∥∇ f ( x ) ∥ ≤ ε and λ min ( ∇ 2 f ( x )) ≥ − δ ).
Throughout the section, we assume that f i is L -smooth, and f is bounded: f ( x 0 ) − f ∗ =:
∆ < ∞ with f ∗ = inf x f ( x ) .
Moreover, we assume Hessian-heterogeneity: ∥∇
f i ( x ) − ∇ 2 f j ( x ) ∥ ≤ ζ (note: ζ ≤ 2 L holds).
For Option I, we suppose ∥∇ f i ( x 0 ) −∇
[Equation]
( x [Equation] ∥ ≤ σ [Equation]
For finding SOSPs, let ∥∇
2 f i ( x ) − ∇ 2 f i ( y ) ∥
≤ ρ ∥ x − y ∥ hold.
The µ -PL condition means 2 µ ( f ( x ) − f ∗ ) ≤ ∥∇ f ( x ) ∥ 2 .
2.1. Algorithm description
We introduce our proposed method SLEDGE for the problem (1).
Note that v t , v ti , and ˜ y ti in parentheses are auxiliary variables with which 1 n (cid:80) ni
=
y ti is updated in O ( b ) time, utilizing the following equation: [-1 n (cid:80) ni =
-][Equation]
The small perturbation ξ t is necessary to escape from saddle points and to find SOSPs, see Appendix D.3.
[-1: Option I: Randomly sample b data I 0 and y 
 i ← 1 b (cid:80) j ∈ I 0 ∇ f j ( x 0 ) for i ∈-]Algorithm
1: Option 2: Option[-II: y 
-]
4: x t ←
[-i ( x 0 ) for i ∈ I 
 = I ( v 0 ← 0 and ˜ y 
-][Equation]
[-i ← y 
 i , v 
 i ← 0 for i ∈ I ) 3: For t = 1 to T do 4: x t ← x t − 1 − ηn-] (cid:80) ni
=1 y t − 1 i + ξ t ( ξ t follows the uniform distribution on the Euclidean ball in R d with radius r ) 5:
Randomly sample b data I t 6: y [Equation] ← (cid:40) ∇ f [Equation]
x t ) for
i ∈
I t [Equation] (cid:80)
[Equation] ∇ f [Equation] x
[Equation]
[Equation]
f [Equation] ( x t
[Equation]
+ y t − 1 i for i / ∈ I t ( v t ← 1 b (cid:80) j [Table]
2.2. Convergence analysis
SLEDGE is designed so that it inherits the best points of SAGA [9, 40] and SARAH [34, 35].
It removes periodic full gradient using stored past gradients as SAGA.
Since SAGA’s gradient com- plexity is suboptimal, we also import SARAH’s recursive update to achieve the near-optimal rate.
The noise ξ t is to guarantee the second-order optimality.
The key observation in the analysis is that, for Option II, the discrepancy between the estimator and the true gradient is decomposed as [Equation] 
[Equation]
Here, ˜ I ts =
[ n ] \ (cid:83) tτ =
s I t is the set of indexes not sampled between s and t .
Conditioning on | ˜
I ts | , the inside of the large parentheses for each s is roughly ζ (cid:113) | ˜
I ts | b ∥ x s − x s − 1 ∥ .
The main technical difficulty is that the inside of the large parentheses is not independent each other since | ˜ I ts | depends on all I s , . . .
, I t .
However, the correlation can be shown to be sufficiently weak.
Now we state the theoretical guarantee for SLEDGE.
For the formal version, see Appendix D.
Theorem 1 We take a step size η and a scale of noise r appropriately, and let µ ∈ (0 , 1) .
Then, with probability 1 − ν , SLEDGE finds ε -first-order stationary points using
[Equation]
stochastic gradients.
For finding SOSPs with probability 1 − ν , SLEDGE requires b ≳ √ n + ζ 2 δ
and ˜ O (cid:18)(cid:0)
[Equation]
( Option II ) stochastic gradients.
Under the µ -PL condition, SLEDGE uses ˜ O (cid:18)(cid:18)
Lb
[Equation]
II
[Equation]
stochastic gradients for finding ε -solutions with f ( x t ) − f ∗ ≤ ε , with probability 1 − ν .
Without periodic computation of full gradient, SLEDGE achieves nearly optimal gradient complexity for first-order stationary points, and can take advantage of less heterogeneity ζ .
We show that the rate of ˜ O ( n + ∆( ζ √ n ∧ Lb ) ε) (option II) is near-optimal, see Appendix F. Moreover, SLEDGE can finds SOSPs as the first single-loop algorithm, and exponentially converge under the PL condition.
3. FLEDGE: Federated Learning Method with Gradient Estimator
Indexing (finite) clients and data by i and j , we consider the following FL problem:
[Equation]
In addition to the previous assumptions on f and f i , we additionally assume that each f i,j is Lsmooth.
We bound the intra-client variance as ∥∇ f i,j ( x ) − ∇ f i ( x ) ∥ ≤ σ .
For finding SOSPs, we suppose that ∥∇ 2 f i,j ( x ) − ∇ 2 f i,j ( y )
∥ ≤ ρ ∥ x − y ∥ for all i, j and x, y .
We combine SLEDGE with local update into a novel FL method, FLEDGE.
In Lines 4-9, we use the SARAH-type estimator to control the error from local minibatch sampling.
Then, we construct the estimator of the global gradient ∇ f ( x t ) using the SLEDGE estimator.
Here, we only introduce FLEDGE with Option II and for the case of ζ > δ .
For others, see Appendix E.
1: for i ∈ I
= I in parallel do 2:
Randomly select minibatch J 0 i with size Kb and let y
i ← 1 bK (cid:80) j ∈ J
i ∇ f i,j ( x 0 ) 3: for t = 1 to T do 4:
Randomly sample one client i t and Send 1 P (cid:80) Pi =1 y t − 1 i and x t − 1 from the server to i t 5: x t, 0 ← x t − 1 , z t, 0 ←6: for k
= 1 to K do 7: x t,k ← x t,k − 1 − η ( 1 P (cid:80) Pi =1 y t − 1 i + z t,k − 1 ) + ξ t,k ( ξ t follows the uniform distribution on the Euclidean ball in R d with radius r ) 8:
Randomly select minibatch J t,ki t with size b 9: z t,k ← z t,k − 1 + 1 b (cid:80) j ∈ J t,kit ( ∇ f i t ,j ( x t,k ) − ∇ f i t ,j ( x t,k − 1 ))
Randomly select p clients I t , send x t,K from i t to I t and let x t ← x t,K 11: for i ∈
I t in parallel do 12:
Randomly select minibatch J t i with size Kb 13:
y ti ← 1 bK (cid:80) j ∈ J ti ∇ f i,j ( x t ) and ∆ y ti ← 1 bK (cid:80) j ∈ J ti ( ∇ f i,j ( x t ) − ∇ f i,j ( x t − 1 )) 14:
Send { ( y ti , ∆ y ti ) }
i ∈
I t from I t to the server 15: y ti ← y t − 1 i + 1 p (cid:80) i ∈ I t ∆ y ti for i / ∈
[Table]
Theorem 2 We take an appropriate step size η and a scale of noise r and let µ ∈ (0 , 1) .
Let b ≳
σ 2 PKε 2 .
Then, with probability 1 − ν , FLEDGE finds ε -first-order stationary points using
[Equation]
communication rounds.
For finding SOSPs, FLEDGE requires p ≳ √ P + ζ 2 δ 2 , b ≥ K , and
[Equation]
communication rounds, with probability 1 − ν .
Under the µ -PL condition, FLEDGE uses
[Equation]
communication rounds for finding ε -first-order stationary points, with probability 1 − ν .
Setting client sample size p ≥ √ P and local minibatch size b ≥ K , FLEDGE requires only ˜ O
( ζKε 2 + ζε 2 ) communication rounds for first-order stationary points.
Compared to FedVarp [14], ours require smaller client sample size and can take advantage of less heterogeneity.
What is more, as far as assumptions on the local minibatch size are satisfied, this rate is better than or equivalent to any existing federated learning algorithm and approaches ˜Θ(1) when ζ → 0 and K → ∞ .
In terms of communication complexity, ˜ O ( P + ζ √ P ε 2 ) is near-optimal in a sense that it almost matches the lower bound of gradient complexity of the finite-sum case.
We will show this in Appendix F.
Moreover, FLEDGE guarantees second-order optimality.
To the best of our knowledge, the only such method is BVR-P-LSGD [33] [Figure] but that synchronizes all clients at every round.
Furthermore, the algorithm requires ˜ O ( ζµ log ∆ ε ) communication rounds under µ -PL condition. Thus [Figure]
even if the condition number Lµ is bad [Figure] the required communication rounds goes to ˜ O
( Pp ) when ζ → 0 .
On the other hand, even in the strongly-convex case and without client sampling [Figure] all existing algorithms require Ω( Lµ log 1 ε ) rounds regardless of the Hessian heterogeneity ζ
[Figure] [Figure]
4. Experiments
Finally, we validate our theories on both SLEDGE and FLEDGE by numerical experiments.
We considered classification of the capital letters using EMNIST dataset [6] for both experiments.
Escaping saddle points with SLEDGE
For the finite-sum problem (1), we prepared each f i by sampling 100 data from one class, employing a four-layer neural network as the training model
[Figure]
and then defining the average of the cross-entropy loss over the data as f i .
We repeated this five times for each class, and thus n = 130 .
We compared SLEDGE and its perturbed version with SARAH [34, 37], SSRGD [26], and ZeroSARAH [28], in terms of the test accuracy.
We set b = 12 , the inner-loop length of SARAH and SSRGD to 10 , and λ = bn ≒ 0 .
092 for ZeroSARAH.
For SSRGD and SLEDGE, we added perturbation of δ = 0 .
We tuned the learning rate for each algorithm individually.
The experiment was repeated with ten different random seeds for each method.
Figure 1 shows the result.
We can observe that (i) SLEDGE and ZeroSARAH require fewer gradient evaluations than SARAH to achieve the same test accuracy, owing to avoidance of periodic full gradient.
Similarly, SLEDGE with small noise is faster than SSRGD.
(ii) Adding small noise helps stable convergence; Although SLEDGE with δ = 0 does not necessarily yield a monotonic increase in the accuracy (see around 10000 - 15000 gradient evaluations), SLEDGE with small noise perturbation makes the accuracy increase almost monotone.
Faster Convergence with FLEDGE
For the federated learning problem (2), we again consider the classification of the capital letters, where each f i consists of 90% data from one class and 10% data from the other classes.
This makes each f i a little less heterogeneous.
We used two-layer neural networks as a training model.
We compared FLEDGE with FedAvg [31], SCAFFOLD [19], MimeMVR [20], and BVR-L-SGD [32].
For each algorithm, we employed P = 104 as the total number of clients and p = 10 as the number of the clients used at each communication (except for BVR-L-SGD, which requires P = p = 104 ).
Then, we set b = 16 and K = 10 .
We tuned the learning rate for each algorithm individually.
The experiment was repeated with five different random seeds for each method
Figure 2 (left) shows that FLEDGE achieves the highest test accuracy with fewer communication, compared to FedAvg, SCAFFOLD and MimeMVR.
In Figure 2 (right), FLEDGE achieves the small gradient norm ∥∇ f ( x t ) ∥ and the linear convergence at the neighborhood of solutions.
Moreover, we observe that FLEDGE performs similarly to BVR-L-SGD, which almost can be seen as a special case of FLEDGE with p = P .
This means that FLEDGE can appropriately correct the errors from sampling of the clients and is about ten times more efficient than BVR-L-SGD in terms of communication complexity by allowing sampling of the clients
Introduction
Contributions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
SLEDGE: Single-Loop Method for Gradient Estimator 2.
Algorithm description . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Convergence analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
FLEDGE:
Federated Learning Method with Gradient Estimator
Experiments
A Further Literature review
A.1 Variance reduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.2 Federated learining . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B Additional experiments
B.1 Details of the experiment for Figure 1 . . . . . . . . . . . . . . . . . . . . . . . .
B.2 Details of the experiment for Figure 2 . . . . . . . . . . . . . . . . . . . . . . . .
B.3 Additional experiments for SLEDGE . . . . . . . . . . . . . . . . . . . . . . . . .
B.4 Additional experiments for FLEDGE . . . . . . . . . . . . . . . . . . . . . . . . .
B.5 Computing infrastructures . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C Assumptions and Tools
C.
Assumptions
C.2 Concentration inequalities . .
C.3 Proof of Proposition 6 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.4 Linear Algebraic Tool . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D Missing Proofs for SLEDGE
D.1 Intuitition behind SLEDGE . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.
Finding First-order Stationary Points (Proof of Theorem 10) . . . . . . . . . . . .
D.3 Finding Second-order Stationary Points (Proof of Theorem 11) . . . . . . . . . . .
D.4 Convergence under PL condition (proof of Theorem 12) . . . . . . . . . . . . . .
E Missing Statements and Proofs for FLEDGE
E.
Finding First-order Stationary Points (Proof of Theorem 22) . . . . . . . . . . . .
E.
Finding Second-order Stationary Points (Proof of Theorem 28) . . . . . . . . . . .
E.
Finding Second-Order Stationary Points When Clients are Homogeneous ( ζ ≪ 1 δ ) .
E.
Convergence under PL condition . . . . . . . . . . . . . . . . . . . . . . . . . . .
F Lower bound
The appendix consists of the following sections.
In Appendix A, we provide further literature review and detailed comparison with existing algorithms.
Appendix B describes details of the experiments and additional experiment results.
Appendix C formalizes the assumptions and prepares mathematical tools for later use.
Appendix D gives formal statements for theoretical guarantees on SLEDGE and the complete proofs.
We also explain intuition behind SLEDGE estimators.
Appendix E applies SLEDGE into its federated learning extension.
Finally, Appendix F shows gradient complexity lower bound under the Hessian heterogeneity assumption, which proves SLEDGE’s gradient complexity under the Hessian heterogeneity of ζ is optimal up to polylogarithmic factors 1 .
For those who do not have enough time to read all the contents, we recommend to look over
Appendix D.1.
We expect this part gives a flavor of the core concepts of the SLEDGE estimator to achieve near-optimal complexity and take advantage of less Hessian heterogeneity without periodic full gradient computation.
Appendix A. Further Literature review A.1. Variance reduction
As explained in the main part, variance reduction is a technique in minibatch sampling to construct a gradient estimator with a smaller variance than vanilla SGD by utilizing gradients at previously obtained anchor points [9, 16, 41].
It is originally developed for (strongly)-convex optimization [9, 16, 41, 42], and thereafter extended to nonconvex settings [1, 39, 40].
One of the difficulties in obtaining an appropriate gradient estimator, especially in nonconvex settings, is that recursive update of a gradient estimator with minibatch gradients easily accumulates the error and eventually buries the correct descent directions.
To address this issue, there have been two major approaches.
The first approach is to explicitly store previously calculated gradients as in SAGA [9].
However, SAGA requires O ( 1 ε 2 ) times of updates with minibatch sample size with
O ( n 23 ) [40] for solving (1), which is still sub-optimal from the lower bound of [10, 27].
The second one is to use double-loop algorithms that periodically compute the full gradient or a gradient with a large minibatch to refresh a gradient estimator.
These algorithms include SARAH, SPIDER, and
NestedSVRG [10, 35, 50], which have the optimal rate, i.e., O ( 1 ε 2 ) times of updates with minibatch sample size with O ( √ n ) .
On the other hand, this approach has an issue that the step of gradientrefreshing slows down practical computational speed and becomes a bottleneck in the application to federated learning since this leads to periodic synchronization and communication between the whole client.
Recent studies have attempted to develop methods that solve this trade-off, or namely that do not require periodic computation of gradients with a large minibatch size to achieve near-optimal rates [3, 7, 23, 28, 30, 36, 45].
Among them, Li et al.
[28] introduced ZeroSARAH as a singleloop algorithm with optimal gradient complexity for nonconvex optimization.
Here, we say an algorithm is single-loop when it does not require periodic full or large minibatch gradients.
However, ZeroSARAH’s estimator cannot take advantage of heterogeneity between f i s. (The reason [Table] will be briefly explained in Appendix D.1.)
To benefit from less Hessian-heterogeneity, we need a path-integrated type estimator like Fang et al.
[10], Nguyen et al.
[34], Zhou et al.
Moreover, these recent single-loop methods have an issue in their versatility.
First, while it is usual to extend an optimization algorithm to ensure second-order optimality [11, 15, 48], and variance reduction methods also have been applied to this [2, 10, 26], no single-loop algorithm cannot find SOSPs.
Since first-order stationary points can include a local maximum or a saddle point in nonconvex optimization, escaping them and finding SOSPs are necessary to guarantee the quality of the solution.
In addition, most of the existing single-loop methods have focused on removing full gradient computation in some specific setting.
Thus none of them achieve both optimal complexity in nonconvex settings and exponential convergence in strongly-convex settings.
Note that he Polyak-Łojasiewicz (PL) condition, which we use instead of strong-convexity, is a generalization of strong convexity to nonconvex settings [38].
One of the recent lines of research is to loosen the conventional assumption of strong convexity and to show exponential convergence under the PL condition [18, 27].
For example, PAGE [27] achieves exponential convergence under the PL condition and the optimal rate in general nonconvex settings, but it should compute the full gradient at a certain probability.
Thus, we do not regard PAGE as a single-loop algorithm in our definition.
We summarize the convergence rates of existing algorithms and ours in Table 1.
Our gradient estimator, SLEDGE, satisfies nearly optimal complexity for general nonconvex settings, speedup under less Hessian heterogeneity, second-order optimality, and exponential convergence under the
PL condition, without requiring periodic full gradient computation.
For more details on SLEDGE, see Appendix D. Also, Appendix F shows that SLEDGE with Option II achieves nearly optimal gradient complexity under less Hessian heterogeneity of ζ .
A.2. Federated learining
In Table 2, we compare the required numbers of communication rounds of existing algorithms and ours. [Table] here [Table]
where full participation of
[Table]
Since inner-loop complexity is different for each algorithm
but it is most usual to compare algorithms in terms of communication rounds, we listed the complexity by ignoring dependency on the intra-client variance σ .
To do so, ours requires b ≳ σ 2 PKε 2 and b ≥ K .
Note that this is quite moderate, since FedAvg [31] requires b ≳ σ 2 pKε 2 [19] (, which is larger than σ 2 PKε 2 ), MimeMVR [20] needs at least one local full gradient between communication rounds, and BVR-L-SGD requires b ≥ K and b ≥ ( √
PmK ∧ σKε ) /ζ
We remark that STEM [21] and FedGLOMO [8] do not have advantage in using local updates, when comparing communication rounds and complexity with centralized algorithms.
Appendix B. Additional experiments B.1. Details of the experiment for Figure
We consider a classification of the capital letters using EMNIST By Class dataset [6].
The original dataset consists of 814 , 255 images of handwritten uppercase and lowercase letters and numbers 09 .
Note that the number of data points in each class is not balanced.
Since the number of images of lowercase letters is relatively small, we only used the images of uppercase letters for the experiment.
To balance the number of data points between each class, we took the following procedure.
We repeatedly sampled 100 data points five times per each uppercase letter, which yields 26 × 5 =groups of sampled data.
For each group i , we define f
i
as the average of the cross-entropy loss between the output of the model and the true class, over the 100 data points belonging to the group.
As a model, we adopted a four-layer fully-connected neural network, following Murata and Suzuki [33].
We added L 2 -regularizer with a regularization parameter of λ = 0 .
01 to the empirical risk.
As competitors, we implemented SARAH [34, 35], SSRGD [26], and ZeroSARAH [28].
We set the minibatch size to b = 12 ≒ √ n = √ 130 for all algorithms, the inner-loop length of SARAH and SSRGD to m =
⌊ nb ⌋ = 10 , and λ = bn ≒ 0 .
092 for ZeroSARAH.
Note that [28] adopted λ = b 2 n , but we found that
λ
= bn was more stable in this setting.
The learning rate for each method was tuned individually, from { 1 . 0 , 0 . 3 , 0 . 1 , 0 . 03 , 0 . 01 , 0 . 003 , 0 . 001 } , so that the test accuracy after 2000 iterations is the highest.
For SSRGD and noisy SLEDGE, we added small noise of r = 0 .
We plotted the mean of the ten trials with different random seeds and the sample variance is also shown in the corresponding (lighter) color for each algorithm.
B.2. Details of the experiment for Figure
We consider a classification of the capital letters using EMNIST By Class dataset [6] as well.
However, here we divided the images in such a way that f i is a little less heterogeneous, but still more heterogeneous than i.i.d.
sampling, as follows.
First, we prepared the same number of data points for each class, and divided them into each client i by the following procedure, setting q = 0 . 9 ;
Then, for each class, we distributed q × 100% of the images into four clients, and the rest into the remaining 100 clients.
This yields that we have 4 × 26 = 104 clients, each of which contains q × 100% of the data from one class, and (1 − q ) × 100% of the data from the other classes.
We call this grouping as a dataset with the heterogeneity parameter of q .
Then, we constructed f i,j with the cross-entropy loss and a two-layer fully-connedted neural network, following Murata and Suzuki [32].
L 2 -regularizer with a scale of λ = 0 .
01 is added to the empirical risk.
We compared FLEDGE with FedAvg [31], SCAFFOLD [19], MimeMVR [20], and BVR-LSGD [32].
For each algorithm, we set p =
10 ≒ √ P = √
as the number of the clients used at each communication (except for BVR-L-SGD, which requires p = P = 104 ).
Then, we set the local minibatch size as b = 16 and the number of local update to K = 10 .
We tuned the learning rate for each algorithm individually from { 1 . 0 , 0 . 3 , 0 . 1 , 0 . 03 , 0 . 01 , 0 . 003 , 0 . 001 } , so that the test accuracy after 2000 outer-loop iterations is the highest.
Here we set the global learning rate of SCAFFOLD to η = 1 , as is done in the original paper [19].
MimeMVR adopted a momentum parameter of a = 0 .
1 as the authors of the paper reported as the best.
We plotted the mean of the five trials with different random seeds and the sample variance is also shown in the corresponding (lighter) color for each algorithm.
B.3. Additional experiments for SLEDGE Comparison with SARAH by changing the learning rate
Here we provide comparison of
SLEDGE with SARAH [34, 35], which is one of the most prevailing variance reduction algorithm with theoretical optimal complexity of O ( √ nε 2 ) .
As is done in the experiment for Figure 1,we prepared f i in the following way.
We repeatedly sampled 100 data points five times per each uppercase letter, which yields 26 × 5 = 130 groups of sampled data.
For each group i , we define f
i
as the average of the cross-entropy loss between the output of the model and the true class over the 100 data points belonging to the group.
As a model, we adopted a two-layer fully-connected neural network, following Murata and Suzuki [32].
We set the minibatch size to b = 12 ≒ √ n = √ 130 for both algorithms, and the innerloop length of SARAH to m =
⌊ nb ⌋ = 10 .
We added L 2 -regularizer to the empirical risk with a fixed regularization parameter of λ = 0 .
We compared SLEDGE with SARAH in terms of the training loss, the norm of the gradient computed by the whole training data, the test loss, and the test accuracy, under the same number of stochastic gradient accesses.
We changed the learning rate η between { 0 . 1 , 0 . 03 , 0 . 01 , 0 . 003 , 0 . 001 } .
We plotted the mean of the five trials with different random seeds and the sample variance is also shown in the corresponding (lighter) color for each algorithm.
Figure 3 shows the result.
We clearly observe that the proposed algorithm SLEDGE slightly faster than SARAH in all range of learning rate η .
The trajectories of SLEDGE are as stable as
SARAH in all settings.
This result shows that we can remove the requirement of periodic full gradient evaluation without hurting the stability during optimization with SLEDGE.
Discrepancy between the gradient estimators and the true gradient
Here we compare the norm between the gradient estimators and the true gradient because this is the most essential measure that quantify the quality of the gradient estimator.
The setting is completely the same as the previous experiment for Figure 3, thus n = 130 .
We compared SLEDGE estimator with SARAH and SAGA [9, 40], taking the minibatch size as b = 12 and the inner-loop length of SARAH to m =
⌊ nb ⌋ = 10 .
We set the learning rate to η = 0 .
01 for all algorithms, since the larger step size tend to increase the discrepancy, meaning that it is not fair to compare algorithms with different step sizes to discuss the discrepancy.
We plotted the mean of the five trials with different random seeds and the sample variance is also shown in the corresponding (lighter) color for each algorithm.
Figure 4 shows the squared norm ∥ v t − ∇ f ( x t ) ∥ 2 between the gradient estimator v t of each algorithm and the true gradient ∇ f ( x t ) at each step t .
The discrepancy of the SLEDGE estimator is clearly smaller than that of SAGA, and close to that of SARAH.
Note that SARAH estimator is refreshed at every m = 10 steps.
Remind the discussion in Subsection 3.1.
The SLEDGE estimator is designed to have as small variance as that of SARAH, while removing the need of periodic full gradient computation.
Therefore, this result validates that our strategy actually works well.
B.4. Additional experiments for FLEDGE Escaping saddle points with FLEDGE
Theorem 28 guarantees second-order optimality of FLEDGE.
To validate this theoretical result, we considered the following experiment.
We first prepared a dataset with the heterogeneity parameter of q = 0 .
7 (see Appendix B.2 for details).
Then, we constructed f i,j with the cross-entropy loss and a three-layer fully-connedted neural network, following
Murata and Suzuki [33].
L 2 -regularizer with a scale of λ = 0 .
01 is added to the empirical risk.
We compared FLEDGE with FedAvg [31], SCAFFOLD [19], MimeMVR [20], BVR-L-SGD [32], and
BVR-L-PSGD [33].
Here, we set P = 104 , p = 10 , b = 16 , and K = 10 .
Note that, according to Theorem 28, setting p ≃ √ P theoretically guarantees that the convergence rate of FLEDGE is not affected by the client sampling and achieves the same number of communication complexity as that of BVR-L-PSGD to find SOSPs.
For FLEDGE and BVR-L-PSGD, we added small noise of r = 0 .
We plotted the mean of the five trials with different random seeds.
We omitted the sample variance for clearer presentation.
The result is shown in Figure 5.
We can clearly observe that FLEDGE with small noise and
BVR-L-PSGD achieve the highest test accuracy.
Note that BVR-L-PSGD is almost the same as
FLEDGE with no client sampling ( p = P ).
Thus, this shows that FLEDGE is not affected by the client sampling with p ≒ √ P , which is consistent with the theory.
Ours is as ten times efficient as BVR-L-PSGD, in terms of communication complexity (the number of gradients communicated between the clients).
Performance under changing heterogeneity
To exhibit how correctly FLEDGE can control the variance between clients, we measured the performance of FLEDGE under changing heterogeneity.
We changed heterogeneity parameter in the range of q ∈ { 0 . 04 (i.i.d.) , 0 . 1 , 0 . 3 , 0 . 5 , 0 .
9 , 1 .(completely heterogeneous) } , and compared FLEDGE with FedAvg, in terms of both train and test [Figure] [Figure] [Figure]
[Figure]
[Figure]
[Figure] [Figure] [Figure] [Figure]
(m) Train loss ( η = 0 . 003 )
[Figure]
i
[Figure]
(q) Train loss ( η = 0 .
[Figure]
(r) Gradient norm (
[Figure]
[Figure]
[Figure]
T e s t a cc u r a c y
[Figure]
accuracy.
All other settings are the same as that of the experiment for Figure 2.
Note that we chose p = 10 ≒ √ P = √ 104 , where the theory says that the convergence is never affected by sampling of clients.
Figure 6 shows the average of five trials with different random seeds.
[Figure]
According to Figure 6, while FedAvg decreases the train and test accuracy as the heterogeneity increases, the performance of FLEDGE with even q = 1 . 0 is only slightly worse than that with q = 0 . 04 .
The fact that FLEDGE is little affected by the strong heterogeneity in this experiment supports our theoretical guarantee (Theorems 22 and 34) on the effect from sampling of clients.
That is, setting p ≥ √ P , our algorithm does not affected by sampling of clients and finds ε -first-order stationary points within ˜ O
( 1 Kε 2 + ζε 2 ) communication rounds.
B.5. Computing infrastructures
• OS: Ubuntu 16.04.
• CPU:
Intel(R) Xeon(R)
CPU E5-
v4 2.40GHz • CPU Memory: 512GB • GPU: Nvidia Tesla V100 (32GB) • Programming language: Python 3.6.
• Deep learning framework: PyTorch 1.7.
Appendix C. Assumptions and Tools
In this section, we formally restate the assumptions and introduce mathematical tools we utilize in the missing proofs.
C.1. Assumptions
First, gradient Lipschitzness and boundedness are assumed as usual.
Assumption 1 (Gradient Lipschitzness)
For all i ∈ [ n ] , f i is L -gradient Lipschitz, i.e., ∥∇ f i ( x )
Assumption 2 (Existence of global infimum) f has the global infimum f ∗ = inf x ∈ R d f ( x ) and ∆ :
= f ( x 0 ) − f ∗ .
Below, (i) inter-client gradient boundedness is assumed for SLEDGE with Option I to remove full gradient even at x 0 , as in ZeroSARAH [28].
(ii) Intra-client gradient boundedness is assumed for FLEDGE.
In the main text, we bounded inter-client and intra-client variances uniformly as ∥∇ f i ( x 0 ) − ∇ f ( x 0 )
∥ ≤ σ c and ∥∇ f i,j ( x ) − ∇ f i ( x ) ∥ ≤ σ .
However, if the variances are small in expectation, we can loosen the uniform boundedness assumptions into the followings.
From now, we assume Assumption 3 instead of ∥∇ f i ( x 0 ) − ∇ f ( x 0 ) ∥ ≤ σ c and ∥∇ f i,j ( x ) − ∇ f i ( x ) ∥ ≤ σ .
Assumption 3 (Boundedness of Gradient) (i) It holds that E i [ ∥∇ f i ( x 0 ) − ∇ f ( x 0 ) ∥ 2 ] ≤ σ 2 c , where the expectation E i is taken over the choice of i .
Moreover, ∥∇ f i ( x 0 ) − ∇ f ( x 0 ) ∥ ≤ G c holds for all i .
(ii)
For all i and x , assume E j [ ∥∇ f i,j ( x ) − ∇
f i ( x ) ∥ 2 ] ≤ σ 2 , here the expectation E j is taken about the choice of j .
For all i, j and x , ∥∇ f i,j ( x ) − ∇ f i ( x ) ∥ 2 ≤ G 2 .
In order to give second-order optimality, Hessian Lipschiteness is usually assumed [12, 26].
Assumption 4 (Hessian Lipschitzness) { f i } ni =1 is ρ -Hessian Lipschitz, i.e., ∥∇
f i ( x ) −∇
2 f i ( y ) ∥
≤ ρ ∥ x − y ∥ , ∀ i ∈ [ n ] and x, y ∈ R d .
For federated learning, we solely assume inter-client Hessian-heterogeneity to show the efficiency of the proposed method in a less heterogeneous setting.
It has previously appeared in Mime [20] (but intra-client Hessian-heterogeneity was assumed at the same time) and BVR-L-SGD [32].
Assumption 5 (Hessian-heterogeneity) { f i } ni =1 is Hessian-heterogeneous with ζ , i.e., for any i, j ∈ [ n ] and x ∈ R d , ∥∇
f i ( x ) − ∇ 2 f j ( x ) ∥ ≤ ζ .
Finally, we explain the PL condition [38].
It is easy to see a µ -strongly convex function satisfies this with µ .
Assumption 6 (PL Condition)
f satisfies PL condition, i.e., ∥∇ f ( x ) ∥ 2 ≥ 2 µ ( f ( x ) − f ∗ ) for any x ∈ R d .
C.2. Concentration inequalities
Here, we prepare concentration inequalities for later use.
We first present Bernstein-type bounds.
Proposition 3 (Matrix Bernstein inequality [47]) Let X 1 , · · · , X k be a finite sequence of independent random matrices with dimension d 1 × d 2 .
Assume each random matrix satisfies
E [ X i ] =and ∥
X i ∥ ≤ R almost surely.
Define
[Equation]
Then, for all t ≥ 0 ,
[Equation]
i
[Equation]
In this paper, we deal only with the vector case.
In that case, the inequality is rewritten for bounds with high probability, as follows.
Proposition 4 (Vector Bernstein inequality) Let x 1 , · · · , x k be a finite sequence of independent, random, d -dimensional vectors and ν ∈ (0 , 1) .
Assume that each vector satisfies
[Equation]
almost surely.
Define
[Equation]
Then, with probability at least 1 − ν/ poly( n, P, T, K ) ,
[Equation]
where C 1 = O (cid:0) log( ν − 1 + n + P + d + T + K ) [-where C 
-](cid:1)
=
˜ O (1) .
Remark
Here we do not specify poly( n, P, T, K ) to apply different polynomials to later.
Whenever we use this inequality with different poly( n, m, T ) , we will reuse C 1 for the notational simplicity.
We also use this constant C 1 in the following parts to denote constants as large as O (log( ν − 1 + n + P + d + T + K )) , with a slight abuse of notations.
Moreover, a similar inequality holds when we consider sampling without replacement.
To our knowledge, Bernstein inequality without replacement for vectors has not been rigorously proven and we attach its complete proof at the next subsection.
Proposition 6 (Vector Bernstein inequality without replacement) Let A =
( a 1 , a 2 , · · · , a k ) be d -dimensional fixed vectors, X = ( x 1 , · · · , x l ) ( l ≤ k ) be a random sample without replacement from A .
Assume that (cid:80) ki
=1 a i = 0 and that each vector satisfies
[Equation]
Define
[Equation]
Then, for each t ≥ 0 and l < k ,
[Equation]
Moreover, for each l < k , with probability at least 1 − ν/ poly( n, P, T, K ) ,
[Equation]
where C 1 =
O (log( n + P + d + T + K ))
=
˜ O (1) .
Finally, we need a high-probability version of Azuma-Hoeffding inequality.
Proposition 7 (Azuma-Hoeffding inequality with high probability [5, 44]) Let { x i } be a d -dimensional vector sequence and martingale with respect to a filtration {F i } .
Assume that each x i satisfies E [ x i |F i − 1 ] = 0 and ∥ x i ∥ ≤ R i with probability 1 −
ν
i for ν i ∈ (0 , 1) ( i = 1 , . . .
, k ) .
Then, with probability at least 1 − ν/ poly( n, P, T, K ) − (cid:80) ki
=
ν i , (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) k (cid:88) i
=
x
i (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ≤ Ck (cid:88) i =
R 2 i , where C 1 = O (cid:0) log( ν − 1 + n + P + d + T + K ) (cid:1) =
˜ O (1) .
C.3. Proof of Proposition
In order to show Proposition 6, we use the Martingale counterpart of Bernstein’s Inequality for random matrix.
The following is a slightly weaker version of Tropp [46].
Proposition 8 (Freedman’s inequality for matrix martingales) Consider a matrix martingale { Y i | i = 0 , 1 , · · · } with respect to a filtration {F i } , whose values are matrices with dimension d 1 × d 2 , and let { X i | i = 1 , 2 , · · · } be the difference sequence.
Assume that each of the difference sequence is uniformly bounded:
[Equation]
almost surely.
Also, assume that each i satisfies
[Equation]
almost surely.
Then, for all t ≥ 0 and for each l ,
[Equation]
Proof of Proposition
First, we consider the case l ≤ k 2 .
Let y i
= (cid:80) ij =1 x j and consider a filtration F i =
σ ( x 1 , · · · , x i ) .
Then, we have
[Equation]
This means that (cid:110) 1 k − i y i (cid:111) l
i =0 is martingale with respect to {F i } .
We have that this martingale satisfies the assumptions of Proposition 8 with R ′ 2 =
R 2 ( k − l ) 2 and σ ′ 2 = 2 σ 2 ( k − l ) 2 . In fact, we have (cid:13)(cid:13) (cid:13) (cid:13) k [Equation] [Equation]
(cid:12)(cid:12) (cid:12) (cid:12)
[Equation]
(cid:21)(cid:13)(cid:13) (cid:13) (cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
where the equality follows since x 1 . . .
[Equation]
x i are F i -measurable, and
[Equation]
Thus
[Equation]
we use Proposition 8 to obtain
[Equation]
What remains is the case of l ≥ k 2 .
Since (cid:80) li =1 x i = − (cid:80) ki
= l +1 x i holds, we can apply the above bound for (cid:80) ki
= l +1 x i .
Thus, we have the first assertion for all l < k .
The second assertion follows by setting t
= O (cid:0) ( lσ 2 + R ) log( ν − 1 + n + m + d + T ) (cid:1) = C 1 · ( lσ 2 + R ) .
C.4. Linear Algebraic Tool
The following lemma is due to Murata and Suzuki [33].
We provide its proof here.
Lemma 9 (Murata and Suzuki [33]) Let A be a d × d symmetric matrix with the smallest and largest eigenvalues λ min < 0 and λ max < 1 , respectively.
Then, for k = 0 , 1 , · · · , it holds that
[Equation]
Proof
Since A is diagonalizable, we write A = (cid:80) di
=1 λ
i e i e ⊤ i , where e 1 , . . .
, e d are normalized eigenvectors and λ min = λ 1 ≤ · · · ≤ λ d = λ max are the corresponding eigenvalues.
Then, it holds that
[Equation]
Thus, the remaining is to evaluate max i | λ i (1 − λ i ) k | .
After some algebra, we get
[Equation]
if λ
[Equation]
0 ; the equality holds with λ = 11+ k
[Equation]
which concludes the proof.
Appendix D. Missing Proofs for SLEDGE
This section provides the missing proofs in Section 2 about the convergence property of SLEDGE.
First, we divide Theorem 1 into the following three formal theorems.
Theorem 10 Under Assumptions 1, 2 and 5, and 3-(i) for Option I, if we choose η = ˜Θ( 1 L ∧ b ζ √ n ) and r ≤ ηε 2 , Algorithm 1 finds ε -first-order stationary points using
[Equation]
stochastic gradients with probability at least 1 − ν .
Theorem 11 Assume Assumptions 1, 2, 4 and 5, 3-(i) for Option I. Let b ≳ √ n + ζ 2 δ 2 , η = ˜Θ( 1 L ) , r ≲ ˜ O (cid:0) εL (cid:1) , and ν ∈ (0 , 1) .
Then, Algorithm 1 finds ( ε, δ )
-SOSPs using
[Equation]
stochastic gradients, with probability at least 1 − ν .
Theorem 12 Assume Assumptions 1, 2, 3-(i), 5, and 6.
If η = ˜Θ( 1 L ∧ b ζ √ n ∧ bµn ) , and r ≤ η (cid:113) εµ 3 ,
Algorithm 1 finds an ε -solution with f ( x t ) − f ∗ ≤ ε using
[Equation]
I
[Equation]
II [Equation] stochastic gradients with probability at least 1 − ν .
˜ O hides at most log 5 .
5 ( n + µ − 1 + ν − 1 ) and polyloglog factors.
Below, we first describe basic idea for designing SLEDGE.
Then,appendix D.2 presents the first-order optimality ( Theorem 10).
Then we prove that SLEDGE can escape saddle points and find SOSPs (Theorem 11) in appendix D.3.
Finally, exponential convergence under the PL condition (Theorem 12) is shown in appendix D.4.
D.1. Intuitition behind SLEDGE
SLEDGE is designed so that it inherits the best points of SAGA [9, 40] and SARAH [34, 35].
We first compare gradient estimators of SAGA and SARAH to explain why SAGA is suboptimal.
Then, we also note SARAH’s estimator gets more accurate under less heterogeneity.
According to SAGA’s update rule, the discrepancy of the gradient estimator from the true gradient at a step t can be decomposed as
[Equation]
where I t is the randomly chosen minibatch with size b and T ( t, i ) is the step when f i is last sampled.
Note that SAGA stores ∇ f i ( x T ( t,i ) ) for each i .
Thus, the first term is a change from the referable gradient of 1 n (cid:80) ni =1 ∇ f i ( x T ( t,i ) ) , and the second term is an approximation of the first term using a minibatch with size b .
Then, the variance of the gradient estimator is roughly bounded by L b ∥ x t − x T ( t ) ∥ [Equation] ≤ L
[Equation] t − T [Equation] b
[Equation] with [Equation]
min [Equation]
On the other hand, the difference between SARAH’s gradient estimator, which computes the full gradient periodically, and the true gradient can be written as
[Equation]
where T ( t ) is the time of the last full gradient evaluation.
We can interpret this scheme as it decomposes ∇ f ( x t ) − ∇ f ( x T ( t ) ) into the sum of ∇ f ( x s ) − ∇ f ( x s − 1 ) , and each term is approximated by an independent minibatch with size b .
Then, the variance is bounded by Lb (cid:80) ts = T ( t )+1 ∥ x s − x s − 1 ∥ 2 , meaning that SARAH’s estimator is better than that of SAGA by the t − T ( t ) factor, which can be as large as O ( nb ) .
Moreover, in the perspective of application to FL, whether an algorithm can take advantage of less heterogeneity becomes important.
We consider the difference of the probabilistic term fo In
SAGA’s update, the variance cannot be bounded by a factor of ζ , since x T ( t,i ) is different for different i .
On the other hand, we can see that, for SARAH’s path-integrated type estimator, the variance can be bounded by ζb (cid:80) ts = T ( t )+1 ∥ x s − x s − 1 ∥ 2 under the Hessian heterogeneity of ζ .
Thus, the path-integrated estimator of SARAH (, which also appears in SPIDER [10] or NestedSVRG [50]), is crucial in taking advantage of less heterogeneity.
This is also the reason for why ZeroSARAH [28] fails to utilize less Hessian heterogeneity, since at step t they estimate ∇ f ( x t − 1 ) based on naive application of SAGA.
Based on the above discussion, we first decompose SAGA’s approximation targetn (cid:80) ni
=1 ∇ ( f i ( x t ) − ∇ f i ( x T ( t,i ) )) into the sum of ∇ f i ( x s ) − ∇ f i ( x s − 1 ) , each of which is approximated in SARAH’s manner.
Namely, the decomposed form is written as follows:
[Equation]
n (cid:88) s =
T ( t )+ n (cid:88) i ∈ ˜
I ts ( f i ( x s ) − f i ( x s − 1 )) .
Here ˜ I ts =
[ n ] \ (cid:83) t τ = s I t , so that ˜ I ts is the set of indexes not sampled between s and t .
Then, we approximate (cid:80) ni ∈ ˜
I ts ( f i ( x s ) − f i ( x s − 1 )) with | ˜ I ts | b (cid:80)
ni ∈
I s ( f i ( x s )
− f i ( x s − 1 )) .
This procedure yields
SLEDGE, and the following error bound on the SLEDGE estimator.
Lemma 13 (Informal) Let ν ∈ (0 , 1) , T 1 =
˜
O
( nb ) , and C =
˜ O (1) .
We have that, ignoring the initialization error, with probability 1 − ν for all t = 1 , · · · , T ,
[Table]
Here, T 1 is defined so that T 1 ≥ t − T ( t ) holds with high probability.
This lemma tells us that our gradient estimator has comparable quality to SARAH without computing full gradient.
Moreover, this lemma explicitly states that the variance of SLEDGE estimator is quadratically bounded with ζ , meaning that we require fewer gradients when ζ ≪ L , which is later exploited for federated learning application.
While the development is intuitively straightforward, we have the technical difficulty to evaluate the error, that | ˜
I ts | depends not only on I s but also on I s +1 , . .
, I t .
In other words, unlike SARAH, the discrepancy cannot be decomposed into completely independent terms about I s , which prevents us from using a usual expectation bound.
To address this, we prepared vector Bernstein inequality without replacement (Proposition 6) to give a high probability bound on the discrepancy.
D.2. Finding First-order Stationary Points (Proof of Theorem 10)
In this subsection, we show that SLEDGE finds first-order stationary points with high probability (Theorem 10).
For the proof of Theorem 10, we use the following classical argument (e.g. Ge et al.
[12], Li [26], Li et al.
[27]), which ensures decrease of the function values.
Lemma 14 Let f be an L -gradient Lipschitz function and x t := x t − 1 − ηv t − 1 + ξ t − 1 with ∥ ξ t − 1 ∥ ≤ r .
Then,
[Equation]
holds.
Proof
Starting from the direct result from L -gradient Lipschitzness
[Equation]
we have
[Equation]
where we used x t − x t − 1 = ηv t − 1 + ξ t − 1 and ⟨ a − b, b ⟩ = 12 ( ∥ a − b ∥ 2 − ∥ a ∥ 2 + ∥ b ∥ 2 ) for (3),
∥ a + b ∥ 2 ≤ 2( ∥ a ∥ 2 + ∥ b ∥ 2 ) for (4), and ∥
ξ t − 1 ∥ ≤ r for (5).
Our algorithm uses v t = 1 n (cid:80) ni
=
y ti as an estimator of ∇ f ( x t ) .
To apply Lemma 14 for our algorithm, we need to evaluate the term ∥ v t − ∇ f ( x t ) ∥ 2 , the variance of the gradient estimator.
The next lemma provides its upper bound that holds with high probability.
Lemma 13 Let v t = 1 n (cid:80) ni =1 y ti and all the other variables be as stated in Algorithm 1.
Then, by taking T 1 = nb C 1 ,
[Equation]
[Equation]
holds for all t
= 1 , · · · , T with probability at least 1 − 3 ν .
We decompose ∥ v t − ∇ f ( x t ) ∥ into three parts to each of which one of the following lemmas is applied.
Below, for each 1 ≤ s ≤ t , we let ˜ I
ts =
[ n ] \ (cid:83) tτ =
s I t , which is a set of indexes that are not selected between s + 1 and t .
Lemma
The following holds uniformly for all 1 ≤ t ≤ T with probability at least 1 − ν :
[Equation]
Lemma
The following inequality holds uniformly for all 1 ≤ t ≤ T with probability at least 1 − ν :
[Equation]
Lemma
The following inequality holds with probability at least 1 − ν uniformly over 1 ≤ t ≤
T : (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)n (cid:88) i ∈ ˜
I t
[Equation]
y
[Equation]
f
[Equation]
x 0 )) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ≤ (cid:40) 4 C 21 b (cid:16) σ 2 c + G 2 c b (cid:17) ( Option I ) ( Option II ) .
Proof of Lemma
First, we have that
[Equation]
−
[Equation] [Equation]
) (cid:13) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) , (7) where we use ∥ (cid:80) mi =1 a i ∥ 2 ≤ m (cid:80) mi =1 ∥ a i ∥ 2 for the first inequality.
For each s , from Assumption 5,
[Equation]
holds for all i ∈ [ n ] .
By vector Bernstein inequality without replacement (Proposition 6), for each t ≥ 1 and s satisfying max { 1 , t − T 1 + 1 } ≤ s ≤ t , we have that
[Equation]
i ∈ ˜
I ts
[Equation]
holds with probability at least 1 − νT 2 .
Thus, in (7), (8) holds uniformly for all t and s with probability at least 1 − ν .
Applying this bound to (7) yields
[Equation]
where the second inequality follows from | ˜
I ts | ≤ n and T 1 = nb C 1 .
Proof of Lemma
Since | ˜ I ts | depends not only on I s but also on I s ,
I s +1 , · · · , I t , the left-hand side of (6) is not a sum of martingale variables with respect to the filtration { σ ( I 1 , · · · , I s ) } ts =1 .
Thus, we consider E [ | ˜ I ts | ] instead of | ˜
I ts | and validate the difference between them later.
We decompose (6) as
[Equation]
n
[Equation]
=
[Equation]
E [ | ˜ I ts | ] b (cid:88)
i ∈
I s ( ∇ f i ( x s ) − ∇ f i ( x s − 1 ) − ( ∇ f ( x s )
+ ∇ f ( x s − 1 )))
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
+ 2 (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
1 n t (cid:88) s
=max { 1 ,t − T 1 +1 } (cid:16) | ˜
I ts | − E [ | ˜ I ts | ] (cid:17) 2 1 b (cid:88) i ∈
I s ( ∇ f i ( x s ) − ∇ f i ( x s − 1 ) − ( ∇ f ( x s )
+ ∇ f ( x s − 1 ))) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ≤ 2 (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 1 n t (cid:88) s =max { 1 ,t − T 1 +1 }
E [ | ˜ I ts | ] b (cid:88)
i ∈
I s ( ∇ f i ( x s ) − ∇ f i ( x s − 1 ) − ( ∇ f ( x s )
+ ∇ f ( x s − 1 )))
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) (9) + 2 Tnt (cid:88) s =max { 1 ,t − T 1 +1 } (cid:16) | ˜
I ts | − E [ | ˜ I ts | ] (cid:17) 2 (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 1 b (cid:88)
i ∈
I s ( ∇ f i ( x s
[Equation]
[Equation]
First, we bound the term (9).
We can see that E [ | ˜ I ts | ] b
[Equation]
i ∈
I s
[Equation]
∇ f
[Equation]
f
[Equation] [Equation]
−
[Equation]
∇ f
[Equation]
x s ) + ∇ f ( x s − 1 ))) is a martingale difference sequence.
Moreover, by the vector Bernstein inequality without replacement (Proposition 6) and Assumption 5, we have
[Equation]
with probability at least 1 − ν 5 T 2 for each t and s in (9).
This allows us to use the Azuma-Hoeffding inequality with high probability (Proposition 7).
Consequently, with probability at least 1 − ν 5 T −
T · ν
[Equation]
= 1 − 2 ν 5 T , it holds that
[Equation] [Equation]
for each t ∈ [ T ] .
Therefore, (11) holds for all t with probability 1 − 2 ν 5 .
As for the term (10), by the Bernstein inequality without replacement (Proposition 6) and Assumption 5, we have
[Equation]
for all s with probability at least 1 − ν 5 .
We move to bound the difference (cid:16) | ˜
I ts | − E [ | ˜ I ts | ] (cid:17) 2 .
For this purpose, we regard this as a function of (at most) T 1 variables
I t , · · · , I s and prepare a “reverse” filtration ˜ F = { ˜ F ts } max { 1 ,t − T 1 +1 } s = t with ˜ F ts =
σ ( I t , I t − 1 , · · · , I s ) .
Then, the sequence {| ˜ I ts |} ts =max { 1 ,t − T 1 +1 } is a measurable process with respect to ˜ F .
We consider the conditional expectation of | ˜ I ts | − | ˜
I ts +1 | with respect to ˜ F .
When samples in ˜ I ts +1 are not chosen between t to s + 1 , each of them is chosen with probability bn for the first time at step s .
Thus, we have
[Equation]
which leads to E s (cid:104) |
˜
I ts | (cid:12)(cid:12)(cid:12) ˜ F ts +1 (cid:105) =
(cid:0) 1 − bn (cid:1) | ˜
I ts +1 | .
Hence, the process { u ts := | ˜
I ts |− (cid:0) 1 − bn (cid:1)
| ˜
I ts +1 | | t > s ≥ t − T 1 + 1 } is a martingale with respect to ˜ F and satisfies E s (cid:104) u ts (cid:12)(cid:12)(cid:12) ˜
F ts +1 (cid:105) =
In addition, let A = { 1 , · · · ,(cid:124) (cid:123)(cid:122) (cid:125) | ˜
I ts +1 | , 0 , · · · ,(cid:124) (cid:123)(cid:122) (cid:125) n −| ˜ I
ts +1 | } and
˜ A
=
(˜ a 1 , · · · , ˜ a b ) be a random sample without replacement from
A , Then, u t s conditioned on ˜ F t s +1 follows the same distribution as that of (cid:80) b l =1 ˜ a
i − E (cid:104)(cid:80) b l =1 ˜ a i (cid:105) .
This means that, using Proposition 6, we have ∥ u ts ∥ ≤ C 1 √ b with probability at least 1 − 15 T 2 .
Finally, we apply Proposition 7 to bound | ˜
I ts | = (cid:80) sτ = t (cid:0) 1 − bn (cid:1) ( τ − s )
u tτ + n (cid:0) 1 − bn (cid:1) ( t − s +1) , which yields that
[Equation]
with probability at least 1 − ν 5 T − T · ν 5 T 2 = 2 ν 5 T .
Combining (12) and (13), with probability 1 − ν 5 , we get
T
[Equation] [Equation]
s
[Equation] [Equation]
[Equation] [Equation]
by letting T 1 = nb C 1 .
Finally, we get the assertion by combining (11) and (14), and applying 2 C 21 b + 2 C 81 b 2 ≤ 4 C 81 b .
Proof of
Lemma
As for the Option II, the assertion directly follows from the definition y
i = ∇ f i ( x 0 )
( i = 1 , · · · , n ) .
Henceforth, we prove the bound (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)n (cid:88) i ∈ ˜
[Equation] [Equation]
b
[Equation]
when we use Option I. To this end, we decompose (cid:13)(cid:13)(cid:13) 1 n (cid:80) i ∈
˜
I
[Equation]
n
[Equation]
where we use the inequality ∥ a + b ∥ 2 ≤ 2 ∥ a ∥ 2 + 2 ∥ b ∥ 2 .
For the first term in (15), Propositionand Assumption 3 imply that (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
[Equation]
holds with probability at least 1 − ν 2 for all t .
For the second term in (15), we have (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) (cid:13)
[Equation]
(cid:88) i ∈ ˜
I t
[Equation]
from Proposition 6 and Assumption 3, with probability at least 1 − ν 2 T for each t and at least 1 − νuniformly over all t .
Substituting (15) and (16) to (14), we obtain the desired bound.
Proof of Lemma
We first observe that v t − ∇ f ( x t ) is written as v t − ∇ f ( x t ) =
[Equation]
We can ensure that if t − s is sufficiently large, every f i is sampled at least once between s +and t with high probability.
Indeed, for each f i and ν > 0 , the probability that f i is not sampled between s + 1 and t is bounded as
[Equation]
where we use (cid:80) ts =max { 1 ,t − T 1 +1 } b ≤ T 1 b = nC 1 in the second inequality.
By taking C 1 = Ω (cid:0) log nTν (cid:1) , the right-hand side of (17) is bounded by νnT .
In other words, ˜ I ts = ∅ with probability at least 1 − ν for every t and s ≤ t − T 1 .
Henceforth, we assume ˜
I ts = ∅ and focus on the errors between max { 1 , t − T 1 + 1 } ≤ s ≤ t .
When ˜ I ts = ∅ holds for s ≤ t − T 1 , the variance term ∥ v t − ∇ f ( x t ) ∥ 2 is decomposed as ∥ v t − ∇ f ( x t ) ∥
[Equation]
n
[Equation]
 
[Equation]
n
[Equation]
≤
[Equation]
t
[Equation]
s =
[Equation]
by the inequality ∥ a + b + c ∥ 2 ≤ 3( ∥ a ∥ 2 + ∥ b ∥ 2 + ∥ c ∥ 2 ) .
Then, we give the bound of (18) for
Option I and Option II, respectively.
Option I
According to Lemmas 15 to 17, we have
[Equation]
b with probability at least 1 − 3 ν uniformly over all t .
Option II
Almost as well as the previous case, we have
[Equation]
with probability at least 1 − 2 ν uniformly over all t .
By replacing C 1 with C 1 ∨ 1 and applying 3 C 31 + 12 C 81 ≤ 15 C 81 , we obtain the desired bound for both cases.
Now, we are ready to prove the first-order convergence of SLEDGE.
Proof of Theorem
Summing up (5) over all t
= 1 , 2 , · · · , T and arranging the terms, we get
[Equation]
Applying Lemma 13 to this, we obtain that
T (cid:88)
[Equation]
) with probability at least 1 − 3 ν .
Option I
We take η = min (cid:110) 12 L ,
[Equation]
(cid:17) , which implies 12 η − L 2 − 15 C 91 ηζ 2 n b≥ 0 .
We have 2 Tr 2 η≤ Tε 2 2 by letting r ≤ ηε 2 .
Also, we have f ( x 0 ) − f ( x t ) ≤ ∆ by Assumption 2.
Summarizing these, we get
T (cid:88)
[Equation]
Thus, by setting T ≥ (cid:16)
[Equation] [Equation]
, we obtain 1 T (cid:80)
Tt =1 ∥∇ f
( x t − 1 ) ∥ 2 ≤ ε 2 with probability at least 1 − 3 ν , which implies that there exists some t such that 0 ≤ t ≤ T − 1 and ∥∇ f ( x t ) ∥ 2 ≤ ε 2 .
From this, the desired conclusion is obtained.
Option II
As well as Option I, we take η = min (cid:110) 12 L , b C 41 ζ √ 60 n (cid:111) = ˜Θ (cid:16)
b Lb ∨ ζ √ n (cid:17) , which implies 12 η − L 2 − 15 C 91 ηζ 2 n b≥ 0 .
In addition, by letting r ≤ ηε 2 , we obtain (cid:80) Tt =1 ∥∇ f ( x t − 1 ) ∥ 2 ≤ 2∆ η
+ Tε 2 2 .
Then, by taking T ≥ 4∆ ηε 2 = ˜
O (cid:0) L ∆ ε 2 (cid:1) , there exists some t such that 0 ≤ t ≤ T − 1 and ∥∇ f ( x t ) ∥ 2 ≤ ε 2 , with probability at least 1 − 3 ν .
Accordingly, we get the assertion for Option II.
D.3. Finding Second-order Stationary Points (Proof of Theorem 11)
The goal of this subsection is to show that SLEDGE is the single-loop algorithm with theoretical guarantee for finding second-order stationary points.
The argument follows that of [12, 15, 26]; Let x τ 0 be a point such as λ min ( ∇ f ( x τ 0 ))
≤ − δ .
Around that point, we consider two points x 1 and x 2 such that ⟨ x 1 , e ⟩ ≈ ⟨ x 2 , e ⟩ , where e is the eigenvector of λ min ( ∇ f ( x τ 0 )) .
Then, two coupled sequences that SLEDGE generates from the two initial points ( x 1 and x 2 ) will be separated exponentially, as long as they are in a small region around the initial points.
This means that if we add some noise to the sequence around a saddle point, then with a certain probability, the algorithm can move away from the saddle point.
We again emphasize that, although this proof outline has been popular, we face the difficulties arising from the single-loop structure of the algorithm.
Many existing algorithms compute periodic full gradient and can refresh their gradient estimators around saddle points.
In contrast, our singleloop algorithm does not use full gradient, meaning that we have to deal with the error accumulated before that point, and it is not trivial whether such errors can be sufficiently small so that the direction of the negative eigenvalue can be found by the gradient estimator.
This is the first difficulty, and we found that taking minibatch size as large as b ≳ √ n + ζ 2 δ 2 is sufficient.
When δ =
O
( √ ε ) , ζ 2 δ 2 is about O ( √ n + 1 ε ) , which is usually assumed in existing literature [12, 26].
Secondly, our estimator is more correlated due to the | ˜
I ts | term, thus requiring more delicate analysis than that for SSRGD [26], which is based on SARAH [34, 35].
We formalize the exponential separation of two sequences in the following lemma.
Lemma 18 (Small stuck region) Let { x t } be a sequence generated by SLEDGE and suppose that there exists a step τ 0 such that − γ :=
λ min ( ∇ 2 f ( x τ 0 )) ≤ − δ holds.
We denote the smallest eigenvector direction of λ min ( ∇ 2 f ( x τ 0 )) by e .
Moreover, we define a coupled sequence { ˜ x t } by running SLEDGE with ˜ x 0 = x 0 and share the same choice of randomness, i.e., minibatches and noises with { x t } , except for the noise at some step τ ( > τ 0 ) :
˜
ξ τ = ξ τ − r e e with r e ≥ rνT √ d .
Let w t = x t − ˜ x t , v t = 1 n (cid:80) n i =
y ti , ˜ v t = 1 n (cid:80) n i =
˜
y ti , and g t = v t − ∇ f ( x t ) − (˜ v t − ∇ f (˜ x t )) .
Here ˜ y ti is the counterpart of y ti and corresponds to { ˜ x t } .
Then, there exists constants C 2 =
˜ O (1) , C
= O (1) , such that if we take b ≥ √ n + C 22 ζ 2 δ 2 , η = ˜Θ
(cid:0) 1 L (cid:1) , and T 2 = C 3 log δ C 2 ρre
[Equation]
≤ ˜ O
(cid:0) Lδ (cid:1) , with probability at least 1 − νT ( ν ∈ (0 , T )) , it holds that
[Equation]
In order to show Lemma 18, we need the following lemma, which is analogous to Lemmas 15 to 17.
Lemma
Under the same assumption as that of Lemma 18, we assume max τ 0 ≤ t ≤ τ + T 2 {∥ x t − x τ ∥ , ∥ ˜ x t − x τ ∥} < δC 2 ρ .
Then, the following holds uniformly for all t ≤ τ + T 2 with probability at least 1 − νT :
[Equation]
( otherwise ) , where T 1 = nb C 1 .
Here C
= ˜ O (1) is a sufficiently large constant.
Proof of Lemma
As for the case t < τ , the assertion directly follows from the definition of { ˜ x t } .
For the proof of the rest cases, we use notations as follows:
[Equation]
Moreover, to simplify the notation
[Equation]
we denote
[Equation]
We have that E i [ u si ] = 0 , where the expectation is taken over the choice of i .
Furthermore, for s ≥ τ + 1 , by using the Hessian-heterogenity (Assumption 5) and the Hessian Lipschitzness (Assumption 4)
[Equation]
we have that
[Equation] [Equation]
where we use max τ 0 ≤ t ≤ τ + T 2 {∥ x t − x τ ∥ , ∥ ˜ x t − x τ ∥} < δC 2 ρ for the last inequality.
For s = τ , by
Assumption 5, we have ∥ u τi ∥ =
∥ ( ∇ f i ( x τ ) − ∇ f i (˜ x τ
[Equation]
− ( ∇ f ( x τ ) − ∇ f
(˜ x τ )
[Equation]
∥ ≤ 2 ζ ∥ x τ − ˜ x τ ∥ = 2 ζr e .
Recall the discussion in Lemma
[Equation]
we have
[Equation]
  
[Equation]
i ∈
I s u si (cid:33) −n t (cid:88) s
=max { τ +1 , t − T 1 +1 } (cid:88) i ∈ ˜
I ts u si ( t ≥ τ + 1) .
As for the first term in both cases, we have
[Equation]
n
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
b (cid:88) i ∈ I τ u
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:88)
i
[Equation]
˜
I
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
| n 2 C 1 ζr e √ b ≤
[Equation]
C 1 ζr e √ b , (20) by using Proposition 6 and ∥ u τi ∥ ≤ 2 ζr e , with probability at least 1 − ν 4 T for all t .
For the second term in the case t ≥ τ + 1 , we follow the same line as the proof of Lemma 16.
We just replace ∇ f i ( x s ) − ∇ f i ( x s − 1 ) − ( ∇ f ( x s ) − ∇ f ( x s − 1 )) by u si and use (19) to obtain that
[Equation]
with probability at least 1 − ν 4 T for all t .
Finally, we bound the last term in the case t ≥ τ + 1 .
By using Proposition 6, we obtain (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:118)(cid:117)(cid:117)(cid:117)(cid:116)
[Table]
s =
[Table]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:88)
[Table]
∈
[Table]
ts
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
22) with probability at least 1 − ν 4 T for all t .
Combining (20), (21), and (22), we have
[Equation]
with probability at least 1 − νT for all t > τ .
Here we take C 4 =
˜ O (1) .
For t = τ , (20) directly implies the desired bound.
Now, we move to prove Lemma 18.
Proof of Lemma
We assume the contrary, i.e., max τ 0 ≤ t ≤ τ + T 2 {∥ x t − x τ ∥ , ∥ ˜ x t − x τ ∥} < δC 2 ρ , and show the following by induction: for τ ≤ t ≤ τ + T 2 ,
[Equation]
− τ r e .
Then, (a) yields contradiction by taking t − τ = T
= O (cid:18) log
[Equation]
(cid:19) since it holds that
[Equation]
It is easy to check (a) and (b) for t = τ .
As for (c), by taking b ≥ ζ 2 δ 2 , ∥ g t ∥ ≤ C 4 δr e ≤ C 4 γ (1 + ηγ ) t − τ r e holds with probability at least 1 − νT by Lemma 19.
Now, we derive that (a), (b), and (c) are true for t
+ 1 if they are true for t
= τ, τ + 1 , · · · , t .
For t ≥ τ + 1 , we can decompose w t as w t
[Equation] [Equation]
where we use the same notation as the proof of Lemma 19.
According to this decomposition, we verify (a), (b), and (c).
Verifying (a)
The first term of (23) satisfies
[Equation]
Thus, it suffices to bound the norm of η (cid:80) t − 1 s =
τ ( I − ηH ) t − 1 − s
(d H s w s + y s ) by 12 (1 + ηγ ) t − τ r e .
We have
[Equation] [Equation]
For (24), we used the facts that the maximum eigenvalue of ηH is at most ηL ≤ 1 when η ≤ 1 L and that the minimum eigenvalue is − ηγ , which imply ∥ I − ηH ∥ ≤ 1 + ηγ .
(25) follows from the assumptions on ∥ w s ∥ .
For (26), we use t ≤ τ + T 2 and
[Equation]
where the first inequality follows from the hessian Lipschitzness (Assumption 4).
The final inequality (27) holds when we take C 2 as C 2 ≤ 4 C 3 log δ C
ρr e (this is satisfied by taking C 2 =
˜ O
( C 3 )
=
˜ O (1) ).
In addition, we have
[Equation]
Note that (28) can be checked by the same argument as (24) and the inductive hypothesis.
(29) holds when we take C 2 sufficiently large such that 3 C 4 C
[Equation]
14 holds.
Combining (27) and (29), we can bound the second term of (23) as desired, which concludes (a) holds for t ≥ τ + 1 .
Verifying (b)
For t ≥ τ + 1 , we have
[Equation]
As for the first term, we can bound its norm as
[Equation]
The norm of the second term can be bounded by using (a) and (b) for τ +1 , · · · , t − 1 and Lemmaas follows:
[Equation]
Since T 2 = ˜ O (cid:16) 1 ηδ (cid:17) and γ ≥ δ , setting C 2 = ˜ O (1) and η = ˜ O (cid:0) 1 L (cid:1) with sufficiently large hidden constants yields ( ηγT 2 + log T 2 )
[Equation]
≤ γ .
Thus, the second term is bounded by ηγ (1 + ηγ ) t − τ r e .
Finally, we consider the third term.
We have ∥ d H t w t ∥ ≤ δC 2 r e (1 + ηγ ) t − τ r e and ∥ g t ∥ ≤ 3 C 121 C 4 γ C
(1+ ηγ ) t − τ r e by the inductive hypothesis.
Thus, taking C 2 sufficiently large, the third term is bounded by ηγ (1 + ηγ ) t − τ r e .
Combining these bounds, we get (b) for t + 1 .
Verifying (c)
By using Lemma 19 and the inductive hypothesis, we have
[Equation]
(cid:88)
[Equation]
t
[Equation]
T
[Equation]
with probability at least 1 − νT for all t .
Taking b ≥ C 22 ζ 2 δ 2 , η = ˜Θ (cid:0) 1 L (cid:1) , and C
= O ( C 1 C [Equation] =
˜ O (1) gives C
[Equation]
γ .
Thus, we obtain that (c) holds for t
+ 1 .
Thus, we complete the induction step, and hence, the assertion follows.
From Lemma 18, we can ensure that SLEDGE escapes saddle points with high probability.
Lemma 20 Let { x t } be a sequence generated by SLEDGE and τ 0 ( ≥ 0) be a step where − γ :=
λ min ( ∇ 2 f ( x τ 0 )) ≤ − δ holds.
We denote the eigenvector with the eigenvalue λ min ( ∇ 2 f ( x τ 0 )) by e .
We take b ≥ √ n + ζ 2 δ 2 , η = ˜Θ (cid:0) 1 L (cid:1) , and T 2 = C 3 log
[Equation]
2 ρre
[Equation]
≲ ˜ O (cid:0) Lδ (cid:1) with a constant C
= O (1) .
Then, for arbitrary τ > τ 0 , it holds that
[Equation]
Proof
Let A be a subset of B (0 , r ) such that each a ∈ A satisfies
[Equation]
Then, no two elements, ξ τ +1 and ˜ ξ τ +1 such that ξ τ +1 − ˜ ξ τ +1 =
r e e with r e ≥ rνT √ d , can be elements of A at the same time since by Lemma 18, it holds that
[Equation]
with probability at least 1 − νT .
Let V d ( r ) be the volume of Euclidean ball with radius r in R d .
Then, we have
[Equation]
This means that A occupies at least 1 − νT of the volumes of B (0 , r ) .
From this fact and the definition of A , we have
[Equation]
which gives the conclusion.
Then, we move to the proof of the main theorem of this subsection, which guarantees that the algorithm finds ( ε, δ ) -second-order stationary point with high probability.
Proof of Theorem
Since T 2 = C 3 log δ C 2 ρre ηγ depends on x τ 0 (since γ depends on ∇ 2 f ( x τ 0 ) ), we take T 2 =
C 3 log δ C 2 ρre ηδ instead from now.
Note that this replacement does not affect whether
Lemma 20 holds.
We devide { t = 0 , 1 , · · · , T − 1 } into ⌈ T 2 T 2 ⌉ phases : P τ = { 2 τT 2 ≤ t < 2( τ + 1) T 2 } (cid:16) τ = 0 , · · · , ⌈ T 2 T 2 ⌉ − 1 (cid:17) .
For each phase, we define a τ as a random variable defined by a τ = 
 (cid:0) if (cid:80) t ∈ P τ 1 [ ∥∇ f ( x t ) ∥ > ε ] > T 2 (cid:1) , (cid:0) if there exists t such that 2 τT 2 ≤ t < (2 τ + 1) T 2 , ∥∇
Note that P [ a τ ∈ { 1 , 2 , 3 } ] = 1 for each τ .
This is because if there does not exist t between 2 τT 2 ≤ t < (2 τ + 1) T 2 such that ∥∇ f ( x t ) ∥
≤ ε (i.e., neither a τ = 2 nor 3 ), then we have (cid:80) t ∈ P
[Equation]
1 [ ∥∇ f ( x t ) ∥ >
[Equation]
] ≥ (cid:80) (2 τ +1) T 2 − 1 t =
[Equation]
1 [ ∥∇ f ( x t )
[Equation]
> ε ]
= T 2 , meaning a τ = 1 .
We denote
N 1 = (cid:80)
[Equation]
According to Lemma 20, with probability 1 − 2 ν over all τ , it holds that if a τ = 2 then that phase successes escaping saddle points; i.e., there exists 2 τT 2 ≤
[Equation]
2 τ + 1) T 2 such that
[Equation] [Equation]
holds.
(30) further leads to
[Equation]
On the other hand, in Theorem 10, we derived that
[Equation]
+  
[Equation]
Option I )
[Equation]
( Option II ) with probability 1 − 3 ν .
By taking η = ˜Θ (cid:0) 1 L (cid:1)
[Equation]
applying b ≥ √ n and f ( x 0 ) − f ( x t ) ≤ ∆
[Equation]
and rearranging terms
[Equation]
we obtain
[Equation]
Option I )
[Equation]
( Option II ) .
From the definition of a τ = 1 and (31)
[Equation]
that the left-hand side is bounded as
T
[Equation]
Thus
[Equation]
it holds that max
[Equation]
By the parameter settings, we have 2 η 2 T 22 C 22 ρ 2 δ= ˜ O (cid:16) ρ 2 δ 4 (cid:17) .
From this, ( N 1 + N 2 ) T 2 ≤ ˜ O (cid:16)
ε 2 + ρ 2 δ 4 (cid:17) ×
( the right-hand side of (32) ) .
Taking T ≥ 2( N 1 + N 2 +
T 2 , there exists τ such that a τ = 3 , which concludes the proof.
Remark
Although our main interest in this paper is to develop a simple algorithm with convergence to second-order stationary points, it can be easily shown that adaptive selection of minibatch size can reduce the gradient complexity.
In Lemma 18, if we carefully check the proof, we can see that the condition b ≳ √ n + ζ 2 δ 2 is needed only for the step τ .
On the other hand, for all τ 0 ≤ t ≤ τ + T 2 except for t = τ , b = √ n is sufficient.
If we take b = √ n + ζ 2 δ 2 only at t = (2 τ + 1) T 2 (cid:16) τ = 0 , · · · , T 2 T 2 − 1 (cid:17) and b = √ n at the other steps, the above argument still holds with a slight modification.
Then, the gradient complexity is reduced to
[Equation]
Option
[Equation]
Option II )
[Equation]
In the classical setting δ = O ( √ ρε ) , this bound is no worse than SPIDER-SFO + (+Neon2) [2, 10], no matter what n and δ are.
Finally, we note that if δ is too small, L 2 δ 2 can be as large as n .
In such case, it is more efficient to replace sampling such number of samples is replaced by full gradient computation.
Then, the complexity gets
[Equation]
II )
[Equation]
When δ =
O ( √ ρε )
, this bound is no worse than NestedSVRG+Neon2 [2, 50].
However, it is unusual to assume
L 2 δ 2 = n in the first place.
In fact, carefully looking the proof of SSRGD [26], we find that they implicitly limits their analysis to the case of L 2 δ 2 ≲ n .
D.4. Convergence under PL condition (proof of Theorem 12)
In this subsection, we provide the proof of the convergence under Assumption 6, i.e., PL condition holds for the objective function.
Proof
According to the descent lemma (Lemma 14) and PL condition (Assumption 6), we have that
[Equation]
1 2 η − L
[Equation]
∥ x t − x t − 1 ∥
[Equation]
r 2 η
[Equation]
Rearranging the terms yields
[Equation]
By applying Lemma 13 to this
[Equation]
we obtain that with probability at least 1 − 3 ν ,
[Equation]
holds for all t = 1 , · · · , T .
Multiplying both sides by (1 − ηµ ) T − t and summing up over all t = 1 , 2 , · · · , T and arranging the terms, we get [Equation] II
[Equation]
Note that T 1 = nb C 1 .
According to this, we take η as
[Equation]
so that 12 η − L 2 − 15 C 91 ηζ 2 n (1 − ηµ ) − T 1 b≥ 0 holds.
Then, we have that
[Equation]
[Equation]
Option
[Equation]
The first term (1 − ηµ ) T ( f ( x 0 ) − f ∗ ) is smaller than ε 3 if we take T
= O (cid:16) 1 ηµ log ∆ ε (cid:17) .
The second term is bounded by r 2 η 2 µ , which is smaller than ε 3 if we take r ≤ η (cid:113) εµ 3 .
The third term for Option I,
[Equation] [Equation]
, is also bounded by ε 3 , if we take T
= T 1 + O (cid:16)
[Equation]
Thus, for Option I, taking
[Equation]
yields the desired bound with probability at least 1 − 3 ν .
And for Option II, taking
[Equation]
yields the desired bound.
Note that T depends on ε − 1 only logarithmically, which means that C 1 depends on ε − 1 in only log log order and C
= O ∗ (log( n + µ − 1 + ν − 1 )) , where O ∗ suppresses log log factors.
Appendix E. Missing Statements and Proofs for FLEDGE
This section provides the missing information of FLEDGE that we abbreviate in section 3 and gives the proofs of the theorems on FLEDGE about the convergence property of FLEDGE.
First, we provide the full version of FLEDGE in the following, including Option I. Note that B (0 , r ) is the uniform distribution on the Euclidean ball in R d with radius r .
As the case of SLEDGE, we divide Theorem 2 into Theorem 22 (first-order optimality), Theorem 28 (second-order optimality),
Theorem 34 (exponential convergence under the PL condition), each of which will be presented in one of the following subsections in order.
E.1. Finding First-order Stationary Points (Proof of Theorem 22)
In this subsection, we show that Algorithm 2 finds first-order stationary points with high probability.
First, we describe the formal statement of Theorem 22.
Theorem 22 Let r ≤ ηε 2 √ 2 and P Kb ≥ ˜Ω (cid:16) σ 2 ε 2 + Gε (cid:17) .
Under Assumptions 1 to 3 and 5, if we choose
[Equation]
1: Option I: 2:
Randomly select p agents I3: for i ∈ I 0 in parallel do 4:
Randomly select bK samples J 0 i 5: y
i ← 1 bK (cid:80) j ∈ J
i ∇ f i,j ( x 0 ) 6:
Communicate { y 0 i } i ∈
I 0 between I7: y
i ← 1 p (cid:80) i ∈ I
y
i ( i = 1 , · · · , n ) //
we do not need to explicitly communicate this between all the clients 8: Option II: 9: for i ∈
I
= I in parallel do 10:
Randomly select bK samples J 0 i 11: y
i ← 1 bK (cid:80) j ∈ J
i ∇ f i,j ( x 0 ) 12: for t = 1 to T do 13:
Randomly sample one agent i t 14:
Communicate { 1 P (cid:80) Pi =1 y t − 1 i } , and x t − 1 between I t − 1 ∪ { i t } and the server 15: x t, 0 ← x t − 1 , z t, 0 ←16: for k = 1 to K do 17: x t,k ← x t,k − 1 − (cid:32)
P
P (cid:88) i =
y t − 1 i + z t,k −(cid:33) + ξ t,k ( ξ t,k ∼ B (0 , r )) 18: randomly select b samples J t,ki t 19: z t,k ← z t,k − 1 + 1 b (cid:80) j ∈ J t,kit ( ∇ f i t ,j ( x t,k ) − ∇ f i t ,j ( x t,k − 1 )) 20: x t ← x t,K 21:
Randomly select p agents I t 22:
Communicate x t between I ∪ { i t } 23: for i ∈
I t in parallel do 24:
Randomly select b samples J ti 25: y ti ← 1 bK (cid:80) j ∈ J ti ∇ f i,j ( x t ) 26:
∆ y ti ← 1 bK (cid:80) j ∈ J ti ( ∇ f i,j ( x t ) − f i,j ( x t − 1 )) 27:
Communicate { ∆ y ti } i ∈
I t between I t and the server 28: y ti ← y t − 1 i + 1 p (cid:80) i ∈ I t ∆ y ti ( for i / ∈
I t ) //
Practically, we update only 1 P (cid:80) Pi =
y ti in the server in O ( p ) time.
Algorithm 2 with Option I finds an ε -first-order stationary point for problem (2) by using ˜ O (cid:32)(cid:32)
[Equation]
stochastic gradients and ˜ O
[Equation]
communication rounds with probability at least 1 − 8 ν .
Moreover, under the same assumptions, Algorithm 2 with Option II finds an ε -first-order stationary point for problem (2) by using
[Equation]
stochastic gradients and
[Equation]
communication rounds with probability at least 1 − 8 ν .
Let v t,k − 1 = 1 P (cid:80) Pi =1 y t − 1 i + z t,k − 1 and K ( t ) be the last inner loop step in the t -th outer loop as stated in the algorithm.
The descent lemma (Lemma 14) also works here: as was discussed for
Algorithm 1, for each t and k (1 ≤ t ≤ T, 1 ≤ k ≤ K ( t )) , it holds that
[Equation]
Our strategy is to bound the variance term ∥ v t,k − ∇ f ( x t,k ) ∥ 2 with high probability, as summarized in the following lemma.
Lemma 23 Let v t,k = 1 P (cid:80) Pi =1 y t − 1 i + z t,k and all the other variables be as stated in algorithm 3.
Then, with taking T 3 =
Pp C 1 , we have (cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
≤
[Equation]
Option I ) ( Option II ) for all t, k (1 ≤ t ≤ T, 0 ≤ k ≤ K − 1) , with probability at least 1 − 8 ν .
For the proof of Lemma 23, we utilize the four following auxiliary lemmas.
Below, we define ˜ y
i by
[Equation]
( Option I )
[Equation]
( Option II ) .
As well as the previous section, we define ˜ I
ts :=
[ n ] \ (cid:83) t τ = s I t for 1 ≤ s ≤ t .
In addition, for each s, t ( s ≤ t ) and i ∈ [ P ] , we let T 4 ( t, i ) as T 4 ( t, i ) :
=
max { s | s = 0 or 1 ≤ s ≤ t with s ∈
I s } , i.e., the last step when y si is updated before t .
We remark that the setting T 3 =
Pp C 1 gives ˜ I ts = ∅ with probability at least 1 − ν for all t and s ≤ t − T 3 .
Lemma 24 With probability at least 1 − ν , the following holds for all t
= 1 , · · · , T :
[Equation] [Equation]
Proof
First, we decompose the left hand side as
[Equation]
[Table]
i ∈
I s (cid:88)
j ∈ J si ( ∇ f i,j ( x s ) − ∇ f i,j ( x s − 1 ) − ( ∇ f i
[Table]
s
[Table]
f i (
[Table]
− 1 )))
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
= 2 (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) t (cid:88) s =max { 1 ,
[Table]
T 3 +1 }
E [ | ˜ I ts | ]
PpKb (cid:88)
i
[Table]
I s (cid:88) j ∈ J si ( ∇ f
[Table]
,j
[Table]
f
[Table]
[Table]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
+t (cid:88) s =max { 1 ,t
[Table]
T
[Table]
+1 }
C
[Equation]
(cid:0)
[Table]
˜
[Equation]
ts |
[Equation]
[ |
[Equation]
ts | ] (cid:1)
[Table]
Pp 3 K
[Equation]
b(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:88)
i ∈
I s (cid:88)
j ∈ J si ( ∇ f i,j ( x s ) − ∇ f i,j ( x s − 1 ) − ( ∇ f i ( x s ) − ∇ f i ( x s − 1 ))) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) .
To bound the first term, by applying Proposition 4 to the choice of I s and J si , we have (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:88)
[Equation] [Equation]
with probability at least 1 − ν 4 T 2 .
Then
[Equation]
we use Proposition 7 to obtain
[Equation]
[Equation]
with probability at least 1 − ν 4 T − T · ν 4 T 2 = 1 − ν 2 T .
For the second term, following the same argument in Lemma 16, we can show that | ˜ I
ts −
E [˜ I ts ] | 2 ≤ C 51 P with probability at least 1 − ν 4 T 2 for each s, t .
Combining this with (35), we have
[Equation] [Equation]
with probability at least 1 − T ·
[Equation] [Equation]
− T
[Equation]
−
[Equation]
T
[Equation]
Finally, substituting (36) and (37) for (34), we obtain the assertion.
Lemma 25 With probability at least 1 − ν , the following holds for all t
= 1 , · · · , T :
[Equation]
P
[Equation]
x
[Equation]
t
[Equation]
i ) )
[Equation]
t
[Equation]
i )
[Equation]
) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ≤ C
P Kb (cid:18) σ 2 + G
P Kb (cid:19) .
Proof
We condition the events on { T 4 ( i, s ) } and apply the Bernstein’s inequality to obtain the desired bound.
Lemma 26 With probability at least 1 − ν , the following holds for all t
= 1 , · · · , T :
[Equation]
Option I )
[Equation]
( Option II )
Proof
Recall the definition of ˜ y
i :
Option I
By conditioning I 0 , Proposition 4 yields that
[Equation] [Equation]
Option II
In this case, Proposition 4 directly yields that
[Equation]
∈ J
i
[Equation] [Equation]
with probability at least 1 − νT for each t .
Lemma 27 With probability at least 1 − ν , the following holds for all t
= 1 , · · · , T − 1 and k = 1 , · · · , K − 1 :
[Equation]
Proof
We decompose the z t,k − ( ∇ f ( x t,k ) − ∇ f ( x t, 0 ))
as ∥ z t,k − ( ∇ f ( x t,k ) − ∇ f ( x t, 0 )) ∥≤ 2 ∥
z t,k − ( ∇ f i t ( x t,k ) −
[Equation]
≤ 2 (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) k (cid:88) l =
[Equation]
(cid:88)
[Equation]
≤ 2 (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) k (cid:88) l =
[Equation]
(38) where we use Assumption 5 for the second inequality.
We apply Proposition 4 to 1 b (cid:80) j ∈ J t,lit ( ∇ f i t ,j ( x t,l ) − ∇ f i t ,j ( x t,l − 1 )) and obtain that (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)b (cid:88) j ∈ J t,lit ( ∇ f i t ,j ( x t,l )
[Equation]
with probability at least 1 − 1 2 TK 2 for each ( t, l ) .
Using Proposition 7, with probability 1 − K ·2 TK 2 − 1 2 TK = 1 − 1 TK , we have
[Equation]
[Equation]
for each ( t, k ) .
By substituting (39) to (38), we obtain the desired bound.
Proof of Lemma
First, we observe that
[Equation]
We first bound (cid:13)(cid:13)(cid:13)
[Equation]
Similarly to Lemma 13, with probability at least 1 − ν , ˜ I ts = ∅ holds for all s ≤ t − T 3 and we can expand 1 P (cid:80) Pi
[Equation]
1 y ti − ∇ f ( x t ) as
P
[Equation]
P
[Equation]
[Table]
for all t .
The norm of the part (a) can be bounded by using Lemma 13, just replacing n by P , i.e.,
[Equation]
for Option I and
[Equation]
for Option II, with probability at least 1 − 3 ν for all t .
For the bound of (b), (c) and (d), we apply
Lemma 24, Lemma 25, and Lemma 26, respectively.
Then, by summarizing all these and using ∥ x s − x s − 1 ∥ 2 ≤ K (cid:80) Kl
=1 ∥ x s,l − x s,l − 1 ∥ 2 , we get
[Equation]
σpKb +
Gp 2 K 2 b 2 + σ 2 c p +
G 2 c p(cid:19) ≤ (cid:18) 60 C 81 ζ 2 K p + 16 C 101 ζpb (cid:19) t (cid:88) s =max { 1 ,t − T 3 +1 }
K (cid:88) l = ∥ x s,l − x s,l − 1 ∥ 2 + 4 C
P Kb (cid:18) σ 2 + G
P Kb (cid:19) + 48 C
[Equation]
1 [ t
[Equation]
T
[Equation]
] (cid:18)
[Equation]
pKb
[Equation]
Gp 2 K 2 b 2 + σ 2 c p +
G 2 c p(cid:19) for Option
I and (cid:13)(cid:13)(cid:13)(cid:13) (cid:13)
P
P (cid:88) i =
y ti
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13) (cid:13)
[Equation]
≤ (cid:18)
[Equation]
p
[Equation]
pKb (cid:19)
[Equation]
max
[Equation] [Equation]
[Equation]
Kb (cid:18)
[Equation]
P Kb (cid:19) ≤ (cid:18)
[Equation]
p
[Equation]
(cid:88)
[Equation]
max
[Equation] [Equation]
t
[Equation]
T
[Equation]
(cid:88) l
[Equation]
P Kb (cid:18)
[Equation]
P Kb (cid:19) for Option II, with probability 1 − 7 ν for all t .
Also, we have ∥
[Equation]
Lemma 27, with probability 1 − ν over all t, k .
By substituting these bound to (40), we obtain the desired bound.
Now
[Equation]
we are ready to prove the first-order convergence of FLEDGE
[Equation]
Proof of Theorem
Summing up (33) over all t and k and rearranging the terms, we get
[Equation]
Applying Lemma 23 to this
[Equation]
we have that
[Equation]
( Option I )
[Equation]
( Option II ) with probability at least 1 − 8 ν .
Option I
We set η as
[Equation]
so that 12 η − L 2 − η (cid:16) 152 C 111 ζ 2 PK p 2 b
[Equation]
L
[Equation]
(cid:17) ≥ 0 holds.
By taking r ≤ ηε 2 √ 2 and P Kb ≥ 128 C 21 σε+ 8 √ 2 C 1 G ε , we obtain 2 TKr 2 η≤ TKε
2and 16 C 21 T Pb (cid:16) σ 2 + G 2 PKb (cid:17) ≤ TKε 2, respectively.
Moreover, we apply f ( x 0 ) − f ( x t ) ≤ ∆ .
Summarizing these, we get
T (cid:88) t =
K (cid:88)
[Equation]
Hence, taking
T K ≥ 4∆ ηε
+
[Equation]
˜ O
[Table]
[Equation]
results in
[Equation]
which implies that FLEDGE can find ε -first order stationary points with probability at least 1 − 8 ν .
Thus, the gradient complexity and the communication complexity are bounded as stated.
Option II
We set η as the same as that for Option I, so that 12 η − L 2 − η (cid:16)
[Equation] [Equation]
≥ 0 holds as well.
We take r ≤ ηε 2 √ 2 and P Kb ≥ 128 C 21 σ 2 ε
[Equation]
1 G ε .
Then, we get
[Equation]
Therefore, by the similar argument to Option I, taking T K ≥ 2∆ ηε 2 =
˜ O (cid:16)(cid:16)
L ∨ ζ
[Equation]
∨ L √ K √ b (cid:17) ∆ ε 2 (cid:17) ensures that FLEDGE finds ε -first order stationary points with probability at least 1 − 8 ν .
E.2. Finding Second-order Stationary Points (Proof of Theorem 28)
Here, we show that FLEDGE can efficiently find second-order stationary points.
With a slight abuse of notations, we sometimes identify ( t, k ) with ( t ′ , k ′ )
when tK + k = t ′ K + k ′ holds.
Moreover, we say ( t 1 , k 1 ) > ( t 2 , k 2 ) when t 1 K + k 1 > t 2 K + k 2 .
First, we state the formal theorem as follows:
Theorem
We assume Assumptions 1 to 5, and δ < 1 ζ .
Let p ≥ √ P + C 25 ζ 2 δ+ C 25 L 2 Kbδ 2 with
C 5 =
˜ O (1) , b ≥ K , K = ˜
O (cid:16) Lζ (cid:17) ,
η = ˜Θ (cid:0) 1 L (cid:1) , r = ˜
O (cid:0) εL (cid:1) , P Kb ≥ O (cid:16) σ 2 ε 2 + Gε (cid:17) and ν ∈ (0 , 1) .
Then, FLEDGE with Option I finds ( ε, δ ) -second-order stationary points using ˜ O (cid:18)(cid:18)
L ∆ +
[Equation]
stochastic gradients and ˜ O (cid:18)(cid:18)
[Equation]
(cid:19)(cid:19) communication rounds , with probability at least 1 − 12 ν .
Moreover, FLEDGE with Option II finds ( ε, δ ) -second-order stationary points using ˜ O (cid:18)
[Equation]
pKb (cid:19) stochastic gradients and ˜ O
[Equation]
(cid:19)(cid:19) communication rounds , with probability at least 1 − 12 ν .
Similarly to the previous section, the key argument is the exponential separation of two coupled trajectories with different initial values.
Lemma 29 (Small Stuck Region) Assume δ < 1 ζ .
Let { x t,k } be a sequence generated by FLEDGE and ( τ 0 , κ 0 ) (0 ≤ κ 0 < K ) be a step where − γ :=
λ min ( ∇ 2 f ( x τ 0 ,κ 0 ))
≤ − δ holds.
We denote the eigenvector with the eigenvalue λ min ( ∇ 2 f ( x τ 0 ,κ 0 )) by e .
Moreover, let { ˜ x t,k } by a coupled sequence that is generated by FLEDGE with ˜ x 0 = x 0 and shares the same choice of randomness with { x t } i.e., client samplings, minibatches and noises, except for the noise at a step ( τ 0 , K ) > ( τ 0 , κ 0 ) :
˜ ξ τ 0 ,K = ξ τ 0 ,K − r e e with r e ≥ rν TK √ d .
Let w t,k = x t,k − ˜ x t,k , w t
= x t − ˜ x t , v t,k = 1 P (cid:80) Pi =1 y t − 1 i + z t,k , ˜ v t,k = 1 P (cid:80) Pi =
˜ y t − 1 i + z t,k , g t = 1 P (cid:80) Pi =
y ti − ∇ f ( x t ) − (cid:16) 1 P (cid:80) Pi =
˜ y ti − ∇ f
(˜ x t ) (cid:17) , and h t,k = ( z t,k − ( ∇ f ( x t,k ) − ∇ f ( x t, 0 )))
− (˜ z t,k − ( ∇ f (˜ x t,k ) − ∇ f (˜ x t, 0 ))) .
Using these notations, v t,k − ∇ f ( x t,k ) − (˜ v t,k − ∇ f (˜ x t,k ))
=
g t − 1 + h t,k holds.
Then, there exists a sufficiently large constants C 5 =
˜ O (1) and C
= O (1) with which the following holds: If we take p ≥ √ P + C 2 5 ζ 2 δ+ C 2 5 L 2 Kbδ 2 , b ≥ K, K
= O ( Lζ ) , η =
˜ Θ (cid:0) 1 L (cid:1) , and
T 5 = C 6 log δ C 5 ρre ηγ ≲ ˜ O (cid:0) Lδ (cid:1) , with probability 1 − 3 ν TK ( ν ∈ (0 , 1)) , we have
[Equation]
In order to show Lemma 29, we prepare the two following lemmas, which bound the difference between gradient estimation errors of the two sequence.
Lemma
Under the same assumption as that of Lemma 29, we assume max ( τ 0 ,κ 0 ) ≤ ( t,k ) <
( τ 0 +1 ,T 5 ) {∥ x τ,k − x τ 0 ,κ 0 ∥ , ∥
˜ x τ,κ − x τ 0 ,κ 0 ∥} < δC 5 ρ .
Then, the following holds uniformly for all ( τ 0 , κ 0 ) ≤ ( t, k ) ≤ ( τ 0 + 1 , T 5 ) with probability at least 1 − νTK :
[Equation]
≤  
[Equation]
where T 3 =
Pp C 1 , and C 7 = ˜
O
(1) is a sufficiently large constant.
Lemma
Under the same assumption as that of Lemma 29, the following holds uniformly for all t ≥ τ 0 + 1 and k ≥ 0 with probability at least 1 − 2 ν TK :
[Equation]
For t < τ 0 + 1 , we have (cid:13)(cid:13) h t,k
(cid:13)(cid:13)
= 0 .
Proof of Lemma
As for the case t < τ 0 , the assertion directly follows from the definition of { ˜ x t,k } .
For the proof of the rest cases, we use notations as follows:
[Equation]
Moreover, we denote
[Equation]
) and
[Equation]
Note that E i [ u si ] = 0 (expectation with respect to the choice of i ) and E j [ u si,j ] = 0 (expectation with respect to the choice of j ) hold.
Using Assumptions 1, 4 and 5 and max ( τ 0 ,κ 0 ) ≤ ( t,k )
<
( τ 0 +1 ,T 5 ) {∥ x τ,k − x τ 0 ,κ 0 ∥ , ∥
˜ x τ,κ − x τ 0 ,κ 0 ∥} < δC 5 ρ , we can derive that
[Equation]
and
[Equation]
for s ≥ τ 0 + 1 , by similar argument to the proof of Lemma 19.
For t = τ 0 , we have ∥ u τ 0 i ∥ =
∥ ( ∇ f i ( x τ 0 ) −∇
f i (˜ x τ 0 )) − ( ∇ f ( x τ 0 )
−∇ f
(˜ x τ 0 ))
∥ ≤ ζ ∥ x τ 0 − ˜ x τ 0 ∥ = ζr e and ∥ u τ 0 i,j ∥ = ∥ ( ∇ f i,j ( x τ 0 ) − ∇ f i,j (˜ x τ 0 )) − ( ∇ f ( x τ 0 ) − ∇ f
(˜ x τ 0 ))
∥ ≤ L ∥ x τ 0 − ˜ x τ 0 ∥ =
Lr e .
As we did in Lemma
[Equation]
for t ≥ τ 0 +
[Equation]
we have
[Equation]
=
[Equation]
 
[Equation]
+
P Kb
[Equation]
with probability 1 − ν 8 TK for all t .
For t = τ 0 , g τ 0 = (a) + (b) holds.
Recall the argument in Lemma 19.
We have that ∥ (a) ∥ ≤ 2 C 1 ζr e √ p and
[Equation]
hold with probability at least 1 − 1 8 TK for all t .
Moreover, observe that (b) and (d) are obtained just by replacing u ti in (a) and (c) by 1 Kb (cid:80) j ∈ J ti u si,j .
Note that 1 Kb (cid:80) j ∈ J τ 0 i u τ 0 i,j is mean-zero and its norm is bounded by C 1 Lr e √ Kb for s = τ 0 , with probability 1 − 1 8 T 2 K .
Thus, Proposition 7 yields that
∥ (b) ∥ ≤ (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Table]
P | ˜ I tτ
| p
[Table]
i ∈
[Table] [Table]
j ∈
[Table]
τ
i
u
τ
i
[Equation]
[Table] [Table]
j ∈ J τ
i
[Equation]
[Table]
i
[Equation]
with probability 1 − 1 8 TK − T · 1 8 T 2 K = 1 − 1 4 TK , where we use | ˜
I t τ 0 | ≤ P and p ≤ P for the last inequality.
For the first term of (d), we first observe that 1 Kb (cid:80) i ∈ I s (cid:80) j ∈ J ti u si,j is mean-zero and its norm is bounded by C 1 √ p √ Kb L
∥ w s − w s − 1 ∥ + 2 C 1 δ √ p C
[Equation]
√ Kb
∥
w s − 1 ∥ (for s ≥ τ 0 + 1 ) with probability at least 1 − 1 8 T 2 K .
Then, we apply the same argument as Lemma 16.
This yields
[Equation]
with probability 1 − T · ν 8 T 2 K − ν 8 TK = 1 − ν 4 TK for all t .
As for the second term of (d), by applying
Proposition 6, we get (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
Kb t (cid:88) s
[Equation]
+
[Equation]
t −
[Equation]
1 +1 } (cid:88)
[Equation]
ts (cid:88)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:118)(cid:117)(cid:117)(cid:117)(cid:116)
[Table]
[Equation]
[Table]
Kb (cid:88)
[Table]
(cid:88)
[Table]
[Equation]
[Table] [Table]
[Equation]
(cid:118)(cid:117)(cid:117)(cid:116)
[Equation]
(cid:88)
[Equation]
max
[Equation] [Equation]
[Equation] [Equation]
pKb (cid:118) (cid:117) (cid:117)(cid:116)
[Equation]
(cid:88) s
[Equation]
with probability 1 − T ·
[Equation]
8 T 2 K −
[Equation]
for all t .
By combining all these, we have
[Equation]
s
[Equation] [Equation]
∥ w
[Equation]
∥
[Equation]
with probability at least 1 − ν TK for all t ≥ τ 0 + 1 .
Here we take C 7 = ˜ O (1) , which is independent of C 5 .
Thus, we get the assertion for t ≥ τ 0 + 1 .
For t = τ 0 , the bounds on (a) and (b) imply the desired bound
[Equation]
Proof of Lemma
Let
[Equation]
and
[Equation]
By their definitions, h t,k = u t,li t + 1 b (cid:80) kl
=1 (cid:80) j ∈ J it u t,li t ,j holds.
We can bound the norm of them as
[Equation]
and
[Equation]
Thus, applying Proposition 4 and Proposition 7 to 1 b (cid:80) kl =1 (cid:80) j ∈ J it u t,li t ,j , we get (cid:13)(cid:13)(cid:13) (cid:13) (cid:13)(cid:13)
b k (cid:88) l
= (cid:88) j ∈ J it
[Equation]
t,l i t ,
[Equation]
(cid:13)(cid:13)(cid:13) (cid:13) (cid:13)(cid:13) ≤
[Equation]
(cid:88)
[Equation]
with probability at least 1 − 1 TK for all t and K .
Substituting (41) and (42) to h t,k = u
[Equation]
kl
[Equation]
J
[Equation]
u t
[Equation]
li t ,j , we get the desired bound.
Now, we are ready to prove Lemma 29.
Proof of Lemma
We assume the contrary and show the following by induction, for ( τ 0 + 1 , 0) ≤ ( t, k ) ≤ ( τ 0 + 1 , T 5 ) :
[Equation]
Here C
= ˜ O (1) is a sufficiently large constant independent of C 5 .
Then, (a) yields contradiction by taking ( t, k ) − ( τ
+ 1 , 0)
=
T
= O (cid:18)
+ log δ
[Equation]
2 ρre ηδK (cid:19) to break the assumption.
It is easy to check (a) and (b) for and t = τ 0 +
and k = 0 .
As for (c), checking the initial condition at ( t, k ) =
( τ 0 + 1 , 0) requires assumption on the size of p .
According to Lemma 30, taking p ≥ ζ 2 δ 2 + L 2 δ 2 Kb , ∥ g τ 0 ∥ ≤ 2 C 7 δr e ≤ 2 C 7 γr e holds.
Now, we derive that (a), (b) and (c) are true for ( t, k + 1) , assuming that they are true for all ( τ 0 + 1 , 0) , · · · , ( t, k ) .
To this end, we consider the decomposition of w t,k as follows:
[Equation]
for ( t, k + 1) ≥ ( τ 0 + 1 , 1) .
Verifying (a)
The first term (1 + ηγ ) ( t − τ 0 − 1) K + k +1 r e e of (43) satisfies ∥ (1 + ηγ ) ( t − τ 0 − 1) K
+ k +1 r e e ∥
=
(1 + ηγ ) ( t − τ 0 − 1) K + k +1 r e .
Then, focus on bounding η (cid:80) t,k
( s,l )=( τ 0 +1 , 0) ( I − ηH ) ( t − s ) K +( k − l ) (d H s,l w s,l + g s − 1 + h s,l ) by 12 (1+ ηγ ) ( t − τ 0 − 1) K
+ k +1 r e .
We have (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) η ( t,k ) (cid:88) ( s,l )
[Equation]
( τ 0 +1 , 0) (
[Equation]
l )
[Equation]
,l
[Equation]
,l (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ≤ η ( t
[Equation]
(cid:88) (
[Equation]
+
[Equation]
, 0) ∥
[Equation]
−
[Equation]
t
[Equation]
+
[Equation]
k
[Equation]
l
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
s,l (cid:13)(cid:13)(cid:13) (cid:13)(cid:13)(cid:13) w s,l (cid:13)(cid:13)(cid:13) ≤ 2 η (1 + ηγ ) ( t − s ) K
+( k − l )+( s − τ 0 − 1) K + l r e t,k (cid:88) ( s,l )=( τ 0 +1 , 0) (cid:13)
[Equation]
The last inequality follows from the definition of T 5 = C 6 log δ C
5 ρre ηγ and sufficiently large C 5 .
In addition, we have
[Equation] [Equation]
For the final inequality, we again use T 5 =
C 6 log
[Equation]
ηγ with sufficiently large C 5 .
Combining (44) and (45), we get (a) for ( t, k + 1) as desired.
Verifying (b)
For ( t, k ) ≥ ( τ 0 + 1 , 0) , we have
[Equation]
As for the first term, we can bound it as
[Equation]
Evaluating the second term requires (a) and (b) for ( τ 0 + 1 , 0) , · · · , ( t, k − 1) and Lemma 9: (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) η
[Table] [Table]
) (cid:88)
[Table]
τ 0 − 1) K + k r e .
Since T 5 = ˜ O (cid:16) 1 ηδ (cid:17) and γ ≥ δ , setting C 5 = ˜ O (1) with sufficiently large C 5 yields ( ηγT 5 + log T 5 ) (cid:16) δC 5 + 2 C 8 γ C 5 (cid:17) ≤ γ .
Thus, the second term is bounded by ηγ (1+ ηγ ) ( t − τ 0 − 1) K + k r e .
Finally, we consider the third term.
We have ∥ d H t,k w t,k ∥ ≤ δC 5 r e (1 + ηγ ) ( t − τ 0 − 1) K + k r e and ∥ g t − 1 + h t,k ∥ ≤ 2 C 8 γ C 5 (1 + ηγ ) ( t − τ 0 − 1) K + k r e .
Thus, by taking C 5 sufficiently large, the third term is bounded by ηγ (1 + ηγ ) ( t − τ 0 − 1) K + k r e .
By combining these bounds, we get (b) for ( t, k + 1) .
Verifying (c)
Using Lemma 30 and assumptions
[Equation]
we have
[Equation]
with probability at least 1 − νTK for all t .
Taking p ≥ √ P + C 25 ζ
[Equation]
C
[Equation]
L
[Equation]
Kb
[Equation]
) , b ≥ K, K =
O (cid:16) Lζ (cid:17) , and ∥ g t +1 ∥
≤ C 8 γ C 5 (1 + ηγ ) ( t − τ 0 ) K with sufficiently large constant C 8 , that only depends on C 1 , C 7 , and sufficiently small η = ˜Θ( 1 L ) .
Moreover, Lemma 31 states that, for k < K ,
[Equation]
holds with probability at least 1 − νTK .
If (a) and (b) hold for all ( s, l ) ≤ ( t, k + 1) , then we have
[Equation]
Taking b ≥ K and K
= O (cid:16) Lζ (cid:17) , with sufficiently large C 8 and sufficiently small η , we have ∥ h t,k +1 ∥ ≤ C 8 γ C 5 (1 + ηγ ) ( t − τ 0 − 1) K + k +1 .
Thus, we obtain that (c) holds for ( t, k + 1) .
Therefore, we have completed the induction step and have 12 (1 + ηγ ) ( t − τ 0 − 1) K + k r e ≤ ∥ w t ∥ for all ( τ 0 + 1 , 0) ≤ ( t, k ) <
( τ 0 + 1 , T 5 ) with T 5 =
C 6 log δ C
5 ρre ηγ .
Taking C 6 sufficiently large, we have 12 (1 + ηγ ) ( τ 0 +1 − τ 0 − 1) K + T 5 r e ≥ δC 5 ρ .
This yields contradiction against the assumption and the desired assertion follows.
From Lemma 29, we can show that FLEDGE escapes saddle points with high probability.
We have the following lemma, and the proof is essentially the same as that of Lemma 20.
Lemma 32 Let { x t,k } be a sequence generated by FLEDGE and ( τ 0 , κ 0 ) (0 ≤ κ 0 < K ) be a step where − γ :=
λ min ( ∇ 2 f ( x τ 0 ,κ 0 ))
≤ − δ holds.
We take p ≥ √ P + C 25 ζ 2 δ+ C 25 L 2 Kbδ 2 , b ≥ √ K and, η = ˜Θ (cid:0) 1 L (cid:1) , and T 5 = C 6 log
[Equation]
[Equation]
≲ ˜ O (cid:0) Lδ (cid:1) , with sufficiently large C 5 , C 6 = ˜ O (1) .
Then,
P
[Equation]
Finally, we show the main theorem of this subsection, which guarantees that the algorithm finds ( ε, δ ) -second-order stationary point with high probability.
Proof
of
Theorem
Since T 5 = C 6 log δ C 5 ρre ηγ depends on x τ 0 , we take T 5 =
C 6 log δ C
5 ρre ηδ from now instead.
This change does not affect whether Lemma 32 holds.
Also, we let T 6 = ⌈ 1 + T 5 K ⌉ .
We divide { t = 0 , 1 , · · · , T − 1 } into the following ⌊ T 2 T 6 ⌋ phases : P τ = { 2 τT 6 ≤ t < 2( τ + 1) T 6 } (cid:16) τ = 0 , · · · , ⌊ T 2 T 6 ⌋ − 1 (cid:17) .
For each phase, we define a τ as a random variable taking values
[Equation]
    (cid:16) if (cid:80) t ∈ P τ (cid:80)
[Equation]
(cid:32)
if there exists t such that (2 τT 6 , 0) ≤ ( t, k ) < ((2 τ + 1) T 6 , 0) , ∥∇ f ( x t,k ) ∥ ≤ ε and λ min ( ∇ 2 f ( x t,k ))
≤ − δ
(cid:33)
(cid:32) if there exists t such that (2 τT 6 , 0) ≤ ( t, k ) <
((2 τ + 1) T 6 , 0) , ∥∇ f ( x t,k ) ∥ ≤ ε and λ min ( ∇ 2 f ( x t,k )) >
− δ
(cid:33) .
Note that P [ a τ = 1 , 2 , 3] = 1 for each τ .
This is because if there does not exist t between (2 τT 6 , 0) ≤ ( t, k ) <
((2 τ + 1) T 6 , 0)
such that ∥∇ f
x t,k ) ∥ ≤ ε (i.e., neither a τ = 2 nor 3 ), then we have (cid:80) t ∈ P τ (cid:80)
Kk =0 1 [ ∥∇ f
[Equation]
meaning a τ = 1 .
We denote N 1 = (cid:80)
⌊ T 2 T 6 ⌋ τ =0
[ a τ = 1] , N 2 = (cid:80)
⌊ T 2 T 6 ⌋ τ =0
[ a τ = 2] , and N 3 = (cid:80) ⌊ T 2 T 6 ⌋ τ =0
[ a τ = 3] .
According to Lemma 32, with probability 1 − 4 ν over all τ , it holds that if a τ = 2 then that phase successes escaping saddle points; i.e., there exists (2 τT 6 , 0) ≤ ( t, k ) <
((2 τ + 1) T 6 , 0) and
[Equation]
holds.
Eq.
(46) further leads to
[Equation]
On the other hand
[Equation]
in Theorem
[Equation]
we derived that
[Equation]
+  
[Equation]
( Option II ) with probability 1 − 8 ν .
Taking η = ˜Θ (cid:0) 1 L (cid:1) sufficiently small, applying p ≥ √ P , K
= O (cid:16) Lζ (cid:17) ,
K ≤ b and f ( x 0 ) − f ( x t ) ≤ ∆ , and arranging terms yields
T (cid:88)
[Equation]
From the definition of a τ = 1 and (47), We know that (48) is bounded as
T
[Equation]
Thus, N 1 T 6 K ≤ 1 ε 2 × ( the RHS of (49) ) and N 2 T 6 K ≤ 2 η 2 C 22 ρ 2 K 2 T 26 δ× ( the RHS of (32) ) holds.
Here, 2 η 2 T 26 K 2 C 25 ρ 2 δ= ˜ O (cid:16)
[Equation]
2 δ 4 + η 2 K 2 δ 2 (cid:17) ≲ ˜ O (cid:16) ρ 2 δ 4 (cid:17) ,
when K =
O (cid:16) Lζ (cid:17) ≤ O (cid:0) Lδ (cid:1) .
From this, ( N 1 + N 2 ) T 6 ≤ ˜ O (cid:16) 1 Kε 2 + ρ 2 Kδ 4 (cid:17) × ( the right-hand side of (49) ) .
Taking T ≥ 2( N 1 + N 2 + 1) T 6 , there exists τ such that a τ = 3 , which concludes the proof.
E.3. Finding Second-Order Stationary Points When Clients are Homogeneous ( ζ ≪ 1 δ )
In the previous subsection, we assumed that ζ ≥ 1 δ .
Here, we introduce a simple trick to remove this assumption and give its convergence analysis.
Let T 7 = ˜Θ (cid:0) Lδ (cid:1) with a sufficiently large hidden constant.
In line 18 - 19 of FLEDGE, when k ≡ T 7 , we randomly select C 25 L 2 δ+ b (not b ) samples J t,ki t , and update z t,k as z t,k ← z t,k − 1 + | J t,kit | (cid:80) j ∈ J t,kit ( ∇ f i t ,j ( x t,k ) − ∇ f i t ,j ( x t,k − 1 )) .
This increases the number of gradient evaluations in each inner-loop by ˜ O
( K/ ( L/δ ))
×
˜ O
( L 2 /δ 2 )
=
˜ O ( KL/δ ) ≲ ˜ O
( K 2 )
≲ ˜ O
( Kb ) .
Thus, this does not affect the inner-loop complexity more than by constant factors.
Then, the following lemma holds, which stands as generalization of Lemma 29.
Lemma 33 (Small stuck region) Let { x t,k } be a sequence generated by FLEDGE and ( τ 0 , κ 0 ) be a step where − γ :=
λ min ( ∇ 2 f ( x τ 0 ,κ 0 ))
≤ − δ holds.
We denote the smallest eigenvector direction of λ min ( ∇ 2 f ( x τ 0 ,κ 0 )) as e .
Moreover, we define a coupled sequence { ˜ x t,k } by running FLEDGE with ˜ x 0 = x 0 and the same choice of all randomness i.e., client samplings, minibatches and noises, but the noise at some step ( τ, κ ) > ( τ 0 , κ 0 ) , satisfying κ ≡ T 7 ; We let ˜ ξ τ,κ = ξ τ,κ − r e e with r e ≥ rν TK √ d .
Let w t,k = x t,k − ˜ x t,k , w t
= x t − ˜ x t , v t,k = 1 P (cid:80) Pi =1 y t − 1 i + z t,k , ˜ v t = 1 P (cid:80) Pi =
˜ y t − 1 i + z t,k , g t = 1 P (cid:80) Pi =
y ti −∇ f ( x t ) − (cid:16) 1 P (cid:80) Pi =
˜ y ti − ∇ f
(˜ x t ) (cid:17) , and h t,k = ( z t,k − ( ∇ f ( x t,k ) −∇ f ( x t, 0 )))
− (˜ z t,k − ( ∇ f (˜ x t,k ) − ∇ f (˜ x t, 0 ))) .
Then, v t,k − ∇ f ( x t,k ) − (˜ v t,k − ∇ f
(˜ x t,k ))
=
g t − 1 + h t,k .
There exists a sufficiently large constants C 5 =
˜ O (1) , C
= O (1) , with which the following holds:
If we take p ≥ √ P + C 25 ζ 2 δ 2 + C 25 L 2 Kbδ 2 , b ≥ √ K and, η = ˜Θ (cid:0) 1 L (cid:1) , with probability 1 − 3 ν TK ( ν ∈ (0 , 1)) , we have
[Equation]
Proof of Lemma
We assume K is at least as large as 3 T 7 .
When K − 2 T 7 ≤ κ 0 < K − 1 , taking T 7 ≥ T 5 yields the assertion, considering the two coupled sequence initialized at ( κ 0 , K ) , according to a slight modification of Lemma 29.
Otherwise, we let ( τ, κ ) as the first step after ( τ 0 , κ 0 ) with κ ≡ T 7 .
Then, it suffice to show that, with probability at least 1 − 3 ν TK ,
[Equation]
Since K ≥ 3 T 7 and κ 0 < K − 2 T 7 imply g t − 1 = 0 for all ( τ, κ ) ≤ ( t, k ) <
( τ, κ + T 7 ) , g t − 1 + h t,k = h t,k holds.
Then, ∥ h τ,κ
[Equation]
using
Proposition 4.
Moreover, for ( τ, k ) > ( τ, κ ) , when we assume max ( τ,κ ) ≤ ( t,k ) <
( τ,κ + T 7 ) {∥ x τ,k −
[Equation]
Assuming that (a) 12 (1 + ηγ ) k − κ r e ≤ ∥ w t,k
∥ ≤ 2(1 + ηγ ) k − κ r e and (b) ∥ w t,k − w t,k − 1 ∥ ≤ 3 ηγ (1 + ηγ ) k − κ r e for ( τ, κ ) < ( t, k ) <
( τ, κ + T 7 ) , we get ∥ h t,k ∥ ≤ 2 C 8 γ C 5 (1 + ηγ ) k − κ .
Thus, following the discussion in Lemma 29 and taking T 7 similarly to T 5 , we have (50).
Previously, we only focused on the noise at the last local step ( κ 0 , K ) .
Thus, if the number of steps required to escape saddle points T 5 = ˜ O ( Lδ ) is smaller than the local steps K =
˜ O
( Lζ ) , the algorithm sometimes have to wait more than O ( T 5 ) steps for the last local step.
Therefore, taking K ≥ T 5 was useless to reduce the number of communication rounds.
On the other hand, based on Lemma 33, when FLEDGE comes to a saddle point, FLEDGE does not need to wait next communication, and can escape the stack region within 2 T 7 local steps, even if T 7 ≪ K .
This allows to us to take K larger than O ( Lδ ) , and leads to removal of the assumption δ < 1 ζ from Theorem 28.
E.4. Convergence under PL condition
Theorem 34 Under Assumptions 1 to 3, 5 and 6, if we choose P Kb ≥ Ω (cid:16) C 21 σ 2 ε
[Equation]
and r ≤ ε √ η
[Equation]
(cid:17) , Algorithm 3 with Option I finds an ε -first-order stationary points for problem (2) using
[Equation]
stochastic gradients and
[Equation]
communication rounds with probability at least 1 − 8 ν .
Moreover, under the same conditions, algorithm 3 with Option II finds an ε -first-order stationary points for problem (2) using
[Equation]
stochastic gradients and
[Equation]
communication rounds with probability at least 1 − 8 ν .
Here ˜
O hides only at most log 6 .
5 ( P + K + µ − 1 + ν − 1 ) and polyloglog factors.
Proof
According to eq.
(33) and PL condition, f ( x t,k )
[Equation]
Rearranging the above yields that
[Equation] [Equation]
holds for all t, k (1 ≤ t ≤ T, 0 ≤ k ≤ K − 1) with probability at least 1 − 8 ν .
Applying Lemma 23 to this, for all t = 1 , · · · , T with probability at least 1 − 8 ν ,
[Equation]
+     η
[Equation]
By using this bound repeatedly, we get
[Equation]
( Option II )
[Equation]
We take η as η = Θ (cid:32)
[Equation]
so that 12 η − L 2 − (cid:80) T 3 s =1 (1 − ηµ 2 ) − ( s +1) K (cid:16) 120 C 81 ζ 2 ηK 2 p + 32 C 101 ζ 2 ηK pb (cid:17) − η (cid:80) Kk
=1 (1 − ηµ 2 ) k − K (cid:16) 4 ζ 2 K + 4 C 21 L 2 b (cid:17) ≥ 0 holds.
We also take r ≤ ε √ η 8 and P Kb ≥ 512 C 21 σ 2 ε+ 64 C 1 G ε , then (cid:80)
Tt =
[Equation]
(cid:80)
[Equation]
=
[Equation] [Equation]
+
[Equation]
(cid:16)
[Equation] [Equation]
8 C 21 η PKb (cid:16) σ 2 + G 2 PKb (cid:17)(cid:17) ≤ ε 2 8 µ holds.
Then, we have that
[Equation]
) ( Option II )
For Option I, the first term (1 − ηµ 2 ) TK ( f ( x 0 ) − f ( x ∗ )) is smaller than ε 2 32 if we take T K =
O (cid:16) 1 ηµ log ∆ ε (cid:17) .
The third term is bounded by ε 2 32 , if we take T
= T 3 + O ( 1 ηµK log C 31 ( σ 2 P p 2 b + G 2 P p 3 Kb 2 + σ 2 cPKp 2 + G 2 cPKp 3 ) ε ) =
O (cid:16) Pp C 1 + C 1 ηµK log σ + G + σ c + G c ε (cid:17) .
Moreover, note that f ( x T ) − f ( x ∗ )+ η 4 (cid:80) T t =1 (cid:80) K k =1 (1 − ηµ 2 ) ( T − t +1) − k ∥∇ f
( x t,k − 1 ) ∥ 2 ≤ 6 32 µ min t,k ∥∇ f ( x t,k − 1 ) ∥ 2 holds when we take T
= O (cid:16) 1 ηµK (cid:17) .
Thus, for Option I, if we take
[Equation]
we obtain the desired bound with probability at least 1 − 8 ν .
For Option II, taking
[Equation]
yields the desired bound.
Note that T depends on ε − 1 only logarithmically, which means that C 1 depends on ε − 1 in only log log order and C
= O ∗ (log( P + K + µ − 1 + ν − 1 ))
(where O ∗ suppresses log log factors).
Remark
In order to find ε -solutions (i.e., f ( x t,k ) − f ∗ ≤ ε ), the same statement holds, except for slight change on the assumptions on P Kb and r : P Kb ≥ Ω (cid:16) C 21 σ 2 µε + C 1 G ε √ εµ (cid:17) and r ≤ η √ εµ 2 .
In fact, we can derive f ( x t,k ) − f ∗ ≤
[Equation]
η similarly to (51) , and using this, we have f ( x t ) − f ( x ∗ ) ≤ (
[Equation]
− ηµ )
[Equation]
f
[Equation]
[Equation]
f
[Equation]
−
[Equation]
Option
[Equation]
( Option II )
Taking η similary to the previous theorem, r ≤ η √ εµand P Kb ≥ Ω (cid:16) C 21 σ 2 µε + C 1 G √ µε (cid:17) yields
[Equation]
Thus, we finally have the following:
[Equation]
Now it is trivial to see that the desired bound holds.
Appendix F. Lower bound
In this section, we provide the gradient complexity lower bound of O ( n + ∆( ζ √ n ∨ L ) ε) under Hessian heterogeneity, which recovers the usual lower bound for L -smooth functions by setting ζ = 2 L .
Note that the gradient complexity, or the total number of gradient communicated, of FLEDGE is ˜ O ( P + ζ √ P ε 2 ) .
Thus, this almost matches the lower bound of gradient complexity of the finite-sum case if we identify P with n .
Note that the lower bound is proven under averaged gradient L -Lipshitzness and averaged Hessian ζ -heterogeneity, while we assume gradient L -Lipshitzness of each f i for the upper bounds.
However, we expect that averaged gradient L -Lipshitzness and averaged Hessian ζ -heterogeneity would suffice for deriving the first-order optimality in expectation.
First, we give a definition of the linear-span first-order algorithms.
Definition 36 (Linear-span first-order algorithm)
Fix some x 0 .
Let A be a (randomized) algorithm with the initial point x 0 , and x t be the point at the t -th iteration.
We assume A select one individual function i t at each iteration t and computes ∇ f i t ( x t ) .
Then A is called a linear-span first-order algorithm if
[Equation]
holds for all t with probability one.
Note that this definition includes minibatch updade, by letting x sb = x
sb +
= , · · · , = x ( s +1) b −with the minibatch size b .
We also define problem classes F Ln, ∆ and F L,ζn, ∆ for (1), as follows.
Definition 37 (A class of finite-sum optimization problems) Fix some x 0 .
For an integer n , L > 0 , we define a problem class F Ln as
[Equation]
d ∈ N , 1 n (cid:80) ni
=1 ∥∇ f i ( x ) − ∇ f i ( y ) ∥ 2 ≤ L 2 ∥ x − y ∥
for all x, y (averaged gradient L -Lipschitzness), and f ( x 0 ) − inf x f ( x ) = ∆ .
 
Moreover, for an integer n , L > 0 , and ζ > 0 , a problem class F L,ζn is defined as
F L,ζn
[Equation]
d
[Equation]
{+=1 ∥∇ f i ( x ) − ∇ f i ( y ) ∥ 2 ≤ L 2 ∥ x − y ∥ 2 for all x, y , 1 n 2 (cid:80) ni,j =1 ∥∇ 2 f i ( x ) − ∇ 
 f+}
i ( y ) ∥ 2 ≤ ζ 2 (averaged
Hessian ζ -heterogenity), and f ( x 0 ) − inf x f ( x ) = ∆ .
  .
We now state our lower bound theorem as follows.
Proposition 35 Assume Assumptions 1, 2 and 5.
For any L > 0 , ∆ > 0 , and ε > 0 , there exists a function f ∈ F L,ζn, ∆ such that any linear-span first-order algorithm requires
[Equation]
stochastic gradient accesses to find ε -first-order stationary points of the problem (1) .
Proposition 35 can be derived by using the bounds of Carmon et al.
[4], Fang et al.
[10], Li et al.
Carmon et al.
[4] proved the following lower bound.
Proposition 38 (Carmon et al.
Fix x 0 .
For any L > 0 , ∆ > 0 , and ε > 0 , there exists a function f ∈ F L 1 , ∆ such that any linear-span first-order algorithm requires Ω (cid:0)
∆ L ε 2 (cid:1) stochastic gradient accesses in order to find ε -first-order stationary points.
Fang et al.
[10], Li et al.
[27] extended this to the lower bound on the finite-sum optimization problem.
Proposition 39 (Fang et al.
[10], Li et al.
Fix x 0 .
For n > 0 , L > 0 , ∆ > 0 , and ε > 0 , there exists a function f ∈ F Ln, ∆ such that any linear-span first-order algorithm requires Ω (cid:16) n + ∆ L √ n ε [Equation] stochastic gradient accesses in order to find ε -first-order stationary points.
Based on these, we give the lower bound under the additional assumption of ζ -Hessian-heterogeneity.
Proof
It is easy to see that the lower bound of Proposition 38 also applies to F Ln, ∆ , by letting f 1 = f 2 = · · · = f n = f ∗ where f ∗ is the function that gives the bound of Proposition 38.
On the other hand, we have F ζ 2 n, ∆ ⊆
F L,ζn, ∆ .
Thus, Proposition 39 yields that there exists a function f ∈ F ζ 2 n, ∆ ⊆
F L,ζn, ∆ that requires Ω (cid:16) n + ∆ ζ √ n ε 2 (cid:17) stochastic gradients to find ε -first-order stationary points.
Therefore, by combining these two bounds, we have the desired lower bound.