A Communication-Efficient Distributed Gradient Clipping Algorithm for Training Deep Neural Networks

( 1 Nϵ4 ) iteration complexity for finding an ϵ-stationary point , where N is the number of machines . This indicates that our algorithm enjoys linear speedup . Our experiments on several benchmark datasets and various scenarios demonstrate that our algorithm indeed exhibits fast convergence speed in practice and validate our theory . 1 INTRODUCTION . Deep learning has achieved tremendous successes in many domains including computer vision ( Krizhevsky et al. , 2012 ; He et al. , 2016 ) and natural language processing ( Devlin et al. , 2018 ) , game ( Silver et al. , 2016 ) . To obtain good empirical performance , people usually need to train large models on a huge amount of data , and it is usually very computationally expensive . To speedup the training process , distributed training becomes indispensable ( Dean et al. , 2012 ) . For example , Goyal et al . ( 2017 ) trained a ResNet-50 on ImageNet dataset by distributed SGD with minibatch size 8192 on 256 GPUs in only one hour , which not only matches the small minibatch accuracy but also enjoys parallel speedup , and hence improves the running time . Recently , there is an increasing interest in an variant of distributed learning , namely Federated Learning ( FL ) ( McMahan et al. , 2017 ) , which focuses on the cases where the training data is non-i.i.d . across devices and only limited communication is allowed . McMahan et al . ( 2017 ) proposed an algorithm named Federated Averaging , which runs multiple steps of SGD on each clients before communicating with other clients . Despite the empirical success of distributed SGD and its variants ( e.g. , Federated Averaging ) in deep learning , they may not exhibit good performance when training some neural networks ( e.g. , Recurrent Neural Networks , LSTMs ) , due to the exploding gradient problem ( Pascanu et al. , 2012 ; 2013 ) . To address this issue , Pascanu et al . ( 2013 ) proposed to use the gradient clipping strategy , and it has become a standard technique when training language models ( Gehring et al. , 2017 ; Peters et al. , 2018 ; Merity et al. , 2018 ) . There are some recent works trying to theoretically explain gradient clipping from nonconvex optimization ’ s perspective ( Zhang et al. , 2019 ; 2020 ) . These works are built upon an important observation made in ( Zhang et al. , 2019 ) : for certain neural networks such as LSTM , the gradient does not vary uniformly over the loss landscape ( i.e. , the gradient is not Lipschitz continuous with a uniform constant ) , and the gradient Lipschitz constant can scale linearly with respect to the gradient norm . This is referred to as the relaxed smoothness condition ( i.e. , ( L0 , L1 ) -smoothness defined in Definition 2 ) , which generalizes but strictly relaxes the usual smoothness condition ( i.e. , L-smoothness defined in Definition 1 ) . Under the relaxed smoothness condition , Zhang et al . ( 2019 ; 2020 ) proved that gradient clipping enjoys polynomial-time iteration complexity for finding the first-order stationary point in the single machine setting , and it can be arbitrarily faster than fix-step gradient descent . In practice , both distributed learning ( or FL ) and gradient clipping are important techniques to accelerate neural network training . However , the theoretical analysis of gradient clipping is only restricted to the single machine setting ( Zhang et al. , 2019 ; 2020 ) . Hence it naturally motivates us to consider the following question : Is it possible that the gradient clipping scheme can take advantage of multiple machines to enjoy parallel speedup in training deep neural networks , with data heterogeneity across machines and limited communication ? . In this paper , we give an affirmative answer to the above question . Built upon the relaxed smoothness condition as in ( Zhang et al. , 2019 ; 2020 ) , we design a communication-efficient distributed gradient clipping algorithm . The key characteristics of our algorithm are : ( i ) unlike naive parallel gradient clipping algorithm which requires averaging model weights and gradients from all machines for every iteration , our algorithm only aggregates weights with other machines after certain number of local updates on each machine ; ( ii ) our algorithm clips the gradient according to the norm of the local gradient on each machine , instead of the norm of the averaged gradients across machines as in the naive parallel version . These key features make our algorithm amenable to the FL setting and it is nontrivial to establish desired theoretical guarantees ( e.g. , linear speedup , reduced communication complexity ) . The main difficulty in the analysis lies at dealing with nonconvex objective function , non-Lipschitz continuous gradient , and skipping communication rounds simultaneously . Our main contribution is summarized as the following : • We design a novel communication-efficient distributed stochastic local gradient clipping algorithm , namely CELGC , for solving a nonconvex optimization problem under the relaxed smoothness condition . The algorithm only needs to clip the gradient according to the local gradient ’ s magnitude and globally averages the weights on all machines periodically . To the best of our knowledge , this is the first work proposing communication-efficient distributed stochastic gradient clipping algorithms under the relaxed smoothness condition . • Under the relaxed smoothness condition , we prove iteration and communication complexity results of our algorithm for finding an ϵ-stationary point . First , comparing with ( Zhang et al. , 2020 ) , we prove that our algorithm enjoys linear speedup , which means that the iteration complexity of our algorithm is reduced by a factor of N ( the number of machines ) . Second , comparing with naive parallel verion of the algorithm of ( Zhang et al. , 2020 ) , we 1In this setting , we assume the gradient norm is upper bounded by M such that the gradient is ( L0+L1M ) Lipschitz . However , we want to emphasize that the original paper of ( Ghadimi & Lan , 2013 ) does not require bounded gradient assumption , instead they require L-Lipschitz gradient and bounded variance σ2 . Under their assumption , their complexity result is O ( ∆Lϵ−2 +∆Lσ2ϵ−4 ) . 2Naive Parallel of ( Zhang et al . 2020 ) is different from our algorithm ( CELGC ) with I = 1 in that the naive version requires averaging gradients across all machines to update the model while CELGC only updates the model using local gradients computed in that machine . This also means that in each iteration , naive version clips the gradient based on the globally averaged gradient while ours only bases on the local gradient . prove that our algorithm enjoys better communication complexity . Specifically , our algorithm ’ s communication complexity is smaller than naive parallel clipping algorithm if the number of machines is not too large ( i.e. , N ≤ O ( 1/ϵ ) ) . The detailed comparison over existing algorithms under the same relaxed smoothness condition is described in Table 1 . Please refer to ( Koloskova et al. , 2020 ) for local SGD complexity results for L-smooth functions . • We empirically verify our theoretical results by conducting experiments on different neural network architectures on benchmark datasets and on various scenarios including small to large batch-sizes , homogeneous and heterogeneous data distributions , and partial participation of machines . The experimental results demonstrate that our proposed algorithm indeed exhibit speedup in practice . 2 RELATED WORK . Gradient Clipping/Normalization Algorithms In deep learning literature , gradient clipping ( normalization ) technique was initially proposed by ( Pascanu et al. , 2013 ) to address the issue of exploding gradient problem in ( Pascanu et al. , 2012 ) , and it has become a standard technique when training language models ( Gehring et al. , 2017 ; Peters et al. , 2018 ; Merity et al. , 2018 ) . Menon et al . ( 2019 ) showed that gradient clipping is robust and can mitigate label noise . Recently gradient normalization techniques ( You et al. , 2017 ; 2019 ) were applied to train deep neural networks on the very large batch setting . For example , You et al . ( 2017 ) designed LARS algorithm to train a ResNet50 on ImageNet with batch size 32k , which utilized different learning rate according to the norm of the weights and the norm of the gradient . In optimization literature , gradient clipping ( normalization ) was used in early days in the field of convex optimization ( Ermoliev , 1988 ; Alber et al. , 1998 ; Shor , 2012 ) . Nesterov ( 1984 ) and Hazan et al . ( 2015 ) considered normalized gradient descent for quasi-convex functions in deterministic and stochastic cases respectively . Gorbunov et al . ( 2020 ) designed an accelerated gradient clipping method to solve convex optimization problem with heavy-tailed noise in stochastic gradients . Mai & Johansson ( 2021 ) established the stability and convergence of stochastic gradient clipping algorithms for convex and weakly convex functions . In nonconvex optimization , Levy ( 2016 ) showed that normalized gradient descent can escape from saddle points . Cutkosky & Mehta ( 2020 ) showed that adding a momentum provably improves the normalized SGD in nonconvex optimization . Zhang et al . ( 2019 ) and Zhang et al . ( 2020 ) analyzed the gradient clipping for nonconvex optimization under the relaxed smoothness condition rather than the traditional L-smoothness condition in nonconvex optimization ( Ghadimi & Lan , 2013 ) . However , all of them only consider the algorithm in the single machine setting or the naive parallel setting , and none of them can apply to FL setting where data on different nodes is heterogeneous and only limited communication is allowed . Communication-Efficient Algorithms in Distributed and Federated Learning In large-scale machine learning , people usually train their model using first-order methods on multiple machines and these machines communicate and aggregate their model parameters periodically . When the function is convex , there is scheme named one-shot averaging ( Zinkevich et al. , 2010 ; McDonald et al. , 2010 ; Zhang et al. , 2013 ; Shamir & Srebro , 2014 ) , in which every machine runs an stochastic approximation algorithm and averages the model weights across machines only at the very last iteration . One-shot averaging scheme is communication-efficient and enjoys statistical convergence with one pass of the data ( Zhang et al. , 2013 ; Shamir & Srebro , 2014 ; Jain et al. , 2017 ; Koloskova et al. , 2019 ) , but the training error may not converge in practice . McMahan et al . ( 2017 ) considered the Federated Learning setting where the data is decentralized and might be non-i.i.d . across devices and communication is expensive . McMahan et al . ( 2017 ) designed the very first algorithm for FL ( a.k.a. , FedAvg ) , which is communication-efficient since every node communicates with other nodes infrequently . Stich ( 2018 ) considered a concrete case of FedAvg , namely local SGD , which runs SGD independently in parallel on different works and averages the model parameters only once in a while . Stich ( 2018 ) also showed that local SGD enjoys linear speedup for stronglyconvex objective function . There are also some works analyzing local SGD and its variants on convex ( Dieuleveut & Patel , 2019 ; Khaled et al. , 2020 ; Karimireddy et al. , 2020 ; Woodworth et al. , 2020a ; b ; Gorbunov et al. , 2021 ; Yuan et al. , 2021 ) and nonconvex smooth functions ( Zhou & Cong , 2017 ; Yu et al. , 2019a ; b ; Jiang & Agrawal , 2018 ; Wang & Joshi , 2018 ; Lin et al. , 2018 ; Basu et al. , 2019 ; Haddadpour et al. , 2019 ; Karimireddy et al. , 2020 ) . Recently , Woodworth et al . ( 2020a ; b ) analyzed advantages and drawbacks of local SGD compared with minibatch SGD for convex objectives . Woodworth et al . ( 2021 ) proved hardness results for distributed stochastic convex optimization . Reddi et al . ( 2021 ) introduced a general framework of federated optimization and designed several federated versions of adaptive optimizers . Zhang et al . ( 2021 ) considered to employ gradient clipping to optimize L-smooth functions and achieve differential privacy . Due to a vast amount of literature of FL and limited space , we refer readers to ( Kairouz et al. , 2019 ) and references therein . However , all of these works either assume the objective function is convex or L-smooth . To the best of our knowledge , our algorithm is the first communication-efficient algorithm which does not rely on these assumptions but still enjoys linear speedup .

Summary: The paper considers the effect of gradient clipping in Federated Learning and how it affects the convergence rate. They focused on the relaxed-smooth loss function. Each worker uses local gradient clipping and runs multiple steps of SGD before communicating and averaging the local models. The authors theoretically analyzed the algorithm and showed that for $N$ workers, the algorithm has $O(1/N\\epsilon^4)$ iteration complexity to find an $\\epsilon$-stationary point. Finally, the theoretical results are experimentally verified on CIFAR-10 (Resent-56 model), Penn Treebank and WikiText (LSTM models).