D ISENTANGLING THE M ECHANISMS B EHIND I MPLICIT R EGULARIZATION IN SGD
Anonymous authors
Paper under double-blind review
A BSTRACT
A number of competing hypotheses have been proposed to explain why smallbatch Stochastic Gradient Descent (SGD) leads to improved generalization over the full-batch regime, with recent work crediting the implicit regularization of various quantities throughout training.
However, to date, empirical evidence assessing the explanatory power of these hypotheses is lacking.
In this paper, we conduct an extensive empirical evaluation, focusing on the ability of various theorized mechanisms to close the small-to-large batch generalization gap.
Additionally, we characterize how the quantities that SGD has been claimed to (implicitly) regularize change over the course of training.
By using micro-batches , i.e. disjoint smaller subsets of each mini-batch, we empirically show that explicitly penalizing the gradient norm or the Fisher Information Matrix trace, averaged over micro-batches, in the large-batch regime recovers small-batch SGD generalization, whereas Jacobian-based regularizations fail to do so.
This generalization performance is shown to often be correlated with how well the regularized model’s gradient norms resemble those of small-batch SGD.
We additionally show that this behavior breaks down as the micro-batch size approaches the batch size.
Finally, we note that in this line of inquiry, positive experimental findings on CIFAR10 are often reversed on other datasets like CIFAR100, highlighting the need to test hypotheses on a wider collection of datasets.
I NTRODUCTION
While small-batch SGD has frequently been observed to outperform large-batch SGD (Geiping et al., 2022; Keskar et al., 2017; Masters and Luschi, 2018; Smith et al., 2021; Wu et al., 2020;
Jastrzebski et al., 2018; Wu et al., 2018; Wen et al., 2020; Mori and Ueda, 2020), the upstream cause for this generalization gap is a contested topic, approached from a variety of analytical perspectives (Goyal et al., 2017;
Wu et al., 2020;
Geiping et al., 2022; Lee et al., 2022).
Initial work in this field has generally focused on the learning rate to batch-size ratio (Keskar et al., 2017; Masters and Luschi, 2018;
Goyal et al., 2017; Mandt et al., 2017; He et al., 2019;
Li et al., 2019) or on recreating stochastic noise via mini-batching (Wu et al., 2020; Jastrzebski et al., 2018; Zhu et al., 2019; Mori and Ueda, 2020; Cheng et al., 2020; Simsekli et al., 2019; Xie et al., 2021), whereas recent works have pivoted focus on understanding how mini-batch SGD may implicitly regularize certain quantities that improve generalization (Geiping et al., 2022; Barrett and Dherin, 2020; Smith et al., 2021; Lee et al., 2022; Jastrzebski et al., 2020).
In this paper, we provide a careful empirical analysis of how these competing regularization theories compare to each other as assessed by how well the prescribed interventions, when applied in the large batch setting, recover SGD’s performance.
Additionally, we study their similarities and differences by analyzing the evolution of the regularized quantities over the course of training.
Our main contributions are the following:
By utilizing micro-batches (i.e. disjoint subsets of each mini-batch), we find that explicitly regularizing either the average micro-batch gradient norm (Geiping et al., 2022; Barrett and Dherin, 2020) or Fisher Information Matrix trace (Jastrzebski et al., 2020) (equivalent to the average gradient norm when labels are drawn from the predictive distribution, detailed in Section 2.2) in the large-batch regime fully recovers small-batch SGD generalization
[Figure]
performance, but using Jacobian-based regularization (Lee et al., 2022) fails to recover small-batch SGD performance (see Figure 1)
We show that the generalization performance is strongly correlated with how well the trajectory of the average micro-batch gradient norm during training mimics that of small-batch
SGD, but that this condition is not necessary for recovering performance in some scenarios.
The poor performance of Jacobian regularization, which enforces either uniform or fully random weighting on each class and example (see Section 2.3), highlights that the beneficial aspects of average
micro-batch gradient norm or Fisher trace regularization may come from the loss gradient’s ability to adaptively weight outputs on the per example and per class basis.
We demonstrate that the generalization benefits of both successful methods no longer hold when the micro-batch size is closer to the actual batch size. We too subsequently show that in this regime the average micro-batch gradient norm behavior of both previously successful methods differs significantly from the small-batch SGD case.
We highlight a high-level issue in modern empirical deep learning research: Experimental results that hold on CIFAR10 do not necessarily carry over to other datasets.
In particular, we focus on a technique called gradient grafting (Agarwal et al., 2020), which has been shown to improve generalization for adaptive gradient methods.
By looking at its behavior for normal SGD and GD, we show that gradient grafting recovers small-batch SGD generalization’s performance on CIFAR10 but fails in CIFAR100, arguing that research in this line should prioritize experiments on a larger and diverse range of benchmark datasets
P RIOR W ORK AND P RELIMINARIES
In neural network training, the choice of batch size (and learning rate) heavily influence generalization.
In particular, researchers have found that opting for small batch sizes (and large learning rates) improve a network’s ability to generalize (Keskar et al., 2017; Masters and Luschi, 2018; Goyal et al., 2017;
Mandt et al., 2017;
He et al., 2019; Li et al., 2019) .
Yet, explanations for this phenomenon have long been debated.
While some researchers have attributed the success of small-batch
SGD to gradient noise introduced by stochasticity and mini-batching (Wu et al., 2020; Jastrzebski et al., 2018; Zhu et al., 2019; Mori and Ueda, 2020; Cheng et al., 2020; Simsekli et al., 2019; Xie et al., 2021), others posit that small-batch SGD finds “flat minima” with low non-uniformity, which in turn boosts generalization (Keskar et al., 2017;
Wu et al., 2018; Simsekli et al., 2019).
Meanwhile, some works credit the implicit regularization of quantities such as loss gradient norm, the
Jacobian norm (i.e., the network output-to-weights gradient norm), and the Fisher Information Matrix (Geiping et al., 2022; Barrett and Dherin, 2020; Smith et al., 2021; Lee et al., 2022; Jastrzebski et al., 2020).
Recent works have shown that one can recover SGD generalization performance by training on a modified loss function that regularizes large loss gradients (Geiping et al., 2022; Barrett and Dherin, 2020; Smith et al., 2021).
While Smith et al.
and Barrett and Dherin (2020) expect that training on this modified loss function with large micro-batch sizes will be unable to recover SGD generalization performance, we empirically verify this is the case.
To our knowledge, we are the first to introduce the “micro-batch“ terminology to denote component disjoint sub-batches used in accumulated mini-batch SGD.
This choice was made to avoid overloading the term “mini-batch“ and thus clarify the work done by Smith et al.
(2021) and Barrett and Dherin (2020).
Note that here and for the rest of this paper, we use Large-Batch SGD as an approximation for full-batch GD due to the computational constraints of using the full training set on each update.
We emphasize that throughout this paper, micro-batches are not meant as any sort of “alternative” to mini-batches, as they are purely an implementation feature of gradient accumulationbased large-batch SGD.
We additionally leverage the work done by Agarwal et al.
(2020), who propose the idea of grafting as a meta-optimization algorithm, though in the paper their focus mostly rests on grafting adaptive optimization algorithms together, not plain mini-batch SGD.
As a whole, our paper is situated as a comparative analysis of multiple proposed regularization mechanisms (Geiping et al., 2022; Barrett and Dherin, 2020; Smith et al., 2021; Lee et al., 2022;
Jastrzebski et al., 2020) in a side-by-side empirical context, with additional ablations over how minor design choices may affect the efficacy of these proposed methods to close the generalization gap.
We now discuss various implicit and explicit regularization mechanisms in more depth.
Setup and Notation
We primarily consider the case of a softmax classifier f : R d → R C (where
C is the number of classes) parameterized by some deep neural network with parameters θ .
We let ℓ ( x , y ; θ ) denote the standard cross entropy loss for example x and label y , and let L B ( θ ) =|B| (cid:80) ( x ,y ) ∈B ℓ ( x , y ; θ ) denote the average loss over some batch B .
Note that throughout this paper, the terms “batch” and “mini-batch” are used interchangeably to refer to B .
2. A VERAGE M ICRO - BATCH G RADIENT N ORM R EGULARIZATION
As proposed by Smith et al. (2021), we attempt to understand the generalization behavior of minibatch SGD by how it implicitly regularizes the norm of the micro-batch gradient, ∥∇L M ( θ ) ∥ for some micro-batch M ⊆ B .
In large-batch SGD, we accomplish this through gradient accumulation (i.e. accumulating the gradients of many small-batches to generate the large-batch update), and thus can add an explicit regularizer (described in Geiping et al.
that penalizes the average microbatch norm.
Formally, for some large-batch B and component disjoint micro-batches M ⊆ B , let ∇ θ L M ( θ ) {+= 
+}
| M | (cid:80) ( x
,y ) ∈
M ∇ θ ℓ ( x , y ; θ ) .
The new loss function is:
[Equation]
While this quantity can be approximated through finite differences, it can also be optimized directly by backpropagation using modern deep learning packages, as we do in this paper.
Note that by definition, we can decompose the regularizer term into the product of the Jacobian of the network and the gradient of the loss with respect to network output.
Formally, for some network f with p parameters, if we let z = f ( x ; θ ) ∈ R C be the model output for some input x and denote its corresponding label as y , then:
[Equation]
Where ∇ θ z ∈ R p × C is the Jacobian of the network and the second term is the loss-output gradient.
We explicitly show this equivalence for the comparison made in section 2.3.
2. A VERAGE M ICRO - BATCH F ISHER T RACE R EGULARIZATION
One noticeable artifact of Equation 1 is its implicit reliance on the true labels y to calculate the regularizer penalty.
Jastrzebski et al.
(2020) shows that we can derive a similar quantity in the mini-batch SGD setting by penalizing the trace of the Fisher Information Matrix F , which is given by Tr ( F ) =
E x ∼X , ˆ y ∼ p
θ ( y | x ) [ ∥∇ θ ℓ ( x , ˆ y ; θ ) ∥ 2 ] , where p θ ( y | x ) is the predictive distribution of the model at the current iteration and X is the data distribution.
We thus extend their work to the accumulated large-batch regime and penalize an approximation of the average Fisher trace over micro-batches: if we let (cid:98) L M ( θ ) =
| M | (cid:80) x [Equation] p [Equation]
ℓ ( x
[Equation]
ˆ y ; θ ) , then our penalized loss is
[Equation] The only difference between the expressions in Equation 1 and Equation 3 is that the latter now uses labels sampled from the predictive distribution, rather than the true labels, to calculate the regularizer term.
As with Equation 1, we can directly backpropagate using this term in our loss equation.
Like, in equation 2, we can decompose the regularizer term as:
[Equation]
Where the second term is another loss-output gradient.
Jastrzebski et al.
(2020) observes that models with poor generalization typically show a large spike in the Fisher Trace during the early phases of training, which they coin as Catastrophic Fisher
Explosion .
In Figure 1, we show that this behavior also occurs when looking at the average MicroBatch gradient norm.
2. J ACOBIAN R EGULARIZATION
Given the decompositions shown in equations 2 and 4, it is unclear in either case whether the Jacobian term is the sole source of possible generalization benefit, or if the loss-output gradient is also needed.
To disentangle this effect, we borrow from Lee et al.
and use the average and unit
Jacobian regularization losses:
[Equation]
where u ∈ R C is randomly sampled from the unit hypersphere (i.e. ∥ u ∥ 2 = 1 ), and u is sampled once per micro-batch.
In words, the average Jacobian case penalizes the Jacobian with equal weighting on every class and every example, while the unit Jacobian case penalizes the Jacobian with different but random weighting on each class and example.
Note that the unit Jacobian penalty is an unbiased estimator of the Frobenius norm of the Jacobian ∥∇ θ z ∥ 2 F , which is an upper bound on its spectral norm ∥∇ θ z ∥ 22 (see Lee et al.
for a more detailed theoretical analysis).
E XPLICIT R EGULARIZATION A ND G RADIENT N ORM T RAJECTORY
These aforementioned explicit regularization mechanisms have previously been investigated in limited empirical settings.
To the best of our knowledge, Jastrzebski et al.
(2020) is the only work that has directly compared some of these regularization mechanisms, but only did so in the context of improving small-batch performance.
Like our work, Geiping et al.
(2022) is centered on the smallto-large batch generalization gap, but they do not focus solely on the explicit regularization they propose and do not include any analysis of the micro-batch gradient norm behavior during training.
In this work, we investigate (i) how these regularizers effect generalization for an array of benchmarks and (ii) how such performance may correlate with the evolution of the micro-batch gradient norm during training. [Table]
3. E XPERIMENTAL S ETUP
We first focus our experiments on the case of using a ResNet-
[Table]
et al
[Table]
), with standard initialization and batch normalization, on the CIFAR10, CIFAR100, Tiny-ImageNet [Table] and SVHN image classification benchmarks (Krizhevsky, 2009; Le and Yang, 2015; Netzer et al., 2011).
Additional experiments on different architectures are detailed in Appendix A.1.
Besides our small-batch ( B = 128 ) and large-batch ( B = 5120 )
SGD baselines, we also train the networks in the large-batch regime using (i) average Micro-batch Gradient Norm Regularization (GN); (ii)
average Micro-batch
Fisher Trace Regularization (FT); and (iii) average and unit Micro-batch Jacobian Regularizations (AJ and UJ).
Note that for all the regularized experiments, we use a component micro-batch size equal to the small-batch size (i.e. 128).
In order to compare the best possible performance within each experimental regime, we separately tune the optimal learning rate η and optimal regularization parameter λ independently for each regime.
Additional experimental details can be found in Appendix A.4.
3. R ESULTS
(i) Average micro-batch Gradient Norm and average Fisher trace Regularization recover SGD generalization
For CIFAR100, Tiny-ImageNet, and SVHN we find that we can fully recover small-batch SGD generalization performance by penalizing the average micro-batch Fisher trace and nearly recover performance by penalizing the average micro-batch gradient norm (with an optimally tuned regularization parameter λ , see Figure 1 and Table 1).
In CIFAR10, neither penalizing the gradient norm nor the Fisher trace completely recovers small-batch SGD performance, but rather come within ≈ 0 .
and ≈ 0 .
4% (respectively)
the small-batch SGD performance and significantly improves over large-batch SGD.
We additionally find that using the micro-batch gradient norm leads to slightly faster per-iteration convergence but less stable training (as noted by the tendency for the model to exhibit random drops in performance), while using the Fisher trace leads to slightly slower per-iteration convergence but much more stable training (see Figure 1).
This behavior may be due to the Fisher trace’s ability to more reliably mimic the small-batch SGD micro-batch gradient norm behavior throughout training, whereas penalizing the gradient norm effectively curbs the initial explosion but collapses to much smaller norm values as training progresses
(ii) Average and Unit Jacobian regularizations do not recover SGD generalization
Observe in Table 1 that we are unable to match SGD generalization performance with either Jacobian regularization.
In Section 2 we showed that each regularization method can be viewed as penalizing the norm of the Jacobian matrix-vector product with some C -dimensional vector.
Crucially, both the gradient norm and Fisher trace regularizers use some form of loss-output gradient, which is datadependent and has no constraint on the weighting of each class, while both Jacobian regularizers use data-independent and comparatively simpler vectors.
Given the noticeable difference in generalization performance between the regularization methods that weight the Jacobian with the loss-output gradient and those that do not, we indicate that the loss-output gradient may be crucial to either applying the beneficial regularization effect itself and/or stabilizing the training procedure.
S HORTCOMINGS AND E XTENSIONS OF G RADIENT N ORM R EGULARIZATION 4. G ENERALIZATION F AILURE AT L ARGE M ICRO -B ATCH S IZES
In both successful regularization regimes, namely the average micro-batch gradient norm and average fisher trace regularizers, there is an implicit hyperparameter: the size of the micro-batch used to calculate the regularization term.
Note that this hyperparameter is a practical artifact of modern
GPU memory limits, as efficiently calculating higher-order derivatives for large batch sizes is not feasible in standard autodifferentiation packages.
Consequently, gradient accumulation (and the use of the average micro-batch regularizer, rather than taking the norm over the entire batch) must be used on most standard GPUs (more detailed hadware specifications can be found in appendix A.4).
This restriction, however, may actually be beneficial, as Geiping et al.
; Barrett and Dherin (2020); Smith et al.
(2021) have noted that they expect the benefits of gradient norm regularizers to break down when the micro-batch size becomes too large.
To test this hypothesis, we return to the
ResNet-18 in CIFAR100 and Tiny-ImageNet settings and increase the micro-batch size to as large as we could reasonably fit on a single GPU at | M | = 2560 in both the gradient norm and Fisher trace experiments.
Additionally, we run experiments using a VGG11 (Simonyan and Zisserman, 2015) on CIFAR10, interpolating the micro-batch size from the small to large regimes.
In both settings, we separately tune the learning rate η and regularization coefficient λ in each experiment to find the best possible generalization performance in the large micro-batch regimes.
[Equation]
[Table]
[Figure]
[Figure]
[Figure]
[Figure]
23 ( ± [Figure]
[Figure]
( ±
[Figure]
[Figure]
[Figure]
( ±
[Figure]
[Table]
Results We successfully show that such hypotheses mentioned in Geiping et al.
; Smith et al.
(2021); Barrett and Dherin (2020) hold true: as the micro-batch size approaches the mini-batch size, both regularization mechanisms lose the ability to recover small-batch SGD performance (see Tables 2 and 3).
Additionally, we note that using large micro-batch sizes no longer effectively mimics the average micro-batch gradient norm behavior of small-batch SGD, thus supporting our claim that matching this quantity throughout training is of key importance to recovering generalization performance (Figure 2) [Figure]
4. S AMPLE M ICRO - BATCH G RADIENT N ORM R EGULARIZATION
One potential practical drawback of these gradient-based regularization terms is the relatively high computation cost needed to calculate the second-order gradients for every component micro-batch.
Instead of penalizing the average micro-batch gradient norm [Figure] we can penalize one micro-batch gradient norm.
For some large batch B and fixed sample micro-batch S from batch B [Figure] we define the modified loss function
[Equation]
[Figure] [Equation]
[Figure]
Results
In Figure 3, we plot the final test accuracy (left column) and the average gradient norm (right column) as a function of λ .
We observe that both a larger λ and a smaller micro-batch size | S | boost test accuracy.
Furthermore, we find that with the “optimal” λ and micro-batch size | S | , the final test accuracy for sample micro-batch gradient norm regularization is close to (and sometimes better) than the final test accuracy for SGD.
Just as we observed with the average Micro-batch Gradient
Norm regularization, generalization benefits diminish as the sample micro-batch size approaches the mini-batch size.
I S MIMICKING SGD G RADIENT N ORM B EHAVIOR NECESSARY FOR GENERALIZATION ?
As seen in Figure 1, the trajectory of the average micro-batch gradient norm during training, and its similarity to that of small-batch SGD especially in the early stages of training, is strongly correlated with generalization performance.
Furthermore, we have observed that models with poor generalization performance tend to exhibit the characteristic “explosion” during the early phase of training and quickly plummet to average micro-batch gradient norm values much smaller than seen in small-batch SGD.
That being said, it is not immediately clear whether recreating the micro-batch norm trajectory of small-batch SGD is necessary for ensuring good generalization performance (i.e. whether good generalization directly implies gradient norm behavior similar to SGD).
To test this hypothesis, we empirically validate an orthogonal vein of optimization methods that do not explicitly regularize the micro-batch gradient norm during training for their ability to close the small-to-large batch generalization gap, and whether they too mimic the average micro-batch norm trajectory of small-batch SGD.
5. E XTERNAL AND I TERATIVE G RAFTING AND N ORMALIZED G RADIENT D ESCENT
Inspired by the work of Agarwal et al.
(2020), we proposed to use gradient grafting in order to control the loss gradient norm behavior during training.
Formally, for any two different optimization algorithms M , D , the grafted updated rule is arbitrarily:
[Equation]
In this sense, M controls the magnitude of the update step and D controls the direction .
We first propose Iterative Grafting , wherein M ( θ k ) =
η ∇L M ( θ k ) and D ( θ k ) =
∇L B ( θ k ) , where M ∈
B is sampled uniformly from the component micro-batches at every update.
In words, at every update step we take the large batch gradient, normalize it, and then rescale the gradient by the norm of one of the component micro-batch gradients.
Additionally, we propose External Grafting , where M ( θ k ) =
η ∇L M ( θ k ′ ) and D ( θ k ) =
∇L B ( θ k ) .
Here, we use ∇L B ( θ k ′ ) to denote the gradient norm at step k from a separate smallbatch SGD training run .
We propose this experiment to make a comparison with the Iterative Grafting case, since here the implicit step length schedule is independent of the current run, while with
Iterative grafting the schedule depends upon the current training dynamics.
Aside from grafting algorithms, which define the implicit step length schedule at every step, we also consider the situation where the step length is fixed throughout training through normalized gradient descent (NGD) (Hazan et al., 2015), wherein M ( θ k ) =
η and D ( θ k ) =
η ∇L B ( θ k ) .
[Table]
[Figure]
Results We find that both forms of grafting and NGD can recover the generalization performance of SGD in some model-dataset combinations (see Table 4).
Namely, though grafting / NGD seems to work quite well in CIFAR10, no amount of hyperparameter tuning was able to recover the SGD performance for either model in CIFAR100.
That being said, in the CIFAR10 case we see (in Figure 4) that the grafting experiments (and NGD, not pictured) do not replicate the same average minibatch gradient norm behavior of small-batch SGD despite sometimes replicating its performance.
This thus gives us solid empirical evidence that while controlling the average mini-batch gradient norm behavior through explicit regularization may aide generalization, it is not the only mechanism in which large-batch SGD can recover performance.
5. W IDER I MPLICATIONS
The stark disparity in performance between the CIFAR10 and CIFAR100 benchmarks are of key importance.
These differences may be explained by the much larger disparity between the mid-stage average micro-batch gradient norm behavior in the CIFAR100 case than in the CIFAR10 case (see
Figure 4).
This situation highlights a possible cultural issue within the deep learning community: there is a concerning trend of papers in the deep learning field that cite desired performance on
CIFAR10 , and no harder datasets, as empirical justification for any posed theoretical results (Orvieto et al., 2022; Geiping et al., 2022; Agarwal et al., 2020; Smith et al., 2021; Barrett and Dherin, 2020;
Cheng et al., 2020).
Given the continued advancement of state-of-the-art deep learning models, we argue that it is imperative that baselines like CIFAR100 and ImageNet are adopted as the main standard for empirical verification, so that possibly non-generalizable results (as the grafting / NGD results would have been had we stopped at CIFAR10) do not fall through the cracks in the larger community (see Appendix A.3 for more information).
D ISCUSSION & C ONCLUSION
In this paper, we provide a holistic account of how the proposed regularization mechanisms (Geiping et al., 2022; Barrett and Dherin, 2020; Smith et al., 2021; Lee et al., 2022; Jastrzebski et al., 2020) compare to each other in performance and gradient norm trajectory, and additionally show the limitations of this analytical paradigm for explaining the root cause of generalization.
Our results with regards to the relative poor performance of the Jacobian-based regularizations somewhat conflict with the results of Lee et al.
(2022), which shows positive results on using the unit Jacobian regularization with respect to improving performance within the same batch-size regime .
We attribute this difference to the fact that Lee et al.
(2022) is not concerned with cases where the small-to-large batch generalization gap exists, which is our main focus.
In light of this prior work, more research should be done to disentangle the exact effect that implicitly regularizing the loss-output gradient has on generalization performance.
Next, given the success of average micro-batch gradient norm and average micro-batch Fisher trace regularization (especially with small micro-batches), future work should leverage these regularization mechanisms to investigate the possibility of ameliorating generalization, while improving time efficiency, by taking advantage of high resource, parallelizable settings.
We also show that experimental findings on CIFAR10 may no longer hold in CIFAR100, which sheds light on a wider implication for the research community.
Namely, we urge researchers to adapt the practice of evaluating empirical hypotheses on a more widespread, complex set of benchmarks.
We acknowledge that performance in each experiment could possibly be improved by progressively finer hyperparameter tuning, though we are confident that our core results would continue to hold in such situations given the extensive hyperparameter searches performed for each experiment.
As a whole, the present research helps to shed light on the mechanisms behind SGD’s generalization properties through implicit regularization, and offers robust fixes to the generalization issue at high batch-sizes.
R EPRODUCIBLITY S TATEMENT
The source code for reproducing the work presented here, including all hyperparameters and random seeds, is available at https://github.com/anon2023ICLR/imp-regularizers .
Additional experimental details are available in Appendix A.4.
A A PPENDIX A. A DDITIONAL R EGULARIZATION E XPERIMENTS
Aside from the main results using a ResNet-18, we additionally ran the regularization experiments with a VGG11 (Simonyan and Zisserman, 2015) without batch normalization on CIFAR10.
Results are shown below:
[Table]
Consistent with our earlier observations (see Section 3), we find that average micro-batch gradient norm and average Fisher trace regularization nearly recover SGD generalization performance, whereas average Jacobian regularization does not.
A. S AMPLE M ICRO - BATCH G RADIENT
[Figure]
[Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
We see that switching from using the average micro-batch gradient norm to using a single sample micro-batch gradient norm as a regularizer does not impact final generalization performance.
However, we do see in the CIFAR100 case that while using the sample-based regularizer is faster in terms of per iteration wall-clock time, the average-based regularizer converges to an optimal performance in considerably fewer gradient update steps (Figure 5).
A. L IMITATIONS OF A NTICORRELATED P ERTURBED G RADIENT D ESCENT
Orvieto et al.
(2022) proposes a method for improving generalization by injecting spherical Gaussian noise (with variance σ 2 as a hyperparameter) at each gradient update step that is anticorrelated between concurrent time steps, which they term Anti-PGD .
They empirically show that on CIFAR10, training a ResNet18 in the large-batch regime with Anti-PGD and then shutting off the noise (to allow for convergence) allows them to beat small-batch SGD generalization performance.
However, when we extended their methodology to the CIFAR100 regime (while removing possible confounding factors such as momentum), no hyperparameter combination in the large-batch regime was able to recover SGD generalization performance.
This thus represents just one example of the possible endemic issue described in 5.2.
Hyperparameter combinations and final test accuracy are show below.
[Table]
. E XPERIMENTAL S ETUP D ETAILS
All experiments run for the present paper were performed using the Pytorch deep learning API, and source code can be found here: https://github.com/anon2023ICLR/ imp-regularizers .
ResNet-18 For all ResNet-18 experiments, we use the standard He initialization (He et al., 2015), and the default Pytorch batch normalization initialization.
Additionally, we use the standard data augmentations for CIFAR10 and CIFAR100; that is, random cropping, horizontal flipping, and whitening.
For SVHN, we performed only whitening on the dataset.
For Tiny-ImageNet, no data augmentations were made aside from rescaling the input images (which are 64 × 64 ) to be 32 × 32 .
Additionally, for Tiny-ImageNet the sample regularization penalties were used rather than the normal average regularization penalties given compute constraints (see Section 4.2 for the documented similarities between the sample and average regularizations).
All experiments were run for 50000 update iterations.
In this case, all models are trained well past the point of reaching 100% training accuracy.
No weight decay or momentum was used in any of the experiments.
We use a large-batch size of 5120 for all experiments, and thus have 40 micro-batches of 128 examples each for the regularization experiments.
We calculate the penalty term at every update step, which is different from the procedure in Jastrzebski et al.
(2020), which recalculates the penalty term only every 10 update steps.
For the external grafting experiments, we use the gradient norm data from a separate run of small-batch SGD with batch size equal to the same micro-batch size used for iterative grafting (i.e. 128).
All experiments were run on a single RTX A6000 NVidia GPU.
VGG
The set-up for the VGG16 experiments are identical to the ResNet-18 experiments, including the usage of batch normalization within the architecture.
VGG11 without batch-normalization For all VGG-11 experiments, we train the network with a fixed learning rate (and no momentum) until we reach 99% train accuracy.
Note that we do not use any form of data augmentations.
We use a small batch size of 100 and a large batch size of 5000.