At Stability's Edge: How to Adjust Hyperparameters to Preserve Minima Selection in Asynchronous Training of Neural Networks?

1 INTRODUCTION . Training deep neural networks ( DNNs ) requires large amounts of computational resources , often using many devices operating over days and weeks . Furthermore , with ever-increasing model size and available data , the amount of compute used was noted to increase exponentially over the past few years ( Amodei & Hernandez , 2018 ) . Fortunately , the operations made in DNN training are highly suitable for parallelization over many devices . Most commonly , the parallelization is done by `` data-parallelism '' in which the data is divided into separate batches which are distributed between different devices during training . By using many devices , with each device highly utilized , one could train on a large number of samples without paying in run-time . However , this training is done synchronously , with synchronization between the different devices done on every iteration ( S-SGD ) . With the growth in the number of devices participating in the training , the synchronization becomes a prominent bottleneck . To overcome the synchronization bottleneck , asynchronous training ( A-SGD ) has been proposed ( Dean et al. , 2012 ; Recht et al. , 2011 ) . The basic idea behind such training is that when a device finishes calculating its batch gradients , an update to the DNN parameters is immediately performed , without waiting for other devices . This approach is known as asynchronous centralized training with a parameter server ( PS ) that updates the parameters ( Li , 2014 ) . The main problem in such training is that , since the PS updates the parameters whenever a device communicates with him , the parameters that are being used in other devices calculation are no longer up-to-date . This phenomena is called gradient staleness or delay as we will refer to it , and it causes a deterioration in generalization performance when using A-SGD . ∗Equal contribution The generalization deterioration can be seen in Fig . 1 . Unless the learning rate is significantly decreased , we can see a generalization gap : a deterioration in the generalization of A-SGD ( with delay τ = 32 ) from the baseline ( S-SGD , τ = 0 ) value near the steady state — i.e. , after we are near full convergence , following many training epochs ( 2000 ) . Due to the generalization gap demonstrated in the figure , A-SGD with a parameter server and large delays is not commonly used — despite it is relatively simple to implement and despite its vast potential to accelerate training . Therefore , our main goal in this work is to shed some theoretical light on the generalization gap problem and use it to improve A-SGD training . Previous theoretical analysis of A-SGD ( Liu et al. , 2018 ; Lian et al. , 2016 ; 2015 ; Dai et al. , 2018 ; Dutta et al. , 2018 ; Arjevani et al. , 2018 ) focused on analyzing the convergence rate , i.e. , the time it takes to reach steady state , and not on the properties of the obtained solution . In contrast , in our paper we focus on understanding how the delay affects the selection of the solution we converge to , and changes in this selection can impact generalization . We tackle these questions from the perspective of dynamical stability . The dynamical stability approach was used in ( Nar & Sastry , 2018 ; Wu et al. , 2018 ) , to study which minima are accessible under a specific choice of optimization algorithm and hyperparameters . In ( Nar & Sastry , 2018 ) , the authors analyzed gradient descent ( GD ) algorithm as a discrete-time nonlinear dynamical system . By analyzing the Lyapunov stability of this system for different minima , they showed a relation between the learning rate choice and the accessible minima set , i.e. , the subset of the local optima that GD can converge to . In ( Wu et al. , 2018 ) , the authors focused on stochastic gradient descent ( SGD ) , defined a criterion to evaluate the stability of a specific minimum and used this criterion to show how the learning rate and batch-size play a role in SGD minima selection process . Here , we use dynamical stability to analyze the dynamics of A-SGD . Using this approach , the main question we try to tackle is : How do learning rate , delay and momentum interact and affect the minima selection process ? Contributions We start by modelling A-SGD as a dynamical system with delay . By analyzing the stability properties of that system , we find : • There exist an inverse linear relation between the delay of A-SGD , and the threshold learning rate i.e. , the critical learning rate in which a minimum loses stability . • This implies that in order to keep a specific minimum stability as we change the delay , we need to change the learning rate inversely proportional to the delay . • Momentum has a crucial role in determining stability properties for A-SGD . Depending on the specific training algorithm , it can either hurt the stability or improve it . We show how to modify momentum to stabilize A-SGD training . With our theoretical findings , we derive closed form rules when training asynchronously . We suggest to scale the learning rate inversely to the delay in order to reach a better generalization performance . In section 3 we provide experiments on popular classification tasks to support our theoretical analysis . It is interesting to contrast this simple linear relation with the more complicated picture observed for the scaling of the learning rate with the batch size . Specifically , it has been observed by Shallue et al . ( 2018 ) that the optimal learning rate scales differently with the batch size when changing the models and datasets , i.e. , there is no general `` rule of thumb '' that applies in all cases . A similar scaling rule for the learning rate was proposed before ( Zhang et al. , 2015 ) , however , there it lacked a theoretical basis or a demonstration of its effectiveness with large effective batch sizes , i.e. , the sum of all the minibatch sizes of all the workers , as we do here . 2 THEORETICAL ANALYSIS . 2.1 PRELIMINARIES AND PROBLEM SETUP . Consider the problem of minimizing the empirical loss f ( x ) = N∑ i=1 fi ( x ) , ( 1 ) where N is the number of samples , x ∈ Rd and ∀i = 1 , ... , N : fi : Rd → R is twice continuously differentiable function . The update rule for minimizing f ( x ) in eq . 1 using asynchronous stochastic gradient descent ( A-SGD ) is given by xt+1 = xt − η∇fn ( t−τ ) ( xt−τ ) , ( 2 ) where η is the learning rate , n ( t ) is a random selection process of a sample from { 1 , 2 , . . . , N } at iteration t , and τ is some delay due the asynchronous nature of our training . For simplicity , we focus on a fixed delay τ . The extension to stochastic delay is discussed in section 2.2.1 . Our main goal is understanding how different factors such as the gradient staleness τ and the learning rate η affect the minima selection process in neural networks . In other words , we would like to understand how the different hyperparameters interact to divide the minima of the loss to two sets : those we can converge to and those we can not converge to . To analyze this we suppose we are in a vicinity of some minimum point x∗ , which implies∇f ( x∗ ) = 0 . Given some values of η and τ , we ask whether or not we can converge to x∗ , using the theory of dynamical stability . As we shall see , if the learning rate η or delay τ are too high , then this minimum loses stability , i.e. , A-SGD can not converge to x∗ , in expectation . Specifically , we examine the expectation of the linearized dynamics around x∗ . We show in appendix A that , to ensure stability , it is sufficient that the following one-dimensional equation is stable : xt+1 − xt + ηaxt−τ = 0 , ( 3 ) where we defined the Hessian H = 1N N∑ i=1 ∇2fi ( x∗ ) , the sharpness term a = λmax ( H ) as the maximal eigenvalue of the Hessian , and xt is the projection of the expectation of the perturbation xt − x∗ on the eigenvector which corresponds to the maximal eigenvalue a. Additionally , we show in appendix A that the characteristic equation of eq . 3 is zτ+1 − zτ + ηa = 0 . ( 4 ) 2.2 STABILITY ANALYSIS . For given ( η , τ ) , we would like to determine if the minimum point x∗ is asymptotically stable ( Oppenheim , 1999 ) in expectation . This will enable us to divide the minima into two sets : those we might converge to since they are stable , and those we can not possibly converge ( except for a measure zero set of initialization ) since they are unstable . We use the characteristic equation ( eq . 4 ) to define the stability criterion : Definition 1 ( Stability criterion ) For given ( η , τ ) , we say that the dynamics in eq . 3 are stable if all the roots of the corresponding characteristic equation ( eq . 4 ) are inside the unit circle . This condition is equivalent to xt → 0 for any initialization . We would like to find conditions on when and how can we change the hyperparameters ( η , τ ) to maintain the same stability criterion . This , for example , will enable us to understand how to compensate for the delay introduced by asynchrony using the learning rate . In order to analyze stability we examine the characteristic equation in eq . 4 and ask : when does the maximal root of this polynomial have unit magnitude ? 2.2.1 THE INTERACTION BETWEEN LEARNING RATE AND DELAY . Recall the characteristic equation zτ+1 − zτ + aη = 0 . ( 5 ) This polynomial has τ +1 roots . In order to ensure stability we require that the root with the maximal magnitude , i.e. , the maximal root , will be inside the unit circle . We want to find the threshold learning rate , i.e. , the learning rate in which the maximal root is exactly on the unit circle , so we are at the threshold of stability . In appendix B , we show that this requires that aη = 2 sin ( π 4τ + 2 ) . ( 6 ) Using Taylor approximation we get aη = 2 ( π 4τ + 2 +O ( 1 τ3 ) ) ⇒ 1 aη = 1 π ( 2τ + 1 ) +O ( 1 τ ) . ( 7 ) We observe numerically that the error between the exact solution ( eq . 6 ) and the linear approximation ( eq . 7 ) is smaller than 0.05 for τ ≥ 1 . Additionally , in the appendix Fig . 6 we demonstrate the high accuracy of the analytic approximation of the relation between 1aη and τ , compared to the numerical evaluation of the value of 1aη for which the maximal root of eq . 5 is on the unit circle . Implications : We can see that to maintain stability for a given minimum point ( or a given sharpness value ) the learning rate should be kept inversely proportional to the delay . This implies that given some minimum x∗ for τ = 0 and a corresponding learning rate η0 for which this minimum is stable , we can evaluate the learning rate which will ensure this minimum remain stable for larger delay values . To do this , we first estimate a using eq . 6 : aη0 = 2 sin ( π 4·0+2 ) = 2⇒ a = 2η0 . Note that we do not use eq . 7 to evaluate a since we are interested in τ = 0 , while the relation in eq . 7 is only a good approximation for τ ≥ 1 . Next , given some delay τ , we substitute a into eq . 7 in order to evaluate a learning rate value for which the stability of all minima is maintained as in τ = 0 η ≈ π a ( 2τ + 1 ) = πη0 4 ( τ + 0.5 ) . ( 8 ) Related empirical results : In section 3.1 we show empirically the existence of a stability threshold for different values of delay and compare it with our theoretical threshold . Stochastic delay : In appendix C we extend our theoretical analysis to derive the characteristic equation for A-SGD with stochastic delay , i.e. , when τ is drawn from some discrete distribution . We analyze two cases : • Discrete uniform distribution : τ ∼ unif { a , b } . • Gaussian distribution discrete approximation : ∀µ , ∀k = 1 , ... , 2µ : Pr ( τ = k ) = Pr ( k − 0.5 ≤ z ≤ k + 0.5 ) where z ∼ N ( µ , 1 ) . In Zhang et al . ( 2015 ) , the authors showed empirically a delay distribution that resembles to Gaussian distribution . For both cases , we show numerically that the inverse relation between the learning rate η and the expected delay Eτ still applies .

Summary: The authors introduce a theoretical model for delayed gradients in asynchronous training. It is a very nice model and solving the corresponding differential equation allows to study its stability. Authors derive stability bounds for pure SGD (learning rate needs to decrease with delay) and for SGD with momentum, where they introduce a nice momentum formulation that improves stability. These are nice insights and good results and they are validated by experiments. More experiments and practical analysis would be welcome though. Some example questions: would introducing some sychronization help? Is the lower learning rate hurting training speed when measures as wall-clock time to accuracy?