Practical Quasi-Newton Methods for Training Deep Neural Networks

We consider in this paper the development of practical stochastic quasi-Newton (QN), and in particular Kronecker-factored block-diagonal BFGS [6, 13, 16, 39] and L-BFGS [27], methods for training deep neural networks (DNNs). Recall that the BFGS method starts each iteration with a symmetric positive definite matrix $B$ (or $H=B^{-1}$) that approximates the current Hessian matrix (or its inverse), computes the gradient $\mathbf{\nabla f}$ of $f$ at the current iterate $\mathbf{x}$ and then takes a step $\mathbf{s}=-\alpha H\mathbf{\nabla f}$, where $\alpha$ is a step length (usually) determined by some inexact line-search procedure, such that $\mathbf{y}^{\top}\mathbf{s}>0$, where $\mathbf{y}=\mathbf{\nabla f}^{+}-\mathbf{\nabla f}$ and $\mathbf{\nabla f}^{+}$ is the gradient of $f$ at the new point $\mathbf{x}^{+}=\mathbf{x}+\mathbf{s}$. The method then computes an updated approximation $B^{+}$ to $B$ (or $H^{+}$ to $H$) that remains symmetric and positive-definite and satisfies the so-called quasi-Newton (QN) condition $B^{+}\mathbf{s}=\mathbf{y}$ (or equivalently, $H^{+}\mathbf{y}=\mathbf{s}$). A consequence of this is that the matrix $B^{+}$ operates on the vector $\mathbf{s}$ in exactly the same way as the average of the Hessian matrix along the line segment between $\mathbf{x}$ and $\mathbf{x}^{+}$ operates on $\mathbf{s}$.

In DNN training, the number of variables and components of the gradient $n$ is often of the order of tens of millions and the Hessian has $n^{2}$ elements. Hence, computing and storing a full $n\times n$ BFGS approximation or storing $p$ $(\mathbf{s},\mathbf{y})$ pairs, where $p$ is approximately $10$ or larger for use in an L-BFGS implementation, is out of the question. Consequently, in our methods, we approximate the Hessian by a block-diagonal matrix, where each diagonal block corresponds to a layer, further approximating them as the Kronecker product of two much smaller matrices, as in [30, 5, 19, 10].

Literature Review on Using Second-order Information for DNN Training. For solving the stochastic optimization problems with high-dimensional data that arise in machine learning (ML), stochastic gradient descent (SGD) [36] and its variants are the methods that are most often used, especially for training DNNs. These variants include such methods as AdaGrad [12], RMSprop [21], and Adam [24], all of which scale the stochastic gradient by a diagonal matrix based on estimates of the first and second moments of the individual gradient components. Nonetheless, there has been a lot of effort to find ways to take advantage of second-order information in solving ML optimization problems. Approaches have run the gamut from the use of a diagonal re-scaling of the stochastic gradient, based on the secant condition associated with quasi-Newton (QN) methods [4], to sub-sampled Newton methods (e.g. see [43], and references therein), including those that solve the Newton system using the linear conjugate gradient method (see [8]).

In between these two extremes are stochastic methods that are based either on QN methods or generalized Gauss-Newton (GGN) and natural gradient [1] methods. For example, a stochastic L-BFGS method for solving strongly convex problems was proposed in [9] that uses sampled Hessian-vector products rather than gradient differences, which was proved in [33] to be linearly convergent by incorporating the variance reduction technique (SVRG [23]) to alleviate the effect of noisy gradients. A closely related variance reduced block L-BFGS method was proposed in [17]. A regularized stochastic BFGS method was proposed in [31], and an online L-BFGS method was proposed in [32] for strongly convex problems and extended in [28] to incorporate SVRG variance reduction. Stochastic BFGS and L-BFGS methods were also developed for online convex optimization in [38]. For nonconvex problems, a damped L-BFGS method which incorporated SVRG variance reduction was developed and its convergence properties was studied in [41].

GGN methods that approximate the Hessian have been proposed, including the Hessian-free method [29] and the Krylov subspace method [40]. Variants of the closely related natural gradient method that use block-diagonal approximations to the Fisher information matrix, where blocks correspond to layers, have been proposed in e.g. [20, 11, 30, 14]. Using further approximation of each of these (empirical) Fisher matrix and GGN blocks by the Kronecker product of two much smaller matrices, the efficient KFAC [30], KFRA [5], EKFAC [15], and Shampoo [19] methods were developed. See also [2] and [10], [37], which combine both Hessian and covariance (Fisher-like) matrix information in stochastic Newton type methods, Also, methods are given in [26, 42] that replace the Kullback-Leibler divergence by the Wasserstein distance to define the natural gradient, but with a greater computational cost.

Our Contributions. The main contributions of this paper can be summarized as follows:

1. 1.

   New BFGS and limited-memory variants (i.e. L-BFGS) that take advantage of the structure of feed-forward DNN training problems;

2. 2.

   Efficient non-diagonal second-order algorithms for deep learning that require a comparable amount of memory and computational cost per iteration as first-order methods;

3. 3.

   A new damping scheme for BFGS and L-BFGS updating of an inverse Hessian approximation, that not only preserves its positive definiteness, but also limits the decrease (and increase) in its smallest (and largest) eigenvalues for non-convex problems;

4. 4.

   A novel application of Hessian-action BFGS;

5. 5.

   The first proof of convergence (to the best of our knowledge) of a stochastic Kronecker-factored quasi-Newton method.

After reviewing the computations used in DNN training, we describe the Kronecker structures of the gradient and Hessian for a single data point, followed by their extension to approximate expectations of these quantities for multiple data-points and give a generic algorithm that employs BFGS (or L-BFGS) approximations for the Hessians.

Deep Neural Networks. We consider a feed-forward DNN with $L$ layers, defined by weight matrices $W_{l}$ (whose last columns are bias vectors $b_{l}$), activation functions $\phi_{l}$ for $l\in\{1\dots L\}$ and loss function $\mathcal{L}$. For a data-point $(x,y)$, the loss $\mathcal{L}\left(a_{L},y\right)$ between the output $a_{L}$ of the DNN and $y$ is a non-convex function of $\theta=\left[\operatorname{vec}\left(W_{1}\right)^{\top},\ldots,\operatorname{vec}\left(W_{L}\right)^{\top}\right]^{\top}$. The network’s forward and backward pass for a single input data point $(x,y)$ is described in Algorithm 1.

For a training dataset that contains multiple data-points indexed by $i=1,...,I$, let $f(i;\theta)$ denote the loss for the $i$th data-point. Then, viewing the dataset as an empirical distribution, the total loss function $f(\theta)$ that we wish to minimize is

Single Data-point: Layer-wise Structure of the Gradient and Hessian. Let ${\mathbf{\nabla f}}_{l}$ and $\nabla^{2}f_{l}$ denote, respectively, the restriction of $\mathbf{\nabla f}$ and $\nabla^{2}f$ to the weights $W_{l}$ in layer $l=1,\ldots,L$. For a single data-point ${\mathbf{\nabla f}}_{l}$ and $\nabla^{2}f_{l}$ have a tensor (Kronecker) structure, as shown in [30] and [5]. Specifically,

where the pre-activation gradient ${\mathbf{g}}_{l}(i)=\frac{\partial f(i)}{\partial{\mathbf{h}}_{l}(i)}$, and the pre-activation Hessian $G_{l}(i)=\frac{\partial^{2}f(i)}{\partial{\mathbf{h}}_{l}(i)^{2}}$. Our algorithm uses an approximation to $(G_{l}(i))^{-1}$, which is updated via the BFGS updating formulas based upon a secant condition that relates the change in ${\mathbf{g}}_{l}(i)$ with the change in ${\mathbf{h}}_{l}(i)$.

Although we focus on fully-connected layers in this paper, the idea of Kronecker-factored approximations to the diagonal blocks $\nabla^{2}f_{l}$, $l=1,\ldots,L$ of the Hessian can be extended to other layers used in deep learning, such as convolutional and recurrent layers.

Multiple Data-points: Kronecker-factored QN Approach. Now consider the case where we have a dataset of $I$ data-points indexed by $i=1,\ldots,I$. By (2), we have

Note that the approximation in (3) that the expectation of the Kronecker product of two matrices equals the Kronecker product of their expectations is the same as the one used by KFAC [30]. Now, based on this structural approximation, we use $H^{l}=H^{l}_{a}\otimes H^{l}_{g}$ as our QN approximation to $\left(\mathbb{E}_{i}[\nabla^{2}f_{l}(i)]\right)^{-1}$, where $H^{l}_{a}$ and $H^{l}_{g}$ are positive definite approximations to $A_{l}^{-1}$ and $G_{l}^{-1}$, respectively. Hence, using our layer-wise block-diagonal approximation to the Hessian, a step in our algorithm for each layer $l$ is computed as

where $\widehat{\mathbf{\nabla f}_{l}}$ denotes the estimate to $\mathbb{E}_{i}[\mathbf{\nabla f}_{l}(i)]$ and $\alpha$ is the learning rate. After computing $W_{l}^{+}$ and performing another forward/backward pass, our method computes or updates $H_{a}^{l}$ and $H_{g}^{l}$ as follows:

1. 1.

   For $H_{g}^{l}$, we use a damped version of BFGS (or L-BFGS) (See Section 3) based on the $({\mathbf{s}},{\mathbf{y}})$ pairs corresponding to the average change in ${\mathbf{h}}_{l}(i)$ and in the gradient with respect to ${\mathbf{h}}_{l}(i)$; i.e.,

   $\displaystyle{\mathbf{s}}_{g}^{l}=\mathbb{E}_{i}[{\mathbf{h}}_{l}^{+}(i)]-\mathbb{E}_{i}[{\mathbf{h}}_{l}(i)],\qquad{\mathbf{y}}_{g}^{l}=\mathbb{E}_{i}[{\mathbf{g}}_{l}^{+}(i)]-\mathbb{E}_{i}[{\mathbf{g}}_{l}(i)].$

   (5)

2. 2.

   For $H_{a}^{l}$ we use the "Hessian-action" BFGS method described in Section 4. The issue of possible singularity of the positive semi-definite matrix $A_{l}$ approximated by $(H_{a}^{l})^{-1}$ is also addressed there by incorporating a Levenberg-Marquardt (LM) damping term.

Algorithm 2 gives a high-level summary of our K-BFGS or K-BFGS(L) algorithms (which use BFGS or L-BFGS to update $H_{g}^{l}$, respectively). See Algoirthm 4 in the appendix for a detailed pseudocode. The use of mini-batches is described in Section 6. Note that, an additional forward-backward pass is used in Algorithm 2 because the quantities in (5) need to be estimated using the same mini-batch.

Damped BFGS Updating. It is well-known that training a DNN is a non-convex optimization problem. As (2) and (3) show, this non-convexity manifests in the fact that $G_{l}\succ 0$ often does not hold. Thus, for the BFGS update of $H_{g}^{l}$, the approximation to $G_{l}^{-1}$, to remain positive definite, we have to ensure that $({\mathbf{s}}_{g}^{l})^{\top}{\mathbf{y}}_{g}^{l}>0$. Due to the stochastic setting, ensuring this condition by line-search, as is done in deterministic settings, is impractical. In addition, due to the large changes in curvature in DNN models that occur as the parameters are varied, we also need to suppress large changes to $H_{g}^{l}$ as it is updated. To deal with both of these issues, we propose a double damping (DD) procedure (Algorithm 3), which is based upon Powell’s damped-BFGS approach [35], for modifying the $({\mathbf{s}}_{g}^{l},{\mathbf{y}}_{g}^{l})$ pair. To motivate Algorithm 3, consider the formulas used for BFGS updating of $B$ and $H$:

where $\rho=\frac{1}{{\mathbf{s}}^{\top}{\mathbf{y}}}>0$. If we can ensure that $0<\frac{{\mathbf{y}}^{\top}H{\mathbf{y}}}{{\mathbf{s}}^{\top}{\mathbf{y}}}\leq\frac{1}{\mu_{1}}$ and $0<\frac{{\mathbf{s}}^{\top}{\mathbf{s}}}{{\mathbf{s}}^{\top}{\mathbf{y}}}\leq\frac{1}{\mu_{2}}$ , then we can obtain the following bounds:

and

Thus, the change in $B$ (and $H$) is controlled if $\frac{{\mathbf{y}}^{\top}H{\mathbf{y}}}{{\mathbf{s}}^{\top}{\mathbf{y}}}\leq\frac{1}{\mu_{1}}$ and $\frac{{\mathbf{s}}^{\top}{\mathbf{s}}}{{\mathbf{s}}^{\top}{\mathbf{y}}}\leq\frac{1}{\mu_{2}}$. Our DD approach is a two-step procedure, where the first step (i.e. Powell’s damping of $H$) guarantees that $\frac{{\mathbf{y}}^{\top}H{\mathbf{y}}}{{\mathbf{s}}^{\top}{\mathbf{y}}}\leq\frac{1}{\mu_{1}}$ and the second step (i.e., Powell’s damping with $B=I$) guarantees that $\frac{{\mathbf{s}}^{\top}{\mathbf{s}}}{{\mathbf{s}}^{\top}{\mathbf{y}}}\leq\frac{1}{\mu_{2}}$. Note that there is no guarantee of $\frac{{\mathbf{y}}^{\top}H{\mathbf{y}}}{{\mathbf{s}}^{\top}{\mathbf{y}}}\leq\frac{1}{\mu_{1}}$ after the second step. However, we can skip updating $H$ in this case so that the bounds on these matrices hold. In our implementation, we always do the update, since in empirical testing, we observed that at least 90% of the pairs satisfy $\frac{{\mathbf{y}}^{\top}H{\mathbf{y}}}{{\mathbf{s}}^{\top}{\mathbf{y}}}\leq\frac{2}{\mu_{1}}$. See Section C in the appendix for more details on damping.

L-BFGS Implementation. L-BFGS can also be used to update $H_{g}^{l}$. However, implementing L-BFGS using the standard "two-loop recursion" (see Algorithm 7.4 in [34]) is not efficient. This is because the main work in computing $H^{l}_{g}\widehat{{\mathbf{\nabla f}}_{l}}H^{l}_{a}$ in line 4 of Algorithm 2 would require $4p$ matrix-vector multiplications, each requiring $O(d_{i}d_{o})$ operations, where $p$ denotes the number of $({\mathbf{s}},{\mathbf{y}})$ pairs stored by L-BFGS. (Recall that $\widehat{{\mathbf{\nabla f}}_{l}}\in R^{d_{o}\times d_{i}}$.) Instead, we use a "non-loop" implementation [7] of L-BFGS, whose main work involves 2 matrix-matrix multiplications, each requiring $O(pd_{i}d_{o})$ operations. When $p$ is not small (we used $p=100$ in our tests), and $d_{i}$ and $d_{o}$ are large, this is much more efficient, especially on GPUs.

In addition to approximating $G_{l}^{-1}$ by $H_{g}^{l}$ using BFGS, we also propose approximating $A_{l}^{-1}$ by $H_{a}^{l}$ using BFGS. Note that $A_{l}$ does not correspond to some Hessian of the objective function. However, we can generate $({\mathbf{s}},{\mathbf{y}})$ pairs for it by "Hessian action" (see e.g. [9, 17, 18]).

Connection between Hessian-action BFGS and Matrix Inversion. In our methods, we choose ${\mathbf{s}}=H^{l}_{a}\cdot\mathbb{E}_{i}[{\mathbf{a}}_{l-1}(i)]$ and ${\mathbf{y}}=A_{l}{\mathbf{s}}$, which as we now show, is closely connected to using the Sherman-Morrison modification formula to invert $A_{l}$. In particular, suppose that $A^{+}=A+{c\cdot}{\mathbf{a}}{\mathbf{a}}^{\top}$; i.e., only a rank-one update is made to $A$. This corresponds to the case where the information of $A$ is accumulated from iteration to iteration, and the size of the mini-batch is 1 or ${\mathbf{a}}$ represents the average of the vectors ${\mathbf{a}}(i)$ from multiple data-points.

In fact, all updates in the Broyden family are equivalent to applying the Sherman-Morrison modification formula to $A^{+}=A+{c\cdot}{\mathbf{a}}{\mathbf{a}}^{\top}$, given $H=A^{-1}$, since after substituting for ${\mathbf{s}}$ and ${\mathbf{y}}$ in (11) and simplifying, one obtains

When using momentum, $A^{+}=\beta A+(1-\beta){\mathbf{a}}{\mathbf{a}}^{\top}$ ($0<\beta<1$). Hence, if we still want Theorem 1 to hold, we have to scale $H$ by $1/\beta$ before updating it. This, however, turns out to be unstable. Hence, in practice, we use the non-scaled version of "Hessian action" BFGS.

Levenberg-Marquardt Damping for $A_{l}$. Since $A_{l}=\mathbb{E}_{i}\left[({\mathbf{a}}_{l-1}(i)({\mathbf{a}}_{l-1}(i))^{\top})\right]\succeq 0$ may not be positive definite, or may have very small positive eigenvalues, we add an Levenberg-Marquardt (LM) damping term to make our "Hessian-action" BFGS stable; i.e., we use ${A_{l}}+\lambda_{A}I_{A}$ instead of ${A_{l}}$, when we update $H_{a}^{l}$. Specifically, "Hessian action" BFGS for $A_{l}$ is performed as

1. 1.

   $A_{l}=\beta\cdot A_{l}+(1-\beta)\cdot\mathbb{E}_{i}\left[{\mathbf{a}}_{l-1}(i){\mathbf{a}}_{l-1}(i)^{\top}\right]$; $A_{l}^{\text{LM}}=A_{l}+\lambda_{A}I_{A}$.

2. 2.

   ${\mathbf{s}}_{a}^{l}=H_{a}^{l}\cdot\mathbb{E}_{i}[{\mathbf{a}}_{l-1}(i)]$, ${\mathbf{y}}_{a}^{l}=A_{l}^{\text{LM}}{\mathbf{s}}_{a}^{l}$; use BFGS with $({\mathbf{s}}_{a}^{l},{\mathbf{y}}_{a}^{l})$ to update $H_{a}^{l}$.

Following the framework for stochastic quasi-Newton methods (SQN) established in [41] for solving nonconvex stochastic optimization problems (see Section B in the appendix for this framework), we prove that, under fairly standard assumptions, for our K-BFGS(L) algorithm with skipping DD and exact inversion on $A_{l}$ (see Algorithm 5 in Section B), the number of iterations $N$ needed to obtain $\frac{1}{N}\sum_{k=1}^{N}\mathbb{E}[\|\mathbf{\nabla f}(\theta_{k})\|^{2}]\leq\epsilon$ is $N=O(\epsilon^{-\frac{1}{1-\beta}})$, for step size $\alpha_{k}$ chosen proportional to $k^{-\beta}$, where $\beta\in(0.5,1)$ is a constant. Our proofs, which are delayed until Section B, make use of the following assumptions, the first two of which, were made in [41].

To use the convergence analysis in [41], we need to show that the block-diagonal approximation of the inverse Hessian used in Algorithm 5 satisfies the assumption that it is bounded above and below by positive-definite matrices. Given the Kronecker structure of our Hessian inverse approximation, it suffices to prove boundness of both $H_{a}^{l}(k)$ and $H_{g}^{l}(k)$ for all iterations $k$. Making the additional assumption AS.3, we are able to prove Lemma 1, and hence Lemma 3, below. Note that many popular activation functions satisfy AS.3, such as sigmoid and tanh.

Using Lemma 3, we can now apply Theorem 2.8 in [41] to prove the convergence of Algorithm 5:

Note: other theorems in [41], namely Theorems 2.5 and 2.6, also apply here under our assumptions.

Before we present some experimental results, we address the use of moving averages, and the computational and storage requirements of the algorithms that we tested.

Mini-batch and Moving Average. Clearly, using the whole dataset at each iteration is inefficient; hence, we use a mini-batch to estimate desired quantities. We use $\overline{X}$ to denote the averaged value of $X$ across the mini-batch for any quantity $X$. To incorporate information from the past as well as reducing the variability, we use an exponentially decaying moving average to estimate desired quantities with decay parameter $\beta\in(0,1)$:

1. 1.

   To estimate the gradient $\mathbb{E}_{i}[\mathbf{\nabla f}(i)]$, at each iteration, we update $\widehat{\mathbf{\nabla f}}={\beta}\cdot\widehat{\mathbf{\nabla f}}+{(1-\beta)}\cdot\overline{\mathbf{\nabla f}}$.

2. 2.

   $H_{a}^{l}$: To estimate $A_{l}$, at each iteration we update $\widehat{A_{l}}={\beta}\cdot\widehat{A_{l}}+{(1-\beta)}\cdot\overline{{\mathbf{a}}_{l-1}{\mathbf{a}}_{l-1}^{\top}}.$ Note that although we compute ${\mathbf{s}}_{a}^{l}$ as $H_{a}^{l}\cdot\overline{{\mathbf{a}}_{l-1}}$, we update $\widehat{A_{l}}$ with $\overline{{\mathbf{a}}_{l-1}{\mathbf{a}}_{l-1}^{\top}}$ (i.e. the average ${\mathbf{a}}_{l-1}(i){\mathbf{a}}_{l-1}(i)^{\top}$ over the minibatch, not $\overline{{\mathbf{a}}_{l-1}}\cdot(\overline{{\mathbf{a}}_{l-1}})^{\top}$).

3. 3.

   $H^{l}_{g}$: BFGS "uses" momentum implicitly incorporated in the matrices $H^{l}_{g}$. To further stabilize the BFGS update, we also use a moving-averaged $({\mathbf{s}}_{g}^{l},{\mathbf{y}}_{g}^{l})$ (before damping); i.e., We update ${\mathbf{s}}_{g}^{l}={\beta}\cdot{\mathbf{s}}_{g}^{l}+{(1-\beta)}\cdot\left(\overline{{\mathbf{h}}_{l}^{+}}-\overline{{\mathbf{h}}_{l}}\right)$, and ${\mathbf{y}}_{g}^{l}={\beta}\cdot{\mathbf{y}}_{g}^{l}+{(1-\beta)}\cdot\left(\overline{{\mathbf{g}}_{l}^{+}}-\overline{{\mathbf{g}}_{l}}\right)$.

Finally, when computing $\overline{{\mathbf{h}}_{l}^{+}}$ and $\overline{{\mathbf{g}}_{l}^{+}}$, we use the same mini-batch as was used to compute $\overline{{\mathbf{h}}_{l}}$ and $\overline{{\mathbf{g}}_{l}}$. This doubles the number of forward-backward passes at each iteration.

Storage and Computational Complexity. Tables 1 and 2 compare the storage and computational requirements, respectively, for a layer with $d_{i}$ inputs and $d_{o}$ outputs for K-BFGS, K-BFGS(L), KFAC, and Adam/RMSprop. We denote the size of mini-batch by $m$, the number of $({\mathbf{s}},{\mathbf{y}})$ pairs stored for L-BFGS by $p$, and the frequency of matrix inversion in KFAC by $T$. Besides the requirements listed in Table 1, all algorithms need storage for the parameters $W_{l}$ and the estimate of the gradient, $\widehat{\bf{\nabla f}_{l}}$, (i.e. $O(d_{i}d_{o})$). Besides the work listed in Table 2, all algorithms also need to do a forward-backward pass to compute $\mathbf{\nabla f}_{l}$ as well as updating $W_{l}$, (i.e. $O(md_{i}d_{o})$). Also note that, even though we use big-$O$ notation in these tables, the constants for all of the terms in each of the rows are roughly at the same level and relatively small.

In Table 2, for K-BFGS and K-BFGS(L), "Additional pass" refers to Line 5 of Algorithm 2; under "Curvature", $O(md_{i}^{2})$ arises from "Hessian action" BFGS to update $H_{a}^{l}$ (see the algorithm at the end of Section 4), $O(md_{o})$ arises from (5), $O(d_{o}^{2})$ arises from updating $H_{g}^{l}$ (only for K-BFGS); and "Step $\Delta W_{l}$" refers to (4). For KFAC, referring to Algorithm 7 (in the appendix), "Additional pass" refers to Line 7; under "Curvature", $O(md_{i}^{2}+md_{o}^{2})$ refers to Line 8, and $O(\frac{1}{T}d_{i}^{3}+\frac{1}{T}d_{o}^{3})$ refers to Line 10; and "Step $\Delta W_{l}$" refers to Line 5.

From Table 1, we see that the Kronecker property enables K-BFGS and K-BFGS(L) (as well as KFAC) to have storage requirements comparable to those of first-order methods. Moreover, from Table 2, we see that K-BFGS and K-BFGS(L) require less computation per iteration than KFAC, since they only involve matrix multiplications, whereas KFAC requires matrix inversions which depend cubically on both $d_{i}$ and $d_{o}$. The cost of matrix inversion in KFAC (and singular value decomposition in [19]) is amortized by performing these operations only once every $T$ iterations; nonetheless, these amortized operations usually become much slower than matrix multiplication as models scale up.

Experimental Results. We tested K-BFGS and K-BFGS(L), as well as KFAC, Adam/RMSprop and SGD-m (SGD with momentum) on three autoencoder problems, namely, MNIST [25], FACES, and CURVES, which are used in e.g. [22, 29, 30], except that we replaced the sigmoid activation with ReLU. See Section D in the appendix for a complete description of these problems and the competing algorithms.

Since one can view Powell’s damping with $B=I$ as LM damping, we write $\mu_{2}=\lambda_{G}$, where $\lambda_{G}$ denotes the LM damping parameter for $G_{l}$. We then define $\lambda=\lambda_{A}\lambda_{G}$ as the overall damping term of our QN approximation. To simplify matters, we chose $\lambda_{A}=\lambda_{G}=\sqrt{\lambda}$, so that we needed to tune only one hyper-parameter (HP) $\lambda$.

To obtain the results in Figure 1, we first did a grid-search on (learning rate, damping) pairs for all algorithms (except for SGD-m, whose grid-search was only on learning rate), where damping refers to $\lambda$ for K-BFGS/K-BFGS(L)/KFAC, and $\epsilon$ for RMSprop/Adam. We then selected the best (learning rate, damping) pairs with the lowest training loss encountered. The range for the grid-search and the best HP values (as well as other fixed HP values) are listed in Section D in the appendix. Using the best HP values that we found, we then made 20 runs employing different random seeds, and plotted the mean value of the 20 runs as the solid line and the standard deviation as the shaded area.¹¹1 Results were obtained on a machine with 8 x Intel(R) Xeon(R) CPU @ 2.30GHz and 1 x NVIDIA Tesla P100. Code is available at https://github.com/renyiryry/kbfgs_neurips2020_public.

From the training loss plots in Figure 1, our algorithms clearly outperformed the first-order methods, except for RMSprop/Adam on CURVES, with respect to CPU time, and performed comparably to KFAC in terms of both CPU time and number of epochs taken. The testing error plots in Figure 1 show that our K-BFGS(L) method and KFAC behave very similarly and substantially outperform all of the first-order methods in terms of both of these measures. This suggests that our algorithms not only optimize well, but also generalize well.

To further demonstrate the robustness of our algorithms, we examined the loss under various HP settings, which showed that our algorithms are stable under a fairly wide range for the HPs (see Section D.4 in the appendix).

We also repeated our experiments using mini-batches of size 100 for all algorithms (see Figures 4, 5, and 6 in the appendix, where HPs are optimally tuned for batch sizes of 100) and our proposed methods continue to demonstrate advantageous performance, both in training and testing. For these experiments, we did not experiment with 20 random seeds. These results show that our approach works as well for relatively small mini-batch sizes of 100, as those of size 1000, which are less noisy, and hence is robust in the stochastic setting employed to train DNNs.

Compared with other methods mentioned in this paper, our K-BFGS and K-BFGS(L) have the extra advantage of being able to double the size of minibatch for computing the stochastic gradient with almost no extra cost, which might be of particular interest in a highly stochastic setting. See Section D.6 in the appendix for more discussion on this and some preliminary experimental results.

We proposed Kronecker-factored QN methods, namely, K-BFGS and K-BFGS(L), for training multi-layer feed-forward neural network models, that use layer-wise second-order information and require modest memory and computational resources. Experimental results indicate that our methods outperform or perform comparably to the state-of-the-art first-order and second-order methods. Our methods can also be extended to convolutional and recurrent NNs.

The research presented in this paper provides a new method for training DNNs that our experimental testing has shown to be more efficient in several cases than current state-of-the-art optimization methods for this task. Consequently, because of the wide use of DNNs in machine learning, this should help save a substantial amount of energy. Our new algorithms simply attempt to minimize the loss function that are given to it and the computations that it performs are all transparent. In machine learning, DNNs can be trained to address many types of problems, some of which should lead to positive societal outcomes, such as ones in medicine (e.g., diagnostics and drug effectiveness), autonomous vehicle development, voice recognition and climate change. Of course, optimization algorithms for training DNNs can be used to train models that may have negative consequences, such as those intended to develop psychological profiles, invade privacy and justify biases. The misuse of any efficient optimization algorithm for machine learning, and in particular our algorithms, is beyond the control of the work presented here.

DG and YR were supported in part by NSF Grant IIS-1838061. DG acknowledges the computational support provided by Google Cloud Platform Education Grant, Q81G-U4X3-5AG5-F5CG.