Title: YELLOWFIN AND THE ART OF MOMENTUM TUNING Anonymous authors Paper under double-blind review

Abstract: Hyperparameter tuning is one of the most time-consuming workloads in deep learning. State-of-the-art optimizers, such as AdaGrad, RMSProp and Adam, reduce this labor by adaptively tuning an individual learning rate for each variable. Recently researchers have shown renewed interest in simpler methods like momentum SGD as they may yield better results. Motivated by this trend, we ask: can simple adaptive methods, based on SGD perform as well or better? We revisit the momentum SGD algorithm and show that hand-tuning a single learning rate and momentum makes it competitive with Adam. We then analyze its robustness to learning rate misspecification and objective curvature variation. Based on these insights, we design YELLOWFIN, an automatic tuner for momentum and learning rate in SGD. YELLOWFIN optionally uses a negative-feedback loop to compensate for the momentum dynamics in asynchronous settings on the fly. We empirically show YELLOWFIN can converge in fewer iterations than Adam on ResNets and LSTMs for image recognition, language modeling and constituency parsing, with a speedup of up to 3.28x in synchronous and up to 2.69x in asynchronous settings.

1. INTRODUCTION:

Accelerated forms of stochastic gradient descent (SGD), pioneered by Polyak (1964) and Nesterov (1983) , are the de-facto training algorithms for deep learning. Their use requires a sane choice for their hyperparameters: typically a learning rate and momentum parameter (Sutskever et al., 2013) . However, tuning hyperparameters is arguably the most time-consuming part of deep learning, with many papers outlining best tuning practices written (Bengio, 2012; Orr and Müller, 2003; Bengio et al., 2012; Bottou, 2012) . Deep learning researchers have proposed a number of methods to deal with hyperparameter optimization, ranging from grid-search and smart black-box methods (Bergstra and Bengio, 2012; Snoek et al., 2012) to adaptive optimizers. Adaptive optimizers aim to eliminate hyperparameter search by tuning on the fly for a single training run: algorithms like AdaGrad (Duchi et al., 2011) , RMSProp (Tieleman and Hinton, 2012) and Adam (Kingma and Ba, 2014) use the magnitude of gradient elements to tune learning rates individually for each variable and have been largely successful in relieving practitioners of tuning the learning rate. Recently some researchers have started favoring simple momentum SGD over the previously mentioned adaptive methods (Chen et al., 2016; Gehring et al., 2017) , often reporting better test scores (Wilson et al., 2017) . Motivated by this trend, we ask the question: can simpler adaptive methods, based on momentum SGD perform as well or better? We empirically show that, with hand-tuned learning rate, Polyak's momentum SGD achieves faster convergence than Adam for a large class of models. We then formulate the optimization update as a dynamical system and study certain robustness properties of the momentum operator. Building on our analysis, we design YELLOWFIN, an automatic hyperparameter tuner for momentum SGD. YELLOWFIN simultaneously tunes the learning rate and momentum on the fly, and can handle the complex dynamics of asynchronous execution. Specifically:

• In Section 2, we show that momentum presents convergence robust to learning rate misspecification and curvature variation in a class of non-convex objectives; this robustness is desirable for deep learning. They stem from a known but obscure fact: the momentum operator's spectral radius is constant in a large subset of the hyperparameter space. • In Section 3, we use these robustness insights and a simple quadratic model analysis to design YELLOWFIN, an automatic tuner for momentum SGD. YELLOWFIN uses on-the-fly measurements from the gradients to tune both a single learning rate and momentum. • In Section 3.3, we discuss common stability concerns related to the phenomenon of exploding gradients (Pascanu et al., 2013) . We present a natural extension to our basic tuner, using adaptive gradient clipping, to stabilize training for objectives with exploding gradients. • In Section 4 we present closed-loop YELLOWFIN, suited for asynchronous training. It uses a novel component for measuring the total momentum in a running system, including any asynchrony-induced momentum, a phenomenon described in Mitliagkas et al. (2016) . This measurement is used in a negative feedback loop to control the value of algorithmic momentum.

We provide a thorough evaluation of the performance and stability of our tuner. In Section 5, we demonstrate empirically that on ResNets and LSTMs YELLOWFIN can converge in fewer iterations compared to: (i) hand-tuned momentum SGD (up to 1.75x speedup); and (ii) default Adam (0.8x to 3.3x speedup). Under asynchrony, the closed-loop control architecture speeds up YELLOWFIN, making it up to 2.69x faster than Adam. Our experiments include runs on 7 different models, randomized over at least 5 different random seeds. YELLOWFIN is stable and achieves consistent performance: the normalized sample standard deviation of test metrics varies from 0.05% to 0.6%.

We released PyTorch and TensorFlow implementations, that can be used as drop-in replacements for any optimizer. YELLOWFIN has also been implemented in other various packages. Its large-scale deployment in industry has taught us important lessons about stability; we discuss those challenges and our solution in Section 3.3. We conclude with related work and discussion in Section 6 and 7.

2. THE MOMENTUM OPERATOR:

In this section we identify the main technical insights guiding the design of YELLOWFIN. We show that momentum presents convergence robust to learning rate misspecification and curvature variation for a class of non-convex objectives; these robustness properties are desirable for deep learning.

2.1. PRELIMINARIES:

We aim to minimize some objective f (x). In machine learning, x is referred to as the model and the objective is some loss function. A low loss implies a well-fit model. Gradient descent-based procedures use the gradient of the objective function, rf (x), to update the model iteratively. Polyak's momentum gradient descent update (Polyak, 1964) is given by

EQUATION

where ↵ denotes the learning rate and µ the value of momentum used. Momentum's main appeal is its established ability to accelerate convergence (Polyak, 1964) . On a strongly convex smooth function with condition number , the optimal convergence rate of gradient descent (µ = 0) is O( 1 +1 ) (Nesterov, 2013) . On the other hand, for certain classes of strongly convex and smooth functions, like quadratics, the optimal momentum value,

EQUATION

yields the optimal accelerated rate O(

p  1 p +1

). 1 This is the smallest value of momentum that ensures the same rate of convergence along all directions. This fact is often hidden away in proofs. We shed light on some of its previously unknown implications in Section 2.2.

2.2. ROBUSTNESS PROPERTIES OF THE MOMENTUM OPERATOR:

In this section we analyze the dynamics of momentum on a simple class of one dimensional, nonconvex objectives. We first introduce the notion of generalized curvature and use it to describe the momentum operator. Then we discuss some properties of the momentum operator.

Definition 1 (Generalized curvature). The derivative of f (x) : R ! R, can be written as

f 0 (x) = h(x)(x x ⇤ )

(3) for some h(x) 2 R, where x ⇤ is the global minimum of f (x). We call h(x) the generalized curvature.

The generalized curvature describes, in some sense, curvature with respect to the optimum, x ⇤ . For quadratic objectives, it coincides with the standard definition of curvature, and is the sole quantity related to the objective that influences the dynamics of gradient descent. For example, the contraction of a gradient descent step is 1 ↵h(x t ). Let A t denote the momentum operator at time t. Using a state-space augmentation, we can express the momentum update as

EQUATION

)

Lemma 2 (Robustness of the momentum operator). As proven in Appendix A, if the generalized curvature, h, and hyperparameters ↵, µ are in the robust region, that is:

(1 p µ) 2  ↵h(x t )  (1 + p µ) 2 ,

(5) then the spectral radius of the momentum operator only depends on momentum:

⇢(A t ) = p µ.

We explain Lemma 2 as robust convergence with respect to learning rate and to variations in curvature. Momentum is robust to learning rate misspecification For a one dimensional strongly convex quadratic objective, we get h(x) = h for all x and Lemma 2 suggests that ⇢(A t ) = p µ as long as

EQUATION

In Figure 2 , we plot ⇢(A t ) for different ↵ and µ. As we increase the value of momentum, the optimal rate of convergence p µ is achieved by an ever-widening range of learning rates. Furthermore, for objectives with large condition number, higher values of momentum are both faster and more robust. This property influences the design of our tuner: as long as the learning rate satisfies (6), we are in the robust region and expect the same asymptotics, e.g. a convergence rate of p µ for quadratics, independent of the learning rate. Having established that, we can just focus on optimally tuning momentum.

Momentum is robust to curvature variation As discussed in Section 2.1, the intuition hidden in classic results is that for strongly convex smooth objectives, the momentum value in (2) guarantees the same rate of convergence along all directions. We extend this intuition to certain non-convex functions. Lemma 2 guarantees a constant, time-homogeneous spectral radius for the momentum operators (A t ) t if ( 5) is satisfied at every step. This motivates an extension of the condition number. Definition 3 (Generalized condition number). We define the generalized condition number (GCN) of a scalar function, f (x) : R ! R, to be the dynamic range of its generalized curvature, h(x):

EQUATION

The GCN captures variations in generalized curvature along a scalar slice. From Lemma 2 we get

EQUATION

as the optimal hyperparameters. Specifically, µ ⇤ is the smallest momentum value that guarantees a homogeneous spectral radius of p µ ⇤ for all (A t ) t . The spectral radius of an operator describes its asymptotic convergence behavior. However, the product of a sequence of operators A t • • • A 1 all with spectral radius p µ does not always follow the asymptotics of p µ t . In other words, we do not provide a convergence rate guarantee. Instead, we provide evidence in support of this intuition. For example, the non-convex objective in Figure 3 (a), composed of two quadratics with curvatures 1 and 1000, has a GCN of 1000. Using the tuning rule of (8), and running the momentum algorithm (Figure 3(b) ) yields a practically constant rate of convergence throughout. In Figures 3(c, d ) we demonstrate that for an LSTM, the majority of variables follow a p µ convergence rate. This property influences the design of our tuner: in the next section we use the tuning rules of (8) in YELLOWFIN, generalized appropriately to handle SGD noise. 

3. THE YELLOWFIN TUNER:

In this section we describe YELLOWFIN, our tuner for momentum SGD. We introduce a noisy quadratic model and work on a local quadratic approximation of f (x) to apply the tuning rule of (8) to SGD on an arbitrary objective. YELLOWFIN is our implementation of that rule.

Noisy quadratic model:

We consider minimizing the one-dimensional quadratic

EQUATION

)

The objective is defined as the average of n component functions, f i . This is a common model for SGD, where we use only a single data point (or a mini-batch) drawn uniformly at random,

S t ⇠ Uni([n]

) to compute a noisy gradient, rf St (x), for step t. Here, C = 1 2n P i hc 2 i denotes the gradient variance. This scalar model is sufficient to study an arbitrary local quadratic approximation: optimization on quadratics decomposes trivially into scalar problems along the principal eigenvectors of the Hessian. Next we get an exact expression for the mean square error after running momentum SGD on a scalar quadratic for t steps. Lemma 4. Let f (x) be defined as in (9), x 1 = x 0 and x t follow the momentum update (1) with stochastic gradients rf St (x t 1 ) for t 2. Let e 1 = [1, 0] T , the expectation of squared distance to the optimum x ⇤ is

EQUATION

) where the first and second term correspond to squared bias and variance, and their corresponding momentum dynamics are captured by operators

EQUATION

)

Even though it is possible to numerically work on (10) directly, we use a scalar, asymptotic surrogate based in (12) on the spectral radii of operators to simplify analysis and expose insights. This decision is supported by our findings in Section 2: the spectral radii can capture empirical convergence speed.

EQUATION

One of our design decisions for YELLOWFIN is to always work in the robust region of Lemma 2. We know that this implies a spectral radius p µ of the momentum operator, A, for the bias. Lemma 5

shows that under the exact same condition, the variance operator B has spectral radius µ. Lemma 5. The spectral radius of the variance operator,

B is µ, if (1 p µ) 2  ↵h  (1 + p µ) 2 .

As a result, the surrogate objective of ( 12), takes the following form in the robust region.

E(x t+1 x ⇤ ) 2 ⇡ µ t (x 0 x ⇤ ) 2 + (1 µ t ) ↵ 2 C 1 µ (13)

We use this surrogate objective to extract a noisy tuning rule for YELLOWFIN.

3.1. TUNING RULE:

Based on the surrogate in (13), we present YELLOWFIN (Algorithm 1). Let D denote an estimate of the current model's distance to a local quadratic approximation's minimum, and C denote an estimate for gradient variance. Also, let h max and h min denote estimates for the largest and smallest generalized curvature respectively. The extremal curvatures h min and h max are meant to capture both curvature variation along different directions (like the classic condition number) and also variation that occurs as the landscape evolves. At each iteration, we solve the following SINGLESTEP problem.

(SINGLESTEP)

µ t , ↵ t = arg min µ µD 2 + ↵ 2 C s.t. µ p h max /h min 1 p h max /h min + 1 ! 2 ↵ = (1 p µ) 2 h min Algorithm 1 YELLOWFIN function YELLOWFIN(gradient gt, ) hmax, hmin CURVATURERANGE(gt, ) C VARIANCE(gt, ) D DISTANCE(gt, ) µt, ↵t SINGLESTEP(C, D, hmax, hmin) return µt, ↵t end function

SINGLESTEP minimizes the surrogate for the expected squared distance from the optimum of a local quadratic approximation (13) after a single step (t = 1), while keeping all directions in the robust region (5). This is the SGD version of the noiseless tuning rule in (8). It can be solved in closed form; we refer to Appendix D for discussion on the closed form solution. YELLOWFIN uses functions CURVATURERANGE, VARIANCE and DISTANCE to measure quantities h max , h min , C and D respectively. These measurement functions can be designed in different ways. We present the implementations we used for our experiments, based completely on gradients, in Section 3.2.

3.2. MEASUREMENT FUNCTIONS IN YELLOWFIN:

This section describes our implementation of the measurement oracles used by YELLOWFIN: CUR-VATURERANGE, VARIANCE, and DISTANCE. We design the measurement functions with the assumption of a negative log-probability objective; this is in line with typical losses in machine learning, e.g. cross-entropy for neural nets and maximum likelihood estimation in general. Under this assumption, the Fisher information matrix-i.e. the expected outer product of noisy gradients-equals the Hessian of the objective (Duchi, 2016; Pascanu and Bengio, 2013) . This allows for measurements purely from minibatch gradients with overhead linear to model dimensionality. These implementations are not guaranteed to give accurate measurements. Nonetheless, their use in our experiments in Section 5 shows that they are sufficient for YELLOWFIN to outperform the state of the art on a variety of objectives. We also refer to Appendix E for implementation details on zero-debias (Kingma and Ba, 2014) , slow start (Schaul et al., 2013) 

g 2 • g 2 + (1 ) • gt gt g • g + (1 ) • gt return kg 2 g 2 k1 end function Algorithm 4 Distance to opt. state: kgk 0, h 0 function DISTANCE(gradient gt, ) kgk • kgk + (1 ) • kgtk h • h + (1 ) • kgtk 2 D • D + (1 ) • kgk/h return D end function

Curvature range Let g t be a noisy gradient, we estimate the range of curvatures in Algorithm 2. We notice that the outer product of g t and g T t has an eigenvalue h t = kg t k 2 with eigenvector g t . Thus under our negative log-likelihood assumption, we use h t to approximate the curvature of Hessian along gradient direction g t . Specifically, we maintain h min and h max as running averages of extreme curvature h min,t and h max,t , from a sliding window of width 20. As gradient directions evolve, we explore curvatures along different directions. Then the distance is approximated using kgk/h.

3.3. STABILITY ON NON-SMOOTH OBJECTIVES:

Neural network objectives can involve arbitrary non-linearities, and large Lipschitz constants (Szegedy et al., 2013) . Furthermore, the process of training them is inherently non-stationary, with the landscape abruptly switching from flat to steep areas. In particular, the objective functions of RNNs with hidden units can exhibit occasional but very steep slopes (Pascanu et al., 2013) . To deal with this issue, we use adaptive gradient clipping heuristics as a very natural addition to our basic tuner. It is discussed with extensive details in Appendix F. In Figure 4 , we present an example of an LSTM that exhibits the 'exploding gradient' issue. The proposed adaptive clipping can stabilize the training process using YELLOWFIN and prevent large catastrophic loss spikes.

We validate the proposed adaptive clipping on the convolutional sequence to sequence learning model (Gehring et al., 2017) for IWSLT 2014 German-English translation. The default optimizer (Gehring et al., 2017) uses learning rate 0.25 and Nesterov's momentum 0.99, diverging to loss overflow due to 'exploding gradient'. It requires, as in Gehring et al. (2017) , strict manually set gradient norm threshold 0.1 to stabilize. In Table 1 , we can see YellowFin, with adaptive clipping, outperforms the default optimizer using manually set clipping, with 0.84 higher validation BLEU4 after 120 epochs.

4. CLOSED-LOOP YELLOWFIN:

Asynchrony is a parallelization technique that avoids synchronization barriers (Niu et al., 2011) . It yields better hardware efficiency, i.e. faster steps, but can increase the number of iterations to a given metric, i.e. statistical efficiency, as a tradeoff (Zhang and Ré, 2014) . Mitliagkas et al. (2016) interpret asynchrony as added momentum dynamics. We design closed-loop YELLOWFIN, a variant of YELLOWFIN to automatically control algorithmic momentum, compensate for asynchrony and accelerate convergence. We use the formula in ( 14) to model the dynamics in the system, where the total momentum, µ T , includes both asynchrony-induced and algorithmic momentum, µ, in (1).

EQUATION

We first use ( 14) to design an robust estimator μT for the value of total momentum at every iteration.

Then we use a simple negative feedback control loop to adjust the value of algorithmic momentum so that μT matches the target momentum decided by YELLOWFIN in Algorithm 1. In Figure 5 , we demonstrate momentum dynamics in an asynchronous training system. As directly using the target value as algorithmic momentum, YELLOWFIN (middle) presents total momentum μT strictly larger than the target momentum, due to asynchrony-induced momentum. Closed-loop YELLOWFIN (right) automatically brings down algorithmic momentum, match measured total momentum μT to target value and, as we will see, significantly speeds up convergence comparing to YELLOWFIN. We refer to Appendix G for details on estimator μT and Closed-loop YELLOWFIN in Algorithm 5.

5. EXPERIMENTS:

In this section, we empirically validate the importance of momentum tuning and evaluate YELLOWFIN in both synchronous (single-node) and asynchronous settings. In synchronous settings, we first demonstrate that, with hand-tuning, momentum SGD is competitive with Adam, a state-of-the-art 

Asynchrony -induced momentum:

Figure 5 : When running YELLOWFIN, total momentum μt equals algorithmic value in synchronous settings (left); μt is greater than algorithmic value on 16 asynchronous workers (middle). Closed-loop YELLOWFIN automatically lowers algorithmic momentum and brings total momentum to match the target value (right). Red dots are measured μt at every step with red line as its running average. adaptive method. Then, we evaluate YELLOWFIN without any hand tuning in comparison to handtuned Adam and momentum SGD. In asynchronous settings, we show that closed-loop YELLOWFIN accelerates with momentum closed-loop control, performing significantly better than Adam.

We evaluate on convolutional neural networks (CNN) and recurrent neural networks (RNN). For CNN, we train ResNet (He et al., 2016) for image recognition on CIFAR10 and CIFAR100 (Krizhevsky et al., 2014) . For RNN, we train LSTMs for character-level language modeling with the TinyShakespeare (TS) dataset (Karpathy et al., 2015) , word-level language modeling with the Penn TreeBank (PTB) (Marcus et al., 1993) , and constituency parsing on the Wall Street Journal (WSJ) dataset (Choe and Charniak). We refer to Table 3 in Appendix H for model specifications. To eliminate influences of a specific random seed, in our synchronous and asynchronous experiments, the training loss and validation metrics are averaged from 3 runs using different random seeds.

5.1. SYNCHRONOUS EXPERIMENTS:

We tune Adam and momentum SGD on learning rate grids with prescribed momentum 0.9 for SGD. We fix the parameters of Algorithm 1 in all experiments, i.e. YELLOWFIN runs without any hand tuning. We provide full specifications, including the learning rate (grid) and the number of iterations we train on each model in Appendix I. For visualization purposes, we smooth training losses with a uniform window of width 1000. For Adam and momentum SGD on each model, we pick the configuration achieving the lowest averaged smoothed loss. To compare two algorithms, we record the lowest smoothed loss achieved by both. Then the speedup is reported as the ratio of iterations to achieve this loss. We use this setup to validate our claims. 9 in Appendix J. We can observe that momentum SGD achieves 1.71x and 1.87x speedup to tuned Adam on CIFAR10 and CIFAR100 respectively. In Figure 6 and Table 2 , with the exception of PTB LSTM, momentum SGD also produces better training loss, as well as better validation perplexity in language modeling and validation F1 in parsing. For the parsing task, we also compare with tuned Vanilla SGD and AdaGrad, which are used in the NLP community. Figure 6 (right) shows that fixed momentum 0.9 can already speedup Vanilla SGD by 2.73x, achieving observably better validation F1. We refer to Appendix J.2 for further discussion on the importance of momentum adaptivity in YELLOWFIN.

YELLOWFIN can match hand-tuned momentum SGD and can outperform hand-tuned Adam:

In our experiments, YELLOWFIN, without any hand-tuning, yields training loss matching hand-tuned momentum SGD for all the ResNet and LSTM models in Figure 6 In this section, we evaluate closed-loop YELLOWFIN with focus on the number of iterations to reach a certain solution. To that end, we run 16 asynchronous workers on a single machine and force them to update the model in a round-robin fashion, i.e. the gradient is delayed for 15 iterations. Figure 7 presents training losses on the CIFAR100 ResNet, using YELLOWFIN in Algorithm 1, closedloop YELLOWFIN in Algorithm 5 and Adam with the learning rate achieving the best smoothed loss in Section 5.1. We can observe closed-loop YELLOWFIN achieves 20.1x speedup to YELLOWFIN, and consequently a 2.69x speedup to Adam. This demonstrates that (1) closed-loop YELLOWFIN accelerates by reducing algorithmic momentum to compensate for asynchrony and (2) can converge in less iterations than Adam in asynchronous-parallel training.

6. RELATED WORK:

Many techniques have been proposed on tuning hyperparameters for optimizers. General hyperparameter tuning approaches, such as random search (Bergstra and Bengio, 2012) and Bayesian approaches (Snoek et al., 2012; Hutter et al., 2011) , directly applies to optimizer tuning. As another trend, adaptive methods, including AdaGrad (Duchi et al., 2011) , RMSProp (Tieleman and Hinton, 2012) and Adam (Chilimbi et al., 2014) , uses per-dimension learning rate. Schaul et al. (2013) use a noisy quadratic model similar to ours to extract learning rate tuning rule in Vanilla SGD. However they do not use momentum which is essential in training modern neural networks. Existing adaptive momentum approach either consider the deterministic setting (Graepel and Schraudolph, 2002; Rehman and Nawi, 2011; Hameed et al., 2016; Swanston et al., 1994; Ampazis and Perantonis, 2000; Qiu et al., 1992) or only analyze stochasticity with O(1/t) learning rate (Leen and Orr, 1994) .

In the contrast, we aim at practical momentum adaptivity for stochastically training neural networks.

7. DISCUSSION:

We presented YELLOWFIN, the first optimization method that automatically tunes momentum as well as the learning rate of momentum SGD. YELLOWFIN outperforms the state-of-the-art adaptive optimizers on a large class of models both in synchronous and asynchronous settings. It estimates statistics purely from the gradients of a running system, and then tunes the hyperparameters of momentum SGD based on noisy, local quadratic approximations. As future work, we believe that more accurate curvature estimation methods, like the bbprop method (Martens et al., 2012) can further improve YELLOWFIN. We also believe that our closed-loop momentum control mechanism in Section 4 could accelerate convergence for other adaptive methods in asynchronous-parallel settings.

A PROOF OF LEMMA 2:

To prove Lemma 2, we first prove a more generalized version in Lemma 6. By restricting f to be a one dimensional quadratics function, the generalized curvature h t itself is the only eigenvalue. We can prove Lemma 2 as a straight-forward corollary. Lemma 6 also implies, in the multiple dimensional correspondence of (4), the spectral radius ⇢(A t ) = p µ if the curvature on all eigenvector directions (eigenvalue) satisfies ( 5). Lemma 6. Let the gradients of a function f be described by

rf (x t ) = H (x t )(x t x ⇤ ), (15) with H (x t ) 2 R n 7 ! R n⇥n .

Then the momentum update can be expressed as a linear operator:

EQUATION

)

where y t , x t

x ⇤ . Now, assume that the following condition holds for all eigenvalues (H (x

EQUATION

)

then the spectral radius of A t is controlled by momentum with ⇢(A t ) = p µ.

Proof. Let t be an eigenvalue of matrix A t , it gives det (A t t I ) = 0. We define the blocks in 

A t as C = I ↵H t + µI t I , D = µI , E = I and F = t I which gives det (A t t I ) = det F det C DF 1 E = 0 assuming generally F is invertible. Note we use H t , H (x t )

(M t ) 2 = (1 ↵ (H t )+µ) 2  4µ which guarantees | t | = p µ for all t .

As the spectral radius is equal to the magnitude of the largest eigenvalue of A t , we have the spectral radius of A t being p µ.

B PROOF OF LEMMA 4:

We first prove Lemma 7 and Lemma 8 as preparation for the proof of Lemma 4. After the proof for one dimensional case, we discuss the trivial generalization to multiple dimensional case. Lemma 7. Let the h be the curvature of a one dimensional quadratic function f and x t = Ex t . We assume, without loss of generality, the optimum point of f is x ? = 0. Then we have the following recurrence ✓ x

EQUATION

Proof. From the recurrence of momentum SGD, we have

Ex t+1 =E[x t ↵rf St (x t ) + µ(x t x t 1 )] =E xt [x t ↵E St rf St (x t ) + µ(x t x t 1 )] =E xt [x t ↵hx t + µ(x t x t 1 )] =(1 ↵h + µ)x t µx t 1

By putting the equation in to matrix form, ( 20) is a straight-forward result from unrolling the recurrence for t times. Note as we set x 1 = x 0 with no uncertainty in momentum SGD, we have

[x 0 , x 1 ] = [x 0 , x 1 ]. Lemma 8. Let U t = E(x t x t ) 2 and V t = E(x t x t )(x t 1

x t 1 ) with x t being the expectation of x t . For quadratic function f (x) with curvature h 2 R, We have the following recurrence

EQUATION

where

EQUATION

and

C = E(rf St (x t ) rf (x t )) 2

is the variance of gradient on minibatch S t .

Proof. We prove by first deriving the recurrence for U t and V t respectively and combining them in to a matrix form. For U t , we have

U t+1 =E(x t+1 x t+1 ) 2 =E(x t ↵rf St (x t ) + µ(x t x t 1 ) (1 ↵h + µ)x t + µx t 1 ) 2 =E(x t ↵rf (x t ) + µ(x t x t 1 ) (1 ↵h + µ)x t + µx t 1 + ↵(rf (x t ) rf St (x t ))) 2 =E((1 ↵h + µ)(x t x t ) µ(x t 1 x t 1 )) 2 + ↵ 2 E(rf (x t ) rf St (x t )) 2 =(1 ↵h + µ) 2 E(x t x t ) 2 2µ(1 ↵h + µ)E(x t x t )(x t 1 x t 1 ) + µ 2 E(x t 1 x t 1 ) 2 + ↵ 2 C (23)

where the cross terms cancels due to the fact E St [rf (x t ) rf St (x t )] = 0 in the third equality. For V t , we can similarly derive 23) and ( 24) can be jointly expressed in the following matrix form

V t =E(x t x t )(x t 1 x t 1 ) =E((1 ↵h + µ)(x t 1 x t 1 ) µ(x t 2 x t 2 ) + ↵(rf (x t ) rf St (x t )))(x t 1 x t 1 ) =(1 ↵h + µ)E(x t 1 x t 1 ) 2 µE(x t 1 x t 1 )(x t 2 x t 2 ) (24) Again, the term involving rf (x t ) rf St (x t ) cancels in the third equality as a results of E St [rf (x t ) rf St (x t )] = 0. (

U t+1 U t V t+1 ! = B U t U t 1 V t ! + 0 @ ↵ 2 C 0 0 1 A = t 1 X i=0 B i 0 @ ↵ 2 C 0 0 1 A + B t U 1 U 0 V 1 ! = (I B t )(I B) 1 0 @ ↵ 2 C 0 0 1 A .

(25) Note the second term in the second equality is zero because x 0 and x 1 are deterministic. Thus

U 1 = U 0 = V 1 = 0.

According to Lemma 7 and 8, we have E(x

t x ⇤ ) 2 = (e > 1 A t [x 1 , x 0 ] > ) 2 and E(x t x t ) 2 = ↵ 2 Ce > 1 (I B t )(I B) 1 e

1 where e 1 2 R n has all zero entries but the first dimension. Combining these two terms, we prove Lemma 4. Though the proof here is for one dimensional quadratics, it trivially generalizes to multiple dimensional quadratics. Specifically, we can decompose the quadratics along the eigenvector directions, and then apply Lemma 4 to each eigenvector direction using the corresponding curvature h (eigenvalue). By summing quantities in (10) for all eigenvector directions, we can achieve the multiple dimensional correspondence of (10).

C PROOF OF LEMMA 5:

Again we first present a proof of a multiple dimensional generalized version of Lemma 5. The proof of Lemma 5 is a one dimensional special case of Lemma 9. Lemma 9 also implies that for multiple dimension quadratics, the corresponding spectral radius ⇢(B) = µ if

(1 p µ) 2 ↵  h  (1+ p µ) 2 ↵

on all the eigenvector directions with h being the eigenvalue (curvature). Lemma 9. Let H 2 R n⇥n be a symmetric matrix and ⇢(B) be the spectral radius of matrix

B = 0 @ (I ↵H + µI ) > (I ↵H + µI ) µ 2 I 2µ(I ↵H + µI ) I 0 0 I ↵H + µI 0 µI 1 A (26)

We have ⇢(B) = µ if all eigenvalues (H ) of H satisfies

EQUATION

)

Proof. Let be an eigenvalue of matrix B, it gives det (B I ) = 0 which can be alternatively expressed as det (B

EQUATION

)

assuming F is invertible, i.e. + µ 6 = 0, where the blocks in 28) can be transformed using straight-forward algebra as det

B C = ✓ M > M I µ 2 I I I ◆ , D = ✓ 2µM 0 ◆ , E = ✓ M 0 ◆ > , F = µI I with M = I ↵H + µI . (

EQUATION

Using similar simplification technique as in (28), we can further simplify into

EQUATION

if 6 = µ, as ( + µ) 2 I M > M is diagonalizable, we have ( + µ) 2 (M ) 2 = 0 with (M ) being an eigenvalue of symmetric M . The analytic solution to the equation can be explicitly expressed as

EQUATION

)

When the condition in ( 27) holds, we have (M ) 2 = (1 ↵ (H ) + µ) 2  4µ. One can verify that

( (M ) 2 2µ) 2 4µ 2 = ( (M ) 2 4µ) (M ) 2 = (1 ↵⇢(H ) + µ) 2 4µ (M ) 2  0 (32)

Thus the roots in (31) are conjugate with | | = µ. In conclusion, the condition in ( 27) can guarantee all the eigenvalues of B has magnitude µ. Thus the spectral radius of B is controlled by µ.

D ANALYTICAL SOLUTION TO (14):

The problem in ( 14) does not need iterative solver but has an analytical solution. Substituting only the second constraint, the objective becomes p(x) = x 2 D 2 + (1 x) 4 /h 2 min C with x = p µ 2 [0, 1). By setting the gradient of p(x) to 0, we can get a cubic equation whose root x = p µ p can be computed in closed form using Vieta's substitution. As p(x) is uni-modal in [0, 1), the optimizer for ( 14 

E PRACTICAL IMPLEMENTATION:

In Section 3.2, we discuss estimators for learning rate and momentum tuning in YELLOWFIN. In our experiment practice, we have identified a few practical implementation details which are important for improving estimators. Zero-debias is proposed by Kingma and Ba (2014) , which accelerates the process where exponential average adapts to the level of original quantity in the beginning. We applied zero-debias to all the exponential average quantities involved in our estimators. In some LSTM models, we observe that our estimated curvature may decrease quickly along the optimization process. In order to better estimate extremal curvature h max and h min with fast decreasing trend, we apply zero-debias exponential average on the logarithmic of h max,t and h min,t , instead of directly on h max,t and h min,t . Except from the above two techniques, we also implemented the slow start heuristic proposed by (Schaul et al., 2013) . More specifically, we use ↵ = min{↵ t , t • ↵ t /(10 • w)} as our learning rate with w as the size of our sliding window in h max and h min estimation. It discount the learning rate in the first 10 • w steps and helps to keep the learning rate small in the beginning when the exponential averaged quantities are not accurate enough.

F ADAPTIVE GRADIENT CLIPPING IN YELLOWFIN:

Gradient clipping has been established in literature as a standard-almost necessary-tool for training such objectives (Pascanu et al., 2013; Goodfellow et al., 2016; Gehring et al., 2017) . However, the classic tradeoff between adaptivity and stability applies: setting a clipping threshold that is too low can hurt performance; setting it to be high, can compromise stability. YELLOWFIN, keeps running estimates of extremal gradient magnitude squares, h max and h min in order to estimate a generalized condition number. We posit that p h max is an ideal gradient norm threshold for adaptive clipping. In order to ensure robustness to extreme gradient spikes, like the ones in Figure 4 , we also limit the growth rate of the envelope h max in Algorithm 2 as follows:

EQUATION

Our heuristics follows along the lines of classic recipes like Pascanu et al. (2013) . However, instead of using the average gradient norm to clip, it uses a running estimate of the maximum norm h max . In Section 3.3, we saw that adaptive clipping stabilizes the training on objectives that exhibit exploding gradients. In Figure 8 , we demonstrate that the adaptive clipping does not hurt performance on models that do not exhibit instabilities without clipping. Specifically, for both PTB LSTM and CIFAR10 ResNet, the difference between YELLOWFIN with and without adaptive clipping diminishes quickly.

G CLOSED-LOOP YELLOWFIN FOR ASYNCHRONOUS TRAINING:

In Section 4, we briefly discuss the closed-loop momentum control mechanism in closed-loop YELLOWFIN. In this section, after presenting more preliminaries on asynchrony, we show with details on the mechanism: it measures the dynamics on a running system and controls momentum with a negative feedback loop.

Preliminaries Asynchrony is a popular parallelization technique (Niu et al., 2011) that avoids synchronization barriers. When training on M asynchronous workers, staleness (the number of model updates between a worker's read and write operations) is on average ⌧ = M 1, i.e., the gradient in the SGD update is delayed by ⌧ iterations as rf St ⌧ (x t ⌧ ). Asynchrony yields faster steps, but can increase the number of iterations to achieve the same solution, a tradeoff between hardware and statistical efficiency (Zhang and Ré, 2014) . Mitliagkas et al. (2016) interpret asynchrony as added momentum dynamics. Experiments in Hadjis et al. (2016) support this finding, and demonstrate that reducing algorithmic momentum can compensate for asynchrony-induced momentum and significantly reduce the number of iterations for convergence. Motivated by that result, we use the model in ( 34), where the total momentum, µ T , includes both asynchrony-induced and algorithmic momentum, µ, in (1).

E[x

EQUATION

) We will use this expression to design an estimator for the value of total momentum, μT . This estimator is a basic building block of closed-loop YELLOWFIN, that removes the need to manually compensate for the effects of asynchrony.

Measuring the momentum dynamics Closed-loop YELLOWFIN estimates total momentum µ T on a running system and uses a negative feedback loop to adjust algorithmic momentum accordingly. Equation ( 14) gives an estimate of μT on a system with staleness ⌧ , based on ( 14).

EQUATION

)

We use ⌧ -stale model values to match the staleness of the gradient, and perform all operations in an elementwise fashion. This way we get a total momentum measurement from each variable; the median combines them into a more robust estimate.

Closing the asynchrony loop Given a reliable measurement of µ T , we can use it to adjust the value of algorithmic momentum so that the total momentum matches the target momentum as decided by YELLOWFIN in Algorithm 1. Closed-loop YELLOWFIN in Algorithm 5 uses a simple negative feedback loop to achieve the adjustment. 

H MODEL SPECIFICATION:

The model specification is shown in Table 3 for all the experiments in Section 5. CIRAR10 ResNet uses the regular ResNet units while CIFAR100 ResNet uses the bottleneck units. Only the convolutional layers are shown with filter size, filter number as well as the repeating count of the units. The layer counting for ResNets also includes batch normalization and Relu layers. The LSTM models are also diversified for different tasks with different vocabulary sizes, word embedding dimensions and number of layers.

I SPECIFICATION FOR SYNCHRONOUS EXPERIMENTS:

In Section 5.1, we demonstrate the synchronous experiments with extensive discussions. For the reproducibility, we provide here the specification of learning rate grids. The number of iterations as 

J.2 IMPORTANCE OF MOMENTUM ADAPTIVITY:

To further emphasize the importance of momentum adaptivity in YELLOWFIN, we run YF on CIFAR100 ResNet and TS LSTM. In the experiments, YELLOWFIN tunes the learning rate. Instead of also using the momentum tuned by YF, we continuously feed prescribed momentum value 0.0 and 0.9 to the underlying momentum SGD optimizer which YF is tuning. In Figure 10 , when comparing to YELLOWFIN with prescribed momentum 0.0 or 0.9, YELLOWFIN with adaptively tuned momentum achieves observably faster convergence on both TS LSTM and CIFAR100 ResNet. It empirically demonstrates the essential role of momentum adaptivity in YELLOWFIN. We conduct experiments on PTB LSTM with 16 asynchronous workers using Adam using the same protocol as in Section 5.2. Fixing the learning rate to the value achieving the lowest smoothed loss in Section 5.1, we sweep the smoothing parameter 1 (Kingma and Ba, 2014) of the first order moment estimate in grid { 0.2, 0.0, 0.3, 0.5, 0.7, 0.9}.

J.3 TUNING MOMENTUM CAN IMPROVE ADAM IN ASYNCHRONOUS-PARALLEL SETTING:

1 serves the same role as momentum in SGD and we call it the momentum in Adam. Figure 11 shows tuning momentum for Adam under asynchrony gives measurably better training loss. This result emphasizes the importance of momentum tuning in asynchronous settings and suggests that state-of-the-art adaptive methods can perform sub-optimally when using prescribed momentum.

:

This guarantee does not generalize to arbitrary strongly convex functions(Lessard et al.,

2016). Nonetheless, acceleration is often observed in practice (cf. Section 2.2).

Figures and Tables:

Figure 1: YELLOWFIN in comparison to Adam on a ResNet (CIFAR100, cf. Section 5).

Figure 2: Momentum operator on scalar quadratic.

Figure 3: (a) Non-convex toy example; (b) constant convergence rate achieved empirically on the objective of (a) tuned according to (8); (c,d) LSTM on MNIST: as momentum increases, more variables (shown in grey) fall in the robust region and follow the robust rate, p µ.

Figure 7: Asynchronous performance on CIFAR100 ResNet.

Figure 8: Training losses on PTB LSTM (left) and CIFAR10 ResNet (right) for YellowFin with and without adaptive clipping.

Figure 9: Training loss for ResNet on 100-layer CIFAR10 ResNet (left) and 164-layer CIFAR100 bottleneck ResNet.

Figure 10: Training loss comparison between YELLOWFIN with adaptive momentum and YEL-LOWFIN with fixed momentum value.

Figure 11: Hand-tuning Adam's momentum under asynchrony.

and smoothing for curvature range estimation.

Table content: <table><tr><td colspan="2">Algorithm 2 Curvature range</td><td>Algorithm 3 Gradient variance</td></tr><tr><td colspan="2">state: hmax, hmin, hi, 8i 2 {1, 2, 3, ...} function CURVATURERANGE(gradient gt, ) ht kgtk 2</td><td>state: g 2 function VARIANCE(gradient gt, ) 0, g 0</td></tr><tr><td colspan="2">hmax,t max t wit hi, hmin,t min t wit hi</td></tr><tr><td>hmax hmin return hmax, hmin • hmax + (1 • hmin + (1 end function</td><td>) • hmax,t ) • hmin,t</td></tr></table>

Gradient variance To estimate the gradient variance in Algorithm 3, we use running averages g and g 2 to keep track of g Distance to optimum In Algorithm 4, we estimate the distance to the optimum of the local quadratic approximation. Inspired by the fact that krf (x )k  kH kkx x ? k for a quadratic f (x) with Hessian H and minimizer x ⇤ , we first maintain h and kgk as running averages of curvature h

Table content: <table><tr><td>Gradient norm</td><td>0k 10 3 10 0 10 3 10 5</td><td>1k Without clipping 2k With clipping Clipping thresh.</td><td>3k</td><td>Training loss</td><td>10 2 0k 10 0 10 1 10 1</td><td>1k Without clipping 2k With clipping</td><td>3k</td><td>Default w/o clip. Default w/ clip. YF</td><td>Loss BLEU4 diverge 2.86 30.75 2.75 31.59</td></tr><tr><td/><td/><td>Iterations</td><td/><td/><td/><td>Iterations</td><td/><td/><td/></tr><tr><td colspan="8">Figure 4: A variation of the LSTM architecture in (Zhu et al.,</td><td colspan="2">Table 1: German-English transla-</td></tr><tr><td colspan="8">2016) exhibits exploding gradients. The proposed adaptive</td><td colspan="2">tion validation performance using</td></tr><tr><td colspan="8">gradient clipping threshold (blue) stabilizes the training loss.</td><td colspan="2">convolutional seq-to-seq learning.</td></tr></table>

Speedup of YELLOWFIN and tuned mom. SGD over tuned Adam.

Table content: <table><tr><td/><td colspan="2">Adam mom.SGD YF</td></tr><tr><td>CIFAR10</td><td>1x</td><td>1.71x 1.93x</td></tr><tr><td colspan="2">CIFAR100 1x</td><td>1.87x 1.38x</td></tr><tr><td>PTB</td><td>1x</td><td>0.88x 0.77x</td></tr><tr><td>TS</td><td>1x</td><td>2.49x 3.28x</td></tr><tr><td>WSJ</td><td>1x</td><td>1.33x 2.33x</td></tr></table>

for simplicity in writing. The equation det C DF 1 E = 0 implies that

Table content: <table><tr><td colspan="2">det 2 t I</td><td colspan="2">t M</td><td colspan="3">t + µI = 0</td><td>(18)</td></tr><tr><td>with M</td><td/><td/><td/><td/><td/><td/><td/></tr><tr><td>t =</td><td>(M</td><td>t ) ±</td><td>p</td><td>2</td><td>(M</td><td>t ) 2 4µ</td><td>(19)</td></tr><tr><td colspan="4">On the other hand, (17) guarantees that (1 ↵ (H</td><td/><td/><td/><td/></tr></table>

xt ⌧ xt ⌧ 1 +↵r S t ⌧ 1 f (xt ⌧ 1 ) xt ⌧ 1 xt ⌧ 2

Table content: <table><tr><td colspan="5">Algorithm 5 Closed-loop YELLOWFIN</td></tr><tr><td colspan="2">1: Input: µ 2: for t</td><td colspan="2">0, ↵ 1 to T do</td><td>0.0001,</td><td>0.01, ⌧ (staleness)</td></tr><tr><td>3: 4:</td><td colspan="2">x µ ⇤ , ↵ t x</td><td colspan="2">t 1 + µ(x YELLOWFIN(r t 1 x t 2 ) ↵r St f (x t ⌧ 1 ), ) St f (x t ⌧ 1 ) ⇣</td><td>⌘</td></tr><tr><td>5:</td><td>μT</td><td/><td>median</td><td>. Measuring total momentum</td></tr><tr><td colspan="2">6: 7: end for µ</td><td colspan="3">µ + • (µ ⇤</td><td>μT )</td><td>. Closing the loop</td></tr></table>

Peer review (iclr2018):
 [Apologies for short review, I got called in late. Marking my review as "educated guess" since i didn*t have time for a detailed review] The paper proposes an algorithm to tune the momentum and learning rate for SGD. While the algorithm does not have a theory for general non-quadratic functions, experimental validation is extensive, making it a worthy contribution in my opinion. I have personally tried the algorithm when the paper came out and can vouch for the empirical results presented here.