E XPLORING CURVATURE NOISE IN LARGE - BATCH STOCHASTIC OPTIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
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 signiﬁcant 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.
I NTRODUCTION
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 sufﬁciently 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 veriﬁed empirically in LeCun et al.
; Keskar et al.
(2016) that we have the opposite effect for test loss and error.
To fully realize the beneﬁts 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 sacriﬁcing 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 ﬁnd 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 efﬁcacy of our proposed method.
There are two primary takeaways from our empirical analysis.
First, we show adding diagonal Fisher noise to large-batch
[Figure]
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-efﬁcient and does not require any lengthening of the optimization process.
P RELIMINARIES 2. B ACKGROUND
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 classiﬁers 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 (cid:96) ( 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 (cid:96) i ( θ k ) to denote (cid:96) ( f ( x i , θ ) , y i ) .
The expected risk and empirical risk functions are deﬁned to be
[Equation]
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 deﬁned as
[Equation]
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 efﬁciently 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 sufﬁcient 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. M OTIVATION
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 modiﬁcation of large-batch SGD updates
[Equation]
where α k is the learning rate, ∇L M L ( θ k ) =
[Equation] M L i =1 (cid:96) 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 modiﬁcation as injecting a Gaussian noise with mean zero and covariance
D ( θ k ) D ( θ k ) (cid:62) to the gradients.
We now focus our attention on ﬁnding 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]
where ∇L ( θ k )
= 1 N (cid:80) Ni =1 ∇ (cid:96) i ( θ k ) is the full-batch gradient.
The difference δ k =
∇L ( θ k ) − ∇L M ( θ k ) is the intrinsic noise stemming from mini-batch gradients.
The covariance of δ k is
[Equation]
This result can be found in Hu et al.
; Hoffer et al.
Suppose that the loss function here is the negative log-likelihood, (cid:96)
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 deﬁned as the expectation of the outer product of log-likelihood gradients,
[Equation]
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 deﬁning 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)
[Equation]
For the rest of this paper, unless otherwise speciﬁed, 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 (cid:112) F ( θ k ) is the square-root of Fisher, (cid:112) F ( θ k ) (cid:112) F ( θ k ) (cid:62)
=
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 CIFARimage classiﬁcation 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.
M ETHOD 3. P ROPOSED ALGORITHM
Using diagonal Fisher.
We consider replacing (cid:112) F ( θ k ) with (cid:112) 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,
[Equation]
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 inﬂuences 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 ﬁgure, 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. C ASE - STUDY : CONVEX QUADRATIC EXAMPLE
In this subsection, we restrict our setting and take the loss function L ( θ ) to be the convex quadratic,
[Equation]
where A is a symmetric and positive-deﬁnite matrix.
Observe that in this setting the Fisher and
Hessian coincide, which is simply the matrix A .
Algorithm {+
+}
Adding diagonal blue
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 ) (cid:15) k
= α k (cid:113) M L − M S M L M S (cid:112) diag ( F ( θ k ))
Z k θ k
+
=
θ k − α k ∇L M L ( θ k ) + (cid:15)
k end for
[Figure]
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 difﬁculty 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
[Equation]
for all k ∈ N > 0 and where γ is chosen such that α k ≤ α 0 .
Then for all k ∈ N ,
E [ L ( θ k )] ≤ ν k + γ where ν = max (cid:18)
Tr( C (cid:62) AC )
λ 2min , γ L ( θ 0 ) (cid:19) .
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 beneﬁts 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 (cid:62) 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 speciﬁc cases which are relevant for us: the ﬁrst 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 = (cid:112) diag ( A ) .
For the ﬁrst, we have
[Equation]
For the latter case, we have
[Equation]
Thus, the difference in training performance between the two cases can be measured by the difference of their respective Frobenius norms.
3. S AMPLING 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 ∇ (cid:96) 1 , . . .
, ∇ (cid:96) 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 σ [ (cid:80) Mi =1 σ i ∇ (cid:96) i ]
= 0 .
The covariance is computed using the model prediction f ( x, θ ) instead of the training target y .
E XPERIMENTS
We ﬁrst 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 ﬁnal 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 toby 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. 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 veriﬁcation 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 ﬁnd that LB with diagonal Fisher noise and LB with full Fisher noise give roughly the same scale of marginal gradients variance.
4. 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
[Figure]
m m agn i t ude
[Figure]
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 ﬁnal 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 ﬁnd that LB with diagonal Fisher gives smaller maximum eigenvalue than LB with GBN.
Computing maximum eigenvalue without any modiﬁcation to the model gives inconsistent estimates even between different runs of the same training conﬁguration.
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 afﬁne transformations (Liao et al., 2018).
To achieve this, we make the following modiﬁcation 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 afﬁne transformation; and (2) For the layer without batch normalization, reparameterization changes the prediction conﬁdence 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-conﬁdent).
In Fig.
3d, we give the error bar of the maximum eigenvalue of the Hessian over different runs, which indicates the modiﬁcation gives a roughly consistent estimate.
4. 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.
We found that this improves the baseline reported in Hoffer et al.
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 (cid:112)
Tr( F ( θ k )) ), LB with GBN and diagonal Fisher noise, and SB .
All regimes were trained for the same number of epochs.
For the ﬁrst 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 ﬁnal 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 deﬁnite 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 beneﬁts over diagonal Fisher noise. We defer to Appendix E for exact details of these supplemental experiments.
[Table]
C ONCLUSION
In this paper
[Table]
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 signiﬁcant 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 P ROOF OF T HEOREM 3.
The proof of this theorem follows the spirit of Bottou et al.
The algorithm
[Equation]
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
[Equation]
where ( F k ) k ∈ N is a increasing ﬁltration 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 ﬂow with random perturbation.
We begin the proof by considering the equality,
[Equation]
Using the fact that ∇L ( θ k )
=
Aθ k , ∇ 2 L ( θ k )
=
A , and from the deﬁnition of θ k +1 , we can rewrite the above equation as
[Equation]
Now, taking the conditional expectation E [ ·|F k ] on both sides of the equality, we obtain by independence of the noise ξ k +1 to F k [Equation]
[Equation]
A simple computation shows
[Equation]
Moreover, we have
[Equation]
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
[Equation]
Furthermore,
[Equation]
Using this above fact and then taking the expectation of Eqn.
9 leads to
[Equation]
We proceed by induction to prove the ﬁnal result.
By deﬁnition of ν , the result is obvious for k = 0 .
For the inductive step, suppose that the induction hypothesis holds for k , i [Equation] e.,
[Equation]
We prove the k + 1 case.
[Equation]
This comes from the deﬁnition of ν and also the inequality ( k + γ − 1)( k + γ +1) ≤
( k + γ ) 2 .
This conclude the proof.
B R ELATIONSHIP 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 inﬂuences 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 deﬁnition.
Deﬁnition B.1 (Uniform stability) .
A randomized algorithm A is (cid:15) -stable
if for all data sets S and
S (cid:48) where S and S (cid:48) differ in at most one sample, we have
[Equation]
where L ( θ S ) and L ( θ S (cid:48) ) 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 (cid:15) - uniformly stable, then
[Equation]
where E gen is the expected generalization error as deﬁned 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]
[Equation]
Here, π t and π (cid:48) t are the distributions on parameters θ trained on the datasets S and S (cid:48) .
The H 2 refers to the Hellinger distance.
We now proceed to mirror the approach of Mou et al.
(2017) for Eqn.
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 (cid:48) , the continuous-time analogue of Eqn.are Ornstein-Uhlenbeck processes (Uhlenbeck & Ornstein, 1930):
[Equation]
The solution is given by
[Equation]
In fact, this yields the Gaussian distribution
[Equation]
where
[Equation] and Σ S ( t ) satisﬁes the Ricatti equation,
[Equation]
Observe that A S is symmetric and positive-deﬁnite 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 (cid:48) in Eqn.
11 to be the Gaussians N ( µ S ( t ) , Σ S ( t )) and N ( µ S (cid:48) ( t ) , Σ S (cid:48) ( t )) respectively, we obtain a uniform stability bound for Eqn.
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,
[Equation]
Suppose that we choose C
= I d .
Proceeding analogously,
[Equation]
We can think of choosing C
= I d 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 = I d .
Our proposed method of this paper was to choose C = (cid:112) 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
[Equation]
(cid:41) and the Hellinger distance get close to one in the limit of small λ (which intuitively corresponds to the large batch situation).
C F ISHER 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. F EED - FORWARD NETWORKS
Consider a feed-forward network with L layers.
At each layer i ∈ { 1 , . . .
, L } , the network computation is given by
[Equation]
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 ,
[Equation]
Furthermore, let D v = ∇ v log p ( y | x, w ) denote the log-likelihood gradient.
Using backpropagation, we have a decomposition of the log-likelihood gradient D W i into the outer product:
[Equation]
where g i =
D
z 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]
⊗ g i g (cid:62) j ] .
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,
[Equation]
where ( a 2 i − 1 )( g 2 i ) (cid:62) is the outer product of a 2 i − 1 and g 2 i .
C. C ONVOLUTIONAL NETWORKS
In order to write down the Fisher matrix for convolutional networks, it sufﬁces 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
[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:
[Equation]
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 ﬂattening 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
[Equation]
Readers familiar with convolutional networks can immediately see that this is the Conv2D operation.
At a speciﬁc 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 ﬁnd ourselves in the exact same setting as the feed-forward case given earlier.
The diagonal is simply
[Equation]
D K RONECKER - 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.
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
[Equation]
Let A i − 1 =
E [ a i − 1 a (cid:62) i − 1 ] and G i =
E [ g i g (cid:62) i ] .
The K-FAC approximation ˆ F of the Fisher matrix F is
[Equation]
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 sacriﬁces the beneﬁts 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.
[Table] We give the training plots over epochs, the experimental results with the square root learning scaling scheme in the following table and plots, [Figure] [Table]
[Figure]
rr o r
Vgg [Figure] cifar [Figure]