Title: BETTER GENERALIZATION BY EFFICIENT TRUST RE-GION METHOD

Abstract: In this paper, we develop a trust region method for training deep neural networks. At each iteration, trust region method computes the search direction by solving a non-convex subproblem. Solving this subproblem is non-trivial-existing methods have only sub-linear convergence rate. In the first part, we show that a simple modification of gradient descent algorithm can converge to a global minimizer of the subproblem with an asymptotic linear convergence rate. Moreover, our method only requires Hessian-vector products, which can be computed efficiently by back-propagation in neural networks. In the second part, we apply our algorithm to train large-scale convolutional neural networks, such as VGG and Mo-bileNets. Although trust region method is about 3 times slower than SGD in terms of running time, we observe it finds a model that has lower generalization (test) error than SGD, and this difference is even more significant in large batch training. We conduct several interesting experiments to support our conjecture that the trust region method can avoid sharp local minima.

1. INTRODUCTION:

Despite having many optimization algorithms in the literature (Nesterov, 2013; Wright & Nocedal, 1999; Boyd & Vandenberghe, 2004) , when talking about neural network optimization, by far the most popular methods are SGD and its variants such as AdaDelta (Zeiler, 2012) , AdaGrad (Duchi et al., 2011) , Adam (Kingma & Ba, 2014) and RMSProp (Hinton et al., 2006) . Compared to these stochastic first order methods, second order methods are rarely used in training deep neural networks due to large data files and high dimensional parameter space. Moreover, neural network has nonconvex loss function and its Hessian is often ill-conditioned, making it difficult to apply second order methods. Prior to our work, there aren't many works on second order methods for neural network optimization (Martens, 2010; Kiros, 2013; Botev et al., 2017) , and most of them are tested on multilayer perceptron so it is unclear whether second order methods can be useful for more advanced networks such as convolutional neural networks.

In this paper, we develop an efficient trust region method for training deep neural networks, where the objective function is approximated by a quadratic term with a norm constraint, known as the trust-region subproblem. The main difficulty in large-scale applications is how to solve this nonconvex subproblem efficiently (Yuan, 2000) . Approximate solvers such as (Steihaug, 1983; Powell, 1970) cannot converge to the exact minimizer, while the best method for computing the global minimizer has sub-linear convergence rate (Hazan & Koren, 2016) . We propose a new solver that is guaranteed to converge to a global minimum of the nonconvex trust region sub-problem with an asymptotic linear convergence rate. What's more, our method only requires Hessian-vector product, which only costs two forward-backward operations in neural networks. Empirically we find this technique making trust region method comparable to SGD in terms of running time.

Apart from convergence rate, generalization ability is another important issue to practitioners. Starting from (Zhang et al., 2016) which proposed the generalization paradox of deep models, intensive works have been done to find ways to explain and close generalization gap in large batch SGD (Keskar et al., 2016; Kawaguchi et al., 2017; Wilson et al., 2017) . Among them, (Keskar et al., 2016) empirically finds the generalization gap between large and small batch and later (Wilson et al., 2017) notices that adaptive methods like Adam, AdaGrad and RMSprop are even worse than SGD concerning large batch generalization. To close the generalization gap, (Chaudhari et al., 2016) changes Algorithm 1 (Stochastic) Trust region method Input: {F (x), ∇F (x), ∇ 2 F (x)}, or their stochastic counterparts; Initial guess x 0 , radius r 0 . Output: Suboptimal solution: x N +1 .

1: for t = 0, 1, • • • , N do 2:

Calculate gradient g = ∇F (x t ), H = ∇ 2 F (x t )

3:

Expand objective function:

EQUATION

Solve (1) to get z t , compute q t = F (xt)-F (xt+zt) f (0)-f (zt)

5:

(Deterministic version) Update iterate and radius:

• If q t ≥ 0.9, then x t+1 = x t + z t , r t+1 = 2r t .

• If q t ≥ 0.1, then x t+1 = x t + z t , r t+1 = r t .

• If q t < 0.1, then x t+1 = x t , r t+1 = 0.5r t .

6:

(Stochastic version) Update iterate and radius:

• If q t ≥ 0.9 and g ≥ r t , then x t+1 = x t + z t , r t+1 = min{2r t , r max }.

• Otherwise, x t+1 = x t , r t+1 = 0.5r t . 7: end for 8: return x N +1 . the loss function to eliminate sharp local minima, (De et al., 2016) selects batch size adaptively and (Hoffer et al., 2017) combines adaptive step size with new batch normalization layer.

In contrast to these works, we find through experiments that using trust region method alone can effectively escape from sharp minima and achieve lower test error than SGD, especially on largebatch training. For CIFAR-10 with VGG network, our trust region based method achieves 86.8% accuracy and our hybrid method have 88.1%, both of them are better than SGD's solution with 86.7% accuracy when batch size is 128. As batch size grows to 2048, our trust region method still has 85.5% accuracy and our hybrid method is 87.0%, in contrast, SGD only has 80.3% accuracy. Arguably, our method is a more natural solution to large batch training because it can match the small batch accuracy simply by switching to trust region method without changing network structure, loss function and batch size.

2. BACKGROUND AND RELATED WORK:



2.1. (STOCHASTIC) TRUST REGION METHOD:

Trust region (TR) method (Conn et al., 2000) (Algorithm 1) is a classical second order method for solving nonlinear optimization problems. At each iteration, TR forms a second order approximation around the current solution and updates the solution by minimizing this quadratic function within the "trust region". The trust region, usually a hyper ball, is chosen to make the second order approximation close enough to the original one and thus by minimizing the subproblem, the objective function will very likely to descent. More formally, suppose F (x) is the objective function, at iteration x t the subproblem can be written as:

EQUATION

Where H = ∇ 2 F (x t ), g = ∇F (x t ), and r > 0 is the radius of trust region. 1

If objective function F has finite sum structure,

F (x) = 1 n n i=1 f i (x) or more generally F (x) = E ξ f (x, ξ)

where ξ is random variable with bounded variance, then we can turn to stochastic trust region method discussed in (Chen et al., 2015) . Both deterministic and stochastic trust region method 1 In general the constraint can be • M for any PSD matrix M , but we only discuss the simpler case here because a change-of-variable z =

√ M z and H = √ M -H √ M -1

can transform it to (2).

have very similar algorithm as one can see in Algorithm 1. As to convergence rate, (Blanchet et al., 2016) shows that the stochastic version is guaranteed to find a stationary point ∇f (x) ≤ in O( -2 ) iterations, while the deterministic version is able to find -suboptimal stationary point within O( -3/2 ) iterations (Curtis et al., 2017) . For large scale problem when we cannot calculate the gradient and Hessian on the whole dataset, stochastic method is the only choice.

To apply trust region method to large-scale non-convex optimization (e.g., training deep networks), there are two difficulties: (a) The Hessian matrix H may not be positive definite, so (2) is a nonconvex problem with constraint and it is hard to compute its global minimizer. (b) H is very large (millions by millions) and cannot be formed explicitly. For (a), several algorithms have been proposed in the literature. Traditionally, dogleg method (Powell, 1970) for H 0 and Steihaugs method (Steihaug, 1983) for indefinite H can be applied to get an approximated solution. However if a more precise solution is desired, we need to find other ways, e.g. (Hazan & Koren, 2016) is able to find an -suboptimal solution in Õ(M/ √ ) time, where M is the matrix size (number of nonzero elements). In this paper, we design an efficient solver based purely on gradient information that can find an -suboptimal solution in asymptotically linear time Õ(M log(1/ )).

For (b), due to high dimensionality of H, matrix decomposition such as SVD will be in-feasible. Therefore, an efficient subproblem solver should only involve computations of Hessian-vector product, which can be computed by back-propagation in deep networks (Pearlmutter, 1994; Baydin et al., 2015) . Furthermore, since machine learning problems can usually be formulated as finite sum structure F (x) = 1 N N i=1 f i (x), instead of computing the exact Hessian we can compute the subsampled Hessian on the selected batch, so the subsampled Hessian-vector product computation is only slower than SGD computation by a constant factor. Our proposed subproblem solver only involves Hessian-vector product, so it can be applied to large-scale problems and replace SGD seamlessly. This subsampled version is justified by the stochastic trust region method described above.

2.2. OTHER SECOND ORDER METHODS FOR TRAINING NEURAL NETWORKS:

Another related technique for non-convex optimization is adaptive regularization with cubic (ARC) (Griewank, 1981) . Similar to TR method, ARC uses cubic term of positive coefficient σ 3 z 3 2 to replace constraint in (2):

EQUATION

)

ARC method can be seen as the "twin" method of TR, since they share many similar properties (e.g. converge rates). Intuitively the radius parameter r in (2) plays the same role as σ in ARC: a larger r and smaller σ both mean larger "trust-region". However, since the ARC subproblem is unconstrained, the treatment is relatively easy: for example, generalized Lanczos-type iterations (Gould et al., 1999) or gradient descent (Carmon & Duchi, 2016 ) can be used. Our proposed subproblem solver for TR is also based on gradient descent, but due to constraint in the TR subproblem our algorithm and proof are different from (Carmon & Duchi, 2016) .

Recently, several works have applied second order methods to neural network optimization, including (Martens, 2010; Botev et al., 2017; Xu et al., 2017b; a) . However, all of them only consider small-scale fully connected networks with few layers. Among them, (Xu et al., 2017b; a) adopt full gradient and subsampled Hessian and solve the trust region subproblem in the subspace spanned by Cauchy point and minimum eigenvector v 1 (H). We are different from this recent paper in two aspects. First, we develop a new trust region subproblem solver with better theoretical guarantee. Second, the full gradient computation in their algorithm is too expensive for large-scale deep network training. In contrast, we implement our algorithm with subsampled gradient and Hessian, thus we are able to run large-scale experiments and observe better generalization compared with SGD.

2.3. GENERALIZATION ABILITY FOR FIRST ORDER METHODS:

SGD and other stochastic first order methods are popular for training large scale neural networks since gradient for each training sample can be computed in one forward-backward cycle. Recently there are many works on the effect of batch size on training time as well as generalization ability. Indeed, a larger batch makes more use of parallel processors, so the program can scale perfectly to more GPUs. Moreover, more samples make the gradient less noisy and thus a larger learning rate and higher converge rate are possible, see (Goyal et al., 2017; You et al., 2017a; b) and references therein.

At the same time, large batch in SGD is harmful to generalization ability, experiments in (Keskar et al., 2016) show that batch size can affect the geometry around the solution that SGD finds. Specifically, SGD with a larger batch size tends to find sharper local minima, so even a tiny change of data from training set to test set will lead to high loss, i.e. bad generalization ability. As to the converge rate to local minima, although (Lee et al., 2016) shows gradient descent always converges to local minima, (Du et al., 2017) proves the converge process can take exponential time. On the other hand by adding extra noise, (Ge et al., 2015) claims SGD is able to escape strict saddle in poly(d/ ) iterations, and further (Jin et al., 2017) design a full gradient with noise method that improves the result to poly-log(d)/ 2 steps.

Notations: Throughout this paper, we use S n-1 to denote the unit sphere in R n . The model parameters x ∈ R d and objective function

F (x) = 1 N N i=1 f i (x)

is in C 2 ; its gradient and Hessian at step x t are denoted as g t = ∇F (x t ) and H t = ∇ 2 F (x t ), and if the index t is unimportant we will simplify them as g and H respectively. For the Hessian matrix H, its eigenvalue decomposition is

H = d i=1 λ i v i v i such that v i v j = δ ij and λ 1 ≤ λ 2 ≤ • • • ≤ λ d , accordingly we decompose any vector a ∈ R d as a = d i=1 a (i) v i . Denote operator norm β = H op = max{|λ 1 |, |λ d |}.

• represents 2-norm if not stated explicitly, I d ∈ R d×d is the identity matrix. Finally, following (Carmon & Duchi, 2016) we use s to represent the global minimizer of trust region subproblem (2).

3. EFFICIENT TRUST REGION SUBPROBLEM SOLVER:

Algorithm 2 Proposed trust region subproblem solver Input: Gradient g, Hessian matrix H Output: Approximated solution z 1: Initialize z 0 and η so that Assumption 1,2 are met. 2: boundary=false 3: for t = 0, 1, . . . , N 1 do 4: 

if z t < 1 then 5: z t+1 = z t -η(

z t +1 = z t -α t (I d -z t z t )(Hz t +g) z t -α t (I d -z t z t )(Hz t +g)

16: end for 17: return z N2+1

In this section, we show an efficient gradient-descent based algorithm with proper initialization can find the global minimum of the trust region subproblem (2). If the global minimum lies inside of the sphere then gradient descent itself is guaranteed to find it; Otherwise we first conduct gradient descent until the iterate hits the spherical constraint, then a manifold gradient descent on the sphere can converge to the solution. We prove this simple procedure can return the global minimum of the non-convex trust region subproblem and has asymptotically linear convergence rate. The details are shown in Algorithm 2.

3.1. PROPERTIES OF GLOBAL MINIMUM:

The necessary and sufficient condition of global minimum comes from KKT condition, see (Wright & Nocedal, 1999) for details. Lemma 1. (Global minimum) s is the global minimum of (2) if and only if s ≤ 1 and there is a scalar λ ≥ 0 such that: gradient condition: (H + λI d )s + g = 0, complementary slackness: λ(1 -s 2 ) = 0, Hessian: H + λI d 0.

(4)

Proposition 1. Based on (4) we can describe the solution(s) as follows:

• λ 1 > 0: only one global minimum. If H -1 g < 1 then s < 1, otherwise s = 1.

• λ 1 = 0: if g (1) = 0, only one solution with s = 1; Otherwise if g (1) = 0 and

n i=2 ( g (i) λi ) 2 ≥ 1, one solution with s = 1; If g (1) = 0 and n i=2 ( g (i)

λi ) 2 < 1, there are multiple solutions.

• λ 1 < 0, g (1) = 0: only one solution.

• λ 1 < 0, g (1) = 0: if n i=2 ( g (i)

λi-λ1 ) 2 ≥ 1 then only one global minimum, otherwise there are multiple ones.

Proof details are postponed to appendix. Since when λ 1 ≥ 0 every stationary points are global minimum, and further if s < 1 then we only need to use gradient descent to solve it (that the case for Line 11 in Algorithm 2), otherwise s = 1 and our manifold based algorithm also applies, since in this case subproblem is strongly convex, we don't need to worry about converging to suboptimal point. Now we only need to consider about λ 1 < 0 case as follows. Notice that in this case we always have s = 1. In the following lemma we try to distinguish global minimum from other stationary points, restricted to g (1) = 0 case (we call it "easy-case", as oppose to "hard-case" in (Wright & Nocedal, 1999) ). As to the "hard-case" g (1) = 0, in theory we can apply a small perturbation to b as (Carmon & Duchi, 2016) does: b = b + ε where ε is a small Gaussian noise. In practice due to rounding error this case is hardly seen, for best efficiency we choose to ignore it.

The following lemma gives a sufficient condition of global minimum: Lemma 2. When g (1) = 0 and λ 1 < 0, among all stationary points if s (1) g (1) ≤ 0 then s is the global minimum.

3.2. CONVERGENCE OF ITERATION:

Now it left to see how the solution found by Algorithm 2 will meet Lemma 2. To this end, we need to enforce following assumptions:

Assumption 1. (Bounded step size) Step size η < 1/β. Assumption 2. (Initialize) z 0 = -α g g , 0 < α < 1.

We remark that Assumption 2 is a good guess of global minimum if ignoring the curvature information, and using this alone as subproblem solution will reduce trust region method to gradient descent. Under these assumptions, Algorithm 2 is guaranteed to find (one of) the global minimum, which is formally stated in the following theorem: Theorem 3. Under proximal gradient descent update:

z t+1 = Prox • z t -η∇f (z t ) , and As- sumption 1 if z (i) t g (i) ≤ 0 then z (i) t+1 g (i) ≤ 0.

Combining with Assumption 2 and Lemma 2, if λ 1 < 0, g (1) = 0 then z t converges to a global minimum s.

To see the converging process more clearly, we divide the iterations into two phases: in the first phase z t stays strictly inside the sphere S n-1 : z t < 1, and during this phase we will show that z t is monotone increasing until it hits the sphere and that is when the second phase begins. In the second phase the iterates adhere to the sphere and converge to the global minimum with asymptotically linear rate. First of all we show the monotone increasing property based on following lemma: Lemma 4. For z t defined above, we have z t H∇f (z t ) ≥ βz t ∇f (z t ) (recall we define β as the operator norm of H). Now imagine there is another iterate zt that does "plain" GD (i.e. without projection): zt+1 = zt -η∇f (z t ), using the same step size as z t and same initialization z0 = z 0 . Such a iteration rule guarantees that as long as z t < 1 for all z ≤ τ then zt = z t . By showing zt is monotone increasing, we actually prove monotone increasing property in the first phase.

Theorem 5. Suppose z t is in the region such that proximal gradient update equals to plain GD: zt+1 = zt -η∇f (z t ), then under this update rule, zt is monotone increasing.

Another key observation from Theorem 5 is the unboundedness of zt , i.e. zt → ∞ as t → ∞. Indeed we have the following lemma: Lemma 6. (Finite phase I) Assuming λ 1 < 0, suppose t * is the index that zt * < 1 and zt * +1 ≥ 1, then t * is bounded by:

t * ≤ log(1 -ηλ 1 ) -1 log 1 η|g (1) | - 1 ηλ 1 -log -z

(1) 0

ηg (1) -

EQUATION

)

So after t * iterations the algorithm will reach phase II. Then we turn to manifold optimization method, which is discussed extensively in (Absil et al., 2009) and for the manifold theory we refer (Do Carmo & Flaherty Francis, 1992 ). Here we list some concepts and its explicit form in our problem: denote M as the smooth manifold, and z ∈ M can be any point on manifold M. The tangent space of z, denoted as T z , is the set of all tangent vectors to M at z.

When M = S n-1 then T z = {ξ ∈ R n : z ξ = 0}; Retraction is a mapping R z (ξ) : ξ ∈ T z → M, in S n-1 mani- fold, one of the commonly used retraction is R z (ξ) = z+ξ z+ξ ; Projection onto T z , denoted as P z ξ is a mapping from R n to T z , for S n-1 the projection is simply P z (ξ) = (I d -zz )ξ.

Based on above concepts, we can write the gradient descent on M as: z t+1 = R zt (-α t gradf (z t )), and α t is tactically selected by line search such that f (z t ) -f (z t+1 ) ≥ σα t gradf (x t ) , where σ ∈ (0, 1).

Because we already know at least one of the global minimum lies on S n-1 , and according to Lemma 6 after at most t * iterations we start to use manifold based optimization, i.e. the subproblem becomes:

EQUATION

Therefore, our method shrinks the search space from z ≤ 1 to z = 1, by doing so we can apply well studied manifold optimization theory (Absil et al., 2009; Udriste, 1994) to our problem. Indeed we have following theorems:

Theorem 7. Let {z t } be an infinite sequence of iterates generated by line search gradient descent, then every accumulation point of {z t } is a stationary point of the cost function f .

This above only guarantees convergence to stationary points, however, according to Theorem 3 if the step size α t is not too large, it actually converges to global minimum. Then it remains to show a linear convergence rate, as guaranteed by the following theorem:

Theorem 8. Let {z t } be an infinite sequence of iterates generated by line search gradient descent, suppose it converges to s. Let λ H,min and λ H,max be the smallest and largest eigenvalues of the Hessian at s. Assume that s is a local minimizer then λ H,min > 0 and given r in the interval (r * , 1) with r * = 1 -min 2σ ᾱλ H,min , 4σ(1 -σ)β λ H,min λ H,max , there exists an integer K such that:

f (z t+1 ) -f (s) ≤ r f (z t ) -f (s) ,

for all t ≥ K.

Remarks: λ H,min and λ H,max are the minimum and maximum eigenvalue of Riemannian Hessian. Specifically the Riemannian Hessian can be calculated by(Proposition 5.5.4, Absil et al. (2009) ):

EQUATION

By direct calculation(in appendix) we shall see: Hessf (s)[ξ], ξ = -s Hs + ξ Hξ -g s, where ξ ∈ Null(s), ξ = 1. By optimal condition (4), we have:

s (H + λI d )s + s g = 0 ⇒ -s Hs -g s = λ, (8) so Hessf (s)[ξ], ξ = λ + ξ Hξ ≥ λ + λ 1 ! > 0.

Where ! > is guaranteed by gradient condition in (4):

EQUATION

in "easy-case", g (1) = 0, so λ 1 + λ = 0 and Hessian condition in (4) can be improved to λ 1 + λ > 0.

Based on above discussion, we know λ H,min ≥ λ + λ 1 > 0 and λ H,max ≤ λ + λ n .

4. EXPERIMENTS:

In this section, we examine the performance of our algorithm. First of all, we use a random generated problem to check the dynamics of our proposed trust region subproblem solver, then we compare the performance of our trust region method with SGD on deep convolutional neural network. After that, we focus on generalization ability of the solution that trust region method returns. We include following algorithms in our experiments:

• SGD: We tune the step size to get the fastest, stable convergence.

• TR: Trust region method with our proposed subproblem solver. For each subproblem, we choose N 1 = N 2 = 1 (defined in Algorithm 2). The trust region radius is set to 1 and the step size of subproblem solver is set to 0.1. • Hybrid: Our proposed hybrid method, it runs trust region method until the accuracy no longer increases, and then uses the last epoch to initialize SGD method. When the training accuracy stabilizes, it returns the final model. Both TR and SGD here follow the same hyper-parameters above.

4.1. SOLVING THE TRUST REGION SUBPROBLEM:

We sample an indefinite random matrix by H = BB -λI n , where B ∈ R n×(n-1) and B ij iid ∼ N (0, 1), obviously λ min (H) = λ 1 = -λ. Afterwards we sample a vector g by g i iid ∼ N (0, 1). By changing the value of λ in {10, 30, 50, 70, 90, 110}, we plot the function value decrement with respect to number of iterations in Figure 1 (left). As we can see, the iterates first stay inside of the sphere (phase I) for a few iterations and then stay on the boundary (phase II). To inspect how λ changes the duration of phase I, we then plot the number of iterations it takes to reach phase II, under different λ values shown in Figure 1 (middle). Recall in (5), number of iterations is bounded as a function of λ, which can be further simplified to:

EQUATION

)

where we set z(1) 0 = 0 to simplify the formula. By fitting the data point with function T (λ) = log(1+c1λ) log(1+c2λ) , we find our bounds given by Lemma 6 is quite accurate.

4.2. OPTIMIZING DEEP CNNS:

Next, we test the performance of our proposed trust region method (actually stochastic is used, denoted as TR) on training deep CNNs, including VGG (Simonyan & Zisserman, 2014) and Mo-bileNets (Howard et al., 2017) . VGG/MobileNets are adopted to classify CIFAR10/STL10 data.

During all experiments, we only compare trust region method with plain SGD. We did not include other SGD variants because (1) We want to focus on whether second order information is useful without considering the effect of momentum/acceleration. (2) It is known that SGD generalize better than adaptive methods in large-batch (Wilson et al., 2017) . For reference, we also list the best accuracy obtained by Adam with small batch size in Table 1 , which shows Adam's accuracy is only slightly better than SGD but still worse than our hybrid method. We have some tricks to make it possible to apply trust region method to a deep convolutional network: First of all, rather than taking ReLU as the activation function, we use a similar yet everywhere continuously differentiable function, namely SReLU (see Figure 1 (right)):

SReLU(x) = √ x 2 + α + x 2 , α > 0.

Although we can also define Hessian on ReLU function, it is not well supported on major platforms (Theano/PyTorch). Likewise, we find max-pooling is also not supported by platforms to calculate higher order derivative, one way to walk around is to change all the max-pooling layers to avgpooling-it hurts accuracy a little bit (we observe 1∼2% accuracy loss compared with (Keskar et al., 2016) ), albeit this is not our primary concern. Secondly, the oracle in TR method is Hessianvector product, or Hv. Theoretically, the cost of Hv is comparable to gradient oracle because Hv = ∂ ∂w (v ∂f ∂w ). However, due to inefficient implementation, in many deep learning frameworks the cost of Hv is ∼7x slower than gradient oracle, making running time comparison unfavorable, so we use the following numerical differentiation to replace chain-rule:

Hv = ∇f (x + εv) -∇f (x) /ε, where ε → 0.

Experiments show the relative error is controllable ( Hv -Hv / Hv ≤ 1%). Nevertheless, the Hessian-vector operation is still a bottleneck to performance, thus if we care about running time then 1∼2 inner iteration(s) for each subproblem is more suitable. Lastly, although in theory we need full gradient and full Hessian to guarantee convergence (in Xu et al. (2017b) they only need full gradient), calculating them in each iteration is not practical, so we calculate both Hessian and gradient on subsampled data to replace the whole dataset.

In the first experiment, we compare the training/testing accuracy w.r.t. running time and epoch on VGG16+CIFAR10/STL10 dataset. For SGD we choose step size η = 0.1, which is the best step size in {10 -n | n = 1, 2, . . . }. For trust region subproblem we choose step size η = 0.1 for all the experiments. Both of them use batch size B = 512(for CIFAR10) and B = 1024(for STL10). Our machine has 4 Titan Xp GPUs and Xeon E5-2620 CPU. The experimental results are shown in Figure 2 . Without surprising, the trust region method is ∼3x slower than SGD, partly because we use subsampled gradient/Hessian (currently there is no theoretical result to guarantee convergence in this situation) and Hv operation is too expensive. Despite the slower training time, we observe TR converges to a solution with better test error.

4.3. BETTER GENERALIZATION ABILITY IN LARGE BATCH TRAINING:

Training neural network with larger batch size has become an important issue for faster training. To fully exploit the computational power of multi-GPU systems, we need to increase the batch size, ). We find TR method has better test accuracy on both datasets, although the running time is 2∼3x longer. at least proportional to number of cores. However, large batch size can affect the generalization ability (Keskar et al., 2016) , so it would be interesting to see if this dilemma can be solved merely by choosing a different optimization method. To do so, we design an experiment that compares testing accuracy under different batch sizes in Figure 3 . We observe that our method (TR) converges to solutions with much better test error when batch size is larger than 128. We postulate this is because SGD has not enough noise to escape from a sharp local minimum, especially when the large batch is used. For smaller batch size, TR method is not as good as SGD method, this is because the approximation of gradient and Hessian is not good enough by only a few samples. However, our Hybrid method combines the benefits of TR and SGD methods to make it both fast and accurate, for both small and large batch size.

4.4. LOSS LANDSCAPE OF DIFFERENT SOLUTIONS:

To explore whether TR/Hybrid methods find the wide local minimum, we follow the experiment techniques in (Keskar et al., 2016) to draw the loss curve that connects two models. Specifically, if we use w SGD , w TR and w Hybrid to denote models of SGD, trust region and hybrid method respectively, then the interpolated model in between is

EQUATION

The loss and accuracy can be evaluated on both training set and test set by parameters w , and the result is displayed in Figure 4 . We also calculate the 2 -distance between model A and B, shown in titles of sub-figures. Figure 4 : Loss(blue lines) and accuracy(red lines) curve of interpolated models, the x-label is α defined in (11), y-label in the left is the averaged loss and y-label in the right is accuracy. We also calculate the distance of models between α = 0 and α = 1 in the titles.

Figure 4 clearly shows the width and relative location of local minima: models of TR and hybrid methods lie in the same wide local minimum, while SGD finds a much sharper local minimum. Moreover, we notice Dist(TR, SGD) ≈ Dist(Hybrid, SGD) Dist(Hybrid, TR), implying that hybrid method is a refinement of TR method (both of them are close to the exact wide local minimum, but hybrid method is even closer). In contrast, SGD converges to a very different solution, where the loss curve is sharp and the model does not generalize as good as the other two.

4.5. PROPERTIES OF LARGE BATCH SGD VS TR:

By far, we can only conclude that trust region method generalize better than SGD, especially when batch size is large. In what follows, we design a novel experiment that help us identifying the differences between those two methods. Rather than running a homogeneous algorithm, we interleave SGD and TR as two tracks, for simplicity we denote SGD(n) as running n epochs SGD, and TR(n) as running n epochs TR. In the first track, we run SGD(50) in the beginning, then we use the last iteration in SGD(50) to initialize TR(100), lastly based on the latest iterate, run SGD(50) again. In the second track we change the order: Begin with TR(100), then SGD(50), and end with TR(50).

The purpose of switching between TR and SGD methods is to detect the character of respective solutions they find. Before running experiment we do such "thought experiment": Imagine if large batch SGD converges to a sharp local minimum, and TR method can successfully escape it, then initialized by the solution that SGD finds, TR method will "climb over" the loss hill and down to a wide minimum, during this process the loss function (as well as accuracy) first increases and then decreases. On the contrary if TR method already finds a wide local minimum, then initialized by that solution SGD shall not escape to other minima, so we don't expect a sharp change in either loss or accuracy. This thought experiment is illustrated in Figure 5 .

Our experimental results are presented in Table 1 , where we run VGG16+CIFAR10 and choose batch size B = 2048. The overall running process is shown in the picture, and we further extract some important data into the table in order to see the differences more clearly.

We explain the results from following angles:

1. As expectation, there is a sharp raise in both train loss and test error at the 50th epoch in SGD(50)-TR(100)-SGD(50) track, meaning trust region method does escape from the solution that SGD(50) converges to. Meanwhile both loss function and test error change smoothly at the 100th epoch in TR(100)-SGD(50)-TR(50) track, which is understandable because SGD can't escape the basin of wide minimum obtained by TR. 100)→SGD( 50), in Track 2 we change the order to TR(100)→SGD( 50)→TR(50), so both tracks run 200 epochs in total. Testing accuracy are reported by the end of each stage. For both tracks we choose batch size B = 2048 and CIFAR-10 dataset. As a reference, we also run Adam algorithm on small batch (B = 64).

2. Warm-started by TR(100), SGD can reach the best testing accuracy, moreover, the accuracy merely change by enlarging batch size. This sets the foundation of our hybrid method in Figure 3 , where the testing accuracy only drops 0.35% from B = 64 to B = 1024, in comparison SGD testing accuracy drops 6.97%.

3. Due to subsampled Hessian and gradient, TR method can only fluctuate around the minimizer, this causes a 1.6% drop in testing accuracy compared to the best result. Another clue is the 1.0% drop in testing accuracy and raise of training loss at 150th epoch in TR(100)-SGD(50)-TR(50) track, when SGD(50)→TR(50) happens.

solving this geometric series, we get:

w (i) t = (1 -ηλ i ) t w (i) 0 - 1 ηλ i + 1 ηλ i (14) suppose at t-th iteration, w (i) t ≥ w (i)

t+1 which is equivalent to:

EQUATION

)

because from Assumption 2 we know w

(i) 0 ≥ 0, if w (i)

0 -1 ηλi ≤ 0, i.e. 0 < ηλ i ≤ 1/w (i) 0 then by ( 15) we know 1 -ηλ i ≥ 1 ⇔ ηλ i ≤ 0 leading to a contradiction, so w (i) 0 -1 ηλi > 0 and 1 -ηλ i ≤ 1. On the other hand, λ j ≥ λ i for j ≥ i, so 1 -ηλ j ≤ 1, together with w

(j) 0 -1 ηλj = z (i) 0 -1 ηλj ≥ w (i)

0 -1 ηλi > 0 we conclude:

w (j) 0 - 1 ηλ j ≥ (1 -ηλ j ) w (j) 0 - 1 ηλ j ⇔ w (j) t ≥ w (j)

t+1 for j ≥ i.

For any t, suppose i * ∈ [n] is the smallest coordinate index such that w

(i * ) t ≥ w (i * )

t+1 , which implies w i t < w i t+1 for any i < i * and w i t ≥ w i t+1 for any i ≥ i * (such a i * may not exist, but it doesn't matter). By analyzing the sign of z t we know:

sgn z i t z (i) t -z (i) t+1 = sgn w (i) t w (i) t -w (i) t+1 = sgn w (i) t -w (i)

t+1 , finally we have:

z t A∇f (z t ) = 1 η i * -1 i=1 λ i z (i) t (z (i) t -z (i) t+1 ) + 1 η n i=i * λ i z (i) t (z (i) t -z (i) t+1 ) ≥ λ i * -1 η i * -1 i=1 z (i) t (z (i) t -z (i) t+1 ) + λ i * η n i=i * z (i) t (z (i) t -z (i) t+1 ) ≥ λ i * η n i=1 z (i) t (z (i) t -z (i) t+1 )

≥ βz t ∇f (z t ).

(16)

A.5 PROOF OF THEOREM 5

Proof. First of all, notice zt+1 2 = zt 2 -2ηz t ∇f (z t ) + η 2 ∇f (z t ) 2 , so it remains to show z t ∇f (z t ) ≤ 0. We prove this by induction rule, suppose z t-1 ∇f (z t-1 ) ≤ 0 and from zt = zt-1 -η∇f (z t-1 ) we know: z t ∇f (z t ) =z t-1 ∇f (z t-1 ) -η ∇f (z t-1 ) 2 -ηz t-1 A∇f (z t-1 )

(1) + η 2 ∇f (z t-1 ) A∇f (z t-1 )

(2)

.

From Lemma 4 we know (1) ≥ β zt-1 ∇f (z t-1 ) and recall β is the operator norm of A, we have (2) ≤ β ∇f (z t-1 ) 2 , so:

EQUATION

naturally, by choosing η < 1/β we have z t ∇f (z t ) ≤ 0 for all t, based on this observation zt+1 is monotone increasing. 

:



Figures and Tables:

Figure 1: Left: Trust region experiment, we use solid lines to indicate iterations inside the sphere and dash lines to indicate iterations on the sphere. By changing λ we can modify the function curvature. Middle: #Iteration it takes to reach sphere under different λ's, we also fit the curve by model T = log(1+c1λ) log(1+c2λ) derived in Lemma 6. Right: Comparing ReLU with SReLU (see definition in Section 4.2), as α closes to zero, SReLU looks more like ReLU.

Figure 2: Training/Testing error with both time and epoch. The first row is CIFAR10 dataset(batch size B = 512) and the second row is STL10 dataset(batch size B = 1024). We find TR method has better test accuracy on both datasets, although the running time is 2∼3x longer.

Figure3: We choose batch size from {64, 128, 256, 512, 1024, 2048}, for each batch size and each algorithm, we independently run 5 times to gather mean accuracy and deviation (in order to show significance). The experiment is conducted on VGG16+CIFAR10.

Figure 5: Illustration of imaginary process of Track 1(left) and Track 2(right). Note that we use subsampled Hessian and gradient, so the iterate of trust region method will fluctuate around local minimum.

Figure 6: Experiments on two tracks as described inTable 2, we display training loss and test error w.r.t epochs, the batch size is B = 2048.

Figure 7: Experiments on two tracks as described inTable 3, we display training loss and test error w.r.t epochs, the batch size is B = 1024.

if boundary = false then return z N1+1 12: end if 13: for t = t, t + 1, . . . , N 2 do

Table content: <table><tr><td>6:</td><td>else</td></tr><tr><td>7:</td><td>boundary = true</td></tr><tr><td>8:</td><td>break</td></tr><tr><td>9:</td><td>end if</td></tr><tr><td colspan="2">10: end for</td></tr><tr><td>11: 14:</td><td>Choose α t by Armijo line search</td></tr><tr><td>15:</td><td/></tr></table>

Testing accuracy in different stages. In Track 1 we run SGD(50)→TR(

Table content: <table><tr><td/><td>Stage I</td><td/><td colspan="2">Stage II</td><td colspan="2">Stage III</td></tr><tr><td/><td>method</td><td>accuracy</td><td>method</td><td>accuracy</td><td>method</td><td>accuracy</td></tr><tr><td>Track 1</td><td>SGD(50)</td><td>80.1%</td><td>TR(100)</td><td>86.2%</td><td>SGD(50)</td><td>87.8%</td></tr><tr><td>Track 2</td><td>TR(100)</td><td>85.6%</td><td>SGD(50)</td><td>87.5%</td><td>TR(50)</td><td>86.5%</td></tr><tr><td>Adam</td><td>ADAM(100)</td><td>86.6%</td><td/><td/><td/><td/></tr></table>

Table 2, we display training loss and test error w.r.t epochs, the batch size is B = 2048. Testing accuracy in different stages.In Track 1 we first run 150 epochs SGD then 300 epochs TR followed by 50 epochs SGD; In Track 2 we change the order to first 300 epochs TR then 50 epochs SGD and end by 150 epochs TR, so both tracks run 500 epochs in total. Testing accuracy are reported by the end of each stage. Here we choose batch size B = 1024.

Table content: <table><tr><td/><td colspan="2">Stage I</td><td colspan="2">Stage II</td><td colspan="2">Stage III</td></tr><tr><td/><td>method</td><td>accuracy</td><td>method</td><td>accuracy</td><td>method</td><td>accuracy</td></tr><tr><td colspan="2">Track 1 SGD(150)</td><td>62.4%</td><td>TR(300)</td><td>64.6%</td><td>SGD(50)</td><td>66.0%</td></tr><tr><td>Track 2</td><td>TR(300)</td><td>63.6%</td><td>SGD(50)</td><td>66.3%</td><td>TR(150)</td><td>65.8%</td></tr></table>

Table 3, we display training loss and test error w.r.t epochs, the batch size is B = 1024.

Table content: <table/>

Peer review (iclr2018):
 The paper proposes training neural networks using a trust region method, in which at each iteration a (non-convex) quadratic approximation of the objective function is found, and the minimizer of this quadratic within a fixed radius is chosen as the next iterate, with the radius of the trust region growing or shrinking at each iteration based on how closely the gains of the quadratic approximation matched those observed on the objective function. The authors claim that this approach is better at avoiding "narrow" local optima, and therefore will tend to generalize better than minibatched SGD. The main novelty seems to be algorithm 2, which finds the minimizer of the quadratic approximation within the trust region by performing GD iterations until the boundary is hit (if it is--it might not, if the quadratic is convex), and then Riemannian GD along the boundary. The paper contains several grammatical mistakes, and in my opinion could explain things more clearly, particularly when arguing that the algorithm 2 will converge. I had particular difficulty accepting that the phase 1 GD iterates would never hit the boundary if the quadratic was strongly convex, although I accept that it is true due to the careful choice of step size and initialization (assumptions 1 and 2). The central claim of the paper, that a trust region method will be better at avoiding narrow basins, seems plausible, since if the trust region is sufficiently large then it will simply pass straight over them. But if this is the case, wouldn*t that imply that the quadratic approximation to the objective function is poor, and therefore that line 5 of algorithm 1 should shrink the trust region radius? Additionally, at some times the authors seem to indicate that the trust region method should be good at escaping from narrow basins (as opposed to avoiding them in the first place), see for example the left plot of figure 4. I don*t see why this is true--the quadratic approximation would be likely to capture the narrow basin only. This skepticism aside, the experiments in figure 2 do clearly show that, while the proposed approach doesn*t converge nearly as quickly as SGD in terms of training loss, it does ultimately find a solution that generalizes better, as long as both SGD and TR use the same batch size (but I don*t see why they should be using the same batch size). How does SGD with a batch size of 1 compare to TR with the batch sizes of 512 (CIFAR10) or 1024 (STL10)? Section 4.3 (Figure 3) contain a very nice experiment that I think directly explores this issue, and seems to show that SGD with a batch size of 64 generalizes better than TR at any of the considered batch sizes (but not as well as the proposed TR+SGD hybrid). Furthermore, 64 was the smallest batch size considered, but SGD was performing monotonically better as the batch size decreased, so one would expect it to be still better for 32, 16, etc. Smaller comments: You say that you base the Hessian and gradient estimates on minibatched samples. I assume that the same is true for the evaluations of F on line 4 of Algorithm 1? Do these all use the same minibatch, at each iteration? On the top of page 3: "M is the matrix size". Is this the number of elements, or the number of rows/columns? Lemma 1: This looks correct to me, but are these the KKT conditions, which I understand to be first order optimality conditions (these are second order)? You cite Nocedal & Wright, but could you please provide a page number (or at least a chapter)? On the top of page 5, "Line 10 of Algorithm 1": I think you mean Line 11 of Algorithm 2.