Year: 2018

Abstract:
We propose a fast second-order method that can be used as a drop-in replacement for current deep learning solvers. Compared to stochastic gradient descent (SGD), it only requires two additional forward-mode automatic differentiation operations per iteration, which has a computational cost comparable to two standard forward passes and is easy to implement. Our method addresses long-standing issues with current second-order solvers, which invert an approximate Hessian matrix every iteration exactly or by conjugate-gradient methods, procedures that are much slower than a SGD step. Instead, we propose to keep a single estimate of the gradient projected by the inverse Hessian matrix, and update it once per iteration with just two passes over the network. This estimate has the same size and is similar to the momentum variable that is commonly used in SGD. No estimate of the Hessian is maintained. We first validate our method, called CURVEBALL, on small problems with known solutions (noisy Rosenbrock function and degenerate 2-layer linear networks), where current deep learning solvers struggle. We then train several large models on CIFAR and ImageNet, including ResNet and VGG-f networks, where we demonstrate faster convergence with no hyperparameter tuning. We also show our optimiser’s generality by testing on a large set of randomly-generated architectures.

heading: 1 INTRODUCTION
text: Stochastic Gradient Descent (SGD) and back-propagation (LeCun et al., 1998) are the algorithmic backbone of current deep network training. The success of deep learning demonstrates the power of this combination, which has been successfully applied on various tasks with large datasets and very deep networks (He et al., 2016).
Yet, while SGD has many advantages, speed of convergence (in terms of number of iterations) is not necessarily one of them. While individual SGD iterations are very quick to compute and lead to rapid progress at the beginning of the optimisation, it soon reaches a slower phase where further improvements are achieved slowly. This can be attributed to entering regions of the parameter space where the objective function is poorly scaled. In such cases, rapid progress would require vastly different step sizes for different directions in parameter space, which SGD cannot deliver.
Second-order methods, such as Newton’s method and its variants, eliminate this issue by rescaling the gradient according to the local curvature of the objective function. For a scalar loss in R, this rescaling takes the form H−1J where H is the Hessian matrix (second-order derivatives) or an approximation of the local curvature in the objective space, and J is the gradient of the objective. They can in fact achieve local scale-invariance (Wright & Nocedal, 1999, p. 27), and make provably better progress in the regions where gradient descent stalls. While they are unmatched in other domains, there are several obstacles to their application to deep models. First, it is impractical to invert or even store the Hessian, since it grows quadratically with the number of parameters, and there are typically millions of them. Second, any Hessian estimate is necessarily noisy and ill-conditioned due to stochastic sampling, to which classic inversion methods such as conjugate-gradient are not robust.
In this paper, we propose a new algorithm that can overcome these difficulties and make second order optimisation practical for deep learning. We show in particular how to avoid the storage of any estimate of the Hessian matrix or its inverse. Instead, we treat the computation of the Newton update, H−1J , as solving a linear system that itself can be solved via gradient descent. The cost of solving this system is amortized over time by interleaving its steps with the parameter update steps. Our
proposed method adds little overhead, since a Hessian-vector product can be implemented for modern networks with just two steps of automatic differentiation. Interestingly, we show that our method is equivalent to momentum SGD (also known as the heavy-ball method) with a single additional term, accounting for curvature. For this reason we named our method CURVEBALL. Unlike other proposals, the total memory footprint is as small as that of momentum SGD.
This paper is structured as follows. We introduce relevant technical background in sec. 2, and present our method in sec. 3. We evaluate our method and show experimental results in sec. 4. Related work is discussed in sec. 5. Finally we summarise our findings in sec. 6.

heading: 2 BACKGROUND
text: In order to make the description of our method self-contained, we succinctly summarise a few standard concepts in optimisation. Our goal is to find the optimal parameters of a model (e.g. a neural network) φ : Rp → Ro, with p parameters w ∈ Rp and o outputs (the notation does not show the dependency on the training data, which is subsumed in φ for compactness). The quality of the outputs is evaluated by a loss function L : Ro → R, so finding w is reduced to the optimisation problem: 1
w∗ = arg min w L(φ(w)) = arg min w f(w). (1)
Perhaps the simplest algorithm to find an optimum (or at least a stationary point) of eq. 1 is gradient descent (GD). GD updates the parameters using the iteration w ← w − βJ(w), where β > 0 is the learning rate and J(w) ∈ Rp is the gradient (or Jacobian) of the objective function f with respect to the parameters w. A useful variant is to augment GD with a momentum variable z (Polyak, 1964), which can be interpreted as a decaying average of past gradients:2
z ← ρz − βJ(w) (2) w ← w + z (3)
with a momentum parameter ρ. Momentum GD, as given by eq. 2-3, can be shown to have faster convergence than GD for convex functions, remaining stable under higher learning rates, and exhibits somewhat better resistance to poor scaling of the objective function (Nesterov, 2013; Goh, 2017). One important aspect is that these advantages cost almost no additional computation and only a modest additional memory, which explains why it is widely used in practice.
In neural networks, GD is usually replaced by its stochastic version (SGD), where at each iteration one computes the gradient not of the model f = L(φ(w)), but of the model ft = Lt(φt(w)) assessed on a small batch of samples, drawn at random from the training set.

heading: 2.1 SECOND-ORDER OPTIMISATION
text: As mentioned in section 1, the Newton method is similar to GD, but steers the gradient by the inverse Hessian matrix, computing H−1J as a descent direction. However, inverting the Hessian may be numerically unstable or the inverse may not even exist. To address this issue, the Hessian is usually regularized with a parameter λ > 0, obtaining what is known as the Levenberg (Moré, 1978) method:
z = −(H + λI)−1J, (4) w ← w + z, (5)
where H ∈ Rp×p, J ∈ Rp and I ∈ Rp×p is the identity matrix. Note that, unlike for momentum GD (eq. 2), the new step z is independent of the previous step. To avoid burdensome notation, we omit the w argument in H(w) and J(w), but they must be recomputed at each iteration. Intuitively, the effect of eq. 4 is to rescale the step appropriately for different directions — directions with high curvature require small steps, while directions with low curvature require large steps to make progress.
Note also that Levenberg’s regularization loses the scale-invariance of the original Newton method, meaning that rescaling the function f changes the scale of the gradient and hence the regularised descent direction chosen by the method. An alternative that alleviates this issue is LevenbergMarquardt, which replaces I in eq. 4 with diag(H). For non-convex functions such as deep networks, these methods only converge to a local minimum when the Hessian is positive-semidefinite (PSD).
1We omit the optional regulariser term for brevity, but this does not materially change our derivations. 2Some sources use an alternative form with β in eq. 3 instead (equivalent by a re-parametrisation of z),

heading: 2.2 AUTOMATIC DIFFERENTIATION AND BACK-PROPAGATION
text: In order to introduce fast computations involving the Hessian, we must take a short digression into how Jacobians are computed. The Jacobian of L(φ(w)) (eq. 1) is generally computed as J = JφJL where Jφ ∈ Rp×o and JL ∈ Ro×1 are the Jacobians of the model and loss, respectively. In practice, a Jacobian is never formed explicitly, but Jacobian-vector products Jv are implemented with the back-propagation algorithm. We define
←− AD(v) = Jv (6)
as the reverse-mode automatic differentiation (RMAD) operation, commonly known as backpropagation. Note that, because the loss is a scalar function, the starting projection vector v typically used in gradient descent is a scalar and we set v = 1. For intermediate computations, however, it is generally a (vectorized) tensor of gradients (see eq. 9).
A perhaps lesser known alternative is forward-mode automatic differentiation (FMAD), which computes a vector-Jacobian product, from the other direction:
−→ AD(v′) = v′J (7)
This variant is less commonly-known in deep learning as RMAD is appropriate to compute the derivatives of a scalar-valued function, such as the learning objective, whereas FMAD is more appropriate for vector-valued functions of a scalar argument. However, we will show later that FMAD is relevant in calculations involving the Hessian.
The only difference between RMAD and FMAD is the direction of associativity of the multiplication: FMAD propagates gradients in the forward direction, while RMAD (or back-propagation) does it in the backward direction. For example, for the composition of functions a ◦ b ◦ c,
−→ ADa◦b◦c(v) = ((vJa) Jb) Jc (8) ←− ADa◦b◦c(v ′) = Ja (Jb (Jcv ′)) (9)
Because of this, both operations have similar computational overhead, and can be implemented similarly.

heading: 2.3 FAST HESSIAN-VECTOR PRODUCTS
text: Since the Hessian of learning objectives involving deep networks is not necessarily positive semidefinite (PSD), it is common to use a surrogate matrix with this property, which prevents second-order methods from being attracted to saddle-points (this problem is discussed by Dauphin et al. (2014)). One of the most widely used is the Gauss-Newton approximation (Martens & Grosse, 2015; Botev et al., 2017; Wright & Nocedal, 1999, p. 254):
Ĥ = JφHLJ T φ , (10)
When HL is PSD, which is the case for all convex losses (e.g. logistic loss, Lp distance), the resulting Ĥ is PSD by construction. For the method that we propose, and indeed for any method that implicitly inverts the Hessian (or its approximation), only computing Hessian-vector products Ĥv is required. As such, eq. 10 takes a very convenient form:
Ĥv = Jφ ( HL ( JTφ v )) (11)
= ←− ADφ ( HL (−→ ADφ(v) )) . (12)
The cost of eq. 12 is thus equivalent to that of two back-propagation operations. This is similar to a classic result (Pearlmutter, 1994; Schraudolph, 2002), but written in terms of common automatic differentiation operations. The intermediate matrix-vector product HLu has negligible cost: for example, for the squared-distance loss, HL = 2I ⇒ HLu = 2u. Similarly, for the multinomial logistic loss we have HL = diag(p) − ppT ⇒ HLu = p u − p(pTu), where p is the vector of predictions from a softmax layer and is the element-wise product. These products thus require only element-wise operations.
Algorithm 1: CURVEBALL (proposed). 1: z0 = 0 2: for t = 0, ..., T − 1 do 3: ∆z = Ĥ(wt)zt + J(wt) 4: zt+1 = ρzt − β∆z 5: wt+1 = wt + zt+1 6: end for
Algorithm 2: Simplified Hessian-free method. 1: for t = 0, ..., T − 1 do 2: z0 = −J(wt) 3: for r = 0, ..., R− 1 (or convergence) do 4: zr+1 = CG(zr, Ĥ(wt)zr, J(wt)) 5: end for 6: wt+1 = wt + zR 7: end for

heading: 3 METHOD
text: This section presents our main contribution: a method that minimizes a second-order Taylor expansion of the objective (like the Newton variants from section 2.1), but at a much reduced computational and memory cost, suitable for very large-scale problems. The result of taking a step z away from a starting point w can be modelled with a second-order Taylor expansion of the objective f :
f(w + z) ' f̂(w, z) = f(w) + zTJ(w) + 12z TH(w)z (13)
Most second-order methods seek the update z that minimizes f̂ , by ignoring the higher-order terms:
z = arg min z′ f̂(z′) = arg min z′ 1 2z ′T Ĥz′ + z′TJ (14)
In general, a step z is found by minimizing eq. 14, either via explicit inversion Ĥ−1J (Martens & Grosse, 2015; Botev et al., 2017) or the conjugate-gradient (CG) method (Martens, 2010). The later approach, called the Hessian-free method (also Truncated Newton or Newton-CG (Wright & Nocedal, 1999, p. 168)) is the most economical in terms of memory, since it only needs access to Hessian-vector products (section 2.3). A high-level view is illustrated in Algorithm 2, where CG stands for one step of conjugate-gradient (a stopping condition, line search and some intermediate variables were omitted for clarity). Note that for every update of w (outer loop), Algorithm 2 must perform several steps of CG (inner loop) to find a single search direction z.
We propose a number of changes in order to eliminate this costly inner loop. The first is to reuse the previous search direction z to warm-start the inner iterations, instead of resetting z each time (Algorithm 2, line 2). If z does not change abruptly, then this should help reduce the number of CG iterations, by starting closer to the solution. The second change is to use this fact to dramatically reduce the inner loop iterations to just one (R = 1). A different interpretation is that we now interleave the updates of the search direction z and parameters w (Algorithm 1, lines 4 and 5), instead of nesting them (Algorithm 2, lines 4 and 6).
Unfortunately, this change loses the guarantees of the CG method, which depend on the starting point being z0 = −J(wt) (Wright & Nocedal, 1999, p. 124). This loss of guarantee was verified in practice, as we found the resulting algorithm extremely unstable. Our third change is then to replace CG with gradient descent, which has no such dependency. Differentiating eq. 14 w.r.t. z yields:
4z = Jf̂(z) = Ĥz + J (15)
Applying these changes to the Hessian-free method (Algorithm 2) results in Algorithm 1. By contrasting it to momentum GD (eq. 2-3), we can see that it is equivalent, except for an extra curvature term Ĥ(w)z. In order to establish this equivalence, we introduced a factor ρ that decays z each step (Algorithm 1, line 4), whereas a pure gradient descent step on z would not include this factor. We can obtain it formally by simply regularizing the quadratic model in eq. 13 with the term (1 − ρ) ‖z‖2, which is small when ρ ' 1 (the recommended setting for the momentum parameter (He et al., 2016)). Due to the addition of curvature to momentum GD, which is also known as the heavy-ball method, we name our algorithm CURVEBALL.
Implementation. Using the fast Hessian-vector products from section 2.3, it is easy to implement eq. 15, including a regularization term λI (section 2.1). We can further improve eq. 15 by grouping
In this way, the total number of passes over the model is two: we compute Jφv and JTφ v ′ products, implemented respectively as one RMAD (back-propagation) and one FMAD operation (section 2.2).
Automatic ρ and β hyper-parameters in closed form. Our proposed method introduces a few hyper-parameters, which just like with SGD, would require tuning for different settings. Ideally, we would like to have no tuning at all. Fortunately, the quadratic minimization interpretation in eq. 14 allows us to draw on standard results in optimisation. At any given step, the optimal ρ and β can be obtained by solving a 2× 2 linear system (Martens & Grosse, 2015, sec. 7):[
β −ρ
] = [ ∆Tz Ĥ∆z z
T Ĥ∆z zT Ĥ∆z z T Ĥz ]−1 [ JT∆z JT z ] (18)
Note that, in calculating the proposed update (eq. 16), the quantities ∆z , JTφ z and JL have already been computed and can now be reused. Together with the fact that Ĥ = JφHLJTφ , this means that the elements of the above 2× 2 matrix can be computed with only one additional forward pass.
Automatic λ hyper-parameter rescaling. The regularization term λI (eq. 4) can be interpreted as a trust-region (Wright & Nocedal, 1999, p. 68). When the second-order approximation holds well, λ can be small, corresponding to an unregularized Hessian and a large trust-region. Conversely, a poor fit requires a correspondingly large λ. We can measure the difference (or ratio) between the objective change predicted by the quadratic fit (f̂ ) and the real objective change (f ), by computing γ = (f(w + z)− f(w)) /f̂(z). This requires one additional evaluation of the objective for f(w+z), but otherwise relies only on previously computed quantities. This makes it a very attractive estimate of the trust region, with γ = 1 corresponding to a perfect approximation. Following (Wright & Nocedal, 1999, p. 69), we evaluate γ every 5 iterations, decreasing λ by a factor of 0.999 when γ > 3/2, and increasing by the inverse factor when γ < 1/2. We noted that our algorithm is not very sensitive to the initial λ. In experiments using batch-normalization (section 4), we simply initialize it to one, otherwise setting it to 10. We plot the evolution of automatically tuned hyper-parameters in fig. 8 (Appendix B).
Convergence. In addition to the usual absence of strong guarantees for non-convex problems, which applies in our setting (deep neural networks), there is an added difficulty due to the recursive nature of our algorithm (the interleaved w and z steps). Our method is a variant of the heavy-ball method (Lessard et al., 2016; Flammarion & Bach, 2015) (by adding a curvature term), which until very recently had resisted establishing global convergence rates that improve on gradient descent without momentum (Loizou & Richtárik (2017), table 1), and even then only for strongly convex or quadratic functions.
For this reason, we present proofs for two more tractable cases (Appendix A). The first is the global linear convergence of our method for convex quadratic functions, which allows a direct inspection
of the region of convergence as well as its rate (Theorem A.1). The second establishes that, for convex non-quadratic functions, CURVEBALL’s steps are always in a descent direction, when using the automatic hyper-parameter tuning of eq. 18 (Theorem A.2). We note that in practice, due to the Gauss-Newton approximation and the trust region (eq. 4), the effective Hessian is guaranteed to be positive-definite.
Similarly to momentum SGD, our main claim as to the method’s suitability for non-convex deep network optimisation is necessarily empirical, based on the extensive experiments in section 4, which show strong performance on several large-scale problems with no hyper-parameter tuning.

heading: 4 EXPERIMENTS
text: Degenerate problems with known solutions. While the main purpose of our optimizer is its application to large-scale deep learning architecture, we begin by applying our method to problems of limited complexity, with the goal of exploring the strengths and weaknesses of our approach in an interpretable domain. We perform a comparison with two popular first order solvers — SGD with momentum and Adam (Kingma & Ba, 2014)3, as well as with more traditional methods such as Levenberg-Marquardt, BFGS (Wright & Nocedal, 1999, p. 136) (with cubic line-search) and Newton’s method with the exact Hessian. The first problem we consider is the search for the minimum of the two-dimensional Rosenbrock test function, which has the useful benefit of enabling us to visualise the trajectories found by each optimiser. Specifically, we use the stochastic variant of this function (Yang & Deb, 2010),R : R2 → R:
R(u, v) = (1− u)2 + 100 i(v − u2)2, (19)
where at each evaluation of the function, a noise sample i is drawn from a uniform distribution U [λ1, λ2] with λ1, λ2 ∈ R (we can recover the deterministic Rosenbrock function with λ1 = λ2 = 1). To assess robustness to noise, we compare each optimiser on the deterministic formulation and two stochastic variants (with differing noise regimes). We also consider a second problem of interest, recently introduced by Rahimi & Recht (2017). It consists of fitting a deep network with only two linear layers to a dataset where sample inputs x are related to sample outputs y by the relation y = Ax, where A is an ill-conditioned matrix (with condition number = 105).
The results are shown in Table 1. We use a grid-search to determine the best hyper-parameters for both SGD and Adam (reported in appendix B.1). We report the number of iterates taken to reach the solution, with a tolerance of τ = 10−4. Statistics are computed over 100 runs of each optimiser. We observe that first-order methods perform poorly in all cases, and moreover show a very high
3We also experimented with RMSProp (Tieleman & Hinton, 2012), AdaGrad (Duchi et al., 2011) and AdaDelta (Zeiler, 2012), but found that these methods consistently underperformed Adam and SGD on these “toy” problems.
variance of results. The Newton method with an exact Hessian4 generally performs best, followed closely by Levenberg-Marquardt (LM), however they are impractical for larger-scale problems. Our method delivers comparable (and sometimes better) performance despite avoiding a costly Hessian inversion. On the other hand, the performance of BFGS, which approximates the Hessian with a buffer of parameter updates, seems to correlate negatively with the level of noise.
Fig. 1 shows example trajectories. The slow, oscillating behaviour of first-order methods is noticeable, as well as the impact of noise on the BFGS steps. On the other hand, CURVEBALL, Newton and LM converge in few iterations.
CIFAR. We now turn to the task of training deep networks on more realistic datasets. Second-order methods are typically not used in such scenarios, due to the large number of parameters and stochastic sampling. We start with a basic 5-layer convolutional neural network (CNN).5 We train this network for 20 epochs on CIFAR-10, with and without batch-normalization (which is known to improve conditioning (Kohler et al., 2018)) using for every experiment a mini-batch size of 128. To assess optimiser performance on larger models, we also train a much larger ResNet-18 model (He et al., 2016). As baselines, we picked SGD (with momentum) and Adam, which we found to outperform the competing first-order optimisers. Their learning rates are chosen from the set 10−k, k ∈ N with a grid search for the basic CNN, while for the ResNet SGD uses the schedule recommended by the authors (He et al., 2016). We focus on the training error, since it is the quantity being optimised by eq. 1 (validation error is discussed below). The results can be seen in Fig. 2 (top row). We observe that in each setting, CURVEBALL outperforms its competitors, in a manner that is robust to normalisation and model type.
ImageNet. To assess the practicality of our method at larger scales, we apply it to the classification task on the large-scale ImageNet dataset. We report results of training on both a medium-scale setting using a subset formed from the images of 100 randomly sampled classes as well as the large-scale setting, by training on the full dataset. Both experiments use the VGG-f architecture with mini-batch size of 256 and follow the settings described by Chatfield et al. (2014). The results are depicted in Fig. 2. We see that our method provides compelling performance against popular first order solvers in both cases, and that interestingly, its margin of improvement grows with the scale of the dataset.
4When the Hessian contains negative eigenvalues, their absolute values are used (Dauphin et al., 2014). 5The basic CNN has 5× 5 filters and ReLU activations, and 3× 3 max-pooling layers (with stride 2) after
each of the first 3 convolutions. The number of output channels are, respectively, (32, 32, 64, 64, 10).
Comparison to other second-order methods on MNIST. In order to compare ours with existing second-order methods, we use the public KFAC (Martens & Grosse, 2015) implementation made available by the authors and run a simple experiment on the MNIST dataset. In this scenario a four layer MLP (with output sizes 128-64-32-10) with hyperbolic tangent activations is trained on this classification task. We closely follow the same protocol as Martens & Grosse (2015) for layer initialisation and data normalisation, with batch size 64. We show results in Fig 2 (bottom row, right) with the best learning rate for each method. On this problem our method performs comparably to first order solvers, while KFAC makes less progress until it has stabilised its Fisher matrix estimation.
Random architecture results. It can be argued that standard architectures are biased to favour SGD, since it was used in the architecture searches, and architectures in which it failed to optimise were discarded (He et al., 2016). It would be useful to assess the optimisers’ ability to generalise across architectures, testing how well they perform regardless of the network model. We make an attempt in this direction by comparing the optimisers on 50 deep CNN architectures that are generated randomly (see appendix B.3 for details). In addition to being more architecture-agnostic, this makes any hand-tuning of hyper-parameters infeasible, which we believe to be a fair requirement for a dependable optimiser. The results on CIFAR10 are shown in figure 3 (left), as the median across all runs (thick lines) and 25th-75th percentiles (shaded regions). CURVEBALL consistently outperforms first-order methods, with the bulk of the achieved errors below those of SGD and Adam.
Wall-clock time. To provide an estimate of the relative efficiency of each model, Fig. 3 shows wall clock time on the basic CIFAR-10 model (without batch norm). Importantly, from a practical perspective, we observe that our method is competitive with first order solvers, while not requiring any tuning. Moreover, our prototype implementation includes custom FMAD operations which have not received the same degree of optimisation as RMAD (back-propagation), and could further benefit from careful engineering. We also experimented with a Hessian-free optimiser (based on conjugate gradients) (Martens, 2010). We show a comparison in logarithmic time in the appendix (Fig. 7). Due to the costly CG operation, which requires several passes through the network, it is an order of magnitude slower than the first-order methods and our own second-order method. This validates our initial motivation of designing a Hessian-free method without the inner CG loop (Section 3).
Overfitting and validation error While the focus of this work is optimisation, it is also of interest to compare the validation errors attained by the trained models – these are reported Table 2. We observe that models trained with the proposed method exhibit better training and validation error on most models, with the exception of ResNet where overfitting plays a more significant role. However, we note that this could be addressed with better regularisation, and we show one such example, by also reporting the validation error with a dropout rate of 0.3 in brackets.

heading: 5 RELATED WORK
text: While second order methods have proved to be highly effective tools for optimising deterministic functions (Moré, 1978; Wright & Nocedal, 1999, p. 164) their application to stochastic optimisation, and in particular to deep neural networks remains an active area of research. Many methods have been developed to improve stochastic optimisation with curvature information to avoid slow progress in ill-conditioned regions (Dauphin et al., 2014), while avoiding the cost of storing and inverting a Hessian matrix. A popular approach is to construct updates from a buffer of parameter gradients and their first-and-second-order moments at previous iterates (e.g. AdaGrad (Duchi et al., 2011), AdaDelta (Zeiler, 2012), RMSProp (Tieleman & Hinton, 2012) or Adam (Kingma & Ba, 2014)). These solvers benefit from needing no additional function evaluations beyond traditional mini-batch stochastic gradient descent. Typically they set adaptive learning rates by making use of empirical estimates of the curvature with a diagonal approximation to the Hessian (e.g. Zeiler (2012)) or a rescaled diagonal Gauss-Newton approximation (e.g. Duchi et al. (2011)). While the diagonal structure decreases the computational cost, their overall efficiency remains limited and in many cases can be matched by a well tuned SGD solver (Wilson et al., 2017).
Second order solvers take a different approach, investing more computation per iteration in the hope of achieving higher quality updates. Trust-region methods (Conn et al., 2000) and cubic regularization (Nesterov & Polyak, 2006; Cartis et al., 2011) are canonical examples. To achieve this higher quality, they invert the Hessian matrix H , or a tractable approximation such as the GaussNewton approximation (Martens, 2010; Martens & Sutskever, 2011; Botev et al., 2017) (described in section 2), or other regularized (Dauphin et al., 2014) or subsampled versions of the Hessian (Kohler & Lucchi, 2017; Xu et al., 2017). Another line of work belonging to the trust-region family (Conn et al., 2000), which has proven effective for tasks such as classification, introduces second order information with natural gradients (Amari, 1998). In this context, it is common to derive a loss function from a Kullback-Leibler (KL) divergence. The natural gradient makes use of the infinitesimal distance induced by the latter to follow the curvature in the Riemannian manifold equipped with this new distance. In practice the natural gradient method amounts to replacing the Hessian H in the modified gradient formula H−1J with the Fisher matrix F , which facilitates traversal of the optimal path in the metric space induced by the KL-divergence. Since the seminal work of Amari (1998) several authors have studied variations of this idea. TONGA (Roux et al., 2008) relies on the empirical Fisher matrix where the previous expectation over the model predictive distribution is replaced by the sample predictive distribution. The works of Pascanu & Bengio (2013) and Martens established a link between Gauss-Newton methods and the natural gradient. More recently Martens & Grosse (2015) introduced the KFAC optimiser which uses a block diagonal approximation of the Fisher matrix. This was shown to be an efficient stochastic solver in several settings, but it remains a computationally challenging approach for larger-scale deep networks problems.
Many of the methods discussed above perform an explicit system inversion that can often prove prohibitively expensive (Yang & Amari, 1998). Consequently, a number of works (Martens, 2010; Martens & Sutskever, 2011; Zhang et al., 2017) have sought to exploit the cheaper computation of Hessian-vector products via automatic differentiation (Pearlmutter, 1994; Schraudolph, 2002), to perform system inversions with conjugate gradients (Hessian-free methods). Other approaches (Carmon et al., 2018; Agarwal et al., 2017) have resorted to rank-1 approximations of the Hessian for efficiency. While these methods have had some success, they have only been demonstrated on single-layer models of moderate scale compared to the state-of-the-art in deep learning. We speculate that the main reason they are not widely adopted is their requirement of several steps (network passes) per parameter update (Sutskever et al., 2013; Koh & Liang, 2017), which would put them at a similar disadvantage w.r.t. first-order methods as the Hessian-free method that we tested (fig. 7 in the appendix). Perhaps more closely related to our approach, Orr (1995) uses automatic differentiation to compute Hessian-vector products to construct adaptive, per-parameter learning rates.
The closest method is LiSSA (Agarwal et al., 2017), which is built around the idea of approximating the Hessian inverse with a Taylor series expansion. This series can be implemented as the recursion H−1(r) = I + (I −H)H −1 (r−1), starting with H −1 (0) = I . Since LiSSA is a type of Hessian-free method, the core of the algorithm is similar to Algorithm 2: it also refines an estimate of the Newton step iteratively, but with a different update rule in line 4. With some simple algebraic manipulations, we can use the Taylor recursion to write this update in a form that is similar to ours: zr+1 = zr−Ĥzr−J . This looks similar to our gradient-descent-based update with a learning rate of β = 1 (Alg. 1, lines 3-4), with some key differences. First, they reset the state of the step estimate for every minibatch (Alg. 2, line 2). Reusing past solutions, like momentum-SGD, is an important factor in the performance of our algorithm, since we only have to perform one update per mini-batch. In contrast, Agarwal et al. (2017) report a typical number of inner-loop updates (R in Alg. 2) equal to the number of samples (e.g. R = 10, 000 for a tested subset of MNIST). While this is not a problem for their tested case of linear Support Vector Machines, since each update only requires one inner-product, the same does not apply to deep neural networks. Second, they invert the Hessian independently for each mini-batch, while our method aggregates the implicit Hessian across all past mini-batches (with a forgetting factor of ρ). Since batch sizes are orders of magnitude smaller than the number of parameters (e.g. 256 samples vs. 60 million parameters for the VGG-f), the Hessian matrix for a mini-batch is a poor substitute for the Hessian of the full dataset in these problems, and severely ill-conditioned. Third, while their method fixes β to 1, we found that setting it correctly can be used to attain convergence on ill-conditioned problems. For example, even on quadratic problems, we show that this parameter needs to be carefully chosen to avoid divergence (Theorem A.1). The gradient descent interpretation used in our work, contrasted to the Taylor series recursion of Agarwal et al. (2017), thus brings an additional degree of freedom that may be useful to relax other assumptions. Finally, while they demonstrate improved performance on convex problems with linear models, we focus on the needs of training deep networks on large datasets (millions of samples and parameters), on which no previous Newton method has been able to surpass the first-order methods that are commonly used by the deep learning community.

heading: 6 CONCLUSIONS AND FUTURE WORK
text: In this work, we have proposed a practical second-order solver that has been specifically tailored for deep-learning-scale stochastic optimisation problems. We showed that our optimiser can be applied to a large range of datasets and reach better training error than first order method with the same number of iterations, with essentially no hyper-parameters tuning. In future work, we intend to bring more improvements to the wall-clock time of our method by engineering the FMAD operation to the same standard as back-propagation, and study optimal trust-region strategies to obtain λ in closed-form.

heading: A ADDITIONAL PROOFS

heading: A.1 DERIVATION OF AUTOMATIC HYPER-PARAMETER TUNING IN CLOSED-FORM
text: We rewrite the problem in eq. 14 as a minimization over ρ and β where z′ = ρz − β∆z:
z = arg min ρ,β
f̂(z′) (20)
= arg min ρ,β
[ ρ −β ]T [ z ∆z ] T J + 1
2
[ ρ −β ]T [ z ∆z ] T H [ z ∆z ] [ ρ −β ] (21)
= arg min ρ,β
[ ρ −β ]T [ zTJ ∆Tz J ] + 1 2 [ ρ −β ]T [ zT Ĥz zT Ĥ∆z zT Ĥ∆z ∆ T z Ĥ∆z ] [ ρ −β ] . (22)
Since f̂ is a quadratic function of ρ and β with PSD Hessian it is therefore convex and we can find its extrema by cancelling the gradient:
∇ρ,β f̂(z′) = 0. (23)
Therefore, we have: [ zTJ −∆Tz J ] + [ zT Ĥz −zT Ĥ∆z −zT Ĥ∆z ∆Tz Ĥ∆z ] [ ρ β ] = 0 (24)[
−ρ β
] = [ zT Ĥz zT Ĥ∆z zT Ĥ∆z ∆ T z Ĥ∆z ]−1 [ zTJ ∆Tz J ] , (25)
where the last equality can be computed by inverting the 2× 2 matrix explicitly.

heading: A.2 PROOF OF CONVERGENCE IN THE QUADRATIC CASE
text: Theorem A.1. Let f be a convex quadratic function, and its hyper-parameters β > 0, ρ > 0 satisfy 3 2βhmax − 1 < ρ < 1 + βhmin, (26)
where hmin and hmax are the smallest and largest eigenvalues of the Hessian H , respectively. Then Algorithm 1 converges linearly to the minimum of f .
Corollary A.1.1. Algorithm 1 converges for any momentum parameter 0 < ρ < 1 with a sufficiently small learning rate β > 0, regardless of the (PSD) Hessian spectrum.
Proof of Theorem A.1. We follow similar derivations on quadratic models by previous work on the heavy-ball method (Goh, 2017; Lessard et al., 2016; Flammarion & Bach, 2015), but including our curvature term in the update. We assume the quadratic model:
f(w) = 1
2 wTHw − bTw, (27)
which has Hessian matrix H , and gradient J(w) = Hw − b. Without loss of generality, we will consider the pure Newton method, where H is not regularized (λ = 0):6
zt+1 = ρzt − β(Hzt + J(wt)) (28) wt+1 = wt + zt+1 (29)
Eq. 28 can be rearranged to zt+1 = (ρI − βH)zt − βJ(wt). (30)
6For the general case, the momentum parameter ρ is simply replaced by the slightly perturbed value ρ− βλ (since ρ βλ), and similar derivations follow.
We now perform a change of variables to diagonalize the Hessian, H = Qdiag(h)QT , with Q orthogonal and h the vector of eigenvalues. Let w∗ = arg minw f(w) = H−1b be the optimal solution of the minimization. Then, replacing wt = Qxt + w∗ in eq. 30:
Qyt+1 = (ρI − βH)Qyt − βHQxt (31)
with J = H(Qxt + w∗)− b = H(Qxt +H−1b)− b = HQxt. Then, expanding H with its eigendecomposition,
Qyt+1 = ρQyt − βQdiag(h)QTQyt − βQdiag(h)QTQxt. (32)
Left-multiplying by QT ,and canceling out Q due to orthogonality,
yt+1 = ρyt − βdiag(h)yt − βdiag(h)xt. (33)
Similarly for eq. 29, replacing zt = Qyt yields
xt+1 = xt + yt+1. (34)
Note that each pair formed by the corresponding element of yt and xt is an independent system with only 2 variables, since the pairs do not interact (eq. 33 and 34 only contain element-wise operations). From now on, we will be working on the ith element of each vector.
We can thus write eq. 33 and 34 (for a single element i of each) as a vector equation:[ 1 0 −1 1 ] [ yt+1,i xt+1,i ] = [ ρ− βhi −βhi 0 1 ] [ yt,i xt,i ] . (35)
The matrix on the left is necessary to express the fact that the yt+1 factor in eq. 34 must be moved to the left-hand side, which corresponds to iteration t+ 1 (xt+1 − yt+1 = xt). Left-multiplying eq. 35 by the inverse,7 [
yt+1,i xt+1,i
] = [ ρ− βhi −βhi ρ− βhi 1− βhi ] [ yt,i xt,i ] . (36)
This is the transition matrix Ri that characterizes the iteration, and taking its power models multiple iterations in closed-form: [
yt,i xt,i
] = Rti [ y0,i x0,i ] . (37)
The two eigenvalues of Ri are given in closed-form by:
eig (Ri) = 1
2
( ρ− 2βhi + 1± √ (ρ− 2βhi)2 − 2ρ+ 1 ) (38)
The series in eq. 37 converges when |eig (Ri)| < 1 simultaneously for both eigenvalues, which is equivalent to:
3 2βhi − 1 < ρ < 1 + βhi, (39)
with ρ > 0 and βhi > 0. Note that when using the Gauss-Newton approximation of the Hessian, hi > 0 and thus the last condition simplifies to β > 0.
Since eq. 39 has to be satisfied for every eigenvalue, we have 3 2βhmax − 1 < ρ < 1 + βhmin, (40)
with hmin and hmax the smallest and largest eigenvalues of the Hessian H , respectively, proving the result.
The rate of convergence is the largest of the two values |eig (Ri)|. When the argument of the square root in eq. 38 is non-negative, it does not admit an easy interpretation; however, when it is negative, eq. 38 simplifies to:
|eig (Ri)| = √ ρ− βhi. (41)
7We have: [
1 0 −1 1
]−1 = [ 1 0 1 1 ]

heading: A.2.1 GRAPHICAL INTERPRETATION
text: The convergence rate for a single eigenvalue is illustrated in Figure 4. Graphically, the regions of convergence for different eigenvalues will differ only by a scale factor along the βhi axis (horizontal stretching of Figure 4). Moreover, the largest possible range of βhi values is obtained when ρ = 1, and that range is 0 < βhi < 43 . We can infer that the intersection of the regions of convergence for several eigenvalues will be maximized with ρ = 1, for any fixed β.

heading: A.3 PROOF OF GUARANTEED DESCENT ON GENERAL NON-CONVEX FUNCTIONS
text: Theorem A.2. Let the Hessian Ĥt+1 be positive definite (which holds when the objective is convex or when Gauss-Newton approximation and trust region are used). Then the update zt+1 in Algorithm 1 is a descent direction when β and ρ are chosen according to eq. 18, and zt+1 6= 0.
Proof. To show that the update represents a descent direction, it suffices to show that JT zt+1 < 0 (where we have written J = J(wt) to simplify notation). Since the surrogate Hessian Ĥt+1 is positive definite (PD) by construction, the update zt+1 = ρzt − β∆zt+1 satisfies zTt+1Ĥt+1zt+1 > 0. It is therefore sufficient to prove that JT zt+1 + zTt+1Ĥt+1zt+1 ≤ 0. It follows from their definition in eq. (18) that ρ and β minimise the RHS of
JT zt+1 + 1
2 zTt+1Ĥt+1zt+1 =[
JT∆zt+1 JT zt ]T [ −β ρ ] + 1 2 [ −β ρ ]T [ ∆Tzt+1Ĥt+1∆zt+1 z T t Ĥt+1∆zt+1
zTt Ĥt+1∆zt+1 z T t Ĥt+1zt
] [ −β ρ ] (42)
In particular, they minimise a quadratic form in (−β, ρ) with the following symmetric Hessian
K =
[ ∆Tzt+1Ĥt+1∆zt+1 z T t Ĥt+1∆zt+1
zTt Ĥ∆zt+1 z T t Ĥt+1zt
] . (43)
Moreover, for any x = (x1, x2) ∈ R2,
xTKx = [ x1 x2 ]T [ ∆Tzt+1Ĥt+1∆zt+1 z T t Ĥt+1∆zt+1
zTt Ĥ∆zt+1 z T t Ĥt+1zt
] [ x1 x2 ] = (x1∆zt+1 + x2zt) T Ĥ(x1∆zt+1 + x2zt). (44)
Consequently, K is guaranteed to be Positive Semidefinite (PSD) and the form is convex with zero gradient at the minimum. Since zt+1 6= 0, it follows that at least one of the following holds: (1) K is
invertible and hence PD (rather than simply PSD); (2) one of factors zt = 0 or ∆zt+1 = 0 is zero; (3) the factors zt = 0 and ∆zt+1 = 0 are colinear. In the first case we have,
JT zt+1 + 1
2 zTt+1Ĥt+1zt+1 =
−1 2 [ JT∆zt+1 JT zt ]T [ ∆Tzt+1Ĥt+1∆zt+1 z T t Ĥt+1∆zt+1
zTt Ĥt+1∆zt+1 z T t Ĥt+1zt
]−1 [ JT∆zt+1 JT zt ] (45)
Since the inverse of a PD matrix is PD, the RHS of eq. (42) is negative. Further, as Ĥt+1 is PD, it follows that final term in eq. (42) is positive, thus K is PD, showing that JT zt+1 < 0.
For the second case in which zt = 0 or ∆zt+1 = 0, the system reduces to a trivial convex second order equation in ρ or β with a negative solution.
Finally, consider the case when zt and ∆zt+1 are colinear but both non-negative. Writing ∆zt+1 = αzt for α ∈ R, we note that at the minimum we have
JT zt+1 + 1
2 zTt+1Ĥt+1zt+1 = −
1
2 [ JT∆zt+1 JT zt ]T [ −β ρ ] = −1 2 (ρ− αβ)JT zt. (46)
Thus at the minimum (46) is negative, closing the proof.
Remark. It follows from the definition of ρ and β that if J(wt) = 0, then zt+1 = 0. Remark. If zt+1 = 0, then zt+2 = −βJ(wt+1), i.e. we reset the momentum variable z. This guarantees that the algorithm takes a strictly descending direction at least every two steps.

heading: B ADDITIONAL RESULTS AND IMPLEMENTATION DETAILS

heading: B.1 CONFIGURATIONS FOR SMALL-SCALE DATASET EXPERIMENTS
text: Here we provide additional details of the small-scale datasets described in sec. 4. As noted in the main paper, to give every method the best chance of working effectively we first perform a grid-search over its hyperparameters. This search is performed for each of the small-scale dataset experiments. For each first order solver listed in Table. 1, we select the configuration which achieves the lowest average error across the final ten iterations of a trajectory. The values included in the search were:
• SGD with momentum: learning rates: Γ, momentum values: 0.9, 0.95, 0.99 • Adam: learning rates Γ, β1 : 0.9, 0.99, β2 : 0.99, 0.999
where Γ = 0.1, 0.05, 0.01, 0.05, 0.001, 0.005, 0.0001, 0.0005.

heading: B.2 HYPER-PARAMETER AND GRADIENT EVOLUTION

heading: B.3 RANDOM ARCHITECTURE EXPERIMENT SETUP
text: Each optimiser is tested on 50 random networks, that are held fixed across all methods. The number of convolutional layers is uniformly sampled between 3 and 10, and the number of channels in each layer is drawn uniformly, in powers of two, between 32 and 256. The kernel size is 3× 3. Following each convolution (except the last one) there is a ReLU activation and batch-normalisation, and 3× 3 max-pooling (stride 2) is placed with 50% chance. Training and evaluation is performed on CIFAR10, with a batch size of 256.

heading: B.4 WALL-CLOCK TIME RESULTS WITH CONJUGATE GRADIENT
text: B.5 EXPERIMENTS WITHOUT A MOMENTUM HYPER-PARAMETER (FIXED ρ = 1)

Review summaries:
Authors propose choosing direction by using a single step of gradient descent "towards Newton step" from an original estimate, and then taking this direction instead of original gradient. This direction is reused as a starting estimate for the next iteration of the algorithm. This can be efficiently implemented since it only relies on Hessian-vector products which are accessible in all major frameworks.

This paper proposes an approximate second-order method with low computational cost. A common pitfall of second-order methods is the computation (and perhaps inversion) of the Hessian matrix. While this can be avoided by instead relying on Hessian-vector products as done in CG, it typically still requires several iterations. Instead, the authors suggest a simpler approach that relies on one single gradient step and a warm start strategy. The authors points out that the resulting algorithm resembles a momentum method. They also provide some simple convergence proofs on quadratics and benchmark their method to train deep neural networks.