Communicate Then Adapt: An Effective Decentralized Adaptive Method for Deep Training

1 INTRODUCTION . Decentralized SGD ( Lopes & Sayed , 2008 ; Nedic & Ozdaglar , 2009 ; Chen & Sayed , 2012 ; Lian et al. , 2017 ; Assran et al. , 2019 ) is an emerging training approach for deep learning known for its much less communication overhead . In contrast to parallel SGD in which a global averaging across all computing nodes is required per iteration , decentralized SGD does not involve any global operations . Building upon partial averaging , in which each node only needs to compute the locally averaged model within its neighborhood , decentralized SGD can save remarkable communications and training time in large-scale distributed deep learning tasks compared to parallel SGD . Although simple to use , the vanilla decentralized SGD sometimes suffers from the slow convergence . Inspired by the well-documented success of adaptive methods such as AdaGrad ( Duchi et al. , 2011a ) , Adam ( Kingma & Ba , 2014 ) and AMSGrad ( Reddi et al. , 2019 ) , several decentralized adaptive methods ( Nazari et al. , 2019 ; Lin et al. , 2021 ) have been proposed to accelerate decentralized SGD training . While these algorithms have achieved remarkable success in several practical applications , they have also been observed to not converge to the desired solution ( i.e. , global optimal solution in the convex scenario or stationary solution in the non-convex scenario ) in some other settings . For example , it has been observed in the convex setting ( see Sec . 3 ) that DAdam ( Nazari et al. , 2019 ) and QG-DAdam ( Lin et al. , 2021 ) do not converge to the global optimal solution . This paper studies this situation in detail . We rigorously uncover the reason why DAdam and QGDAdam fail to achieve the desired solution , and propose a novel decentralized adaptive method to resolve the convergence issue . In particular , we make the following key contributions : • We find the algorithms DAdam and QG-DAdam , while different in concrete recursions , share a similar structure : each node scales its gradient with the past squared gradients ( which is referred to as the adaptive step ) before or while it communicates with neighbors . We identify the limitation of such adapt-then/while-communicate structure : it will make the developed algorithms highly sensitive to data heterogeneity , and hence deviate their limiting points from the desired solution . • To overcome these limitations , we propose a novel communicate-then-adapt algorithm structure , in which each node will conduct the adaptive step after all neighborhood communications . It is not just trivially switching the order of communication step and adaptive step . The key component to guarantee the effectiveness of the communicate-then-adapt structure is the utilization of the decentralized augmented gradient ( DAG ) rather than the standard stochastic gradient in algorithm development . The newly proposed algorithm , which is coined as DAG-Adam , can provably converge to the desired solution . While this paper mainly focuses on the vanilla Adam algorithm , the core idea behind DAG-Adam can be easilty extended to AMSGrad or other adaptive methods . • Experimental results on a variety of computer vision and natural language processing tasks show that DAG-Adam outperforms various existing state-of-the-art decentralized adaptive baselines under different network and data configurations . Furthermore , the performance of our proposed algorithm is persistently competitive to the centralized counterpart . Related work on decentralized optimization and deep training . Decentralized optimization was extensively studied in the control and signal processing community . The first decentralized algorithms on general optimization problems include decentralized gradient descent ( Nedic & Ozdaglar , 2009 ) , diffusion ( Lopes & Sayed , 2008 ; Chen & Sayed , 2012 ; Sayed , 2014 ) , and dual averaging ( Duchi et al. , 2011b ) . After that , various primal-dual algorithms come out to further speed up the convergence , and they are based on alternating direction method of multipliers ( Adam ) ( Shi et al. , 2014 ) , explicit bias-correction ( Shi et al. , 2015 ; Yuan et al. , 2019 ; Li et al. , 2019 ) , gradient tracking ( Xu et al. , 2015 ; Di Lorenzo & Scutari , 2016 ; Nedic et al. , 2017 ; Qu & Li , 2018 ; Lu et al. , 2019 ) , and dual acceleration ( Scaman et al. , 2017 ; Uribe et al. , 2020 ) . In deep learning tasks , decentralize SGD , which was established in ( Lian et al. , 2017 ) to achieve the same linear speedup as parallel SGD in convergence rate , has attracted a lot of attentions . Many efforts have been made to extend the algorithm to directed topologies ( Assran et al. , 2019 ) , time-varying topologies ( Koloskova et al. , 2020 ) , asynchronous settings ( Lian et al. , 2018 ) , and data-heterogeneous scenarios ( Tang et al. , 2018 ; Xin et al. , 2020 ; Lin et al. , 2021 ; Yuan et al. , 2021 ) . With careful consensus control ( Kong et al. , 2021 ) or periodic global averaging ( Chen et al. , 2021b ) , decentralize SGD can achieve 1.3 ∼ 2× training time speedup without severe performance degradation . Techniques such as quantization/compression ( Alistarh et al. , 2017 ; Bernstein et al. , 2018 ; Koloskova et al. , 2019a ; b ; Tang et al. , 2019 ; Liu et al. , 2020 ) , periodic updates ( Stich , 2019 ; Koloskova et al. , 2020 ; Yu et al. , 2019 ) , and lazy communication ( Chen et al. , 2018a ; Liu et al. , 2019b ) were also integrated into decentralized SGD to further reduce communication overheads . There are also a few studies on accelerated variants of decentralized SGD , and most of them are on ( static ) momentum acceleration . ( Assran et al. , 2019 ; Gao & Huang , 2020 ) propose to run a local momentum SGD step first before the partial averaging is conducted . Another work ( Yu et al. , 2019 ) imposes an additional partial averaging over momentum to increase stability . Recent works ( Lin et al. , 2021 ; Yuan et al. , 2021 ) developed strategies that can remove the momentum-incurred bias in decentralized momentum SGD . All these methods are not with adaptive strategies to scale gradients . Related work on adaptive gradient method . Adaptive gradient methods , with AdaGrad ( Duchi et al. , 2011a ) and Adam ( Kingma & Ba , 2014 ) as two representatives , have shown strong performance in training deep neural networks . With adaptive adjustment in the gradient direction and automatic tune in the learning rate , adaptive gradient methods can boost the performance of SGD training significantly . In spite of its remarkable empirical success , Adam suffers from a convergence issue : it may not converge to the desired solution with a fixed mini-batch size ( Reddi et al. , 2019 ; Zaheer et al. , 2018 ) . Many algorithms have been proposed to resolve this issue . AMSGrad ( Reddi et al. , 2019 ) preserves the long-term memory of past gradients to improve convergence , and ( Zaheer et al. , 2018 ) suggests the usage of increasing mini-batch sizes in convergence guarantees . ( Chen et al. , 2018b ) has studied the convergence of a family of Adam-type algorithms . ( Zhou et al. , 2018 ) provides a better convergence rate for AMSGrad . From the empirical side , RAdam ( Liu et al. , 2019a ) and AdamW ( Loshchilov & Hutter , 2018 ) can significantly improve Adam performance . The exploration in decentralized adaptive methods are rather limited . DAdam ( Nazari et al. , 2019 ) is the first consensus-based adaptive methods for distributed optimization to our knowledge . ( Lin et al. , 2021 ) proposes QG-DAdam , which utilizes quasi-global momentum to locally approximates the globally descent direction . A recent work ( Chen et al. , 2021a ) proposes a unified framework that incorporates various adaptive methods into decentralized setting . While these methods have shown strong empirical performance in several practical applications , they either suffer from unstable convergence or heavy communications . We leave more discussion on their limitations in Sec . 2 . Adaptive gradient methods have also been extended to the federated learning setting in which multiple clients cooperate to learn a model under the supervision of a central server ( McMahan et al. , 2017 ) . Useful references in this direction can be found in ( Xie et al. , 2019 ; Reddi et al. , 2020 ) . 2 ADAPT-THEN/WHILE-COMMUNICATE STRUCTURE AND ITS LIMITATION . Problem . Suppose n computing nodes cooperate to solve the distributed optimization problem : min x∈Rd f ( x ) = 1 n n∑ i=1 fi ( x ) , where fi ( x ) : = Eξi∼DiF ( x ; ξi ) ( 1 ) In the above problem , fi ( x ) is local to node i , and random variable ξi denotes the local data that follows distribution Di . We do not assume each distribution Di is the same across all nodes . Network topology and weights . We assume all computing nodes are connected by a network topology . We define wij , the weight to scale information flowing from node j to node i , as follows : wij { > 0 if node j is connected to i , or i = j ; = 0 otherwise . ( 2 ) Ni : = { j|wij > 0 } is defined as the set of neighbors of node i which also includes node i itself and the weight matrix W : = [ wij ] ni , j=1 ∈ Rn×n are denoted as a matrix that stacks the weights of all nodes . This matrix W characterizes the sparsity and connectivity of the underlying topology . Partial averaging . Decentralized methods are based on partial averaging within neighborhood defined by the network topology . With weights { wij } and the set of neighbors Ni , the partial averaging operation of node i can be expressed as Partial averaging : x+i ← ∑ j∈Ni wijxj . ( 3 ) Partial averaging requires much less communication than global averaging on sparse topologies . Assumptions . We will make the following standard assumptions throughout the paper : A.1 [ SMOOTHNESS ] Each fi ( x ) isL-smooth , i.e. , ‖∇fi ( x ) −∇fi ( y ) ‖ ≤ L‖x−y‖ for any x , y ∈ Rd . A.2 [ GRADIENT NOISE ] The random sample ξ ( t ) i is independent of each other for any k and i . We also assume E [ ∇F ( x ; ξi ) ] = ∇fi ( x ) and E‖∇F ( x ; ξi ) −∇fi ( x ) ‖2 ≤ σ2 . A.3 [ WEIGHT MATRIX ] The weight matrix W is symmetric and doubly-stochastic , i.e . W1 = 1 and 1TW = 1T , where 1 is an all-ones vector . Further , we assume ‖W − 1n11 T ‖2 ≤ ρ ∈ ( 0 , 1 ) . A.4 [ BOUNDED GRADIENT ] The loss function F has bounded gradient , i.e. , the maximum norm ‖∇F ( x ; ξi ) ‖∞ ≤ G for all x and ξi , and i ∈ [ n ] . Notation . Given a constant n , we let [ n ] : = { 1 , · · · , n } . Given a vector x ∈ Rd , we let diag ( x ) = diag { x1 , · · · , xd } ∈ Rd×d be a diagonal matrix . Given two vectors x ∈ Rd and y ∈ Rd , we let x y be the element-wise product between x and y .

Summary: This paper proposes a decentralized adaptive method for distributed deep learning, termed DAG-Adam. Convergence results are provided for smooth non-convex objectives under a bounded gradient assumption. Numerical experiments are conducted on Image Classification (CIFAR10, ImageNet-1k) and Language Modelling (fine-tuning pre-trained BERT models on SQuAD).