Parallel Restarted SGD with Faster Convergence and Less Communication: Demystifying Why Model Averaging Works for Deep Learning

Introduction

Consider the distributed training of deep neural networks over multiple workers (?), where all workers can access all or partial training data and aim to find a common model that yields the minimum training loss. Such a scenario can be modeled as the following distributed parallel non-convex optimization

One classical parallel method to solve problem (1) is to sample each worker’s local stochastic gradient in parallel, aggregate all gradients in a single server to obtain the average, and update each worker’s local solution using the averaged gradient in its SGD stepEquivalently, we can let the server update its solution using the averaged gradient and broadcast this solution to all local workers. Another equivalent implementation is to let each worker take a single SGD step using its own gradient and send the updated local solution to the server; let the server calculate the average of all workers’ updated solutions and refresh each worker’s local solution with the averaged version. (?) (?). Such a classical method, called parallel mini-batch SGD in this paper, is conceptually equivalent to a single node Stochastic Gradient Descent (SGD) with a batch size $N$ times large and achieves $O(1/\sqrt{NT})$ convergence with a linear speed-up with respect to (w.r.t.) the number of workers (?). Since every iteration of parallel mini-batch SGD requires exchanging of local gradient information among all workers, the corresponding communication cost is quite heavy and often becomes the performance bottleneck.

There have been many attempts to reduce communication overhead in parallel mini-batch SGD. One notable method called decentralized parallel SGD (D-PSGD) is studied in (?)(?) (?). Remarkably, D-PSGD can achieve the same $O(1/\sqrt{NT})$ convergence rate as parallel mini-batch SGD, i.e., the linear speed-up w.r.t. the number of workers is preserved, without requiring a single server to collect stochastic gradient information from local workers. However, since D-PSGD requires each worker to exchange their local solutions/gradients with its neighbors at every iteration, the total number of communication rounds in D-PSGD is the same as that in parallel mini-batch SGD. Another notable method to reduce communication overhead in parallel mini-batch SGD is to let each worker use compressed gradients rather than raw gradients for communication. For example, quantized SGD studied in (?)(?)(?) or sparsified SGD studied in (?)(?)(?) allow each worker to pass low bit quantized or sparsified gradients to the server at every iteration by sacrificing the convergence to a mild extent. Similarly to D-PSGD, such gradient compression based methods require message passing at every iteration and hence their total number of communication rounds is still the same as that in parallel mini-batch SGD.

Recall that parallel mini-batch SGD can be equivalently interpreted as a procedure where at each iteration each local worker first takes a single SGD step and then replaces its own solution by the average of individual solutions. With a motivation to reduce the number of inter-node communication rounds, a lot of works suggest to reduce the frequency of averaging individual solutions in parallel mini-batch SGD. Such a method is known as model averaging and has been widely used in practical training of deep neural networks. Model averaging can at least date back to (?) (?) where individual models are averaged only at the last iteration before which all workers simply run SGD in parallel. The method in (?) (?), referred to as one-shot averaging, uses only one single communication step at the end and is numerically shown to have good solution quality in many applications. However, it is unclear whether the one-shot averaging can preserve the linear speed-up w.r.t. the number of workers. In fact, (?) shows that one-shot averaging can yield inaccurate solutions for certain non-convex optimization. As a remedy, (?) suggests more frequent averaging should be used to improve the performance. However, the understanding on how averaging frequency can affect the performance of parallel SGD is quite limited in the current literature. Work (?) proves that by averaging local worker solutions only every $I$ iterations, parallel SGD has convergence rate $O(\sqrt{I}/\sqrt{NT})$ for non-convex optimization.In this paper, we shall show that if $I$ is chosen as $I=O(T^{1/4}/N^{3/4})$, parallel SGD for non-convex optimization does not lose any factor in its convergence rate. That is, the convergence slows down by a factor of $I$ by saving $I$ times inter-node communication. A recent exciting result reported in (?) proves that for strongly-convex minimization, model averaging can achieve a linear speed-up w.r.t. $N$ as long as the averaging (communication) step is performed once at least every $I=O(\sqrt{T}/\sqrt{N})$ iterations. Work (?) provides the first theoretical analysis that demonstrates the possibility of achieving the same linear speedup attained by parallel mini-batch SGD with strictly less communication for strongly-convex stochastic optimization. However, it remains as an open question in (?) whether it is possible to achieve $O(1/\sqrt{NT})$ convergence for non-convex optimization, which is the case of deep learning.

On the other hand, many experimental works (?) (?) (?) (?) (?) (?) observe that model averaging can achieve a superior performance for various deep learning applications. One may be curious whether these positive experimental results are merely coincidences for special case examples or can be attained universally. In this paper, we shall show that model averaging indeed can achieve $O(1/\sqrt{NT})$ convergence for non-convex optimization by averaging only every $I=O(T^{1/4}/N^{3/4})$ iterations. That is, the same $O(1/\sqrt{NT})$ convergence is preserved for non-convex optimization while communication overhead is saved by a factor of $O(T^{1/4}/N^{3/4})$. To our knowledge, this paper is the firstAfter the preprint (?) of this paper is posted on ArXiv in July 2018, another work (?) subsequently analyzes the convergence rate of model averaging for non-convex optimization. Their independent analysis relaxes our bounded second moment assumption but further assumes all $f_{i}(\mathbf{x})$ in formulation (1) are identical, i.e, all workers must access a common training set when training deep neural networks. to present provable convergence rate guarantees (with the linear speed-up w.r.t. number of workers and less communication) of model averaging for non-convex optimization such as deep learning and provide guidelines on how often averaging is needed without losing the linear speed-up.

Besides reducing the communication cost, the method of model averaging also has the advantage of reducing privacy and security risks in the federated learning scenario recently proposed by Google in (?). This is because model averaging only passes deep learning models, which are shown to preserve good differential privacy, and does not pass raw data or gradients owned by each individual worker.

Parallel Restarted SGD and Its Performance Analysis

Throughout this paper, we assume problem (1) satisfies the following assumption.

Smoothness: Each function $f_{i}(\mathbf{x})$ is smooth with modulus $L$.

Bounded variances and second moments: There exits constants $\sigma>0$ and $G>0$ such that

Consider the simple parallel SGD described in Algorithm 1. If we divide iteration indices into epochs of length $I$, then in each epochs all $N$ workers are running SGD in parallel with the same initial point $\overline{\mathbf{y}}$ that is the average of final individual solutions from the previous epoch. This is why we call Algorithm 1 “Parallel Restarted SGD”. The “model averaging” technique used as a common practice for training deep neural networks can be viewed as a special case since Algorithm 1 calculates the model average to obtain $\overline{\mathbf{y}}$ every $I$ iterations and performs local SGDs at each worker otherwise. Such an algorithm is different from elastic averaging SGD (EASGD) proposed in (?) which periodically drags each local solution towards their average using a controlled weight. Note that synchronization (of iterations) across $N$ workers is not necessary inside each epoch of Algorithm 1. Furthermore, inter-node communication is only needed to calculate the initial point at the beginning of each epoch and is longer needed inside each epoch. As a consequence, Algorithm 1 with $I>1$ reduces its number of communication rounds by a factor of $I$ when compared with the classical parallel mini-batch SGD. The linear speed-up property (w.r.t. number of workers) with $I>1$ is recently proven only for strongly convex optimization in (?). However, there is no theoretical guarantee on whether the linear speed-up with $I>1$ can be preserved for non-convex optimization, which is the case of deep neural networks.

as the average of local solution $\mathbf{x}_{i}^{t}$ over all $N$ nodes. It is immediate that

where $\overline{\mathbf{x}}^{t}$ is defined in (4) and $G$ is the constant defined in Assumption 1.

Fix $t\geq 1$ and $i\in\{1,2,\ldots,N\}$. Note that Algorithm 1 calculates the node average $\overline{\mathbf{y}}\overset{\Delta}{=}\frac{1}{N}\sum_{i=1}^{N}\mathbf{x}_{i}^{t-1}$ every $I$ iterations. Consider the largest $t_{0}\leq t$ such that $\overline{\mathbf{y}}=\overline{\mathbf{x}}^{t_{0}}$ at iteration $t_{0}$ in Algorithm 1. (Note that such $t_{0}$ must exist and $t-t_{0}\leq I$.) We further note, from the update equations (2) and (3) in Algorithm 1, that

where (a)-(c) follows by using the inequality $\|\sum_{i=1}^{n}\mathbf{z}_{i}\|^{2}\leq n\sum_{i=1}^{n}\|\mathbf{z}_{i}\|^{2}$ for any vectors $\mathbf{z}_{i}$ and any positive integer $n$ (using $n=2$ in (a), $n=t-t_{0}$ in (b) and $n=N$ in (c)); and (d) follows from Assumption 1. ∎

Consider problem (1) under Assumption 1. If $0<\gamma\leq\frac{1}{L}$ in Algorithm 1, then for all $T\geq 1$, we have

where $f^{\ast}$ is the minimum value of problem (1).

Fix $t\geq 1$. By the smoothness of $f$ , we have

where (a) follows from (5); (b) follows because

where (b) follows from $0<\gamma\leq\frac{1}{L}$ and (a) follows because

where the first inequality follows by using $\|\sum_{i=1}^{N}\mathbf{z}_{i}\|^{2}\leq N\sum_{i=1}^{N}\|\mathbf{z}_{i}\|^{2}$ for any vectors $\mathbf{z}_{i}$; the second inequality follows from the smoothness of each $f_{i}$ by Assumption 1; and the third inequality follows from Lemma 1.

Dividing (11) both sides by $\frac{\gamma}{2}$ and rearranging terms yields

Summing over $t\in\{1,\ldots,T\}$ and dividing both sides by $T$ yields

where (a) follows because $f^{\ast}$ is the minimum value of problem (1). ∎

The next corollary follows by substituting suitable $\gamma,I$ values into Theorem 1.

Consider problem (1) under Assumption 1. Let $T\geq N$.

For non-convex optimization, it is generally impossible to develop a convergence rate for objective values. In Theorem 1 and Corollary 1, we follow the convention in literature (?) (?) (?) to use the (average) expected squared gradient norm to characterize the convergence rate. Note that the average can be attained in expectation by taking each $\overline{\mathbf{x}}^{t-1}$ with an equal probability $1/T$.

From Theorem 1 and Corollary 1, we have the following important observations:

Linear Speedup: By part (1) of Corollary 1, Algorithm 1 with any fixed constant $I$ has convergence rate $O(\frac{1}{\sqrt{NT}}+\frac{N}{T})$. If $T$ is large enough, i.e., $T>N^{3}$, then the term $\frac{N}{T}$ is dominated by the term $\frac{1}{\sqrt{NT}}$ and hence Algorithm 1 has convergence rate $O(\frac{1}{\sqrt{NT}})$. That is, our algorithm achieves a linear speed-up with respect to the number of workers. Such linear speedup for stochastic non-convex optimization was previously attained by decentralized-parallel stochastic gradient descent (D-PSGD) considered in (?) by requiring at least $T>N^{5}$. See, e.g., Corollary 2 in (?).In fact, for a ring network considered in Theorem 3 in (?), D-PSGD requires a even larger $T$ satisfying $T>N^{9}$ since its implementation depends on the network topology. In contrast, the linear speedup of our algorithm is irrelevant to the network topology.

Communication Reduction: Note that Algorithm 1 requires inter-node communication only at the iterations that are multiples of $I$. By Corollary 1, it suffices to choose any $I\leq\frac{T^{1/4}}{N^{3/4}}$ to ensure the $O(\frac{1}{\sqrt{NT}})$ convergence of our algorithm. That is, compared with parallel mini-batch SGD or the D-PSGD in (?), the number of communication rounds in our algorithm can be reduced by a factor $\frac{T^{1/4}}{N^{3/4}}$. Although Algorithm 1 does not describe how the node average $\overline{\mathbf{y}}$ is obtained at each node, in practice, the simplest way is to introduce a parameter server that collects all local solutions and broadcasts their average as in parallel mini-batch SGD (?). Alternatively, we can perform an all-reduce operation on the local models(without introducing a server) such that all nodes obtain $\overline{\mathbf{y}}$ independently and simultaneously. (Using an all-reduce operation among all nodes to obtain gradients averages has been previously suggested in (?) for distributed training of deep learning.)

Extensions

Note that Corollary 1 assumes time horizon $T$ is known and uses a constant learning rate in Algorithm 1. In this subsection, we consider the scenario where the time horizon $T$ is not known beforehand and develop a variant of Algorithm 1 with time-varying rates to achieve the same computation and communication complexity. Compared with Algorithm 1, Algorithm 2 has the advantage that its accuracy is being improved automatically as it runs longer.

Although Algorithm 2 introduces the concept of epoch for the convenience of description, we note that it is nothing but a parallel restarted SGD where each worker restarts itself every epoch using the node average of the last epoch’s final solutions as the initial point. If we sequentially reindex $\{\mathbf{x}_{i}^{s,k}\}_{s\in\{1,\ldots,S\},k\in\{1,\ldots,K^{s}\}}$ as $\mathbf{x}_{i}^{t}$ (note that all $\mathbf{x}_{i}^{s,0}$ are ignored since $\mathbf{x}_{i}^{s,0}=\mathbf{x}_{i}^{s-1,K^{s-1}}$), then Algorithm 2 is mathematically equivalent to Algorithm 1 except that time-varying learning rates $\gamma^{s}$ are used in different epochs. Similarly to (4), we can define $\overline{\mathbf{x}}^{s,k}$ via $\overline{\mathbf{x}}^{s,k}\overset{\Delta}{=}\frac{1}{N}\sum_{i=1}^{N}\mathbf{x}_{i}^{s,k}$ and have

Consider problem (1) under Assumption 1. If we choose $K^{s}=\lceil\frac{s^{1/3}}{N}\rceil$ and $\gamma^{s}=\frac{N}{s^{2/3}}$ in Algorithm 2, then for all $S\geq 1$, we haveA logarithm factor $\log(NT)$ is hidden in the notation $\widetilde{O}(\cdot)$.

Asynchronous Implementations in Heterogeneous Networks

Algorithm 1 requires all workers to compute the average of individual solutions every $I$ iterations and synchronization among local workers are not needed before averaging. However, the fastest worker still needs to wait until all the other workers finish $I$ iterations of SGD even if it finishes its own $I$ iteration SGD much earlier. (See Figure 1 for a $2$ worker example where one worker is significantly faster than the other. Note that orange “syn” rectangles represent the procedures to compute the node average.) As a consequence, the computation capability of faster workers is wasted. Such an issue can arise quite often in heterogeneous networks where nodes are equipped with different hardwares.

Intuitively, if one worker finishes its $I$ iteration local SGD earlier, to avoid wasting its computation capability, we might want to let this worker continue running its local SGD until all the other workers finish their $I$ iteration local SGD. However, such a method can drag the node average too far towards the local solution at the fastest worker. Note that if $f_{i}(\cdot)$ in (1) are significantly different from each other such that the minimizer of $f_{i}(\cdot)$ at the $i$-th worker, which is the fastest one, deviates the true minimizer of problem (1) too much, then dragging the node average towards the fastest worker’s local solution is undesired. In this subsection, we further assume that problem (1) satisfies the following assumption:

Note that Assumption 2 is satisfied if all local workers can access a common training data set or each local training data set is obtained from uniform sampling from the global training set. Consider the restarted local SGD for heterogeneous networks described in Algorithm 3. Note that if $I_{i}\equiv I,\forall i$ for some fixed constant $I$, then Algorithm 3 degrades to Algorithm 1.

In practice, if the hardware configurations or measurements (from previous experiments) of each local worker are known, we can predetermine the value of each $I_{i}$, i.e., if worker $i$ is two times faster than worker $j$, then $I_{i}=2I_{j}$. Alternatively, under a more practical implementation, we can set a fixed time duration for each epoch and let each local worker keep running its local SGD until the given time elapses. By doing so, within the same time duration, the faster a worker is, the more SGD iterations it runs. In contrast, if we apply Algorithm 1 in this setting, then all local workers have to run the same number of SGD iterations as that can be run by the slowest worker within the given time interval. This subsection shows that, under Assumption 2, Algorithm 3 can achieve a better performance than Algorithm 1 in heterogeneous networks where some workers are much faster than others.

Without loss of generality, this subsection always indexes local workers in a decreasing order of speed. That is, worker $1$ is the fastest while worker $N$ is the slowest. If we run Algorithm 3 by specifying a fixed wall clock time duration for each epoch, during which each local worker keeps running its local SGD, then we have $I_{1}\geq I_{2}\geq\cdots\geq I_{N}$. Fix epoch index $s$, note that for $i\neq 1$, variables $\mathbf{x}_{i}^{s,k}$ with $k>I_{i}$ is never used. However, for the convenience of analysis, we define

Conceptually, the above equation can be interpreted as assuming worker $i$, which is slower than worker $1$, runs extra $I_{1}-I_{i}$ iterations of SGD by using as an imaginary stochastic gradient (with no computation cost). See Figure 2 for a $2$ worker example where $I_{1}=16$ and $I_{2}=8$. Using the definition $\overline{\mathbf{x}}^{s,k}\overset{\Delta}{=}\frac{1}{N}\sum_{i=1}^{N}\mathbf{x}_{i}^{s,k}$, we have

Consider problem (1) under Assumptions 1 and 2. Suppose all workers are indexed in a decreasing order of their speed, i.e., worker $1$ is the fastest and worker $N$ is the slowest. If $0<\gamma\leq\frac{1}{L}$ in Algorithm 3, then for all $S\geq 1$,

where $j_{k}$ for each given $k$ is the largest integer in $\{1,2,\ldots,N\}$ such that $k\leq I_{j_{k}}$(That is, for each fixed $k$, $j_{k}$ is the number of workers that are still using sampled true stochastic gradients to update their local solutions at iteration $k$.); and $f^{\ast}$ is the minimum value of problem (1).

The next corollary shows that Algorithm 3 in heterogeneous networks can ensure the convergence and preserve the same $O(1/\sqrt{NT})$ convergence rate with the same $O(\frac{T^{1/4}}{N^{3/4}})$ communication reduction.

Consider problem (1) under Assumptions 1 and 2. Let $T\geq N$. If we use $\gamma=\Theta(\frac{\sqrt{N}}{\sqrt{T}})$ such that $\gamma\leq\frac{1}{L}$, I_{i}=\Theta\big{(}\frac{T^{1/4}}{N^{3/4}}\big{)},\forall i and $S=\frac{T}{I_{N}}$ in Algorithm 3, then

where $j_{k}$ for each given $k$ is the largest integer in $\{1,2,\ldots,N\}$ such that $k\leq I_{j_{k}}$.

This simply follows by substituting values of $\gamma,I_{i},S$ into (16) in Theorem 3. ∎

Note that once $I_{i}$ values are known, then $j_{k}$ for any $k$ in Theorem 3 and Corollary 2 are also available by its definition. To appreciate the implication of Theorem 2, we recall that Algorithm 1 can be interpreted as a special case of Algorithm 3 with $I_{i}\equiv I_{N},\forall i$, i.e., all workers can only run the same number (determined by the slowest worker) of SGD iterations in each epoch. In this perspective, Theorem 1 (with $I=I_{N}$) implies that the performance of Algorithm 1 is given by

Note that the left sides of (16) and (17) (weighted average expressions) can be attained by taking each $\overline{\mathbf{x}}^{s,k-1}$ randomly with a probability equal to the normalized weight in the summation. The first error term in (16) is strictly smaller than that in (17) while the second error term in (16) is larger than that in (17). Note that the constant factor $f(\overline{\mathbf{x}}^{0})-f^{\ast}$ in the first error term in (17) is large when a poor initial point $\overline{\mathbf{x}}^{0}$ is chosen (and dominates the second error term if $f(\overline{\mathbf{x}}^{0})-f^{\ast})\geq 2\gamma^{3}I_{N}^{3}G^{2}L^{2}S$). So the main message of Theorem 3 is that if a poor initial point is selected, Algorithm 2 can possibly converges faster than Algorithm 1 (at least for the first few epochs) in a heterogeneous network.

Experiment

The superior training speed-up performance of model averaging has been empirically observed in various deep learning scenarios, e.g., CNN for MNIST in (?)(?)(?); VGG for CIFAR10 in (?); DNN-GMM for speech recognition in (?) (?); and LSTM for language modeling in (?). A thorough empirical study of ResNet over CIFAR and ImageNet is also available in the recent work (?). In Figures 3 and 4, we compare model averaging, i.e., PR-SGD (Algorithm 1) with $I\in\{4,8,16,32\}$) with the classical parallel mini-batch SGDThe classical parallel mini-batch SGD is equivalent to Algorithm 1 with $I=1$. Our implementation with Horovod uses the more efficient “all-reduce“method rather than the “parameter server” method to synchronize information between workers. by training ResNet20 with CIFAR10 on a machine with $8$ P100 GPUs. Our implementation uses Horovod (?) for inter-worker communication and uses PyTorch 0.4 for algorithm implementations.

Conclusion

This paper studies parallel restarted SGD, which is a theoretical abstraction of the “model averaging” practice widely used in training deep neural networks. This paper shows that parallel restarted SGD can achieve $O(1/\sqrt{NT})$ convergence for non-convex optimization with a number of communication rounds reduced by a factor $O(T^{1/4}$) compared with that required by the classical parallel mini-batch SGD.

References

Supplement

where $G$ is the constant defined in Assumption 1.

This lemma trivially extends Lemma 1 by restricting our attentions to each particular epoch $s$. ∎

Main Proof of Theorem 2: Fix $s$ and $k$. Following the lines in the proof of Theorem 1 until (10) (by replacing $\overline{\mathbf{x}}^{t},\gamma$ with $\overline{\mathbf{x}}^{s,k},\gamma^{s}$, respectively), we have

Note that when $s$ is sufficiently large, $\gamma^{s}=\frac{N}{s^{2/3}}$ is eventually between $(0,\frac{1}{L}]$. That is, for large $s$, e.g., $s>\lceil(NL)^{3/2}\rceil$, we have

Recall that $f(\cdot)$ has stochastic gradients with bounded second order moments by Assumption 1. There exists a constant $C$ such that for $s\leq\lceil(NL)^{3/2}\rceil$, we have

Summing (19) over $s\in\{1,\ldots,\lceil(NL)^{3/2}\rceil\},k\in\{1,2,\ldots,K^{s}\}$ and (18) over $s\in\{\lceil(NL)^{3/2}\rceil+1,\ldots,S\},k\in\{1,2,\ldots,K^{s}\}$ (noting that $\overline{\mathbf{x}}^{s,T_{s}}=\overline{\mathbf{x}}^{s+1,0},\forall s$ and $\overline{\mathbf{x}}^{1,0}=\overline{\mathbf{x}}^{0}$) yields

where (a) follows because $f^{\ast}$ is the minimum value of problem (1).

Note that there exists constant $c_{1}>0$ such that

and there exits constant $c_{2}>0$ such that

and there exits constant $c_{3}>0$ such that

where (a) in the above three equations (21)-(23) follows by noting that if $\sum_{s=1}^{S}K^{s}=\sum_{s=1}^{S}\lceil\frac{s^{1/3}}{N}\rceil=T$, then $S=\Theta((NT)^{3/4})$

Dividing both sides of (20) by $\sum_{s=1}^{S}\sum_{k=1}^{K^{s}}\frac{\gamma^{s}}{2}$ and substituting (21)-(23) into it yields

Proof of Theorem 3

where $G$ is the constant defined in Assumption 1.

Fix $i,s$ and $k$. Let $\overline{\mathbf{y}}=\frac{1}{N}\sum_{i=1}^{N}\mathbf{x}_{i}^{s-1,I_{i}}$, which is the common initial point of epoch $s$. Note that

where (a) follows by using the inequality $\|\sum_{i=1}^{n}\mathbf{z}_{i}\|^{2}\leq n\sum_{i=1}^{n}\|\mathbf{z}_{i}\|^{2}$ for any vectors $\mathbf{z}_{i}$ and any positive integer $n$; and (b) follows by using the same inequality (noting that there are less than $N$ terms in the summation $\sum_{l:\tau\leq I_{l}}\mathbf{G}_{l}^{s,\tau}$) and applying Assumption 1. ∎

Main Proof of Theorem 3: Fix $s$ and $k\in\{1,2,\ldots,I_{1}\}$. Recall that $j_{k}$ is the largest integer in $\{1,2,\ldots,N\}$ such that $k\leq I_{j_{k}}$. That is, at iteration $k$ in epoch $s$, only the first $j_{k}$ workers are using true stochastic gradients to update and the other workers stop updating (since they are too slow and do not have the chance to perform the $k$-th update in this epoch). Note that as $k$ increases to $I_{1}$, $j_{k}$ decreases to $1$. By the definition of $j_{k}$, we have

where (a)-(c) follows by using the same arguments used in showing (9).

Substituting (25) and (26) into (24) yields

where (a) follows because $0<\gamma\leq\frac{1}{L}\leq\frac{N}{j_{k}L}$ (noting that $j_{k}\leq N,\forall k$).

where (a) follows because $f_{i}(\cdot)$ are identical by Assumption 2; (b) follows by using the inequality $\|\sum_{i=1}^{n}\mathbf{z}_{i}\|^{2}\leq n\sum_{i=1}^{n}\|\mathbf{z}_{i}\|^{2}$ for any vectors $\mathbf{z}_{i}$ and any integer $n$; (c) follows from the smoothness of each $f_{i}$ by Assumption 1; and (d) follows from Lemma 3.

Dividing both sides by $\frac{\gamma}{2}$ and rearranging terms yields

As a sanity check of our analysis, we note that if $I_{1}=I_{2}=\cdots=I_{N}$, then $j_{k}=N,\forall k$ and (30) reduces to (11). (Recall that Algorithm 2 is identical to Algorithm 1 in this case.)

Summing (30) over $k\in\{1,\ldots,I_{1}\}$, $s\in\{1,2,\ldots,S\}$ (noting that $\overline{\mathbf{x}}^{s,I_{1}}=\overline{\mathbf{x}}^{s+1,0},\forall s$ and $\overline{\mathbf{x}}^{1,0}=\overline{\mathbf{x}}^{0}$) and dividing both sides by $\sum_{s=1}^{S}\sum_{k=1}^{I_{1}}\frac{j_{k}}{N}=S\frac{1}{N}\sum_{i=1}^{N}I_{i}$ yields

where (a) follows because $f^{\ast}$ is the minimum value of problem (1).