A Quadratic Synchronization Rule for Distributed Deep Learning
Authors: Xinran Gu
Kaifeng Lyu
Sanjeev Arora
Jingzhao Zhang
Longbo Huang
Venue: ICLR 2024 poster (2024)
Keywords: distributed training
Local SGD
local gradient methods
generalization
implicit bias
sharpness
Primary area: optimization
TL;DR: We propose Quadratic Synchronization Rule (QSR) to dynamically set the synchronization period in local gradient methods based on the learning rate, reducing wall-clock training time and improving test accuracy of ResNet and ViT on ImageNet.

Abstract: In distributed deep learning with data parallelism, synchronizing gradients at each training step can cause a huge communication overhead, especially when many nodes work together to train large models.
  Local gradient methods, such as Local SGD, address this issue by allowing workers to compute locally for $H$ steps without synchronizing with others, hence reducing communication frequency.
  While $H$ has been viewed as a hyperparameter to trade optimization efficiency for communication cost, recent research indicates that setting a proper  $H$ value can lead to generalization improvement. Yet, selecting a proper $H$ is elusive. This work proposes a theory-grounded method for determining $H$, named the Quadratic Synchronization Rule (QSR), which recommends dynamically setting $H$ in proportion to $\frac{1}{\eta^2}$ as the learning rate $\eta$ decays over time.
  Extensive ImageNet experiments on ResNet and ViT show that local gradient methods with QSR consistently improve the test accuracy over other synchronization strategies. Compared to the standard data parallel training, QSR enables Local AdamW to cut the training time on 16 or 64 GPUs down from 26.7 to 20.2 hours or from 8.6 to 5.5 hours and, at the same time, achieves 1.16% or 0.84% higher top-1 validation accuracy.

Published as a conference paper at ICLR 2024
A QUADRATIC SYNCHRONIZATION RULE FOR DISTRIBUTED DEEP LEARNING
Xinran Gu Kaifeng Lyu"* Sanjeev Arora Jingzhao Zhang" 122.3t Zhangl.2.3f Longbo Huang'
Institute for Interdisciplinary Information Sciences, Tsinghua University "Shnngaai Qizhi Institute $Shangaai Al Laboratory
*Department of Computer Science & Princeton Language and Intelligence, Princeton University gxr21@mails ts tsiinguua... edu. en {j Ljingzaozz longbohuang @t singhua edu. on
(klyu, arora @cs. wrwnceton. edu
ABSTRACT
In distributed deep learning with data parallelism, synchronizing gradients at each training step can cause a huge communication overhead, especially when many nodes work together to train large models. Local gradient methods, such as Local SGD, address this issue by allowing workers to compute locally for H steps with- out synchronizing with others, hence reducing communication frequency. While H has been viewed as a hyperparameter to trade optimization efficiency for com- munication cost, recent research indicates that setting a proper H value can lead generalization improvement. Yet, selecting proper H is elusive. This work proposes a theory-grounded method for determining H, named the Quadratic Syn- chronization Rule (QSR), which recommends dynamically setting H in proportion to + as the learning rate vi decays over time. Extensive ImageNet experiments on ResNet and ViT show that local gradient methods with QSR consistently improve the test accuracy over other synchronization strategies. Compared with the stan- dard data parallel training, QSR enables Local AdamW on ViT-B to cut the train- ing time on 16 or GPUs down from 26.7 to 20.2 hours or from 8.6 to 5.5 hours and, at the same time, achieves 1.12% or 0.84% higher top-1 validation accuracy.
INTRODUCTION
The growing scale of deep learning necessitates distributed training to reduce the wall-clock time. Data parallel training is a foundational technique that distributes the workload of gradient computa- tion K workers, also serving as a key building block of more advanced parallel strategies. At each step of this method, each worker first computes gradients on their own local batches of data. Then, they take an average over local gradients, which typically involves a costly All-Reduce operation. Finally, they update the model parameter with the averaged gradient and gradient-baas based optimizer OPT, c.g., SGD, AdamW. W. In this paper, we term the data parallel implementation of optimizer OPT as "Parallel OPT". See Algorithm 1 for the pseudocode. The cost for this data parallelism is obvi- ous. Frequent gradient synchronization can induce huge communication overhead as the number of workers and model size grow, severely hindering the scalability of distributed training (Tang al., 2021; Li et al., 2022; Xu et al., 2023).
One approach reducing this communication overhead is Local SGD (Stich, 2018; Zhou & Cong. 2018; Woodworth et al., 2020). Rather than synchronizing gradients at every step, Local SGD allows workers independently train their local replicas using their own local batches with SGD updates. It is only after completing H 1 local steps that these workers synchronize, where the model parameters get averaged over all replicas. Notably, while we mention SGD, this approach can be readily adapted to other popular optimizers. In this paper, if gradient-based optimizer OPT is used for local updates, we term the variant as "Local OPT" (e.g.. Local SGD, Local AdamW), and collectively refer to this class of approaches as local gradient methods. We provide pseudocode for local gradient methods in Algorithm 2.
The main focus of this paper is to study the best strategies to set the synchronization period H (i.e., the number of local steps per communication round) in local gradient methods. While setting to larger value reduces communication, a very large H can hinder the training loss from decreasing
"quual contribution Corresponning authors

---

Published as a conference paper at ICLR 2024
Scheciile and Communcation Volume
H Schodule and Communcation Volume
Parallel SGD, 100%
H ~ 19.3%
Parallel AdamW, 100%
1f n. 15.1%
Post-local SGD, 56.3%
Constant Period, 25% QSR, 20.1%
Constant Period, 25%
QSR, 10.4%
80
80-
75-
Top-4 Val. Ave. (N)
70
660.5
70
81
80.0
65
80
65-
79.5
79.0
100
160
110)
200
250
300
00
0
50
100
130
200
100 200
300
Epucts
Epocht
(a) Local SGD on ResNet-152
(b) Local AdamW on ViT-B
Figure 1: When training 152 and ViT-B on ImageNet with cosine learning rate decay, Local SGD/AdamW with QSR consistently outperforms data parallel methods or Local SGD/AdamW with other synchronization strategies in terms of top-1 validation accuracy, while only requiring 20.1% and 10.4% of the communication volume, respectively. With QSR, Local SGD on ResNet Local AdamW on ViT the training time from 20.7 to hours or 26.7 to 20.2 hours on GPUs, when compared with data parallel meth- ods. We report the mean and the standard deviation over 3 runs. See Appendix C for training details.
normal speed, since the local replicas may significantly diverge from each other before averaging. Indeed, it has been observed empirically that larger H leads to higher training loss after the same number of steps (Wang& & Joshi, 2021; Ortiz al., 2021), and efforts analyze the convergence of local gradient methods in theory usually end up with loss bounds increasing with H (Khaled et al., 2020; Stich, 2018; Haddadpour al., 2019; Yu al., 2019). To better trade-off between commu- nication cost and optimization speed, Kamp et al. (2014): Wang & Joshi (2019); Haddadpour et al. (2019); Shen et al. (2021) proposed adaptive synchronization schemes, such as linearly increasing H as the iteration goes on (Haddadpour et al., 2019), adjusting H/ based on the variance in model parameters (Kamp et al., 2014). Nonetheless, their effectiveness has only been validated on linear models or small-scale datasets, c.g., CIFAR-10/100.
All these strategies are developed to avoid sacrificing too much training loss, but training loss is never the final evaluation metric that one cares about in deep learning. Due to the overparameterized nature of modern neural networks, reaching the same training loss does not correspond to the same performance test data. It has also been long known that the choice of optimizers hyperpa- rameters can change not only the optimization speed of the training loss but also their implicil bias towards solutions with different test accuracies.
The presence of this implicit bias indeed complicates the picture of setting H in local gradient methods. Though a large I might be harmful for training loss, it has been observed empirically that setting !/ properly can sometimes improve rather than hurt the final test accuracy. Lin et al. (2020) are the first to report this phenomenon. Comparing with running just the standard data parallel SGD (equivalent to H = 1), they observed that switching from SGD to Local SGD (H > 1) halfway through consistently leads to higher final test accuracy. Local SGD with this specific schedule of H is designated as Post-local SGD. Lin et al. (2020)'s work opens up a new angle in setting H in local gradient methods, yet, the proposed schedule in Post-local SGD, referred to as the post-local schedule in this paper, is suboptimal in improving test accuracy. It was later reported by Ortiz et al. (2021) that Post-local SGD does not improve much on ImageNet. For both stepwise decay and cosine decay learning rate schedules, the test accuracy improvement of Post-local SGD diminishes as learning rate decreases. Further, it remains unclear whether the generalization benefit continues to appear when the optimizer is changed from SGD to adaptive gradient methods such as Adam/Adea Adam W. which are now indispensable for training large models.
Our Contributions. In this paper, we aim propose general and effective H schedule that can be readily applied to various optimizers and neural network models. Specifically, we introduce a simple yet effective strategy, called Quadratic Synchronization Rule (QSR), for dynamically adjust- ing the synchronization period according to the learning rate: given a learning rate schedule, we set H proportional to 17 -2 as the learning rate 7) decays. This rule is largely inspired by a previous theo- retical work (Gu et al., 2023), which shows that the generalization benefits arise only if H 2 when 17 0, but did not make any recommendation on how to set H.
Our main contributions are:
2

---

Published as a conference paper at ICLR 2024
1. We propose the Quadratic Synchronization Rule (QSR) to simultaneously reduce the wall-clock time and improve the final test accuracy of local gradient methods. Based on the theoretical insights in Theorem 3.1, we provide theoretical separation among data parallel SGD, Local SGI) with H ~ y and Local SGD with QSR in terms of SDE approximations. We show that QSR can help reduce sharpness faster and hence improve generalization.
2. We demonstrate with ImageNet experiments that QSR consistently improve the final test accuracy of ResNet-152 and ViT-B over other synchronization strategies, including constant- period and post-local schedules, and also H ~ 1 which one will expect to be optimal from the optimization perspective (Figure 1).
3. We thoroughly validate the efficacy of QSR not only for Local SGD but also for Local Adam W, which is arguably more suitable for training large models. We also validate its efficacy for cosine, linear and step decay learning rate schedules that are commonly used in practice.
4. We evaluate the communication efficiency of QSR on 64-GPU NVIDIA GeForce RTX 3090 cluster. As an illustrative example, the standard data parallel Adam takes 8.6 hours to train VIT-B for 300 epochs. With QSR, Local AdamW cuts the training time down to 5.5 hours with even higher test accuracy.
2 OUR METHOD: QUADRATIC SYNCHRONIZATION RULE
Below formulate the local gradient methods and present our Quadratic Synchronization Rule.
Local Gradient Methods. Given any gradient-based optimizer OPT, the corresponding local gra- dient method consists of multiple communication rounds. At the s-th round, each of the K workers
(say the k-t gets a local copy of the global iterate ((+),, i.c., {(*) < g(*), and then performs H steps of local updates. At the h-th local step of the s-th round, which corresponds to the (sH + h)-th
iteration globally, each worker gets batch of Bioe samples (5k,h,1 Bt.c from globally shared dataset D, computes the gradient on that batch, and updates the model with optimizer OPT and learning rate +/s
c(e)
1 Bioe
k + (PT(( k,h 7.H Ik.h) where Ik.h Bioc 1=1
(I)
After finishing steps of local updates, all workers average their local models to generate the ((a) next global iterate: g(++1) + k Ex= K. Note that conventional local gradient methods set the synchronization period as a constant, denoted as H. throughout training. See also Algorithm 2.
Quadratic Synchronization Rule. Given learning rate schedule Tr. f E {0..... T 1} that decays with time, instead of keeping H constant, we propose to dynamically increase the synchro- nization period (+) at each round 8 as the learning rate decreases. More specifically, if at the global iteration we need start new communication round, then we set
2
(+)) = max
Hinses
Here Hhase is a constant indicating the minimum number of local steps one would like to use for each round, which should be set according to the relative cost of computation and communication. The coefficient a, termed the "growth coefficient" henceforth, is hyperparameter controlling how fast (+) increases as 17, decreases.
As suggested by our later theorem 3. a should be set as a small constant. In our experiments, we tune properly between 0.01 and 0.5 and test the effectiveness of QSR with H 2, 4, 8. Note that the last communication round may not finish exactly at the last iteration the learning rate schedule. If this is the case, we force synchronization the last step by setting (+)) ==
A surprising part of our method is that we use the power 2 in the above formula (2). This choice of power 2 is inspired by the analysis in Gu al. (2023), which suggests that setting H 2 2
is beneficial for reducing the sharpness of the local landscape. Indeed, (*)) could have been set to (x) max { of our method, and we will provide theoretical justification and empirical evidence for this choice in Section 3. We also visualize the H schedule for QSR in Figure 5 in the appendix.
Hbsse, (a)
for any 7. However, using 7 2 is crucial for the success
3

---

Published as a conference paper at ICLR 2024
Dealing with Learning Rate Warmup. Many learning rate schedules use warmup phase where the learning rate increases linearly from 0 to Mmaax and then decays monotonically. This warmup phase is often used to avoid the instability caused by the initial large learning rate (Goyal et al., 2017). Our rule is not directly compatible with the warmup phase, since it is designed for a decaying learning rate, but the learning rate increases rather than decreases in this phase. Practically, we recommend setting H(*) as the value to be used in the communication round right after the warmup.
3 THEORETICAL MOTIVATIONS OF QUADRATIC SYNCHRONIZATION RULE
To justify our choice of power 2, we build on the same theoretical setup as Gu et al. (2023) to analyze the Stochastic Differential Equation (SDE) approximation of SGD and Local SGD using different scalings of H with respect to 7/. Though the learning rate continuously decays over time in most of our experiments, it does not usually change much within a couple of epochs. Inspired by this, we take a quasistatic viewpoint: consider a significant period of time where the learning rate is relatively constant, and directly treat the learning rate as real constant 7j. First, we recap Gu et al. (2023)'s 's theory that applies to Local SGD with H ~ 17 11 then we show how to generalize the result to our rule where H ~ .7- leading to stronger implicit bias towards flatter minima.
Setup. Consider optimizing the loss function (()) := Eg E)l, where E I is the pa- rameter vector and {(0;{) {) is the loss function for a single data sample & drawn from a training set/training distribution D. We use Σ(0) == Cov co VE(())) to denote the covariance matrix of the stochastic gradient vε(0:() at 0. Following Gu et al. (2023), we make regularity assumptions on (()). (() and [|v l)ll in Assumption E. and we assumee that L has a manifold I of minimizers in Assumption E.2. Our analysis is based on SDE approximations near I, providing clean view of how different choices of H affect the selection of minimizers by Local SGD.
SDE approximations of SGD and Local SGD. SDE is a powerful tool to precisely character- ize the effect of noise in SGD, leading to many applications such as Linear Scaling Rule (Goyal et al., 2017). The SDE (1) = (V)((())) + 1 VB (/)) is conventionally used in the literature (Jastrzębski et al., 2017; Smith et al., 2020; Li et al., 2021b), where W, is the standard Wiener process. In this SDE, each discrete step corresponds to a continuous time interval of length 77 and the expected gradient and gradient noise become a deterministic drift term and a stochastic diffusion term, respectively. When the training proceeds to a point (L) near a minimizer Ço on the manifold I, the gradient VZ()(t)) almost zero but the gradient noise 1
vB [((())) drives
the parameter to diffuse locally. This can be captured by a careful first-order approximation of the dynamics, leading
process (Zhu al., 2019; Li et al., 2019a; Izmailov et al., 2018). However, these rough approximations only hold for about O(n-i) steps, whereas neural networks in practice are usually trained for much longer.
Recently, series of works (Blanc et al., 2020; Damian et al., 2021; Li et al., 2021c) study the dynamics of SGD on a longer horizon. They show that higher-order terms can accumulate over time and drive this local diffusion to gradually move on the manifold r. Among them, Li et al. (2021c) precisely characterized this with an SDE tracking the gradient flow projection of ((t) on I, denoted as & (((()) (see Definition 1). Here, ((((()) can be thought of as natural "center" of the local diffusion. This SDE, termed as Slow SDE, tracks the dynamics of SGD over 02) steps, which is much longer than the 11) horizon for conventional SDEs.
To provide a theoretical understanding of why Local SGD generalizes better than SGD, Gu et al. (2023) derived the Slow SDEs for Local SGD using the scaling H ~ n7_ By comparing the Slow SDEs, they argued that Local SGD drifts faster to flatter minima than SGD. However, their analysis does not encompass the more aggressive scaling H 17 recommended by our QSR. Recognizing this gap, we derive the Slow SDE for this scaling, enriching the theoretical framework for the generalization behavior of Local SGD. Below, first present the Slow SDEs for SGD and Local SGD with H ~ 17- and H ~ -2, then we interpret why H ~ - may generalize better. Definition 3.1 (Slow SDE for SGD, informal, (Li et al., 2021c; Gu et al., 2023)). Given So E r, define ((t) as the solution to the following SDE with initial condition ((0) Co:
dc(1) Pe(
1 D /B
/2 (C)dW, 2B 3 ((())) (()))
(3)
(a) difftstion on I
(b) drift I
4

---

Published as a conference paper at ICLR 2024
Here, Pç is projection operator of differential forms ensure that taking infinitesimal step from Ç T remains on the manifoldT. B is the total batch size. En((S) and Eg (C) are certain PSD matrices related to gradient noise and Hessian. See Definition E.2 for the full definition.
Definition 3.2 (Slow SDE for Local SGD with H ~ 17 .informal (Gu et al., 2023)). Consider the scaling H = 3/n for some constant 3. Given Ço E r, define ((t) the solution to the following SDE with initial condition Ç(0) = Soo
(c(t)= ( JB 1/2 (C)dW,
2B V D(G))
R- K- L(C) 2B
(n)jatt)
(4)
(a) diffasion on I
(b) drift I, same SGD
(e) extro drift term on T
where K is the number of workers, B, B and EC(C) are the same as Definition 3.1.
Here, (C: (3) is a PSD matrix depending on gradient noise and Hessian. It scales with 3 as lim V ((; 3) = 0, limg---oo ((; (S). See Desinition E.3 for the fall definition.
Definition 3.3 (Slow SDE for Local SGD with QSR). Given Ço E r, define ((t) as the solution to the following SDE with initial condition ((0) = Ço:
dc(t) Ps
1 VB
(S)
K 2B
S ((C))
(5)
(a) diffustion I
(b) drift on F, times larger
where K, B, (i) and Eg(C) are defined in Definitions 3.1 and 3.2.
The following approximation theorem indicates that when the learning rate and the growth co- efficient for QSR are small, the above Slow SDEs closely track their discrete counterparts. The approximation theorem for QSR is new, and we defer the proof to Appendix E.2.
Theorem 3.1 (Weak Approximations), Let T > O be constant and ((t) be the solution to of the above Slow SDES with the initial condition ((0) +()()) E T. Let g(0) be any CA- -mooohh function.
1. (Gu et al., 2023) For SGD, let ((t) be the solution to (3). Then, max B[g(
B[g({(sn?)))] = 215
2. (Gu et al., 2023) For Local SGD with H = f/r some constant 3, let ((t) be the solution to (4). Then, maxoses + [B[g( (b ((()))) -Ig)([s 25)
3. For Local SGD with H = ( where the positive constant a is small but larger than 2(1))) for all 7 > 0, let ((t) be the solution to (5). Then, max T E(g(((>Hn = O(a2).
[ (Fgg((((())))
Hn
Here, OC.) and Ō(-) hide constants that are independent of a and n but can depend on and T.
(.) also hides log terms.
By comparing the Slow SDEs, we can predict the generalization order for different scaling QSR (constant H}, which we explain in detail below.
(H
Interpretation of the Slow SDEs. We first focus on the Slow SDE for SGD (3). The key com- ponent of this Slow SDE is the drift term (b), which comes from higher-order approximations of the aforementioned local diffusion that happens in Onn steps. Viewing V V²) [5)
semi-gradient of (V²C(C) Eg (C)) that discards the dependence of 0 in Eg (C). we can interpret the Slow SDE as a continuous version of a semi-gradient method for reducing (V²L(C), (C)) on T. Since the Hessian matrix L(() determines the local curvature of the loss landscape, we conclude from the Slow SDE that SGD tends to reduce sharpness and move towards flatter minimiz- in (n) steps. Reduced sharpness has been shown to yield better sample complexity bounds in specific theoretical settings. For details, we refer the readers to Li et al. (2021c).
!Given Ç, (C: is monotonically increasing function of / in the eigenspace of the Hessian matrix VZC(C).
5

---

Published as a conference paper at ICLR 2024
0.3
0.55)
0.2
Top-1 Val. Acc. (%)
81.00
(.013)
Top-1 Val. Acz.
80.0
80.5
(100)
(100))
A LE)
60.0
(8
If
(A 0.64)
Paralle SCDD
79.5
79.5
17.5
20.0
22.5
25.0 100
10
100
Communication Volume (%)
Communication Volume (%)
(a) Local SGD on ResNet-152
(b) Local A ViT-B
Figure 2: Empirical results on Local SGD and Local AdamW validate the generalization performance order predicted by our theory: QSR {H n > (constant H). For SGD, we additionally have (constant H = (parallel SGDJ since the latter is equivalent Local SGD with I. Here, a and / are tuned to maximize the test accuracy QSR and ~ r7" ", respectively.
Now, we turn to the Slow SDE for QSR. Compared with the SDE for SGD, possesses a K times larger drift term, leading to much faster sharpness reduction than SGD. An intuitive explanation for why this extra drift arises is as follows. Since the local batch size is K times smaller than the global one, this local diffusion at each worker is much more significant than that in parallel SGD, thereby leading to an extra drift term in Slow SDE accumulated from higher-order terms.
The case of Local SGD with H B/r is somewhere in between QSR and SGD. Compared with the SDE for SGD, it has an extra drift term (c), where 3 serves as the knob to control the magnitude of the drift term. For small B, (C) diminishes to zero, yielding the same SDE as SGD. By contrast, as A goes to infinity, T (C) approximates Eg (C). leading to the Slow SDE for QSR.
Comparison of different scalings. Based on the interpretation, keeping constant as 1) dimin- ishes is equivalent to setting a small 33 for H = -3/1,, making the extra drift term negligible and thus yielding nearly no generalization benefit over SGD. Conversely, the SDE for = B/n converges to the SDE of QSR in the limit B + oo, maximizing the drift term. But in practice, & cannot be arbi- trarily large. In Theorem 3.3 Gu et al. (2023), the distance between the iterate and blows up as Ó( (B))). suggesting that setting a very large 3 for not-so-small r can blow up the loss. Therefore, the generalization performance of H~n is expected to be worse than QSR. In summary, the order of generalization performance predicted by our theory is QSR > ( リ > (constant H).
Experimental results in Figure validate that this order of generalization performance for different scalings holds not only for Local SGD but also for Local AdamW. For Local SGD we additionally have (constant H} s (parallel SGD) since parallel SGD is mathematically equivalent Local SGD with H = 1. Apart from H ~ ソーー and H ~ n we also tried a more aggressive scaling, H 7 3, but does not provide consistent improvements over QSR. See Appendix G for more discussion.
EXPERIMENTS
In this section, we empirically demonstrate that QSR not only improves the test accuracy of local gradient methods but also reduces the wall-clock time of standard data parallel training, with a focus on the ImageNet classification task (Russakovsky et al., 2015). Our experiments include Local SGD on ResNet- 152 (He et al., 2016), and Local AdamW on ViT-B with patch size 16x16 (Dosovitskiy et al., 2021). We briefly outline our training configuration below. See Appendix C for full details.
Baselines. For QSR with base synchronization period Hasee we benchmark their performance against two baselines running the same number of epochs: ① Local SGD/AdamW with constant synchronization period H Hasee and parallel SGD/Adam W. When comparing with these baselines, we mainly focus on validating that (a) QSR maintains or sometimes outperforms the communication efficiency of @n thus communicating much less than 2, and (b) QSR improves the generalization performance of O, even surpassing 2 in test accuracy.
Comparison with other synchronization strategies. Besides the above two baselines, other po- tential baselines include @ Post-local SGD, @ the scaling of Hl ~ 17 and @ large batch training with batch size H x B, which we discuss below. is proposed for the same purpose as QSR: to improve communication efficiency and generalization together. However, it is less communication efficient than our QSR because it starts with parallel SGD and sustains this for a significant frac- tion of the training duration, leading limited reduction in communication. Also, shown by

---

Published as a conference paper at ICLR 2024
Table I: QSR enhances the test accuracy of local gradient methods, even outperforming the communication- intensive data parallel approach. The experiments below use batch size 4096. We report the validation accuracy and train loss averaged over 3 runs, along with the standard deviation.
(a) Local SGD on ResNet-152
(b) Local AdamW on ViT-B
Method
Val. acc. (%) Train loss Comm.
Method
Val. acc. (%) Train loss Comm.
Parallel SGD
79.53 (0.07) 1.57 (0.01) 100%
Parallel AdamW
79.86 (0.03) 1.09 (0.00) 100%
Local SGD (H=2)
79.54 (0.07) 1.58 (0.00) 50%
Local AdamW (H=4)
79.32 (0.06) 1.01 (0.02) 25%
+ QSR (Huuas 22) 80.30 (0.04) 1.67 (0.01) 39.7% %
+ QSR (H)ms==4)
80.98 (0.05) 1.32 (0.00)
10.4%
Local SGD (H=4) 79.48 (0.12) 1.62 (0.02) 25%
Local AdamW (H=8)
78.93 (0.10) 1.06 (0.00) 12.5%
QSR (H (iman=4) 80.27 (0.05) 1.69 (0.01) 20.1 %
QSR (Hhmm=88) 80.56 (0.10) 1.35 (0.01) 6.9%
our comparison in Figure I(a) (also observed in Ortiz et al. 2021), its generalization benefits over SGD appear shortly after switching and diminish in the end. @ is inspired by Gu et al. (2023) and may also improve generalization while reducing communication, but we have conducted a thorough comparison between QSR and @ in Figure 2, demonstrating the superiority of QSR. G) has the same communication efficiency as Local SGD with the same constant H (Q), but it has been observed to have worse test accuracy than parallel SGD/Adam without scaling up the batch size (2), which we also observe in Table 2. For the above reasons, we mainly compare with baselines ① and @.
Hardware.
We conduct the experiments on Tencent Cloud, where each machine is equipped with 8 NVIDIA GeForce RIX 3090 GPUs The machines are interconnected by a 25Gbps network. Since intra-machine communication speed is not substantially faster than inter-machine speed on our specific hardware, we treat each GPU as an independent worker and set the batch size on each GPU as Bioc = 256. In this paper, we use axb GPUs to denote a machines with GPUs each.
Training Setup. Our experiments on ResNet-152 152 follow the 200-epoch recipe in Foret et al. (2021b) except that we use 5 epochs of linear learning rate warmup. For experiments on ViT-B, we follow the simple and effective 300-epoch recipe proposed in Beyer et al. (2022) with RandAug- ment and Mixup. We use the cosine decay unless otherwise stated. The hyperparameters (primarily learning rate and weight decay) are optimally tuned for all baselines. We explore H = 2, 4 for ResNet- 152 and H base = 4, 8 for ViT-B. This choice stems from the observation that the communi- cation overhead for ResNet-152 is smaller than VIT-B (see Table 4). To tune the growth coefficient a for QSR, we first fix the learning rate schedule and then search among few values of a. The a values we explore typically allow the training to start with Husse maintain H Hhase for initial period to optimize the training loss, and gradually increase H as 17 decays in the late phase.
4.1 QSR IMPROVES GENERALIZATION
Through experiments spanning various batch sizes and learning rate schedules, in this subsection, we illustrate that QSR consistently enhances the generalization of gradient methods, even outper- forming the communication-intensiis data parallel approach.
Main results. We first present our main results for batch size B = 4096 on 2x8 GPUs, covering Local SGD on ResNet-152 and Local AdamW on VIT-B. As shown in Table 1, QSR significantly improves the validation accuracy of local gradient methods by up to 0.8% on ResNet-152 and 1.7% on ViT-B, despite inducing higher training loss. The results support the thesis that the improvement in generalization is due to the implicit regularization of local gradient noise instead of better opti- mization. Noticeably, QSR surpasses the data parallel approach in validation accuracy by 0.7% on ResNet- 152 and by 1.1% on VIT-B while cutting the communication volume to less than 25%. As an added benefit of increasing the synchronization interval in line with the decaying learning rate, QSR further reduces communication overhead, even halving the communication volume compared to Local Adam W with a fixed synchronization period on ViT-B.
The advantages of QSR are more pronounced for ViT-B compared to 155 This is probably because vision transformers are general-purpose architectures with less imagc-specific inductive bias than CNNs (Dosovitskiy et al., 2021; Chen et al., 2021). As a result, they may benefit more from external regularization effects, such as those induced by adding local steps.
Scaling up the batch size. In Table 2, when scaling the training up to 8x8 GPUs with total batch size = 16384, we observe a drop in test accuracy for both data parallel approach and local gradient methods. This generalization degradation for large batch training, which has been widely observed
7

---

Published as a conference paper at ICLR 2024
Table 2: QSR mitigates the generalization degradation in barge-batch training. Here the batch size is 16384.
(a) Local SGD on ResNet-152
(b) Local Adam on ViT-B
Method
Val. Acc.(%) Comm. (%)
Method
Val. Acc. (%) Comm. (%)
Parallel SGD
79.20
100
Parallel AdamW
78.52
100
Local SGD (H=2)
78.67
50
Local AdamW (H=4)
77.83
25
+ QSR (Hase = 2)
79.27
42.8
+QSR (Hbare = 4)
79.36
16.1
Local SGD (H=4)
78,34
25
Local AdamW (H=8)
77.62
12.5
+ QSR (Hasse = 4)
78.65
21.9
+QSR (Hbase = 8)
78.26
9.8
Table 3: QSR also exhibits strong generalization performance on the step-decay learning rate schedule. (b) Local Adam ViT-B.
(a) Local SGD ResNet-152.
Method
Val. Acc. (%) Comm. (%)
Method
Val. Acc.(%) Comm
Parallel SGD
79.68
100
Parallel AdamW
79.91
100
Local SGD (H=2)
79.58
50
Local AdamW (H=4)
79.36
25
+QSR (Has 2)
80.40
40.3
+ QSR base
80.9
12.7
Local SGD (H=4)
79.53
25
Local AdamW (H=8)
79.23
12.5
+QSR (Hbase 4)
80.11
20.5
QSR(Hbase 8)
80.65
7.2
in the literature (Shallue et al., 2019; Jastrzebski et al., 2017; You et al., 2018), probably arises from reduced level of gradient noise associated with increased batch size (Keskar et al., 2017b; Smith et al., 2021). While the Linear Scaling Rule for SGD (Krizhevsky, 2014; Goyal et al., 2017) and the Square Root Scaling Rule (Malladi al., 2022; Granziol al., 2022) for adaptive gradient methods which increase the learning rate in proportion to the total batch size or its square root can mitigate this degradation, they cannot fully bridge the gap. Table 2, the test accuracy drop persists even when we tune the learning rate for all baselines. Applying QSR to local gradient methods can help reduce this generalization gap. improves the validation accuracy of local gradient methods by up to 0.6% on 152 and 1.5% on ViT-B, This enables local gradient methods to achieve comparable validation accuracy as the data parallel approach on ResNet or outperform by 0.8% on VIT while communicating considerably less.
Method Communcation Volume
AdamtW 100%
(3i... 4)
Other learning rate schedules. So far, our experi- ments are conducted with the cosine learning rate sched- ule, which is a common choice for training modern deep neural nets (Liu et al., 2021; 2022; Brown et al., 2020). To further validate the efficacy of QSR. we now investigate other popular learning rate schedules, including linear (Li et al., 2020a; Izsak et al., 2021; Leclere et al., 2023) and step decay (He al., 2016; Huang et al., 2017: Ma et al., 2019). See Figure 4 for visualization of these sched- ules. Figure 3 presents the results for Local Adam on VTT-B with linear decay, where the peak learning rates for baselines are tuned optimally. QSR improves the test ac- curacy of Local Adam by significant margin of 1.4%, even outperforming parallel AdamW by 0.6% while cut- ting the communication volume only 9.3%. The step decay scheduler divides the learning rate by factors such as 2 or 10 at some specified epochs. Given the absence of standard recipes determine the decay points in our training setup, we derive step decay schedule from the cosine decay by rounding its learning rate to powers of 2, which is defined as J/step (t) := 2rognd (()) As shown in Table 3, QSR exhibits strong generalization performance with this decay schedule, enhancing the test accuracy of local gradient methods by up to 0.8% on ResNet- 152 and 1.5% on ViT-B. It even surpasses the communication-intensii parallel SGD by 0.7% on ResNet and parallel AdamW by 1% on ViT.
Local AannW (1f 4), 25%
80
Tap-1 Val. Ace. (%)
76-
20
280
290
300
100
3001
Epochs
Figure 3: For linear decay, QSR improves the test accuracy of Local Adam on ViT- B, even outperforming the communication- intensive parallel AdamW.
4.2 QSR REDUCES WALL-CLOCK TIME
In addition to improving generalization, our original motivation for adopting local steps is to reduce communication overhead and hence reduce the wall-clock time. In this section, confirm this for training with 2x8 and 8x8 GPUs, as shown in Table 4. See also Appendix for our method of measuring the communication time. In setup, scaling the training from 2x8 8x8 GPUs
8

---

Published as a conference paper at ICLR 2024
Table 4: QSR reduces the wall-clock time of data parallel training. The following tables present wall-clock time for the entire training process on 2x8 GPUs and 8x8 GPUs, with batch sizes 4096 and 16384, respectively. We highlight the wall-clock time of QSR when it matches or outperforms the data parallel baseline in test accuracy. "Ratio" represents communication time divided by total time, reflecting the communication overhead. We also include local gradient methods with constant synchronization period for reference.
(a) ResNet- 152 (200 epochs) 2x8 GPUs
(b) ViT-B (300 epochs) on 2x8 GPUs
Method
Comm. (h)
Total (h)
Ratio (%)
Method
Comm. (h)
Total (h) Ratio(%)
Parallel SGD
3,3
20.7
15.9
Parallel AdamW
7.3
26.7
27.3
QSR (Huase 2)
1.3
18.7
7.0
QSR (Huase = 4)
0.8
20.2
4.0
QSR (Hasee = 4)
0.7
18.0
3.9
QSR (Haase = 8)
0.5
20.0
2.5
Local SGD (H=2)
1.6
19.0
8.4
Local AdamW (H=4)
1.8
21.2
8.4
Local SGD (H=4)
0.8
18.0
4.4
Local AdamW (H=8)
0.9
20.5
4.4
(c) ResNet-152 (200 epochs) on 8x8 GPUs
(d) ViT-B (300 epochs) on 8x8 GPUs
Method
Comm. (h) Total (h)
Ratio (%)
Method
Comm. (h) Total (h)
Ratio (%)
Parallel SGD
1.3
5.7
22.8
Parallel Adam
3.7
8.6
43.0
QSR (Hbaee 2)
0.6
5.0
12.0
QSR (Hae 1)
0.6
5,5
10.9
QSR bare 1)
0.3
4.7
6.4
QSR hasse 8)
0.4
5.3
7.5
Local SGD (H==2
0.7
5.1
13.7
Local AdamW (H==4
0.9
5.8
15.5
Local SGD (H==)
0.3
4.8
6.3
Local AdamW (H=8)
0.5
5.3
9.4
increases the communication overhead for both models. Notably, on 8x8 GPUs, communication accounts for almost half of the total training time for ViT-B. Since communication makes up larger portion of the total time for ViT-B compared to ResNet-152, the speedup from QSR is more significant VIT-B: the time is cut from 26.7 20.2 hours 2x8 GPUs, and 8.6 to 5.5 hours on 8x8 GPUs As discussed in Section 4.1, compared to the constant period local gradient method, QSR further reduces the communication cost by increasing the synchronization period in the late phase. For example, applying QSR to Local AdamW with H = further reduces the time by I hour for ViT training on 2x8 GPUs.
Discussion on the choice of H As elaborated in Section 2, H indicates the minimum syn- chronization period and should be determined based on the communication overhead. For ResNet- 152, given that communication only accounts for 3.3 out of 20.7 hours on 2x8 GPUs and 1.3 out of 5.7 hours on 8x8 GPUs, setting Hasse as inconsequential amount. By contrast, the communication overhead for VIT-B is more prominent, motivating us to consider larger values of H such as 4 and 8. As shown in Tables and 2, Hbase introduces tradeoff between communication efficiency and final test accuracy. For instance, when training ResNet- 152 with batch size 16384, one can either choose Hbsse ble test accuracy parallel SGD, or Hbsse = 4 to further halve the communication volume at the expense of a 0.6% drop in test accuracy. One probable explanation for this accuracy drop for larger Haase can be worse optimization in the early training phase, where the learning rate is large.
suffices to reduce the communication time to an
to achieve compara-
5 DISCUSSIONS AND FUTURE DIRECTIONS
This paper primarily focuses on relatively large models trained with long horizons, and proposes the Quadratic Synchronization Rule (QSR). As validated by our experiments, QSR effectively im- proves test accuracy and communication efficiency simultaneously for training large vision models (ResNet-152 152 and ViT-B) with quite a few hundred epochs. However, on the downside, for smaller models trained with shorter horizons, QSR may not consistently deliver noticeable generalization improvements (see Table 5). Nonetheless, training in this regime is not costly, either, making it less of critical concern. Another limitation of our work that the effectiveness of QSR relies on the im- plicit regularization effects of noise, but regularization techniques become less important in bridging the gap between the training and population loss (Vyas et al., 2023) in pertaining large model with unsupervised learning, where the training is done on massive data with only few epochs. Still, cer- tain implicit/explicit regularization effects have been found to be effective in improving downstream performance despite the same pertaining loss (Liu et al., 2023; Panigrahi et al.. 2024). We leave it to future work to explore and design communication-efin -efitiiioneffiinnttt methods for unsupervised learning, particularly language model pretraining, that improve models' transferability downstream tasks.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENT AND DISCLOSURE OF FUNDING
The work of Xinran Gu and Longbo Huang is supported by the Technology and Innovation Major Project of the Ministry of Science and Technology of China under Grant 2020AAA0108400 and 2020AAA0108403. The work of Kaifeng Lyu and Sanjeev Arora is partly supported by NSF and ONR.
REFERENCES
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix fac- torization. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché Buc, E. Fox, and R. Gamett (eds.), Advances in Neural Information Processing Systems 32, pp. 741 1-4422. Curran Asso- ciates, Inc., 2019.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (cds.), Proceedings of the 39th International Conference on Machine Learning. volume 162 of Proceedings of Machine Learning Research, pp. 948-1024. PMLR, 17-23 Jul 2022.
Debraj Basu, Deepesh Data, Can Karakus, and Suhas Diggavi. Qsparse-local-SGiD Distributed SGD with quantization, sparsification and local computations. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-But Baché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Lucas Beyer, Xiaohua Zhai, and Alexander Kolesnikov. Better plain vit baselines for imagenet- -kk arXiv preprint arXiv:2205.0f 5800 2022.
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by ornstein-uhlenbeck like process. In Jacob Abernethy and Shivani Agarwal (eds.), Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pp. 483-513. PMLR, 09-12 Jul 2020.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askeli, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. CORR, abs/2005. 14165, 2020. URL https: / wwwaa.iivvv..... www.i.og.......o............oooo 20t.//200...1..
Kai Chen and Qiang Huo. Scalable training of deep learning machines by incremental block train- ing with intra-block parallel optimization and blockwise model-update filtering. In 2016 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pp. 5880-5884, 2016. doi: 10. (WW0//CAASSP221..77722222.
Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre- preitraining or strong data augmentations. In International Conference on Learning Rep- resentations, 2021.
Lénaïc Chizat and Francis Bach. Implicit bias gradient descent for wide two-layer neural networks trained with the logistic loss. In Jacob Abernethy and Shivani Agarwal (eds.), Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pp. 1305-1338. PMLR, (9- Jul 2020.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar, Gradient descent on neural networks typically occurs at the edge of stability, In International Conference on Learning Representations, 2020.
Aditya Cowsik, Tankut Can, and Paolo Glorioso. Flatter, faster: scaling momentum for optimal speedup of sgd. arXiv preprint arXiv:22100 16400, 2022.
10

---

Published as a conference paper at ICLR 2024
Alex Damian, Tengyu Ma, and Jason D Lee. Label noise SGD provably prefers flat global mini- mizers. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 27449-27444 Curran As- sociates, Inc., 2021.
Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization:: The implicit bias of gradient descent at the edge of stability, In The Eleventh International Conference on Learning Represen- lations, 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image worth 16x 16 words: Transformers for image recognition scale. In International Conference on Learning Representations, 2021.
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. International conference on machine learning, pp. 1309- 1318. PMLR. 2018.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. In International Conference on Learning Represen- tations, 2021a.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. In International Conference Leurning Represen- tations, 2021b.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode con- nectivity and the lottery ticket hypothesis. In International Conference on Machine Learning, pp. PMLR, 2020.
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. Advances in neural information pro- cessing systems, 31, 2018.
Rong Ge, Yunwei Ren, Xiang Wang, and Mo Zhou Understanding deflation process in over- parametrized tensor decomposition. Advances in Neural Information Processing Systems, 34, 2021.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training imagenet in hour. arXiv preprint arXiv: 706.02677, 2017.
Diego Granziol, Stefan Zohren, and Stephen Roberts. Learning rates as a function of batch size: A random matrix theory approach to neural network training. The Journal of Machine Learning Research, 23(1):7795-788 7859, 2022.
Xinran Gu, Kaifeng Lyu. Longbo Huang, and Sanjeev Arora. Why (and when) does local SGD gen- eralize better than SGD? In The Eleventh International Conference on Learning Representations, 2023. URL https::////o wpw.reveeeeeeeeeeeeeeeeeee fot//forrmiddeeveuui 6Drl.
Vipul Gupta, Santiago Akle Serrano, and Dennis DeCoste. Stochastic weight averaging in parallel: Largh- training that generalizes well. In International Conference on Learning Representa- tions, 2020.
Farzin Haddadpour, Mohammad Mahdi Kamani, Mehrdad Mahdavi, and Viveck Cadambe. Local SGD with periodic averaging: Tighter analysis and adaptive synchronization. Advances in Neural Information Processing Systems, 32, 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1):1-42, 1997.

---

Published as a conference paper at ICLR 2024
Wenqing Hu, Chris Junchi Li, Lei Li, and Jian-Guo Liu. On the diffusion approximation of noncon- vex stochastic gradient descent. arXiv preprint arXiv:1705.07550 2017.
Gao Huang. Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinherger. Densely connected convolutional networks. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 4700-4708, 2017.
Hikaru Ibayashi and Masaaki Imaizumi. Exponential escape efficiency of SGD from sharp minima in non-stationary regime. arXiv preprint arXiv:2111.00 2021.
P Izmailov, AG Wilson, D Podoprikhin, D Vetrov, and T Garipov. Averaging weights leads to wider optima and better generalization. In 34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018, pp. 876-885, 2018.
Peter Izsak, Moshe Berchansky, and Omer Levy. How to train bert with an academic budget. arXiv preprint drX1v:2104.07755, 2021.
Stanislaw Jastrzębski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey, Three factors influencing minima in SGD, arXiv preprint 17X1v:1711.04623, 2017.
Ziwei and Matus Telgarsky. Directional convergence and alignment in deep learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 17176-17186. Curran Associates, Inc., 2020.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where to find them. In International Conference on Learning Representations, 2020.
Jikai Jin, Zhiyuan Li, Kaifeng Lyu, Simon Shaolei Du, and Jason D. Lee. Understanding incremental learning of gradient descent: A fine-graiee ane-ggainedd analysis of matrix sensing. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Pro- ceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 15200-15253 PMLR, 23-29 Jul 2023.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Ad- vances and open problems in federated learning. Foundations and Trends@ in Machine Learning, 14(1-2):1-210, 2021.
Michael Kamp, Mario Boley, Daniel Keren, Assaf Schuster, and Izchak Sharfman. Communication- efficient distributed online prediction by dynamic model synchronization. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2014, Nancy, France, September 15-19, 2014. Proceedings, Part 14, pp. 633.639. Springer, 2014.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning, pp. 5132-5143, PMLR, 2020.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On larce-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017a.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference Learning Representations, 2017b.
Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. Tighter theory for local SGD iden- tical and heterogeneous data. In International Conference on Artificial Intelligence and Statistics, pp. 4519-4529. PMLR, 2020.
12

---

Published as a conference paper at ICLR 2024
Bobby Kleinberg, Yuanzhi Li, and Yang Yuan. An alternative view: When does SGD escape lo- cal minima? In Jennifer Dy and Andreas Krause (eds.). Proceedings of the 35th International Conference Machine Learning. volume 80 of Proceedings of Machine Learning Research, pp- 2698-2707. PMLR, 10-15 Jul 2018.
Jakub Konečný, H Brendan McMahan, Felix X Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. arXiv preprint (ttp::11100054222 2016.
Alex Krizhevsky. One weird trick for parallelizing convolutional neural networks. arXiv preprint arXiv: Ieoe..44499 2014.
Guillaume Leclere, Andrew Ilyas, Logan Engstrom, Sung Min Park, Hadi Salman, and Aleksander Madry. ffev. https: /tt::///ittti gttt....... com/ wwtffliffffffooooooot/eeooo 2022.
Guillaume Leclere, Andrew Ilyas, Logan Engstrom, Sung Min Park, Hadi Salman, and Alek- sander Mądry. Ffev: Accelerating training by removing data bottlenecks. In Proceedings of the JEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 12011-12020, June 2023.
Conglong Li, Ammar Ahmad Awan, Hanlin Tang, Samyam Rajbhandari, and Yuxiong He. 1-bit lamb: communication efficient large-scale large-batch training with lamb's convergence speed. In 2022 IEEE 29th International Conference on High Performance Computing. Data, and Analytics (HiPC), pp. 272-281. IEEE, 2022.
Mengtian Li, Ersin Yumer, and Deva Ramanan. Budgeted training: Rethinking deep neural network training under resource constraints. In International Conference on Learning Representations, 2020a.
Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and dynamics of stochastic gradient algorithms Mathematical foundations. Journal of Machine Learning Research, 20(40): 2019a.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and sys- tems, 2:429-451 2020b.
Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and 7hihua Zhang. On the convergence of fedavg on non-iid data. In International Conference on Learning Representations, 2019b.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over- over-parameterinze matrix sensing and neural networks with quadratic activations. In Sébastien Bubeck, Vianney Perchet, and Philippe Rigollet (eds.), Proceedings of the 31 Conference On Learning Theory, volume 75 of Proceedings of Machine Learning Research, pp. 2-47. PMLR, 06-09 Jul 2018.
Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. International Conference on Learning Representations, 2021a.
Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling SGD with stochastic differential equations (sdes). Advances in Neural Information Processing Systems, 34:12712- 12725, 2021b.
Zhiyuan Li. Tianhao Wang. and Sanjeev Arora. What happens after SGD reaches zero loss?-a mathematical framework. In International Conference on Learning Representations, 2021c.
Tao Lin, Sebastian U. Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't use large mini- bani-batches, use Local SGD. In International Conference on Learning Representations, 2020.
Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better down- stream: Implicit bias matters for language models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning. volume 202 of Proceedings of Ma- chine Learning Research, pp. 22118-22224,2.111111111111111148 PMLR, 23-29 Jul 2023.
13

---

Published as a conference paper at ICLR 2024
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference computer vision. pp. 10012-10022, 2021.
Zhuang Liu, Hanzi Mao, Chao- Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF conference on computer vision and paltern recognition, pp. 11976- 1986, 2022.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations, 2020.
Kaifeng Lyu, Zhiyuan Li, Runzhe Wang. and Sanjeev Arora. Gradient descent on two-layer nets:
Margin maximization and simplicity bias. Advances in Neural Information Processing Systems, 34, 2021.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normal- ization layers: Sharpness reduction, 2022.
Chao Ma and Lexing Ying. On linear stability of SGD and imput-smoothness of neural networks. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 16805-16817. Curran Associates, Inc., 2021.
Chao Ma, Daniel Kunin, Lei Wu, and Lexing Ying. Beyond the quadratic approximation: The multiscale structure of neural network loss landscapes. Journal of Machine Learning, 1(3):247- 267, 2022. ISSN 2790-2048.
Wei-Chiu Ma, Shenlong Wang, Rui Hu, Yuwen Xiong, and Raquel Urtasun. Deep rigid instance scene flow. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recog- nition, pp. 3614-3622, 2019.
Sadhika Malladi. Kaifeng Lyu, Abhishek Panigrahi. and Sanjeev Arora. On the SDEs and scaling rules for adaptive gradient algorithms. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Gideon Mann, Ryan T. McDonald, Mehryar Mohri, Nathan Silberman, and Dan Walker. Efficient large-scale distributed training of conditional maximum entropy models. In Advances in Neural Information Processing Systems 22, pp. 1231-1239, 2009.
Brendan McMahan, Eider Moore, Daniel Ramage. Seth Hampson, and Blaise Aguera y Arcas. Communication-eficeee learning of deep networks from decentralized data. In Artificial intelli- gence and statisticy, pp. 1273-1282. PMLR, 2017.
Mor Shpigel Nacson, Suriya Gunasekar, Jason Lee, Nathan Srebro, and Daniel Soudry. Lexico- graphic and depth-sensitee margins in homogeneous and non- non-homogeneous deep models. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp.
Long Beach, California, USA, 09-15 15 Jun 2019. PMLR.
Giorgi Nadiradze, Amirmojtaba Sabour, Peter Davies, Shigang Li, and Dan Alistarh. Asynchronous decentralized sgd with quantized and local updates. Advances in Neural Information Processing Systems,
2021.
Behnam Neyshabur, Srinadh Bhojanapalli, David Mcallester, and Nati Srebro. Exploring gener- alization in deep learning. In Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Jose Javier Gonzalez Ortiz, Jonathan Frankle, Mike Rabbat, Ari Morcos, and Nicolas Ballas. Trade- offs of Local SGD scale: An empirical study. arXiv preprint 10.00:2110.08/11 2021.
Abhishek Panigrahi, Nikunj Saunshi, Kaifeng Lyu, Sobhan Miryoosefi, Sashank Reddi, Satyen Kale, and Sanjiv Kumar. Efficient stagewise pretraining via progressive subnetworks. arXiv preprint arXiv:2402.0591. 2024.
14

---

Published as a conference paper at ICLR 2024
Daniel Povey, Xiaohui Zhang. and Sanjeev Khudanpur. Parallel training of dnns with natural gradi- ent and parameter averaging. arXiv preprint arXiv:1410.7455, 2014.
Noam Razin and Nadav Cohen. Implicit regularization in deep learning may not be explainable by norms. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 21174-21187. Curran Associates, Inc., 2020.
Noam Razin, Asaf Maman, and Nadav Cohen. Implicit regularization in hierarchical tensor fac- torization and deep convolutional neural networks. In Kamalika Chaudhuri, Stefanic Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th Inter- national Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 18422. PMLR, 17-23 Jul 2022.
Sashank J Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Representations, 2020.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma. Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJC)), 1 115(3)2211225) 252 2015. doi: 10.1007/s1 1200077/11223----08116-66-
Christopher J. Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E. Dahl. Measuring the effects of data parallelism on neural network training. Journal of Machine Learning Research, 20( 12):1-49, 2019.
Shuheng Shen, Yifei Cheng, Jingchang Liu, and Linli Xu. Stl-sgd: Speeding up local sgd with stage- wise communication period. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 9576-9584, 2021.
Samuel Smith, Erich Elsen, and Soham De. On the generalization benefit of noise in stochastic gradient descent. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 9058-9067. PMLR, 13-18 Jul 2020.
Samuel L Smith, Benoit Dherin, David Barrett, and Soham De. On the origin of implicit regular- ization stochastic gradient descent. International Conference on Learning Representations, 2021.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The im- plicit bias of gradient descent on separable data, Journal of Machine Learning Research, 19(70): 1-57, 2018a.
Daniel Soudry, Elad Hoffer, and Nathan Srebro. The implicit bias of gradient descent on separable data. In International Conference on Learning Representations, 201 8b.
Sebastian U Stich. Local SGD converges fast and communicates little. In International Conference on Learning Representations, 2018.
Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruc- tion. Advances in Neural Information Processing Systems, 34, 2021.
Hang Su and Haoyu Chen. Experiments on parallel training of deep neural network using model averaging. arXiv preprint 1.X.vv1507.01230 2015.
Hanlin Tang, Shaoduo Gan, Ammar Ahmad Awan, Samyam Rajbhandari. Conglong Li, Xiangru Lian, Ji Liu, Ce Zhang, and Yuxiong He. 1-bit adam: Communication efficient large-scale training with adam's 's convergence speed. In Proceedings of the 38th International Conference on Machine Learning, 2021.
15

---

Published as a conference paper at ICLR 2024
Nikhil Vyas, Depen Morwani, Rosie Zhao, Gal Kaplun, Sham Kakade, and Boaz Barak. Beyond im- plicit bias: The insignificance of SGD noise in online learning. arXiv preprint arXiv:2306.0883 2023.
Jianyu Wang and Gauri Joshi. Adaptive communication strategies to achieve the best etror-runtime trade-off in local-update SGD. Proceedings of Machine Learning and Systems, 1:212-229, 2019.
Jianyu Wang and Gauri Joshi. Cooperative SGD: unified framework for the design and analysis of local-update SGD algorithms. Journal of Machine Learning Research, 22(213):1-50, 2021.
Jianyu Wang, Vinayak Tantia, Nicolas Ballas, and Michael Rabbat. Slowmo: Improving communication-effee distributed SGD with slow momentum. In International Conference on Learning Representutions, 2019.
Runzhe Wang, Sadhika Malladi, Tianhao Wang, Kaifeng Lyu, and Zhiyuan Li. The marginal value of momentum for small learning rate sgd. arXiv preprint urXiv:2307.15166 2023.
Blake Woodworth, Kumar Kshitij Patel, Sebastian Stich, Zhen Dai, Brian Bullins, Brendan Mcma- han, Ohad Shamir, and Nathan Srebro. Is local sgd better than minibatch sgd? In International Conference on Machine Learning, pp. 10334-10343, PMLR, 2020.
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lppes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learning, pp. 2022.
PMLR,
Mitchell Wortsman, Suchin Gururangan, Shen Li, Ali Farhadi, Ludwig Schmidt, Michael Rabbat, and Ari S. Morcos. lo-fi: distributed fine-tuning without communication. Transactions on Ma- chine Learning Research, 2023. ISSN 2835-8856.
Lei Wu, Chao Ma, and Weinan E. How sgd selects the global minima over-parameterized learning: A dynamical stability perspective. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. In International Conference on Learning Representations, 2021.
Hang Xu, Wenxuan Zhang, Jiawei Fei, Yuzhe Wu, Tingwen Xie, Jun Huang, Yuchen Xie, Mohamed Elhoseiny, and Panos Kalnis. SLAMB: Accelerated large batch training with sparse communica- tion. In Proceedings of the 40th International Conference on Machine Learning, 2023.
Yang You, Zhao Zhang, Cho-Jui Hsieh, James Demmel, and Kurt Keutzer. Imagenet training in minutes. In Proceedings of the 47th International Conference on Parallel Processing, pp. 1-10, 2018.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learn- ing: Training BERT in 76 minutes. In International Conference on Learning Represenations, 2020.
Hao Yu, Sen Yang, and Shenghuo Zhu. Parallel restarted SGD with faster convergence and less communication: Demystifying why model averaging works for deep learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 5693-5700, 2019.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Rep- resentations, 2017.
16

---

Published as a conference paper at ICLR 2024
Xiaohui Zhang, Jan Trmal, Daniel Povey, and Sanjeev Khudanpur. Improving deep neural network acoustic models using generalized maxout networks. 2014 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 215-219, 2014. doi: 10.1109/CASSS..
2014.6853589.
Fan Zhou and Guojing Cong. On the convergence properties of k-step averaging stochastic gra- dient descent algorithm for nonconvex optimization. In Proceedings of the Twenty-Seventh Inter- national Joint Conference on Artificial Intelligence, LJCAI-18, pp. 3219-3227. International Joint Conferences on Artificial Intelligence Organization, 2018. doi: 10. 2i..2496/jeaai.20188// URL https: /tw.///o..////// /ergg/10 1 2.3963/////. joai 2018/ 447.
Tongtian Zhu, Fengxiang He, Kaixuan Chen, Mingli Song, and Dacheng Tao. Decentralized SGD and average-direction SAM are asymptotically equivalent. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 43005-43036. PMLR, 23-29 Jul 2023.
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In Ka- malika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Con- ference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp- 7654-7663. PMLR, 09-15 Jun 2019.
Martin Zinkevich, Markus Weimer, Lihong Li, and Alex Smola. Parallelized stochastic gradient descent. In J. Lafferty, C. Williams, J. Taylor, R. Zemel, and A. Culotta (eds.), Advances in Neural Infommtion Processing Systems, volume 23. Curran Associates, Inc., 2010.
17

---

Published as a conference paper at ICLR 2024
CONTENTS
Introduction
Our Method: Quadratic Synchronization Rule
3 Theoretical Motivations of Quadratic Synchronization Rule
4 Experiments
4.1 QSR Improves Generalization
4.2 QSR Reduces Wall-clock Time
5 Discussions and Future Directions
A Additional Related Works
19
B PseudoCode
21
Experimental Details
22
C.I Training Details for ResNet-152
22
C.2 Training Details For ViT-B
D Additional Experiments on ResNet-50
24
E Supplementary Materials for Section 3
25
E.1 Missing Definitions and Assumptions
25
E.2 Proof for Theorem 3.1
26
Details for Communication Time Measurement
32
Discussion More Aggressive Scalings
33
Comparison with Local SGD/AdamW + SWAP
35
18

---

Published as a conference paper at ICLR 2024
A ADDITIONAL RELATED WORKS
Advances in local gradient methods. Local gradient methods are a class of communication- efficient algorithms for distributed training. In this approach, workers update their models locally and average the model parameters every time they finish H steps of updates. Dating back to Mann et al. (2009) and Zinkevich et al. (2010), local gradient methods have been widely used to improve communication efficiency in both datacenter distributed training Zhang et al. (2014); Povey et al. (2014); Su & Chen (2015); Chen & Huo (2016) and Federated Learning (Kairouz et al., 2021; McMahan et al., 2017; Li et al., 2019b; Konečný et al., 2016). Many variants have been proposed to facilitate the convergence speed. Examples include using control variates (Karimireddy et al., 2020), adding proximal terms to local loss functions (Li et al., 2020b), and applying adaptivity on top of each communication round ( (Wang et al., 2019; Reddi et al., 2020). Local gradient methods can also be readily combined with orthogonal approaches like communication compression (Basu al., 2019) and asynchronous updates (Nadiradze et al., 2021) for further communication cost reduction.
Optimization perspectives on selecting H. Extensive prior research has been devoted op- timizing the selection of the synchronization period H from an optimization perspective. The conventional approach sets H as constant throughout training. In this setup, series of studies (e.ge. et al. (2020); Stich (2018); Haddadpour et al. (2019); Yu et al. (2019)) established con- vergence bounds for the training loss, which typically degrade as H gets larger. leading to a trade-off between communication efficiency and model accuracy. Drawing upon these theoretical results, H should be set as the smallest value that reduces the communication cost to an acceptable level to minimize the negative impact optimization. To better trade-off between optimization and gen- cralization, researchers introduced various adaptive communication strategies. Kamp et al. (2014) designed synchronization protocol controlled by the variance in model parameters. Haddadpour et al. (2019) suggested linearly increasing H as the iteration goes on. Shen et al. (2021) introduced stagewise communication scheme that halves the learning rate 77 while doubles H every time the training has finished a predefined stage. Aimed at optimizing the convergence of training loss with respect to wall-clock time, Wang & Joshi (2019) proposed a strategy that starts with infrequent communication and gradually decreases H as training progresses. Nonetheless, the effectiveness of these adaptive communication strategies has only been empirically validated on linear models or small-scale datasets like (IFAR-10/100
Generalization perspectives on selecting H. While larger H usually hurts optimization, it can sometimes improve generalization. Apart from Lin et al. (2020) that has been discussed in detail in Section 1, similar observations have been reported by Gupta et al. (2020) and Wortsman et al. (2023). Specifically, Gupta et al. (2020) introduced the Stochastic Weight Averaging in Parallel (SWAP) algorithm, which runs parallel SGD until a target training accuracy, then lets workers per- form local updates with a final model averaging. Their empirical results validate SWAP's superior generalization performance over parallel SGD. When using LAMB (You et al., 2020) as the op- timizer, Wortsman et al. (2023) find that complete local fine-tuning, followed by single model averaging in the end (equivalent to setting H as the total number of iterations). outperforms the standard parallel LAMB in test accuracy under distribution shifts. Another relevant method is the "model soup" (Wortsman et al., 2022), which averages multiple models fine-tuned with different hyperparameters and turns out beat the single model in test accuracy. Our paper focuses on designing the synchronization scheme best for generalization.
Implicit bias of optimizers. The success of deep learning lies in its remarkable ability to general- ize to unseen data, though it possesses the capacity to fit randomly labeled data (Zhang et al., 2017). A significant contributing factor to this success the implicit bias inherent in popular optimizers like Gradient Descent (GD) and Stochastic Gradient Descent (SGD). Specifically, these optimizers favor minima that exhibit good generalization, without explicitly encoding such bias into the train- ing loss. A lot of studies have been devoted to characterizing this implicit bias, some through the lens of margin maximization (Soudry et al., 20;8b;;; Lyu & Li, 2020; Ji & Telgarsky, 2020; Chizat & Bach, 2020; Nacson et al., 2019), and some others focus on the simplicity bias from small ini- tialization (Li et al., 2018; Razin & Cohen, 2020; Arora et al., 2019 Li et al., 2021a; Lyu et al., 2021: Razin et al., 2022; Stöger & Soltanolkotabi, 2021; Ge et al., 2021; Jin al., 2023). The line of work most closely related to our paper interprets the implicit bias via sharpness reduction. The connection between flatter minima and better generalization is a commonly held belief that has
19

---

Published as a conference paper at ICLR 2024
been investigated both theoretically (Hochreiter & Schmidhuber, 1997; Neyshabur et al., 2017) and empirically (Keskar et al., 2017a; Jiang et al., 2020). Drawing on this insight, Foret et al. (2021a) introduced SAM optimizer, which delivers superior generalization performance by explicitly penal- izing sharpness. Recent theoretical studies (Arora et al., 2022; Lyu et al., 2022; Damian et al., 2023; Ma et al., 2022) elucidate that GD inherently biases towards flatter regions the loss landscape. Specifically, under some regularity conditions, they show that GD will eventually enter the "Edge of Stability"(Cohne
et al., 2020), where the maximum eigenvalue of the loss Hessian stays around 2/leannnng rate, and then constantly moves towards flatter minima. Going beyond GD, another line of work studies how gradient noise SGD helps reduce sharpness. Wu al. (2018); Hu et al. (2017); Ma & Ying (2021) showed that gradient noise can cause training instability around sharp minima, and hence, the iterate can only settle around flat minima. Kleinberg et al. (2018); Zhu et al. (2019); Xie et al. (2021); Ibayashi & Imaizumi (2021) analyzed the escaping behavior of SGD from sharp minima. Motivated by recent empirical observations that low-loss solutions on the loss landscape are path-connected (Garipov et al., 2018; Draxler et al., 2018: Frankle et al., 2020) rather than isolated, Blanc et al. (2020); Damian et al. (2021); Li et al. (2021c) assume the existence of minimizer manifold and show that gradient noise provably drives the iterate towards flatter minima on this manifold. Cowsik et al. (2022); Wang et al. (2023) discuss how momentum preserves or strengthens this effect. Also through the lens of sharpness reduction, the recent work by Gu et al. (2023) explains the generalization benefit of Local SGD, as discussed in Section 3. Zhu et al. (2023) elucidate that similar implicit bias also manifests decentralized training by making connections to certain variants of SAM.
20

---

Published as a conference paper at ICLR 2024
B PSEUDOCODE
We present the pseudocodes for standard data parallel methods and local gradient methods below.
Algorithm 1: Parallel OPT: Data Parallel Methods on K Workers
Input: loss function ((0:5) E), initial parameter (()) Hyperparameters: total number of iterations
Hyperparameters: learning rate schedule T/t ,T,.,,TT) local batch size Bioc
t 0; for t = 0 ...... R do
// initialize the global iteration number
for each worker l: do in parallel (4)
(5)
(5k.1.1
Btn Sample ()
sample local batch
(1)
Dar ( :
c())
9k
Bice
// computing the local gradient All-Reduce aggregation Local gradlents
end
k
9 (1)
10
K kk=1
11
(c++1) OPT(A Te: g (1);
1/ update the model with optimizer OPT
12 end
Algorithm 2: Local OPT: Local Gradient Methods on K Workers
Input: loss function ((0: E), initial parameter (0) Hyperparameters: total number of rounds R
Hyperparameters: learning rate schedule T), f E ( ,T), local batch size Bioc
++0 0; for s s= 0, R do
initialize the gicbal iteration
((a) t Gst
// get synchronization peried for the curzent round
for each worker k do in parallel
(10)
k.0 + g())
1/ maintain local copy of the globsl model parameter
for h 0, ..... (+)) do
10
h5
Sk.. Bloe t Sample
1/ sample local batea 1/ computing the local gradient k,h Meth, phh) // update the local model with optimizer OPT
(8)
Blse EBe
Biee Ve(O R.N. ((th.))
11
(a))
k,h+1 f- OPT(( (s)
end
14
end
15
( 1)
K
Ex=1
(s))
kHH
1/ All Beduce aggregation of local model parameters // update the global iteration number
t- + H()
16
17 end
Sampling local batches. In Algorithms and 2, Sample( returns a local batch for each worker. In our experiments, local batches are sampled without replacement at each epoch, which standard for distributed training (Goyal et al., 2017; Lin et al., 2020; Ortiz et al., 2021). More specifically, the beginning of each epoch, all the workers use the same random seed draw shared random permutation of train data points, and partition the data points evenly among the K workers. Then at each local step of each worker, Sample() sequentially takes samples from its own partition. Once there are too few remaining samples to form complete batch, new permutation is sampled and a new epoch starts. For our theoretical analysis, following Gu et al. (2023), we assume Samplel] takes samples with replacement, i.e. the K workers are taking i.i.d. samples from the globally shared dataast/ddsstriuution. See Appendix B in Gu et al. (2023) for pseudocodes of sampling with and without replacement.
21

---

Published as a conference paper at ICLR 2024
Setting synchronization periods. In Algorithm 2, Get:((s) is function that returns the syn- chronization period H(s) for the current round. Conventionally, H is chosen as a fixed value, so Ge):((s) always returns a constant. In this paper, we study how H*s) should change as training goes on, e.g., in QSR, GetH(s) works as specified in Section
C EXPERIMENTAL DETAILS
This section lists the additional experimental details omitted in the main text.
0.0075
cosine decay step decay
Learning Rate
Software and platform. We use Pytorch Dis- tributed with NCCL backend to support multinode distributed training and use FFCV (Leclere et al., 2022) to accelerate data loading of ImageNet.
0.0050
linear decay
0.0025
Sampling scheme. We employ the "sampling without replacement" scheme, as described in Ap- pendix B.
0.0000
100
200
300
Epochs
C.I TRAINING DETAILS FOR RESNET-151
Figure 4:
A visualization of the learning rate schedules we investigate.
We generally follow the recipe in Foret et al. (2021b) to train ResNet-152. Specifically, we set the momentum as 0.9 and the weight decay À as 0.0001. For data augmentation, we employ random resized crop and random horizontal flip. We additionally use label smoothing 0.1. We adopt local batch size Bic 256 through 8 gradient accumulations. Therefore, the batch size for BatchNorm 32. This choice stems from our observation that a smaller batch size for BatchNorm enhances the test accuracy of parallel SGD, Since the BatchNorm statistics on each worker are estimated on the local model parameter, we pass 100 batches, each of size 32, to estimate the BatchNorm statistics on the global parameter before evaluation.
Training details for batch size 4096. We search the optimal peak learning rate Maxx of the cosine learning rate schedule among (0.4, 0.8, 1.6) for all baseline algorithms, i.e., parallel SGD and Local SGD with constant synchronization period H = 2 and H = 4. The learning rate yielding the highest final test accuracy is selected. We find that mmma 0.8 is optimal for all the baseline algorithms. For QSR with Hhase = 2 and Hasse 4, we directly set mmax = 0.8. We search a and 4 respectively.
among {0.2, 0.25, 0.3), and choose a 0.2 and 0.25 for QSR with Whsse Regarding other communication strategies in Figure 1(a), we set the switching point at epoch 100 and employ H 8 for Post-local SGD. For H 3/n, we search 8 among (0.6, 0.8. 1, 1.2), finally selecting 3 = 1.
Training details for batch size 16384. The hyperparameter tuning procedure for B 16384 is similar to that of B 4096. We search 1mmax among (0.8, 1.6,3.2) for all baseline algorithms, including SGD and Local SGD with constant synchronization period Hosee
and Hasse 4. We find that 7mmx 3.2 yields the highest final test accuracy for all of them. However, for QSR, we find that peak learning rate 7max = 3.2 is excessively large, causing the dynamic scheduling to be triggered too late in the training process. This late triggering leaves insufficient training time for the training to fully leverage the generalization benefits introduced by local steps. Consequently, we and 4. We search a among {0.2, 0.25, 0.3). and choose = 0.2 for both QSR with H Basse = 2 and 4.
set 7/max 1.6 for QSR with H
Training details for the step decay scheduler. In our experiments with step decay, we employ batch size of 4096. Given that our step decay scheduler is derived from the cosine decay, we only need to specify the weight decay X, and peak learning rate Tmmax. These are identically to the values used in our cosine decay experiments. For QSR, we search the growth coefficient a among {0.2, 0.3) and choose 0.2 for both Hbase 2 and 4.
Training details for experiments in Appendix H. For Local SGD + SWAP experiments in Ap- pendix H. we use the cosine learning rate schedule with peak learning rate Mmmax 0.8. We start with Local SGD with a constant synchronization period H = 4 and explore the switching point to from (175, 180, 185, 190).
22

---

Published as a conference paper at ICLR 2024
C.2 TRAINING DETAILS FOR VIT-B
For training VIT-B, we primarily follow the 300- epoch recipe proposed by Beyer et al. (2022). Specifically, we replace the [cls] token of the orig- inal VIT token with global average pooling and use fixed 2D sin-cos position rather than learned po- sitional embeddings. Our implementation of the model architecture follows the high-starred reposi- tory by Phil Wang. Apart from random resized crop and random horizontal flip, we employ Ran- dAugment with parameters (2, 10) and MixUp with a coefficient of 0.2 for data augmentation. Different from Beyer al. (2022), we use larger batch size (B = 4096 or 16384 as opposed to their 1024) and use AdamW instead of Adam.
10*
Schedale and Communcarion Volume
H (log scale)
10a
Constant Period, 25%
QSR, 10.4%
10a
10
100
200
300
Epochs
Figure 5: A visualization of the H schedule for Local AdamW with a constant synchronization period H = 4 and with QSR Hbase 4, o 0.0175. The corresponding lcarning rate schedule is cosine decay with peak learning rate of 0.008. Adopting QSR improves the top-1 validation curacy Local Adam on ViT-B from 79.32% to 80.95 80.98%..
As for gradient clipping, we set it as 1 for standard AdamW following Beyer al. (2022); Dosovitskiy et al. (2021) and Chen et al. (2021). However, for Local AdamW, the smaller batch size locally leads to larger gradient noise and, hence larger gradient norm for local updates. This calls for an increase in the gradient clipping threshold. We find that the training process remains stable even when we remove gradient clipping (cquivalent to setting the clipping threshold to +x) for most of the hyperparameter configurations we tested. For ease of tuning, we choose to turn off gradient clipping for Local Adam unless otherwise stated.
Training details for batch size 4096. We use 10k iterations for learning rate warmup following (Beyer et al., 2022: Dosovitskiy et al., 2021; Chen et al., 2021). For parallel AdamW and Local AdamW (H- 4), explore combinations of Tmax and weight decay À from the grid (0.05, 0.1) x (0.004, 0.008, 0.016). To optimize the final test accuracy, we select 7mma = 0.008, X = 0.1 for parallel Adam W and 7max 0.008,
0.05 for Local AdamW (H 4). For Local AdamW (H1 = 8). keeping X= = 0.05, we conduct a grid search for 7max among {0.004, 0.008, 0.016] and choose 7max 0.008 For QSR with Hasse and 8, we directly use Mmax = 0.008 and X 0.05. To optimize we search among (0.015, 0.0175, 0.02] and find a = 0.0175 works best for both QSR with Hhase = 4 and 8. Regarding the communication strategy of H = 3/11 in Figure 1(b), we explore 3 among (0.025, 0.03, 0.035, 0.04), settling on 3 = 0.03. In Figure 5, we also visualize the H schedule for Local Adam W with a constant synchronization period and with QSR.
Training details for batch size 16384. To keep the same portion of the total budget for learn- ing rate warmup as B = 4096, we set the warmup iterations to 2.5k. We set À as 0.1 and 0.05 for parallel AdamW and Local AdamW, respectively. We search for the optimal Thma among {0.004, 0.008, 0.016} and select 7/nax = 0.004 for parallel Adam W, 7max = 0.016 for Local AdamW with H 4 and 8. We adopt the same À and Mmax as Local AdamW for QSR. For QSR with Haase = 4, we search for the optimal a among {0.015, 0.01 75, 0.02) and choose a = 0.0175. For QSR with Hasne = 8, we search for the optimal a among (0.01, 0.0175), finally picking
0.01.
Training details for linear and step decay schedulers. For both step and linear decay schedulers, we employ a batch size of 4096. For the step decay scheduler, the peak learning rate 7/max and weight decay À are set identically to the values used in our cosine decay experiments. We search the growth coefficient a for QSR among (0.015,0.0155} and choose 0.015 for both Hbsse 4 and 8. For linear decay, we use the same weight decay as our cosine decay experiments. We explore mmax values from {0.004, 0.008, 0.016) for baselines, finally picking 7mmax = 0.008 for parallel AdamW and mmax 0.016 for Local Adam W. For QSR, we adopt the same 7/ma and as in our cosine decay experiments. Additionally, we add a gradient clipping threshold of for Local Adam W with constant synchronization period to stabilize training.
23

---

Published as a conference paper at ICLR 2024
Training details for experiments in Appendix G. For the experiments in Table 6, we employ the same weight decay À and peak learning rate Mmax as used in the cosine schedule. Specifically, set À 0.1, Mmax 0.008 for parallel Adam W and X 0.05, Tmax = 0.008 for Local AdamW. In Table 6(a), for the cubic rule, we search among (0.0025, 0.005, 0075, 0.01) and opt for 0.0075, which gives the highest test accuracy. For QSR, we adopt the same a value, 0.0175, as in our cosine decay experiments. In Table 6(b), we set p = 0.0075 and a 0.0175 for the cubic rule and QSR, respectively, which optimal for the original cosine decay schedule, indicated by Figure 6. As mentioned in Section 2, the final synchronization period may be truncated. Specifically, workers are forced to synchronize at the last iteration if the last synchronization period exceeds the remaining iterations. However, the modified cosine schedule experiments seek to validate that the cubic rule can produce an overly large when the learning rate constant. To prevent the truncation from distorting the results, we present the test accuracy at the conclusion of the last full synchronization period, which is not truncated, for both scalings.
Training details for experiments in Appendix H. For Local AdamW + SWAP experiments in Figure 9(b,, we use the cosine learning rate schedule with peak learning rate mmaa 0.008 and weight decay À 0.05. We start with Local Adam with constant synchronization period H and explore the switching point to from {220, 240, 280).
D ADDITIONAL EXPERIMENTS ON RESNET-50
Our paper primarily focuses on training relatively large models with long horizons and proposes QSR to effectively improve the generalization while saving communication. However, the flip side, QSR may not always yield noticeable generalization benefits for smaller models trained with shorter horizons. As shown in Table 5, for the 90-epoch training of ResNet-50 with cosine learning rate decay, the generalization benefit of QSR over Local SGD with a constant communication period negligible. Nonetheless, training in this regime is not costly, either, making less of critical concern. Specifically, completing this 90-epoch training of ResNet-50 requires only 6.6 hours on single machine equipped with NVIDIA GeForce RTX 3090 GPUs. In comparison, the 300-epoch training of ViT investigated in the main text necessitates over 50 hours on the same setup.
Table 5: QSR does not yield noticeable improvement test accu- racy for the 90-epoch training of ResNet-50.
Method
Val. Acc. (%)
Parallel SGD
76.84
Local SGD (H = 2)
76.60
+QSR (Haase =
76.65
24

---

Published as a conference paper at ICLR 2024
E SUPPLEMENTARY MATERIALS FOR SECTION 3
E. MISSING DEFINITIONS AND ASSUMPTIONS
For function F TRd R.d. we use OF( to denote its Jacobian and use 02 F()) to denote the second order derivative at O. For any matrix M E Rd [(F = EiEld where e, the i-th vector of the standard basis. For convenience, we write 8 (O) as [M]
8 M)
80
Assumption E.1. Following Gu el al. (2023). we assume that L(0) and E(0) are -mooth on R2. We also assume that [[EE(0; E)Il2 is uniformly bounded for all A and &.
Assumption E.2. I is a C -mooth, (d m) dimennsiil compact submanifold of Rd such that any ( E I is local minimizer of C and (V2))()) = Additionally, there exists open neighhorhood U of such that r = argminozu
Assumption E.2 is motivated by recent empirical observations that low-loss solutions on the loss landscape are not isolated but path-connected (Garipov et al., 2018; Draxler et al., 2018; Frankle et al., 2020). also adopted by Li et al. (2021c); Lyu et al. (2022); Gu et al. (2023).
Definition E.1 (Gradient Flow Projection). Fix Aull r. For E IRd, the gradient flow starting from x is the solution to dæ(t) (((((t)) with the initial conditiona( = x. The gradient lim, + sx) if the limit exists and belongs r. Otherwise,
dt projection of is defined as Φ (xx) == Ouill-
Definition E.2 (Slow SDE for SGD, formal), Given Ço E T. define C(t) as the solution to the following SDE with initial condition ((0) = Co:
dc(t)= Pc
1 VB
E (S)dW, 2B (=(()))
(6)
(a) diffiution on T
(b) drift
Here, for any I, Pe is projection operator that maps any differential form AdW, + bdtt in Itō calculus to AdW) + (Øb(S)b + (C)A(C)AAATT)) which guarantees to remain on the manifold after taking such an infmitesimal step. B is the total batch size. Σn({) == (+(())((())(() is the covariance matrix of gradient noise projected onto the tangent space of at I. and Eg (C) is the noise covariance in the rest, with coordinates rescaled in the eigenbasis {(A,, ((A,,,)) b of C(C)::
E = E EAS(AA:OVVOAAHO ActA, (E()) (S), v.0)
Definition E.3 (Slow SDE for Local SGD with H n formal). Consider the scaling H = B/n for some constant B. Given So E I, define ((t) as the solution to the following SDE with initial condition ((0) Co:
dc(t)= Pc
1 E VB /
(C)dW, 2B V (C) -AC--KK-- L (,) (7)
2B
(a) diffasion on r
(b) drift
SGD
extra drift
where K is the number of workers, B, E (S) and Eo(s) the same in Definition 3.1. Here, (C: is PSD) matrix depending on gradient noise and Hessian defined as follows:
T (c) EnsH() (0nH-Ox+2211 (Eg(S), u,0,) v.0
Acthy
where {u: a set of eigenvectors of V²L(C) that forms an orthonormal eigenbasis, and dy. A da the corresponding eigenvalues. Additionally, vb(x) =-112 for z + 0 and (0) =
Notice that (xx)) monotonically increases in and has the limit lim (xx) = 0 and lim,--oc (x) Therefore, given c, (C; B) is a monotonically increasing function of 3 in the eigenspace of the Hessian matrix 2 L(C).
25

---

Published as a conference paper at ICLR 2024
E.2 PROOF FOR THEOREM 3.1
We consider the asymptoties that 17 + 0, a 0 and a (2()²)) for all y > We use big-O nota- tion to hide constants independent of 77, a, and use big-O notations to hides constants independent of v) a and also polylog factors of 77, a. We define <b(*) (((())) and let Reo 1777 be the total number of rounds.
Proof outline. The general framework of our proof follows (Li et al., 2019a) which demonstrates the close tracking between SGD iterates and the conventional SDE by examining the moments of parameter changes over small observation interval 1). However, their analysis is not directly appli- cable to our case. Their SDE approximation is only valid for O(n steps while our QSR involves multiple communication rounds, each containing 2) steps. To tackle this challenge, we treat each round as a continuous- ments of changes in the manifold projection of Local SGD and the corresponding slow SDE (5), specifically the moments of ((++1) (b( and Ç((s + 1)a - <(sa²), are closely aligned.
observation interval of length a², and then establish that the mo-
Notably, though the results in (Gu et al., 2023) serve as a building block compute the moments of ((++1) (5(*) in Lemmas 1. E.3, their analysis not trivially extendable to QSR. This is because their analysis depends on the condition Нп (1), and many bounds therein explode as Hn co, e.g., Theorem 3.3, Lemmas 1.14 and 1.16 therein. In the context of QSR, where Hn a² goes to infinity as approaches 0, the condition Hn = (1) is violated, rendering the analysis in Gu et al. (2023) ineffective for QSR.
"
In the following lemma, we present equivalent forms of (3), (4) and (5) that are less intuitive but more friendly to mathematical analysis.
Theorem E.1. Equations (3), (4), (5) can be rewritten as the following SDEs, respectively:
dc(() = 04 (C)E) +
02 (()())((())
(9)
VB
2B
d((t) 1 ((
1
8 (()(E()((()))))))))))))))))))))) + (K
(10)
VB
2B
K
dc(() = (0(())(())
0 (C)
(11)
VB B
2B
Proof. Directly apply Lemmas I.1 to 1.5 of Gu et al. (2023), and we have this theorem.
Based on Gu et al. (2023)'s analysis, below we compute the moments of g(++1) Φ through series of lemmas. Then, we follow Gu et al. (2023)'s method of moments derive the SDE approximation.
Lemma E.1. For any round S Rio and any worker ks [K]. if ((()) E r, then it holds with probability at least 1 - 6, where = Opooly())) that (((())) E and (0(0) ((((() 2l2))12 =
k.H)
k.H
O n log n
Proof. The key insight is that the dynamics of each worker before averaging in each round is just the standard SGD with a smaller batch size, Bloe- Since the distance bound I, Theorem 3.3 in Gu et al. (2023), also applies to SGD by taking K' = 1 and H' = n we can apply this result to obtain that (x) - (((())
((Ok))l )l2 = Oly n log )
Before computing the moments of the change in manifold projection for each worker ((0(.)) ((s), introduce Preliminary Lemmas E. and E.2. Specifically, the Itô-Taylor expansion lemma E. is a straightforward application of Lemma B.7 of (Malladi et al., 2022) on a bounded set. Preliminary Lemma E.2 is adapted from Lemma 26 of (Li al., 2019a).
26

---

Published as a conference paper at ICLR 2024
Let X be the solution to the SDE (t) ((X (X
where ( R R
the drift function and (-)) R + Rax is the diffusion matrix. Both b(.) and o belong to
Ct. Let S be a bounded invariant set of the SDE. That is, if (() E S. for any t 2 0, X (t) E S almost surely. Let be the "effective learning rate", which can be viewed as the length of the continuous-time observation interval for (t). Then we have the following lemma.
Preliminary Lemma E.1 (Itô-Taylor expansion). Let R2 Define
R be any C"-smooth function.
Ag(æ:) Σ (1(x))0g((a) vE(D)
E
Σ ((l(())oss())) [E[D]
Rgg(ar)
(12)
(.jE(D]
Given X(1) = x E S, there exists constan! C independent of the such that
(X (t + ne)) g(æ) (g(A)()))] & Cni
Proof. WLOG, we prove the case for t = 0. Due to the Markovian property Itō processes, the same proof can be done for any >> 0 by time shift. Give X(0) E S, by Itô's lemma,
g(X(nn)) = g(x) + / Ao(X(s))ds+ 1/" (A)(X
where A(x) (②) T Vg(x).
Further apply Itô's lemma to Ag( (s) and we have
g(X(n))) = g(x) + Th
fa (A((z) + L ( J/ (A))(
+ Lo (Ag(X(s)),
g(x) neAl)) J/
['" A ggg(X
[ (AAo(X (r)), ds 1 (Ao(X
Take expectation on both sides, and the last two terms become zero:
Eg(X(ne) g(x) + Ag(x)) + I." no J (()((()))dds...
Since X() belongs to the bounded set S, there exists constant C independent of be such that
[A?g(y)) S C' for all y E S. Therefore,
[[g(( (n=) g(xr))))))))(()) Ag())) S Cnje
Preliminary Lemma E.2 (Adaptation of Lemma 26 in (Li et al., 2019a)). Given X(t) E S, denote the change in X(s) over time interval Tee À Á(x, me X(t + ne) Then, for all x and 2 0, there exists a constant C' independent of The such that
A (a. (a, 1, ne)) C'n t+1
4I 4 ... Ant1 S d,
where n 2 1.
Proof. WLOG, we prove the case for 0. Due to the Markovian property of Itô processes, the same proof can be done for any > by time shift. Denote A(x) = A [x, ne) for brevity. By definition,
Å ((x) L/ ((x(s)) f The
((X(s))
27

---

Published as a conference paper at ICLR 2024
By triangle inequality, for all i E [d].
(xxx) - $11% b(x b(x()))ds
+ +11%** a(x())aw.)
Therefore,
Th
I o (X(s)) ! + +
1+1
A (x) =
b (s)
==1
n+1
n+1
The ((X(s))dL
%
S 2"
E
+ 2
E
o (X(s))
T2
By triangle inequality,
[ T
Ti S E
(1(((())l2d))
By Cauchy-Schwand inequality and Itô's isometry,
1/- The o (X
T2 Σ
(x() (X())))
Since X(s) E S almost surely and S is bounded set, there exists constants C1 and C2 such that Ti S Cynl, Ta < C2ne. Substituting the bounds for Ti and To back, we have the lemma.
Lemma E.2. For any round Reut and any worker lo E [K], given (6)) E I, then
a²
( (60) 1
O² ( E(
(13)
2B0cc
a²
E
(#()))) (,
(0))) k.)) -b(1)) [b00]
-
(14)
Bloc
(3) 03 (6)) = =O(at) ( k.H
(15)
E
(+(0()) k.H) g [(0] =O(a).
(16)
Proof. Again, the key insight is that the dynamics each worker before averaging each round just the standard SGD with smaller batch size, Bloe- Since the SDE approximation theorem for Local SGD, Theorem 3.2 in Gu et al. (2023), also applies to SGD by taking K' and H' "
can apply this result to obtain that, for any -mooohh function g it holds for defined in (9) with the initial condition (0) (5()) that
E (lg((((0)))) ((())) - E[o(C(T)))] = ((,
(17)
where 7" a² is the continuous-time observation interval.
To establish a connection between the moments of (<(o(t)) (++) and those of C(T') - (b()) we
can let the function g(0) to take specific forms, each returning a single coordinate of (1a)
(0 - (0 (())) (0- d(1))83 and ||0 ( I|5. For example, to relate (0 (x) ((()
k.H
<(T') ((0), let g(0) (e1, 0) where e₁ (1, 0, .....,..,00 T. Substitute into (17), and we
get (el, F. (e,,E,
(((0) k. (()) EE(((((T)) ((+)) = (125.))). We can obtain the same results
for all coordinates by letting g()) (e:, 0) for all i E [D]. Therefore, [0()(
k.
E(((T') (0)) = (1,0.20) Similarly, we can show that the LHS of (14) to (16) are only changed by 0.55 o(poly(a)) when replacing (()(.) with <(T").
28

---

Published as a conference paper at ICLR 2024
Then, it suffices to compute the moments for <(T") and verify that they match the RHS of (13) to (16). Since I' is compact and invariant for the SDE (11) (Lemma L.39 in (Gu et al., 2023)), apply the Itô- Tayloy expansion in Preliminary Lemma E.1 with Te a², X (t) ((t),
b(() 2B 8
x (() (()()( and (S) v 2/3 ( (S)
To obtain the first moment (13), let g(C) (el, - x(1)) and substitute it into (12). By Preliminary Lemma E.I, we have
Fl (T') by (cb(t)) -0(T2)=((0')
We can repeat this process for all coordinates of ((T") to obtain
a²
Bl((T") - (+))
b =
8 [X([
(18)
and thus (13).
For the second moment (14), define (C) (Ma (S - ((*))(( ()( where My
[1, - (1.5), Since O. (C) = 0 for all i', the first term of Ag (C) vanishes. It suf-
(id)
(0, otherwise
fices to compute the second term, When i = j. 8% g (3.00 (C) = (0.
= (2, 00 otherwise
Therefore,
(g( Σ (,,(())),)())
Vi [D].
(19)
IE(D]
1, ((' E ((.)))
When 8.
(1,0 (C) 1gg00)
Therefore,
0, otherwise
(1)
Σ (o)()) it IE|D
(20)
Combining (19) and (20) and noticing that ((+3)((()))) = 0 for all i, we have
a²
B[(C(T)) (
E ( + O(a*),
(21)
B
and thus (14).
For the third moment (15), define (8.5 (C) ( e, @ e, 8 e,, (0()(_))
( Noticing that
(()()))))))))))))))))))))))))( for all i' and 0g y g (Q)))) = 0 for all (1',)). we have
E - (())) -
{ = Oaa))
(22)
and thus (15).
Finally, by directly applying Preliminary Lemma E.2, we have
E [[<(T') o6 111 r Oaa))
(23)
and thus (16).
Now are ready to compute the moments for (f++1) - -b))) at each round:
Lemma E.3. For any round Reot, given (s) T, then
E El - -b(0)
((,0]
= a²
(T (5) (E((0))) O(o+),
(24)
2B
a²
E - ( -6(x) (
T [s))
B (0)()) - 0(a4),
(25)
[1 - 1b10)11 (600 = -aaa))
E
(26)
29

---

Published as a conference paper at ICLR 2024
Proof. Let A
((Ok,1)) E
Φl By Lemma E.2,
K
1
B
1 =
E
- -
K
a
[E +0
OBc
K
B =
(0 j.H 2 -(00)
-
K2
k=1
= K a²
B (1)))))) +O(aa) + K(K 11) -(aa)
= a²
( + +((a).
Bloc
Let 2 R Σ k
(
Φ(0). Then A A Σx= ( ((0
)). Finally, let ((5). By Lemma E.1 holds with probability at least - that
A3
2 K
E 0
KHH
|10 2
O log 1) and thus || | A l22 = Oly 1 log H)) Let d 17100_ Since
|100(-) 1l1 is always bounded by O(1), we can always add an error of ((6) to our bounds for the moments and ignore the possibility that this event does not happen. To prove (24), we do Taylor
expansion of @ at (+)), then
E A
( A2) -01
A x 0 1 1 +6)
E
(9) A + ( [B)] + +0 11 +0(y +5)
0
[ + (aa))
200c
The last equation uses the fact that ( for r, is a projection matrix the tangent space of I at 0 (Lemma 4.3 of (Li et al., 2021c)).
To prove (25), again do Taylor expansion of t at (+() connect ((gh()) A2) with ((b(*)) A1) and obtain:
A I]
+ A2) - (
A2) - T
E
(
+ A1) ( + A1) - ((0)]// + +6)
Applying the second-order Taylor expansion gives
E
A
E
(BD(
( (})()))))
+ (8 (1)
A A 10 of
A + A (((()))))))))))))))))))))))))))) AilT)
+ + ()) +o (o mlog 1+44 +6)
By (15) and the fact that (02 (*) I|2 is bounded, the above equation can be simplified
1 + 6) o(a*) + o 1) log n A ($(6!*) (+) A (8()))) + O((a))
2
- (+())) + O(at).
Finally, for (26), we can repeat the above process to bound [[|A,118 and then conclude that
E Ag O)a)).
30

---

Published as a conference paper at ICLR 2024
Now are ready to prove main theorem.
Proof for Theorem 3.1. Let ((t) be the solution (11). Let r be some integer greater than 8. Φ(s) E r, define Sa. as the random variable sampled from the distribution of ((a²r) conditioned on ((a2s) Φ(s), If (6(*)) mull I, define Č..-
If 0 E I, define u(0,4, (,, .22) the expected value of 9({(62) conditioned on {(4)) = 0. If 0 = 0mulls define u 00, 11 (2, 0. That
((,,t1,t2) = [[(((((12) ((h)) =0], E
0.
0 $ r.
For all & Reot, have [ [1)] Elg(S(na2))] (l Eig(Cu,n)]l
B -
E [Bg((G+1,n)) / (b(+t))) [[g((a,))) sb())
By the law of total expectation and the Markovian property of Itô process,
Tg = E E[g(((na²)) ((sa²)= = xb1))]
= E E[g(5(na()) (((s + 1)a²) 1)()))((88)) c(sa)) - 6(1)
Bluu(S,,411 (s+ 1)a² na²)]
Therefore,
E[g(6() Big((((na))]]]
Σ
Elu(gbot1), (s 1)a², naa)] -E(u(S Elu((S,0+1:: +s+ 1)a? na
T
By the law of total expectation,
T B[u(p((+1), (s 1)a?, na?) ul u(Ga,041, (s 1)a? na²) 1111),(0-111 E
A
+ |P(O))(((()) or ((++)) L)),
where the latter term comes from the definition of u(O, t t2) and č,,.. By Lemma 1.11 in (Gu et al., 2023), there exists constant such that if minøer ||e - Φ112 S then ( E r. Therefore, substituting J = y7 100 into Lemma 1. we can conclude that the latter term is at most O 100).
For A.. notice that the two terms differ only in the first position. By Lemma E.3 and (18) to (23),
the moments of (x++1) ((s) and Č,,0+1 ((()) are close to each other. Therefore, it suffices to discuss the smoothness of and perform Taylor expansion. By Proposition 25 of (Li et al., 2019a),
since g E Ct,, u(Φ, f,, t2) satisfies the compatibility condition for the Whitney Extension Theorem for Φ E T. Therefore, there exists function i(Φ, f1, t2) that is C4 in & for Φ IR and satisfies ū(Φ,,, t,. t2) u(Φ, t to) for all E Г. Denote ū(Ø sa²,, na?) as (Φ) for brevity. Now, we can safely substitute in A, with and perform Taylor expansion:
A, E (((( (0(0+1) (()))), (s 1)a², na²) ( ( ( I
(0) + (S..0t1 (l)))), (s + 1)a² na?) 1 ET)
ET
A
31

---

Published as a conference paper at ICLR 2024
A, = -()))))) + + (
B(olo)))) Φ6 / ((6-1) ETl)
8 ( -
x (0()) xl++1) Erl)
+ -((1 Φ ((+t1) ET))
A" i()l (0 R( x5()) /
+ 2
Uz (
- - 1 Eri)
+ 0 (EI1 11,11
1 ( r])
Substituting in = 17 into Lemma E.I, we can conclude that, given (5()) E I, the event { E
} happens with probability at least 17 We can replace the condition sl()) Φls+1) T with Φ E I' in A with an error of only (17 Therefore,
A. = ( (aūs+Im(((()) [[([((((1)1)))))) b()))= (6,,0+1- - -bb(0) 26)) E erl)
1 O ūs Bottin (1) B[((((()) 66( (G - (
+
2
- 165,011-1110 (5 - - Φ ET]
+ +0(1111((5t) 1 165) E rJ) O Č 112 Φ6 rJ) O(1
Since (5() E I where r is a compact set, both [ (b)))) and (. are bounded. Substituting Lemma E.3 and (18) to (23) to the expression of A., we have Ag = (aa))
and thus [T]l = (aa²). Summing [7,] up, we have [B]g( - Elg(((na²))] S O(na*) S (a))), which completes the proof.
F DETAILS FOR COMMUNICATION TIME MEASUREMENT
straightforward to measure the time duration for the entire training, but it is hard directly measure the communication time due to the asynchronous nature of CUDA computation. Hence, in our experiments, we derive the communication time from the difference total training time across runs with various communication frequencies.
Specifically, let Ttot Thot be the total time durations of data parallel approaches and local gradi- ent methods with H H1, respectively. Also let Teomm, Too.m be their communication times, and para coomp be their computation time. Ideally, setting the synchronization period to H/ duces the communication volume exactly by factor of H. so these variables satisfy the following relationships:
T Tromp
7
comm
H.
H1 para
7room
+
comp
T
H1
T
Tramp
Ttat
para
Then we can express the communication and computation times in terms the total time duration para and Thot.
Thot
H1
c.omm
H1 T.oom
(Ttot
para
[H, )
H1
H1 H1 -
1
coom 11,
7 para
Tccomm para
7ot
tot para
para
H
32

---

Published as a conference paper at ICLR 2024
Therefore, empirically measure the total time duration Tpat and H2 for some H1, then the following formulas to obtain estimates of the communication and computation times:
H1
Teamm
(T
para
That
(27)
H1
H1
T
Tist
(28)
para
H1 -
H1
para
These estimates are very predictive for the total time duration of local gradient methods with different H. For example, when H H2, we can predict the total time duration Ttot as follows:
H2
Troom
Tcomm
(29)
H2
para
T
Tcoom para
(30)
H
para
We find that the relative error
Hig
Tho
X 100%. where Thot denotes the measured total time, is
Tiat
only ~ 1% across all configurations in Table 4, where we set H1 = 2, H2 = for ResNet-152 and H 4, H2
for ViT-B. The small relative error suggests that our method offers close approximation to the actual time. For this reason, in Table 4, we report the communication time estimated by (27) and (29) for data-parallel approaches and local gradient methods with constant synchronization period.
For QSR, since its communication volume relative data parallel approaches, denoted as fOSR. can be easily computed given the learning rate schedule, the growth coefficient a and the base synchronization period Hasee we can estimate its communication time Tromm in a similar vein to (27) and (29):
T QSR
fQSR Tromm.
(31)
We report the communication time estimated by (31) in Table 4 for QSR.
G DISCUSSION ON MORE AGGRESSIVE SCALINGS
(00058)
Top-1 Val. Acc. (%)
041
0.15
80.0
0.2
0.2
[J000]
80.0-
80.5
0.00)
If ((E17)
27 UEE))
0.25)
80.0-
N m
Top-1
Paratiel (GD/Ad
79.5
79.5
20.0
22.5
25.0
100
10
15
20
20
100
Communication Volume (%)
Communication
Volume (%)
(a) Local SGD ResNet-152
(b) Local AdamW ViT-B
Figure 6: The cubic rule H [(())) with properly tuned p can either outperform underperform the QSR in test accuracy, depending on the training scenarios.
Apart from the scalings discussed in Section 3, one can consider more aggressive scalings, e.g., H 33]. Compared with QSR this cubic synchronization rule communicates more frequently at carlier stages but much less at later stages. Our theory in Theorem 3.1 suggests that taking ~ nj blows up the approximation error, but as shown in Figure 6, this cubic rule with properly tuned p either outperform or underperform the QSR in test accuracy, depending on the training scenarios.
(a/n)) -2] that uses the same amount of communication,
We argue that this is because where the learning rate decays so fast that 7/, sees significant decay within a single communication round. As example where the cubic rule performs better, plot in Figure 8 the test accuracy
quasistatic view may break in the very late phase of cosine decay,
33

---

Published as a conference paper at ICLR 2024
Table We validate the higher test accuracy achieved by 17 relies rapid decaying learning rate within synchronization period via ablation studies on ViT-B. In Table 6(a), we replace the cosine decay schedule with a variant of the step decay schedule in Smith et al. (2020), In Table 6(b). we run both scalings with a modified cosine decay schedule that ceases to decay some epoch t". QSR consistently outperforms H 17
in both cases.
(b) Local AdamW with modified cosine decay. Both scalings use Hbase
(a) Local AdamW with step decay.
Method
t"
Val. Acc. (%)
Method
Val. Acc. (%) Comm. (%)
QSR
260
80.75
Parallel AdamW
78.51
100
H 17
260
80.51 1
Local AdamW (H=4)
78.70
25
QSR
250
80.34
+QSR (Husse 4)
80.99
13.2
H ~ 17
250
79.91
+H ~ (Hbase = 4)
80.86
14.4
QSR
240
79.85
~ 11 3
240
79.72
curves of the QSR and cubic rule for training ViT-B with Local Adam W and batch size 4096. Same as our experiment setup in Section 4, the learning rate peaks the value 0.008 and then decays to nearly zero (10-9) following cosine decay schedule. Setting H = (0000005/))))) consistently worse test accuracy than QSR (with the same communication volume) before epoch 265. However, during the final communication round, which spans from epoch 265 to 300, the cubic rule catches up with QSR. During this period, the learning rate dramatically decreases from 3.5 10-4 to nearly zero, but quasistatic view assumes that the learning rate T) should remain relatively constant for at least one communication round.
results in
Based on the above observation, rule offers benefits over QSR only for certain schedules that have a rapid tail of learning rate decay near the end of training. To validate this view, we replace the cosine decay schedule with variant of the step decay schedule in Smith et al. (2020). In our step decay schedule, given peak
argue that the cubic
cosine
step decay similar to Smith al.(20
modified cosine decay
0.008
total of 300 epochs, the learning rate remains until epoch 1 150, after which it is divided by 2 every 30 epochs. See Figure for an illustration. Unlike the cosine schedule, this step decay schedule maintains constant learning rate for a significant amount of time. As shown in Table 6(a), the cubic rule yields inferior generalization performance compared with our QSR, even after careful tuning of p. See Appendix C.2 for training details.
Learning Rate
0.006
0.004
0.002
0.000
1001
200
300
Epochs
Figure 7: An illustration of the learning rate schedules.
Another way to corroborate our view to run both seal- ings with modified cosine learning rate schedule, which ceases to decay after a specific epoch !" and remains con- stant until training ends. See Figure 7 for illustration of this modified cosine schedule. As shown in Table 6(b), QSR consistently outperforms the cubic rule across various choices of 1", Further training details can be found in Appendix C.2. The probable reason is that when the learning rate is held constant, the cubic rule results in an excessively large H, negatively impacting optimization.
Given these failure cases of the cubic rule, generally recommend using the QSR and leave it to future work to design a better rule to deal with schedules that have rapid tail of learning rate decay.
34

---

Published as a conference paper at ICLR 2024
Schodale Cummancatioo Volume
~ 1 14.9%
~ 1 15.4%
80
28 755
Ar
81
05:
80
Topa
73
60
200
250
300
35
100
200
300
Figure 8: Test accuracy curves for QSR (a = 0.01) and the cubic rule (p 0.0075). The cubic rule results in consistently worse test accuracy than QSR (with the communication volume) before last communication round.
H COMPARISON WITH LOCAL SGD/ADAMW SWAP
In this section, compare QSR with the modified Stochastic Weight Averaging in Parallel (SWAP) algorithm, termed "Local SGD/AdamW + SWAP"" Specifically, the original SWAP proposed by (Gupta et al., 2020) uses SGD for the majority of the training process and only switches to local updates at some to near the end, thus saving less communication than QSR. To compare SWAP with QSR at similar level of communication volume, experiment with the modified SWAP, which starts with Local SGD/Adam using a constant communication period Hase and, after some time to, lets workers perform local updates with a final model averaging. As shown in Figure 9, QSR outperforms Local SGD/Adam W SWAP though we have tuned to carefully for the latter.
(125)
C.017
BOIS)
to 0.0LY
80.0
Acc.
Top-1 Val. Acc.
80
R SWAP with ta
79.5
79-
Top-1
Porailel SGD/AdamW
79.0
0000
17.5
20.0
22.5
25.0
100
16
15
25
100
Communication Volume (%)
Communication Volume (%)
(a) Local SGD ResNet- 152
(b) Local AdamW on VIT-B
Figure 9: QSR outperforms Local SGD/AdamW SWAP on both models. See Appendix C training
details.
35


Reviews:
[{"review_id": "pxzpK4TCH8", "reviewer": "Reviewer_u85z", "summary": "This paper proposes a scheduler for the synchronization interval of local SGD/ADAM, a.k.a. the Quadratic Synchronization Rule (QSR), which recommends dynamically setting such intervals in proportion to the inverse of the square of learning rate. The proposed algorithm is supported by theoretical analysis based on SDE. The empirical results show that QSR can achieve better validation accuracy with less communication overhead.", "questions": "1. Although QSR is supported by theoretical analysis, I still wonder how the other options work, such as setting $H$ proportion to $\\frac{1}{\\eta}$ or $\\frac{1}{\\sqrt{\\eta}}$. Is there any corresponding ablation experiments?\n\n2. Many theoretical analysis suggests a learning rate scheduler $\\eta \\propto \\frac{1}{\\sqrt{T}}$, where $T$ is the number of steps. What if we detach the choice of $H$ from the learning rate, and simply use $H \\propto T$ (with some rounding of course)?\n\n3. Is it possible to extend the theoretical analysis to local Adam?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper proposes a scheduler for the synchronization interval of local SGD/ADAM, a.k.a. the Quadratic Synchronization Rule (QSR), which recommends dynamically setting such intervals in proportion to the inverse of the square of learning rate. \n\n2. The proposed algorithm is supported by theoretical analysis based on SDE. \n\n3. The empirical results show that QSR can achieve better validation accuracy with less communication overhead.", "weaknesses": "1. The experiments focuses on vision tasks and models. Although ViT model uses the transformer-like architecture, I'm still very interested in how QSR performs in NLP tasks with transformer architectures.\n\n2. The theoretical analysis only applies to local SGD. I wonder whether it could be extended to local Adam.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sa2zG2B1uy", "reviewer": "Reviewer_tFqJ", "summary": "This paper builds upon the theoretical insights from [[1]](https://openreview.net/pdf?id=svCcui6Drl) to introduce the \"Quadratic Synchronization Rule\" (QSR), a theoretically grounded method to set the parameter $H$ in Local-SGD. By analyzing a \"Slow SDE for Local SGD with QSR\", it is shown that scaling $H \\propto \\frac 1 {\\eta^2}$ (with $\\eta$ the Local-SGD decaying learning rate) allows to regularize the objective function towards flatter minima, increasing generalization compared to other schemes (constant values of $H$, or $H \\propto \\frac 1 {\\eta}$). Extensive experiments using ResNet-152 and ViT-B on ImageNet are made to confirm the advantages of QSR compared to previous methods both in terms of test accuracy and wall-clock time for the training.\n\n[1] Xinran Gu and Kaifeng Lyu and Longbo Huang and Sanjeev Arora. *Why (and When) does Local {SGD} Generalize Better than {SGD}?*, In The Eleventh International Conference on Learning Representations, 2023.", "questions": "* A network speed of 25Gbps is used for the experiments. Would your algorithm still help in HPC settings where fast communications can happen over 100 Gbps Infiniband ?\n* Are both the model size **and** number of epochs important factors for QSR to work ? For example, using a smaller model (such as ResNet18) but with a decent amount of epochs (e.g., 300 epochs on CIFAR10), would QSR have an impact ?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* **Promising experimental results**: In the presented settings, the introduction of QSR seems to lead to noticable gains in terms of validation accuracy compared to all the chosen baselines, as well as some mild savings in wall-clock time compared to methods using a constant $H$.\n* **Theoretical analysis of Slow SDE for Local SGD with QSR**: Efforts are made to make the reader intuitively understand the changes the introduction of QSR makes to the Slow SDE and how the $K$ times larger drift term favors generalization.", "weaknesses": "* **No confidence intervals in experiments**: given the extensive hyper-parameter search Appendix C claims, I would have expected that the  experimental results in the main paper would come with confidence intervals, and were presented as average over several runs to confirm the reproducibility of the results.\n* **No visualization of the effective $H$ scheduler**: Tab. 4 hints that using constant $H$ and QSR actually leads to very similar training times, suggesting that using the QSR scheduler for the values of $\\alpha$ retained does not lead to major deviations compared to using constant $H$. In effect, what does the evolution of $H$ look like throughout training, and how different is it from using a constant $H$ ?\n* **Theoretical reasons for the focus on large model and long horizon unclear**: While the main paper focuses on fairly large models (ResNet-152 and ViT-B) and long horizon (200-300 epochs on ImageNet), Appendix D shows that for training a ResNet-50 on 90 epochs, QSR has no effect. Are there any theoretical reasons why QSR should only work for large model and training procedures ?\n* **Unclear why $\\gamma =2$ is optimal:** Section 2 claims that *\"$H^{(s)}$ could have been set to $H^{(s)}:= \\max \\left \\\\{ H_\\text{base},  \\left \\lfloor \\left ( \\frac{\\alpha}{\\eta_t}  \\right )^{\\gamma} \\right \\rfloor \\right \\\\}$ for any $\\gamma$\"*. However, it still remains unclear to me why setting $\\gamma = 2$ is optimal.\n* **Comparison with related work lacking:** Appendix A cites the *\"Stochastic Weight Averaging in Parallel (SWAP) algorithm\"* [[2]](https://openreview.net/pdf?id=rygFWAEFwS) as another method that varies the value of $H$ during training, leading to a decrease in training time and better generalization. Yet, no comparison with the method is performed in the paper.\n* **Additional hyper-parameter to tune**: *\"choosing the right value of $H$\"* has been traded for *\"choosing the right value of $\\alpha$\"*, thus QSR does not \"automatically\" set $H$ and no hyper-parameter tuning has been saved.\n\n[2] Vipul Gupta and Santiago Akle Serrano and Dennis DeCoste. *Stochastic Weight Averaging in Parallel: Large-Batch Training That Generalizes Well*, In International Conference on Learning Representations, 2020.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "48q2Cv7Jka", "reviewer": "Reviewer_ydD5", "summary": "The paper examines the optimal approach for dynamically determining H, a hyperparameter in Local SGD, allowing workers to compute locally for H steps without synchronizing with others.\n\nThe goal is to strike a balance between communication costs and the final test accuracy. While a large 'H' might degrade accuracy, a small 'H' could intensify communication overhead. Drawing inspiration from the stochastic differential equation introduced by [Gu et al.](https://arxiv.org/pdf/2303.01215.pdf), this paper introduces a new Quadratic Synchronization Rule for setting the H dynamically given a predetermined learning rate schedule. Furthermore, the research provides proof delineating the approximate convergence bounds concerning the distance between the iteration and the manifold. The paper also presents experimental results in various settings.", "questions": "(1) Have you considered comparing the performance of $H\\sim \\eta^{-3}$, $H\\sim \\eta^{-4}$? While prior work [Gu et al.](https://arxiv.org/pdf/2303.01215.pdf) already shows that the generalization benefits arise only if H = Ω( 1/η ) when η → 0, this work specifically sets H to $η^{-2}$. This results seem to be quite heuristic and not a thorough investigation of the problem domain. What understanding about the limitations or insights of the method has been derived?\n\n(2) The proof provided for Lemma E.2 appears to be too succinct. Furthermore, in Lemma E.2, it is stated that 'we can apply the Itô-Taylor expansion from Lemma B.7 of Malladi et al. (2022) to derive (16), (17), and (18).' Is there a typo in this reference to equations? I suspect you meant to refer to equations (11), (12), and (13). Also, could you please elucidate how you derived equations 11, 12, and 13?\n\n(3) Proof for Theorem 3.1 is overly concise. What is Lemma I.43? What is the definition of r? How do you get \n$$E[u(\\hat{\\zeta}_{s, s+1},(s+1) \\alpha^2, n \\alpha^2)]$$\n\nfrom\n \n$$E[g(\\hat{\\zeta}_{s,n})]$$\n\n(4) How do you calculate the communication time in Table 4, why don't you opt for using communication volume instead?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "S1. The concept presented is intuitive—convergence can be improved by dynamically tuning H the number of steps of local updates before syncing with others. The introduced Quadratic Synchronization Rule holds potential for application in distributed machine learning.\n\nS2. The paper provides a theoretical analysis of their idea and conducts experiments to support their theoretical results.", "weaknesses": "W1. The motivation of finding 'H’  given a learning rate schedule is questionable. Both learning rate and H should be tunable hyperparameters in distributed learning.\n\"Specifically, we introduce a simple yet effective strategy, called Quadratic Synchronization Rule (QSR), for dynamically adjusting the synchronization period according to the learning rate: given a learning rate schedule, we set H proportional to η^−2 as the learning rate η decays. This rule is largely inspired by a previous theoretical work (Gu et al., 2023), which shows that the generalization benefits arise only if H = Ω( 1/η ) when η → 0, but did not make any recommendation on how to set H.\"\n\nSolely determining H given the learning rate schedule, yet without asking about what the optimal learning rate is is a less valuable question. While prior work [1] already shows that the generalization benefits arise only if H = Ω( 1/η ) when η → 0, this work specifically sets H to 1/η^2, which is incremental, reinforcing the same results. In fact, the learning rate is of greater significance in deep learning and distributed learning; for many problems, one can boost generalization and convergence performance solely by carefully tuning the learning rate. A more valid question that I believe should be asked is what is the optimal dynamic learning rate schedule under a given schedule of H, e.g. H=1 or 2 or Post-local SGD with variable H. There is no comparison to this learning rate question under a fixed H schedule, which further saves communication and is easier to implement in reality. Therefore, I think the problem studied here to recommend how to set H is kind of artificial.\n\nW2. The baseline for comparison is not clear. For instance, what is the Parallel AdamW? How does the Parallel AdamW work? I can’t find the rigorous definition of it given by the paper. Because of the second momentum terms, Parallel AdamW that assumably aggregates the model parameters every H=1 step is not equivalent to Centralized AdamW when only setting H = 1. The second momentum can not be simply averaged. Even the definition of Parallel SGD is not provided. The paper simply says that refers to Local SGD when H=1. At least a reference or a definition should be provided. \n\nW3. Some experimental settings are questionable. What has motivated the adoption of only SGD for ResNet experiments and AdamW only for VIT experiments? Furthermore, while all experiments employ momentum-based optimizers and underscore the effectiveness of QSR with such optimizers, the paper doesn't deeply explore the role of momentum in this methodology.\n\nW4. The theoretical analysis and proof overlap too much with previous work [1]. Both rely on the same SDE technique (which is the novelty of the prior work) for analyzing analogous issues, with the current work lacking additional theoretical insights and development. Additionally, some proof steps are missing. See questions for details.\n\nW5. There is no comparison with other practical distributed learning systems and methods, given the vast literature on distributed learning. The comparison is more of a simple self-comparison between different hyperparameter scalings of H, e.g., constant H vs. the QSR of H,  assuming the same learning rate schedule is used. This comparison is not convincing because 1) for different H setups, the optimal learning rate schedule will vary; 2) it didn’t compare to other possibilities of H = Ω( 1/η ) making it hard to understand the limitations of the problem domain. This makes the work more of hyperparameter tuning heuristics and lacks insights. The experimental results are more of validating the theorems and lack practical value. \n\nW6. The exploration of the choice for $H_{base}$ seems insufficient. Providing a graph that illustrates the relationship between final test accuracy, communication volume, training time reduction, and the selection of $H_{base}$ would be more persuasive.\n\nFor the above reasons, I would recommend rejection.\n\nReference:\n\n[1] Gu, Xinran, et al. \"Why (and When) does Local SGD Generalize Better than SGD?.\" in ICLR 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yZ1HLLU2YC", "reviewer": "Reviewer_duf3", "summary": "This work proposes a new synchronization rule for local gradient methods, where synchronization among processors is performed after multiple rounds of local updates. In the synchronization rule, the number of local updates is proportional to the inverse of the learning rate square. Theoretical analysis is provided to motivate the proposed synchronization rule. Experimental results on ResNet and ViT show that local gradient methods with the proposed synchronization rule yield better test accuracy and less training time.", "questions": "n/a", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed quadratic synchronization rule is simple but effective, backed up by strong theoretical analysis.\n2. Experimental results show both efficiency gain and better generalizations.\n3. The experiments are comprehensive, including various architectures, learning rate decay strategies, and batch sizes.\n\nOverall, I think this is a good paper.", "weaknesses": "n/a", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)