Title: EXPLORING CURVATURE NOISE IN LARGE-BATCH STOCHASTIC OPTIMIZATION

Abstract: Using stochastic gradient descent (SGD) with large batch-sizes to train deep neural networks is an increasingly popular technique. By doing so, one can improve parallelization by scaling to multiple workers (GPUs) and hence leading to significant reductions in training time. Unfortunately, a major drawback is the socalled "generalization gap": large-batch training typically leads to a degradation in generalization performance of the model as compared to small-batch training. In this paper, we propose to correct this generalization gap by adding diagonal Fisher curvature noise to large-batch gradient updates. We provide a theoretical analysis of our method in the convex quadratic setting. Our empirical study with state-of-the-art deep learning models shows that our method not only improves the generalization performance in large-batch training, but furthermore, does so in a way where the training convergence remains desirable and the training duration is not elongated. We additionally connect our method to recent works on loss surface landscape in the experimental section.
Modern datasets and network architectures in deep learning are becoming increasingly larger which results in longer training time. An ongoing challenge for both researchers and practitioners is how to scale up deep learning while keeping training time manageable. Using stochastic gradient descent (SGD) with large batch-sizes offers a potential avenue to address scalability issues. Increasing the batch-size used during training improves data parallelization (Goyal et al., 2017; You et al., 2018) ; it ensures multiple processors (GPUs) have sufficiently useful workload at each iteration hence reducing communication-to-cost ratio.
Unfortunately, a severe limitation to employing large-batch training in practice is the so-called "generalization gap". While large-batch yields considerable advantages over small-batch on training loss and error per parameter update, it has been verified empirically in LeCun et al. (2012) ; Keskar et al. ( 2016) that we have the opposite effect for test loss and error. To fully realize the benefits of using large-batches in the distributed synchronous setup, it is necessary to engineer large-batch training in a way such that this generalization gap can be closed without sacrificing too much the training performance. This is precisely the central objective of our paper.
In this paper, we propose to add diagonal Fisher curvature noise to large-batch gradient updates. We discuss the motivations behind our approach. Under the typical log-likelihood loss assumption, the difference of large-batch and small-batch gradients can be modeled as a Fisher noise. We can expect that adding this noise directly to large-batch gradients gives small-batch performance. While this remedies the generalization issues associated with large-batch training, the resulting convergence performance will be undesirable. To attain our end goal of designing a method which enjoys both desirable convergence and generalization performance, we reduce the amplitude of the noise by changing the covariance structure from full Fisher to diagonal Fisher. We find that in practice, this improves convergence considerably while maintaining good generalization.
We give a theoretical analysis of our proposed method in Section 3. This is done over the convex quadratic setting which often serves as an excellent "proxy" for neural network optimization (Martens, 2010). We additionally provide numerical experiments on standard deep-learning tasks in Section 4 to demonstrate the efficacy of our proposed method. There are two primary takeaways from our empirical analysis. First, we show adding diagonal Fisher noise to large-batch

:

Under review as a conference paper at ICLR 2019 Observe that the full Fisher noise direction is perpendicular to the contours of the loss surface. Moreover, full Fisher exhibits slower convergence than diagonal Fisher; we refer to Section 3.2 for a more detailed analysis.

gradients does not hinder much the training performance; while at the same time giving a validation accuracy comparable to small-batch training, thereby correcting the "generalization gap". Second, this validation accuracy was attained in the same number of epochs used for small-batch training. This indicates that our method is data-efficient and does not require any lengthening of the optimization process.

2. PRELIMINARIES:

2.1 BACKGROUND Excess risk decomposition. Let S = {(x 1 , y 1 ), . . . , (x N , y N )} be a training set of N input-target samples drawn i.i.d. from an unknown joint probability distribution D. The family of classifiers of interest to us are neural network outputs f (x i , θ), where θ ∈ Θ ⊂ R d are parameters of the network and Θ here denotes the parameter space of the network. Let (f (x i , θ), y i ) be the loss function measuring the disagreement between outputs f (x i , θ) and targets y i . For convenience, we use the notation i (θ k ) to denote (f (x i , θ), y i ). The expected risk and empirical risk functions are defined to be L (θ) ≡ E (x,y)∼D [ (f (x, θ) 

, y)], L(θ) ≡ 1 N N i=1 i (θ).

The standard technique to analyze the interplay between optimization and generalization in statistical learning theory is through excess risk decomposition. After k iterations, the excess risk is defined as

L(θ k ) ≡ L (θ k ) -inf θ∈Θ L (θ).

From Bottou & Bousquet (2008) ; Chen et al. (2018) , the expected excess risk is upper-bounded by

EQUATION

where θ * = argmin θ L(θ) here is the empirical risk minimizer. The terms E gen and E opt are the expected generalization and optimization errors respectively. In the machine learning literature, optimization algorithms are often studied from one perspective: either optimization or generalization.

The excess risk decomposition suggests that both aspects should be analyzed together (Chen et al., 2018; Bottou & Bousquet, 2008) since the goal of a good optimization-generalization algorithm in machine learning is to minimize the excess risk in the least amount of iterations.

Related Work. In the context of large-scale learning, stochastic algorithms are very popular compared to full-batch methods due to lower computational overhead (Bottou et al., 2018; Bottou, 1991) . Despite a decay in optimization performance (slower rate of convergence in the convex case (Moulines & Bach, 2011; Bottou et al., 2018) ), stochastic algorithms possess good generalization properties due to the inherent noise of their gradients. There are vast bodies of research literature devoted to understanding the connection between inherent noise and generalization; for example, scaling the learning rate or batch-size (Smith & Le, 2017; Goyal et al., 2017; Hoffer et al., 2017) or studying the covariance structure of mini-batch gradients (Jastrzebski et al., 2017; Xing et al., 2018; Zhu et al., 2018; Li et al., 2015) . In the non-convex setting, the inherent noise allows SGD to efficiently escape from saddle points or shallow local minima which tends to give poor generalization (Zhu et al., 2018; Daneshmand et al., 2018; Jin et al., 2017) . The trade-off between optimization and generalization was also observed for large-batch versus small-batch training in deep learning: both large-batch and small-batch training can reach similar training loss after sufficient number of iterations. Since large-batch gradients have smaller variance, large-batch typically requires less iterations to reach an optimum. However, large-batch training usually has worse test performance compared to small-batch (Hoffer et al., 2017; Masters & Luschi, 2018; Keskar et al., 2016) . It was recently observed that generalization and optimization should not be decoupled in deep learning (Neyshabur et al., 2014; 2017a; b) . The choice of the batch-size has a direct impact on the trade-off between optimization and generalization.

In this paper, we aim to design a novel algorithm for large-batch training by increasing the variance of its gradients such that the resulting generalization performance matches small-batch gradient descent. Moreover, we require that this generalization performance be achieved within a number of iterations comparable to original large-batch.

2.2. MOTIVATION:

We tackle the large-batch generalization problem through gradient noise injection. Let B L denote large-batch and M L = |B L | denote size of the large-batch. Consider the following modification of large-batch SGD updates

EQUATION

)

where α k is the learning rate,

∇L M L (θ k ) = 1 M L M L i=1 i (θ k )

is the large-batch gradient and N (0, I d ) is the multivariate Gaussian distribution with mean zero and identity convariance. We can interpret this modification as injecting a Gaussian noise with mean zero and covariance D(θ k )D(θ k ) to the gradients. We now focus our attention on finding a suitable matrix D(θ k ) such that the algorithm in Eqn. 2 has desirable convergence and generalization performance.

Intrinsic SGD noise. Let B ⊂ S be a mini-batch drawn uniformly and without replacement from S and M = |B| be the size of this chosen mini-batch. We can write the SGD update rule here as

EQUATION

This result can be found in Hu et al. (2017) ; Hoffer et al. (2017) . Suppose that the loss function here is the negative log-likelihood, i (θ k ) = -log p(y i |x i , θ k ), where p(y|x, θ) is the density function for the model's predictive distribution.

Fisher information matrix. Given an input-target pair (x, y), the Fisher information matrix is defined as the expectation of the outer product of log-likelihood gradients,

E Px,P y|x [∇ log p(y|x, θ)∇ log p(y|x, θ) ].

(4) The expectation here is taken with respect to the data distribution P x for inputs x and the model's predictive distribution P y|x for targets y. The Fisher matrix can be realized as the second-order Taylor expansion of the KL-divergence on the space of predictive distributions, hence defining a Riemannian metric on this space (Amari, 1998; Martens, 2014) . The Fisher matrix for neural networks captures the local curvature information of the parameter space. We give a detailed description of Fisher matrices for feed-forward and convolutional network architectures in Appendix C.

In practical situations, we often sample y from the empirical training distribution rather than the predictive distribution. Therefore, we have the empirical Fisher (Martens, 2014) 1

N (xi,yi)∈S ∇ log p(y i |x i , θ)∇ log p(y i |x i , θ) .

For the rest of this paper, unless otherwise specified, all mentions of "Fisher matrix" or F (θ) refers to the empirical Fisher.

Naive method to close "generalization gap". Having modeled the covariance of the mini-batch SGD noise as a Fisher matrix, we now consider a naive approach to close "generalization gap" in large-batch training. Let B S denote small-batch and M S = |B S | denote the size of small-batch. Let us choose D(θ k ) in Eqn. 2 to be

EQUATION

where

F (θ k ) is the square-root of Fisher, F (θ k ) F (θ k ) = F (θ k ).

The motivation here is that the difference between large-batch and small-batch gradients can be approximated as a Gaussian noise with mean zero and covariance M L -M S M L M S F (θ k ) (Jastrzebski et al., 2017; Xing et al., 2018; Zhu et al., 2018; Li et al., 2015) . However, there is an immediate issue with this naive approach: if the Gaussian approximation is reasonable, then this algorithm should have similar behavior as smallbatch training which implies poor convergence performance. Indeed, as shown in the 2D convex example in Fig. 1c , adding full Fisher noise recovers SGD behavior. Furthermore, on the CIFAR-10 image classification task trained using ResNet44, we see in Fig. 3c that the training performance is almost identical to small-batch during the entire optimization process. Thus, this choice of D(θ k ) does not satisfy our requirement of maintaining desirable convergence.

3. METHOD:



3.1. PROPOSED ALGORITHM:

Using diagonal Fisher. We consider replacing F (θ k ) with diag (F (θ k )) in Eqn. 5 for D(θ k ). Our proposed algorithm to correct the "generalization gap" in large-batch training is to inject diagonal Fisher noise to large-batch gradients,

θ k+1 = θ k -α k ∇L M L (θ k ) + α k M L -M S M L M S diag (F (θ k ))ξ k+1 , ξ k+1 ∼ N (0, I d ).

The formal statement is given in Algorithm 1. Most of the empirical analysis of Algorithm 1 in Section 4 later uses feed-forward and convolutional networks. For completeness, we provide explicit expressions of diagonal Fisher for these architectures in Appendix C.

Changing the covariance structure in this way has important implications with regards to convergence and generalization performance. In the next subsection, we analyze our algorithm in a simple quadratic setting and compare it to the case where the covariance is given by the Fisher matrix.

Even in this simple case, it is not obvious to compare the generalization performance between the two algorithms. In Appendix B, we provide an extensive theoretical discussion regarding how the choice of covariance structure influences generalization performance of Eqn. 2.

Working with the assumption that the generalization error is comparable between diagonal Fisher and full Fisher, the excess risk in Eqn. 1 can be minimized by focusing only on the optimization error. We prove in Theorem 3.1 that the discrepancy in convergence behavior can be measured by the difference of their respective Frobenius norms. In Fig. 2 , we illustrate this behavior on a 2D toy problem where we compare the training trajectory of adding full Fisher noise versus diagonal full Fisher noise to the true gradient. As shown in the figure, diagonal Fisher converges faster than full Fisher with respect to the same number of iterations. In our experiments in Section 4, we observe that this phenomena carries over to the deep learning setting.

3.2. CASE-STUDY: CONVEX QUADRATIC EXAMPLE:

In this subsection, we restrict our setting and take the loss function L(θ) to be the convex quadratic,

L(θ) ≡ 1 2 θ Aθ,

where A is a symmetric and positive-definite matrix. Observe that in this setting the Fisher and Hessian coincide, which is simply the matrix A.

Algorithm 1 Adding diagonal Fisher noise to large-batch SGD. Differences from standard SGD are highlighted in blue We used a learning rate of 0.1 for 500 iterations (plotting every 10 iterations). Observe that adding diagonal Fisher to the true gradient achieves faster convergence than full Fisher.

Require: Number of iterations K, initial step-size α 0 , large-batch B L of size M L , small-batch B S of size M S , initial condition θ 0 ∈ Θ ⊂ R d for k = 1 to K do Z k ∼ N (0, I d ) k = α k M L -M S M L M S diag (F (θ k ))Z k θ k+1 = θ k -α k ∇L M L (θ k ) + k end for

We stress here that approximating the loss surface of a neural network with a quadratic often serves as a fertile "testing ground" when introducing new methodologies in deep learning. Analyzing the toy quadratic problem has led to important advances; for example, in learning rate scheduling (Schaul et al., 2013) and formulating SGD as approximate Bayesian inference (Mandt et al., 2017) . With regards to optimization, Martens (2010) observed that much of the difficulty in neural network optimization can be captured using quadratic models.

We focus now on the algorithm in Eqn. 2 and choose a constant d × d covariance matrix C. The following theorem, adapted from Bottou et al. (2018) , analyzes the convergence of this optimization method. The proof is relegated to Appendix A.

Theorem 3.1. Let λ max and λ min denote the maximum and minimum eigenvalue of A respectively. For a chosen α 0 ≤ λ -1 max , suppose that we run the algorithm in Eqn. 2 according to the decaying step-size sequence

α k = 2 (k + γ)λ min ,

for all k ∈ N >0 and where γ is chosen such that α k ≤ α 0 . Then for all k ∈ N,

E[L(θ k )] ≤ ν k + γ where ν = max 2 Tr(C AC) λ 2 min , γL(θ 0 ) .

We make a few observations regarding this bound. First, the convergence rate is optimal when C = 0. However, in this case, there is no noise and so we obtain no regularization benefits which leads to poor generalization. A more formal discussion is given at the end of Appendix B where if we use a scaling factor C λ := λC; as λ → 0, the expected generalization error becomes worse.

The second observation to note is that the term of importance in this theorem is Tr(C AC). While the overall convergence rate of the algorithm is O(1/k), the discrepancy in convergence performance for different choices of the matrix C rests entirely on this term. We analyze two specific cases which are relevant for us: the first case where C is square-root of A, C = √ A, and the second case where C is the square-root of the diagonal of A, C = diag (A). For the first, we have

Tr(C AC) = Tr(A 2 ) = A 2 Frob .

For the latter case, we have

Tr(C AC) = Tr(diag (A) 2 ) = diag (A) 2 Frob .

Thus, the difference in training performance between the two cases can be measured by the difference of their respective Frobenius norms.

3.3. SAMPLING RANDOM VECTOR WITH FISHER COVARIANCE:

We describe a method to sample a random vector with mean zero and approximately Fisher covariance to avoid computing directly the square root of the empirical Fisher. Let M be the size of the mini-batch and from the M -forward passes, we obtain prediction f (x, θ) and then from the backward passes, we obtain the back-propagated gradients ∇ 1 , . . . , ∇ M for each data-point. Consider independent random variables σ 1 , . . . , σ M drawn from Rademacher distribution, i.e., P(σ i = 1) = P(σ i = -1) = 1/2. Then, the mean

E σ [ M i=1 σ i ∇ i ] = 0.

The covariance is empirical Fisher.

4. EXPERIMENTS:

We first empirically compute the Frobenius norm of full Fisher and diagonal Fisher, followed by the marginal variance of the gradients for a variety of training regimes. We also compare the convergence speed of diagonal Fisher noise and full Fisher noise; the convergence speed here is measured with respect to the number of parameter updates. In Section 4.2, we compute the maximum eigenlue of the Hessian with respect to the final model parameters of different training regimes. The purpose of this is to empirically relate our methodology with the curvature of the loss surface landscape. In Section 4.3, we give the generalization performance of each method discussed previously.

Throughout our experiments, large-batch size LB is set to 4096 and small-batch size SB is set to 128 by default. The network architectures we use are fully-connected networks, shallow convolutional networks (LeNet (LeCun et al., 1998) , AlexNet (Krizhevsky et al., 2012) ), and deep convolutional networks (VGG16 (Simonyan & Zisserman, 2014) , ResNet44 (He et al., 2016) ). These models are evaluated on the standard deep-learning datasets: MNIST, Fashion-MNIST (LeCun et al., 1998; Xiao et al., 2017) , CIFAR-10 and CIFAR-100 (Krizhevsky & Hinton, 2009) .

4.1. CONVERGENCE EXPERIMENTS:

In the convex quadratic setting in Section 3.2, we observed that full-batch gradient descent with diagonal Fisher noise enjoyed faster convergence than full Fisher noise. This was characterized by the difference of their Frobenius norms. We now give an empirical verification of this phenomena in the non-convex setting of deep neural networks. We compute the Frobenius norms during the training of a ResNet44 network on CIFAR-10. Fig. 3a shows that the full Fisher matrix has much larger Frobenius norm than the diagonal Fisher matrix, which suggests that using diagonal Fisher noise should have faster convergence than full Fisher noise in the deep neural network setting. To justify this intuition, we analyze the training loss (per parameter update) of ResNet44 (CIFAR-10) on the following four regimes: LB, SB, LB with diagonal Fisher noise and LB with full Fisher noise. For a fair comparison, we use the same learning rate schedule for all four regimes. Fig. 3c shows that LB with diagonal Fisher noise trains much faster than LB with full Fisher. More interestingly, LB with full Fisher matches the convergence performance of SB, indicating that the intrinsic noise of SB is accurately modeled by the Fisher. In contrast, LB with diagonal Fisher attains a convergence similar to LB, showing that our approach preserves the desired convergence behavior of LB.

Next, we give an estimation of the marginal variance of gradients for the four regimes mentioned above as well as LB with K-FAC noise. K-FAC (Martens & Grosse, 2015) is a block-diagonal approximation of the Fisher matrix used as a second-order optimization method in deep learning. The purpose of this experiment is to show despite the fact that LB with diagonal Fisher trains much faster than LB with full Fisher, they share roughly the same marginal variance of the gradients. This suggests that the off-diagonal elements of the Fisher matrix is the key reason for slow convergence. The experiment is performed as follows: we freeze a partially-trained network and compute Monte-Carlo estimates of gradient variance with respect to each parameter over different mini-batches. This variance is then averaged over the parameters within each layer. The results are presented in Fig. 3b . We find that LB with diagonal Fisher noise and LB with full Fisher noise give roughly the same scale of marginal gradients variance.

4.2. MAXIMUM EIGENVALUE COMPARISON:

While the relationship between loss surface curvature and generalization is not completely explicit, numerous works have suggested that the maximum eigenvalue of Hessian is possibly correlated with generalization performance (Keskar et al., 2016; Chaudhari et al., 2016; Chaudhari & Soatto, 2017; Yoshida & Miyato, 2017; Xing et al., 2018) . To situate our method with this line of research, we compute the maximum eigenvalue of the Hessian of the final model for the following three regimes: LB with diagonal Fisher noise, LB with ghost batch-normalization, and SB. Ghost batchnormalization (GBN) is an adaptation of usual batch-normalization introduced in Hoffer et al. (2017) to close generalization gap. In Fig. 3b , we find that LB with diagonal Fisher gives smaller maximum eigenvalue than LB with GBN.

Computing maximum eigenvalue without any modification to the model gives inconsistent estimates even between different runs of the same training configuration. To make the maximum eigenvalue of the Hessian matrix comparable over different training trajectories, the Hessian needs to be invariant under typical weight reparameterizations such as affine transformations (Liao et al., 2018) . To achieve this, we make the following modification to the trained model: (1) For the layer with batch normalization, we can just push the batch-norm layer parameters and running statistics into the layerwise parameter space so that the the layer is invariant under affine transformation; and (2) For the layer without batch normalization, reparameterization changes the prediction confidence while the prediction remains the same. So we train a temperature parameter on the cross-entropy loss on the test data set, this encourages the model to make a calibrated prediction (prevent it from being over-confident). In Fig. 3d , we give the error bar of the maximum eigenvalue of the Hessian over different runs, which indicates the modification gives a roughly consistent estimate.

Under review as a conference paper at ICLR 2019

4.3. RESULTS:

In this subsection, we provide the validation accuracy for a number of training regimes. We point out that we do not experiment LB with full Fisher due to its exceedingly long training time. This can be seen from Fig. 3c where on ResNet44 (CIFAR-10), full Fisher does not achieve good convergence even after 8000 iterations. We use the baseline in Hoffer et al. (2017) , where LB is trained using GBN with learning rate adaptation. However, instead of using square-root scaling of the learning rate in Hoffer et al. (2017) , we use linear scaling in conjunction with a warmup scheme suggested by Goyal et al. (2017) . We found that this improves the baseline reported in Hoffer et al. (2017) .

We experimented with the following regimes: LB with GBN (our baseline), LB with GBN and trace Fisher noise (where D(θ k ) in Eqn. 2 is taken to be the square-root of the trace of Fisher Tr(F (θ k ))), LB with GBN and diagonal Fisher noise, and SB. All regimes were trained for the same number of epochs. For the first three regimes involving LB, we terminate the noise at epoch 50 and use standard LB for the remainder of training. The reasoning behind this comes from the work of Smith et al. (2017) which suggests that using noise only matters in the early stages of the optimization procedure. We give further explanation in Appendix E.

The final validation accuracies of each method are reported in Table 1 . While it is true that adding diagonal Fisher noise cannot completely close the "generalization gap" in some cases, we can see from Table 1 that doing so yields definite improvements. We point out that we explored other regimes such as injecting multiplicative Guassian noise with constant diagonal covariance (Hoffer et al., 2017) , but found that they all perform no better than our baseline (LB with GBN). In addition, we experimented with K-FAC noise but this did not give additional benefits over diagonal Fisher noise. We defer to Appendix E for exact details of these supplemental experiments. Table 1 : Validation accuracy results on classification tasks. GBN stands for Ghost-BN. Fisher Trace+GBN stands for LB + Ghost-BN with isotropic Gaussian noise scaled by square-root of trace of Fisher. Diag-F+GBN stands for LB + Ghost-BN with diagonal Fisher noise. All methods in each row are trained with the same number of epochs. Confidence interval is computed over 3 random seeds. LB with full Fisher requires same number of updates as SB, which renders it impractical for most of the models for which we work with. While this is indeed infeasible to compute for all the models, we note that LB with full Fisher reaches roughly the same validation accuracy (93.22) as SB in the case of ResNet44 (CIFAR-10). 

Dataset:



5. CONCLUSION:

In this paper, we explored in depth the relationship between curvature noise and stochastic optimization. We proposed a method to engineer large-batch training such that we retain fast convergence performance while achieving significant gains in generalization. In addition, we highlighted the importance of noise covariance structure on optimization and generalization. An interesting future direction would be to further understand both empirically and theoretically how different noise covariance structures impact optimization and generalization performance of a learning algorithm.

A PROOF OF THEOREM 3.1

The proof of this theorem follows the spirit of Bottou et al. (2018) . The algorithm

θ k+1 = θ k -α k Aθ k + α k Cξ k+1 , ξ k+1 ∼ N (0, I d ).

falls into the Robbins-Monro setting where the true gradient is perturbed by random noise. This perturbation can be considered as a martingale difference in the sense that

E[Cξ k+1 |F k ] = 0

where (F k ) k∈N is a increasing filtration generated by the sequence of parameters (θ k ) k∈N . When the step size is constant α k = α for all k, it corresponds to the Euler discretization of a gradient flow with random perturbation. We begin the proof by considering the equality,

L(θ k+1 ) = L(θ k ) + L(θ k ), θ k+1 -θ k + 1 2 (θ k+1 -θ k ) ∇ 2 L(θ k )(θ k+1 -θ k ).

Using the fact that ∇L(θ k ) = Aθ k , ∇ 2 L(θ k ) = A, and from the definition of θ k+1 , we can rewrite the above equation as

L(θ k+1 ) = L(θ k ) + Aθ k , -α k Aθ k + α k Cξ k+1 + 1 2 α k Aθ k -α k Cξ k+1 2 A .

Now, taking the conditional expectation E[•|F k ] on both sides of the equality, we obtain by independence of the noise

EQUATION

A simple computation shows

E[ Cξ k+1 2 A ] = E[(Cξ k+1 ) A(Cξ k+1 )] = E[ξ k+1 C ACξ k+1 ] = Tr(C AC) (7) Moreover, we have Aθ k 2 A = (θ k A)A(Aθ k ) = A Aθ k 2 2 ≤ λ max Aθ k 2 2 .

(8)

Using the results in Eqns. 7 and 8 as well as the assumption on the step-size schedule for all k: α k < α 0 < 1 λmax , we rewrite Eqn. 6 as

E[L(θ k+1 )|F k ] ≤ L(θ k ) + α k 2 λ max -1 α k Aθ k 2 2 + α 2 k 2 Tr(C AC) ≤ L(θ k ) - α k 2 Aθ k 2 2 + α 2 k 2 Tr(C AC). (9) Furthermore, Aθ k 2 2 = A(θ k Aθ k ) ≥ λ min θ k 2 A = 2λ min L(θ k )

Using this above fact and then taking the expectation of Eqn. 9 leads to

E[L(θ k+1 )] ≤ (1 -α k λ min )E[L(θ k )] + α 2 k 2 Tr(C AC).

We proceed by induction to prove the final result. By definition of ν, the result is obvious for k = 0.

For the inductive step, suppose that the induction hypothesis holds for k, i.e.,

α k = 2 (k + γ)λ min , E[L(θ k )] ≤ ν k + γ .

We prove the k + 1 case.

E[L(θ k+1 )] ≤ 1 - 2 k + γ ν k + γ + 2 (k + γ) 2 λ 2 min Tr(C AC) ≤ ν (k + γ + 1)

This comes from the definition of ν and also the inequality (k + γ -1)(k + γ + 1) ≤ (k + γ) 2 . This conclude the proof.

B RELATIONSHIP BETWEEN CHOICE OF COVARIANCE STRUCTURE AND:



GENERALIZATION:

As in Section 3.2, we work entirely in the convex quadratic setting. Consider the algorithm

EQUATION

Our aim in this section is to provide some theoretical discussions on how the choice of covariance structure C influences the generalization behavior of the above algorithm.

Uniform stability. Uniform stability (Bousquet & Elisseeff, 2002) is one of the most common techniques used in statistical learning theory to study generalization of a learning algorithm. Intuitively speaking, uniform stability measures how sensitive an algorithm is to perturbations of the sampling data. The more stable an algorithm is, the better its generalization will be. Recently, the uniform stability of several algorithms has been investigated for stochastic gradient methods (Hardt et al., 2015) or stochastic gradient Langevin dynamics algorithm (Mou et al., 2017; Raginsky et al., 2017) . We present the precise definition. Definition B.1 (Uniform stability). A randomized algorithm A is -stable if for all data sets S and S where S and S differ in at most one sample, we have

sup (x,y) |E A [L(θ S ) -L(θ S )]| ≤ ,

where L(θ S ) and L(θ S ) highlight the dependence of parameters on sampling datasets. The supremum is taken over input-target pairs (x, y) belonging to the sample domain.

The following theorem from Bousquet & Elisseeff (2002) shows that uniform stability implies generalization. Theorem B.2 (Generalization in expectation). Let A be a randomized algorithm which isuniformly stable, then

|E A [E gen ]| ≤

, where E gen is the expected generalization error as defined in Eqn. 1.

Continuous-time dynamics. We like to use the uniform stability framework to analyze generalization properties of Eqn. 10. To do this, we borrow ideas from the recent work of Mou et al. (2017) which give uniform stability bounds for Stochastic Gradient Langevin Dynamics (SGLD) in non-convex learning. While the authors in that work give uniform stability bounds in both the discrete-time and continuous-time setting, we work with the continuous setting since this conveys relevant ideas while minimizing technical complications. The key takeaway from Mou et al. (2017) is that uniform stability of SGLD may be bounded in the following way

EQUATION

Here, π t and π t are the distributions on parameters θ trained on the datasets S and S . The H 2 refers to the Hellinger distance.

We now proceed to mirror the approach of Mou et al. (2017) for Eqn. 10. Our usage of stochastic differential equations will be very soft but we refer to reader to Gardiner (2009) ; Pavliotis (2014) for necessary backgrounds. For the two datasets S and S , the continuous-time analogue of Eqn. 10 are Ornstein-Uhlenbeck processes (Uhlenbeck & Ornstein, 1930) :

dθ S (t) = -A S θ S (t) + √ αC S dW (t) dθ S (t) = -A S θ S (t) + √ αC S dW (t).

The solution is given by

θ S (t) = e -A S t θ S (0) + √ α t 0 e -A S (t-u) C S dW (u),

In fact, this yields the Gaussian distribution

θ S (t) ∼ N (µ S (t), Σ S (t)),

where µ S (t) = e -A S t θ S (0)

and Σ S (t) satisfies the Ricatti equation,

d dt Σ S (t) = -(A S Σ S (t) + Σ S (t)A S ) + αC S C S .

Observe that A S is symmetric and positive-definite which means that it admits a diagonalization A S = P S D S P -1 S . Solving the equation for the covariance matrix gives

EQUATION

We are in the position to directly apply the framework of (Mou et al., 2017) . Choosing π t and π t in Eqn. 11 to be the Gaussians N (µ S (t), Σ S (t)) and N (µ S (t), Σ S (t)) respectively, we obtain a uniform stability bound for Eqn. 10. We compute the right-hand side of the bound to get derive insights on generalization. Using the standard formula for Hellinger distance between two Gaussians, we have

EQUATION

Choosing the noise covariance. From Eqn. 13 above, it is evident that to ensure good generalization error for Eqn. 10, we want to choose a covariance C S such that the Hellinger distance H is minimized. Since we are working within the uniform stability framework, a good choice of C S should be one where Eqn. 10 becomes less data-dependent. This is intuitive after all -the less data-dependent an algorithm is; the better suited it should be for generalization.

We study Eqn. 13. Note that as time t → ∞, the exponential term goes to 1. Hence, we focus our attention on the ratio of the determinants. Suppose that we choose C S = √ A S . Since A S is the Fisher for this convex quadratic example, Eqn. 10 is essentially the naive method given in Section 2.2. Simplifying the determinant of Σ S (t) in this case, A to be extreme cases and it is interesting to observe that the Hellinger distance is more sensitive to dataset perturbation when C = I d . Our proposed method of this paper was to choose C = diag (A) and experiments in Section 4 seem to suggest that choosing the square-root of diagonal captures much of the generalization behavior of full Fisher. Understanding precisely why this is the case poses an interesting research direction to pursue in the future.

A simple scaling argument also highlights the importance of the trade-off between optimization and generalization. Consider C λ = λC. Then Theorem 3.1 suggests to take λ small to reduce the variance and improve convergence. However, in that case Σ λ = λ 2 Σ where Σ is given by the Eqn. 12 for C and

H 2 (π t , π t ) = 1 - det(Σ S ) 1 4 det(Σ S ) 1 4 det( Σ S +Σ S 2 ) 1 2 exp - 1 8λ 2 (µ S -µ S ) Σ S + Σ S 2 -1 (µ S -µ S )

and the Hellinger distance get close to one in the limit of small λ (which intuitively corresponds to the large batch situation).

C FISHER INFORMATION MATRIX FOR DEEP NEURAL NETWORKS:

In this section, we give a formal description of the true Fisher information matrix, rather than the empirical version, for both feed-forward networks and convolutional networks. In addition, we give the diagonal expression for both networks.

C.1 FEED-FORWARD NETWORKS:

Consider a feed-forward network with L layers. At each layer i ∈ {1, . . . , L}, the network computation is given by

z i = W i a i-1 a i = φ i (z i ),

where a i-1 is an activation vector, z i is a pre-activation vector, W i is the weight matrix, and φ i : R → R is a nonlinear activation function applied coordinate-wise. Let w be the parameter vector of network obtained by vectorizing and then concatenating all the weight matrices W i ,

w = [vec(W 1 ) vec(W 2 ) . . . vec(W L ) ] .

Furthermore, let Dv = ∇ v log p(y|x, w) denote the log-likelihood gradient. Using backpropagation, we have a decomposition of the log-likelihood gradient DW i into the outer product:

DW i = g i a i-1 ,

where g i = Dz i are pre-activation derivatives. The Fisher matrix F (w) of this feed-forward network is a L × L matrix where each (i, j) block is given by

EQUATION

Diagonal version. We give an expression for the diagonal of F i,i (w) here. The diagonal of F (w) follows immediately afterwards. Let a 2 i-1 and g 2 i be the element-wise product of a i-1 and g i respectively. Then, in vectorized form,

diag (F i,i (w)) = E[vec((a 2 i-1 )(g 2 i ) )],

where (a 2 i-1 )(g 2 i ) is the outer product of a 2 i-1 and g 2 i .

C.2 CONVOLUTIONAL NETWORKS:

In order to write down the Fisher matrix for convolutional networks, it suffices to only consider convolution layers as the pooling and response normalization layers typically do not contain (many) trainable weights. We focus our analysis on a single layer. Much of the presentation here follows (Grosse & Martens, 2016; Luk & Grosse, 2018) .

A convolution layer l takes as input a layer of activations a j,t where j ∈ {1, . . . , J} indexes the input map and t ∈ T indexes the spatial location. T here denotes the set of spatial locations, which we typically take to be a 2D-grid. We assume that the convolution here is performed with a stide of 1 and padding equal to the kernel radius R, so that the set of spatial locations is shared between the input and output feature maps. This layer is parameterized by a set of weights w i,j,δ , where i ∈ {1, . . . , I} indexes the output map and δ ∈ ∆ indexes the spatial offset. The numbers of spatial locations and spatial offsets are denoted by |T | and |∆| respectively. The computation of the convolution layer is given by z

EQUATION

The pre-activations z i,t are then passed through a nonlinear activation function φ l . The loglikelihood derivatives of the weights are computed through backpropagation:

Dw i,j,δ = t∈T a j,t+δ Dz i,t .

Then, the Fisher matrix here is

EQUATION

Diagonal version. To give the diagonal version, it will be convenient for us to express the computation of the convolution layer in matrix notation. First, we represent the activations a j,t as a J × |T | matrix A l-1 , the pre-activations z i,t as a I × |T | matrix Z l , and the weights w i,j,δ as a I × J|∆| matrix W l . Furthermore, by extracting the patches surrounding each spatial location t ∈ T and flattening these patches into column vectors, we can form a J|∆| × |T | matrix A exp l-1 which we call the expanded activations. Then, the computation is Eqn. 15 can be reformulated as the matrix multiplication Z l = W l A exp l-1 . Readers familiar with convolutional networks can immediately see that this is the Conv2D operation.

At a specific spatial location t ∈ T , consider the J|∆|-dimensional column vectors of A exp l-1 and I-dimensional column vectors of Z l . Denote these by a (:,t) l-1 and z (t) l respectively. The matrix W l maps a (:,t) l-1 to z (t) l . In this case, we find ourselves in the exact same setting as the feed-forward case given earlier. The diagonal is simply

E vec (a (:,t) l-1 ) 2 (Dz (t) l ) 2 D KRONECKER-FACTORED APPROXIMATE CURVATURE (K-FAC)

In Section 4, we compared the diagonal approximation of the Fisher matrix to the Kroneckerfactored approximate curvature (K-FAC) (Martens & Grosse, 2015) approximation of the Fisher matrix. We give a brief overview of the K-FAC approximation in the case of feed-forward networks.

Recall that the Fisher matrix for a feed-forward network is a L × L matrix where each of the (i, j) blocks are given by Eqn. 14. Consider the diagonal (i, i) blocks. If we approximate the activations a i-1 and pre-activation derivatives g i as statistically independent, we have

F i,i (w) = E[vec(DW i ) vec(DW i ) ] = E[a i-1 a i-1 ⊗ g i g i ] ≈ E[a i-1 a i-1 ] ⊗ E[g i g i ].

Let A i-1 = E[a i-1 a i-1 ] and G i = E[g i g i ]. The K-FAC approximation F of the Fisher matrix F is

F =     A 0 ⊗ G 1 0 A 1 ⊗ G 2 . . . 0 A L-1 ⊗ G L     .

The K-FAC approximation of the Fisher matrix can be summarized in the following way: (1) keep only the diagonal blocks corresponding to individual layers, and (2) make the probabilistic modeling assumption where the activations and pre-activation derivatives are statistically independent.

E SUPPLEMENTARY EXPERIMENTAL RESULTS:

In this section, we give additional validation accuracy results to complement Table 1 . The additional regimes we experiment with are: BatchChange, Multiplicative, and K-FAC. For BatchChange, we use SB for 50 epochs and LB for 150 epochs. We perform this experiment to verify the argument given in Smith et al. (2017) that adapting from a small-batch regime to large-batch regime can close generalization gap. While Table 2 shows that BatchChange attains good validation accuracy numbers comparable to SB, this is not a preferred solution as it is not strictly LB training during the whole process which sacrifices the benefits of LB training in the early stages of training.

The other experiments we report here are Multiplicative and K-FAC. Multiplicative stands for multiplying gradients with Gaussian noise with constant diagonal convariance structure. This idea comes from (Hoffer et al., 2017) . K-FAC means we use the K-FAC approximation instead of diagonal Fisher noise in Algorithm 1. The reason why 4K with diagonal Fisher converges faster than the 4K GBN baseline in the beginning is because we need to use a larger learning rate for 4K GBN to get the best generalization performance, which makes it actually converge a bit slower in the beginning. If they share the same learning rate scheme, 4K GBN would converge a bit faster because it is less noisy.



Figures and Tables:

Figure 1: Noise structure in a simple two-dimensional regression problem. a) One-step SGD update. b) Onestep SGD update with isotropic Gaussian (σ = 0.1) noise. c) One-step SGD update with full Fisher noise. d) One-step SGD update with diagonal Fisher noise. The full Fisher noise almost recovers the SGD noise.Observe that the full Fisher noise direction is perpendicular to the contours of the loss surface. Moreover, full Fisher exhibits slower convergence than diagonal Fisher; we refer to Section 3.2 for a more detailed analysis.

Figure2: Trajectory using full Fisher noise versus diagonal Fisher noise for the algorithm in Eqn. 2 used to minimize a two-dimensional quadratic function. Blue dot indicates the initial parameter value and the green dot shows the final parameter value. We used a learning rate of 0.1 for 500 iterations (plotting every 10 iterations). Observe that adding diagonal Fisher to the true gradient achieves faster convergence than full Fisher.

Maximum Eigenvalue of the Hessian.

Figure 3: a) Frobenius norm of full Fisher matrix and diagonal Fisher matrix. The model is trained on ResNet44 with CIFAR-10. b) Estimation of gradient variances for randomly selected convolutional layers on ResNet44 (CIFAR-10). fullF: LB with full Fisher noise. diagF: LB with diagonal Fisher noise. KfacF: LB with K-FAC noise. c) Training error with respect to iterations between SB, LB, LB with full Fisher noise and diagonal Fisher noise on ResNet44 (CIFAR-10). All of the above are trained with the same learning rate. d) The maximum eigenvalue of the Hessian matrix at the end of training. 4Kgbn: LB with Ghost-BN. 4KdiagF: LB with diagonal Fisher noise.

d -e -2D S t ) Suppose that we choose C = I d . Proceeding analogously, det(Σ S (t)) = α 2 d det(I d -e -2D S t ) det(D S )We can think of choosing C = I d or C = √

Figure 4: Validation error over number of epochs. Left: Training on CIFAR-10 with VGG-16. Right: Training on CIFAR-100 with resnet44.The reason why 4K with diagonal Fisher converges faster than the 4K GBN baseline in the beginning is because we need to use a larger learning rate for 4K GBN to get the best generalization performance, which makes it actually converge a bit slower in the beginning. If they share the same learning rate scheme, 4K GBN would converge a bit faster because it is less noisy.

Validation accuracy results on classification tasks using BatchChange, Multiplicative, K-FAC. For reader's convenience, we report again the result of Diag-F+GBN We give the training plots over epochs, the experimental results with the square root learning scaling scheme in the following table and plots,

Table content: <table><tr><td>Dataset</td><td>Network</td><td>SB</td><td colspan="4">BatchChange Multiplicative K-FAC Diag-F+GBN</td></tr><tr><td>CIFAR-10</td><td>VGG16</td><td>93.25</td><td>93.18</td><td>90.98</td><td>93.06</td><td>93.15 ± 0.05</td></tr><tr><td>CIFAR-100</td><td>VGG16</td><td>72.83</td><td>72.44</td><td>68.77</td><td>71.86</td><td>71.94 ± 0.14</td></tr><tr><td>CIFAR-10</td><td>ResNet44</td><td>93.42</td><td>93.02</td><td>91.28</td><td>92.81</td><td>92.72 ± 0.15</td></tr><tr><td colspan="3">CIFAR-100 ResNet44x2 75.55</td><td>75.16</td><td>71.98</td><td>73.84</td><td>74.10 ± 0.18</td></tr></table>

Validation accuracy results on classification tasks where Diag F+GBN uses square-root scaling learning rate.

Table content: <table><tr><td>Dataset</td><td>Network</td><td>SB</td><td colspan="2">LB+GBN Diag-F+GBN</td></tr><tr><td>CIFAR-10</td><td>VGG16</td><td>93.25</td><td>91.6</td><td>92.9</td></tr><tr><td>CIFAR-100</td><td>VGG16</td><td>72.83</td><td>69.1</td><td>71.5</td></tr><tr><td>CIFAR-10</td><td>ResNet44</td><td>93.42</td><td>91.7</td><td>92.6</td></tr><tr><td colspan="3">CIFAR-100 ResNet44x2 75.55</td><td>72.8</td><td>73.6</td></tr></table>

Peer review (iclr2019):
 It has previously been observed that training deep networks using large batch-sizes leads to a larger generalization gap compared to the gap when training with a relatively small batch-size. This paper proposes to add noise sampled from diagonal "empirical" Fisher matrix to the large batch gradient as a method for closing the generalization gap. The authors motivate the use of empirical Fisher for sampling noise by arguing that the covariance of gradients from small batch-sizes can be seen as approximately equal to a scaled version of the Fisher matrix. It is then pointed out that using the Fisher matrix directly to sample noise could in principle close the generalization gap but would lead to slow converegence similar to SGD with a small batch-size. The authors then claim that the convergence speed is better when noise is sampled from the diagonal Fisher matrix instead of the full Fisher matrix. This claim is proven in theory for a convex quadratic loss surface and experiments are conducted to empirically verify this claim both in the quadratic setting are for realistic deep networks. Finally an efficient method for sampling noise from the diagonal empirical Fisher matrix is proposed. Comments: I think the paper is very well written and the results are presented clearly. In terms of novelty, I found the argument about convergence using diagonal Fisher being faster compared with full Fisher quite interesting, and its application for large batch training to be insightful. As a minor comment, for motivating theorem 3.1, it is pointed out by the authors that the diagonal Fisher acts as an approximation of the full Fisher and hence their regularization effects should be similar while convergence should be faster for diagonal Fisher. As a caveat, I think the authors should also point out that the convergence rate would be best when C is set to 0 in the result of the theorem. This implies no noise is used during SGD updates. However, this would imply the regularization effect from the noise will also vanish which would lead to poor generalization. However, there is a crucial detail that makes the main argument of the paper weak. In the main experiments in section 4.3, for the proposed large batch training method, the authors mention that they use a small batch-size of 128 for the first 50 epochs similar to Smith et al (2017) and then switch to the large batch-size of 4096, at which point, the learning rate is linearly scaled proportional to the large batch-size with a warmup scheme similar to Goyal et al (2017) and ghost batch normalization is used similar to Hoffer et al (2017). The former two tricks have individually been shown on their own to close the generalization gap for large batch-size training on large datasets like ImageNet. This paper combines these tricks and adds noise sampled from the diagonal Fisher matrix on top when switching to large batch-size after epoch 50 and reports experiments on smaller datasets-- MNIST, Fashion MNIST and the CIFAR datasets. Finally, the accuracy numbers for the proposed method is only marginally better than the baseline where isotropic noise is added to the large batch-size gradient. For these reasons, I do not consider the proposed method a significant improvement over existing techniques for closing the generalization gap for large batch training. There is also a statement in the paper that is problematic but can be fixed by re-writing. In the paper, empirical Fisher matrix, as termed by the authors in the paper, refers to the Fisher matrix where the target values in the dataset is used as the output of the model rather than sampling it from the model itself as done for computing the true Fisher matrix. This empirical (diagonal) Fisher matrix is used to sample noise which is added to the large batch gradient in the proposed method. It is mentioned that the covariance of the noise in small batch SGD is exactly same as the empirical Fisher matrix. This claim is premised on the argument that the expected gradient (over dataset) is unconditionally roughly 0, i.e., throughout the training. This is absolutely false. If this was the case, gradient descent (using full dataset) should not be able to find minima and this is far from the truth. Even if we compare the scale of expected gradient to the mini-batch gradient (for small batch-size), the scale of these two gradients at any point during training (using say small batch-size SGD) is of the same order. I am saying the latter statement from my personal experience. The authors can verify this as well. Overall, while I found the theoretical argument of the paper to be mostly interesting, I was dissapointed by the experimental details as they make the gains from the proposed method questionable when considered in isolation from the existing methods that close the generalization gap for large batch training.