Asymptotic Behaviors of Projected Stochastic Approximation: A Jump Diffusion Perspective
Abstract
In this paper we consider linearly constrained stochastic approximation problems with federated learning as a special case.
We propose a loopless projection stochastic approximation algorithm (LPSA) to ensure feasibility by performing the projection with probability p n at the n -th iteration.
Considering a specific family of the probability p n and step size η n , we analyze our algorithm from an asymptotic and continuous perspective.
Using a novel jump diffusion approximation, we show that the trajectories connecting those properly rescaled last iterates weakly converge to the solution of specific stochastic differential equations (SDEs).
By analyzing SDEs, we identify the asymptotic behaviors of LPSA for different choices of ( p n , η n ) .
We find the algorithm presents an intriguing asymptotic bias-variance trade-off according to the relative magnitude of p n w.r.t.
η n .
It brings insights on how to choose appropriate { ( p n , η n ) } n ≥ 1 to minimize the projection complexity.
Introduction
Recently, a novel distributed computing paradigm that called Federated Learning (FL) has been proposed for collaboratively training a global model from data that remote clients hold [31].
As a standard optimization algorithm in FL, Local SGD alternates between running stochastic gradient descent (SGD) independently in parallel on different clients and averaging the sequences only once in a while.
Put simply, it learns a shared global model via infrequent communication.
Empirical investigation finds its superior performance in communication efficiency [30] and theoretical analysis toward it has already provided a complete picture [26, 1, 15, 42, 43, 15].
Among them, Li et al.
[27] establishes a functional CLT that Local SGD with Polyak-Ruppert averaging simultaneously achieves the optimal asymptotic variance and diminishing average communication frequency.
However, they are all derived from a discrete perspective.
The use of a continuous-time stochastic process to characterize the entire trajectory of a discrete stochastic algorithm has been witnessed progresses in recent years, and we call it diffusion approximation.
The continuous approach has advantages in its rich toolbox and can provide intuitive explanation for uncanny phenomena that are intractable to analyze in discrete cases.
It can also motivate new optimization algorithms and statistical inference methods.
Current works applying diffusion approximation to stochastic optimization algorithms can be roughly divided into two classes.
The first one is to interest the optimization algorithm as a numerical discretization of a specific stochastic differential equation (SDE) [14] in a finite time interval [0 , T ] .
When the step size η is sufficiently small and the length T (= nη ) of the interval is fixed ( n is the total iterations), such approximation is of high accuracy, and it is easy to analyze the geometric properties of our target algorithms [41, 23, 13, 7, 38, 34, 8].
However, this avenue is difficult to capture the convergence behaviors around the optimal point due to the fixed T .
The second class comes up to solve the issue.
It instead considers the iterates divided by a proper power function of step sizes.
Under certain conditions, as n goes to infinity, the rescaled iterates would weakly converge to the stationary solution of corresponding SDEs [21, 35, 6, 9, 10].
In FL, to the best of our knowledge, no work considers analyzing Local SGD via the aspect, which is our focus here.
However, it is not easy to serialize Local SGD iterates due to its double-loop nature.
Recent researchers developed a new technique named as ‘loopless’ to simplify the two-loop structure for
SVRG and Katyusha [16].
The key is to replace the hard loop with a probabilistic loop.
Specifically, we will independently toss a (possibly biased) coin ω n with head probability p n at iteration n .
When getting the head ω n = 1 , we start a new loop and update the outer-loop intermediate variables; when getting the tail ω n = 0 , we stay in the same loop and keep the intermediate variables.
In this way, we obtain a loopless counterpart algorithm and do not need to distinguish inner and outer loops anymore.
It facilitates theoretical analysis and typically does not deteriorate the convergence rate [12, 29, 28, 11].
It is worth mentioning that Hanzely and Richtárik [12] first introduced the loopless technique to
FL and obtained many efficient FL algorithms.
Li [28] used a dynamic p n (which varies with n ) to generalize the scope of original methods.
We are then motivated to analyze a loopless version of
Local SGD with decreasing p n , but from an asymptotic and continuous perspective.
1.1 Contribution
Our work is motivated by Local SGD but beyond it.
In particular, for a general optimization problems with linear constrains (of which FL is a special case), we develop a loopless projection stochastic approximation method (LPSA) as a generalization of Local SGD (see Appendix A for more details).
Such generality renders us the possibility to transfer our techniques and results to other linearly constrained problems.
LPSA is affected by two important hyperparameters, namely the step size { η n } and the projection probability { p n } .
For the choices of η n ∝ n − α and p n ∝ min { η βn , 1 } , we derive a non-asymptotic convergence rate for different α ∈ (0 , 1] and β ∈ (0 , 1) in Theorem 3.1.
We observe a phase transition for the convergence rate O ( n − α min { 1 , 2 − 2 β } ) when β crosses 0 .
To derive asymptotic results, we obtain two sequences { u n } and { v n } by orthogonal decomposition for the optimized sequence of LPSA.
We then construct two sequences of processes which pass through the appropriately rescaled u n and v n , respectively.
We show rigorously, when the iteration goes to infinity, these two sequences of stochastic processes weakly converge to the solutions of specific SDEs that are driven by either a Brownian motion or a Poisson process.
As a corollary, the rescaled last iterate of u n (which we mainly care about) has a known asymptotic distribution (either Gaussian distribution in Theorem 3.3 or Dirac in Corollary 1).
And the phase transition we mentioned above evolves into a trade-off between the bias caused by the low frequency projection and the fluctuation resulting from the gradient noise (see Section 3.2.3).
Moreover, according to different convergence rate for every { ( η n , p n ) } pair, we consider a selection scheme at the end of Section 3.2.3, which makes the algorithm have the same nonasymptotic convergence order as the conventional stochastic approximation and spend as little as possible on the projection operation which is usually expensive in practice.
At the end, we conduct numerical experiments to confirm the theoretical results.
From a technical level, we propose a novel proof technique to analyze the discontinuity brought by probabilistic projection.
In particular, we borrow tools from jump diffusion and verify necessary conditions (e.g., stochastic tightness) to apply it.
See the paragraph after Theorem 3.4 for a main idea.
We believe our technique can extend to and help analyze other stochastic approximation algorithms which can be approximated by a jump diffusion.
Problem Formulation 2.1 Loopless Projected Stochastic Approximation
Notice that distributed optimization such as FL can be formulated as a global consensus problem which is a linearly constrained problem [5].
For the sake of simplicity and generality, we aim to solve the following problem
[Equation]
via a randomly (and infrequently) projected stochastic approximation algorithm.
In particular, at iteration n , we first perform one step of SGD via
[Equation]
where f ( x ) =
E ζ ∼D f ( x , ζ ) and ξ n = ∇ f ( x n ) − ∇ f ( x n , ζ n ) .
Here { ξ n } is a martingale difference sequence (m.d.s.) under the natural filtration F n +1 :
= σ ( ζ k , ω k ; k ≤ n +1) .
We then use the loopless trick introduced in the introduction, i.e., we independently cast a coin with the head probability p n and obtain the result ω n ∼ Bernoulli( p n ) .
If ω n = 1 , we perform one step of projection to ensure x n +1 fall into the feasible region: x n +1 =
P
A ⊥ ( x n + 12 ) where
P A ⊥ denotes the projection onto the null space of A ⊤ .
If ω n = 0 , we assign x n +1 as the same value of x n + 12 , i.e., x n +1 = x n + 12 .
It is clear this algorithm (2) mimics the behavior of Local SGD in FL settings (see Appendix A for the equivalence).
2.2 Assumptions
For the linearly constrained convex optimization problem (1), we make the following assumptions which are quite common in the literature.
Without special clarification, ∥ · ∥ denotes the Euclidean norm for vectors and the spectral norm for matrices.
Assumption 1 (Smoothness) .
We assume that f : R d → R is L -smooth, that is,
[Equation]
Assumption 2 (Strong convexity) .
We assume that f : R d → R is µ -strongly convex, that is,
[Equation]
Assumption 3 (Continuous Hessian matrix) .
We assume that f : R d → R is Hessian Lipschitz, that is, there is a constant ˜ L such that
[Equation]
Assumption 4 (Continuous covariance matrix) .
Given an m.d.s.
{ ξ t } , we denote the conditional covariance as E [ ξ t ξ ⊤ t |F t ] = Σ( x t ) and assume it is L -Lipschitz continuous in the sense that
[Equation]
Assumption 5.
For the m.d.s.
{ ξ n } , we assume there exists a p > 2 such that the p -th moment of every element in { ξ n } is uniformly bounded, that is,
[Equation]
The first three assumptions imply we consider the strongly convex case.
The last two assumptions help us identify the asymptotic variance.
Especially, the assumption of uniformly bounded p ( p > 2) moments is typically required to establish central limit theorems [9, 10, 27].
Finally, we want to emphasize that the stationary condition for the problem (1) is different from unconstrained ones; ∇ f ( x ⋆ ) is not necessarily zero, however, its projection into the null space of A ⊤ must be zero.
Proposition 1 ([25], Corollary 2.1) .
Let P A be the projection onto the column space of A and P A ⊥ the projection onto the null space of A ⊤ .
Under Assumption 2, the solution of (1) is unique (denoted x ⋆ ). Moreover, we have P A ⊥ ( ∇ f ( x ⋆ )) = 0 .
2.3 Jump Diffusion
Jump diffusion is a stochastic Lévy process that involves jumps and diffusion.
Typically, the former is modeled by a Poisson process, while the latter is modeled as a Brownian motion.
It has wide and important applications in physics, finance[36], and computer vision.
We say a function f defined on R is càdlàg when f is right-continuous and has left limits everywhere.
For a càdlàg process ( V s ) s ≥ 0 , we denote V t − as the left limit of V · at time t .
Let N γ ( t ) denote the
Poisson process with γ the intensity, which quantifies the number of jumps up to the time t and is clearly càdlàg.
We use N γ ( dt ) =
N γ ( t ) − N γ ( t − ) ∈ { 0 , 1 } to indicate whether N γ jumps at time t and (cid:82) T 0 g ( t − ) N ( dt ) =
(cid:80) { t : N γ ( t ) ̸
= N γ ( t − ) } g ( t − ) to denote the integral that drives for a measurable function g ( · ) .
We will consider a special class of jump diffusion in the following form
[Equation]
When the coefficient functions α ( t, X t ) and β ( t, X t ) satisfy conditions like linear growth and
Lipschitz continuity, there exists a solution for the jump diffusion (3) (e.g., Theorem 1.19 in [33]).
Main Results
In the section, we are going to capture the convergence behaviors of our projected stochastic approximation method (2) from both non-asymptotic and asymptotic perspectives.
We consider a specific family of step size η n and projection probability p n , namely, η n = η 0 n − α and p n = min { η βn , 1 } indexed by 0 < α ≤ 1 and 0 ≤ β < 1 , respectively.
The choice of step sizes η n has been used to establish CLTs [37, 27], while the choice of p n is quite novel.
To provide a complete picture of convergence, we will consider almost all combinations of α and β .
3.1 Non-asymptotic Analysis
To provide the convergence rate, it is natural to focus on the projection of x n into the column space of
A (since it is the easiest feasible solution one can obtain from x n ).
Hence, we decompose the iterated x n into two orthogonal components x n := u n + v n where u n =
P
A ⊥ ( x n ) and v n = P A ( x n ) .
2 We specify the the convergence rate of E ∥ u n − x ⋆ ∥ 2 in terms of α, β and n in the following theorem.
Theorem 3.1.
Suppose that Assumptions 1, 2 and 4 hold.
Let η n = η 0 n − α and p n = min { η βn , 1 } with 0 ≤ β < 1 .
Then for (i) 0 < α < 1 or (ii) α = 1 with η 0 > 2 /µ ( µ is the strong convexity parameter of the objective function f ), we have
[Equation]
From Theorem 3.1, as β decreases, that is, the projection happens more frequently, E ∥ u n − x ⋆ ∥converges faster.
The rate is O ( n − α ) when β < 0 .
5 , while the rate is O ( n − 2 α (1 − β ) )
when β > 0 .
Thus there exists a phase transition when β goes across 0 .
5 , which implies we should analyze asymptotic performances for these two phases respectively.
As an extreme, when β = 1 , the algorithm is possible to disconverge in an artifact quadratic loss with a specific A (see Theorem 3.2).
Though for a specific A , it could apply to FL (see Corollary 2 in Appendix A.2.1 for the detail).
Theorem 3.2.
If η n = η 0 n − α and p n = min { p 0 η n , 1 } with 0 < α ≤ 1 , for a specific A , there exists a quadratic function f ( x ) so that ∇ 2 f ( x ) ⪰ I d and E ∥
u n − x ⋆ ∥ 2 does not converge to 0 .
Here
I d ∈ R d × d is the identity matrix, and ∇ 2 f ( x ) ⪰
I d means ∇ 2 f ( x ) − I d is positive semidefinite.
3.2 Asymptotic Behavior of the Rescaled Trajectory
In this section, we want to derive an asymptotic convergence for (2).
Recall that there exists a phase transition for the convergence rate of E ∥
u t − x ⋆ ∥ 2 when β crosses 0 .
5 , when the projection probability is set as p n = η βn .
In the following, we will analyze the asymptotic behaviors of LSPA for the two cases β ∈ [0 , 1 / 2) and β ∈ (1 / 2 , 1) .
3.2.1 Case 1: Frequent Projection where β ∈ [0 , 1 / 2)
From an asymptotic perspective, the typical central limit theorem (CLT) claims ˇ u n := u n − x ∗ √ η n −would weakly converge to a rescaled standard distribution [21].
It helps us capture the large-sample convergence behaviors and provide ways for future statistical inference.
However, we can provide a stronger result that captures the asymptotic behavior of the whole trajectory.
In particular, we serialize the sequence { ˇ u n } by constructing a continuous random function (denoted ¯ u ( n ) t ) such that it starts from ¯ u ( n )= ˇ u n , and as t increases it will pass through ˇ u n +1 , ˇ u n +2 and so on.
We will show that such a random function ¯ u ( n ) t will weakly converge to the solution of a specific SDE.
From the
SDE, we can derive asymptotic variance of ˇ u n and the whole trajectory evolution.
Since ¯ u ( n ) t should pass all { ˇ u k } k ≥ n , we can connect these discrete points by piecewise linear functions.
To that end, we first derive the one-step relation between ˇ u n and ˇ u n +1 .
In particular,
[Equation]
where R n stands for a high-order residual error, and ξ (1) n denotes the component of noise ξ n on the null space of A .
One can find the derivation of (4) in Appendix C.1.
Roughly speaking, (4) can be viewed as a one-step Euler Maruyama discretization with timescale η n for an SDE, which starts at ˇ u n with local drift coefficient b n and local diffusion coefficient var( ξ (1) n ) .
Definition 1 ( Time interpolation ) .
Let a positive sequence γ = { γ n } ∞ n decrease to zero.
For n ∈ N , t ≥ 0 , define N ( n, t, γ [Equation] =
[Equation]
m
[Equation]
[Equation]
k
[Equation]
γ k , and t n ( γ ) =
Γ N ( n,t,γ ) − Γ n .
We introduce a time interpolation for the formal description of the continuous function and further analysis.
Intuitively, N ( n, t, γ ) is the number of iterations m at which the sum of step sizes (cid:80) m
+ k = n +1 η k is just larger than t and t n ( γ ) is the approximation of t when we only use step sizes { γ k } k ≥ n .
Since γ n → 0 , t n ( γ ) → t as n goes to infinity.
A property of Definition 1 is that
N ( n, Γ m ( γ ) − Γ n ( γ ) , γ ) = m for any m ≥ n .
By now, we are ready to construct ¯ u ( n ) t .
For a given n ∈ N , let ¯ u ( n )=
ˇ u n and define for t ≥ 0 ,
[Equation]
From the construction, we can see that ¯ u ( n ) t k ( η ) =
ˇ u n + k .
Theorem 3.3 ( Diffusion Approximation ) .
Let Assumptions 1-5 hold.
The following family of continuous stochastic processes { ¯ u ( n ) t : t ≥ 0 } ∞ n =1 weakly converges to the stationary weak solution of the following SDE:
[Equation]
Further, the rescaled sequence { ˇ u n } ∞ n =1 converges weakly to the invariant distribution of the dynamics (7) , i [Equation] e., N ( 0 , ˜Σ) .
Here the variance ˜Σ satisfies the Lyapunov equation
[Equation]
⊥
[Equation]
Remark 1.
By using the continuous time version of the Lyapunov theorem (Lemma 1 in [40]), the
Lyapunov equation has a unique positive semidefinite solution (denoted ˜Σ ).
From Theorem 3. and Theorem 4.1.1 in [9], we can tell that when β ∈ [0 , 12 ) our algorithm LPSA achieves the same asymptotic variance as SGD that also uses the same step size.
The typical projected SGD corresponds to the case β = 0 , while LPSA allows β to vary in [0 , 12 ) .
One can reduce the projection frequency by increasing β (equivalently decreasing the probability p n ).
Hence, when projection is expensive, LPSA is more efficient in performing projections due to its flexible and moderate projection frequency.
Proof Idea of Theorem 3.3.
We shed light on the proof idea of Theorem 3.3.
From a high level, we leverage the general theory for operator semigroups, which are developed by Trotter and Kurtz [39, 17– 19] and are used to analyze stochastic optimization algorithms in [9].
Our diffusion approximation results are built on it, but generalize it in the sense that we use the celebrated Prokhorov’s theorem to extend to the whole trajectory.
One difficulty is to prove the stochastic tightness of { u ( n ) t } .
To that end, we make use of a classic result (e.g. Theorem 7.3 of [4]).
3.2.2 Case 2: Occasional Projection where β ∈ (1 / 2 , 1)
When we step into the low-frequency regime where β ∈ (cid:0) 12 , 1 (cid:1) , the situation totally changes.
Intuitively, when LPSA performs much less frequent projection, we will frequently use infeasible x t to update parameters, which accumulates residual errors.
These errors would not only dominate and slow down the non-asymptotic convergence rate (see Theorem 3.1), but also change the asymptotic behavior.
In this case, we should not only find the right timescale, but also need to figure out how these errors are accumulated.
To solve the issue, we develop a new analysis routine.
In the following, we consider p t =
[Equation] with γ > 0 .
Our solution is to monitor another random process that is related with { v n } , which serves as a bridge to derive the asymptotic behavior of { u n } .
The right scale should make the scaled sequence have non-vanishing expected L 2 norm.
From Theorem 3.1, it should be ˇ v n = η β − 1 n − 1 v n .
In addition, given ˇ v n , the candidate value of ˇ v n +1 before tossing the coin ω n , can be derived from LPSA’s Algorithmin Appendix A, and we denote this candidate as ˇ v ( n +1) − .
[Equation]
where d n = ∇ f ( x ∗ )
+ η − β n S n with S n a residual error which satisfies η − β n S n = o P (1) (see
Appendix C.2 for more details) and ξ (2) n stands for the component of noise ξ n on the orthogonal complementary space A ⊥ .
Due to the probabilistic projection, ˇ v n +1 takes value ˇ v ( n +1) − with probability 1 − γη βn and takes value zero with probability γη βn .
Similar to the previous section, we then construct a càdlàg random process ¯ v ( n ) t which starts from ˇ v n and will pass through { ˇ v ( k ) − } k ≥ n .
We can connect these discrete points with a step function.
It results in the following construction
[Equation]
From (9), we can claim that ¯ v ( n ) t n [Equation] − =
[Equation] n,t,η β ) − for any t ≥ 0 . With probability p N ( n,t,η β ) , ˇ v N ( n,t,η β ) takes value zero, which causes the process ¯ v ( n ) · to change abruptly at the time t n ( η β ) .
These discontinuities about ¯ v ( n ) · prevent the diffusion process from working on ¯ v ( n ) · as the result of
Theorem 3.3.
Even so, the following theorem shows that we can still find a suitable process in the broader jump diffusion class to approximate ¯ v ( n ) · .
Theorem 3.4 ( Jump Approximation ) .
Let Assumptions 1, 2, 4 and 5 hold.
The following family of càdlàg stochastic processes { ¯ v ( n ) t : t ≥ 0 } ∞ n =1 weakly converges to the stationary weak solution of the following SDE
[Equation]
Here N γ ( t ) represents Poisson process with intensity γ , and N γ ( dt )
[Equation]
N γ ( t ) − N γ ( t − ) .
Further, the rescaled sequence { ˇ v n } ∞ n =1 weakly converges to the invariant distribution of the dynamics (10) , i.e., − ∇ f ( x ⋆ ) ∥∇ f ( x ⋆ ) ∥
· E (cid:16) ∥∇ f ( x ⋆ ) ∥ γ (cid:17) .
Here E ( θ ) represents the exponential distribution with intensity 1 θ .
Theorem 3.4 shows that the sequence { ¯ v ( n ) t } constructed by shifting initial points will finally approximate a jump process with a constant drift as n goes to infinity.
The SDE (10) sheds light on how ¯ v ( n ) t (equivalently a rescaled version of v n ) move as t increases.
As the error incurred by infrequent projections, ¯ v ( n ) t will move towards the direction of ∇ f ( x ∗ ) (due to the drift term −∇ f ( x ⋆ ) dt ) and be periodically forced to set as zero vector ( due to the correcting term − Y t − · N γ ( dt ) ).
From a qualitative perspective, the SDE (10) captures the periodical behavior of v n , hence it shows without projection the residual error will accumulate along the direction of ∇ f ( x ⋆ ) .
As argued in
Proposition 1, ∇ f ( x ⋆ ) is unlikely to be zero in our constrained problems.
The remaining issue is how to link { ¯ v ( n ) t } to our target { u n } .
Similarly, we should consider a rescaled u n , that is, ˆ u n := ( u n − x ⋆ ) /η 1 − β n − 1 .
The following corollary, which is based on Theorem 3.4, shows when β ∈ (cid:0) 12 , 1 (cid:1) , ˆ u n converges to a non-zero vector.
Recall that ˇ u n = u n − x ∗ √ η n − 1 = η 0 .
5 − β n − 1 ˆ u n .
The equation together with Corollary 1 implies ∥ E ˇ u n ∥ =
η 0 .
5 − β n − 1 ∥ E ˆ u n ∥ → ∞ .
As a result, the bias in
Corollary 1 instead of the Gaussian fluctuation in Theorem 3.3 becomes the leading term hindering the convergence.
Corollary 1.
Let Assumptions
1- 3 hold.
Then ˆ u n := 1 η 1 − β n − 1 ( u n − x ⋆ ) converges to a non-zero vector
[Equation]
−
[Equation]
(cid:111) † (cid:0) P A ⊥ ∇ 2 f ( x ⋆ ) ∇ f ( x ⋆ ) (cid:1) in the L 2 as n → ∞ .
Where G † denotes the pseudoinverse of the symmetric matrix G .
Proof Idea of Theorem 3.4 The main proof idea is similar to that of Theorem 3.4 except that we need to handle the jump diffusion which introduces additional discontinuity.
As a result, for each n ≥ 0 , { ¯ v ( n ) t } t ≥ n is càdlàg rather than continuous.
We then use the approximation result for jump diffusions developed by Kushner [20] instead of Trotter and Kurtz’s theories.
Furthermore, the tool for proving tightness also needs to change.
We replace the classic tool in [4] with a generalized determination method, the latter used to establish the stochastic tightness for càdlàg processes (e.g.,
Theorem 4.1 in [19]).
The remaining issue is to figure out properties (e.g., the mixing nature) of (10).
To that end, we establish the geometric ergodicity of Eq.
(10) by combining the coupling method with the Itô’s formula for jump diffusions, and show that its invariant distribution exists uniquely.
3.2.3 Summary and Discussion
From Sections 3.2.1 and 3.2.2, for the choice p n ∝ η βn , when β varies, our algorithm has an interesting bias-variance tradeoff.
In fact, Theorems 3.3 and 3.4 reveal that the fluctuation of u n is of order
O ( η 12 n ) and the bias is of order O
η 1 − β n ) .
When β ∈ [0 , 1 / 2) the fluctuation caused by the randomness of gradient queries in every iteration dominates the optimization accuracy.
And when β ∈ (1 / 2 , 1) , this indicator is manipulated by the biases formed by the accumulation of skewed updates in the unconstrained state within each ‘inner loop’.
In practice, projection is expensive to perform.
Hence, it is important to tune α and β so that the projection complexity is minimized as much as possible.
We use the average projection complexity (APC) to quantify the projection efficiency.
For a target accuracy ϵ > 0 , APC is defined as the number of projections required to obtain an ϵ -accuracy feasible solution.
We summarize the derived results and the corresponding APC in Table 1.
We can see that APC is minimized when α → 1 and β → 0 . 5 .
In this case, APC is approaching 1 √ ϵ .
We find an interesting parallelism between LPSA and Local SGD.
In the case of FL, projection complexity corresponds to communication complexity, because a synchronization in FL is essentially a projection in linearly constrained problems (see Appendix A for the equivalence).
In [27], the authors analyzed the averaged communication complexity (ACC) for Local SGD with Polyak-Ruppert averaging.
3 They considered a general case where the length of the m -th inner loop could be up to E m := m ν with ν ∈ [0 , 1) .
After E m steps of inner loop, communication would perform to synchronize local models.
Hence, E m plays a role similar to p n in our paper.
Li et al.
[27] found that when ν ∈ [0 , 1) , the averaged Local SGD iterates enjoy an optimal asymptotic normality up [Table] to a known constant scale and its ACC is (cid:0) 1 ϵ (cid:1) 11+ ν .
When ν → 1 , ACC is approaching 1 √ ϵ , similar to our case where APC converges to 1 √ ϵ when α → 1 and β → 0 . 5 .
Actually, the 1 √ ϵ average communication complexity is actually optimal for any first-order oracle distributed algorithms, as shown in [43].
Hence, it implies our LSPA is efficient and near optimal in projection, because we can always reduce FL as a special of (1).
Experiments
In this section [Table] we validate our theoretical results through comprehensive experiments.
Due to space limitations [Table] we only show some representative results on synthetic datasets under FL settings.
For the results on general linearly constrained problems, please refer to Appendix D.
Experimental Setup We focus on classification problems with cross entropy loss, and ℓ 22 regularization is imposed to ensure the strong convexity of the objective function.
The synthetic datasets are generated by following [24].
There are K clients and the sample ( x k , z k ) on the k -th client is modeled as x k ∼ N ( ν k , Λ) and z k = argmax(softmax( W k x k + b k )) where Λ ∈ R d × d is diagonal with the entry ( j, j ) equal to j − 1 . 2 , W k ∈ R C ×
d and b k ∈ R C .
We consider two specific datasets.
The first one is denoted by IID , where all the clients share the same W k and b k , and ν k ∼ N ( 0 , I d ) .
For this one, we set K = 100 , d = 60 and C = 10 .
The second one is denoted by Synthetic ( a, b ), where a and b control the heterogeneity across clients.
Specifically, each entry of W k and b k is modeled as N ( µ k , 1) with µ k ∼ N (0 , a ) and ν k ∼ N ( ζ k , I ) with ζ k ∼ N ( 0 , b I d ) .
For this dataset, we set K = 20 , d = 10 and C = 5 .
We find that the results on IID are intuitive enough to demonstrate the convergence rates of the mean squared error (MSE) E ∥
u n − x ⋆ ∥ 2 and the asymptotic behavior of ˇ u n for β ∈ [0 , 1 / 2) .
The results on Synthetic ( a, b ), a dataset with fewer parameters and more heterogeneity, are more appropriate to illustrate the asymptotic biased of ˆ u n for β ∈ (1 / 2 , 1) .
The full results on both the datasets are deferred to Appendix D.
Convergence Rate We plot the log-log scale graphs of averaged MSEs over 5 repetitions on
IID vs iterations in Figure 1.
The value of α is set as { 1 , 0 . 8 , 0 . 6 } and the value of β is from { 0 , 0 . 2 , 0 . 4 , 0 . 6 , 0 . 8 } .
For each repetition, we run 2000 steps of LPSA.
By Theorem 3.1, the slope of the line in the log-log scale graph should be − α min { 1 , 2 − 2 β } .
This is in accordance with Figurewhen the iteration is larger than 100 .
For β ∈ [0 , 1 / 2) , the value of β does not affect the slope, while for β ∈ (1 / 2 , 1) , larger β and smaller α both lead to smoother lines. [Figure] [Figure]
[Figure] Frequent Projection For α = 1 and β ∈ { 0 , 0 . 2 } , we run 2000 steps of LPSA over 100 repetitions on IID and pick up the last 200 iterates.
For these iterates, we compute the rescaled vectors ˇ u n and project them into a two-dimensional random subspace.
Then we plot the heatmaps across the two dimensions in Figure 2.
We observe that the cells near the origin have the lightest colors, and as we move away from the origin the cell color becomes darker.
Since the cells with lighter colors imply more frequencies, these phenomenons agree with Theorem 3.3, where the limiting distribution of ˇ u n is Gaussian.
The results with other values of α and β are deferred to Appendix D.4.
Occasional Projection For α = 0 .
8 and β = 0 .
6 , we run 50000 steps of LPSA over 5 repetition on Synthetic ( 1 , 1 ).
Then we compute the rescaled sequence ˆ u n and project them along two random directions e 1 and e 2 .
The trajectories depicted in Figure 3 show that the limits of ⟨ ˆ u n , e 1 ⟩ and ⟨ ˆ u n , e 2 ⟩ are nonzero and verify the asymptotic biased of ˆ u n mentioned in Corollary 1.
The results with other values of α and β are deferred to Appendix D.6.
Concluding Remarks
In this paper we study the linearly constrained optimization problem.
We propose the LPSA algorithm that is inspired by Local SGD.
The probabilistic projection in LPSA follows the spirit of loopless methods [16, 12, 28] and simplifies the double-loop structure of original Local SGD, facilitating theoretical analysis.
We thoroughly analyze the (non-)asymptotic properties of properly scaled trajectories obtained from { u n } and discover an interesting phase transition where { u n } changes from asymptotically normal to asymptotically biased as the projection frequency decreases.
From a technical level, we generalize jump diffusion approximations to accommodate the particularity and discontinuity of LPSA.
There are also some open problems.
It is unclear about the asymptotic behavior of u n when β = 0 .
5 , i.e., p n = Θ( √ η n ) .
The jump diffusion approach fails because we can’t analyze { u n } via the length of { v n } anymore.
It accounts for failure that { ˇ u n } and { ˇ v n } are incompatible in the sense that they use different time scales and the time interpolation.
However, we speculate ˇ u n would finally converge weakly to a non-centred Gaussian distribution.
In addition, it is also interesting to analyze the performance of projection complexity of LPSA.
From Corollary 1, to achieve a better convergence rate at lower projection frequencies, we must overcome the asymptotically biased nature of u n .
One feasible approach is to build a ‘de-biasing’ algorithm which attenuates the effect of v n during the update of u n .
We leave them as future work.
Acknowledgments and Disclosure of Funding
This work has been supported by the National Natural Science Foundation of China (No. 12271011).