Year: 2018

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.

heading: 1 INTRODUCTION
text: 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
(a) SGD step
(b) Isotropic Gaussian
(c) Full Fisher
(d) Diagonal Fisher
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.

heading: 2 PRELIMINARIES

heading: 2.1 BACKGROUND
text: Excess risk decomposition. Let S = {(x1, y1), . . . , (xN , yN )} 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(xi, θ), where θ ∈ Θ ⊂ Rd are parameters of the network and Θ here denotes the parameter space of the network. Let `(f(xi, θ), yi) be the loss function measuring the disagreement between outputs f(xi, θ) and targets yi. For convenience, we use the notation `i(θk) to denote `(f(xi, θ), yi). 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
ES [L̂(θk)] ≤ ES [L (θk)− L(θk)]︸ ︷︷ ︸ Egen +ES [L(θk)− L(θ∗)]︸ ︷︷ ︸ Eopt , (1)
where θ∗ = argminθL(θ) here is the empirical risk minimizer. The terms Egen and Eopt 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.

heading: 2.2 MOTIVATION
text: We tackle the large-batch generalization problem through gradient noise injection. Let BL denote large-batch and ML = |BL| denote size of the large-batch. Consider the following modification of large-batch SGD updates
θk+1 = θk − αk∇LML(θk) + αkD(θk)ξk+1, ξk+1 ∼ N (0, Id). (2) where αk is the learning rate, ∇LML(θk) = 1ML ∑ML i=1 `i(θk) is the large-batch gradient and N (0, Id) 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
θk+1 = θk − αk∇LM (θk) = θk − αk∇L(θk) + αk(∇L(θk)−∇LM (θk)︸ ︷︷ ︸
δk
)
where ∇L(θk) = 1N ∑N i=1∇`i(θk) is the full-batch gradient. The difference δk = ∇L(θk) − ∇LM (θk) is the intrinsic noise stemming from mini-batch gradients. The covariance of δk is
Cov(δk, δk) = ( N −M NM ) · 1 N N∑ i=1 (∇`i(θk)−∇L(θk))(∇`i(θk)−∇L(θk))>, (3)
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(yi|xi, θ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,
EPx,Py|x [∇ log p(y|x, θ)∇ log p(y|x, θ) >]. (4)
The expectation here is taken with respect to the data distribution Px for inputs x and the model’s predictive distribution Py|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(yi|xi, θ)∇ log p(yi|xi, θ)>.
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 BS denote small-batch and MS = |BS | denote the size of small-batch. Let us choose D(θk) in Eqn. 2 to be
D(θk) = √ ML −MS MLMS √ F (θk), (5)
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 ML−MSMLMS 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.

heading: 3 METHOD

heading: 3.1 PROPOSED ALGORITHM
text: 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∇LML(θk) + αk √ ML −MS MLMS √ diag (F (θk))ξk+1, ξk+1 ∼ N (0, Id).
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.

heading: 3.2 CASE-STUDY: CONVEX QUADRATIC EXAMPLE
text: 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
Require: Number of iterations K, initial step-size α0, large-batch BL of size ML, small-batch BS of size MS , initial condition θ0 ∈ Θ ⊂ Rd for k = 1 to K do Zk ∼ N (0, Id) k = αk √ ML−MS MLMS √ diag (F (θk))Zk
θk+1 = θk − αk∇LML(θk) + k end for
2 0 2
4
2
0
2
4
Full Fisher
2 0 2
Diagonal Fisher
Figure 2: 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.
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 ≤ λ−1max, 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)
λ2min , γ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 isO(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(A2) = ‖A‖2Frob .
For the latter case, we have
Tr(C>AC) = Tr(diag (A) 2 ) = ‖diag (A)‖2Frob .
Thus, the difference in training performance between the two cases can be measured by the difference of their respective Frobenius norms.

heading: 3.3 SAMPLING RANDOM VECTOR WITH FISHER COVARIANCE
text: 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.

heading: 4 EXPERIMENTS
text: 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).

heading: 4.1 CONVERGENCE EXPERIMENTS
text: 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 MonteCarlo 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.

heading: 4.2 MAXIMUM EIGENVALUE COMPARISON
text: 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.

heading: 4.3 RESULTS
text: 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.

heading: 5 CONCLUSION
text: 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.

heading: A PROOF OF THEOREM 3.1
text: The proof of this theorem follows the spirit of Bottou et al. (2018). The algorithm
θk+1 = θk − αkAθk + αkCξk+1, ξk+1 ∼ N (0, Id).
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|Fk] = 0
where (Fk)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)>∇2L(θk)(θk+1 − θk).
Using the fact that ∇L(θk) = Aθk, ∇2L(θk) = A, and from the definition of θk+1, we can rewrite the above equation as
L(θk+1) = L(θk) + 〈Aθk,−αkAθk + αkCξk+1〉+ 1
2 ‖αkAθk − αkCξk+1‖2A .
Now, taking the conditional expectation E[·|Fk] on both sides of the equality, we obtain by independence of the noise ξk+1 to Fk
E[L(θk+1)|Fk] = L(θk)− αk ‖Aθk‖22 + α2k 2 ‖Aθk‖2A + α2k 2 E[‖Cξk+1‖2A] (6)
A simple computation shows
E[‖Cξk+1‖2A] = E[(Cξk+1) >A(Cξk+1)]
= E[ξ>k+1C>ACξk+1] = Tr(C>AC)
(7)
Moreover, we have ‖Aθk‖2A = (θ > k A)A(Aθk)
= A ‖Aθk‖22 ≤ λmax ‖Aθk‖22 .
(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)|Fk] ≤ L(θk) + (αk
2 λmax − 1
) αk ‖Aθk‖22 + α2k 2 Tr(C>AC)
≤ L(θk)− αk 2 ‖Aθk‖22 + α2k 2 Tr(C>AC).
(9)
Furthermore, ‖Aθk‖22 = A(θ > k Aθk) ≥ λmin ‖θk‖ 2 A = 2λminL(θ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)] + α2k 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λ2min 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.

heading: B RELATIONSHIP BETWEEN CHOICE OF COVARIANCE STRUCTURE AND GENERALIZATION
text: As in Section 3.2, we work entirely in the convex quadratic setting. Consider the algorithm
θk+1 = θk − αk∇L(θk) + αkCξk+1, ξk+1 ∼ N (0, Id). (10)
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)
|EA[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 is - uniformly stable, then |EA[Egen]| ≤ , where Egen 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
SGLD ≤ sup S,S′
√ H2(πt, π′t). (11)
Here, πt and π′t are the distributions on parameters θ trained on the datasets S and S ′. TheH2 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) = −ASθS(t) + √ αCSdW (t) dθS′(t) = −AS′θS′(t) + √ αCS′dW (t).
The solution is given by
θS(t) = e −AStθS(0) + √ α ∫ t 0 e−AS(t−u)CSdW (u),
In fact, this yields the Gaussian distribution
θS(t) ∼ N (µS(t),ΣS(t)),
where µS(t) = e −AStθS(0)
and ΣS(t) satisfies the Ricatti equation,
d dt ΣS(t) = −(ASΣS(t) + ΣS(t)AS) + αCSC>S .
Observe that AS is symmetric and positive-definite which means that it admits a diagonalization AS = PSDSP −1 S . Solving the equation for the covariance matrix gives
ΣS(t) = αPS (∫ t 0 e−DS(t−u)P−1S CSC > S PSe −DS(t−u)du ) P−1S . (12)
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
H2(πt, π ′ t) = 1−
det(ΣS) 1 4 det(ΣS′) 1 4
det(ΣS+ΣS′2 ) 1 2
exp { −1
8 (µS − µS′)>
( ΣS + ΣS′
2
)−1 (µS − µS′) } (13)
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 CS such that the Hellinger distance H
is minimized. Since we are working within the uniform stability framework, a good choice of CS 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 CS = √ AS . Since AS 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,
det(ΣS(t)) = (α
2
)d det(Id − e−2DSt)
Suppose that we choose C = Id. Proceeding analogously, det(ΣS(t)) = (α
2 )d det(Id − e−2DSt) det(DS)
We can think of choosing C = Id or C = √ A to be extreme cases and it is interesting to observe that the Hellinger distance is more sensitive to dataset perturbation when C = Id. 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
H2(π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).

heading: C FISHER INFORMATION MATRIX FOR DEEP NEURAL NETWORKS
text: 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
zi = Wiai−1
ai = φi(zi),
where ai−1 is an activation vector, zi is a pre-activation vector, Wi 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 Wi,
w = [vec(W1) > vec(W2) > . . . vec(WL) >]>.
Furthermore, letDv = ∇v log p(y|x,w) denote the log-likelihood gradient. Using backpropagation, we have a decomposition of the log-likelihood gradient DWi into the outer product:
DWi = gia>i−1,
where gi = Dzi 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
Fi,j(w) = E[vec(DWi) vec(DWj)>] = E[ai−1a>j−1 ⊗ gig>j ]. (14)
Diagonal version. We give an expression for the diagonal of Fi,i(w) here. The diagonal of F (w) follows immediately afterwards. Let a2i−1 and g 2 i be the element-wise product of ai−1 and gi respectively. Then, in vectorized form,
diag (Fi,i(w)) = E[vec((a2i−1)(g2i )>)],
where (a2i−1)(g 2 i ) > is the outer product of a2i−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 aj,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 wi,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
zi,t = ∑ δ∈∆ wi,j,δaj,t+δ. (15)
The pre-activations zi,t are then passed through a nonlinear activation function φl. The loglikelihood derivatives of the weights are computed through backpropagation:
Dwi,j,δ = ∑ t∈T aj,t+δDzi,t.
Then, the Fisher matrix here is
E[Dwi,j,δDwi′,j′,δ′ ] = E [(∑ t∈T aj,t+δDzi,t )(∑ t′∈T aj′,t′+δ′Dzi′,t′ )] .
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 aj,t as a J × |T | matrix Al−1, the pre-activations zi,t as a I × |T | matrix Zl, and the weights wi,j,δ as a I×J |∆|matrixWl. 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 Aexpl−1 which we call the expanded activations. Then, the computation is Eqn. 15 can be reformulated as the matrix multiplication
Zl = WlA 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 Aexpl−1 and I-dimensional column vectors of Zl. Denote these by a (:,t) l−1 and z (t) l respectively. The matrix Wl 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 )]

heading: D KRONECKER-FACTORED APPROXIMATE CURVATURE (K-FAC)
text: 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 ai−1 and pre-activation derivatives gi as statistically independent, we have
Fi,i(w) = E[vec(DWi) vec(DWi)>] = E[ai−1a>i−1 ⊗ gig>i ] ≈ E[ai−1a>i−1]⊗ E[gig>i ].
Let Ai−1 = E[ai−1a>i−1] and Gi = E[gig>i ]. The K-FAC approximation F̂ of the Fisher matrix F is
F̂ =  A0 ⊗G1 0 A1 ⊗G2 . . .
0 AL−1 ⊗GL  . 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.

heading: E SUPPLEMENTARY EXPERIMENTAL RESULTS
text: 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.
We give the training plots over epochs, the experimental results with the square root learning scaling scheme in the following table and plots,
Table 3: Validation accuracy results on classification tasks where Diag F+GBN uses square-root scaling learning rate.
Dataset Network SB LB+GBN Diag-F+GBN
CIFAR-10 VGG16 93.25 91.6 92.9 CIFAR-100 VGG16 72.83 69.1 71.5 CIFAR-10 ResNet44 93.42 91.7 92.6 CIFAR-100 ResNet44x2 75.55 72.8 73.6
0 50 100 150 200 2500
10
20
30
40
50
60
Va lid
at io
n er
ro r
Vgg on cifar10 4K gbn Baseline 128 Baseline 4K diag Fisher
0 50 100 150 200 25020
30
40
50
60
70
80
90 Resnet on cifar100
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.

Review summaries:
In this paper, the authors propose a method to close the generalization gap that arises in training DNNs with large batch. The author reasons about the effectiveness in SGD small batch training by looking at the curvature structure of the noise. Instead of using the naïve empirical fisher matrix, the authors propose to use diagonal fisher noise for large batch SGD training for DNNs. The proposed method is shown empirically to achieve both comparable generalization and the training speedup compared to small batch training. A convergence analysis is provided for the proposed method under convex quadratic setting. 

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.